diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index ef653842ac3..3b8bb47c387 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -80,6 +80,12 @@ class BaseServiceConfig(_StrictModel): container: str # Resolved to the sole compute resource name at validation time when not set. placement: str | None = None + # Name of a node pool in the placed compute's `node_pools`. Pins this service to + # that pool's slice of the allocation instead of letting it land wherever srun + # starts, which is how two services get nodes of their own -- a scorer or judge + # that cannot share a GPU with the policy, or a prefill/decode split. Pools take + # contiguous node ranges in declaration order. None means the whole allocation. + node_pool: str | None = None health_check: HealthCheckConfig | None = None # Values may be prefixed `lit:` (literal), `host:VAR` (read from the submitting # machine's env), or `runtime:VAR` (resolved from the job's own env at run time). @@ -152,6 +158,34 @@ def effective_ray_serve(service: "VllmServiceConfig", total_nodes: int, gpus_per class RayServiceConfig(BaseServiceConfig): type: Literal["ray"] + # Node pools the cluster spans. The head starts on the driver's node and every + # other node of these pools joins it as a worker. Empty: a single-node cluster. + node_pools: list[str] = [] + port: int = 6379 + # Custom Ray resources advertised by each spanned pool's nodes, keyed by pool, + # e.g. {"aux": {"extra_gpu": 4}}. A benchmark asks for these by name when it + # manages device placement itself. + resources: dict[str, dict[str, float]] = {} + num_cpus: int | None = None + num_gpus: int | None = None + # Raw extra flags appended verbatim to `ray start` on every node (e.g. fixed ports). + extra_args: str = "" + + @model_validator(mode="after") + def _validate_shape(self) -> "RayServiceConfig": + if self.node_pool is not None: + raise ValueError( + f"A ray service spans `node_pools`, not a single `node_pool`; use node_pools: [{self.node_pool}]." + ) + if len(set(self.node_pools)) != len(self.node_pools): + raise ValueError(f"A ray service's node_pools lists a pool more than once: {self.node_pools}.") + unspanned = sorted(set(self.resources) - set(self.node_pools)) + if unspanned: + raise ValueError( + f"A ray service sets resources for {', '.join(unspanned)}, which it does not span " + f"(node_pools: {self.node_pools}). Add the pool to node_pools or drop its resources." + ) + return self # Discriminated union keyed on `type`; Pydantic rejects unknown type values at parse time. @@ -304,12 +338,8 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": total_nodes = ( sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1 ) - is_multi_node = total_nodes > 1 - gpus_per_node_values = ( - [p.gpus_per_node for p in compute.node_pools.values() if p.gpus_per_node is not None] - if isinstance(compute, SlurmComputeConfig) - else [] - ) + + pool_names = set(compute.node_pools) if isinstance(compute, SlurmComputeConfig) else set() for service_name, service in self.services.items(): if service.placement is None: @@ -320,25 +350,50 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) + if isinstance(service, RayServiceConfig): + unknown = [pool for pool in service.node_pools if pool not in pool_names] + if unknown: + raise ValueError( + f"Service '{service_name}' node_pools {unknown} do not match any node pool of compute " + f"'{service.placement}' ({', '.join(sorted(pool_names)) or 'none declared'})." + ) + + if service.node_pool is not None and service.node_pool not in pool_names: + raise ValueError( + f"Service '{service_name}' node_pool '{service.node_pool}' does not match any node pool of " + f"compute '{service.placement}' ({', '.join(sorted(pool_names)) or 'none declared'})." + ) + if not isinstance(service, VllmServiceConfig): continue - is_ray_serve = effective_ray_serve(service, total_nodes, gpus_per_node_values) + # A pinned service is sized against its own pool, not the whole job: one node of a + # ten-node allocation is a single-node deployment with that pool's GPUs, and judging + # it by the allocation total both mis-builds the command and mis-reports idle GPUs. + service_pools = ( + {service.node_pool: compute.node_pools[service.node_pool]} + if service.node_pool is not None and isinstance(compute, SlurmComputeConfig) + else (compute.node_pools if isinstance(compute, SlurmComputeConfig) else {}) + ) + service_nodes = sum(p.nodes for p in service_pools.values()) or total_nodes + service_gpus = [p.gpus_per_node for p in service_pools.values() if p.gpus_per_node is not None] + + is_ray_serve = effective_ray_serve(service, service_nodes, service_gpus) if ( - is_multi_node + service_nodes > 1 and service.number_of_instances > 1 - and service.number_of_instances % total_nodes != 0 + and service.number_of_instances % service_nodes != 0 and not is_ray_serve ): raise ValueError( f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must " - f"be evenly divisible by the number of nodes ({total_nodes}) for multi-node data-parallel " + f"be evenly divisible by the number of nodes ({service_nodes}) for multi-node data-parallel " "deployment - each node hosts an equal share of the data-parallel replicas." ) self._validate_vllm_gpu_footprint( - service_name, service, total_nodes, compute, gpus_per_node_values, is_ray_serve + service_name, service, service_nodes, service_pools, service_gpus, is_ray_serve ) if self.driver.policy_model is not None: @@ -370,7 +425,7 @@ def _validate_vllm_gpu_footprint( service_name: str, service: "VllmServiceConfig", total_nodes: int, - compute: "SlurmComputeConfig", + node_pools: dict[str, "NodePool"], gpus_per_node_values: list[int], is_ray_serve: bool, ) -> None: @@ -383,9 +438,7 @@ def _validate_vllm_gpu_footprint( if total_nodes > 1 and is_ray_serve: # Ray Serve's placement-group scheduler packs the aggregate footprint across the cluster. gpus_needed = tp_pp * service.number_of_instances - gpus_available = sum( - pool.nodes * pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node - ) + gpus_available = sum(pool.nodes * pool.gpus_per_node for pool in node_pools.values() if pool.gpus_per_node) footprint = ( f"tensor_parallel_size={service.tensor_parallel_size} x " f"pipeline_parallel_size={service.pipeline_parallel_size} x " @@ -408,9 +461,7 @@ def _validate_vllm_gpu_footprint( elif total_nodes > 1: # Single instance's TP/PP footprint spans the whole allocation via the ray backend. gpus_needed = tp_pp - gpus_available = sum( - pool.nodes * pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node - ) + gpus_available = sum(pool.nodes * pool.gpus_per_node for pool in node_pools.values() if pool.gpus_per_node) footprint = ( f"tensor_parallel_size={service.tensor_parallel_size} x " f"pipeline_parallel_size={service.pipeline_parallel_size}" diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index ce9fb5d52c1..f30029ddf26 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -134,11 +134,13 @@ def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_fla ) -def render_health_check(name: str, port: int, path: str, timeout: int) -> str: +def render_health_check(name: str, port: int, path: str, timeout: int, host: str = "localhost") -> str: + """`host` is emitted verbatim so it may be a shell expansion: a service pinned to + a node pool answers on that pool's head, not on the node running this script.""" return _HEALTH_WAIT_MULTI.format( name=name, name_upper=bash_var(name), - url=f"http://localhost:{port}", + url=f"http://{host}:{port}", path=path, max_attempts=timeout // 5, ) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 399b899f539..1d6abfb5b69 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import re import shlex from pathlib import Path @@ -78,21 +79,54 @@ def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benc lines.append(f"#SBATCH --chdir={remote_bench_dir}") for key, val in compute.extra_args.items(): lines.append(f"#SBATCH --{key}={val}") - for pool_name, pool in compute.node_pools.items(): - lines.extend(_render_pool_directives(pool_name, pool)) + lines.extend(_render_pool_directives(compute.node_pools)) return "\n".join(lines) -def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: +def _pool_directive(pools: dict[str, NodePool], attribute: str) -> 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. + """ + values = {getattr(pool, attribute) for pool in pools.values()} + if len(values) > 1: + named = ", ".join(f"{name}={getattr(pool, attribute)!r}" for name, pool in pools.items()) + raise ValueError( + f"Node pools disagree on {attribute} ({named}). One Slurm job takes a single value for the whole " + "allocation; split the run or make the pools agree." + ) + return next(iter(values)) + + +def _render_pool_directives(pools: dict[str, NodePool]) -> 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. + """ + if not pools: + return [] lines = [ - f"#SBATCH --partition={pool.partition} # pool: {pool_name}", - f"#SBATCH --nodes={pool.nodes}", - f"#SBATCH --ntasks-per-node={pool.ntasks_per_node}", + f"#SBATCH --partition={_pool_directive(pools, 'partition')}", + f"#SBATCH --nodes={sum(pool.nodes for pool in pools.values())}", + f"#SBATCH --ntasks-per-node={_pool_directive(pools, 'ntasks_per_node')}", ] - if pool.gpus_per_node is not None: - lines.append(f"#SBATCH --gpus-per-node={pool.gpus_per_node}") - for key, val in pool.extra_args.items(): - lines.append(f"#SBATCH --{key}={val}") + gpus_per_node = _pool_directive(pools, "gpus_per_node") + if gpus_per_node is not None: + lines.append(f"#SBATCH --gpus-per-node={gpus_per_node}") + extra_args: dict[str, str] = {} + for name, pool in pools.items(): + for key, val in pool.extra_args.items(): + if extra_args.setdefault(key, val) != val: + raise ValueError( + f"Node pool {name!r} sets extra_args[{key!r}]={val!r}, which conflicts with another pool's " + f"{extra_args[key]!r}. #SBATCH directives apply to the whole allocation." + ) + lines.extend(f"#SBATCH --{key}={val}" for key, val in extra_args.items()) return lines @@ -133,6 +167,7 @@ def _render_service_command( pre_command: str = "", workdir: str | None = None, single_node: bool = False, + nodelist: str | None = None, ) -> str: var = bash_var(name) env_prefix = _resolve_env(env) if env else "" @@ -141,6 +176,10 @@ def _render_service_command( node_flags = " --nodes=1 --ntasks=1" else: node_flags = f" --nodes={nodes} --ntasks={ntasks}" if (nodes is not None and nodes > 1) else "" + # --nodelist names the exact hosts. --relative is only a starting point Slurm may + # move off when that node's resources are taken, landing a service on the wrong node. + if nodelist is not None: + node_flags = f' --nodelist="${{{nodelist}}}" --nodes={nodes} --ntasks={ntasks}' mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else "" workdir_flag = f" --container-workdir={shlex.quote(workdir)}" if workdir else "" if pre_command: @@ -315,8 +354,170 @@ def _build_vllm_ray_serve_command( ) -def _build_ray_command(_service: RayServiceConfig) -> str: - return "ray start --head" +# Flags `ray start` refuses on a worker. They reach a worker through the service's +# shared extra_args; Gym decides which nodes are workers, so Gym drops them there. +_RAY_HEAD_ONLY_FLAG = re.compile(r"\s--(?:port|redis-shard-ports|include-dashboard)(?:[= ]\S+)?") + + +def _build_ray_command( + service: RayServiceConfig, *, 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). + """ + # --block keeps the srun step alive. `ray start` daemonises and returns, so + # without it the step exits the moment the node is up and Slurm tears the + # service down again. + cmd = "ray start --block" + if worker: + cmd += f' --address "${address_var}"' + else: + cmd += f" --head --port {service.port}" + if address_var is not None: + # Bind the head to the address the workers and the driver dial. + cmd += f' --node-ip-address "${{{address_var}%:*}}"' + if service.num_cpus is not None: + cmd += f" --num-cpus {service.num_cpus}" + if service.num_gpus is not None: + cmd += f" --num-gpus {service.num_gpus}" + resources = service.resources.get(pool, {}) if pool is not None else {} + if resources: + # Ray takes fractional custom resources, so the field is float-typed, but a + # whole number is written as one: {"extra_gpu": 4}, not 4.0. + resources = {k: int(v) if v.is_integer() else v for k, v in resources.items()} + cmd += " --resources=" + shlex.quote(json.dumps(resources, sort_keys=True)) + extra_args = _RAY_HEAD_ONLY_FLAG.sub("", " " + service.extra_args).strip() if worker else service.extra_args + if extra_args: + cmd += " " + extra_args + return cmd + + +_NODE_ARRAY = 'gym_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))' + + +def pool_nodes_var(pool: str) -> str: + """The env var build_sbatch_script exports with a node pool's comma-separated hosts.""" + return f"GYM_POOL_{bash_var(pool)}_NODES" + + +def _places_services(config: SubmitConfig, is_multi_node: bool) -> bool: + """Whether the script declares the allocation's host array to place services with.""" + return is_multi_node or any(s.node_pool for s in config.services.values()) + + +def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig, is_multi_node: bool) -> str: + """Declare the allocation's hosts, and export each node pool's, for --nodelist.""" + if not _places_services(config, is_multi_node): + return "" + lines = [_NODE_ARRAY] + for name, (start, count) in _pool_offsets(compute).items(): + if count: + lines.append(f'export {pool_nodes_var(name)}="$(IFS=,; echo "${{gym_nodes[*]:{start}:{count}}}")"') + return "\n".join(lines) + + +def ray_head_address_var(ray_service: str) -> str: + """The env var build_sbatch_script exports with a ray service's head host:port.""" + return f"GYM_RAY_ADDRESS_{bash_var(ray_service)}" + + +def 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.""" + return f"GYM_RAY_{bash_var(ray_service)}_{bash_var(pool)}_WORKERS" + + +def _ray_worker_ranges( + service: RayServiceConfig, compute: SlurmComputeConfig, 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. + """ + offsets = _pool_offsets(compute) + ranges: dict[str, tuple[int, int]] = {} + for pool in service.node_pools: + start, count = offsets[pool] + if start == head_node: + start, count = start + 1, count - 1 + if count > 0: + ranges[pool] = (start, count) + return ranges + + +def _render_ray_head_addresses(config: SubmitConfig, compute: SlurmComputeConfig) -> 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. + """ + services = {n: s for n, s in config.services.items() if isinstance(s, RayServiceConfig)} + if not services: + return "" + head_node = _driver_node(config, compute) + lines = [] if _places_services(config, _node_totals(compute)[0] > 1) else [_NODE_ARRAY] + for name, service in services.items(): + ip = f"$(getent hosts ${{gym_nodes[{head_node}]}} | awk '{{print $1}}')" + lines.append(f'export {ray_head_address_var(name)}="{ip}:{service.port}"') + for pool, (start, count) in _ray_worker_ranges(service, compute, head_node).items(): + lines.append(f'export {ray_workers_var(name, pool)}="$(IFS=,; echo "${{gym_nodes[*]:{start}:{count}}}")"') + return "\n".join(lines) + + +def _pool_of(compute: SlurmComputeConfig, node: int) -> str | None: + for name, (start, count) in _pool_offsets(compute).items(): + if start <= node < start + count: + return name + return None + + +def _render_ray_service( + name: str, + service: RayServiceConfig, + config: SubmitConfig, + compute: SlurmComputeConfig, + 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. + """ + total_nodes, total_ntasks = _node_totals(compute) + head_node = _driver_node(config, compute) + head_pool = _pool_of(compute, head_node) + address_var = ray_head_address_var(name) + steps = [ + _render_service_command( + name, + service.container, + _build_ray_command( + service, pool=head_pool if head_pool in service.node_pools else None, address_var=address_var + ), + service.env or None, + service.mounts or None, + nodes=_srun_nodes(service, compute, total_nodes), + ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks), + pre_command=service.pre_command, + nodelist=_service_nodelist(service, driver_node, total_nodes), + ) + ] + # A worker only joins a running head, and the head may still be installing. + wait_for_head = f'until ray status --address "${address_var}" >/dev/null 2>&1; do sleep 5; done' + for pool, (_, count) in _ray_worker_ranges(service, compute, head_node).items(): + steps.append( + _render_service_command( + f"{name}_{pool}_workers", + service.container, + _build_ray_command(service, pool=pool, address_var=address_var, worker=True), + service.env or None, + service.mounts or None, + nodes=count, + ntasks=count, + pre_command=f"{service.pre_command.rstrip()}\n{wait_for_head}".lstrip(), + nodelist=ray_workers_var(name, pool), + ) + ) + return "\n\n".join(steps) _BUILDERS = { @@ -344,7 +545,9 @@ def _build_service_command( return _BUILDERS[type(service)](service) -def _render_collector_service(config: SubmitConfig, remote_bench_dir: Path, *, is_multi_node: bool) -> str: +def _render_collector_service( + config: SubmitConfig, remote_bench_dir: 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. @@ -370,13 +573,22 @@ def _render_collector_service(config: SubmitConfig, remote_bench_dir: Path, *, i "SLURM_JOB_ID": f"{RUNTIME_ENV_PREFIX}SLURM_JOB_ID", }, single_node=is_multi_node, + # Beside the driver: it scrapes the policy API on localhost, as the driver calls it. + nodelist=f"gym_nodes[{driver_node}]" if driver_node is not None else None, + nodes=1 if driver_node is not None else None, + ntasks=1 if driver_node is not None else None, **container_kwargs, ) -def _render_collector_health_check(config: SubmitConfig) -> str: +def _render_collector_health_check(config: SubmitConfig, driver_node: int | None = None) -> str: + # The collector runs beside the driver (see _render_collector_service). return render_health_check( - COLLECTOR_SERVICE_NAME, COLLECTOR_HEALTH_PORT, "/", config.otel.health_check_timeout_seconds + COLLECTOR_SERVICE_NAME, + COLLECTOR_HEALTH_PORT, + "/", + config.otel.health_check_timeout_seconds, + _probe_host(driver_node or 0), ) @@ -403,6 +615,88 @@ def _render_collector_shutdown(config: SubmitConfig, remote_bench_dir: Path) -> ) +def _pool_offsets(compute: SlurmComputeConfig) -> 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. + """ + offsets: dict[str, tuple[int, int]] = {} + start = 0 + for name, pool in compute.node_pools.items(): + offsets[name] = (start, pool.nodes) + start += pool.nodes + return offsets + + +def _service_nodes( + service: VllmServiceConfig | RayServiceConfig, compute: SlurmComputeConfig, 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. + """ + if service.node_pool is None: + return total_nodes + return compute.node_pools[service.node_pool].nodes + + +def _driver_node(config: SubmitConfig, compute: SlurmComputeConfig) -> 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. + """ + policy = config.services.get(config.driver.policy_model or "") + if policy is not None and policy.node_pool is not None: + return _pool_offsets(compute)[policy.node_pool][0] + return 0 + + +def _service_nodelist( + service: VllmServiceConfig | RayServiceConfig, 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. + """ + if service.node_pool is not None: + return pool_nodes_var(service.node_pool) + if driver_node is not None and not _vllm_spans_multiple_nodes(service, total_nodes): + return f"gym_nodes[{driver_node}]" + return None + + +def _srun_nodes( + service: VllmServiceConfig | RayServiceConfig, compute: SlurmComputeConfig, total_nodes: int +) -> int | None: + nodes = _service_nodes(service, compute, total_nodes) + if service.node_pool is not None: + return nodes + if _vllm_spans_multiple_nodes(service, total_nodes): + return total_nodes + return 1 if total_nodes > 1 else None + + +def _srun_ntasks( + service: VllmServiceConfig | RayServiceConfig, + compute: SlurmComputeConfig, + total_nodes: int, + total_ntasks: int, +) -> int | None: + if service.node_pool is not None: + pool = compute.node_pools[service.node_pool] + return pool.nodes * pool.ntasks_per_node + if _vllm_spans_multiple_nodes(service, total_nodes): + return total_ntasks + return 1 if total_nodes > 1 else None + + def _node_totals(compute: SlurmComputeConfig) -> tuple[int, int]: total_nodes = sum(pool.nodes for pool in compute.node_pools.values()) total_ntasks = sum(pool.nodes * pool.ntasks_per_node for pool in compute.node_pools.values()) @@ -424,6 +718,28 @@ def _with_default_capture_dir(run: dict[str, Any], remote_bench_dir: Path) -> di return run +def _probe_host(node: int) -> str: + """How the batch script, which runs on the allocation's first node, reaches `node`.""" + return "localhost" if node == 0 else f"${{gym_nodes[{node}]}}" + + +def _health_check_host( + service: VllmServiceConfig | RayServiceConfig, + config: SubmitConfig, + compute: SlurmComputeConfig, + 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.""" + if isinstance(service, RayServiceConfig): + return _probe_host(_driver_node(config, compute)) + if service.node_pool is not None: + return _probe_host(_pool_offsets(compute)[service.node_pool][0]) + if driver_node is not None and not _vllm_spans_multiple_nodes(service, total_nodes): + return _probe_host(driver_node) + return "localhost" + + def build_sbatch_script( config: SubmitConfig, benchmark_name: str, @@ -439,39 +755,62 @@ def build_sbatch_script( pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None ] - ray_prelude = ( - render_ray_prelude() - if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values()) - else "" + ray_prelude = "\n".join( + block + for block in ( + render_ray_prelude() + if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values()) + else "", + _render_pool_nodes(config, compute, is_multi_node), + _render_ray_head_addresses(config, compute), + ) + if block ) observed = otel_active(config) + driver_node = _driver_node(config, compute) if is_multi_node else None service_commands = "\n\n".join( - ([_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node)] if observed else []) + ( + [_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node, driver_node=driver_node)] + if observed + else [] + ) + [ _render_service_command( name, service.container, - _build_service_command(service, total_nodes, gpus_per_node_values), + _build_service_command(service, _service_nodes(service, compute, total_nodes), gpus_per_node_values), service.env or None, service.mounts or None, - # Only services that actually span multiple nodes need the whole allocation's --nodes/ - # --ntasks - not every service in a multi-node job (e.g. a plain Ray head service runs - # on a single node regardless of how many nodes the overall job spans). - nodes=total_nodes if _vllm_spans_multiple_nodes(service, total_nodes) else None, - ntasks=total_ntasks if _vllm_spans_multiple_nodes(service, total_nodes) else None, + # Only services that actually span multiple nodes need --nodes/--ntasks - not every + # service in a multi-node job (e.g. a plain Ray head service runs on a single node + # regardless of how many nodes the overall job spans). A pinned service always gets + # them, so the node list and the step size agree. + nodes=_srun_nodes(service, compute, total_nodes), + ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks), pre_command=service.pre_command, + nodelist=_service_nodelist(service, driver_node, total_nodes), ) for name, service in config.services.items() + if not isinstance(service, RayServiceConfig) + ] + + [ + _render_ray_service(name, service, config, compute, driver_node) + for name, service in config.services.items() + if isinstance(service, RayServiceConfig) ] ) health_checks = "\n\n".join( - ([_render_collector_health_check(config)] if observed else []) + ([_render_collector_health_check(config, driver_node)] if observed else []) + [ render_health_check( - name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds + name, + service.health_check.port, + service.health_check.path, + service.health_check.timeout_seconds, + _health_check_host(service, config, compute, driver_node, total_nodes), ) for name, service in config.services.items() if service.health_check @@ -502,7 +841,9 @@ def build_sbatch_script( ) prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" - driver_node_flags = " --nodes=1 --ntasks=1" if is_multi_node else "" + driver_node_flags = ( + f' --nodelist="${{gym_nodes[{driver_node}]}}" --nodes=1 --ntasks=1' if driver_node is not None else "" + ) # The driver writes everything relative to the job directory -- `output_path` # above is `artifacts/rollouts.jsonl`. `#SBATCH --chdir` sets the cwd of the # BATCH script on the host, but inside a Pyxis container the cwd is whatever diff --git a/tests/unit_tests/test_orchestration_otel.py b/tests/unit_tests/test_orchestration_otel.py index 3254f2d2115..c1334c09940 100644 --- a/tests/unit_tests/test_orchestration_otel.py +++ b/tests/unit_tests/test_orchestration_otel.py @@ -309,6 +309,8 @@ def test_script_pins_the_collector_to_one_node_of_a_multi_node_job(): ) line = _collector_line(_script(multi)) assert " --nodes=1 --ntasks=1" in line + # Named, not left to Slurm: the same node the driver is placed on. + assert '--nodelist="${gym_nodes[0]}"' in line assert " --nodes=1 --ntasks=1" not in _collector_line(_script(_config())) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index ac78c03f865..74b0bfa245b 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -21,13 +21,14 @@ import pytest -from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.api import NodePool, RayServiceConfig, SubmitConfig from nemo_gym.orchestration.executors.script_templates import ( render_driver_entrypoint, render_gym_cmd, ) from nemo_gym.orchestration.executors.slurm_script import ( _RAY_SERVE_GATEWAY_SOURCE_PATH, + _build_ray_command, _build_service_command, _build_vllm_command, _build_vllm_multi_instance_multi_node_command, @@ -87,24 +88,46 @@ def test_deeply_nested(): def test_render_pool_directives_basic(pool): - lines = _render_pool_directives("main", pool) - assert "#SBATCH --partition=batch # pool: main" in lines + lines = _render_pool_directives({"main": pool}) + assert "#SBATCH --partition=batch" in lines assert "#SBATCH --nodes=1" in lines assert "#SBATCH --ntasks-per-node=4" in lines def test_render_pool_directives_gpus(pool): pool.gpus_per_node = 4 - lines = _render_pool_directives("main", pool) + lines = _render_pool_directives({"main": pool}) assert "#SBATCH --gpus-per-node=4" in lines def test_render_pool_directives_extra_args(pool): pool.extra_args["gres"] = "shard:8" - lines = _render_pool_directives("main", pool) + lines = _render_pool_directives({"main": pool}) assert "#SBATCH --gres=shard:8" in lines +def test_pools_ask_for_the_sum_of_their_nodes_once(pool): + # One #SBATCH --nodes for the allocation, not one per pool. Emitting them per + # pool made every pool but the last a no-op while the rest of the executor + # sized itself on the sum, so a two-pool job asked for one pool's nodes. + aux = NodePool(partition="batch", nodes=3, ntasks_per_node=4) + lines = _render_pool_directives({"main": pool, "aux": aux}) + assert [line for line in lines if line.startswith("#SBATCH --nodes")] == ["#SBATCH --nodes=4"] + + +def test_pools_that_disagree_on_a_whole_allocation_directive_are_refused(pool): + other = NodePool(partition="interactive", nodes=1, ntasks_per_node=4) + with pytest.raises(ValueError, match="disagree on partition"): + _render_pool_directives({"main": pool, "other": other}) + + +def test_pools_that_disagree_on_extra_args_are_refused(pool): + pool.extra_args["gres"] = "shard:8" + other = NodePool(partition="batch", nodes=1, ntasks_per_node=4, extra_args={"gres": "shard:4"}) + with pytest.raises(ValueError, match="conflicts with another pool"): + _render_pool_directives({"main": pool, "other": other}) + + # --------------------------------------------------------------------------- # _render_directives # --------------------------------------------------------------------------- @@ -1418,7 +1441,8 @@ def test_build_sbatch_script_non_vllm_service_omits_node_flags_in_multi_node_job vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) ray_line = next(line for line in script.splitlines() if "ray:latest" in line) assert "--nodes=4" in vllm_line - assert "--nodes=" not in ray_line + # Once, on the driver's node -- not once per node of the allocation. + assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in ray_line # --------------------------------------------------------------------------- @@ -1570,3 +1594,325 @@ def test_gym_install_runs_from_the_install_root(): assert entrypoint.index('cd "$GYM_SRC/gym"') < entrypoint.index('exec "$@"') # The only `cd` is into the clone -- nothing else may move cwd. assert entrypoint.count("cd ") == entrypoint.count('cd "$GYM_SRC/gym"') + + +# --------------------------------------------------------------------------- +# node_pool placement +# --------------------------------------------------------------------------- + + +def _placement_config(tmp_path, services, pools): + return SubmitConfig.model_validate( + { + "services": services, + "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": pools}}, + "driver": {"container": "gym:latest", "benchmarks": {"b": {"run": {}}}}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + + +def _vllm(port, pool, **extra): + return { + "type": "vllm", + "container": "img", + "model": "/ckpt", + "port": port, + "node_pool": pool, + **extra, + } + + +_TWO_POOLS = { + "gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}, + "aux": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}, +} + + +def _render(tmp_path, services, pools=None): + config = _placement_config(tmp_path, services, pools or _TWO_POOLS) + return build_sbatch_script( + config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "jobs" / "b" + ) + + +def test_pinned_services_take_contiguous_node_ranges(tmp_path): + # --nodelist names each pool's hosts; --relative is only a hint Slurm can move a + # step off (seen on a cluster run: the head landed on the worker's node). Without it both steps + # start at node 0 and the second one shares the first one's GPUs. + script = _render( + tmp_path, + { + "policy": _vllm(8000, "gpu", tensor_parallel_size=4), + "scorer": _vllm(8001, "aux", tensor_parallel_size=4), + }, + ) + assert 'export GYM_POOL_GPU_NODES="$(IFS=,; echo "${gym_nodes[*]:0:1}")"' in script + assert 'export GYM_POOL_AUX_NODES="$(IFS=,; echo "${gym_nodes[*]:1:1}")"' in script + assert '--nodelist="${GYM_POOL_GPU_NODES}" --nodes=1 --ntasks=1' in script.split("# service: policy")[1] + assert '--nodelist="${GYM_POOL_AUX_NODES}" --nodes=1 --ntasks=1' in script.split("# service: scorer")[1] + + +def test_a_pinned_single_node_service_is_not_built_as_multi_node(tmp_path): + # The job spans two nodes, but each service owns one. Sizing them by the + # allocation total would build both as multi-node Ray deployments. + script = _render( + tmp_path, + { + "policy": _vllm(8000, "gpu", tensor_parallel_size=4), + "scorer": _vllm(8001, "aux", tensor_parallel_size=4), + }, + ) + assert "ray start" not in script + assert script.count("vllm serve") == 2 + + +def test_an_unpinned_service_keeps_the_whole_allocation(tmp_path): + script = _render( + tmp_path, + {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}}, + ) + policy_line = script.split("# service: policy")[1].splitlines()[1] + assert "--nodelist" not in policy_line + assert "--nodes=2" in policy_line + + +def test_the_driver_runs_on_node_0_of_a_multi_node_job(tmp_path): + # A multi-node policy serves its API from node 0, and the driver reaches it on + # localhost. Left unpinned, Slurm may start the driver on any node. + script = _render( + tmp_path, + {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}}, + ) + driver_line = next(line for line in script.splitlines() if "logs/driver.log" in line) + assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in driver_line + assert 'gym_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))' in script + + +def test_a_single_node_job_places_nothing(tmp_path): + one = {"gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}} + script = _render( + tmp_path, + {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 4}}, + one, + ) + assert "--nodelist" not in script + assert "gym_nodes=" not in script + + +def test_an_unknown_node_pool_is_named(tmp_path): + with pytest.raises(ValueError, match="node_pool 'nope' does not match any node pool"): + _placement_config(tmp_path, {"policy": _vllm(8000, "nope", tensor_parallel_size=4)}, _TWO_POOLS) + + +# --------------------------------------------------------------------------- +# ray services +# --------------------------------------------------------------------------- + + +def test_a_ray_head_blocks_so_slurm_keeps_the_step_alive(): + # `ray start` daemonises and returns; without --block the srun step exits the + # moment the node is up and the service is torn down again. + command = _build_ray_command(RayServiceConfig(type="ray", container="img")) + assert command == "ray start --block --head --port 6379" + + +def test_a_ray_worker_joins_the_head_and_advertises_its_pools_resources(): + service = RayServiceConfig( + type="ray", container="img", node_pools=["aux"], num_gpus=0, resources={"aux": {"extra_gpu": 4}} + ) + command = _build_ray_command(service, pool="aux", address_var="GYM_RAY_ADDRESS_RAY", worker=True) + assert command == ( + 'ray start --block --address "$GYM_RAY_ADDRESS_RAY" --num-gpus 0 --resources=\'{"extra_gpu": 4}\'' + ) + + +def test_a_ray_worker_drops_the_flags_ray_allows_only_on_a_head(): + service = RayServiceConfig( + type="ray", container="img", extra_args="--include-dashboard=false --port 1 --node-manager-port=8366" + ) + command = _build_ray_command(service, address_var="A", worker=True) + assert command == 'ray start --block --address "$A" --node-manager-port=8366' + assert _build_ray_command(service).endswith("--include-dashboard=false --port 1 --node-manager-port=8366") + + +def test_the_old_head_and_worker_fields_are_gone(): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + RayServiceConfig(type="ray", container="img", mode="worker", address="10.0.0.1:6379") + + +def test_a_ray_service_names_node_pools_not_a_node_pool(): + with pytest.raises(ValueError, match=r"use node_pools: \[aux\]"): + RayServiceConfig(type="ray", container="img", node_pool="aux") + + +def test_ray_resources_for_a_pool_the_service_does_not_span_are_refused(): + with pytest.raises(ValueError, match="sets resources for aux, which it does not span"): + RayServiceConfig(type="ray", container="img", node_pools=["gpu"], resources={"aux": {"extra_gpu": 4}}) + + +def test_an_unknown_ray_node_pool_is_named(tmp_path): + with pytest.raises(ValueError, match=r"node_pools \['nope'\] do not match any node pool"): + _placement_config(tmp_path, {"ray": {"type": "ray", "container": "img", "node_pools": ["nope"]}}, _TWO_POOLS) + + +def _comet_like(tmp_path, pools=None, **ray): + return _render( + tmp_path, + { + "policy": _vllm(8000, "gpu", tensor_parallel_size=4), + "ray": { + "type": "ray", + "container": "img", + "node_pools": ["gpu", "aux"], + "port": 6380, + "resources": {"aux": {"extra_gpu": 4}}, + **ray, + }, + }, + pools, + ) + + +def test_a_ray_service_starts_its_head_beside_the_driver(tmp_path): + script = _comet_like(tmp_path) + assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[0]}' in script + assert ':6380"' in script + head = script.split("# service: ray\n")[1].split("\n\n")[0] + assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in head + assert '--head --port 6380 --node-ip-address "${GYM_RAY_ADDRESS_RAY%:*}"' in head + # The head's own pool advertises nothing here, so it offers no extra_gpu. + assert "extra_gpu" not in head + + +def test_every_other_spanned_node_joins_as_a_worker(tmp_path): + pools = {**_TWO_POOLS, "aux": {**_TWO_POOLS["aux"], "nodes": 2}} + script = _comet_like(tmp_path, pools) + assert 'export GYM_RAY_RAY_AUX_WORKERS="$(IFS=,; echo "${gym_nodes[*]:1:2}")"' in script + # The gpu pool's only node is the head, so it gets no worker step. + assert "GYM_RAY_RAY_GPU_WORKERS" not in script + workers = script.split("# service: ray_aux_workers\n")[1].split("\n\n")[0] + assert '--nodelist="${GYM_RAY_RAY_AUX_WORKERS}" --nodes=2 --ntasks=2' in workers + assert '--address "$GYM_RAY_ADDRESS_RAY"' in workers + assert "extra_gpu" in workers + + +def test_a_ray_service_runs_every_step_in_its_one_container(tmp_path): + script = _comet_like(tmp_path, container="ray-img:1", mounts=["/x:/x"], pre_command="setup") + for name in ("ray", "ray_aux_workers"): + step = script.split(f"# service: {name}\n")[1].split("\n\n")[0] + assert "--container-image=ray-img:1" in step + assert "--container-mounts=/x:/x" in step + assert "setup" in step + + +def test_a_worker_waits_for_the_head_after_its_own_setup(tmp_path): + script = _comet_like(tmp_path, pre_command="install-ray") + workers = script.split("# service: ray_aux_workers\n")[1].split("\n\n")[0] + assert workers.index("install-ray") < workers.index("until ray status") < workers.index("exec ray start") + + +def test_the_ray_head_follows_the_driver_to_a_later_pool(tmp_path): + config = SubmitConfig.model_validate( + { + "services": { + "policy": _vllm(8000, "aux", tensor_parallel_size=4), + "ray": {"type": "ray", "container": "img", "node_pools": ["gpu", "aux"]}, + }, + "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}}, + "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b") + assert '--nodelist="${gym_nodes[1]}"' in script.split("# service: ray\n")[1].split("\n\n")[0] + assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[1]}' in script + assert 'export GYM_RAY_RAY_GPU_WORKERS="$(IFS=,; echo "${gym_nodes[*]:0:1}")"' in script + assert "GYM_RAY_RAY_AUX_WORKERS" not in script + + +def test_a_single_node_ray_service_is_just_a_head(tmp_path): + one = {"gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}} + script = _render(tmp_path, {"ray": {"type": "ray", "container": "img", "node_pools": ["gpu"]}}, one) + assert "_workers" not in script + assert "--nodelist" not in script + assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[0]}' in script + + +def test_a_rendered_ray_service_is_valid_bash(tmp_path): + script = _comet_like(tmp_path, pre_command="echo 'quoted'", extra_args="--temp-dir=/tmp/ray-$SLURM_JOB_ID") + result = subprocess.run(["bash", "-n"], input=script, text=True, capture_output=True) + assert result.returncode == 0, result.stderr + + +def test_the_driver_goes_to_the_policys_node_when_services_are_pinned(tmp_path): + # The driver reaches the policy on localhost. Unpinned, Slurm put it on the aux + # node of a real two-node job, where neither the policy nor a local raylet was. + config = SubmitConfig.model_validate( + { + "services": { + "scorer": _vllm(8001, "gpu", tensor_parallel_size=4), + "policy": _vllm(8000, "aux", tensor_parallel_size=4), + }, + "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}}, + "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b") + driver = next(line for line in script.splitlines() if "--output=logs/driver.log" in line) + assert '--nodelist="${gym_nodes[1]}" --nodes=1 --ntasks=1' in driver + + +# --------------------------------------------------------------------------- +# health probes +# --------------------------------------------------------------------------- + + +def _driver_on_aux(tmp_path, services): + config = SubmitConfig.model_validate( + { + "services": {"policy": _vllm(8000, "aux", tensor_parallel_size=4), **services}, + "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}}, + "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}}, + "job": {"output_path": str(tmp_path / "jobs")}, + } + ) + return build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b") + + +_RAY_WITH_PROBE = { + "type": "ray", + "container": "img", + "node_pools": ["gpu", "aux"], + "health_check": {"port": 8011, "path": "/"}, +} + + +def test_a_pinned_service_on_a_later_pool_is_probed_where_it_runs(tmp_path): + # The probe runs on node 0; the aux pool's service answers on node 1. + script = _render(tmp_path, {"policy": _vllm(8000, "gpu"), "scorer": _vllm(8001, "aux")}) + assert "Waiting for scorer at http://${gym_nodes[1]}:8001" in script + assert "Waiting for policy at http://localhost:8000" in script + + +def test_a_ray_head_is_probed_beside_a_non_zero_driver_node(tmp_path): + script = _driver_on_aux(tmp_path, {"ray": _RAY_WITH_PROBE}) + assert "Waiting for ray at http://${gym_nodes[1]}:8011" in script + + +def test_the_collector_is_probed_beside_a_non_zero_driver_node(tmp_path): + script = _driver_on_aux(tmp_path, {}) + assert "Waiting for otel_collector at http://${gym_nodes[1]}:13133" in script + + +def test_services_on_node_0_are_probed_locally(tmp_path): + script = _render(tmp_path, {"policy": _vllm(8000, "gpu", tensor_parallel_size=4), "ray": _RAY_WITH_PROBE}) + assert "Waiting for ray at http://localhost:8011" in script + assert "Waiting for policy at http://localhost:8000" in script + + +def test_an_unpinned_multi_node_service_is_probed_on_node_0(tmp_path): + # It spans the allocation and serves its API from node 0, where the probe runs. + script = _driver_on_aux(tmp_path, {"judge": {"type": "vllm", "container": "img", "model": "/j", "port": 9000}}) + assert "Waiting for judge at http://localhost:9000" in script