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`	-
`NodePool`	-
`OtelConfig`	An OpenTelemetry collector beside every benchmark job: scrapes each model service’s
`RayServiceConfig`	-
`ResumeConfig`	-
`SlurmComputeConfig`	-
`SubmitConfig`	-
`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`	-
`_StrictModel`	-

Functions

Name	Description
`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

ComputeConfig

RUNTIME_ENV_PREFIX

ServiceConfig

_DECODE_NIXL_PORT_OFFSET

_ENV_VAR_NAME_RE

_PD_SERVICE_ONLY_FIELDS

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: _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]
) -> dict[str, str]

classmethod

class nemo_gym.orchestration.api.BenchmarkRunConfig()

Bases: _StrictModel

command

str | None = None

prepare

dict[str, Any] = {}

resumable

bool | ResumeConfig = False

resume_config

ResumeConfig | None

run

dict[str, Any] = {}

nemo_gym.orchestration.api.BenchmarkRunConfig._validate_command() -> nemo_gym.orchestration.api.BenchmarkRunConfig

class nemo_gym.orchestration.api.DriverConfig()

Bases: _StrictModel

benchmarks

dict[str, BenchmarkRunConfig]

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]
) -> 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

output_path

str

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; endpoint and service_name come from the deployment’s own config (a cluster fragment, typically) and are required while enabled.

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

scrape_interval_seconds

int = 15

service_name

str | None = None

token_env

str = 'OTEL_TOKEN'

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 = 3

max_walltime

str | None = None

class nemo_gym.orchestration.api.SlurmComputeConfig()

Bases: BaseComputeConfig

account

str

extra_args

dict[str, str] = {}

hostname

str | None = None

node_pools

dict[str, NodePool] = {}

type

Literal['slurm']

walltime

str | None = None

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() -> 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.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._StrictModel()

Bases: BaseModel

model_config

= ConfigDict(extra='forbid')

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]
) -> 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.

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')], Discriminator('type')]

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

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

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')