Skip to navigation

nemo_gym.rollouts.model_calls

View as Markdown

Read saved model exchanges, resolve ownership, and attach evidence to rollouts.

Module Contents

Functions

NameDescription
_read_capture_records-
_store_for_rollout-
aggregate_model_call_metricsAggregate model-call metrics for one rollout id.
await_model_call_capture_quiescenceWait until no claimed capture write for this rollout is outstanding.
clear_model_call_captures_for_rolloutsRemove stale per-rollout capture files for these records before dispatch.
merge_model_call_capture_into_recordAttach captured model-call observability data to a rollout record in place.
model_call_capture_dirs_from_configReturn the single run-wide capture directory when capture is enabled.
observability_enabled_from_configReturn the run-wide observability_enabled flag directly, without going through capture dirs.
read_available_model_call_recordsRead valid call records and count damaged records.
read_model_call_recordsRead captured exchanges in durable append order.

Data

CAPTURE_DRAIN_TIMEOUT_S

_CAPTURE_DRAIN_POLL_S

__getattr__

logger

API

nemo_gym.rollouts.model_calls._read_capture_records(
rollout_id: str
nemo_gym.rollouts.model_calls._store_for_rollout(
rollout_id: str,
capture_dirs: list[pathlib.Path]
nemo_gym.rollouts.model_calls.aggregate_model_call_metrics(
rollout_id: str
) -> dict[str, typing.Any]

Aggregate model-call metrics for one rollout id.

nemo_gym.rollouts.model_calls.await_model_call_capture_quiescence(
rollout_id: str,
capture_dirs: list[pathlib.Path],
timeout_s: float = CAPTURE_DRAIN_TIMEOUT_S,
poll_s: float = _CAPTURE_DRAIN_POLL_S
) -> bool
async

Wait until no claimed capture write for this rollout is outstanding.

A cancelled model call finishes its capture write detached (see the shield in _CaptureMiddleware), so the row can land after the rollout record is snapshotted and be missing from ng_model_call_capture while present on disk. Returns False on timeout, which the merge turns into a model_call_capture_pending gap rather than a silent omission.

nemo_gym.rollouts.model_calls.clear_model_call_captures_for_rollouts(
records: list[typing.Any],
capture_dirs: list[pathlib.Path]
) -> None

Remove stale per-rollout capture files for these records before dispatch.

Capture files are keyed by a deterministic rollout id (task-rollout-attempt), so without this a fresh run or a kill-shaped retry would append onto the previous attempt’s capture for the same id. The caller passes only rows about to be dispatched, after assigning any retry suffix.

nemo_gym.rollouts.model_calls.merge_model_call_capture_into_record(
record: dict[str, typing.Any],
capture_dirs: list[pathlib.Path],
include_payloads: bool = False
) -> dict[str, typing.Any]

Attach captured model-call observability data to a rollout record in place.

Keyed by the rollout id derived from the record’s task/rollout/attempt indices, so the attached shape is identical for every agent harness. Adds ng_model_call_capture = {rollout_id, metrics, calls} where calls are derived observability records. Raw request and response payloads remain in the capture store and are omitted from the attachment unless include_payloads is true. Capture/read/join failures are attached as gaps. The harness output and reward are not modified.

nemo_gym.rollouts.model_calls.model_call_capture_dirs_from_config(
global_config_dict: typing.Any
) -> list[pathlib.Path]

Return the single run-wide capture directory when capture is enabled.

nemo_gym.rollouts.model_calls.observability_enabled_from_config(
global_config_dict: typing.Any
) -> bool

Return the run-wide observability_enabled flag directly, without going through capture dirs.

nemo_gym.rollouts.model_calls.read_available_model_call_records(
rollout_id: str

Read valid call records and count damaged records.

nemo_gym.rollouts.model_calls.read_model_call_records(
rollout_id: str

Read captured exchanges in durable append order.

nemo_gym.rollouts.model_calls.CAPTURE_DRAIN_TIMEOUT_S = 10.0
nemo_gym.rollouts.model_calls._CAPTURE_DRAIN_POLL_S = 0.05
nemo_gym.rollouts.model_calls.__getattr__ = moved_names_getattr(__name__)
nemo_gym.rollouts.model_calls.logger = logging.getLogger(__name__)