Skip to content

feat(slurm): harden distributed allocation runtime #914

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nabinchha merged 4 commits into from
Sep 10, 2026
Merged

feat(slurm): harden distributed allocation runtime #914

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3681e04
feat(slurm): compose distributed runtime launch
nabinchha Sep 10, 2026
3492cad
map distributed model paths
nabinchha Sep 10, 2026
5222dd5
adapt retry bootstrap coverage
nabinchha Sep 10, 2026
0a86e11
fix distributed launch preflight
nabinchha Sep 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
219 changes: 36 additions & 183 deletions packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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(
Expand All @@ -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(
Expand All @@ -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,
Expand All @@ -172,6 +71,7 @@ def build_runtime_manifest(
endpoints,
runtime_container_root,
log_directory,
layout,
)
)
secret_names = set(collect_secret_environment_names(plan))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),),
)


Expand All @@ -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",
Expand All @@ -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,
),
),
)

Expand All @@ -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,
Expand All @@ -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,
Expand Down
Loading
Loading