test: add deterministic Slurm fakes

packages/data-designer-slurm/tests/conftest.py

@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import cast
+
+import pytest
+from slurm_test_fakes import (
+    FakeClock,
+    FakeCommandResponse,
+    FakeLogicalEndpoint,
+    FakeSlurmArray,
+    FakeSlurmRunner,
+    FakeSlurmTask,
+    FakeVllmBackend,
+)
+
+from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport
+from data_designer.slurm.client import ClientResult
+from data_designer.slurm.config import (
+    DataDesignerSlurmBenchmarkConfig,
+    DataDesignerSlurmConfig,
+    ImageInspectionRecord,
+    SlurmProfileCatalog,
+)
+from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan
+from data_designer.slurm.state import (
+    AttemptManifest,
+    AttemptReadiness,
+    CandidateOutputManifest,
+    CollectionPlan,
+    RunManifest,
+    SchedulerIdentity,
+    SchedulerObservation,
+    ShardManifest,
+    ShardWinner,
+)
+
+TEST_DIRECTORY = Path(__file__).parent
+CONTRACT_GOLDEN_DIRECTORY = TEST_DIRECTORY / "contracts" / "golden"
+INTEGRATION_GOLDEN_PATH = TEST_DIRECTORY / "integration" / "golden" / "finalization_chain.json"
+SLURM_GOLDEN_DIRECTORY = TEST_DIRECTORY / "slurm_test_fakes" / "golden" / "slurm"
+STATE_GOLDEN_DIRECTORY = TEST_DIRECTORY / "state" / "golden"
+
+
+@pytest.fixture
+def authored_run() -> DataDesignerSlurmConfig:
+    return DataDesignerSlurmConfig.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "authored_run.json").read_text())
+
+
+@pytest.fixture
+def authored_run_single() -> DataDesignerSlurmConfig:
+    return DataDesignerSlurmConfig.model_validate_json(
+        (CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json").read_text()
+    )
+
+
+@pytest.fixture
+def profile_catalog() -> SlurmProfileCatalog:
+    return SlurmProfileCatalog.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "profile_catalog.json").read_text())
+
+
+@pytest.fixture
+def dependency_lock() -> ResolvedDependencyLock:
+    return ResolvedDependencyLock.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "dependency_lock.json").read_text())
+
+
+@pytest.fixture
+def dependency_lock_single() -> ResolvedDependencyLock:
+    return ResolvedDependencyLock.model_validate_json(
+        (CONTRACT_GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text()
+    )
+
+
+@pytest.fixture
+def single_node_plan() -> ResolvedSlurmRunPlan:
+    return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text())
+
+
+@pytest.fixture
+def multi_node_plan() -> ResolvedSlurmRunPlan:
+    return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "multi_node_plan.json").read_text())
+
+
+@pytest.fixture
+def client_result() -> ClientResult:
+    return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.json").read_text())
+
+
+@pytest.fixture
+def benchmark_config() -> DataDesignerSlurmBenchmarkConfig:
+    return DataDesignerSlurmBenchmarkConfig.model_validate_json(
+        (CONTRACT_GOLDEN_DIRECTORY / "benchmark_config.json").read_text()
+    )
+
+
+@pytest.fixture
+def benchmark_manifest() -> BenchmarkManifest:
+    return BenchmarkManifest.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_manifest.json").read_text())
+
+
+@pytest.fixture
+def benchmark_report() -> BenchmarkReport:
+    return BenchmarkReport.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_report.json").read_text())
+
+
+@pytest.fixture
+def client_image_inspection() -> ImageInspectionRecord:
+    return ImageInspectionRecord.model_validate_json(
+        (CONTRACT_GOLDEN_DIRECTORY / "client_image_inspection.json").read_text()
+    )
+
+
+@pytest.fixture
+def serving_image_inspection() -> ImageInspectionRecord:
+    return ImageInspectionRecord.model_validate_json(
+        (CONTRACT_GOLDEN_DIRECTORY / "serving_image_inspection.json").read_text()
+    )
+
+
+@pytest.fixture
+def collection_plan() -> CollectionPlan:
+    return CollectionPlan.model_validate_json((STATE_GOLDEN_DIRECTORY / "collection_plan.json").read_text())
+
+
+@pytest.fixture
+def accounting_lag_observation() -> SchedulerObservation:
+    return SchedulerObservation.model_validate_json((STATE_GOLDEN_DIRECTORY / "accounting_lag.json").read_text())
+
+
+@pytest.fixture
+def finalization_chain_payload() -> dict[str, object]:
+    return cast(dict[str, object], json.loads(INTEGRATION_GOLDEN_PATH.read_text()))
+
+
+@pytest.fixture
+def run_manifest(finalization_chain_payload: dict[str, object]) -> RunManifest:
+    return RunManifest.model_validate_json(json.dumps(finalization_chain_payload["run"]))
+
+
+@pytest.fixture
+def shard_manifests(finalization_chain_payload: dict[str, object]) -> tuple[ShardManifest, ...]:
+    return tuple(
+        ShardManifest.model_validate_json(json.dumps(payload))
+        for payload in cast(list[object], finalization_chain_payload["shards"])
+    )
+
+
+@pytest.fixture
+def attempt_manifest(finalization_chain_payload: dict[str, object]) -> AttemptManifest:
+    return AttemptManifest.model_validate_json(json.dumps(finalization_chain_payload["attempt"]))
+
+
+@pytest.fixture
+def attempt_readiness(finalization_chain_payload: dict[str, object]) -> AttemptReadiness:
+    return AttemptReadiness.model_validate_json(json.dumps(finalization_chain_payload["readiness"]))
+
+
+@pytest.fixture
+def finalization_client_result(finalization_chain_payload: dict[str, object]) -> ClientResult:
+    return ClientResult.model_validate_json(json.dumps(finalization_chain_payload["client_result"]))
+
+
+@pytest.fixture
+def candidate_output_manifest(finalization_chain_payload: dict[str, object]) -> CandidateOutputManifest:
+    return CandidateOutputManifest.model_validate_json(json.dumps(finalization_chain_payload["candidate"]))
+
+
+@pytest.fixture
+def shard_winner(finalization_chain_payload: dict[str, object]) -> ShardWinner:
+    return ShardWinner.model_validate_json(json.dumps(finalization_chain_payload["winner"]))
+
+
+@pytest.fixture
+def fake_clock() -> FakeClock:
+    """Return an isolated explicitly controlled clock."""
+    return FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=100.0)
+
+
+@pytest.fixture
+def fake_slurm_runner() -> FakeSlurmRunner:
+    """Return an isolated Slurm runner with one array and one bounded sinfo query."""
+    return FakeSlurmRunner(
+        arrays=(
+            FakeSlurmArray(
+                tasks=(
+                    FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),
+                    FakeSlurmTask(
+                        SchedulerIdentity(array_job_id=4101, array_task_id=1),
+                        queue_state="RUNNING",
+                    ),
+                )
+            ),
+        ),
+        sinfo_responses={
+            ("sinfo", "--noheader", "--format=%G"): FakeCommandResponse(
+                stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text()
+            )
+        },
+    )
+
+
+@pytest.fixture
+def fake_plugin_overlay() -> Path:
+    """Return the installed-layout fake Data Designer plugin overlay."""
+    return TEST_DIRECTORY / "fixtures" / "fake_plugin_overlay"
+
+
+@pytest.fixture
+def fake_logical_endpoint() -> FakeLogicalEndpoint:
+    """Return an isolated two-backend logical endpoint."""
+    return FakeLogicalEndpoint(
+        "http://127.0.0.1:31000",
+        (
+            FakeVllmBackend("http://127.0.0.1:31001", rank=0),
+            FakeVllmBackend("http://127.0.0.1:31002", rank=1),
+        ),
+    )

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

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

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

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

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,
+)

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",
+]

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