diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py new file mode 100644 index 000000000..a46fcb160 --- /dev/null +++ b/packages/data-designer-slurm/tests/conftest.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast + +import pytest +from slurm_test_fakes import ( + FakeClock, + FakeCommandResponse, + FakeLogicalEndpoint, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, + FakeVllmBackend, +) + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageInspectionRecord, + SlurmProfileCatalog, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerIdentity, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + +TEST_DIRECTORY = Path(__file__).parent +CONTRACT_GOLDEN_DIRECTORY = TEST_DIRECTORY / "contracts" / "golden" +INTEGRATION_GOLDEN_PATH = TEST_DIRECTORY / "integration" / "golden" / "finalization_chain.json" +SLURM_GOLDEN_DIRECTORY = TEST_DIRECTORY / "slurm_test_fakes" / "golden" / "slurm" +STATE_GOLDEN_DIRECTORY = TEST_DIRECTORY / "state" / "golden" + + +@pytest.fixture +def authored_run() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "authored_run.json").read_text()) + + +@pytest.fixture +def authored_run_single() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json").read_text() + ) + + +@pytest.fixture +def profile_catalog() -> SlurmProfileCatalog: + return SlurmProfileCatalog.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "profile_catalog.json").read_text()) + + +@pytest.fixture +def dependency_lock() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "dependency_lock.json").read_text()) + + +@pytest.fixture +def dependency_lock_single() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text() + ) + + +@pytest.fixture +def single_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text()) + + +@pytest.fixture +def multi_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "multi_node_plan.json").read_text()) + + +@pytest.fixture +def client_result() -> ClientResult: + return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.json").read_text()) + + +@pytest.fixture +def benchmark_config() -> DataDesignerSlurmBenchmarkConfig: + return DataDesignerSlurmBenchmarkConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "benchmark_config.json").read_text() + ) + + +@pytest.fixture +def benchmark_manifest() -> BenchmarkManifest: + return BenchmarkManifest.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_manifest.json").read_text()) + + +@pytest.fixture +def benchmark_report() -> BenchmarkReport: + return BenchmarkReport.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_report.json").read_text()) + + +@pytest.fixture +def client_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "client_image_inspection.json").read_text() + ) + + +@pytest.fixture +def serving_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "serving_image_inspection.json").read_text() + ) + + +@pytest.fixture +def collection_plan() -> CollectionPlan: + return CollectionPlan.model_validate_json((STATE_GOLDEN_DIRECTORY / "collection_plan.json").read_text()) + + +@pytest.fixture +def accounting_lag_observation() -> SchedulerObservation: + return SchedulerObservation.model_validate_json((STATE_GOLDEN_DIRECTORY / "accounting_lag.json").read_text()) + + +@pytest.fixture +def finalization_chain_payload() -> dict[str, object]: + return cast(dict[str, object], json.loads(INTEGRATION_GOLDEN_PATH.read_text())) + + +@pytest.fixture +def run_manifest(finalization_chain_payload: dict[str, object]) -> RunManifest: + return RunManifest.model_validate_json(json.dumps(finalization_chain_payload["run"])) + + +@pytest.fixture +def shard_manifests(finalization_chain_payload: dict[str, object]) -> tuple[ShardManifest, ...]: + return tuple( + ShardManifest.model_validate_json(json.dumps(payload)) + for payload in cast(list[object], finalization_chain_payload["shards"]) + ) + + +@pytest.fixture +def attempt_manifest(finalization_chain_payload: dict[str, object]) -> AttemptManifest: + return AttemptManifest.model_validate_json(json.dumps(finalization_chain_payload["attempt"])) + + +@pytest.fixture +def attempt_readiness(finalization_chain_payload: dict[str, object]) -> AttemptReadiness: + return AttemptReadiness.model_validate_json(json.dumps(finalization_chain_payload["readiness"])) + + +@pytest.fixture +def finalization_client_result(finalization_chain_payload: dict[str, object]) -> ClientResult: + return ClientResult.model_validate_json(json.dumps(finalization_chain_payload["client_result"])) + + +@pytest.fixture +def candidate_output_manifest(finalization_chain_payload: dict[str, object]) -> CandidateOutputManifest: + return CandidateOutputManifest.model_validate_json(json.dumps(finalization_chain_payload["candidate"])) + + +@pytest.fixture +def shard_winner(finalization_chain_payload: dict[str, object]) -> ShardWinner: + return ShardWinner.model_validate_json(json.dumps(finalization_chain_payload["winner"])) + + +@pytest.fixture +def fake_clock() -> FakeClock: + """Return an isolated explicitly controlled clock.""" + return FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=100.0) + + +@pytest.fixture +def fake_slurm_runner() -> FakeSlurmRunner: + """Return an isolated Slurm runner with one array and one bounded sinfo query.""" + return FakeSlurmRunner( + arrays=( + FakeSlurmArray( + tasks=( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)), + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + ), + ) + ), + ), + sinfo_responses={ + ("sinfo", "--noheader", "--format=%G"): FakeCommandResponse( + stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + ) + }, + ) + + +@pytest.fixture +def fake_plugin_overlay() -> Path: + """Return the installed-layout fake Data Designer plugin overlay.""" + return TEST_DIRECTORY / "fixtures" / "fake_plugin_overlay" + + +@pytest.fixture +def fake_logical_endpoint() -> FakeLogicalEndpoint: + """Return an isolated two-backend logical endpoint.""" + return FakeLogicalEndpoint( + "http://127.0.0.1:31000", + ( + FakeVllmBackend("http://127.0.0.1:31001", rank=0), + FakeVllmBackend("http://127.0.0.1:31002", rank=1), + ), + ) diff --git a/packages/data-designer-slurm/tests/contracts/conftest.py b/packages/data-designer-slurm/tests/contracts/conftest.py deleted file mode 100644 index 44a74d99c..000000000 --- a/packages/data-designer-slurm/tests/contracts/conftest.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from pathlib import Path - -import pytest - -from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog -from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan - -GOLDEN_DIR = Path(__file__).parent / "golden" - - -@pytest.fixture -def authored_run() -> DataDesignerSlurmConfig: - return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run.json").read_text()) - - -@pytest.fixture -def authored_run_single() -> DataDesignerSlurmConfig: - return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run_single.json").read_text()) - - -@pytest.fixture -def profile_catalog() -> SlurmProfileCatalog: - return SlurmProfileCatalog.model_validate_json((GOLDEN_DIR / "profile_catalog.json").read_text()) - - -@pytest.fixture -def dependency_lock() -> ResolvedDependencyLock: - return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock.json").read_text()) - - -@pytest.fixture -def dependency_lock_single() -> ResolvedDependencyLock: - return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock_single.json").read_text()) - - -@pytest.fixture -def single_node_plan() -> ResolvedSlurmRunPlan: - return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "single_node_plan.json").read_text()) - - -@pytest.fixture -def multi_node_plan() -> ResolvedSlurmRunPlan: - return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "multi_node_plan.json").read_text()) diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA new file mode 100644 index 000000000..1e52ff779 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA @@ -0,0 +1,4 @@ +Metadata-Version: 2.3 +Name: fake-data-designer-plugin +Version: 1.0.0 +Requires-Dist: data-designer diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt new file mode 100644 index 000000000..19e752123 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[data_designer.plugins] +fake-slurm-column = fake_data_designer_plugin.plugin:plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt new file mode 100644 index 000000000..7ca87bde3 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt @@ -0,0 +1 @@ +fake_data_designer_plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py new file mode 100644 index 000000000..aa059a7a1 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal Data Designer plugin used by Slurm client-environment tests.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py new file mode 100644 index 000000000..2dc7726b3 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Literal + +from data_designer.config.base import SingleColumnConfig +from data_designer.plugins import Plugin, PluginType + + +class FakePluginConfig(SingleColumnConfig): + """Minimal custom column configuration.""" + + column_type: Literal["fake-slurm-column"] = "fake-slurm-column" + + +class FakePluginImplementation: + """Minimal loadable implementation for entry-point verification.""" + + def generate(self, data: dict[str, object]) -> dict[str, object]: + """Return the provided record unchanged.""" + return data + + +plugin = Plugin( + config_qualified_name="fake_data_designer_plugin.plugin.FakePluginConfig", + impl_qualified_name="fake_data_designer_plugin.plugin.FakePluginImplementation", + plugin_type=PluginType.COLUMN_GENERATOR, +) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py new file mode 100644 index 000000000..a73536ea8 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic test doubles for the Slurm integration.""" + +from __future__ import annotations + +from slurm_test_fakes.clock import FakeClock +from slurm_test_fakes.dependencies import FakeDependencyInstaller, FakeDependencyResolver +from slurm_test_fakes.serving import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend +from slurm_test_fakes.slurm import ( + FakeCommandResponse, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, +) + +__all__ = [ + "FakeClock", + "FakeCommandResponse", + "FakeDependencyInstaller", + "FakeDependencyResolver", + "FakeLogicalEndpoint", + "FakeServingState", + "FakeSlurmArray", + "FakeSlurmRunner", + "FakeSlurmTask", + "FakeVllmBackend", +] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py new file mode 100644 index 000000000..45485aae0 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timedelta + + +@dataclass +class FakeClock: + """Clock advanced only by explicit test input.""" + + current_time: datetime + monotonic_time: float = 0.0 + sleep_calls: list[float] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.current_time.tzinfo is None or self.current_time.utcoffset() != timedelta(0): + raise ValueError("current_time must be in UTC") + if self.monotonic_time < 0: + raise ValueError("monotonic_time must not be negative") + + def now(self) -> datetime: + """Return the controlled wall-clock time.""" + return self.current_time + + def monotonic(self) -> float: + """Return the controlled monotonic time.""" + return self.monotonic_time + + def sleep(self, seconds: float) -> None: + """Record a sleep and advance both clocks without blocking.""" + self.advance(seconds) + self.sleep_calls.append(seconds) + + def advance(self, seconds: float) -> None: + """Advance both clocks by a non-negative duration.""" + if seconds < 0: + raise ValueError("seconds must not be negative") + self.current_time += timedelta(seconds=seconds) + self.monotonic_time += seconds diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py new file mode 100644 index 000000000..ac0a8a4cb --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from pathlib import Path +from typing import Generic, TypeVar + +RequestT = TypeVar("RequestT") +ResolutionT = TypeVar("ResolutionT") +LockT = TypeVar("LockT") +InstallationT = TypeVar("InstallationT") + + +class FakeDependencyResolver(Generic[RequestT, ResolutionT]): + """Return or raise exact scripted dependency-resolution outcomes.""" + + def __init__(self, responses: Iterable[tuple[RequestT, ResolutionT | BaseException]]) -> None: + self._responses = deque(responses) + self.calls: list[RequestT] = [] + + def resolve(self, request: RequestT) -> ResolutionT: + """Resolve one expected request without consulting package indexes.""" + self.calls.append(request) + if not self._responses: + raise AssertionError("unexpected dependency resolution") + expected, response = self._responses.popleft() + if request != expected: + raise AssertionError(f"expected dependency request {expected!r}, got {request!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted resolution was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency resolutions remain") + + +class FakeDependencyInstaller(Generic[LockT, InstallationT]): + """Return or raise exact scripted dependency-install outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[tuple[LockT, Path], InstallationT | BaseException]], + ) -> None: + self._responses = deque(responses) + self.calls: list[tuple[LockT, Path]] = [] + + def install(self, lock: LockT, target: Path) -> InstallationT: + """Install one expected lock without invoking an installer process.""" + call = (lock, target) + self.calls.append(call) + if not self._responses: + raise AssertionError("unexpected dependency installation") + expected, response = self._responses.popleft() + if call != expected: + raise AssertionError(f"expected dependency installation {expected!r}, got {call!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted installation was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency installations remain") diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch new file mode 100644 index 000000000..5a78c97b7 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +#SBATCH --job-name=dd-two-model +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=3 +#SBATCH --time=03:55:00 +#SBATCH --array=0-1%2 +#SBATCH --gres=gpu:8 +set -Eeuo pipefail + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" +readonly DD_PLAN_SHA256="dbcae1da6ce8ef6799faf2add8e08b1b093390ae08a37dd13e336da522b9a3fd" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + printf '%s %s\n' "$1" "$2" | sha256sum --check --status - +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +install -d -m 0700 "${DD_ATTEMPT_DIR}" +readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" + +source "${DD_RUNTIME_DIR}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch new file mode 100644 index 000000000..e10478f7d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +#SBATCH --job-name=data-designer +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --time=03:55:00 +#SBATCH --array=0 +#SBATCH --gres=gpu:8 +set -Eeuo pipefail + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" +readonly DD_PLAN_SHA256="d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + printf '%s %s\n' "$1" "$2" | sha256sum --check --status - +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +install -d -m 0700 "${DD_ATTEMPT_DIR}" +readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" + +source "${DD_RUNTIME_DIR}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt new file mode 100644 index 000000000..7ac0e7937 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt @@ -0,0 +1,6 @@ +4101_0|TIMEOUT|0:125 +4101_1|NODE_FAIL|1:0 +4101_2|PREEMPTED|0:0 +4101_3|REQUEUED|0:0 +4101_4|OUT_OF_MEMORY|0:125 +4101_5|CANCELLED by 1234|0:15 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt new file mode 100644 index 000000000..17686ec3f --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt @@ -0,0 +1,2 @@ +4101_0|COMPLETED|0:0 +4101_1|FAILED|1:0 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt new file mode 100644 index 000000000..25175f090 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt @@ -0,0 +1 @@ +gpu:2 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt new file mode 100644 index 000000000..80e935355 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt @@ -0,0 +1,2 @@ +4101_0|PENDING +4101_1|RUNNING diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt new file mode 100644 index 000000000..670c5234d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt @@ -0,0 +1 @@ +malformed scheduler output diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py new file mode 100644 index 000000000..0bc05aed3 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from enum import Enum + +import httpx + + +class FakeServingState(str, Enum): + """Lifecycle states shared by the serving test doubles.""" + + CREATED = "created" + STARTING = "starting" + READY = "ready" + FAILED = "failed" + STOPPED = "stopped" + + +class FakeVllmBackend: + """In-memory vLLM boundary with explicit lifecycle and responses.""" + + def __init__( + self, + endpoint: str, + *, + rank: int, + responses: Iterable[httpx.Response | BaseException] = (), + ) -> None: + if rank < 0: + raise ValueError("rank must not be negative") + self.endpoint = endpoint + self.rank = rank + self.state = FakeServingState.CREATED + self.failure_reason: str | None = None + self.requests: list[httpx.Request] = [] + self.start_calls = 0 + self.cleanup_calls = 0 + self._responses = deque(responses) + + def start(self) -> None: + """Enter the explicit startup state.""" + self.start_calls += 1 + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"backend cannot start from {self.state.value}") + self.state = FakeServingState.STARTING + + def mark_ready(self) -> None: + """Mark startup complete.""" + if self.state is not FakeServingState.STARTING: + raise RuntimeError("backend must be starting before it becomes ready") + self.state = FakeServingState.READY + + def fail(self, reason: str) -> None: + """Inject a backend or rank failure.""" + if not reason: + raise ValueError("failure reason must not be empty") + self.failure_reason = reason + self.state = FakeServingState.FAILED + + def queue_response(self, response: httpx.Response | BaseException) -> None: + """Queue one explicit generation response or failure.""" + self._responses.append(response) + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle a health or generation request without opening a socket.""" + self.requests.append(request) + if request.url.path == "/health": + status_code = 200 if self.state is FakeServingState.READY else 503 + return httpx.Response(status_code, request=request) + if self.state is not FakeServingState.READY: + return httpx.Response(503, json={"error": "backend_unavailable"}, request=request) + if not self._responses: + return httpx.Response( + 200, + json={"backend": self.endpoint, "rank": self.rank}, + request=request, + ) + response = self._responses.popleft() + if isinstance(response, BaseException): + raise response + response.request = request + return response + + def cleanup(self) -> None: + """Stop the backend once while allowing cleanup re-entry.""" + self.cleanup_calls += 1 + self.state = FakeServingState.STOPPED + + +class FakeLogicalEndpoint: + """In-memory logical endpoint over one or more fake vLLM backends.""" + + def __init__(self, endpoint: str, backends: Iterable[FakeVllmBackend]) -> None: + self.endpoint = endpoint + self.backends = tuple(backends) + if not self.backends: + raise ValueError("logical endpoints require at least one backend") + self.state = FakeServingState.CREATED + self.published_endpoint: str | None = None + self.requests: list[httpx.Request] = [] + self.publish_calls = 0 + self.cleanup_calls = 0 + self.failure_reason: str | None = None + self._next_backend = 0 + self._publication_failure: Exception | None = None + + def start(self) -> None: + """Start every backend and the readiness aggregator.""" + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"logical endpoint cannot start from {self.state.value}") + for backend in self.backends: + backend.start() + self.state = FakeServingState.STARTING + + def refresh_readiness(self) -> FakeServingState: + """Aggregate backend readiness with coordinated failure semantics.""" + if self.state is FakeServingState.STOPPED: + return self.state + if self._publication_failure is not None and self.publish_calls: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + return self.state + backend_states = {backend.state for backend in self.backends} + if FakeServingState.FAILED in backend_states: + self.state = FakeServingState.FAILED + self.failure_reason = "backend_failed" + elif backend_states == {FakeServingState.READY}: + self.state = FakeServingState.READY + else: + self.state = FakeServingState.STARTING + return self.state + + def publish(self) -> str: + """Publish the logical endpoint after every backend is ready.""" + if self.refresh_readiness() is not FakeServingState.READY: + raise RuntimeError("logical endpoint is not ready") + self.publish_calls += 1 + if self._publication_failure is not None: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + raise self._publication_failure + self.published_endpoint = self.endpoint + return self.endpoint + + def script_publication_failure(self, error: Exception) -> None: + """Inject one explicit endpoint-publication failure.""" + self._publication_failure = error + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle an aggregated health request or forward in round-robin order.""" + self.requests.append(request) + state = self.refresh_readiness() + if request.url.path == "/health": + status_code = 200 if state is FakeServingState.READY and self.published_endpoint else 503 + return httpx.Response(status_code, request=request) + if state is FakeServingState.FAILED: + return httpx.Response(503, json={"error": self.failure_reason}, request=request) + if state is not FakeServingState.READY or self.published_endpoint is None: + return httpx.Response(503, json={"error": "endpoint_unavailable"}, request=request) + response: httpx.Response | None = None + for _ in self.backends: + backend = self.backends[self._next_backend] + self._next_backend = (self._next_backend + 1) % len(self.backends) + response = backend.handle(request) + if response.status_code != 429: + return response + assert response is not None + return response + + def cleanup(self) -> None: + """Stop every backend and unpublish the process-local endpoint idempotently.""" + self.cleanup_calls += 1 + if self.state is FakeServingState.STOPPED: + return + for backend in self.backends: + backend.cleanup() + self.published_endpoint = None + self.state = FakeServingState.STOPPED diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py new file mode 100644 index 000000000..785af0630 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import re +import subprocess +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.state import SchedulerIdentity + +_JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--format=%i|%T") +_SACCT_REQUIRED_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") + + +@dataclass(frozen=True) +class FakeCommandResponse: + """One deterministic subprocess response.""" + + stdout: str = "" + stderr: str = "" + returncode: int = 0 + + +@dataclass +class FakeSlurmTask: + """Mutable scheduler views for one canonical array-task identity.""" + + scheduler: SchedulerIdentity + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + +@dataclass +class FakeSlurmArray: + """A deterministic array submission exposed through Slurm command output.""" + + tasks: tuple[FakeSlurmTask, ...] + + def __post_init__(self) -> None: + if not self.tasks: + raise ValueError("fake Slurm arrays require at least one task") + job_ids = {task.scheduler.array_job_id for task in self.tasks} + task_ids = [task.scheduler.array_task_id for task in self.tasks] + if len(job_ids) != 1: + raise ValueError("fake Slurm array tasks must share one array job ID") + if len(task_ids) != len(set(task_ids)): + raise ValueError("fake Slurm array task IDs must be unique") + + @property + def array_job_id(self) -> int: + """Return the canonical job ID shared by the array tasks.""" + return self.tasks[0].scheduler.array_job_id + + +class FakeSlurmRunner: + """Stateful fake that copies its configured arrays before exposing Slurm commands.""" + + def __init__( + self, + arrays: Iterable[FakeSlurmArray] = (), + *, + sinfo_responses: Mapping[tuple[str, ...], FakeCommandResponse] | None = None, + ) -> None: + self._pending_arrays = deque(copy.deepcopy(tuple(arrays))) + self._submitted_arrays: dict[int, FakeSlurmArray] = {} + self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} + self._sinfo_responses = dict(sinfo_responses or {}) + self.calls: list[tuple[str, ...]] = [] + + def run(self, command: Sequence[str], *, check: bool = False) -> subprocess.CompletedProcess[str]: + """Run one fake Slurm command and optionally raise on failure.""" + if not command: + raise ValueError("command must not be empty") + argv = tuple(command) + self.calls.append(argv) + command_name = Path(argv[0]).name + scripted = self._scripted_responses.get(command_name) + if scripted: + response = scripted.popleft() + else: + response = self._dispatch(command_name, argv) + + completed = subprocess.CompletedProcess( + args=list(argv), + returncode=response.returncode, + stdout=response.stdout, + stderr=response.stderr, + ) + if check and response.returncode: + raise subprocess.CalledProcessError( + response.returncode, + list(argv), + output=response.stdout, + stderr=response.stderr, + ) + return completed + + def script_next(self, command_name: str, response: FakeCommandResponse) -> None: + """Inject one explicit response before the command's stateful behavior.""" + self._scripted_responses.setdefault(command_name, deque()).append(response) + + def set_task_state( + self, + scheduler: SchedulerIdentity, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable queue and accounting states.""" + task = self._find_task(scheduler) + task.queue_state = queue_state + task.accounting_state = accounting_state + task.exit_code = exit_code + + def assert_scripts_consumed(self) -> None: + """Assert that no scripted command response remains.""" + remaining = sum(len(responses) for responses in self._scripted_responses.values()) + if remaining: + raise AssertionError(f"{remaining} scripted Slurm responses remain") + + def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResponse: + handlers = { + "sacct": self._run_sacct, + "sbatch": self._run_sbatch, + "scancel": self._run_scancel, + "sinfo": self._run_sinfo, + "squeue": self._run_squeue, + } + try: + handler = handlers[command_name] + except KeyError: + raise AssertionError(f"unexpected command {command_name!r}") from None + return handler(argv) + + def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: + if not self._pending_arrays: + return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) + array = self._pending_arrays.popleft() + self._submitted_arrays[array.array_job_id] = array + if "--parsable" in argv[1:]: + return FakeCommandResponse(stdout=f"{array.array_job_id}\n") + return FakeCommandResponse(stdout=f"Submitted batch job {array.array_job_id}\n") + + def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SQUEUE_REQUIRED_ARGUMENTS) + rows = [ + f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.queue_state}" + for task in self._selected_submitted_tasks(argv) + if task.queue_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + + def _run_sacct(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SACCT_REQUIRED_ARGUMENTS) + rows = [ + (f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.accounting_state}|{task.exit_code}") + for task in self._selected_submitted_tasks(argv) + if task.accounting_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + + def _run_scancel(self, argv: tuple[str, ...]) -> FakeCommandResponse: + targets = tuple(argument for argument in argv[1:] if not argument.startswith("-")) + if len(targets) != 1: + return FakeCommandResponse(stderr="expected one cancellation target\n", returncode=1) + try: + tasks = self._tasks_for_target(targets[0]) + except (KeyError, ValueError): + return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) + for task in tasks: + task.queue_state = None + task.accounting_state = "CANCELLED" + task.exit_code = "0:15" + return FakeCommandResponse() + + def _run_sinfo(self, argv: tuple[str, ...]) -> FakeCommandResponse: + key = ("sinfo", *argv[1:]) + try: + return self._sinfo_responses[key] + except KeyError: + raise AssertionError(f"unexpected sinfo query {key!r}") from None + + @staticmethod + def _require_arguments(argv: tuple[str, ...], required: tuple[str, ...]) -> None: + missing = tuple(argument for argument in required if argument not in argv[1:]) + if missing: + raise AssertionError(f"missing required arguments {missing!r} in {argv!r}") + + def _tasks_for_target(self, target: str) -> tuple[FakeSlurmTask, ...]: + if "_" not in target: + return self._submitted_arrays[int(target)].tasks + job_id, task_id = target.split("_", 1) + scheduler = SchedulerIdentity(array_job_id=int(job_id), array_task_id=int(task_id)) + return (self._find_task(scheduler),) + + def _find_task(self, scheduler: SchedulerIdentity) -> FakeSlurmTask: + array = self._submitted_arrays[scheduler.array_job_id] + for task in array.tasks: + if task.scheduler == scheduler: + return task + raise KeyError(scheduler) + + def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: + return sorted( + (task for array in self._submitted_arrays.values() for task in array.tasks), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) + + def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: + selectors: list[str] = [] + for index, argument in enumerate(argv[1:]): + if argument in {"-j", "--jobs"}: + try: + selectors.extend(argv[index + 2].split(",")) + except IndexError: + raise AssertionError(f"missing job selector in {argv!r}") from None + elif argument.startswith("--jobs="): + selectors.extend(argument.partition("=")[2].split(",")) + if not selectors: + return self._sorted_submitted_tasks() + + selected: dict[SchedulerIdentity, FakeSlurmTask] = {} + for selector in selectors: + if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: + raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") + try: + tasks = self._tasks_for_target(selector) + except KeyError: + continue + selected.update((task.scheduler, task) for task in tasks) + return sorted( + selected.values(), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py new file mode 100644 index 000000000..b160d5c21 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime, timedelta, timezone + +import pytest + +from slurm_test_fakes import FakeClock + + +def test_fake_clock_advances_without_blocking(fake_clock: FakeClock) -> None: + fake_clock.sleep(2.5) + fake_clock.advance(1.5) + + assert fake_clock.now() == datetime(2026, 8, 18, 12, 0, 4, tzinfo=timezone.utc) + assert fake_clock.monotonic() == 104.0 + assert fake_clock.sleep_calls == [2.5] + + +@pytest.mark.parametrize( + "clock", + ( + pytest.param(lambda: FakeClock(datetime(2026, 8, 18, 12)), id="naive-datetime"), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 13, tzinfo=timezone(timedelta(hours=1)))), + id="non-utc-offset", + ), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=-1), + id="negative-monotonic", + ), + ), +) +def test_fake_clock_rejects_ambient_or_invalid_time(clock: Callable[[], FakeClock]) -> None: + with pytest.raises(ValueError): + clock() + + +def test_fake_clock_rejects_negative_advances(fake_clock: FakeClock) -> None: + with pytest.raises(ValueError, match="must not be negative"): + fake_clock.advance(-0.1) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py new file mode 100644 index 000000000..d32aefa5b --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, ImageInspectionRecord +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + + +def test_shared_contract_fixtures_load_canonical_records( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, + client_result: ClientResult, + run_manifest: RunManifest, + shard_manifests: tuple[ShardManifest, ...], + attempt_manifest: AttemptManifest, + attempt_readiness: AttemptReadiness, + finalization_client_result: ClientResult, + candidate_output_manifest: CandidateOutputManifest, + shard_winner: ShardWinner, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, + client_image_inspection: ImageInspectionRecord, + serving_image_inspection: ImageInspectionRecord, + collection_plan: CollectionPlan, + accounting_lag_observation: SchedulerObservation, +) -> None: + assert single_node_plan.run_id == "run-single" + assert multi_node_plan.run_id == "run-001" + assert client_result.run_id == "run-001" + assert run_manifest.run_id == "run-single" + assert len(shard_manifests) == 1 + assert attempt_manifest.attempt_id == "attempt-0001" + assert attempt_readiness.revision == 1 + assert finalization_client_result.outcome.value == "complete" + assert candidate_output_manifest.winner_eligible + assert shard_winner.attempt_id == "attempt-0001" + assert benchmark_config.name == "generator-scaling" + assert benchmark_manifest.benchmark_id == benchmark_report.benchmark_id + assert client_image_inspection.inspection.kind == "client" + assert serving_image_inspection.inspection.kind == "serving" + assert collection_plan.collection_id == "collection-0001" + assert accounting_lag_observation.state.value == "accounting_lag" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py new file mode 100644 index 000000000..324a8948d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +from importlib.metadata import entry_points +from pathlib import Path + +import pytest + +from data_designer.plugins import PluginType +from data_designer.plugins.registry import PluginRegistry +from data_designer.slurm.config import ClientDependencies, DataDesignerSlurmConfig +from data_designer.slurm.planning import ResolvedDependencyLock +from slurm_test_fakes import FakeDependencyInstaller, FakeDependencyResolver + + +def test_fake_dependency_resolver_scripts_compatible_incompatible_and_missing_cases( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, +) -> None: + compatible = authored_run.client.dependencies + incompatible = ClientDependencies(requirements=["fake-data-designer-plugin==2.0.0"]) + missing = ClientDependencies(requirements=["missing-plugin==1.0.0"]) + resolver: FakeDependencyResolver[ClientDependencies, ResolvedDependencyLock] = FakeDependencyResolver( + ( + (compatible, dependency_lock), + (incompatible, ValueError("incompatible Data Designer version")), + (missing, FileNotFoundError("dependency artifact is missing")), + ) + ) + + assert resolver.resolve(compatible) is dependency_lock + with pytest.raises(ValueError, match="incompatible"): + resolver.resolve(incompatible) + with pytest.raises(FileNotFoundError, match="missing"): + resolver.resolve(missing) + + assert resolver.calls == [compatible, incompatible, missing] + resolver.assert_complete() + + +def test_fake_dependency_installer_scripts_success_and_digest_mismatch( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + first_target = tmp_path / "compatible" + second_target = tmp_path / "digest-mismatch" + mismatched_lock = dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + installer = FakeDependencyInstaller[ResolvedDependencyLock, tuple[object, ...]]( + ( + ((dependency_lock, first_target), dependency_lock.overlay_packages), + ((mismatched_lock, second_target), ValueError("lock digest mismatch")), + ) + ) + + assert installer.install(dependency_lock, first_target) == dependency_lock.overlay_packages + with pytest.raises(ValueError, match="digest mismatch"): + installer.install(mismatched_lock, second_target) + + assert installer.calls == [(dependency_lock, first_target), (mismatched_lock, second_target)] + installer.assert_complete() + + +def test_dependency_fakes_raise_cancellation_signals( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + request = ClientDependencies(requirements=[]) + resolver = FakeDependencyResolver(((request, KeyboardInterrupt()),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path), SystemExit(2)),)) + + with pytest.raises(KeyboardInterrupt): + resolver.resolve(request) + with pytest.raises(SystemExit, match="2"): + installer.install(dependency_lock, tmp_path) + + +def test_dependency_fakes_reject_unexpected_calls( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + expected = ClientDependencies(requirements=[]) + unexpected = ClientDependencies(requirements=["unexpected==1.0.0"]) + resolver = FakeDependencyResolver(((expected, dependency_lock),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path / "expected"), "installed"),)) + + with pytest.raises(AssertionError, match="expected dependency request"): + resolver.resolve(unexpected) + with pytest.raises(AssertionError, match="expected dependency installation"): + installer.install(dependency_lock, tmp_path / "unexpected") + + +def test_fake_plugin_overlay_supports_real_entry_point_discovery( + fake_plugin_overlay: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(str(fake_plugin_overlay)) + importlib.invalidate_caches() + monkeypatch.setattr("data_designer.plugins.registry.PLUGINS_DISABLED", False) + PluginRegistry.reset() + try: + discovered = tuple( + entry_point + for entry_point in entry_points(group="data_designer.plugins") + if entry_point.name == "fake-slurm-column" + ) + assert len(discovered) == 1 + assert discovered[0].dist is not None + assert discovered[0].dist.version == "1.0.0" + + plugin = discovered[0].load() + registry = PluginRegistry() + + assert plugin.name == "fake-slurm-column" + assert plugin.plugin_type is PluginType.COLUMN_GENERATOR + assert registry.get_plugin("fake-slurm-column") == plugin + finally: + PluginRegistry.reset() diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py new file mode 100644 index 000000000..c7d119389 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import posixpath +from pathlib import Path + +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "rendered" + + +def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + _assert_script_matches_plan( + single_node_plan, + "single_node.sbatch", + expected_fixture_sha256="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", + ) + _assert_script_matches_plan( + multi_node_plan, + "multi_node.sbatch", + expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + ) + + +def test_rendered_scripts_contain_no_site_or_user_data() -> None: + contents = "\n".join(path.read_text().casefold() for path in GOLDEN_DIRECTORY.glob("*.sbatch")) + + for forbidden in ( + "nvidia", + "cluster", + "login", + "lustre", + "/home/", + "/users/", + ): + assert forbidden not in contents + + +def _assert_script_matches_plan( + plan: ResolvedSlurmRunPlan, + filename: str, + *, + expected_fixture_sha256: str, +) -> None: + script = (GOLDEN_DIRECTORY / filename).read_text() + node_indices = ( + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + ) + node_count = max(node_indices) + 1 + array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") + run_root = posixpath.dirname(plan.authored_config.path) + + assert f"#SBATCH --job-name={plan.submission.job_name}\n" in script + assert f"#SBATCH --account={plan.submission.account}\n" in script + assert f"#SBATCH --partition={plan.submission.partition}\n" in script + assert f"#SBATCH --nodes={node_count}\n" in script + assert f"#SBATCH --time={plan.submission.time_limit}\n" in script + assert f"#SBATCH --array={array}\n" in script + assert f"#SBATCH --gres=gpu:{plan.resolved_gpus_per_node}\n" in script + assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"\n' in script + assert f'readonly DD_RUNTIME_SHA256="{plan.runtime_bundle.sha256}"\n' in script + assert f'readonly DD_PLAN="{plan_path}"\n' in script + assert f'readonly DD_PLAN_SHA256="{plan.compute_sha256()}"\n' in script + assert f'readonly DD_RUN_ROOT="{run_root}"\n' in script + assert script.count("dd_slurm_run_allocation") == 1 + assert script.endswith("\n") + assert hashlib.sha256(script.encode()).hexdigest() == expected_fixture_sha256 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py new file mode 100644 index 000000000..cf92ca840 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.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 httpx +import pytest + +from slurm_test_fakes import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend + + +def _client(endpoint: FakeLogicalEndpoint) -> httpx.Client: + return httpx.Client(transport=httpx.MockTransport(endpoint.handle), base_url=endpoint.endpoint) + + +def _start_and_publish(endpoint: FakeLogicalEndpoint) -> None: + endpoint.start() + for backend in endpoint.backends: + backend.mark_ready() + endpoint.publish() + + +def test_single_backend_startup_readiness_and_endpoint_publication() -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + + endpoint.start() + with _client(endpoint) as client: + assert client.get("/health").status_code == 503 + + backend.mark_ready() + assert endpoint.publish() == endpoint.endpoint + assert client.get("/health").status_code == 200 + assert client.post("/v1/chat/completions", json={"prompt": "fixture"}).json() == { + "backend": backend.endpoint, + "rank": 0, + } + + +def test_multiple_backends_receive_requests_in_deterministic_order( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + with _client(fake_logical_endpoint) as client: + responses = [client.post("/v1/chat/completions").json() for _ in range(3)] + + assert [response["rank"] for response in responses] == [0, 1, 0] + assert [len(backend.requests) for backend in fake_logical_endpoint.backends] == [2, 1] + + +def test_partial_multi_backend_startup_is_not_ready_and_cleans_up_once( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + fake_logical_endpoint.backends[0].mark_ready() + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.STARTING + with pytest.raises(RuntimeError, match="not ready"): + fake_logical_endpoint.publish() + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +@pytest.mark.parametrize("retry_after", ("2", None)) +def test_logical_endpoint_preserves_overload_response(retry_after: str | None) -> None: + headers = {"Retry-After": retry_after} if retry_after is not None else {} + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers=headers, json={"error": "overloaded"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers.get("Retry-After") == retry_after + assert response.json() == {"error": "overloaded"} + + +def test_logical_endpoint_retries_overload_against_another_backend() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, json={"error": "first overloaded"}),), + ) + second = FakeVllmBackend("http://127.0.0.1:31002", rank=1) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 200 + assert response.json()["rank"] == 1 + assert [len(backend.requests) for backend in endpoint.backends] == [1, 1] + + +def test_logical_endpoint_returns_the_final_overload_response() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers={"Retry-After": "1"}),), + ) + second = FakeVllmBackend( + "http://127.0.0.1:31002", + rank=1, + responses=(httpx.Response(429, headers={"Retry-After": "3"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers["Retry-After"] == "3" + + +def test_backend_rank_failure_fails_the_logical_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + fake_logical_endpoint.backends[1].fail("rank exited") + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert response.status_code == 503 + assert response.json() == {"error": "backend_failed"} + + +def test_endpoint_publication_failure_is_explicit_and_cleanup_remains_idempotent( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + for backend in fake_logical_endpoint.backends: + backend.mark_ready() + fake_logical_endpoint.script_publication_failure(RuntimeError("publication failed")) + + with pytest.raises(RuntimeError, match="publication failed"): + fake_logical_endpoint.publish() + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert fake_logical_endpoint.failure_reason == "endpoint_publication_failed" + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + assert response.json() == {"error": "endpoint_publication_failed"} + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +def test_backend_failure_injection_is_explicit() -> None: + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(RuntimeError("scripted backend failure"),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client, pytest.raises(RuntimeError, match="scripted backend failure"): + client.post("/v1/chat/completions") + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_backend_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + if terminal_state is FakeServingState.FAILED: + backend.fail("rank exited") + else: + backend.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + backend.start() + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_logical_endpoint_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + endpoint = FakeLogicalEndpoint( + "http://127.0.0.1:31000", + (FakeVllmBackend("http://127.0.0.1:31001", rank=0),), + ) + if terminal_state is FakeServingState.FAILED: + endpoint.start() + endpoint.backends[0].fail("rank exited") + endpoint.refresh_readiness() + else: + endpoint.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + endpoint.start() + + +def test_cleanup_is_idempotent_and_unpublishes_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + + assert fake_logical_endpoint.state is FakeServingState.STOPPED + assert fake_logical_endpoint.published_endpoint is None + assert fake_logical_endpoint.cleanup_calls == 2 + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py new file mode 100644 index 000000000..8bfb3de0e --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from data_designer.slurm.state import SchedulerIdentity +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" +SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") +SACCT_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") + + +def _submit(runner: FakeSlurmRunner) -> None: + completed = runner.run(("sbatch", "--parsable", "run.sbatch"), check=True) + assert completed.stdout == "4101\n" + + +def test_fake_slurm_runner_models_array_submission_and_active_states( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + observed = fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS), check=True) + + assert observed.stdout == (GOLDEN_DIRECTORY / "squeue_active.txt").read_text() + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "run.sbatch"), + ("squeue", *SQUEUE_ARGUMENTS), + ] + + +def test_fake_slurm_runner_exposes_terminal_accounting_precedence( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=0), + queue_state="RUNNING", + accounting_state="COMPLETED", + ) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ) + + assert "4101_1|RUNNING" in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert ( + fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + == (GOLDEN_DIRECTORY / "sacct_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_models_accounting_lag_and_later_terminal_state( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=0) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state=None, + ) + + assert "4101_0" not in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout == "" + + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + ) + assert "4101_0|COMPLETED|0:0" in fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + + +def test_fake_slurm_runner_filters_job_scoped_queries() -> None: + first = FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + second = FakeSlurmArray( + tasks=( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4201, array_task_id=0), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ), + ) + ) + runner = FakeSlurmRunner((first, second)) + runner.run(("sbatch", "first.sbatch"), check=True) + runner.run(("sbatch", "second.sbatch"), check=True) + + first_queue = runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs", "4101")).stdout + second_accounting = runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=4201")).stdout + + assert first_queue == "4101_0|PENDING\n" + assert second_accounting == "4201_0|FAILED|1:0\n" + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", *SQUEUE_ARGUMENTS, "--jobs", "--noheader"), + ("sacct", *SACCT_ARGUMENTS, "--jobs="), + ), +) +def test_fake_slurm_runner_rejects_malformed_job_selectors( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="malformed job selector"): + fake_slurm_runner.run(command) + + +@pytest.mark.parametrize("target", ("4101", "4101_1")) +def test_fake_slurm_runner_models_cancellation( + fake_slurm_runner: FakeSlurmRunner, + target: str, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("scancel", target), check=True).returncode == 0 + accounting = fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + + expected_tasks = (0, 1) if target == "4101" else (1,) + for task_id in expected_tasks: + assert f"4101_{task_id}|CANCELLED|0:15" in accounting + + +def test_fake_slurm_runner_supports_malformed_output_and_command_failures( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + malformed = (GOLDEN_DIRECTORY / "squeue_malformed.txt").read_text() + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout=malformed)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + + assert fake_slurm_runner.run(("squeue",)).stdout == malformed + assert fake_slurm_runner.run(("sacct",)).returncode == 2 + with pytest.raises(subprocess.CalledProcessError) as error: + fake_slurm_runner.run(("sacct",), check=True) + + assert error.value.stderr == "accounting unavailable\n" + fake_slurm_runner.assert_scripts_consumed() + + +@pytest.mark.parametrize( + ("arguments", "expected_stdout"), + ( + (("run.sbatch",), "Submitted batch job 4101\n"), + (("--parsable", "run.sbatch"), "4101\n"), + ), +) +def test_fake_slurm_runner_matches_sbatch_parsable_mode( + arguments: tuple[str, ...], + expected_stdout: str, +) -> None: + runner = FakeSlurmRunner( + (FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + + assert runner.run(("sbatch", *arguments), check=True).stdout == expected_stdout + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", "--noheader"), + ("sacct", "--noheader", "--format=%i|%State|%ExitCode"), + ), +) +def test_fake_slurm_runner_rejects_underspecified_state_queries( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="missing required arguments"): + fake_slurm_runner.run(command) + + +def test_fake_slurm_runner_exposes_retry_terminal_spellings() -> None: + states = ( + ("TIMEOUT", "0:125"), + ("NODE_FAIL", "1:0"), + ("PREEMPTED", "0:0"), + ("REQUEUED", "0:0"), + ("OUT_OF_MEMORY", "0:125"), + ("CANCELLED by 1234", "0:15"), + ) + runner = FakeSlurmRunner( + ( + FakeSlurmArray( + tasks=tuple( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=task_id), + queue_state=None, + accounting_state=state, + exit_code=exit_code, + ) + for task_id, (state, exit_code) in enumerate(states) + ) + ), + ) + ) + _submit(runner) + + assert ( + runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + == (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_bounds_sinfo_queries(fake_slurm_runner: FakeSlurmRunner) -> None: + response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%G"), check=True) + + assert response.stdout == (GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + with pytest.raises(AssertionError, match="unexpected sinfo query"): + fake_slurm_runner.run(("sinfo", "--all")) + + +def test_fake_slurm_runner_rejects_unscripted_commands_and_submissions( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("sbatch", "second.sbatch")).returncode == 1 + with pytest.raises(AssertionError, match="unexpected command"): + fake_slurm_runner.run(("srun", "hostname")) diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json new file mode 100644 index 000000000..2b5f14067 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "failed" + } + ], + "revision": 4, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:02:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json new file mode 100644 index 000000000..ca41b2d6f --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": null, + "model_alias": "primary", + "ready_backends": 0, + "state": "pending" + } + ], + "revision": 1, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "pending", + "updated_at": "2026-08-18T12:00:02Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json new file mode 100644 index 000000000..5001bc9fb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "failed", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:30Z", + "outcome": "failure", + "reason_code": "endpoint_publication_failed", + "redacted_message": "Logical endpoint publication failed" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "failed" + } + ], + "revision": 3, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:01:30Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json new file mode 100644 index 000000000..7929011bb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:00Z", + "outcome": "failure", + "reason_code": "backend_starting", + "redacted_message": "Backend is starting" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "starting" + } + ], + "revision": 2, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "starting", + "updated_at": "2026-08-18T12:01:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json new file mode 100644 index 000000000..894019149 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "stopped" + } + ], + "revision": 5, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "stopped", + "updated_at": "2026-08-18T12:02:01Z" +} diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index e029b99b4..d7b5d1adf 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -25,7 +25,12 @@ ("run_manifest.json", RunManifest), ("shard_manifest.json", ShardManifest), ("successful_attempt.json", AttemptManifest), + ("single_node_pending_readiness.json", AttemptReadiness), + ("single_node_starting_readiness.json", AttemptReadiness), ("single_node_readiness.json", AttemptReadiness), + ("single_node_publication_failed_readiness.json", AttemptReadiness), + ("single_node_failed_readiness.json", AttemptReadiness), + ("single_node_stopped_readiness.json", AttemptReadiness), ("multi_node_readiness.json", AttemptReadiness), ("failed_attempt.json", AttemptManifest), ("stale_readiness.json", AttemptReadiness), diff --git a/packages/data-designer-slurm/tests/state/test_transition_fixtures.py b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py new file mode 100644 index 000000000..4bc1c7188 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from itertools import pairwise +from pathlib import Path + +from data_designer.slurm.state import AttemptReadiness, validate_readiness_transition + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +TRANSITION_FILES = ( + "single_node_pending_readiness.json", + "single_node_starting_readiness.json", + "single_node_readiness.json", + "single_node_failed_readiness.json", + "single_node_stopped_readiness.json", +) + + +def test_single_node_readiness_transition_fixture_is_canonical() -> None: + records = tuple( + AttemptReadiness.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) for filename in TRANSITION_FILES + ) + + assert tuple(record.revision for record in records) == (1, 2, 3, 4, 5) + for previous, current in pairwise(records): + assert validate_readiness_transition(previous, current) is current + + +def test_endpoint_publication_failure_transition_fixture_is_canonical() -> None: + starting = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_starting_readiness.json").read_text() + ) + failed = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_publication_failed_readiness.json").read_text() + ) + + assert validate_readiness_transition(starting, failed) is failed