nemo_gym.core.model_call_sampling

Sampling values that a training framework sets for a rollout’s model calls, and how they reach each call.

A training framework that serves the policy samples the calls of the agents it trains with its own values: its training values, or its validation values for a validation rollout. It sets them on each rollout’s row under _ng_model_sampling, next to _ng_rollout_id. Gym carries them into the URL of every model call the rollout’s agent makes, as /ng-sampling/<values> after any other prefix, and the Model Server writes them into the request. Gym never interprets them.

In such a run, a call from a caller the framework does not train, such as a verifier’s LLM judge, carries /ng-sampling/null: it keeps the sampling its caller chose, and its own Model Server’s sampling_overrides. A run whose rows carry no values, such as Gym-only evaluation, adds no segment.

Module Contents

Classes

Name	Description
`ModelCallSampling`	Sampling values for the model calls of an agent that a training framework trains.
`ModelCallSamplingConfig`	The global `model_call_sampling` block. A training framework sets it when its rows carry values.

Functions

Name	Description
`_unique_keys`	-
`current_model_sampling_segment`	Return the segment value for this request’s model calls, or `None` when they carry no segment.
`decode_model_sampling`	Decode a segment value: `null` to `None`, anything else to values.
`encode_model_sampling`	Encode values as the segment’s unpadded base64url JSON, and `None` as `null`.
`model_call_sampling_config`	Return the run’s `model_call_sampling` block, or `None` when the run sets none.
`model_sampling_context`	Give the model calls made inside the block the segment value `segment`.
`model_sampling_from_run_body`	Return the values a run request or row carries under `_ng_model_sampling`, or `None`.
`model_sampling_path_segment`	Return `/ng-sampling/<segment>`, or `""` for no segment.
`with_model_sampling_segment`	Insert `/ng-sampling/<segment>` into a Model Server generation path, where its correlation prefix puts it.

Data

MODEL_SAMPLING_HEADER

MODEL_SAMPLING_NULL

_ENCODED

_GENERATION_API_PATHS

_MAX_ENCODED_LENGTH

_MODEL_SAMPLING_SEGMENT

API

class nemo_gym.core.model_call_sampling.ModelCallSampling()

Bases: BaseModel

Sampling values for the model calls of an agent that a training framework trains.

The Model Server sends the smaller of max_output_tokens and the call’s own limit, so a harness that reserves fewer tokens, near the end of its context or for a summary, keeps its limit. An unset field leaves the call’s own value.

max_output_tokens

Optional[int] = Field(default=None, gt=0, le=(2 ** 53 - 1))

model_config

temperature

Optional[float]

top_p

Optional[float]

class nemo_gym.core.model_call_sampling.ModelCallSamplingConfig()

Bases: BaseModel

The global model_call_sampling block. A training framework sets it when its rows carry values.

Its presence marks the run: resources servers then send null with the model calls they make through Gym’s server client while handling a request, so their judges keep their own sampling.

model_config

= ConfigDict(extra='forbid')

required_model_servers

list[str] = Field(default_factory=list)

nemo_gym.core.model_call_sampling._unique_keys(
    pairs: list[tuple[str, typing.Any]]
) -> dict[str, typing.Any]

nemo_gym.core.model_call_sampling.current_model_sampling_segment() -> typing.Optional[str]

Return the segment value for this request’s model calls, or None when they carry no segment.

nemo_gym.core.model_call_sampling.decode_model_sampling(
    encoded: str
) -> typing.Optional[nemo_gym.core.model_call_sampling.ModelCallSampling]

Decode a segment value: null to None, anything else to values.

Raises ValueError for a malformed or ambiguous value, so a call never runs with values it did not send.

nemo_gym.core.model_call_sampling.encode_model_sampling(
    sampling: typing.Optional[nemo_gym.core.model_call_sampling.ModelCallSampling]
) -> str

Encode values as the segment’s unpadded base64url JSON, and None as null.

nemo_gym.core.model_call_sampling.model_call_sampling_config(
    global_config_dict: typing.Optional[collections.abc.Mapping[str, typing.Any]]
) -> typing.Optional[nemo_gym.core.model_call_sampling.ModelCallSamplingConfig]

Return the run’s model_call_sampling block, or None when the run sets none.

Raises ValueError when required_model_servers names a server that is not a Model Server of the run: a misspelled name would otherwise leave the trained agent’s Model Server unenforced, without an error.

nemo_gym.core.model_call_sampling.model_sampling_context(
    segment: typing.Optional[str]
) -> collections.abc.Iterator[None]

Give the model calls made inside the block the segment value segment.

nemo_gym.core.model_call_sampling.model_sampling_from_run_body(
    body: pydantic.BaseModel | collections.abc.Mapping[str, typing.Any] | None
) -> typing.Optional[nemo_gym.core.model_call_sampling.ModelCallSampling]

Return the values a run request or row carries under _ng_model_sampling, or None.

nemo_gym.core.model_call_sampling.model_sampling_path_segment(
    segment: typing.Optional[str]
) -> str

Return /ng-sampling/<segment>, or "" for no segment.

nemo_gym.core.model_call_sampling.with_model_sampling_segment(
    url_path: str,
    segment: typing.Optional[str]
) -> str

Insert /ng-sampling/<segment> into a Model Server generation path, where its correlation prefix puts it.

Other paths, and paths that already carry a segment or a malformed prefix, are returned unchanged.

nemo_gym.core.model_call_sampling.MODEL_SAMPLING_HEADER = 'x-ng-sampling'

nemo_gym.core.model_call_sampling.MODEL_SAMPLING_NULL = 'null'

nemo_gym.core.model_call_sampling._ENCODED = re.compile('[A-Za-z0-9_-]+')

nemo_gym.core.model_call_sampling._GENERATION_API_PATHS = frozenset({'/v1/chat/completions', '/v1/responses', '/v1/messages'})

nemo_gym.core.model_call_sampling._MAX_ENCODED_LENGTH = 1024

nemo_gym.core.model_call_sampling._MODEL_SAMPLING_SEGMENT: ContextVar[Optional[str]] = ContextVar('nemo_gym_model_sampling_segment', default=None)