nemo_gym.orchestration.api

Module Contents

Classes

Name	Description
`BaseComputeConfig`	-
`BaseModelServiceConfig`	Base for services that serve a model and can be wired as the policy model.
`BaseServiceConfig`	-
`BenchmarkRunConfig`	-
`DriverConfig`	-
`GymInstallConfig`	-
`HealthCheckConfig`	-
`JobConfig`	-
`KubernetesComputeConfig`	-
`NodePool`	-
`OtelConfig`	An OpenTelemetry collector beside every benchmark job: scrapes each model service’s
`RayServiceConfig`	-
`ResumeConfig`	-
`SlurmComputeConfig`	-
`SubmitConfig`	-
`VllmNodeServiceConfig`	-
`VllmPDServiceConfig`	Prefill/decode disaggregated vLLM: two tiers behind a vllm-router.
`VllmPDTierConfig`	One tier of a `vllm_pd` service as deployed. Built by the service, never written by a user.
`VllmServiceConfig`	-
`_HostEnvRefs`	Remembers which `env` values came from `host:` so a dump can write the reference back.
`_StrictModel`	-

Functions

Name	Description
`_shown_policy_value`	A policy run value as a conflict message shows it, with no credential: an API key is never printed, and a
`_without_url_credentials`	`value` with the user-info password and every query value of each URL in it masked, also inside a list.
`effective_ray_serve`	Whether the Ray Serve gateway manages this service’s instances/routing instead of vLLM’s own DP.
`resolve_env_dict`	Resolve `lit:`/`host:`/`runtime:` prefixes on `env` values. Every value must use one

Data

_DECODE_NIXL_PORT_OFFSET

_ENV_VAR_NAME_RE

_PD_SERVICE_ONLY_FIELDS

_URL_QUERY_VALUE

API

class nemo_gym.orchestration.api.BaseComputeConfig()

Bases: _StrictModel

class nemo_gym.orchestration.api.BaseModelServiceConfig()

Bases: BaseServiceConfig

Base for services that serve a model and can be wired as the policy model.

model

str

port

int = 8000

served_model_name

str | None = None

class nemo_gym.orchestration.api.BaseServiceConfig()

Bases: _HostEnvRefs, _StrictModel

container

str

env

dict[str, str] = {}

health_check

HealthCheckConfig | None = None

mounts

list[str] = []

node_pool

str | None = None

placement

str | None = None

pre_command

str = ''

nemo_gym.orchestration.api.BaseServiceConfig._resolve_env_prefixes(
    v: dict[str, str],
    info: pydantic.ValidationInfo
) -> dict[str, str]

classmethod

class nemo_gym.orchestration.api.BenchmarkRunConfig()

Bases: _StrictModel

command

str | None = None

export

dict[str, dict[str, Any] | None] = {}

prepare

dict[str, Any] = {}

resumable

bool | ResumeConfig = False

resume_config

ResumeConfig | None

run

dict[str, Any] = {}

nemo_gym.orchestration.api.BenchmarkRunConfig._validate_command(
    info: pydantic.ValidationInfo
) -> nemo_gym.orchestration.api.BenchmarkRunConfig

class nemo_gym.orchestration.api.DriverConfig()

Bases: _HostEnvRefs, _StrictModel

benchmarks

dict[str, BenchmarkRunConfig] = Field(min_length=1)

container

str = 'python:3.12'

env

dict[str, str] = {}

gym_install

GymInstallConfig | None = None

mounts

list[str] = []

policy_host

str = 'localhost'

policy_model

str | None = None

policy_model_type

str = 'openai_model'

nemo_gym.orchestration.api.DriverConfig._resolve_env_prefixes(
    v: dict[str, str],
    info: pydantic.ValidationInfo
) -> dict[str, str]

classmethod

class nemo_gym.orchestration.api.GymInstallConfig()

Bases: _StrictModel

ref

str

repo

str = 'https://github.com/NVIDIA-NeMo/gym'

class nemo_gym.orchestration.api.HealthCheckConfig()

Bases: _StrictModel

path

str = '/health'

port

int | None = None

timeout_seconds

int = 60

class nemo_gym.orchestration.api.JobConfig()

Bases: _StrictModel

export

dict[str, dict[str, Any] | None] = {}

output_path

str

class nemo_gym.orchestration.api.KubernetesComputeConfig()

Bases: BaseComputeConfig

active_deadline_seconds

int | None = None

context

str | None = None

extra_args

dict[str, str] = {}

gpus_per_node

int | None = None

memory_per_gpu

str = '32Gi'

namespace

str = 'default'

node_selector

dict[str, str] = {}

pvc_name

str

service_account

str | None = None

ttl_seconds_after_finished

int = 60 * 60 * 24

type

Literal['kubernetes']

class nemo_gym.orchestration.api.NodePool()

Bases: _StrictModel

extra_args

dict[str, str] = {}

gpus_per_node

int | None = None

nodes

int = 1

ntasks_per_node

int = 1

partition

str

class nemo_gym.orchestration.api.OtelConfig()

Bases: _StrictModel

An OpenTelemetry collector beside every benchmark job: scrapes each model service’s Prometheus /metrics, receives OTLP from the job’s own processes on :4317/:4318, and ships both to an OTLP/HTTP backend while keeping a copy under <job dir>/otel/. On by default, so a run is observable unless it opts out. Remote export needs endpoint, service_name and a token from the submitting shell; without the token, only the local copies are written.

binary

str = 'otelcol-contrib'

component

str = 'gym-vllm'

container

str | None = None

enabled

bool = True

endpoint

str | None = None

gpu_metrics_port

int | None = 9400

gym_logs

bool = True

gym_span_groups

str = 'default,verify'

gym_telemetry

bool = True

health_check_timeout_seconds

int = 300

node_metrics_port

int | None = 9100

remote_export

bool = True

scrape_interval_seconds

int = 15

service_name

str | None = None

token_env

str = 'OTEL_TOKEN'

vllm_traces

bool = False

nemo_gym.orchestration.api.OtelConfig._validate_positive(
    v: int
) -> int

classmethod

nemo_gym.orchestration.api.OtelConfig._validate_token_env(
    v: str
) -> str

classmethod

class nemo_gym.orchestration.api.RayServiceConfig()

Bases: BaseServiceConfig

extra_args

str = ''

node_pools

list[str] = []

num_cpus

int | None = None

num_gpus

int | None = None

port

int = 6379

resources

dict[str, dict[str, float]] = {}

type

Literal['ray']

nemo_gym.orchestration.api.RayServiceConfig._validate_shape() -> nemo_gym.orchestration.api.RayServiceConfig

class nemo_gym.orchestration.api.ResumeConfig()

Bases: _StrictModel

max_retries

int = Field(default=3, ge=0)

max_walltime

str | None = None

recovery_outcome_policy

dict[str, Any] | None = None

nemo_gym.orchestration.api.ResumeConfig._omit_unset_recovery_outcome_policy(
    handler: pydantic.SerializerFunctionWrapHandler
) -> dict[str, typing.Any]

nemo_gym.orchestration.api.ResumeConfig._validate_recovery_outcome_policy(
    value: dict[str, typing.Any] | None
) -> dict[str, typing.Any] | None

classmethod

class nemo_gym.orchestration.api.SlurmComputeConfig()

Bases: BaseComputeConfig

account

str

extra_args

dict[str, str] = {}

forward_env

list[str] = []

hostname

str | None = None

node_pools

dict[str, NodePool] = {}

type

Literal['slurm']

walltime

str | None = None

nemo_gym.orchestration.api.SlurmComputeConfig._validate_forward_env(
    names: list[str]
) -> list[str]

classmethod

class nemo_gym.orchestration.api.SubmitConfig()

Bases: _StrictModel

compute

dict[str, ComputeConfig]

deployed_services

dict[str, ServiceConfig]

services with each vllm_pd service expanded into its two tiers and its router.

driver

DriverConfig

job

JobConfig

otel

OtelConfig = OtelConfig()

services

dict[str, ServiceConfig]

nemo_gym.orchestration.api.SubmitConfig._resolve_and_validate_placements(
    info: pydantic.ValidationInfo
) -> nemo_gym.orchestration.api.SubmitConfig

nemo_gym.orchestration.api.SubmitConfig._validate_exports() -> nemo_gym.orchestration.api.SubmitConfig

nemo_gym.orchestration.api.SubmitConfig._validate_pd_services(
    compute: nemo_gym.orchestration.api.ComputeConfig
) -> None

nemo_gym.orchestration.api.SubmitConfig._validate_vllm_gpu_footprint(
    service_name: str,
    service: nemo_gym.orchestration.api.VllmServiceConfig,
    total_nodes: int,
    node_pools: dict[str, nemo_gym.orchestration.api.NodePool],
    gpus_per_node_values: list[int],
    is_ray_serve: bool
) -> None

class nemo_gym.orchestration.api.VllmNodeServiceConfig()

Bases: VllmServiceConfig

server_per_node

bool = False

nemo_gym.orchestration.api.VllmNodeServiceConfig.node_host_var(
    index: int
) -> str

The env var the batch script exports with the host of one node of this service’s pool.

nemo_gym.orchestration.api.VllmNodeServiceConfig.node_urls(
    count: int
) -> list[str]

One base URL per node, with each host left for the driver to resolve from its env var.

class nemo_gym.orchestration.api.VllmPDServiceConfig()

Bases: BaseModelServiceConfig

Prefill/decode disaggregated vLLM: two tiers behind a vllm-router.

Deployed as three services named <name>-prefill, <name>-decode and <name> (the router); driver.policy_model names the router.

_tiers

dict[str, VllmPDTierConfig] = PrivateAttr(default_factory=dict)

decode

VllmServiceConfig

decode_policy

str = 'cache_aware'

intra_node_data_parallel_size

int = 1

kv_connector

str = 'NixlConnector'

kv_load_failure_policy

str = 'fail'

log_level

str = 'error'

nixl_side_channel_port

int = 5600

prefill

VllmServiceConfig

prefill_policy

str = 'cache_aware'

request_timeout_secs

int = 86400

server_per_node

bool = False

type

Literal['vllm_pd']

nemo_gym.orchestration.api.VllmPDServiceConfig._build_tiers() -> nemo_gym.orchestration.api.VllmPDServiceConfig

nemo_gym.orchestration.api.VllmPDServiceConfig._fill_tier_defaults(
    data: typing.Any
) -> typing.Any

classmethod

nemo_gym.orchestration.api.VllmPDServiceConfig.tiers(
    name: str
) -> dict[str, nemo_gym.orchestration.api.VllmPDTierConfig]

class nemo_gym.orchestration.api.VllmPDTierConfig()

Bases: VllmServiceConfig

One tier of a vllm_pd service as deployed. Built by the service, never written by a user.

kv_transfer_config

dict[str, str]

nixl_side_channel_port

int

server_per_node

bool

class nemo_gym.orchestration.api.VllmServiceConfig()

Bases: BaseModelServiceConfig

data_parallel_rpc_port

int = 13345

extra_args

str = ''

number_of_instances

int = 1

pipeline_parallel_size

int = 1

tensor_parallel_size

int = 1

trust_remote_code

bool = False

type

Literal['vllm']

use_ray_serve

bool = False

nemo_gym.orchestration.api.VllmServiceConfig._default_health_check() -> nemo_gym.orchestration.api.VllmServiceConfig

nemo_gym.orchestration.api.VllmServiceConfig._validate_number_of_instances(
    v: int
) -> int

classmethod

class nemo_gym.orchestration.api._HostEnvRefs()

Bases: BaseModel

Remembers which env values came from host: so a dump can write the reference back.

_host_env_refs

dict[str, str] = PrivateAttr(default_factory=dict)

nemo_gym.orchestration.api._HostEnvRefs._keep_host_env_refs(
    data: typing.Any,
    handler: typing.Any
) -> typing.Any

classmethod

nemo_gym.orchestration.api._HostEnvRefs._serialize_env(
    env: dict[str, str],
    info: pydantic.SerializationInfo
) -> dict[str, str]

class nemo_gym.orchestration.api._StrictModel()

Bases: BaseModel

model_config

= ConfigDict(extra='forbid')

nemo_gym.orchestration.api._shown_policy_value(
    key: str,
    value: typing.Any
) -> str

A policy run value as a conflict message shows it, with no credential: an API key is never printed, and a URL keeps neither its password nor its query values.

nemo_gym.orchestration.api._without_url_credentials(
    value: typing.Any
) -> typing.Any

value with the user-info password and every query value of each URL in it masked, also inside a list.

nemo_gym.orchestration.api.effective_ray_serve(
    service: nemo_gym.orchestration.api.VllmServiceConfig,
    total_nodes: int,
    gpus_per_node_values: list[int]
) -> bool

Whether the Ray Serve gateway manages this service’s instances/routing instead of vLLM’s own DP.

nemo_gym.orchestration.api.resolve_env_dict(
    env: dict[str, str],
    saved_config: bool = False
) -> dict[str, str]

Resolve lit:/host:/runtime: prefixes on env values. Every value must use one of these prefixes; a missing or misspelled prefix raises rather than being guessed at. saved_config reads a value without a prefix as lit: (see SAVED_CONFIG).

lit:VALUE -> literal VALUE.
host:VAR -> read from os.environ[VAR] on the machine running gym eval submit; raises if VAR isn’t set there.
runtime:VAR -> left unresolved; canonicalized to runtime:VAR for executors to pick up and reference from the job’s own environment at run time.

nemo_gym.orchestration.api.ComputeConfig = Annotated[Annotated[SlurmComputeConfig, Tag('slurm')] | Annotated[KubernetesComp...

nemo_gym.orchestration.api.HOST_ENV_REFS = 'host_env_refs'

nemo_gym.orchestration.api.RUNTIME_ENV_PREFIX = 'runtime:'

nemo_gym.orchestration.api.SAVED_CONFIG = 'saved_config'

nemo_gym.orchestration.api.ServiceConfig = Annotated[Annotated[VllmNodeServiceConfig, Tag('vllm')] | Annotated[RayServiceCo...

nemo_gym.orchestration.api._COMMAND_AND_RUN = 'sets both `command` and `run`, but `run` only configures `gym eval run`, which ...

nemo_gym.orchestration.api._DECODE_NIXL_PORT_OFFSET = 1

nemo_gym.orchestration.api._ENV_VAR_NAME_RE = re.compile('^[A-Za-z_][A-Za-z0-9_]*$')

nemo_gym.orchestration.api._PD_SERVICE_ONLY_FIELDS = ('server_per_node', 'nixl_side_channel_port')

nemo_gym.orchestration.api._URL_QUERY_VALUE = re.compile('(?P<name>[?&][^=&#\\s]*=)[^&#\\s]*')