nemo_gym.orchestration.executors.slurm_script

View as Markdown

Module Contents

Functions

NameDescription
_build_ray_commandray start for one node of a ray service: the head, or with worker a node joining it.
_build_service_command-
_build_vllm_command-
_build_vllm_multi_instance_multi_node_command-
_build_vllm_ray_command-
_build_vllm_ray_serve_command-
_build_vllm_single_instance_multi_node_command-
_driver_nodeThe node the driver runs on in a multi-node job: the policy’s first node.
_escape_for_double_quoted_bashEscape text for safe embedding inside a double-quoted bash string (”…”).
_health_check_hostWhere this service answers its health probe: the node it runs on, as _service_nodelist places it.
_node_totals-
_places_servicesWhether the script declares the allocation’s host array to place services with.
_pool_directiveThe one value every pool agrees on for attribute.
_pool_of-
_pool_offsetsEach pool’s (first node index, node count) within the allocation.
_probe_hostHow the batch script, which runs on the allocation’s first node, reaches node.
_ray_worker_rangesEach spanned pool’s worker nodes as (first index, count): every node but the head.
_render_collector_health_check-
_render_collector_serviceThe collector’s srun step. Started before the model services so the scrape covers their
_render_collector_shutdownRun after the driver: one more scrape interval so the final counters are seen, then a
_render_directives-
_render_pool_directivesOne set of directives for the whole allocation, not one per pool.
_render_pool_nodesDeclare the allocation’s hosts, and export each node pool’s, for —nodelist.
_render_ray_head_addressesExport every ray service’s head address, and its worker hosts, from the nodes Slurm gave it.
_render_ray_serviceA ray service’s srun steps: the head beside the driver, then one step per spanned pool’s workers.
_render_service_command-
_resolve_envReturn an ‘env K=V …’ prefix string (trailing space) scoped to a single command, or ” if empty.
_service_nodelistWhat a service’s —nodelist names, as the shell variable it expands.
_service_nodesHow many nodes this service actually runs on.
_srun_nodes-
_srun_ntasks-
_strip_headless_incompatible_flags-
_validate_env_key-
_vllm_base_flags-
_vllm_spans_multiple_nodes-
_with_default_capture_dirEnable capture for submitted evaluations unless explicitly disabled.
build_sbatch_script-
pool_nodes_varThe env var build_sbatch_script exports with a node pool’s comma-separated hosts.
ray_head_address_varThe env var build_sbatch_script exports with a ray service’s head host:port.
ray_workers_varThe env var build_sbatch_script exports with a ray service’s worker hosts in one pool.

Data

_BUILDERS

_HEADLESS_INCOMPATIBLE_FLAG

_NODE_ARRAY

_RAY_HEAD_ONLY_FLAG

_RAY_SERVE_GATEWAY_SOURCE_PATH

_SCRIPT_TEMPLATE

_VALID_ENV_KEY

API

nemo_gym.orchestration.executors.slurm_script._build_ray_command(
pool: str | None = None,
address_var: str | None = None,
worker: bool = False
) -> str

ray start for one node of a ray service: the head, or with worker a node joining it.

address_var names the env var holding the head’s host:port (see ray_head_address_var).

nemo_gym.orchestration.executors.slurm_script._build_service_command(
total_nodes: int,
gpus_per_node_values: list[int]
) -> str
nemo_gym.orchestration.executors.slurm_script._build_vllm_command(
) -> str
nemo_gym.orchestration.executors.slurm_script._build_vllm_multi_instance_multi_node_command(
total_nodes: int
) -> str
nemo_gym.orchestration.executors.slurm_script._build_vllm_ray_command(
total_nodes: int
) -> str
nemo_gym.orchestration.executors.slurm_script._build_vllm_ray_serve_command(
total_nodes: int,
gpus_per_node_values: list[int]
) -> str
nemo_gym.orchestration.executors.slurm_script._build_vllm_single_instance_multi_node_command(
total_nodes: int
) -> str
nemo_gym.orchestration.executors.slurm_script._driver_node(
) -> int

The node the driver runs on in a multi-node job: the policy’s first node.

The driver reaches the policy on localhost. A multi-node policy serves its API from node 0, and a pinned one from its pool’s first node.

nemo_gym.orchestration.executors.slurm_script._escape_for_double_quoted_bash(
text: str
) -> str

Escape text for safe embedding inside a double-quoted bash string (”…”).

nemo_gym.orchestration.executors.slurm_script._health_check_host(
driver_node: int | None,
total_nodes: int
) -> str

Where this service answers its health probe: the node it runs on, as _service_nodelist places it.

nemo_gym.orchestration.executors.slurm_script._node_totals(
) -> tuple[int, int]
nemo_gym.orchestration.executors.slurm_script._places_services(
is_multi_node: bool
) -> bool

Whether the script declares the allocation’s host array to place services with.

nemo_gym.orchestration.executors.slurm_script._pool_directive(
attribute: str
) -> typing.Any

The one value every pool agrees on for attribute.

A plain sbatch job takes a single —partition/—ntasks-per-node/—gpus-per-node for the whole allocation, so pools that disagree cannot both be honoured. Slurm would silently apply whichever directive came last; say so instead.

nemo_gym.orchestration.executors.slurm_script._pool_of(
node: int
) -> str | None
nemo_gym.orchestration.executors.slurm_script._pool_offsets(
) -> dict[str, tuple[int, int]]

Each pool’s (first node index, node count) within the allocation.

Pools are laid out contiguously in declaration order, which is the order the single #SBATCH —nodes total is built from, so pool i owns the nodes after every pool before it.

nemo_gym.orchestration.executors.slurm_script._probe_host(
node: int
) -> str

How the batch script, which runs on the allocation’s first node, reaches node.

nemo_gym.orchestration.executors.slurm_script._ray_worker_ranges(
head_node: int
) -> dict[str, tuple[int, int]]

Each spanned pool’s worker nodes as (first index, count): every node but the head.

The head is the driver’s node, which is always the first node of some pool.

nemo_gym.orchestration.executors.slurm_script._render_collector_health_check(
driver_node: int | None = None
) -> str
nemo_gym.orchestration.executors.slurm_script._render_collector_service(
remote_bench_dir: pathlib.Path,
is_multi_node: bool,
driver_node: int | None = None
) -> str

The collector’s srun step. Started before the model services so the scrape covers their startup.

Runs on exactly one node, the batch host. That is the first node of the allocation, which is also where the Ray prelude puts the head of a multi-node vLLM service and where the driver runs, so localhost:<port> reaches the API server and the OTLP endpoints from the same node. In a container the job directory is mounted for the config and the local otel/*.jsonl output; on the node it is simply there, so no mounts or workdir are passed.

nemo_gym.orchestration.executors.slurm_script._render_collector_shutdown(
remote_bench_dir: pathlib.Path
) -> str

Run after the driver: one more scrape interval so the final counters are seen, then a graceful stop, keeping the driver’s exit code.

The TERM goes to the collector process itself, matched by its unique --config path. Sent to srun instead, TERM makes Slurm kill the step outright and INT is treated as a console interrupt; neither reaches the collector, so its final batch would be lost.

nemo_gym.orchestration.executors.slurm_script._render_directives(
remote_bench_dir: pathlib.Path,
benchmark_name: str
) -> str
nemo_gym.orchestration.executors.slurm_script._render_pool_directives(
) -> list[str]

One set of directives for the whole allocation, not one per pool.

Pools divide an allocation between services (see _pool_offsets); they are not separate Slurm requests. Emitting —nodes per pool made every pool but the last a no-op, so a two-pool job asked for one pool’s nodes while the rest of the executor sized itself on the sum.

nemo_gym.orchestration.executors.slurm_script._render_pool_nodes(
is_multi_node: bool
) -> str

Declare the allocation’s hosts, and export each node pool’s, for —nodelist.

nemo_gym.orchestration.executors.slurm_script._render_ray_head_addresses(
) -> str

Export every ray service’s head address, and its worker hosts, from the nodes Slurm gave it.

The head runs beside the driver. Workers and the driver join it without knowing, when the config is written, which host the job will land on.

nemo_gym.orchestration.executors.slurm_script._render_ray_service(
name: str,
driver_node: int | None
) -> str

A ray service’s srun steps: the head beside the driver, then one step per spanned pool’s workers.

Every step uses the service’s one container, env and mounts.

nemo_gym.orchestration.executors.slurm_script._render_service_command(
name: str,
container: str | None,
command: str,
env: dict[str, str] | None = None,
mounts: list[str] | None = None,
nodes: int | None = None,
ntasks: int | None = None,
pre_command: str = '',
workdir: str | None = None,
single_node: bool = False,
nodelist: str | 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.

A runtime:VAR value (see resolve_env_dict in api.py) is emitted as an unquoted K=$VAR shell reference instead of a literal, so it’s resolved from the job’s own environment when the command actually runs on the compute node, rather than baked in at script-generation time.

nemo_gym.orchestration.executors.slurm_script._service_nodelist(
driver_node: int | None,
total_nodes: int
) -> str | None

What a service’s —nodelist names, as the shell variable it expands.

A pinned service names its pool. In a multi-node job an unpinned service that fits on one node joins the driver, which reaches it on localhost; otherwise Slurm could start it on any node.

nemo_gym.orchestration.executors.slurm_script._service_nodes(
total_nodes: int
) -> int

How many nodes this service actually runs on.

A service pinned to a pool sees only that pool, so a single-node pool inside a ten-node job is a single-node deployment and must not be built as a multi-node Ray one.

nemo_gym.orchestration.executors.slurm_script._srun_nodes(
total_nodes: int
) -> int | None
nemo_gym.orchestration.executors.slurm_script._srun_ntasks(
total_nodes: int,
total_ntasks: int
) -> int | None
nemo_gym.orchestration.executors.slurm_script._strip_headless_incompatible_flags(
cmd: str
) -> str
nemo_gym.orchestration.executors.slurm_script._validate_env_key(
key: str
) -> None
nemo_gym.orchestration.executors.slurm_script._vllm_base_flags(
) -> str
nemo_gym.orchestration.executors.slurm_script._vllm_spans_multiple_nodes(
total_nodes: int
) -> bool
nemo_gym.orchestration.executors.slurm_script._with_default_capture_dir(
run: dict[str, typing.Any],
remote_bench_dir: pathlib.Path
) -> dict[str, typing.Any]

Enable capture for submitted evaluations unless explicitly disabled.

Auto-derive model_call_capture_dir from this benchmark’s own real output directory when observability is on and the caller didn’t set one.

Hydra interpolation resolves before remote_bench_dir exists (it’s computed here, in build_sbatch_script, well after SubmitConfig validation), so there’s no way for a YAML value to reference it — this has to happen in Python, once the real path is known. An explicit model_call_capture_dir in run always wins over this default.

nemo_gym.orchestration.executors.slurm_script.build_sbatch_script(
benchmark_name: str,
remote_bench_dir: pathlib.Path
) -> str
nemo_gym.orchestration.executors.slurm_script.pool_nodes_var(
pool: str
) -> str

The env var build_sbatch_script exports with a node pool’s comma-separated hosts.

nemo_gym.orchestration.executors.slurm_script.ray_head_address_var(
ray_service: str
) -> str

The env var build_sbatch_script exports with a ray service’s head host:port.

nemo_gym.orchestration.executors.slurm_script.ray_workers_var(
ray_service: str,
pool: str
) -> str

The env var build_sbatch_script exports with a ray service’s worker hosts in one pool.

nemo_gym.orchestration.executors.slurm_script._BUILDERS = {VllmServiceConfig: _build_vllm_command, RayServiceConfig: _build_ray_command}
nemo_gym.orchestration.executors.slurm_script._HEADLESS_INCOMPATIBLE_FLAG = re.compile('\\s--api-server-count(?:[= ]\\S+)?')
nemo_gym.orchestration.executors.slurm_script._NODE_ARRAY = 'gym_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))'
nemo_gym.orchestration.executors.slurm_script._RAY_HEAD_ONLY_FLAG = re.compile('\\s--(?:port|redis-shard-ports|include-dashboard)(?:[= ]\\S+)?')
nemo_gym.orchestration.executors.slurm_script._RAY_SERVE_GATEWAY_SOURCE_PATH = Path(__file__).resolve().parent.parent / 'ray_serve_gateway.py'
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_]*$')