`nemo_rl.algorithms.rollout_checkpointing.trace`#

A per-rollout event log for checking what a stop and restore cost.

With rollout_checkpointing.trace, the controller appends one JSON line per event to rollout_trace.jsonl in the log directory: every rollout attempt’s dispatch, seal, refusal, abandonment (with its reason) and resume, each prompt group’s reservation, admission, seal and discard, each training step and weight sync, and each checkpoint published or restored. Every line carries the wall-clock time t and a run ID that is new for each process, so runs that resume into the same log directory stay apart.

Events are queued as they happen and a background thread appends them, so writing never blocks the controller’s event loop; a run that is killed loses only the events not yet written, typically the last few milliseconds’ worth. Rows match across runs by prompt_id (the dataset index) and generation_index; group and attempt IDs are new in every run.

With token capture, each published group counts its trained tokens by the weight version that generated them. vLLM keeps running calls across a refit, so RefitCuts records how far each running call got when a refit paused it and splits the call there. batch_age_by_step turns those counts into each step’s batch age: the token-weighted number of weight updates between generating a token and training on it. The controller logs the same number as batch_age every step; run this module on a trace file to print it per step.

Module Contents#

Classes#

`RolloutTrace`	Append rollout events to a JSONL file off the caller’s thread, or drop them.
`RefitCuts`	Where each running captured call switched weights, until its group publishes.

Functions#

`read_events`	Yield the events of a trace file in the order they were written.
`tokens_by_version`	Each published group’s trained tokens, by the weight version that made them.
`batch_age_group_metrics`	A published group’s share of its step’s batch age, as two sums.
`pop_batch_age`	Remove the step’s batch-age sums and return its batch age, if it has tokens.
`batch_age_by_step`	Each step’s token-weighted age: policy updates from generating to training.
`main`

Data#

`ROLLOUT_TRACE_FILENAME`
`BATCH_AGE_TOKENS_METRIC`
`BATCH_AGE_VERSION_SUM_METRIC`

API#

nemo_rl.algorithms.rollout_checkpointing.trace.ROLLOUT_TRACE_FILENAME#: ‘rollout_trace.jsonl’

nemo_rl.algorithms.rollout_checkpointing.trace.BATCH_AGE_TOKENS_METRIC#: ‘batch_age/trained_tokens’

nemo_rl.algorithms.rollout_checkpointing.trace.BATCH_AGE_VERSION_SUM_METRIC#: ‘batch_age/token_version_sum’

class nemo_rl.algorithms.rollout_checkpointing.trace.RolloutTrace(path: Optional[pathlib.Path])[source]#

Append rollout events to a JSONL file off the caller’s thread, or drop them.

Initialization

Open path for appending; None makes every call a no-op.

Parameters:: path – The JSONL file to append to, or None to disable tracing.

property enabled: bool#

emit(event: str, **fields: Any) → None[source]#: Queue one event, stamped with the current wall-clock time.

flush() → None[source]#: Wait until every queued event has been written.

close() → None[source]#: Write every queued event, then stop the writer and close the file.

_write_lines() → None[source]#: Append queued events in batches until close queues None.

nemo_rl.algorithms.rollout_checkpointing.trace.read_events( path: pathlib.Path, ) → collections.abc.Iterator[dict[str, Any]][source]#: Yield the events of a trace file in the order they were written.

class nemo_rl.algorithms.rollout_checkpointing.trace.RefitCuts[source]#

Where each running captured call switched weights, until its group publishes.

vLLM keeps running calls across a refit, so a call’s tokens after the refit come from the new weights.

Initialization

record( replaced_version: int, new_version: int, calls: collections.abc.Iterable[collections.abc.Sequence[Any]], ) → None[source]#: Record how many tokens each call had when a refit paused it.

forget_before(oldest_trainable_version: int) → None[source]#

Drop calls whose group cannot be trained any more.

A call began on or before the version its first refit replaced, and its group no later than the call, so the group is too old to train when that version is.

tokens_by_version( call_versions: collections.abc.Iterable[tuple[str, int, int]], ) → dict[int, int][source]#

A published group’s trained tokens by the weight version that made them.

Parameters:: call_versions – Each trained call’s ID, the version it began on and the tokens it generated.
Returns:: Token counts by weight version. A call’s tokens count under the version it began on until a refit paused it, and under the version that refit installed from there on.

nemo_rl.algorithms.rollout_checkpointing.trace.tokens_by_version( events: collections.abc.Iterable[dict[str, Any]], ) → dict[str, dict[int, int]][source]#: Each published group’s trained tokens, by the weight version that made them.

nemo_rl.algorithms.rollout_checkpointing.trace.batch_age_group_metrics(counts: dict[int, int]) → dict[str, int][source]#: A published group’s share of its step’s batch age, as two sums.

nemo_rl.algorithms.rollout_checkpointing.trace.pop_batch_age( per_group_metrics: dict[str, list[Any]], trained_version: int, ) → Optional[float][source]#

Remove the step’s batch-age sums and return its batch age, if it has tokens.

Parameters:

per_group_metrics – The step’s rollout metrics, one value per group.
trained_version – The policy version the step updates.

nemo_rl.algorithms.rollout_checkpointing.trace.batch_age_by_step( events: collections.abc.Iterable[dict[str, Any]], ) → dict[int, float][source]#

Each step’s token-weighted age: policy updates from generating to training.

A step a resumed run trained again keeps the later run’s value.

nemo_rl.algorithms.rollout_checkpointing.trace.main() → None[source]#

nemo_rl.algorithms.rollout_checkpointing.trace#