From b67821c42eb56ae53780ea9a630f46004090caff Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 24 Aug 2026 14:36:40 -0300 Subject: [PATCH 1/3] test: add deterministic Slurm fakes Add reusable scheduler, dependency, plugin, serving, clock, and state-transition fixtures for local and CI Slurm integration tests. Keep the harness deterministic and isolated from production lane implementations. Signed-off-by: Andre Manoel --- .../data-designer-slurm/tests/conftest.py | 74 ++++++ .../METADATA | 4 + .../entry_points.txt | 2 + .../top_level.txt | 1 + .../fake_data_designer_plugin/__init__.py | 6 + .../fake_data_designer_plugin/plugin.py | 30 +++ .../tests/slurm_test_fakes/__init__.py | 29 +++ .../tests/slurm_test_fakes/clock.py | 42 ++++ .../tests/slurm_test_fakes/dependencies.py | 68 ++++++ .../golden/slurm/sacct_terminal.txt | 2 + .../golden/slurm/sinfo_gres.txt | 1 + .../golden/slurm/squeue_active.txt | 2 + .../golden/slurm/squeue_malformed.txt | 1 + .../tests/slurm_test_fakes/serving.py | 181 ++++++++++++++ .../tests/slurm_test_fakes/slurm.py | 231 ++++++++++++++++++ .../tests/slurm_test_fakes/test_clock.py | 44 ++++ .../slurm_test_fakes/test_dependencies.py | 102 ++++++++ .../tests/slurm_test_fakes/test_serving.py | 185 ++++++++++++++ .../tests/slurm_test_fakes/test_slurm.py | 165 +++++++++++++ .../golden/single_node_failed_readiness.json | 25 ++ .../golden/single_node_pending_readiness.json | 20 ++ ...gle_node_publication_failed_readiness.json | 25 ++ .../single_node_starting_readiness.json | 25 ++ .../golden/single_node_stopped_readiness.json | 25 ++ .../tests/state/test_state_golden_records.py | 5 + .../tests/state/test_transition_fixtures.py | 39 +++ 26 files changed, 1334 insertions(+) create mode 100644 packages/data-designer-slurm/tests/conftest.py create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/clock.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/serving.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/test_transition_fixtures.py diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py new file mode 100644 index 000000000..aea6becbf --- /dev/null +++ b/packages/data-designer-slurm/tests/conftest.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from slurm_test_fakes import ( + FakeClock, + FakeCommandResponse, + FakeLogicalEndpoint, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, + FakeVllmBackend, +) + +from data_designer.slurm.state import SchedulerIdentity + +TEST_DIRECTORY = Path(__file__).parent +SLURM_GOLDEN_DIRECTORY = TEST_DIRECTORY / "slurm_test_fakes" / "golden" / "slurm" + + +@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_array() -> FakeSlurmArray: + """Return a two-task array copied when the fake runner is constructed.""" + return 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", + ), + ) + ) + + +@pytest.fixture +def fake_slurm_runner(fake_slurm_array: FakeSlurmArray) -> FakeSlurmRunner: + """Return an isolated Slurm runner with one array and one bounded sinfo query.""" + return FakeSlurmRunner( + arrays=(fake_slurm_array,), + 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/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/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..75eb9038b --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py @@ -0,0 +1,181 @@ +# 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 FakeServingState.CREATED: + 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: + return + 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 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..7ef9ee142 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -0,0 +1,231 @@ +# 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]+)?$") + + +@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: + del argv + 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 + return FakeCommandResponse(stdout=f"{array.array_job_id}\n") + + def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: + 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: + 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 + + 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_dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py new file mode 100644 index 000000000..2bfbe5e58 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -0,0 +1,102 @@ +# 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 slurm_test_fakes import FakeDependencyInstaller, FakeDependencyResolver + + +def test_fake_dependency_resolver_scripts_compatible_incompatible_and_missing_cases() -> None: + compatible = ("fake-data-designer-plugin==1.0.0",) + incompatible = ("fake-data-designer-plugin==2.0.0",) + missing = ("missing-plugin==1.0.0",) + resolver = FakeDependencyResolver( + ( + (compatible, "resolved-lock"), + (incompatible, ValueError("incompatible Data Designer version")), + (missing, FileNotFoundError("dependency artifact is missing")), + ) + ) + + assert resolver.resolve(compatible) == "resolved-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) -> None: + first_target = tmp_path / "compatible" + second_target = tmp_path / "digest-mismatch" + installer = FakeDependencyInstaller( + ( + (("resolved-lock", first_target), ("fake-data-designer-plugin==1.0.0",)), + (("mismatched-lock", second_target), ValueError("lock digest mismatch")), + ) + ) + + assert installer.install("resolved-lock", first_target) == ("fake-data-designer-plugin==1.0.0",) + with pytest.raises(ValueError, match="digest mismatch"): + installer.install("mismatched-lock", second_target) + + assert installer.calls == [("resolved-lock", first_target), ("mismatched-lock", second_target)] + installer.assert_complete() + + +def test_dependency_fakes_raise_cancellation_signals(tmp_path: Path) -> None: + resolver = FakeDependencyResolver((("cancelled", KeyboardInterrupt()),)) + installer = FakeDependencyInstaller(((("lock", tmp_path), SystemExit(2)),)) + + with pytest.raises(KeyboardInterrupt): + resolver.resolve("cancelled") + with pytest.raises(SystemExit, match="2"): + installer.install("lock", tmp_path) + + +def test_dependency_fakes_reject_unexpected_calls(tmp_path: Path) -> None: + resolver = FakeDependencyResolver((("expected", "lock"),)) + installer = FakeDependencyInstaller(((("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("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_serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py new file mode 100644 index 000000000..cd1411f76 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py @@ -0,0 +1,185 @@ +# 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") + + +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..815bbc7a9 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -0,0 +1,165 @@ +# 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" + + +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", "--noheader"), check=True) + + assert observed.stdout == (GOLDEN_DIRECTORY / "squeue_active.txt").read_text() + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "run.sbatch"), + ("squeue", "--noheader"), + ] + + +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",)).stdout + assert fake_slurm_runner.run(("sacct",)).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",)).stdout + assert "4101_0" not in fake_slurm_runner.run(("sacct",)).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",)).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", "--jobs", "4101")).stdout + second_accounting = runner.run(("sacct", "--jobs=4201")).stdout + + assert first_queue == "4101_0|PENDING\n" + assert second_accounting == "4201_0|FAILED|1:0\n" + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", "--jobs", "--noheader"), + ("sacct", "--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",), 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() + + +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 From 300f0fc1334c7246ead9879e75c30b328b0eb160 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 24 Aug 2026 18:46:47 -0300 Subject: [PATCH 2/3] test: bind Slurm fakes to contracts Promote the merged planning and finalization goldens to package-wide fixtures, exercise dependency doubles with canonical records, and add site-neutral rendered script goldens for downstream lanes. Signed-off-by: Andre Manoel --- .../data-designer-slurm/tests/conftest.py | 118 +++++++++++++++++- .../tests/contracts/conftest.py | 48 ------- .../golden/rendered/multi_node.sbatch | 31 +++++ .../golden/rendered/single_node.sbatch | 31 +++++ .../test_contract_fixtures.py | 39 ++++++ .../slurm_test_fakes/test_dependencies.py | 66 ++++++---- .../slurm_test_fakes/test_rendered_scripts.py | 71 +++++++++++ 7 files changed, 331 insertions(+), 73 deletions(-) delete mode 100644 packages/data-designer-slurm/tests/contracts/conftest.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index aea6becbf..ac15be2a8 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -3,8 +3,10 @@ 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 ( @@ -17,12 +19,126 @@ FakeVllmBackend, ) -from data_designer.slurm.state import SchedulerIdentity +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + RunManifest, + SchedulerIdentity, + 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" +@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 single_node_render_plan(single_node_plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: + return single_node_plan.model_copy( + update={"submission": single_node_plan.submission.model_copy(update={"account": None, "partition": None})} + ) + + +@pytest.fixture +def multi_node_render_plan(multi_node_plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: + return multi_node_plan.model_copy( + update={"submission": multi_node_plan.submission.model_copy(update={"account": None, "partition": None})} + ) + + +@pytest.fixture +def client_result() -> ClientResult: + return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.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.""" 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/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..f4f8b9625 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +#SBATCH --job-name=dd-two-model +#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="e101840b0ba888fd0f2d9feef70cb6bb2a668dc4dff4c2198a0f8ba2e8a1eb0a" +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..07a49a458 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +#SBATCH --job-name=data-designer +#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="541e532f7e668e862fc7732decef887d404d15290bb943f7cf1920746cde8720" +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/test_contract_fixtures.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py new file mode 100644 index 000000000..f61e85cc4 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_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 data_designer.slurm.client import ClientResult +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + RunManifest, + 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, +) -> 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" 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 index 2bfbe5e58..324a8948d 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -11,22 +11,27 @@ 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() -> None: - compatible = ("fake-data-designer-plugin==1.0.0",) - incompatible = ("fake-data-designer-plugin==2.0.0",) - missing = ("missing-plugin==1.0.0",) - resolver = 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, "resolved-lock"), + (compatible, dependency_lock), (incompatible, ValueError("incompatible Data Designer version")), (missing, FileNotFoundError("dependency artifact is missing")), ) ) - assert resolver.resolve(compatible) == "resolved-lock" + assert resolver.resolve(compatible) is dependency_lock with pytest.raises(ValueError, match="incompatible"): resolver.resolve(incompatible) with pytest.raises(FileNotFoundError, match="missing"): @@ -36,42 +41,55 @@ def test_fake_dependency_resolver_scripts_compatible_incompatible_and_missing_ca resolver.assert_complete() -def test_fake_dependency_installer_scripts_success_and_digest_mismatch(tmp_path: Path) -> None: +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" - installer = FakeDependencyInstaller( + mismatched_lock = dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + installer = FakeDependencyInstaller[ResolvedDependencyLock, tuple[object, ...]]( ( - (("resolved-lock", first_target), ("fake-data-designer-plugin==1.0.0",)), - (("mismatched-lock", second_target), ValueError("lock digest mismatch")), + ((dependency_lock, first_target), dependency_lock.overlay_packages), + ((mismatched_lock, second_target), ValueError("lock digest mismatch")), ) ) - assert installer.install("resolved-lock", first_target) == ("fake-data-designer-plugin==1.0.0",) + assert installer.install(dependency_lock, first_target) == dependency_lock.overlay_packages with pytest.raises(ValueError, match="digest mismatch"): - installer.install("mismatched-lock", second_target) + installer.install(mismatched_lock, second_target) - assert installer.calls == [("resolved-lock", first_target), ("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) -> None: - resolver = FakeDependencyResolver((("cancelled", KeyboardInterrupt()),)) - installer = FakeDependencyInstaller(((("lock", tmp_path), SystemExit(2)),)) +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("cancelled") + resolver.resolve(request) with pytest.raises(SystemExit, match="2"): - installer.install("lock", tmp_path) + installer.install(dependency_lock, tmp_path) -def test_dependency_fakes_reject_unexpected_calls(tmp_path: Path) -> None: - resolver = FakeDependencyResolver((("expected", "lock"),)) - installer = FakeDependencyInstaller(((("lock", tmp_path / "expected"), "installed"),)) +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") + resolver.resolve(unexpected) with pytest.raises(AssertionError, match="expected dependency installation"): - installer.install("lock", tmp_path / "unexpected") + installer.install(dependency_lock, tmp_path / "unexpected") def test_fake_plugin_overlay_supports_real_entry_point_discovery( 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..be6b880e8 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -0,0 +1,71 @@ +# 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_scripts_are_deterministic_and_bound_to_canonical_plans( + single_node_render_plan: ResolvedSlurmRunPlan, + multi_node_render_plan: ResolvedSlurmRunPlan, +) -> None: + _assert_script_matches_plan( + single_node_render_plan, + "single_node.sbatch", + expected_sha256="ba4c0d07f9771242e33f29e0bc82dcb83b5502ed12258231548b7037ad8754fd", + ) + _assert_script_matches_plan( + multi_node_render_plan, + "multi_node.sbatch", + expected_sha256="d4ac042c1eac5be34dc467686dc93e9c77c7ec5a4c0a16656b5892a1ca47b4c6", + ) + + +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", + "account", + "partition", + "lustre", + "/home/", + "/users/", + ): + assert forbidden not in contents + + +def _assert_script_matches_plan( + plan: ResolvedSlurmRunPlan, + filename: str, + *, + expected_sha256: str, +) -> None: + script = (GOLDEN_DIRECTORY / filename).read_text() + node_count = max(index for deployment in plan.deployments for index in deployment.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 --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_sha256 From 14f5dee3d89e5d505a67febdd5e0e51220436495 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 24 Aug 2026 18:59:03 -0300 Subject: [PATCH 3/3] test: complete post-F3 Slurm fake coverage Bind rendered scripts and shared fixtures to the merged contracts, tighten scheduler command behavior, and cover retry terminal states and invalid serving restarts. Signed-off-by: Andre Manoel --- .../data-designer-slurm/tests/conftest.py | 82 +++++++++++----- .../golden/rendered/multi_node.sbatch | 4 +- .../golden/rendered/single_node.sbatch | 4 +- .../golden/slurm/sacct_retry_terminal.txt | 6 ++ .../tests/slurm_test_fakes/serving.py | 9 +- .../tests/slurm_test_fakes/slurm.py | 15 ++- .../test_contract_fixtures.py | 17 ++++ .../slurm_test_fakes/test_rendered_scripts.py | 28 +++--- .../tests/slurm_test_fakes/test_serving.py | 29 ++++++ .../tests/slurm_test_fakes/test_slurm.py | 96 ++++++++++++++++--- 10 files changed, 233 insertions(+), 57 deletions(-) create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index ac15be2a8..a46fcb160 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -19,15 +19,23 @@ FakeVllmBackend, ) +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport from data_designer.slurm.client import ClientResult -from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog +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, ) @@ -36,6 +44,7 @@ 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 @@ -78,22 +87,49 @@ def multi_node_plan() -> ResolvedSlurmRunPlan: @pytest.fixture -def single_node_render_plan(single_node_plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: - return single_node_plan.model_copy( - update={"submission": single_node_plan.submission.model_copy(update={"account": None, "partition": None})} +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 multi_node_render_plan(multi_node_plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: - return multi_node_plan.model_copy( - update={"submission": multi_node_plan.submission.model_copy(update={"account": None, "partition": None})} +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 client_result() -> ClientResult: - return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.json").read_text()) +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 @@ -146,24 +182,20 @@ def fake_clock() -> FakeClock: @pytest.fixture -def fake_slurm_array() -> FakeSlurmArray: - """Return a two-task array copied when the fake runner is constructed.""" - return 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", - ), - ) - ) - - -@pytest.fixture -def fake_slurm_runner(fake_slurm_array: FakeSlurmArray) -> FakeSlurmRunner: +def fake_slurm_runner() -> FakeSlurmRunner: """Return an isolated Slurm runner with one array and one bounded sinfo query.""" return FakeSlurmRunner( - arrays=(fake_slurm_array,), + 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() 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 index f4f8b9625..5a78c97b7 100644 --- 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 @@ -1,5 +1,7 @@ #!/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 @@ -9,7 +11,7 @@ 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="e101840b0ba888fd0f2d9feef70cb6bb2a668dc4dff4c2198a0f8ba2e8a1eb0a" +readonly DD_PLAN_SHA256="dbcae1da6ce8ef6799faf2add8e08b1b093390ae08a37dd13e336da522b9a3fd" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" 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 index 07a49a458..e10478f7d 100644 --- 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 @@ -1,5 +1,7 @@ #!/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 @@ -9,7 +11,7 @@ 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="541e532f7e668e862fc7732decef887d404d15290bb943f7cf1920746cde8720" +readonly DD_PLAN_SHA256="d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" 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/serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py index 75eb9038b..0bc05aed3 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py @@ -44,8 +44,9 @@ def __init__( def start(self) -> None: """Enter the explicit startup state.""" self.start_calls += 1 - if self.state is FakeServingState.CREATED: - self.state = FakeServingState.STARTING + 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.""" @@ -110,7 +111,7 @@ def __init__(self, endpoint: str, backends: Iterable[FakeVllmBackend]) -> None: def start(self) -> None: """Start every backend and the readiness aggregator.""" if self.state is not FakeServingState.CREATED: - return + raise RuntimeError(f"logical endpoint cannot start from {self.state.value}") for backend in self.backends: backend.start() self.state = FakeServingState.STARTING @@ -171,7 +172,7 @@ def handle(self, request: httpx.Request) -> httpx.Response: return response def cleanup(self) -> None: - """Stop every backend and unpublish the endpoint idempotently.""" + """Stop every backend and unpublish the process-local endpoint idempotently.""" self.cleanup_calls += 1 if self.state is FakeServingState.STOPPED: return diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 7ef9ee142..785af0630 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -14,6 +14,8 @@ 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) @@ -139,14 +141,16 @@ def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResp return handler(argv) def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: - del argv 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 - return FakeCommandResponse(stdout=f"{array.array_job_id}\n") + 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) @@ -155,6 +159,7 @@ def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: 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) @@ -183,6 +188,12 @@ def _run_sinfo(self, argv: tuple[str, ...]) -> FakeCommandResponse: 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 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 index f61e85cc4..d32aefa5b 100644 --- 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 @@ -3,13 +3,17 @@ 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, ) @@ -26,6 +30,13 @@ def test_shared_contract_fixtures_load_canonical_records( 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" @@ -37,3 +48,9 @@ def test_shared_contract_fixtures_load_canonical_records( 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_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index be6b880e8..c7d119389 100644 --- 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 @@ -12,19 +12,19 @@ GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "rendered" -def test_rendered_scripts_are_deterministic_and_bound_to_canonical_plans( - single_node_render_plan: ResolvedSlurmRunPlan, - multi_node_render_plan: ResolvedSlurmRunPlan, +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_render_plan, + single_node_plan, "single_node.sbatch", - expected_sha256="ba4c0d07f9771242e33f29e0bc82dcb83b5502ed12258231548b7037ad8754fd", + expected_fixture_sha256="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", ) _assert_script_matches_plan( - multi_node_render_plan, + multi_node_plan, "multi_node.sbatch", - expected_sha256="d4ac042c1eac5be34dc467686dc93e9c77c7ec5a4c0a16656b5892a1ca47b4c6", + expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", ) @@ -35,8 +35,6 @@ def test_rendered_scripts_contain_no_site_or_user_data() -> None: "nvidia", "cluster", "login", - "account", - "partition", "lustre", "/home/", "/users/", @@ -48,15 +46,21 @@ def _assert_script_matches_plan( plan: ResolvedSlurmRunPlan, filename: str, *, - expected_sha256: str, + expected_fixture_sha256: str, ) -> None: script = (GOLDEN_DIRECTORY / filename).read_text() - node_count = max(index for deployment in plan.deployments for index in deployment.node_indices) + 1 + 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 @@ -68,4 +72,4 @@ def _assert_script_matches_plan( 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_sha256 + 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 index cd1411f76..cf92ca840 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py @@ -171,6 +171,35 @@ def test_backend_failure_injection_is_explicit() -> None: 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: 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 index 815bbc7a9..8bfb3de0e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -12,6 +12,8 @@ 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: @@ -24,12 +26,12 @@ def test_fake_slurm_runner_models_array_submission_and_active_states( ) -> None: _submit(fake_slurm_runner) - observed = fake_slurm_runner.run(("squeue", "--noheader"), check=True) + 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", "--noheader"), + ("squeue", *SQUEUE_ARGUMENTS), ] @@ -49,8 +51,11 @@ def test_fake_slurm_runner_exposes_terminal_accounting_precedence( exit_code="1:0", ) - assert "4101_1|RUNNING" in fake_slurm_runner.run(("squeue",)).stdout - assert fake_slurm_runner.run(("sacct",)).stdout == (GOLDEN_DIRECTORY / "sacct_terminal.txt").read_text() + 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( @@ -64,15 +69,15 @@ def test_fake_slurm_runner_models_accounting_lag_and_later_terminal_state( accounting_state=None, ) - assert "4101_0" not in fake_slurm_runner.run(("squeue",)).stdout - assert "4101_0" not in fake_slurm_runner.run(("sacct",)).stdout + 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",)).stdout + 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: @@ -91,8 +96,8 @@ def test_fake_slurm_runner_filters_job_scoped_queries() -> None: runner.run(("sbatch", "first.sbatch"), check=True) runner.run(("sbatch", "second.sbatch"), check=True) - first_queue = runner.run(("squeue", "--jobs", "4101")).stdout - second_accounting = runner.run(("sacct", "--jobs=4201")).stdout + 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" @@ -101,8 +106,8 @@ def test_fake_slurm_runner_filters_job_scoped_queries() -> None: @pytest.mark.parametrize( "command", ( - ("squeue", "--jobs", "--noheader"), - ("sacct", "--jobs="), + ("squeue", *SQUEUE_ARGUMENTS, "--jobs", "--noheader"), + ("sacct", *SACCT_ARGUMENTS, "--jobs="), ), ) def test_fake_slurm_runner_rejects_malformed_job_selectors( @@ -123,7 +128,7 @@ def test_fake_slurm_runner_models_cancellation( _submit(fake_slurm_runner) assert fake_slurm_runner.run(("scancel", target), check=True).returncode == 0 - accounting = fake_slurm_runner.run(("sacct",), check=True).stdout + 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: @@ -147,6 +152,73 @@ def test_fake_slurm_runner_supports_malformed_output_and_command_failures( 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)