diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 264bfdd4a..2f425b3f7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -1,120 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed one-node step manifest produced inside the sealed client image.""" +"""Typed allocation step manifest produced inside the sealed client image.""" from __future__ import annotations -import os from collections.abc import Mapping from pathlib import Path -from typing import Literal - -from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator from data_designer.slurm.config.environment import ( - LiteralEnvironmentBinding, SecretRef, collect_secret_environment_names, ) -from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path -from data_designer.slurm.runtime.backpressure import ( - MAX_WAITING_REQUESTS_ENVIRONMENT, - RETRY_AFTER_SECONDS_ENVIRONMENT, -) -from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.manifest import RuntimeBootstrapManifest, RuntimeProbeSpec, RuntimeStepSpec from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStepRole from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_deployments -from data_designer.slurm.runtime.steps import ( - build_client_command, - build_endpoint_command, - build_vllm_command, -) +from data_designer.slurm.runtime.preflight import AllocationLayout, validate_allocation_layout +from data_designer.slurm.runtime.server_manifest import build_server_steps +from data_designer.slurm.runtime.steps import build_client_command, build_endpoint_command from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.vllm import ResolvedVllmProcess -from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest - - -class RuntimeProbeSpec(ContractValue): - """One loopback readiness target monitored by the Bash controller.""" - - host: Literal["127.0.0.1"] = "127.0.0.1" - port: NetworkPort - path: str - deadline_seconds: PositiveInt - - @field_validator("path") - @classmethod - def validate_path(cls, value: str) -> str: - if not value.startswith("/") or any(ord(character) < 32 or ord(character) == 127 for character in value): - raise ValueError("runtime probe path is invalid") - return value - - -class RuntimeStepSpec(ContractValue): - """Container command and environment consumed by the Bash step runner.""" - - step_id: Identifier - role: RuntimeStepRole - image_path: str - command: tuple[str, ...] = Field(min_length=1) - cpus: PositiveInt - gpu_indices: tuple[NonNegativeInt, ...] = () - literal_environment: dict[EnvironmentName, str] = Field(default_factory=dict) - secret_environment: dict[EnvironmentName, EnvironmentName] = Field(default_factory=dict) - environment_prefixes: dict[EnvironmentName, str] = Field(default_factory=dict) - container_environment: tuple[EnvironmentName, ...] = () - stdout_path: str - stderr_path: str - launch_delay_seconds: NonNegativeInt = 0 - readiness: RuntimeProbeSpec | None = None - - _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) - _stdout_path_is_absolute = field_validator("stdout_path")(validate_absolute_path) - _stderr_path_is_absolute = field_validator("stderr_path")(validate_absolute_path) - - @model_validator(mode="after") - def validate_step(self) -> RuntimeStepSpec: - if any(not argument or "\0" in argument for argument in self.command): - raise ValueError("runtime command is invalid") - if self.gpu_indices != tuple(sorted(set(self.gpu_indices))): - raise ValueError("runtime GPU indices must be sorted and unique") - if self.stdout_path == self.stderr_path or Path(self.stdout_path).parent != Path(self.stderr_path).parent: - raise ValueError("runtime log paths must be distinct siblings") - if set(self.environment_prefixes) - (set(self.literal_environment) | set(self.secret_environment)): - raise ValueError("environment prefixes require a materialized variable") - container_names = set(self.container_environment) - if container_names - (set(self.literal_environment) | set(self.secret_environment)): - raise ValueError("container environment contains an unavailable variable") - if self.role is RuntimeStepRole.SERVER and not self.gpu_indices: - raise ValueError("server runtime steps require GPUs") - if self.role is not RuntimeStepRole.SERVER and self.gpu_indices: - raise ValueError("non-server runtime steps cannot request GPUs") - return self - - -class RuntimeBootstrapManifest(ContractRecord): - """Secret-free one-node allocation command manifest.""" - - run_id: Identifier - shard_id: Identifier - attempt_id: Identifier - plan_sha256: Sha256Digest - all_secret_environment_names: tuple[EnvironmentName, ...] - steps: tuple[RuntimeStepSpec, ...] = Field(min_length=4) - - @model_validator(mode="after") - def validate_steps(self) -> RuntimeBootstrapManifest: - step_ids = tuple(step.step_id for step in self.steps) - if len(step_ids) != len(set(step_ids)): - raise ValueError("runtime step identifiers must be unique") - roles = tuple(step.role for step in self.steps) - if roles.count(RuntimeStepRole.CLIENT_PREFLIGHT) != 1 or roles.count(RuntimeStepRole.CLIENT) != 1: - raise ValueError("runtime manifest requires one preflight and generation step") - if RuntimeStepRole.SERVER not in roles or RuntimeStepRole.ENDPOINT not in roles: - raise ValueError("runtime manifest requires server and endpoint steps") - return self def build_runtime_manifest( @@ -123,9 +28,11 @@ def build_runtime_manifest( *, runtime_root: Path, log_directory: Path, + layout: AllocationLayout, ) -> RuntimeBootstrapManifest: - """Build the secret-free one-node command handoff for the Bash controller.""" + """Build the secret-free command handoff for the Bash controller.""" plan = context.plan + validate_allocation_layout(plan, layout) runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) deployments = resolve_allocation_deployments(context, environment) endpoints = tuple( @@ -147,21 +54,13 @@ def build_runtime_manifest( endpoints, runtime_container_root, log_directory, + layout, ) ] - for deployment in deployments: - steps.extend( - _build_server_step( - deployment, - process, - context, - runtime_root, - runtime_container_root, - log_directory, - ) - for process in deployment.processes - ) - steps.extend(_build_endpoint_step(deployment, context, runtime_root, log_directory) for deployment in deployments) + steps.extend(build_server_steps(deployments, context, runtime_container_root, log_directory, layout)) + steps.extend( + _build_endpoint_step(deployment, context, runtime_root, log_directory, layout) for deployment in deployments + ) steps.append( _build_client_step( RuntimeStepRole.CLIENT, @@ -172,6 +71,7 @@ def build_runtime_manifest( endpoints, runtime_container_root, log_directory, + layout, ) ) secret_names = set(collect_secret_environment_names(plan)) @@ -201,6 +101,7 @@ def _build_client_step( endpoints: tuple[RuntimeEndpoint, ...], runtime_container_root: str, log_directory: Path, + layout: AllocationLayout, ) -> RuntimeStepSpec: plan = context.plan retry_resume_mode = None if context.retry_plan is None else context.retry_plan.effective_resume_mode @@ -261,70 +162,7 @@ def _build_client_step( environment_prefixes={}, container_environment=tuple(sorted((*secret_names, *allocation_environment, "PYTHONPATH"))), log_directory=log_directory, - ) - - -def _build_server_step( - deployment: ResolvedVllmServerDeployment, - process: ResolvedVllmProcess, - context: AllocationContext, - runtime_root: Path, - runtime_container_root: str, - log_directory: Path, -) -> RuntimeStepSpec: - if process.pipeline_parallel != 1 or process.node_index != 0 or process.http_port is None: - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.INVALID_CONTEXT, - "one-node runtime received a distributed vLLM process", - ) - literal_environment: dict[str, str] = {"LC_ALL": "C", "PYTHONPATH": runtime_container_root} - secret_environment: dict[str, str] = {} - environment_prefixes: dict[str, str] = {} - for name, binding in deployment.launch_policy.environment.items(): - if isinstance(binding, LiteralEnvironmentBinding): - literal_environment[name] = binding.value - elif isinstance(binding, SecretRef): - secret_environment[name] = binding.environment - else: # pragma: no cover - persisted contracts reject unknown bindings - raise AssertionError(f"unhandled environment binding: {type(binding)!r}") - if "PYTHONPATH" in secret_environment: - literal_environment.pop("PYTHONPATH") - environment_prefixes["PYTHONPATH"] = runtime_container_root - elif "PYTHONPATH" in deployment.launch_policy.environment: - literal_environment["PYTHONPATH"] = os.pathsep.join((runtime_container_root, literal_environment["PYTHONPATH"])) - policy = deployment.launch_policy.queue_backpressure - literal_environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(policy.max_waiting_requests) - literal_environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( - "" if policy.retry_after_seconds is None else str(policy.retry_after_seconds) - ) - probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) - return _step( - step_id=process.process_id, - role=RuntimeStepRole.SERVER, - image_path=deployment.image.path, - command=build_vllm_command(deployment, process, context.plan), - cpus=context.plan.client.authored.cpus, - gpu_indices=tuple(process.gpu_indices), - literal_environment=literal_environment, - secret_environment=secret_environment, - environment_prefixes=environment_prefixes, - container_environment=tuple( - sorted( - { - *deployment.launch_policy.environment, - "PYTHONPATH", - MAX_WAITING_REQUESTS_ENVIRONMENT, - RETRY_AFTER_SECONDS_ENVIRONMENT, - } - ) - ), - log_directory=log_directory, - launch_delay_seconds=process.launch_delay_seconds, - readiness=RuntimeProbeSpec( - port=probe.port, - path=probe.path, - deadline_seconds=probe.deadline_seconds, - ), + node_hosts=(layout.get_host(plan.client.host_node_index),), ) @@ -333,13 +171,20 @@ def _build_endpoint_step( context: AllocationContext, runtime_root: Path, log_directory: Path, + layout: AllocationLayout, ) -> RuntimeStepSpec: proxy_path = runtime_root / "data_designer/slurm/runtime/proxy.py" + backend_hosts = ( + tuple(layout.get_host(backend.node_index) for backend in deployment.backend_endpoints) + if len(layout.node_hosts) > 1 + else None + ) command = build_endpoint_command( deployment, context.plan, proxy_path, deployment.logical_endpoint.port, + backend_hosts=backend_hosts, ) return _step( step_id=f"{deployment.deployment_id}-endpoint", @@ -353,10 +198,14 @@ def _build_endpoint_step( environment_prefixes={}, container_environment=(), log_directory=log_directory, - readiness=RuntimeProbeSpec( - port=deployment.logical_endpoint.port, - path="/health", - deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + node_hosts=(layout.get_host(context.plan.client.host_node_index),), + readiness=( + RuntimeProbeSpec( + host="127.0.0.1", + port=deployment.logical_endpoint.port, + path="/health", + deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + ), ), ) @@ -369,13 +218,15 @@ def _step( command: tuple[str, ...], cpus: int, gpu_indices: tuple[int, ...], + node_hosts: tuple[str, ...], literal_environment: dict[str, str], secret_environment: dict[str, str], environment_prefixes: dict[str, str], container_environment: tuple[str, ...], log_directory: Path, + kill_on_bad_exit: bool = False, launch_delay_seconds: int = 0, - readiness: RuntimeProbeSpec | None = None, + readiness: tuple[RuntimeProbeSpec, ...] = (), ) -> RuntimeStepSpec: return RuntimeStepSpec( step_id=step_id, @@ -384,6 +235,8 @@ def _step( command=command, cpus=cpus, gpu_indices=gpu_indices, + node_hosts=node_hosts, + kill_on_bad_exit=kill_on_bad_exit, literal_environment=literal_environment, secret_environment=secret_environment, environment_prefixes=environment_prefixes, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py new file mode 100644 index 000000000..88236233f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose resolved deployments into node-worker specifications.""" + +from __future__ import annotations + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.node_spec import NodeProcessSpec, NodeSpec, NodeWorkerSpec +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess + + +def build_node_worker_spec( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + layout: AllocationLayout, +) -> NodeWorkerSpec: + """Build the validated work assigned to each node in one deployment.""" + model = _resolve_model(deployment, plan) + nodes = tuple( + NodeSpec( + node_index=node_index, + host=layout.get_host(node_index), + ports=_get_node_ports(deployment, node_index), + processes=tuple( + NodeProcessSpec( + process_id=process.process_id, + command=build_vllm_process_command(deployment, process, plan, layout), + gpu_indices=tuple(process.gpu_indices), + launch_delay_seconds=process.launch_delay_seconds, + ) + for process in deployment.processes + if process.node_index == node_index + ), + ) + for node_index in deployment.node_indices + ) + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=deployment.gpus_per_node, + required_model_path=model if deployment.model.startswith("/") else None, + nodes=nodes, + ) + + +def build_vllm_process_command( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + plan: ResolvedSlurmRunPlan, + layout: AllocationLayout, +) -> tuple[str, ...]: + """Build one shell-free vLLM lane command at its resolved host placement.""" + backend = deployment.backend_endpoints[process.deployment_replica_index] + model = _resolve_model(deployment, plan) + command: tuple[str, ...] = ( + deployment.executable_path, + "serve", + model, + "--served-model-name", + deployment.served_model_name, + "--host", + "0.0.0.0", + "--port", + str(backend.port), + "--tensor-parallel-size", + str(process.tensor_parallel), + "--distributed-executor-backend", + "uni" if process.tensor_parallel * process.pipeline_parallel == 1 else "mp", + "--data-parallel-backend", + "mp", + "--middleware", + "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware", + *_distributed_arguments(process, layout), + ) + if deployment.launch_policy.enable_expert_parallel: + command += ("--enable-expert-parallel",) + return command + deployment.launch_policy.extra_args + + +def _resolve_model(deployment: ResolvedVllmServerDeployment, plan: ResolvedSlurmRunPlan) -> str: + return get_container_path(plan, deployment.model) if deployment.model.startswith("/") else deployment.model + + +def _distributed_arguments( + process: ResolvedVllmProcess, + layout: AllocationLayout, +) -> tuple[str, ...]: + if process.pipeline_parallel == 1: + return () + rendezvous = process.rendezvous + if rendezvous is None: # pragma: no cover - resolved contracts enforce this + raise AssertionError("distributed process has no rendezvous") + arguments = ( + "--pipeline-parallel-size", + str(process.pipeline_parallel), + "--nnodes", + str(process.pipeline_parallel), + "--node-rank", + str(process.pipeline_rank), + "--master-addr", + layout.get_host(rendezvous.master_node_index), + "--master-port", + str(rendezvous.port), + "--distributed-timeout-seconds", + str(rendezvous.timeout_seconds), + ) + return arguments + (("--headless",) if process.pipeline_rank > 0 else ()) + + +def _get_node_ports(deployment: ResolvedVllmServerDeployment, node_index: int) -> tuple[int, ...]: + http_ports = tuple(endpoint.port for endpoint in deployment.backend_endpoints if endpoint.node_index == node_index) + rendezvous_ports = tuple( + process.rendezvous.port + for process in deployment.processes + if process.node_index == node_index and process.pipeline_rank == 0 and process.rendezvous is not None + ) + return tuple(sorted((*http_ports, *rendezvous_ports))) + + +__all__ = ["build_node_worker_spec", "build_vllm_process_command"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index dd43a06f2..28a0ea265 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -21,7 +21,11 @@ from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_deployments -from data_designer.slurm.runtime.preflight import SystemAllocationPreflight +from data_designer.slurm.runtime.preflight import ( + AllocationLayout, + SystemAllocationPreflight, + validate_allocation_layout, +) from data_designer.slurm.runtime.records import load_complete_client_candidate from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment from data_designer.slurm.state import ( @@ -76,6 +80,7 @@ def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: prepare = subparsers.choices["prepare"] prepare.add_argument("--runtime-root", required=True, type=Path) prepare.add_argument("--manifest", required=True, type=Path) + prepare.add_argument("--node-host", action="append", required=True) client = subparsers.add_parser("client") _add_context_arguments(client) client.add_argument("--endpoint", action="append", default=[]) @@ -93,6 +98,8 @@ def _add_context_arguments(parser: argparse.ArgumentParser) -> None: def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: context, writer = _load_context(arguments, environment) _validate_attempt_is_executable(context.attempt) + layout = AllocationLayout(tuple(arguments.node_host)) + validate_allocation_layout(context.plan, layout) SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) SystemAllocationPreflight.verify_ports(context, environment) readiness = _begin_attempt(context, writer, environment) @@ -104,6 +111,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N environment, runtime_root=arguments.runtime_root, log_directory=log_directory, + layout=layout, ) expected_manifest = context.attempt_directory / "runtime-manifest.json" if arguments.manifest.as_posix() != get_container_path( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh index cbba85438..4015b4f34 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -47,9 +47,15 @@ dd_slurm_run_allocation() { trap 'exit 130' INT TERM DD_RUNTIME_PREPARED=1 + local host + local -a host_arguments=() + for host in "${DD_ALLOCATION_HOSTS[@]}"; do + host_arguments+=(--node-host "${host}") + done dd_run_bound_control_phase prepare \ --runtime-root "${DD_RUNTIME_DIR}" \ - --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" + --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" \ + "${host_arguments[@]}" dd_verify_runtime_manifest \ "${DD_RUNTIME_MANIFEST}" \ "${DD_PLAN_SHA256}" \ @@ -61,6 +67,7 @@ dd_slurm_run_allocation() { ((${#DD_STEP_IDS[@]} == 1)) dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + dd_run_role_steps server_preflight dd_start_servers dd_wait_for_role_readiness server dd_start_endpoints @@ -100,7 +107,11 @@ dd_verify_host_context() { done [[ -d ${DD_ATTEMPT_PATH} && ! -L ${DD_ATTEMPT_PATH} ]] [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] - [[ ${SLURM_JOB_NUM_NODES:-} == 1 && ${SLURM_NODEID:-} == 0 ]] + [[ ${DD_EXPECTED_NODES} =~ ^[1-9][0-9]*$ ]] + [[ ${DD_CLIENT_NODE_INDEX} =~ ^[0-9]+$ && ${DD_CLIENT_NODE_INDEX} -lt ${DD_EXPECTED_NODES} ]] + [[ ${SLURM_JOB_NUM_NODES:-} == "${DD_EXPECTED_NODES}" ]] + [[ ${SLURM_NODEID:-} == "${DD_CLIENT_NODE_INDEX}" ]] + dd_resolve_allocation_hosts dd_verify_gpu_count dd_read_artifacts "${DD_PLAN_PATH}" "${SLURM_ARRAY_TASK_ID}" local index path digest actual @@ -115,6 +126,24 @@ dd_verify_host_context() { done } +dd_resolve_allocation_hosts() { + local node_list=${SLURM_JOB_NODELIST:-} + [[ -n ${node_list} && ${node_list} != -* && ${#node_list} -le 4096 ]] || return 65 + [[ ${node_list} != *[$'\t\r\n ']* ]] || return 65 + local expanded + expanded=$(scontrol show hostnames "${node_list}") || return 65 + DD_ALLOCATION_HOSTS=() + local host + while IFS= read -r host; do + DD_ALLOCATION_HOSTS+=("${host}") + done <<<"${expanded}" + ((${#DD_ALLOCATION_HOSTS[@]} == DD_EXPECTED_NODES)) || return 65 + for host in "${DD_ALLOCATION_HOSTS[@]}"; do + [[ ${host} =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ && ${#host} -le 253 ]] || return 65 + done + DD_CLIENT_HOST=${DD_ALLOCATION_HOSTS[DD_CLIENT_NODE_INDEX]} +} + dd_verify_gpu_count() { local visible=${CUDA_VISIBLE_DEVICES:-${SLURM_JOB_GPUS:-}} local count=0 item @@ -156,6 +185,15 @@ dd_start_servers() { done } +dd_run_role_steps() { + local role=$1 + local step_id + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" "${role}" + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_run_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + done +} + dd_start_endpoints() { dd_read_step_ids "${DD_RUNTIME_MANIFEST}" endpoint local step_id @@ -168,19 +206,25 @@ dd_start_endpoints() { dd_wait_for_role_readiness() { local role=$1 - local step_id deadline + local step_id index host port path deadline_seconds deadline dd_read_step_ids "${DD_RUNTIME_MANIFEST}" "${role}" for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" - deadline=$((SECONDS + DD_STEP_PROBE_DEADLINE)) - until curl --fail --silent --max-time 1 \ - "http://${DD_STEP_PROBE_HOST}:${DD_STEP_PROBE_PORT}${DD_STEP_PROBE_PATH}" >/dev/null 2>&1; do - dd_require_running - ((SECONDS < deadline)) || { - printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 - return 70 - } - dd_sleep 0.5 + ((${#DD_STEP_PROBE_FIELDS[@]} % 4 == 0 && ${#DD_STEP_PROBE_FIELDS[@]} > 0)) || return 65 + for ((index = 0; index < ${#DD_STEP_PROBE_FIELDS[@]}; index += 4)); do + host=${DD_STEP_PROBE_FIELDS[index]} + port=${DD_STEP_PROBE_FIELDS[index + 1]} + path=${DD_STEP_PROBE_FIELDS[index + 2]} + deadline_seconds=${DD_STEP_PROBE_FIELDS[index + 3]} + deadline=$((SECONDS + deadline_seconds)) + until curl --fail --silent --max-time 1 "http://${host}:${port}${path}" >/dev/null 2>&1; do + dd_require_running + ((SECONDS < deadline)) || { + printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 + return 70 + } + dd_sleep 0.5 + done done done } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py new file mode 100644 index 000000000..728c7d8b9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated command manifest consumed by the allocation shell runtime.""" + +from __future__ import annotations + +from pathlib import Path + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path +from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.runtime.network import validate_host_name +from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest + + +class RuntimeProbeSpec(ContractValue): + """One readiness target monitored from the allocation client host.""" + + host: str + port: NetworkPort + path: str + deadline_seconds: PositiveInt + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + if not value.startswith("/") or any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("runtime probe path is invalid") + return value + + @field_validator("host") + @classmethod + def validate_host(cls, value: str) -> str: + return validate_host_name(value) + + +class RuntimeStepSpec(ContractValue): + """Container command and placement consumed by the Bash step runner.""" + + step_id: Identifier + role: RuntimeStepRole + image_path: str + command: tuple[str, ...] = Field(min_length=1) + cpus: PositiveInt + gpu_indices: tuple[NonNegativeInt, ...] = () + node_hosts: tuple[str, ...] = Field(min_length=1) + kill_on_bad_exit: bool = False + literal_environment: dict[EnvironmentName, str] = Field(default_factory=dict) + secret_environment: dict[EnvironmentName, EnvironmentName] = Field(default_factory=dict) + environment_prefixes: dict[EnvironmentName, str] = Field(default_factory=dict) + container_environment: tuple[EnvironmentName, ...] = () + stdout_path: str + stderr_path: str + launch_delay_seconds: NonNegativeInt = 0 + readiness: tuple[RuntimeProbeSpec, ...] = () + + _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) + _stdout_path_is_absolute = field_validator("stdout_path")(validate_absolute_path) + _stderr_path_is_absolute = field_validator("stderr_path")(validate_absolute_path) + + @model_validator(mode="after") + def validate_step(self) -> RuntimeStepSpec: + if any(not argument or "\0" in argument for argument in self.command): + raise ValueError("runtime command is invalid") + if self.gpu_indices != tuple(sorted(set(self.gpu_indices))): + raise ValueError("runtime GPU indices must be sorted and unique") + if self.node_hosts != tuple(dict.fromkeys(self.node_hosts)): + raise ValueError("runtime node hosts must be unique") + for host in self.node_hosts: + validate_host_name(host) + if self.stdout_path == self.stderr_path or Path(self.stdout_path).parent != Path(self.stderr_path).parent: + raise ValueError("runtime log paths must be distinct siblings") + if set(self.environment_prefixes) - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("environment prefixes require a materialized variable") + container_names = set(self.container_environment) + if container_names - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("container environment contains an unavailable variable") + self._validate_placement() + self._validate_readiness() + return self + + def _validate_placement(self) -> None: + server_roles = {RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER} + if self.role in server_roles and not self.gpu_indices: + raise ValueError("server runtime steps require GPUs") + if self.role not in server_roles and self.gpu_indices: + raise ValueError("non-server runtime steps cannot request GPUs") + if len(self.node_hosts) > 1 and self.role not in server_roles: + raise ValueError("only server runtime steps may span nodes") + if self.kill_on_bad_exit and self.role not in server_roles: + raise ValueError("kill-on-bad-exit is only valid for server runtime steps") + if len(self.node_hosts) > 1 and not self.kill_on_bad_exit: + raise ValueError("multi-node server runtime steps require kill-on-bad-exit") + + def _validate_readiness(self) -> None: + if self.role is RuntimeStepRole.SERVER_PREFLIGHT and self.readiness: + raise ValueError("server preflight steps cannot have readiness probes") + if self.role in {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} and not self.readiness: + raise ValueError("long-running runtime steps require readiness probes") + if self.role in {RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.CLIENT} and self.readiness: + raise ValueError("client runtime steps cannot have readiness probes") + + +class RuntimeBootstrapManifest(ContractRecord): + """Secret-free allocation command manifest.""" + + run_id: Identifier + shard_id: Identifier + attempt_id: Identifier + plan_sha256: Sha256Digest + all_secret_environment_names: tuple[EnvironmentName, ...] + steps: tuple[RuntimeStepSpec, ...] = Field(min_length=4) + + @model_validator(mode="after") + def validate_steps(self) -> RuntimeBootstrapManifest: + step_ids = tuple(step.step_id for step in self.steps) + if len(step_ids) != len(set(step_ids)): + raise ValueError("runtime step identifiers must be unique") + roles = tuple(step.role for step in self.steps) + if roles.count(RuntimeStepRole.CLIENT_PREFLIGHT) != 1 or roles.count(RuntimeStepRole.CLIENT) != 1: + raise ValueError("runtime manifest requires one preflight and generation step") + if RuntimeStepRole.SERVER not in roles or RuntimeStepRole.ENDPOINT not in roles: + raise ValueError("runtime manifest requires server and endpoint steps") + return self + + +__all__ = ["RuntimeBootstrapManifest", "RuntimeProbeSpec", "RuntimeStepSpec"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py index 2aa93ab16..c91370b38 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -25,6 +25,7 @@ class RuntimeStepRole(str, Enum): """Lifecycle role of one allocation-local Slurm step.""" CLIENT_PREFLIGHT = "client_preflight" + SERVER_PREFLIGHT = "server_preflight" SERVER = "server" ENDPOINT = "endpoint" CLIENT = "client" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py new file mode 100644 index 000000000..d17e4311a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared validation for scheduler-derived runtime network identities.""" + +from __future__ import annotations + +import re + +_HOST_NAME_PATTERN = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,251}[A-Za-z0-9])?$") + + +def validate_host_name(host: str) -> str: + """Return a safe scheduler host name or raise ``ValueError``.""" + if type(host) is not str or _HOST_NAME_PATTERN.fullmatch(host) is None: + raise ValueError("allocation host identity is invalid") + return host + + +def validate_network_port(port: int) -> int: + """Return a valid TCP port or raise ``ValueError``.""" + if type(port) is not int or not 1 <= port <= 65535: + raise ValueError("allocation network port is invalid") + return port + + +__all__ = ["validate_host_name", "validate_network_port"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py new file mode 100644 index 000000000..b225328dd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict transport records for one coordinated deployment's node workers.""" + +from __future__ import annotations + +import base64 +import binascii +import json +from dataclasses import dataclass + +from data_designer.slurm.contracts import validate_absolute_path +from data_designer.slurm.runtime.network import validate_host_name, validate_network_port + +_MAXIMUM_SPEC_BYTES = 64 * 1024 + + +@dataclass(frozen=True, slots=True) +class NodeProcessSpec: + """One resolved lane process owned by a node worker.""" + + process_id: str + command: tuple[str, ...] + gpu_indices: tuple[int, ...] + launch_delay_seconds: int + + def __post_init__(self) -> None: + if ( + type(self.process_id) is not str + or not self.process_id + or any(ord(value) < 32 or ord(value) == 127 for value in self.process_id) + ): + raise ValueError("node process identity is invalid") + if ( + type(self.command) is not tuple + or not self.command + or any(type(value) is not str or not value or "\0" in value for value in self.command) + ): + raise ValueError("node process command is invalid") + if ( + type(self.gpu_indices) is not tuple + or not self.gpu_indices + or self.gpu_indices != tuple(sorted(set(self.gpu_indices))) + or any(type(value) is not int or value < 0 for value in self.gpu_indices) + ): + raise ValueError("node process GPU indices are invalid") + if type(self.launch_delay_seconds) is not int or self.launch_delay_seconds < 0: + raise ValueError("node process launch delay is invalid") + + +@dataclass(frozen=True, slots=True) +class NodeSpec: + """Resolved work and ports for one node in a coordinated step.""" + + node_index: int + host: str + ports: tuple[int, ...] + processes: tuple[NodeProcessSpec, ...] + + def __post_init__(self) -> None: + if type(self.node_index) is not int or self.node_index < 0: + raise ValueError("node index is invalid") + validate_host_name(self.host) + if type(self.ports) is not tuple or type(self.processes) is not tuple: + raise ValueError("node ports and processes must be tuples") + for port in self.ports: + validate_network_port(port) + if self.ports != tuple(sorted(set(self.ports))): + raise ValueError("node ports must be sorted and unique") + if not self.processes or len({process.process_id for process in self.processes}) != len(self.processes): + raise ValueError("node process identities must be present and unique") + + +@dataclass(frozen=True, slots=True) +class NodeWorkerSpec: + """Strict, bounded input shared by every task in one deployment step.""" + + schema_version: int + resolved_gpus_per_node: int + required_model_path: str | None + nodes: tuple[NodeSpec, ...] + + def __post_init__(self) -> None: + if type(self.schema_version) is not int or self.schema_version != 1: + raise ValueError("node worker schema version is unsupported") + if type(self.resolved_gpus_per_node) is not int or self.resolved_gpus_per_node <= 0: + raise ValueError("resolved GPU count is invalid") + if self.required_model_path is not None and type(self.required_model_path) is not str: + raise ValueError("required model path is invalid") + if self.required_model_path is not None: + validate_absolute_path(self.required_model_path) + if ( + type(self.nodes) is not tuple + or not self.nodes + or tuple(node.node_index for node in self.nodes) != tuple(sorted({node.node_index for node in self.nodes})) + ): + raise ValueError("node worker nodes must use sorted unique indices") + process_ids = tuple(process.process_id for node in self.nodes for process in node.processes) + if len({node.host for node in self.nodes}) != len(self.nodes) or len(set(process_ids)) != len(process_ids): + raise ValueError("node worker hosts and process identities must be unique") + for node in self.nodes: + if any(index >= self.resolved_gpus_per_node for process in node.processes for index in process.gpu_indices): + raise ValueError("node process GPU index is outside the allocation") + assigned_gpus = tuple(index for process in node.processes for index in process.gpu_indices) + if len(set(assigned_gpus)) != len(assigned_gpus): + raise ValueError("node processes must not share GPUs") + + +def encode_node_worker_spec(spec: NodeWorkerSpec) -> str: + """Serialize one validated node-worker specification for an argv boundary.""" + payload = { + "schema_version": spec.schema_version, + "resolved_gpus_per_node": spec.resolved_gpus_per_node, + "required_model_path": spec.required_model_path, + "nodes": [ + { + "node_index": node.node_index, + "host": node.host, + "ports": list(node.ports), + "processes": [ + { + "process_id": process.process_id, + "command": list(process.command), + "gpu_indices": list(process.gpu_indices), + "launch_delay_seconds": process.launch_delay_seconds, + } + for process in node.processes + ], + } + for node in spec.nodes + ], + } + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + return base64.urlsafe_b64encode(serialized).decode() + + +def decode_node_worker_spec(encoded: str) -> NodeWorkerSpec: + """Decode and validate an untrusted node-worker specification.""" + if type(encoded) is not str or not encoded or len(encoded) > 2 * _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification is invalid") + try: + serialized = base64.b64decode(encoded, altchars=b"-_", validate=True) + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + payload = json.loads(serialized) + except (binascii.Error, json.JSONDecodeError, UnicodeEncodeError) as error: + raise ValueError("node worker specification is invalid") from error + return _parse_worker_spec(payload) + + +def _parse_worker_spec(payload: object) -> NodeWorkerSpec: + root = _require_mapping(payload, {"schema_version", "resolved_gpus_per_node", "required_model_path", "nodes"}) + nodes = tuple(_parse_node(value) for value in _require_list(root["nodes"])) + return NodeWorkerSpec( + schema_version=_require_integer(root["schema_version"]), + resolved_gpus_per_node=_require_integer(root["resolved_gpus_per_node"]), + required_model_path=_require_optional_string(root["required_model_path"]), + nodes=nodes, + ) + + +def _parse_node(payload: object) -> NodeSpec: + value = _require_mapping(payload, {"node_index", "host", "ports", "processes"}) + return NodeSpec( + node_index=_require_integer(value["node_index"]), + host=_require_string(value["host"]), + ports=tuple(_require_integer(port) for port in _require_list(value["ports"])), + processes=tuple(_parse_process(process) for process in _require_list(value["processes"])), + ) + + +def _parse_process(payload: object) -> NodeProcessSpec: + value = _require_mapping(payload, {"process_id", "command", "gpu_indices", "launch_delay_seconds"}) + return NodeProcessSpec( + process_id=_require_string(value["process_id"]), + command=tuple(_require_string(argument) for argument in _require_list(value["command"])), + gpu_indices=tuple(_require_integer(index) for index in _require_list(value["gpu_indices"])), + launch_delay_seconds=_require_integer(value["launch_delay_seconds"]), + ) + + +def _require_mapping(payload: object, keys: set[str]) -> dict[str, object]: + if not isinstance(payload, dict) or set(payload) != keys or not all(type(key) is str for key in payload): + raise ValueError("node worker object shape is invalid") + return payload + + +def _require_list(payload: object) -> list[object]: + if not isinstance(payload, list): + raise ValueError("node worker list value is invalid") + return payload + + +def _require_string(payload: object) -> str: + if type(payload) is not str: + raise ValueError("node worker string value is invalid") + return payload + + +def _require_optional_string(payload: object) -> str | None: + if payload is None: + return None + return _require_string(payload) + + +def _require_integer(payload: object) -> int: + if type(payload) is not int: + raise ValueError("node worker integer value is invalid") + return payload + + +__all__ = [ + "NodeProcessSpec", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py new file mode 100644 index 000000000..d8e8a0eff --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the resolved vLLM lanes assigned to one physical allocation node.""" + +from __future__ import annotations + +import argparse +import os +import signal +import socket +import stat +import subprocess +import sys +import time +from collections.abc import Mapping, Sequence + +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec as NodeProcessSpec, +) +from data_designer.slurm.runtime.node_spec import ( + NodeSpec as NodeSpec, +) +from data_designer.slurm.runtime.node_spec import ( + NodeWorkerSpec as NodeWorkerSpec, +) +from data_designer.slurm.runtime.node_spec import ( + decode_node_worker_spec as decode_node_worker_spec, +) +from data_designer.slurm.runtime.node_spec import ( + encode_node_worker_spec as encode_node_worker_spec, +) + +_POLL_INTERVAL_SECONDS = 0.1 +_TERMINATION_GRACE_SECONDS = 10.0 + + +class NodeProcessSupervisor: + """Own every lane child on one node and stop them as a unit.""" + + def __init__(self, *, environment: Mapping[str, str]) -> None: + self._environment = dict(environment) + self._children: list[subprocess.Popen[bytes]] = [] + self._cleanup_started = False + self._cleanup_complete = False + self._stop_requested = False + + @property + def cleanup_complete(self) -> bool: + """Return whether every registered lane child has exited.""" + return self._cleanup_complete + + def run(self, node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + """Launch local lanes and fail the node task when any lane exits.""" + started_at = time.monotonic() + try: + for process in node.processes: + self._wait_for_launch(process.launch_delay_seconds, started_at) + if self._stop_requested: + return 1 + self._children.append(self._start_process(process, visible_gpus)) + return self._wait_for_first_exit() + finally: + self.cleanup() + + def request_stop(self) -> None: + """Ask the poll loop to enter cleanup at its next boundary.""" + self._stop_requested = True + + def cleanup(self) -> None: + """Idempotently terminate every registered lane child.""" + if self._cleanup_complete: + return + self._cleanup_started = True + live = tuple(child for child in reversed(self._children) if child.poll() is None) + _signal_children(live, signal.SIGTERM) + deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS + for child in live: + _wait_until_deadline(child, deadline) + _signal_children(tuple(child for child in live if child.poll() is None), signal.SIGKILL) + for child in live: + try: + child.wait(timeout=_TERMINATION_GRACE_SECONDS) + except (OSError, subprocess.SubprocessError): + pass + self._cleanup_complete = all(child.poll() is not None for child in self._children) + if not self._cleanup_complete: + raise RuntimeError("node process cleanup is incomplete") + + def _start_process(self, process: NodeProcessSpec, visible_gpus: tuple[str, ...]) -> subprocess.Popen[bytes]: + if self._cleanup_started: + raise RuntimeError("cannot launch a node process after cleanup") + environment = dict(self._environment) + environment["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpus[index] for index in process.gpu_indices) + return subprocess.Popen( + process.command, + stdin=subprocess.DEVNULL, + env=environment, + shell=False, + start_new_session=True, + close_fds=True, + ) + + def _wait_for_launch(self, delay_seconds: int, started_at: float) -> None: + while not self._stop_requested: + remaining = delay_seconds - (time.monotonic() - started_at) + if remaining <= 0: + return + time.sleep(min(remaining, _POLL_INTERVAL_SECONDS)) + + def _wait_for_first_exit(self) -> int: + while True: + if self._stop_requested: + return 1 + for child in self._children: + returncode = child.poll() + if returncode is not None: + return returncode if returncode != 0 else 1 + time.sleep(_POLL_INTERVAL_SECONDS) + + +def main(arguments: Sequence[str] | None = None) -> int: + """Preflight or serve the node-local slice of a coordinated deployment.""" + parser = argparse.ArgumentParser(prog="data-designer-slurm-node-worker") + parser.add_argument("operation", choices=("preflight", "serve")) + parser.add_argument("--spec", required=True) + parsed = parser.parse_args(arguments) + try: + spec = decode_node_worker_spec(parsed.spec) + node = _select_node(spec, os.environ) + visible_gpus = _parse_visible_gpus(os.environ.get("CUDA_VISIBLE_DEVICES")) + _verify_node(spec, node, visible_gpus, os.environ) + if parsed.operation == "preflight": + _verify_required_model_path(spec.required_model_path) + _verify_ports(node.ports) + return 0 + return _run_node(node, visible_gpus) + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError): + print("node worker failed at a validated runtime boundary", file=sys.stderr) + return 70 + + +def _run_node(node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + supervisor = NodeProcessSupervisor(environment=os.environ) + interrupted: list[int] = [] + + def handle_termination(signum: int, frame: object) -> None: + del frame + if not interrupted: + interrupted.append(signum) + supervisor.request_stop() + + previous = {selected: signal.signal(selected, handle_termination) for selected in (signal.SIGINT, signal.SIGTERM)} + try: + return _run_until_exit_or_signal(supervisor, node, visible_gpus, interrupted) + finally: + try: + supervisor.cleanup() + finally: + for selected, handler in previous.items(): + signal.signal(selected, handler) + + +def _run_until_exit_or_signal( + supervisor: NodeProcessSupervisor, + node: NodeSpec, + visible_gpus: tuple[str, ...], + interrupted: list[int], +) -> int: + if interrupted: + return 128 + interrupted[0] + result = supervisor.run(node, visible_gpus) + if interrupted: + return 128 + interrupted[0] + return result + + +def _select_node(spec: NodeWorkerSpec, environment: Mapping[str, str]) -> NodeSpec: + process_id = _parse_non_negative_integer(environment.get("SLURM_PROCID"), "SLURM_PROCID") + if process_id >= len(spec.nodes): + raise ValueError("Slurm process identity is outside the deployment") + return spec.nodes[process_id] + + +def _verify_node( + spec: NodeWorkerSpec, + node: NodeSpec, + visible_gpus: tuple[str, ...], + environment: Mapping[str, str], +) -> None: + if len(visible_gpus) != spec.resolved_gpus_per_node: + raise ValueError("node GPU visibility does not match the resolved plan") + task_count = _parse_non_negative_integer(environment.get("SLURM_NTASKS"), "SLURM_NTASKS") + if task_count != len(spec.nodes): + raise ValueError("Slurm task count does not match the deployment") + scheduler_host = environment.get("SLURMD_NODENAME") + if scheduler_host is not None and scheduler_host != node.host: + raise ValueError("Slurm node identity does not match the deployment") + + +def _verify_ports(ports: tuple[int, ...]) -> None: + reservations: list[socket.socket] = [] + try: + for port in ports: + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + reservations.append(reservation) + reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) + reservation.bind(("0.0.0.0", port)) + finally: + for reservation in reservations: + reservation.close() + + +def _verify_required_model_path(path: str | None) -> None: + if path is None: + return + status = os.stat(path) + if not stat.S_ISREG(status.st_mode) and not stat.S_ISDIR(status.st_mode): + raise OSError("required model path has an unsupported type") + access_mode = os.R_OK | os.X_OK if stat.S_ISDIR(status.st_mode) else os.R_OK + if not os.access(path, access_mode): + raise OSError("required model path is not readable") + + +def _parse_visible_gpus(value: str | None) -> tuple[str, ...]: + if value is None or not value.strip(): + return () + values = tuple(item.strip() for item in value.split(",")) + if not all(values) or len(values) != len(set(values)) or any("\0" in item for item in values): + raise ValueError("node GPU visibility is invalid") + return values + + +def _parse_non_negative_integer(value: str | None, name: str) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise ValueError(f"{name} is unavailable or invalid") + return int(value) + + +def _signal_process_group(process: subprocess.Popen[bytes], selected: signal.Signals) -> None: + try: + os.killpg(process.pid, selected) + except ProcessLookupError: + pass + + +def _signal_children(children: tuple[subprocess.Popen[bytes], ...], selected: signal.Signals) -> None: + for child in children: + try: + _signal_process_group(child, selected) + except OSError: + continue + + +def _wait_until_deadline(process: subprocess.Popen[bytes], deadline: float) -> None: + while process.poll() is None and time.monotonic() < deadline: + time.sleep(_POLL_INTERVAL_SECONDS) + + +if __name__ == "__main__": # pragma: no cover - exercised as a managed step + raise SystemExit(main()) + + +__all__ = [ + "NodeProcessSpec", + "NodeProcessSupervisor", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", + "main", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh index d9424183d..a16c54392 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -18,13 +18,15 @@ dd_read_control_plan() { local plan=$1 jq -e ' .schema_version == 1 - and .client.host_node_index == 0 - and ([.deployments[].node_indices] | all(. == [0])) + and ([.deployments[].node_indices[]] | length > 0) + and ([.deployments[].node_indices[]] | all(type == "number" and . >= 0)) and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) ' "${plan}" >/dev/null DD_CLIENT_IMAGE=$(jq -er '.client.image.path' "${plan}") DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | tostring' "${plan}") + DD_EXPECTED_NODES=$(jq -er '[.deployments[].node_indices[]] | max + 1 | tostring' "${plan}") + DD_CLIENT_NODE_INDEX=$(jq -er '.client.host_node_index | tostring' "${plan}") DD_GPU_REQUEST_MODE=$(jq -er '.selected_profile.profile.gpu_request_mode' "${plan}") DD_CONTAINER_MOUNTS=$(jq -jr ' [.container_mounts[] | .source + ":" + .target + (if .read_only then ":ro" else "" end)] @@ -140,10 +142,7 @@ dd_read_step() { .stdout_path, "\u0000", .stderr_path, "\u0000", (.launch_delay_seconds | tostring), "\u0000", - (.readiness.host // ""), "\u0000", - (.readiness.port // "" | tostring), "\u0000", - (.readiness.path // ""), "\u0000", - (.readiness.deadline_seconds // "" | tostring), "\u0000" + (.kill_on_bad_exit | tostring), "\u0000" ' "${manifest}" ) DD_STEP_IMAGE=${DD_STEP_FIELDS[0]} @@ -151,10 +150,7 @@ dd_read_step() { DD_STEP_STDOUT=${DD_STEP_FIELDS[2]} DD_STEP_STDERR=${DD_STEP_FIELDS[3]} DD_STEP_DELAY=${DD_STEP_FIELDS[4]} - DD_STEP_PROBE_HOST=${DD_STEP_FIELDS[5]} - DD_STEP_PROBE_PORT=${DD_STEP_FIELDS[6]} - DD_STEP_PROBE_PATH=${DD_STEP_FIELDS[7]} - DD_STEP_PROBE_DEADLINE=${DD_STEP_FIELDS[8]} + DD_STEP_KILL_ON_BAD_EXIT=${DD_STEP_FIELDS[5]} dd_read_null_values DD_STEP_COMMAND < <( jq -j --arg step_id "${step_id}" '.steps[] | select(.step_id == $step_id) | .command[] | ., "\u0000"' \ "${manifest}" @@ -163,6 +159,21 @@ dd_read_step() { jq -j --arg step_id "${step_id}" \ '.steps[] | select(.step_id == $step_id) | .gpu_indices[] | tostring, "\u0000"' "${manifest}" ) + dd_read_null_values DD_STEP_NODE_HOSTS < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .node_hosts[] | ., "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_PROBE_FIELDS < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .readiness[] + | .host, "\u0000", + (.port | tostring), "\u0000", + .path, "\u0000", + (.deadline_seconds | tostring), "\u0000" + ' "${manifest}" + ) dd_read_null_values DD_STEP_CONTAINER_ENV < <( jq -j --arg step_id "${step_id}" \ '.steps[] | select(.step_id == $step_id) | .container_environment[] | ., "\u0000"' "${manifest}" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py index 78366c390..1b3204d41 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -11,19 +11,70 @@ import socket import stat from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path from typing import Protocol from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.network import validate_host_name from data_designer.slurm.runtime.paths import get_container_path -from data_designer.slurm.runtime.ports import allocation_ports +from data_designer.slurm.runtime.ports import resolve_allocation_plan _DIGEST_CHUNK_SIZE = 1024 * 1024 _GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") +@dataclass(frozen=True, slots=True) +class AllocationLayout: + """Verified allocation host identities in planner-index order.""" + + node_hosts: tuple[str, ...] + + def __post_init__(self) -> None: + if ( + type(self.node_hosts) is not tuple + or not self.node_hosts + or len(self.node_hosts) != len(set(self.node_hosts)) + ): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation node identities are invalid") + try: + for host in self.node_hosts: + validate_host_name(host) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, str(error)) from error + + def get_host(self, node_index: int) -> str: + """Return the host assigned to one planner node index.""" + if type(node_index) is not int or node_index < 0: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) + try: + return self.node_hosts[node_index] + except (IndexError, TypeError): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) from None + + +def validate_allocation_layout(plan: ResolvedSlurmRunPlan, layout: AllocationLayout) -> None: + """Require one scheduler host for every contiguous planner node index.""" + node_indices = { + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + } + if node_indices != set(range(len(layout.node_hosts))): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation host identities do not match the resolved plan", + ) + + class AllocationPreflight(Protocol): """Verify one allocation without starting package-managed processes.""" @@ -59,16 +110,17 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] context.plan.client.host_node_index, *(index for deployment in context.plan.deployments for index in deployment.node_indices), } - if node_indices != {0}: + node_count = max(node_indices) + 1 + if node_indices != set(range(node_count)): raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one-node runtime received a multi-node execution plan", + "resolved allocation node indices are not complete", ) expected = { "SLURM_ARRAY_JOB_ID": context.attempt.scheduler.array_job_id, "SLURM_ARRAY_TASK_ID": context.shard.array_task_index, - "SLURM_JOB_NUM_NODES": 1, - "SLURM_NODEID": 0, + "SLURM_JOB_NUM_NODES": node_count, + "SLURM_NODEID": context.plan.client.host_node_index, } for name, value in expected.items(): if _parse_non_negative_integer(environment.get(name), name) != value: @@ -125,19 +177,26 @@ def _verify_artifacts(context: AllocationContext) -> None: @staticmethod def verify_ports(context: AllocationContext, environment: Mapping[str, str]) -> None: - """Verify that every planned one-node port is currently bindable.""" - ports = allocation_ports(context, environment) + """Verify ports owned by the local client host before nested steps start.""" + plan = resolve_allocation_plan(context.plan, environment) + local_node_index = plan.client.host_node_index + ports = tuple(port.port for port in plan.client.ports if port.node_index == local_node_index) + tuple( + port.port + for deployment in plan.deployments + for port in deployment.ports + if port.node_index == local_node_index + ) reservations: list[socket.socket] = [] try: for port in ports: reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) reservations.append(reservation) reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) - reservation.bind(("127.0.0.1", port)) + reservation.bind(("0.0.0.0", port)) except OSError as error: raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one or more resolved allocation ports are unavailable", + "one or more local allocation ports are unavailable", ) from error finally: for reservation in reservations: @@ -195,4 +254,9 @@ def _parse_gpu_count(value: str | None) -> int: return len(values) -__all__ = ["AllocationPreflight", "SystemAllocationPreflight"] +__all__ = [ + "AllocationLayout", + "AllocationPreflight", + "SystemAllocationPreflight", + "validate_allocation_layout", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py index 64dd51fe2..402677767 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -13,6 +13,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import SplitResult, urlsplit +from data_designer.slurm.runtime.network import validate_host_name + _MAXIMUM_REQUEST_BYTES = 64 * 1024 * 1024 _MAXIMUM_RESPONSE_BYTES = 64 * 1024 * 1024 _HOP_HEADERS = frozenset( @@ -191,9 +193,14 @@ def main(arguments: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="data-designer-slurm-proxy") parser.add_argument("--listen-port", required=True, type=int) parser.add_argument("--backend", action="append", required=True) + parser.add_argument("--allowed-host", action="append") parser.add_argument("--retry-after-seconds", type=int) parsed = parser.parse_args(arguments) - backends = tuple(_parse_backend(value) for value in parsed.backend) + try: + allowed_hosts = frozenset(validate_host_name(value) for value in (parsed.allowed_host or ("127.0.0.1",))) + except ValueError as error: + parser.error(str(error)) + backends = tuple(_parse_backend(value, allowed_hosts) for value in parsed.backend) if not 1 <= parsed.listen_port <= 65535: parser.error("listen port must be between 1 and 65535") if parsed.retry_after_seconds is not None and parsed.retry_after_seconds <= 0: @@ -204,15 +211,18 @@ def main(arguments: Sequence[str] | None = None) -> int: return 0 -def _parse_backend(value: str) -> _Backend: +def _parse_backend(value: str, allowed_hosts: frozenset[str] = frozenset({"127.0.0.1"})) -> _Backend: parsed: SplitResult = urlsplit(value) + host = parsed.hostname + normalized_allowed_hosts = frozenset(allowed_host.casefold() for allowed_host in allowed_hosts) try: port = parsed.port except ValueError as error: raise argparse.ArgumentTypeError("backend port is invalid") from error if ( parsed.scheme != "http" - or parsed.hostname != "127.0.0.1" + or host is None + or host.casefold() not in normalized_allowed_hosts or parsed.username is not None or parsed.password is not None or parsed.path not in {"", "/"} @@ -220,8 +230,8 @@ def _parse_backend(value: str) -> _Backend: or parsed.fragment or port is None ): - raise argparse.ArgumentTypeError("backends must be loopback HTTP origins") - return _Backend(parsed.hostname, port) + raise argparse.ArgumentTypeError("backends must be allowed HTTP origins") + return _Backend(host, port) def _retry_after_headers(headers: dict[str, str], retry_after_seconds: int | None) -> dict[str, str]: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py new file mode 100644 index 000000000..4cd6b4401 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose local and distributed server entries for the runtime manifest.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Literal + +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.distributed import build_node_worker_spec, build_vllm_process_command +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.manifest import RuntimeProbeSpec, RuntimeStepSpec +from data_designer.slurm.runtime.models import AllocationContext, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import encode_node_worker_spec +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.steps import build_vllm_command +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess + + +def build_server_steps( + deployments: tuple[ResolvedVllmServerDeployment, ...], + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> tuple[RuntimeStepSpec, ...]: + """Build remote preflight entries followed by all serving entries.""" + preflight_steps = tuple( + build_distributed_server_step( + deployment, + context, + runtime_container_root, + log_directory, + layout, + operation="preflight", + ) + for deployment in deployments + if _requires_server_preflight(deployment, context) + ) + serving_steps = tuple( + step + for deployment in deployments + for step in _build_deployment_serving_steps( + deployment, + context, + runtime_container_root, + log_directory, + layout, + ) + ) + return preflight_steps + serving_steps + + +def build_distributed_server_step( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, + *, + operation: Literal["preflight", "serve"], +) -> RuntimeStepSpec: + """Build one coordinated node-worker step for a deployment.""" + literal_environment, secret_environment, environment_prefixes = _server_environment( + deployment, + runtime_container_root, + ) + worker_spec = build_node_worker_spec(deployment, context.plan, layout) + readiness = ( + tuple( + RuntimeProbeSpec( + host=layout.get_host(probe.node_index), + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ) + for probe in deployment.readiness_probes + ) + if operation == "serve" + else () + ) + step_id = f"{deployment.deployment_id}-{operation}" + return RuntimeStepSpec( + step_id=step_id, + role=RuntimeStepRole.SERVER if operation == "serve" else RuntimeStepRole.SERVER_PREFLIGHT, + image_path=deployment.image.path, + command=( + "python3", + f"{runtime_container_root}/data_designer/slurm/runtime/node_worker.py", + operation, + "--spec", + encode_node_worker_spec(worker_spec), + ), + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple(range(deployment.gpus_per_node)), + node_hosts=tuple(layout.get_host(index) for index in deployment.node_indices), + kill_on_bad_exit=True, + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=_server_container_environment(deployment), + stdout_path=(log_directory / f"{step_id}.out").as_posix(), + stderr_path=(log_directory / f"{step_id}.err").as_posix(), + readiness=readiness, + ) + + +def _build_deployment_serving_steps( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> tuple[RuntimeStepSpec, ...]: + if len(deployment.node_indices) > 1: + return ( + build_distributed_server_step( + deployment, + context, + runtime_container_root, + log_directory, + layout, + operation="serve", + ), + ) + return tuple( + _build_local_server_step( + deployment, + process, + context, + runtime_container_root, + log_directory, + layout, + ) + for process in deployment.processes + ) + + +def _build_local_server_step( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> RuntimeStepSpec: + if process.pipeline_parallel != 1 or process.http_port is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "single-node runtime received a distributed vLLM process", + ) + literal_environment, secret_environment, environment_prefixes = _server_environment( + deployment, + runtime_container_root, + ) + probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) + probe_host = "127.0.0.1" if len(layout.node_hosts) == 1 else layout.get_host(process.node_index) + command = ( + build_vllm_command(deployment, process, context.plan) + if len(layout.node_hosts) == 1 + else build_vllm_process_command(deployment, process, context.plan, layout) + ) + return RuntimeStepSpec( + step_id=process.process_id, + role=RuntimeStepRole.SERVER, + image_path=deployment.image.path, + command=command, + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple(process.gpu_indices), + node_hosts=(layout.get_host(process.node_index),), + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=_server_container_environment(deployment), + stdout_path=(log_directory / f"{process.process_id}.out").as_posix(), + stderr_path=(log_directory / f"{process.process_id}.err").as_posix(), + launch_delay_seconds=process.launch_delay_seconds, + readiness=( + RuntimeProbeSpec( + host=probe_host, + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ), + ), + ) + + +def _server_environment( + deployment: ResolvedVllmServerDeployment, + runtime_container_root: str, +) -> tuple[dict[str, str], dict[str, str], dict[str, str]]: + literal_environment: dict[str, str] = {"LC_ALL": "C", "PYTHONPATH": runtime_container_root} + secret_environment: dict[str, str] = {} + environment_prefixes: dict[str, str] = {} + for name, binding in deployment.launch_policy.environment.items(): + if isinstance(binding, LiteralEnvironmentBinding): + literal_environment[name] = binding.value + elif isinstance(binding, SecretRef): + secret_environment[name] = binding.environment + else: # pragma: no cover - persisted contracts reject unknown bindings + raise AssertionError(f"unhandled environment binding: {type(binding)!r}") + if "PYTHONPATH" in secret_environment: + literal_environment.pop("PYTHONPATH") + environment_prefixes["PYTHONPATH"] = runtime_container_root + elif "PYTHONPATH" in deployment.launch_policy.environment: + literal_environment["PYTHONPATH"] = os.pathsep.join((runtime_container_root, literal_environment["PYTHONPATH"])) + policy = deployment.launch_policy.queue_backpressure + literal_environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(policy.max_waiting_requests) + literal_environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if policy.retry_after_seconds is None else str(policy.retry_after_seconds) + ) + return literal_environment, secret_environment, environment_prefixes + + +def _server_container_environment(deployment: ResolvedVllmServerDeployment) -> tuple[str, ...]: + return tuple( + sorted( + { + *deployment.launch_policy.environment, + "PYTHONPATH", + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + } + ) + ) + + +def _requires_server_preflight( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, +) -> bool: + return any(node_index != context.plan.client.host_node_index for node_index in deployment.node_indices) + + +__all__ = ["build_distributed_server_step", "build_server_steps"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh index 7dc5aeae3..34e5c6aeb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh @@ -66,12 +66,17 @@ dd_materialize_step_environment() { } dd_build_srun_command() { - local index gpu_mask=0 container_names= + local index gpu_mask=0 container_names= node_list= + ((${#DD_STEP_NODE_HOSTS[@]})) || return 64 + printf -v node_list '%s,' "${DD_STEP_NODE_HOSTS[@]}" + node_list=${node_list%,} DD_STEP_VISIBLE_GPUS= DD_SRUN_COMMAND=( srun - --nodes=1 - --ntasks=1 + "--nodes=${#DD_STEP_NODE_HOSTS[@]}" + "--ntasks=${#DD_STEP_NODE_HOSTS[@]}" + --ntasks-per-node=1 + "--nodelist=${node_list}" --exact --overlap --unbuffered @@ -79,6 +84,9 @@ dd_build_srun_command() { "--cpus-per-task=${DD_STEP_CPUS}" "--container-image=${DD_STEP_IMAGE}" ) + if [[ ${DD_STEP_KILL_ON_BAD_EXIT} == true ]]; then + DD_SRUN_COMMAND+=(--kill-on-bad-exit=1) + fi if ((${#DD_STEP_GPU_INDICES[@]})); then if [[ ${DD_GPU_REQUEST_MODE} == gres ]]; then for index in "${DD_STEP_GPU_INDICES[@]+"${DD_STEP_GPU_INDICES[@]}"}"; do @@ -138,6 +146,7 @@ dd_run_control_phase() { srun --nodes=1 --ntasks=1 + "--nodelist=${DD_CLIENT_HOST}" --exact --overlap --unbuffered diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index 74ca803ea..555c0a5db 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -274,8 +274,24 @@ def build_endpoint_command( plan: ResolvedSlurmRunPlan, runtime_proxy_path: Path, port: int, + *, + backend_hosts: tuple[str, ...] | None = None, ) -> tuple[str, ...]: - backends = tuple(f"http://127.0.0.1:{backend.port}" for backend in deployment.backend_endpoints) + selected_hosts = backend_hosts or ("127.0.0.1",) * len(deployment.backend_endpoints) + if len(selected_hosts) != len(deployment.backend_endpoints): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "endpoint backend hosts do not match the resolved deployment", + ) + backends = tuple( + f"http://{host}:{backend.port}" + for host, backend in zip(selected_hosts, deployment.backend_endpoints, strict=True) + ) + allowed_host_arguments = ( + tuple(argument for host in dict.fromkeys(selected_hosts) for argument in ("--allowed-host", host)) + if backend_hosts is not None + else () + ) retry_after_seconds = deployment.launch_policy.queue_backpressure.retry_after_seconds retry_arguments = ("--retry-after-seconds", str(retry_after_seconds)) if retry_after_seconds is not None else () return ( @@ -284,6 +300,7 @@ def build_endpoint_command( "--listen-port", str(port), *retry_arguments, + *allowed_host_arguments, *(argument for backend in backends for argument in ("--backend", backend)), ) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 10fb6c277..c9cedca44 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -4,11 +4,19 @@ from __future__ import annotations from dataclasses import replace +from pathlib import Path -from conftest import RuntimeCase +from conftest import RuntimeCase, relocate_plan +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest -from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.runtime.distributed import build_vllm_process_command +from data_designer.slurm.runtime.models import AllocationContext, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import decode_node_worker_spec +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.serving.vllm import ResolvedVllmProcess from data_designer.slurm.state import RetryPlan, RetryShard @@ -22,6 +30,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru {"SLURM_JOB_GPUS": "0"}, runtime_root=runtime_root, log_directory=log_directory, + layout=AllocationLayout(("compute-001",)), ) reloaded = RuntimeBootstrapManifest.model_validate_json(manifest.serialize_json()) @@ -44,6 +53,8 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert "--attempt-id" not in manifest.steps[-1].command assert "--plan" in manifest.steps[-1].command assert "--attempt-dir" in manifest.steps[-1].command + assert all(step.node_hosts == ("compute-001",) for step in manifest.steps) + assert all(step.role is not RuntimeStepRole.SERVER_PREFLIGHT for step in manifest.steps) def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runtime_case: RuntimeCase) -> None: @@ -71,6 +82,7 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti {"SLURM_JOB_GPUS": "0"}, runtime_root=context.attempt_directory / "runtime", log_directory=context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), ) preflight = manifest.steps[0].command @@ -80,3 +92,107 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti assert ("--retry-plan-sha256", retry.compute_sha256()) == client[client.index("--retry-plan-sha256") :][:2] assert ("--effective-resume-mode", "never") == client[client.index("--effective-resume-mode") :][:2] assert "--resume-mode" not in client + + +def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( + runtime_case: RuntimeCase, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployments = tuple( + deployment.model_copy( + update={ + "authored": deployment.authored.model_copy( + update={ + "model": f"/workspace/primary/models/model-{deployment_index}", + "served_model_name": deployment.served_model_name, + } + ), + "model": f"/workspace/primary/models/model-{deployment_index}", + } + ) + for deployment_index, deployment in enumerate(multi_node_plan.deployments) + ) + context = _replace_plan(runtime_case, multi_node_plan.model_copy(update={"deployments": deployments})) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, + runtime_root=context.attempt_directory / "runtime", + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=layout, + ) + + distributed = next(step for step in manifest.steps if step.step_id == "deployment-00000-serve") + preflight = next(step for step in manifest.steps if step.step_id == "deployment-00000-preflight") + worker_spec = decode_node_worker_spec(distributed.command[-1]) + endpoint = next(step for step in manifest.steps if step.step_id == "deployment-00000-endpoint") + remote_server = next(step for step in manifest.steps if step.step_id == "deployment-00001-replica-00000-rank-00000") + remote_preflight = next(step for step in manifest.steps if step.step_id == "deployment-00001-preflight") + + assert distributed.node_hosts == ("compute-001", "compute-002") + assert distributed.kill_on_bad_exit + assert preflight.role is RuntimeStepRole.SERVER_PREFLIGHT + assert tuple(node.host for node in worker_spec.nodes) == distributed.node_hosts + assert "--master-addr" in worker_spec.nodes[0].processes[0].command + assert "compute-001" in worker_spec.nodes[0].processes[0].command + assert worker_spec.nodes[0].processes[0].command[2] == "/workspace/primary/models/model-0" + assert worker_spec.required_model_path == "/workspace/primary/models/model-0" + assert "--headless" in worker_spec.nodes[1].processes[0].command + assert tuple(probe.host for probe in distributed.readiness) == ("compute-001",) + assert "http://compute-001:" in " ".join(endpoint.command) + assert endpoint.node_hosts == ("compute-001",) + assert remote_preflight.node_hosts == ("compute-003",) + remote_worker_spec = decode_node_worker_spec(remote_preflight.command[-1]) + assert remote_worker_spec.required_model_path == "/workspace/primary/models/model-1" + assert remote_server.node_hosts == ("compute-003",) + assert remote_server.command[2] == "/workspace/primary/models/model-1" + executor_index = remote_server.command.index("--distributed-executor-backend") + assert remote_server.command[executor_index + 1] == "uni" + assert tuple(probe.host for probe in remote_server.readiness) == ("compute-003",) + + +def test_pipeline_parallel_process_uses_multi_process_executor( + runtime_case: RuntimeCase, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + context = _replace_plan(runtime_case, multi_node_plan) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + deployment = resolve_allocation_deployments( + context, + {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, + )[0] + payload = deployment.processes[0].model_dump(mode="python") + payload.update({"gpu_indices": (0,), "tensor_parallel": 1}) + process = ResolvedVllmProcess.model_validate(payload) + + command = build_vllm_process_command(deployment, process, context.plan, layout) + + executor_index = command.index("--distributed-executor-backend") + assert process.tensor_parallel == 1 + assert process.pipeline_parallel == 2 + assert command[executor_index + 1] == "mp" + + +def _replace_plan(runtime_case: RuntimeCase, source_plan: ResolvedSlurmRunPlan) -> AllocationContext: + plan = relocate_plan(source_plan, runtime_case.workspace) + shard = plan.shards[0] + plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + attempt = runtime_case.context.attempt.model_copy( + update={ + "run_id": plan.run_id, + "shard_id": shard.shard_id, + "resolved_plan": ArtifactReference(path=plan_path.as_posix(), sha256=plan.compute_sha256()), + "scheduler": runtime_case.context.attempt.scheduler.model_copy( + update={"array_task_id": shard.array_task_index} + ), + } + ) + attempt_directory = plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt.attempt_id + attempt_directory.mkdir(parents=True, mode=0o700) + return AllocationContext( + plan=plan, + shard=shard, + attempt=attempt, + attempt_directory=attempt_directory, + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 561588810..3c27e87c9 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -60,6 +60,8 @@ def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.Capt "runtime", "--manifest", "runtime-manifest.json", + "--node-host", + "compute-001", ) ) == 64 @@ -465,5 +467,13 @@ def _phase_arguments( ) if operation == "prepare": assert runtime_root is not None and manifest_path is not None - return (*arguments, "--runtime-root", runtime_root.as_posix(), "--manifest", manifest_path.as_posix()) + return ( + *arguments, + "--runtime-root", + runtime_root.as_posix(), + "--manifest", + manifest_path.as_posix(), + "--node-host", + "compute-001", + ) return arguments diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py new file mode 100644 index 000000000..a9511f1d8 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from data_designer.slurm.runtime import node_worker as runtime_node_worker +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec, + NodeSpec, + NodeWorkerSpec, + decode_node_worker_spec, + encode_node_worker_spec, +) +from data_designer.slurm.runtime.node_worker import NodeProcessSupervisor + + +def test_node_worker_spec_round_trips_and_rejects_untrusted_payload() -> None: + spec = _worker_spec((NodeProcessSpec("lane-0", ("true",), (0,), 0),)) + assert decode_node_worker_spec(encode_node_worker_spec(spec)) == spec + with pytest.raises(ValueError, match="invalid"): + decode_node_worker_spec("not-base64") + + +def test_node_preflight_rejects_missing_required_model_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + model_path = tmp_path / "missing-model" + spec = _worker_spec( + (NodeProcessSpec("lane-0", ("true",), (0,), 0),), + required_model_path=model_path.as_posix(), + ) + _set_node_environment(monkeypatch) + + assert runtime_node_worker.main(("preflight", "--spec", encode_node_worker_spec(spec))) == 70 + assert capsys.readouterr().err == "node worker failed at a validated runtime boundary\n" + + +def test_node_preflight_rejects_unreadable_required_model_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + model_path = tmp_path / "model" + model_path.mkdir() + spec = _worker_spec( + (NodeProcessSpec("lane-0", ("true",), (0,), 0),), + required_model_path=model_path.as_posix(), + ) + _set_node_environment(monkeypatch) + monkeypatch.setattr(runtime_node_worker.os, "access", lambda path, mode: False) + + assert runtime_node_worker.main(("preflight", "--spec", encode_node_worker_spec(spec))) == 70 + assert capsys.readouterr().err == "node worker failed at a validated runtime boundary\n" + + +def test_partial_startup_failure_cleans_already_started_lane(monkeypatch: pytest.MonkeyPatch) -> None: + process = _FakeProcess(pid=41) + calls = 0 + signals: list[int] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + nonlocal calls + del arguments, keywords + calls += 1 + if calls == 2: + raise OSError("injected partial startup") + return process + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + + def terminate(child: _FakeProcess, selected: int) -> None: + signals.append(selected) + child.returncode = -15 + + monkeypatch.setattr(runtime_node_worker, "_signal_process_group", terminate) + supervisor = NodeProcessSupervisor(environment={}) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("lane-0", ("first",), (0,), 0), + NodeProcessSpec("lane-1", ("second",), (1,), 0), + ), + ) + + with pytest.raises(OSError, match="partial startup"): + supervisor.run(node, ("0", "1")) + + assert signals + assert supervisor.cleanup_complete + + +def test_follower_failure_terminates_sibling_without_orphan(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + long_running = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + failing = "import time; time.sleep(0.2); raise SystemExit(23)" + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("head", (sys.executable, "-c", long_running), (0,), 0), + NodeProcessSpec("follower", (sys.executable, "-c", failing), (1,), 0), + ), + ) + supervisor = NodeProcessSupervisor(environment={}) + + assert supervisor.run(node, ("0", "1")) == 23 + supervisor.cleanup() + + assert ready.exists() + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_request_stops_all_node_lanes(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + child = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("lane", (sys.executable, "-c", child), (0,), 0),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + _wait_for_file(ready) + supervisor.request_stop() + assert result.result(timeout=5) == 1 + + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_interrupts_a_pending_launch_delay() -> None: + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("delayed", ("must-not-start",), (0,), 30),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + time.sleep(0.05) + supervisor.request_stop() + assert result.result(timeout=2) == 1 + + assert supervisor.cleanup_complete + + +@dataclass(slots=True) +class _FakeProcess: + pid: int + returncode: int | None = None + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + del timeout + self.returncode = -15 + return self.returncode + + +def _worker_spec( + processes: tuple[NodeProcessSpec, ...], + *, + required_model_path: str | None = None, +) -> NodeWorkerSpec: + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=8, + required_model_path=required_model_path, + nodes=(NodeSpec(node_index=0, host="compute-001", ports=(18000,), processes=processes),), + ) + + +def _set_node_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("SLURM_NTASKS", "1") + monkeypatch.setenv("SLURMD_NODENAME", "compute-001") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + + +def _wait_for_file(path: Path) -> None: + deadline = time.monotonic() + 3 + while not path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert path.exists() diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 4cc0890e6..604e3e737 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -6,6 +6,8 @@ import hashlib import os from pathlib import Path +from types import SimpleNamespace +from typing import cast import pytest from conftest import RuntimeCase @@ -15,7 +17,40 @@ from data_designer.slurm.runtime import preflight as runtime_preflight from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.runtime.models import AllocationContext -from data_designer.slurm.runtime.preflight import SystemAllocationPreflight, _verify_artifact +from data_designer.slurm.runtime.preflight import AllocationLayout, SystemAllocationPreflight, _verify_artifact + + +def test_allocation_layout_rejects_invalid_hosts_and_indices() -> None: + layout = AllocationLayout(("compute-001", "compute-002")) + + with pytest.raises(SlurmRuntimeError, match="outside"): + layout.get_host(-1) + with pytest.raises(SlurmRuntimeError, match="identity"): + AllocationLayout(("compute-001", "--invalid")) + + +def test_scheduler_preflight_accepts_complete_multi_node_plan() -> None: + context = cast( + AllocationContext, + SimpleNamespace( + plan=SimpleNamespace( + client=SimpleNamespace(host_node_index=0), + deployments=(SimpleNamespace(node_indices=(0, 1)), SimpleNamespace(node_indices=(2,))), + resolved_gpus_per_node=8, + ), + attempt=SimpleNamespace(scheduler=SimpleNamespace(array_job_id=4101)), + shard=SimpleNamespace(array_task_index=0), + ), + ) + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "3", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", + } + + SystemAllocationPreflight._verify_scheduler(context, environment) def test_scheduler_preflight_accepts_exact_one_node_gpu_shape(runtime_case: RuntimeCase) -> None: diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py index 0a600f1bf..5387aa4ec 100644 --- a/packages/data-designer-slurm/tests/runtime/test_proxy.py +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -189,6 +189,12 @@ def test_proxy_rejects_non_loopback_or_malformed_backend(value: str) -> None: _parse_backend(value) +def test_proxy_matches_allowed_backend_hosts_case_insensitively() -> None: + backend = _parse_backend("http://compute-001:8000", frozenset({"Compute-001"})) + + assert backend == _Backend("compute-001", 8000) + + def test_pool_selects_least_active_backend() -> None: pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) first = pool.acquire(frozenset()) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index ef877ff6a..555b003d3 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -35,7 +35,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( _write_executable(fake_bin / "python3", f"#!/usr/bin/env bash\nprintf ran > {marker_path}\nexit 99\n") _write_executable(fake_bin / "curl", "#!/usr/bin/env bash\nexit 0\n") _write_executable(fake_bin / "getent", "#!/usr/bin/env bash\nexit 0\n") - _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nprintf 'compute-001\\n'\n") _write_executable(fake_bin / "srun", _fake_srun()) command = f""" set -Eeuo pipefail @@ -57,6 +57,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "SLURM_ARRAY_JOB_ID": "4101", "SLURM_ARRAY_TASK_ID": "0", "SLURM_JOB_NUM_NODES": "1", + "SLURM_JOB_NODELIST": "compute-001", "SLURM_JOB_GPUS": "0", "SLURM_NODEID": "0", } @@ -134,6 +135,62 @@ def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: assert completed.returncode == 0, completed.stderr +def test_step_runner_builds_one_coordinated_srun_across_selected_nodes() -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "step_runner.sh").as_posix())} +DD_STEP_NODE_HOSTS=(compute-001 compute-002) +DD_STEP_GPU_INDICES=(0 1) +DD_STEP_CONTAINER_ENV=() +DD_STEP_CPUS=4 +DD_STEP_IMAGE=/images/server.sqsh +DD_STEP_KILL_ON_BAD_EXIT=true +DD_GPU_REQUEST_MODE=gres +DD_CONTAINER_MOUNTS= +dd_build_srun_command +command=${{DD_SRUN_COMMAND[*]}} +[[ $command == *--nodes=2* ]] +[[ $command == *--ntasks=2* ]] +[[ $command == *--ntasks-per-node=1* ]] +[[ $command == *--nodelist=compute-001,compute-002* ]] +[[ $command == *--kill-on-bad-exit=1* ]] +[[ $command == *--gpus-per-task=2* ]] +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True) + + assert completed.returncode == 0, completed.stderr + + +def test_shell_resolves_scheduler_hosts_in_planner_order(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nprintf 'compute-001\\ncompute-002\\n'\n") + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +DD_EXPECTED_NODES=2 +DD_CLIENT_NODE_INDEX=0 +SLURM_JOB_NODELIST=compute-[001-002] +dd_resolve_allocation_hosts +[[ ${{DD_ALLOCATION_HOSTS[*]}} == 'compute-001 compute-002' ]] +[[ $DD_CLIENT_HOST == compute-001 ]] +DD_EXPECTED_NODES=3 +! dd_resolve_allocation_hosts +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "PATH": f"{fake_bin}:{os.environ['PATH']}"}, + ) + + assert completed.returncode == 0, completed.stderr + + def _artifact(root: Path, name: str) -> tuple[str, str]: path = root / f"{name}.artifact" content = name.encode() @@ -245,6 +302,8 @@ def _step( "command": [command], "cpus": 1, "gpu_indices": gpu_indices or [], + "node_hosts": ["compute-001"], + "kill_on_bad_exit": False, "literal_environment": {"LC_ALL": "C"}, "secret_environment": secret_environment or {}, "environment_prefixes": {}, @@ -252,7 +311,7 @@ def _step( "stdout_path": (log_root / f"{step_id}.out").as_posix(), "stderr_path": (log_root / f"{step_id}.err").as_posix(), "launch_delay_seconds": 0, - "readiness": readiness, + "readiness": [readiness] if readiness is not None else [], }