nemo_gym.orchestration.executors.slurm_script

Module Contents

Functions

Name	Description
`_build_ray_command`	-
`_build_service_command`	-
`_build_vllm_command`	-
`_build_vllm_multi_instance_multi_node_command`	-
`_build_vllm_ray_command`	-
`_build_vllm_single_instance_multi_node_command`	-
`_node_totals`	-
`_render_directives`	-
`_render_pool_directives`	-
`_render_service_command`	-
`_resolve_env`	Return an ‘env K=V …’ prefix string (trailing space) scoped to a single command, or ” if empty.
`_validate_env_key`	-
`_vllm_base_flags`	-
`_vllm_spans_multiple_nodes`	-
`build_sbatch_script`	-

Data

_BUILDERS

_SCRIPT_TEMPLATE

_VALID_ENV_KEY

API

nemo_gym.orchestration.executors.slurm_script._build_ray_command(
    _service: nemo_gym.orchestration.api.RayServiceConfig
) -> str

nemo_gym.orchestration.executors.slurm_script._build_service_command(
    service: nemo_gym.orchestration.api.VllmServiceConfig | nemo_gym.orchestration.api.RayServiceConfig,
    total_nodes: int
) -> str

nemo_gym.orchestration.executors.slurm_script._build_vllm_command(
    service: nemo_gym.orchestration.api.VllmServiceConfig
) -> str

nemo_gym.orchestration.executors.slurm_script._build_vllm_multi_instance_multi_node_command(
    service: nemo_gym.orchestration.api.VllmServiceConfig,
    total_nodes: int
) -> str

nemo_gym.orchestration.executors.slurm_script._build_vllm_ray_command(
    service: nemo_gym.orchestration.api.VllmServiceConfig,
    total_nodes: int
) -> str

nemo_gym.orchestration.executors.slurm_script._build_vllm_single_instance_multi_node_command(
    service: nemo_gym.orchestration.api.VllmServiceConfig,
    total_nodes: int
) -> str

nemo_gym.orchestration.executors.slurm_script._node_totals(
    compute: nemo_gym.orchestration.api.SlurmComputeConfig
) -> tuple[int, int]

nemo_gym.orchestration.executors.slurm_script._render_directives(
    compute: nemo_gym.orchestration.api.SlurmComputeConfig,
    remote_bench_dir: pathlib.Path,
    benchmark_name: str
) -> str

nemo_gym.orchestration.executors.slurm_script._render_pool_directives(
    pool_name: str,
    pool: nemo_gym.orchestration.api.NodePool
) -> list[str]

nemo_gym.orchestration.executors.slurm_script._render_service_command(
    name: str,
    container: str,
    command: str,
    env: dict[str, str] | None = None,
    mounts: list[str] | None = None,
    nodes: int | None = None,
    ntasks: int | None = None
) -> str

nemo_gym.orchestration.executors.slurm_script._resolve_env(
    env: dict[str, str]
) -> str

Return an ‘env K=V …’ prefix string (trailing space) scoped to a single command, or ” if empty.

nemo_gym.orchestration.executors.slurm_script._validate_env_key(
    key: str
) -> None

nemo_gym.orchestration.executors.slurm_script._vllm_base_flags(
    service: nemo_gym.orchestration.api.VllmServiceConfig
) -> str

nemo_gym.orchestration.executors.slurm_script._vllm_spans_multiple_nodes(
    service: nemo_gym.orchestration.api.VllmServiceConfig | nemo_gym.orchestration.api.RayServiceConfig,
    total_nodes: int
) -> bool

nemo_gym.orchestration.executors.slurm_script.build_sbatch_script(
    config: nemo_gym.orchestration.api.SubmitConfig,
    benchmark_name: str,
    benchmark: nemo_gym.orchestration.api.BenchmarkRunConfig,
    compute: nemo_gym.orchestration.api.SlurmComputeConfig,
    remote_bench_dir: pathlib.Path
) -> str

nemo_gym.orchestration.executors.slurm_script._BUILDERS = {VllmServiceConfig: _build_vllm_command, RayServiceConfig: _build_ray_command}

nemo_gym.orchestration.executors.slurm_script._SCRIPT_TEMPLATE = '#!/bin/bash\n{directives}\n\n{ray_prelude}\n\n{service_commands}\n\n{health_che...

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