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`	-
`RayServiceConfig`	-
`SlurmComputeConfig`	-
`SubmitConfig`	-
`VllmServiceConfig`	-
`_StrictModel`	-

Data

ComputeConfig

ServiceConfig

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

class nemo_gym.orchestration.api.BaseServiceConfig()

Bases: _StrictModel

container

str

env

dict[str, str] = {}

health_check

HealthCheckConfig | None = None

mounts

list[str] = []

placement

str | None = None

class nemo_gym.orchestration.api.BenchmarkRunConfig()

Bases: _StrictModel

prepare

dict[str, Any] = {}

run

dict[str, Any] = {}

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_model

str | None = None

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

Bases: BaseServiceConfig

type

Literal['ray']

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]

driver

DriverConfig

job

JobConfig

services

dict[str, ServiceConfig]

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

class nemo_gym.orchestration.api.VllmServiceConfig()

Bases: BaseModelServiceConfig

pipeline_parallel_size

int = 1

tensor_parallel_size

int = 1

trust_remote_code

bool = False

type

Literal['vllm']

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

class nemo_gym.orchestration.api._StrictModel()

Bases: BaseModel

model_config

= ConfigDict(extra='forbid')

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

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