diff --git a/fern/versions/latest.yml b/fern/versions/latest.yml index 94b9af2e3..88c14721c 100644 --- a/fern/versions/latest.yml +++ b/fern/versions/latest.yml @@ -79,6 +79,10 @@ navigation: path: ./latest/pages/notebooks/5-generating-images.mdx - page: Image-to-Image Editing path: ./latest/pages/notebooks/6-editing-images-with-image-context.mdx + - section: Slurm + contents: + - page: Benchmarks + path: ./latest/pages/slurm/benchmarks.mdx - section: Recipes contents: - page: Recipe Cards diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx new file mode 100644 index 000000000..f1161462d --- /dev/null +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -0,0 +1,94 @@ +# Slurm benchmarks + +Use a benchmark to compare concurrency and deployment topology while keeping each case as an ordinary Data Designer Slurm run. No benchmark controller stays resident after submission. + +Install the Slurm extension: + +```bash +pip install "data-designer[slurm]" +``` + +## Define cases + +Create `benchmark.yaml` next to an existing `run.yaml`: + +```yaml +schema_version: 1 +name: generator-scaling +base_run: run.yaml +model_aliases: + - generator +concurrency_values: + - 32 + - 64 +deployment_cases: + - name: two-independent-replicas + deployments: + generator: + nodes: 2 + nodes_per_replica: 1 + - name: one-two-node-replica + deployments: + generator: + nodes: 2 + nodes_per_replica: 2 +record_policy: + type: adaptive + base_records: 1000 + max_records: 5000 + records_per_concurrency: 1.0 +analysis: + target_total_records: 1000000 + target_runtime: 4h +``` + +The compiler expands the authored order deterministically. For this example, the first cases are `two-independent-replicas-c32` and `one-two-node-replica-c32`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. + +## Run + +```bash +data-designer slurm benchmark run benchmark.yaml \ + --profile-file ~/.data-designer-slurm-profile.yml \ + --cluster primary +``` + +The command writes `benchmarks//config.json`, the resolved `base-run.json`, and `benchmark.json` before submitting any child. The manifest is an immutable ordered mapping from case IDs to normal run IDs. Each child keeps its authored config, state, attempts, scheduler evidence, and results under `runs/`. + +If submission stops partway through, rerun the same command. Children with scheduler submission evidence are not submitted again, and missing children are attempted in the original order. If a child has initialized inputs but no submission evidence, rerun with `--force` to resume it. Force never replaces benchmark or child metadata. Partial scheduler evidence across a child's shards is preserved as a conflict for operator investigation and cannot be overwritten with force. + +## Observe and analyze + +Inspect any child with the normal run command: + +```bash +data-designer slurm status +``` + +Write a point-in-time benchmark report from persisted child state: + +```bash +data-designer slurm benchmark analyze --refresh-state +``` + +`--refresh-state` performs one fresh-process scheduler reconciliation per child before analysis. Add `--fail-if-incomplete` when automation should return a conflict after the report is persisted if any case is not successful. + +Reports retain every manifest case in order. Outcomes include `pending`, `accounting_lag`, `succeeded`, `failed`, `incomplete`, `missing`, `stale`, and `scheduler_inconsistent`. Successful cases can be infeasible when boot time consumes the runtime budget. Boot and wall timing begin when execution starts inside each allocation, so queue wait does not affect topology comparisons. For array runs, `rows_per_second` sums the independently measured shard rates, while target jobs and GPU hours count individual task allocations. `generation_seconds` is the target runtime remaining after the slowest shard boot. Only successful feasible cases participate in Pareto, minimum-job, and minimum-GPU-hour recommendations. + +## Python API + +```python +from pathlib import Path + +from data_designer.slurm.config import load_benchmark_config +from data_designer.slurm.services import create_slurm_benchmark_service + +benchmark_file = Path("benchmark.yaml").resolve() +config = load_benchmark_config(benchmark_file) +service = create_slurm_benchmark_service( + profile_file="~/.data-designer-slurm-profile.yml", + cluster="primary", +) + +manifest = service.run(config, source_root=benchmark_file.parent) +report = service.analyze(manifest.benchmark_id, refresh_state=True) +``` diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index f691c26dd..326e052ab 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "data-designer=={{ version }}", "packaging>=25,<27", "pip>=25,<27", - "pydantic>=2.9.2,<3", + "pydantic>=2.12,<3", "pyyaml>=6.0.1,<7", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py index 40acb8f5b..52a7ed22f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py @@ -5,6 +5,12 @@ from __future__ import annotations +from data_designer.slurm.benchmark.compiler import ( + BenchmarkCompiler, + CompiledBenchmark, + CompiledBenchmarkCase, + resolve_requested_records, +) from data_designer.slurm.benchmark.records import ( BenchmarkCaseResult, BenchmarkChildRun, @@ -16,6 +22,7 @@ ) __all__ = [ + "BenchmarkCompiler", "BenchmarkCaseResult", "BenchmarkChildRun", "BenchmarkManifest", @@ -23,4 +30,7 @@ "BenchmarkRecommendation", "BenchmarkRecommendationKind", "BenchmarkReport", + "CompiledBenchmark", + "CompiledBenchmarkCase", + "resolve_requested_records", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py new file mode 100644 index 000000000..28eaec966 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Point-in-time analysis of ordinary benchmark child runs.""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol + +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.records import ( + BenchmarkCaseResult, + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkRecommendation, + BenchmarkRecommendationKind, + BenchmarkReport, +) +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig +from data_designer.slurm.config.utils import convert_duration_to_seconds +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan + + +@dataclass(frozen=True, slots=True) +class BenchmarkShardMeasurements: + """Complete timing and output facts for one successful shard.""" + + actual_records: int + boot_seconds: float + generation_seconds: float + wall_seconds: float + + +@dataclass(frozen=True, slots=True) +class BenchmarkMeasurements: + """Complete timing and output facts for one successful ordinary run.""" + + shards: tuple[BenchmarkShardMeasurements, ...] + + @property + def actual_records(self) -> int: + return sum(shard.actual_records for shard in self.shards) + + @property + def boot_seconds(self) -> float: + return max(shard.boot_seconds for shard in self.shards) + + @property + def generation_seconds(self) -> float: + rate = sum(shard.actual_records / shard.generation_seconds for shard in self.shards) + return self.actual_records / rate + + @property + def wall_seconds(self) -> float: + return max(shard.wall_seconds for shard in self.shards) + + +@dataclass(frozen=True, slots=True) +class BenchmarkRunObservation: + """Verified ordinary run inputs and one status snapshot.""" + + authored_config: DataDesignerSlurmConfig + resolved_plan: ResolvedSlurmRunPlan + outcome: BenchmarkOutcome + measurements: BenchmarkMeasurements | None = None + + +class BenchmarkRunObserver(Protocol): + """Observe one ordinary child run from durable state.""" + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + """Return one verified child snapshot or raise an observation failure.""" + + +class BenchmarkObservationFailure(Exception): + """Preserve a non-success child outcome without aborting the report.""" + + def __init__(self, outcome: BenchmarkOutcome) -> None: + if outcome not in { + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + }: + raise ValueError("invalid benchmark observation failure outcome") + self.outcome = outcome + super().__init__(outcome.value) + + +class BenchmarkManifestMismatchError(ValueError): + """Raised when persisted benchmark children do not match compilation.""" + + +class BenchmarkAnalyzer: + """Analyze one immutable benchmark manifest in manifest order.""" + + def __init__(self, observer: BenchmarkRunObserver, clock: Callable[[], datetime]) -> None: + self._observer = observer + self._clock = clock + + def analyze( + self, + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, + manifest_reference: ArtifactReference, + *, + refresh_state: bool, + ) -> BenchmarkReport: + requested_records = _validate_and_get_requested_records(config, base_run, manifest) + results = tuple( + self._analyze_case( + config, + child, + requested_records[index], + refresh_state=refresh_state, + ) + for index, child in enumerate(manifest.children) + ) + recommendations = _recommend(results) + created_at = self._clock() + analysis_digest = compute_canonical_json_sha256( + { + "benchmark_id": manifest.benchmark_id, + "created_at": created_at.isoformat(), + "cases": [case.model_dump(mode="json") for case in results], + "recommendations": [item.model_dump(mode="json") for item in recommendations], + } + ) + return BenchmarkReport( + schema_version=1, + benchmark_id=manifest.benchmark_id, + analysis_id=f"analysis-{analysis_digest[:32]}", + benchmark_manifest=manifest_reference, + created_at=created_at, + cases=results, + recommendations=recommendations, + ) + + def _analyze_case( + self, + config: DataDesignerSlurmBenchmarkConfig, + child: BenchmarkChildRun, + requested_records: int, + *, + refresh_state: bool, + ) -> BenchmarkCaseResult: + fallback = _fallback_facts(child, requested_records) + try: + observed = self._observer.observe(child.child_run_id, refresh_state=refresh_state) + except BenchmarkObservationFailure as error: + return _update_result(fallback, outcome=error.outcome) + if ( + observed.authored_config.compute_sha256() != child.child_authored_config.sha256 + or observed.resolved_plan.run_id != child.child_run_id + ): + return _update_result(fallback, outcome=BenchmarkOutcome.STALE) + + result = _resolved_facts(child, observed.authored_config, observed.resolved_plan) + if observed.outcome is not BenchmarkOutcome.SUCCEEDED: + return _update_result(result, outcome=observed.outcome) + if observed.measurements is None: + return _update_result(result, outcome=BenchmarkOutcome.INCOMPLETE) + measurements = observed.measurements + actual_records = measurements.actual_records + if ( + not measurements.shards + or len(measurements.shards) != len(observed.resolved_plan.shards) + or any( + shard.actual_records <= 0 + or shard.boot_seconds < 0 + or shard.generation_seconds <= 0 + or shard.wall_seconds <= 0 + for shard in measurements.shards + ) + ): + return _update_result( + result, + outcome=BenchmarkOutcome.INCOMPLETE, + actual_records=actual_records, + ) + boot_seconds = measurements.boot_seconds + generation_seconds = measurements.generation_seconds + wall_seconds = measurements.wall_seconds + requested_records = observed.authored_config.invocation.num_records + if actual_records != requested_records: + return _update_result( + result, + outcome=BenchmarkOutcome.INCOMPLETE, + actual_records=actual_records, + ) + gpus_per_job = result.gpus_per_job + if gpus_per_job is None: + raise ValueError("resolved benchmark case has no GPU count") + rows_per_second = actual_records / generation_seconds + budget_seconds = convert_duration_to_seconds(config.analysis.target_runtime) + effective_generation_seconds = max(0.0, budget_seconds - boot_seconds) + feasible = boot_seconds < budget_seconds + gpu_hours_per_job = gpus_per_job * budget_seconds / 3600 + if not feasible: + return _update_result( + result, + outcome=BenchmarkOutcome.SUCCEEDED, + actual_records=actual_records, + boot_seconds=boot_seconds, + generation_seconds=effective_generation_seconds, + wall_seconds=wall_seconds, + rows_per_second=rows_per_second, + gpu_hours_per_job=gpu_hours_per_job, + feasible=False, + ) + records_per_job = sum( + shard.actual_records / shard.generation_seconds * (budget_seconds - shard.boot_seconds) + for shard in measurements.shards + ) / len(measurements.shards) + target_jobs = math.ceil(config.analysis.target_total_records / records_per_job) + return _update_result( + result, + outcome=BenchmarkOutcome.SUCCEEDED, + actual_records=actual_records, + boot_seconds=boot_seconds, + generation_seconds=effective_generation_seconds, + wall_seconds=wall_seconds, + rows_per_second=rows_per_second, + gpu_hours_per_job=gpu_hours_per_job, + total_gpu_hours=gpu_hours_per_job * target_jobs, + target_jobs=target_jobs, + feasible=True, + ) + + +def _fallback_facts( + child: BenchmarkChildRun, + requested_records: int, +) -> BenchmarkCaseResult: + return BenchmarkCaseResult( + case_id=child.case_id, + child_run_id=child.child_run_id, + outcome=BenchmarkOutcome.MISSING, + requested_records=requested_records, + ) + + +def _resolved_facts( + child: BenchmarkChildRun, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> BenchmarkCaseResult: + nodes = sum(len(deployment.node_indices) for deployment in plan.deployments) + gpus = sum(len(deployment.node_indices) * deployment.gpus_per_node for deployment in plan.deployments) + topology_digest = compute_canonical_json_sha256( + [ + { + "deployment_id": deployment.deployment_id, + "model_alias": deployment.authored.model_alias, + "node_indices": deployment.node_indices, + "gpus_per_node": deployment.gpus_per_node, + "topology": deployment.topology.model_dump(mode="json"), + } + for deployment in plan.deployments + ] + ) + return BenchmarkCaseResult( + case_id=child.case_id, + child_run_id=child.child_run_id, + outcome=BenchmarkOutcome.INCOMPLETE, + topology_digest=topology_digest, + requested_records=authored.invocation.num_records, + gpus_per_job=gpus, + nodes_per_job=nodes, + ) + + +def _validate_and_get_requested_records( + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, +) -> tuple[int, ...]: + try: + compiled = BenchmarkCompiler.compile(config, base_run) + except ValueError as error: + raise BenchmarkManifestMismatchError("benchmark configuration cannot reproduce its children") from error + expected = tuple( + (case.case_id, case.child_run_id, case.child_run_config.compute_sha256()) for case in compiled.cases + ) + actual = tuple( + (child.case_id, child.child_run_id, child.child_authored_config.sha256) for child in manifest.children + ) + if manifest.benchmark_id != compiled.benchmark_id or actual != expected: + raise BenchmarkManifestMismatchError("benchmark manifest does not match deterministic child expansion") + return tuple(case.child_run_config.invocation.num_records for case in compiled.cases) + + +def _update_result(result: BenchmarkCaseResult, **updates: object) -> BenchmarkCaseResult: + return BenchmarkCaseResult.model_validate(result.model_dump(mode="python") | updates) + + +def _recommend(cases: tuple[BenchmarkCaseResult, ...]) -> tuple[BenchmarkRecommendation, ...]: + eligible = tuple( + case + for case in cases + if case.outcome is BenchmarkOutcome.SUCCEEDED + and case.feasible is True + and case.target_jobs is not None + and case.total_gpu_hours is not None + ) + if not eligible: + return () + pareto = tuple( + case + for case in eligible + if not any( + other.case_id != case.case_id + and other.target_jobs <= case.target_jobs + and other.total_gpu_hours <= case.total_gpu_hours + and (other.target_jobs < case.target_jobs or other.total_gpu_hours < case.total_gpu_hours) + for other in eligible + ) + ) + minimum_jobs = min(eligible, key=lambda case: (case.target_jobs, case.total_gpu_hours, case.case_id)) + minimum_gpu_hours = min(eligible, key=lambda case: (case.total_gpu_hours, case.target_jobs, case.case_id)) + return tuple( + BenchmarkRecommendation(kind=BenchmarkRecommendationKind.PARETO, case_id=case.case_id) for case in pareto + ) + ( + BenchmarkRecommendation(kind=BenchmarkRecommendationKind.MINIMUM_JOBS, case_id=minimum_jobs.case_id), + BenchmarkRecommendation( + kind=BenchmarkRecommendationKind.MINIMUM_GPU_HOURS, + case_id=minimum_gpu_hours.case_id, + ), + ) + + +__all__ = [ + "BenchmarkAnalyzer", + "BenchmarkMeasurements", + "BenchmarkShardMeasurements", + "BenchmarkManifestMismatchError", + "BenchmarkObservationFailure", + "BenchmarkRunObservation", + "BenchmarkRunObserver", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py new file mode 100644 index 000000000..568533af9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic expansion of benchmark intent into ordinary run configs.""" + +from __future__ import annotations + +import math +import posixpath + +from pydantic import Field + +from data_designer.slurm.config import ( + AdaptiveRecordPolicy, + BenchmarkDeploymentCase, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, +) +from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 + + +class CompiledBenchmarkCase(ContractValue): + """One stable benchmark case and its ordinary authored run config.""" + + case_id: Identifier + child_run_id: Identifier + child_run_config: DataDesignerSlurmConfig + + +class CompiledBenchmark(ContractValue): + """Ordered immutable result of benchmark expansion.""" + + benchmark_id: Identifier + cases: tuple[CompiledBenchmarkCase, ...] = Field(min_length=1) + + +class BenchmarkCompiler: + """Compile a benchmark declaration without filesystem or scheduler access.""" + + @staticmethod + def compile( + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + ) -> CompiledBenchmark: + if not isinstance(config, DataDesignerSlurmBenchmarkConfig): + raise TypeError("config must be a DataDesignerSlurmBenchmarkConfig") + if not isinstance(base_run, DataDesignerSlurmConfig): + raise TypeError("base_run must be a DataDesignerSlurmConfig") + + aliases = tuple(deployment.model_alias for deployment in base_run.deployments) + selected_aliases = aliases if config.model_aliases == "all" else tuple(config.model_aliases) + unknown_aliases = set(selected_aliases).difference(aliases) + if unknown_aliases: + raise ValueError(f"benchmark references unknown model aliases: {', '.join(sorted(unknown_aliases))}") + for deployment_case in config.deployment_cases: + unknown_deployments = set(deployment_case.deployments).difference(aliases) + if unknown_deployments: + raise ValueError( + f"benchmark deployment case references unknown aliases: {', '.join(sorted(unknown_deployments))}" + ) + + identity_digest = compute_canonical_json_sha256( + { + "benchmark": config.model_dump(mode="json"), + "base_run": base_run.model_dump(mode="json"), + } + ) + benchmark_id = f"benchmark-{identity_digest[:32]}" + expanded: list[tuple[str, DataDesignerSlurmConfig]] = [] + for concurrency in config.concurrency_values: + for deployment_case in config.deployment_cases: + case_id = derive_benchmark_case_id(deployment_case.name, concurrency) + expanded.append( + ( + case_id, + _compile_child( + base_run, + benchmark_id=benchmark_id, + case_id=case_id, + concurrency=concurrency, + selected_aliases=selected_aliases, + deployment_case=deployment_case, + requested_records=resolve_requested_records(config, concurrency), + ), + ) + ) + + cases = tuple( + CompiledBenchmarkCase( + case_id=case_id, + child_run_id=f"run-{identity_digest[:16]}-{index:04d}-{child.compute_sha256()[:12]}", + child_run_config=child, + ) + for index, (case_id, child) in enumerate(expanded) + ) + return CompiledBenchmark(benchmark_id=benchmark_id, cases=cases) + + +def _compile_child( + base_run: DataDesignerSlurmConfig, + *, + benchmark_id: str, + case_id: str, + concurrency: int, + selected_aliases: tuple[str, ...], + deployment_case: BenchmarkDeploymentCase, + requested_records: int, +) -> DataDesignerSlurmConfig: + model_concurrency = dict(base_run.invocation.model_concurrency) + model_concurrency.update(dict.fromkeys(selected_aliases, concurrency)) + invocation = base_run.invocation.model_copy( + update={ + "num_records": requested_records, + "model_concurrency": model_concurrency, + } + ) + deployments = [] + for deployment in base_run.deployments: + override = deployment_case.deployments.get(deployment.model_alias) + if override is None: + deployments.append(deployment) + continue + deployments.append( + deployment.model_copy( + update={ + "resources": deployment.resources.model_copy(update={"nodes": override.nodes}), + "topology": deployment.topology.model_copy( + update={"nodes_per_replica": override.nodes_per_replica} + ), + } + ) + ) + + output_root = base_run.output.root + output = base_run.output + if output_root is not None: + output = output.model_copy(update={"root": posixpath.join(output_root, benchmark_id, case_id)}) + return DataDesignerSlurmConfig.model_validate( + base_run.model_copy( + update={ + "name": _bounded_identifier(f"{base_run.name}-{case_id}"), + "invocation": invocation, + "deployments": deployments, + "output": output, + } + ).model_dump(mode="python") + ) + + +def resolve_requested_records(config: DataDesignerSlurmBenchmarkConfig, concurrency: int) -> int: + """Resolve the deterministic record count for one concurrency value.""" + policy = config.record_policy + if not isinstance(policy, AdaptiveRecordPolicy): + return policy.records + scaled = math.ceil(concurrency * policy.records_per_concurrency) + return min(policy.max_records, max(policy.base_records, scaled)) + + +def derive_benchmark_case_id(case_name: str, concurrency: int) -> str: + """Return the stable bounded identity for one cross-product case.""" + return _bounded_identifier(f"{case_name}-c{concurrency}") + + +def _bounded_identifier(value: str) -> str: + if len(value) <= 128: + return value + digest = compute_canonical_json_sha256(value)[:12] + return f"{value[:115]}-{digest}" + + +__all__ = [ + "BenchmarkCompiler", + "CompiledBenchmark", + "CompiledBenchmarkCase", + "derive_benchmark_case_id", + "resolve_requested_records", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py new file mode 100644 index 000000000..6841a453c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark orchestration through ordinary public Slurm run services.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkAnalyzer, + BenchmarkManifestMismatchError, + BenchmarkRunObserver, +) +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.records import ( + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkReport, +) +from data_designer.slurm.benchmark.store import ( + BenchmarkConflictError, + BenchmarkNotFoundError, + BenchmarkStore, + BenchmarkStoreError, +) +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + SlurmConfigLoadError, + load_run_config, +) +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.services.errors import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) +from data_designer.slurm.services.run import SlurmRunService +from data_designer.slurm.state import SlurmStateError, SlurmStateWriter, StateNotFoundError +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +ChildRunServiceFactory = Callable[[Identifier], SlurmRunService] +ChildConfigLoader = Callable[[Identifier], DataDesignerSlurmConfig | None] +ChildSubmissionLoader = Callable[[Identifier], bool] + + +class SystemBenchmarkBackend: + """Persist benchmark intent and submit each case as an ordinary run.""" + + def __init__( + self, + workspace_root: str | Path, + run_service_factory: ChildRunServiceFactory, + observer: BenchmarkRunObserver, + clock: Callable[[], datetime], + child_config_loader: ChildConfigLoader | None = None, + child_submission_loader: ChildSubmissionLoader | None = None, + ) -> None: + self._workspace_root = Path(workspace_root) + self._run_service_factory = run_service_factory + self._observer = observer + self._clock = clock + self._child_config_loader = child_config_loader or self._load_child_config + self._child_submission_loader = child_submission_loader or self._load_child_submission + + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: + try: + base_run, child_source_root = _resolve_base_run(config, source_root) + compiled = BenchmarkCompiler.compile(config, base_run) + except (SlurmConfigLoadError, ValueError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark configuration cannot be compiled", + ) from None + store = BenchmarkStore(self._workspace_root, compiled.benchmark_id) + children = tuple( + BenchmarkChildRun( + case_id=case.case_id, + child_run_id=case.child_run_id, + child_authored_config=ArtifactReference( + path=(self._workspace_root / "runs" / case.child_run_id / "authored-config.json").as_posix(), + sha256=case.child_run_config.compute_sha256(), + ), + ) + for case in compiled.cases + ) + manifest = store.build_manifest(config, children) + try: + store.publish(config, base_run, manifest) + except BenchmarkConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark metadata conflicts with persisted state", + ) from None + except BenchmarkStoreError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark metadata cannot be persisted", + ) from None + + failures: list[SlurmServiceErrorCode] = [] + for case in compiled.cases: + try: + with StateStorage(self._workspace_root, case.child_run_id).acquire_submission_lock(): + if self._child_exists(case.child_run_id, case.child_run_config, force=force): + continue + execution = self._run_service_factory(case.child_run_id).execute( + case.child_run_config, + source_root=child_source_root, + dry_run=False, + force=False, + ) + if execution.run_id != case.child_run_id or execution.state != "submitted": + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.RUN_BENCHMARK, + "ordinary run service returned an invalid benchmark child", + ) + except SlurmServiceError as error: + failures.append(error.code) + except Exception: + failures.append(SlurmServiceErrorCode.INTERNAL) + if failures: + code = _partial_failure_code(failures) + raise SlurmServiceError( + code, + SlurmServiceOperation.RUN_BENCHMARK, + f"{len(failures)} of {len(compiled.cases)} benchmark child runs could not be submitted", + ) + return manifest + + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: + store = BenchmarkStore(self._workspace_root, benchmark_id) + try: + config, base_run, manifest = store.load() + report = BenchmarkAnalyzer(self._observer, self._clock).analyze( + config, + base_run, + manifest, + store.manifest_reference(manifest), + refresh_state=refresh_state, + ) + store.publish_report(report) + except BenchmarkNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata was not found", + ) from None + except BenchmarkConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata conflicts with persisted state", + ) from None + except BenchmarkStoreError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata cannot be read or persisted", + ) from None + except BenchmarkManifestMismatchError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark manifest does not match its configuration", + ) from None + if fail_if_incomplete: + incomplete = sum(case.outcome is not BenchmarkOutcome.SUCCEEDED for case in report.cases) + if incomplete: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + f"benchmark analysis contains {incomplete} incomplete cases", + ) + return report + + def _child_exists( + self, + run_id: Identifier, + expected: DataDesignerSlurmConfig, + *, + force: bool, + ) -> bool: + persisted = self._child_config_loader(run_id) + if persisted is None: + return False + if persisted != expected: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state contains different inputs", + ) + if self._child_submission_loader(run_id): + return True + if force: + return False + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child inputs exist without submission evidence; rerun with force", + ) + + def _load_child_config(self, run_id: Identifier) -> DataDesignerSlurmConfig | None: + try: + return SlurmStateWriter(self._workspace_root, run_id).load_authored_config() + except StateNotFoundError: + return None + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state is invalid", + ) from None + + def _load_child_submission(self, run_id: Identifier) -> bool: + try: + storage = StateStorage(self._workspace_root, run_id) + reader = StateReader(storage, run_id) + run, plan, shards = reader.load_context() + attempts = reader.load_validated_attempts(run, plan, shards) + submitted = tuple( + bool(attempts[shard.shard_id] and attempts[shard.shard_id][0].scheduler is not None) for shard in shards + ) + if all(submitted): + return True + if any(submitted): + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child contains partial submission evidence", + ) + return False + except SlurmServiceError: + raise + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state is invalid", + ) from None + + +def _resolve_base_run( + config: DataDesignerSlurmBenchmarkConfig, + source_root: Path, +) -> tuple[DataDesignerSlurmConfig, Path]: + if config.base_run.inline is not None: + return config.base_run.inline, source_root + assert config.base_run.source is not None + source_path = source_root / config.base_run.source + return load_run_config(source_path), source_path.resolve().parent + + +def _partial_failure_code(codes: list[SlurmServiceErrorCode]) -> SlurmServiceErrorCode: + for code in ( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceErrorCode.INTERNAL, + ): + if code in codes: + return code + return SlurmServiceErrorCode.INTERNAL + + +__all__ = [ + "ChildConfigLoader", + "ChildRunServiceFactory", + "ChildSubmissionLoader", + "SystemBenchmarkBackend", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py new file mode 100644 index 000000000..3ae540698 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process observation of ordinary benchmark child state.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkMeasurements, + BenchmarkObservationFailure, + BenchmarkRunObservation, + BenchmarkShardMeasurements, +) +from data_designer.slurm.benchmark.records import BenchmarkOutcome +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + EffectiveRunState, + ProbeOutcome, + ReadinessState, + RunManifest, + RunStatus, + SchedulerObservationClient, + SchedulerState, + SchedulerStateConflictError, + ShardManifest, + SlurmStateError, + SlurmStateReconciler, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_MeasurementRecord = tuple[ + AttemptManifest, + AttemptReadiness | None, + tuple[ClientResult, CandidateOutputManifest] | None, +] + + +class PersistedBenchmarkRunObserver: + """Reconstruct benchmark child facts from the ordinary run state tree.""" + + def __init__( + self, + workspace_root: str | Path, + scheduler: SchedulerObservationClient, + clock: Callable[[], datetime], + ) -> None: + self._workspace_root = Path(workspace_root) + self._scheduler = scheduler + self._clock = clock + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + try: + storage = StateStorage(self._workspace_root, run_id) + reader = StateReader(storage, run_id) + run, plan, shards = reader.load_context() + authored = reader.load_authored_config(run) + if refresh_state: + status = SlurmStateReconciler( + self._workspace_root, + run_id, + self._scheduler, + ).refresh(observed_at=self._clock()) + outcome = _status_outcome(status.effective_state) + measurements = _measure_status(status) if outcome is BenchmarkOutcome.SUCCEEDED else None + else: + outcome, measurements = _load_persisted_outcome(storage, reader, run, plan, shards) + return BenchmarkRunObservation( + authored_config=authored, + resolved_plan=plan, + outcome=outcome, + measurements=measurements, + ) + except StateNotFoundError: + raise BenchmarkObservationFailure(BenchmarkOutcome.MISSING) from None + except SchedulerStateConflictError: + raise BenchmarkObservationFailure(BenchmarkOutcome.SCHEDULER_INCONSISTENT) from None + except (StateConflictError, StateCorruptionError): + raise BenchmarkObservationFailure(BenchmarkOutcome.STALE) from None + except (OSError, SlurmStateError): + raise BenchmarkObservationFailure(BenchmarkOutcome.STALE) from None + + +def _load_persisted_outcome( + storage: StateStorage, + reader: StateReader, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shards: tuple[ShardManifest, ...], +) -> tuple[BenchmarkOutcome, BenchmarkMeasurements | None]: + attempts_by_shard = reader.load_validated_attempts(run, plan, shards) + finalizer = WinnerFinalizer(storage, reader) + winning_records: list[_MeasurementRecord] = [] + latest_states: list[AttemptLifecycleState] = [] + scheduler_states: list[SchedulerState] = [] + for shard in shards: + attempts = attempts_by_shard[shard.shard_id] + if not attempts: + return BenchmarkOutcome.INCOMPLETE, None + winner = finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + latest = attempts[-1] + latest_states.append(latest.state) + observation = reader.load_optional_scheduler_observation(latest) + if observation is not None: + scheduler_states.append(observation.state) + continue + attempt = next(item for item in attempts if item.attempt_id == winner.attempt_id) + observation = reader.load_optional_scheduler_observation(attempt) + if observation is not None and observation.state in { + SchedulerState.FAILED, + SchedulerState.CANCELLED, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.OUT_OF_MEMORY, + }: + raise SchedulerStateConflictError("persisted winner conflicts with scheduler evidence") + result = reader.load_optional_attempt_result(plan, shard, attempt) + readiness = reader.load_optional_readiness(plan, attempt) + winning_records.append((attempt, readiness, result)) + + if len(winning_records) == len(shards): + return BenchmarkOutcome.SUCCEEDED, _measure_records(tuple(winning_records)) + if SchedulerState.ACCOUNTING_LAG in scheduler_states: + return BenchmarkOutcome.ACCOUNTING_LAG, None + if any( + state in {SchedulerState.UNKNOWN, SchedulerState.PREEMPTED, SchedulerState.REQUEUED} + for state in scheduler_states + ): + return BenchmarkOutcome.STALE, None + if any( + state + in { + AttemptLifecycleState.CREATED, + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + } + for state in latest_states + ): + return BenchmarkOutcome.PENDING, None + if all(state is AttemptLifecycleState.FAILED for state in latest_states): + return BenchmarkOutcome.FAILED, None + return BenchmarkOutcome.INCOMPLETE, None + + +def _status_outcome(state: EffectiveRunState) -> BenchmarkOutcome: + if state in {EffectiveRunState.PENDING, EffectiveRunState.RUNNING}: + return BenchmarkOutcome.PENDING + if state is EffectiveRunState.ACCOUNTING_LAG: + return BenchmarkOutcome.ACCOUNTING_LAG + if state is EffectiveRunState.SUCCEEDED: + return BenchmarkOutcome.SUCCEEDED + if state is EffectiveRunState.FAILED: + return BenchmarkOutcome.FAILED + return BenchmarkOutcome.STALE + + +def _measure_status(status: RunStatus) -> BenchmarkMeasurements | None: + records: list[_MeasurementRecord] = [] + for shard in status.shards: + winner = next((attempt for attempt in shard.attempts if attempt.is_winner), None) + if winner is None: + return None + result = ( + None + if winner.client_result is None or winner.candidate_output is None + else (winner.client_result, winner.candidate_output) + ) + records.append((winner.attempt, winner.readiness, result)) + return _measure_records(tuple(records)) + + +def _measure_records(records: tuple[_MeasurementRecord, ...]) -> BenchmarkMeasurements | None: + measurements: list[BenchmarkShardMeasurements] = [] + for attempt, readiness, result in records: + if result is None: + return None + client_result, candidate = result + if ( + client_result.outcome is not ClientOutcome.COMPLETE + or candidate.actual_records != client_result.actual_records + ): + return None + if readiness is None or readiness.state not in {ReadinessState.READY, ReadinessState.STOPPED}: + return None + if readiness.started_at is None: + return None + probes = tuple(deployment.last_probe for deployment in readiness.deployments) + if any(probe is None or probe.outcome is not ProbeOutcome.SUCCESS for probe in probes): + return None + started_at = readiness.started_at + ready_at = max(probe.observed_at for probe in probes if probe is not None) + completed_at = client_result.completed_at + stopped_at = attempt.updated_at + if ready_at < started_at or completed_at <= ready_at or stopped_at < completed_at or stopped_at <= started_at: + return None + measurements.append( + BenchmarkShardMeasurements( + actual_records=client_result.actual_records or 0, + boot_seconds=(ready_at - started_at).total_seconds(), + generation_seconds=(completed_at - ready_at).total_seconds(), + wall_seconds=(stopped_at - started_at).total_seconds(), + ) + ) + return BenchmarkMeasurements(shards=tuple(measurements)) + + +__all__ = ["PersistedBenchmarkRunObserver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py index 11c797527..40d98b232 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -57,13 +57,16 @@ class BenchmarkOutcome(str, Enum): SUCCEEDED = "succeeded" FAILED = "failed" INCOMPLETE = "incomplete" + MISSING = "missing" + STALE = "stale" + SCHEDULER_INCONSISTENT = "scheduler_inconsistent" class BenchmarkCaseResult(ContractValue): case_id: Identifier child_run_id: Identifier outcome: BenchmarkOutcome - topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] | None = None requested_records: PositiveInt actual_records: NonNegativeInt | None = None boot_seconds: NonNegativeFloat | None = None @@ -72,8 +75,8 @@ class BenchmarkCaseResult(ContractValue): rows_per_second: NonNegativeFloat | None = None request_count: NonNegativeInt | None = None token_count: NonNegativeInt | None = None - gpus_per_job: PositiveInt - nodes_per_job: PositiveInt + gpus_per_job: PositiveInt | None = None + nodes_per_job: PositiveInt | None = None gpu_hours_per_job: NonNegativeFloat | None = None total_gpu_hours: NonNegativeFloat | None = None target_jobs: PositiveInt | None = None @@ -89,9 +92,9 @@ def validate_metrics(self) -> BenchmarkCaseResult: self.generation_seconds, self.wall_seconds, self.rows_per_second, + self.gpus_per_job, + self.nodes_per_job, self.gpu_hours_per_job, - self.total_gpu_hours, - self.target_jobs, self.feasible, ) if self.outcome is BenchmarkOutcome.SUCCEEDED and any(value is None for value in required): @@ -99,8 +102,14 @@ def validate_metrics(self) -> BenchmarkCaseResult: if self.outcome is BenchmarkOutcome.SUCCEEDED: if self.actual_records != self.requested_records: raise ValueError("successful benchmark cases require the requested record count") - if self.generation_seconds == 0 or self.wall_seconds == 0 or self.rows_per_second == 0: - raise ValueError("successful benchmark generation, wall time, and throughput must be positive") + if self.wall_seconds == 0 or self.rows_per_second == 0: + raise ValueError("successful benchmark wall time and throughput must be positive") + if self.feasible and self.generation_seconds == 0: + raise ValueError("feasible benchmark generation time must be positive") + if self.feasible and (self.total_gpu_hours is None or self.target_jobs is None): + raise ValueError("feasible benchmark cases require target jobs and total GPU hours") + if not self.feasible and (self.total_gpu_hours is not None or self.target_jobs is not None): + raise ValueError("infeasible benchmark cases cannot have target jobs or total GPU hours") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py new file mode 100644 index 000000000..fa7027d2a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable benchmark metadata beneath the selected Slurm workspace.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TypeVar + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.benchmark.records import BenchmarkChildRun, BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig +from data_designer.slurm.contracts import ArtifactReference, ContractRecord, Identifier, validate_absolute_path +from data_designer.slurm.images.filesystem import ensure_private_directory +from data_designer.slurm.state.filesystem import open_verified_directory, publish_immutable_text, read_regular_text + +_CONFIG_FILENAME = "config.json" +_BASE_RUN_FILENAME = "base-run.json" +_MANIFEST_FILENAME = "benchmark.json" +_REPORTS_DIRECTORY = "reports" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_RecordT = TypeVar("_RecordT", bound=ContractRecord) + + +class BenchmarkStoreError(Exception): + """Base error for benchmark metadata persistence.""" + + +class BenchmarkNotFoundError(BenchmarkStoreError): + """A benchmark metadata record is missing.""" + + +class BenchmarkConflictError(BenchmarkStoreError): + """Benchmark metadata is unsafe, corrupt, or conflicts with immutable state.""" + + +class BenchmarkStore: + """Persist one benchmark config, manifest, and point-in-time reports.""" + + def __init__(self, workspace_root: str | Path, benchmark_id: Identifier) -> None: + try: + root = Path(validate_absolute_path(Path(workspace_root).as_posix())) + normalized_id = _IDENTIFIER_ADAPTER.validate_python(benchmark_id, strict=True) + except (ValidationError, ValueError) as error: + raise BenchmarkStoreError("invalid benchmark location") from error + self.workspace_root = root + self.benchmark_id = normalized_id + self.benchmarks_root = root / "benchmarks" + self.benchmark_root = self.benchmarks_root / normalized_id + self.config_path = self.benchmark_root / _CONFIG_FILENAME + self.base_run_path = self.benchmark_root / _BASE_RUN_FILENAME + self.manifest_path = self.benchmark_root / _MANIFEST_FILENAME + self.reports_root = self.benchmark_root / _REPORTS_DIRECTORY + + def build_manifest( + self, + config: DataDesignerSlurmBenchmarkConfig, + children: tuple[BenchmarkChildRun, ...], + ) -> BenchmarkManifest: + """Build a manifest bound to this store's immutable locations.""" + return BenchmarkManifest( + schema_version=1, + benchmark_id=self.benchmark_id, + benchmark_config=ArtifactReference(path=self.config_path.as_posix(), sha256=config.compute_sha256()), + children=children, + ) + + def publish( + self, + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, + ) -> BenchmarkManifest: + """Convergently publish resolved inputs then the manifest commit record.""" + self._validate_manifest(config, manifest) + try: + ensure_private_directory(self.benchmark_root, parents=True) + with open_verified_directory(self.benchmark_root, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + _BASE_RUN_FILENAME, + base_run.serialize_json(), + self.base_run_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + publish_immutable_text( + descriptor, + _CONFIG_FILENAME, + config.serialize_json(), + self.config_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + publish_immutable_text( + descriptor, + _MANIFEST_FILENAME, + manifest.serialize_json(), + self.manifest_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + return manifest + except FileExistsError as error: + raise BenchmarkConflictError("benchmark already contains different immutable metadata") from error + except OSError as error: + raise BenchmarkStoreError("benchmark metadata cannot be persisted") from error + + def load(self) -> tuple[DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, BenchmarkManifest]: + """Load and verify the committed benchmark metadata.""" + try: + with open_verified_directory(self.benchmark_root, require_private=True) as descriptor: + config = self._read_record( + descriptor, + _CONFIG_FILENAME, + self.config_path, + DataDesignerSlurmBenchmarkConfig, + ) + base_run = self._read_record( + descriptor, + _BASE_RUN_FILENAME, + self.base_run_path, + DataDesignerSlurmConfig, + ) + manifest = self._read_record( + descriptor, + _MANIFEST_FILENAME, + self.manifest_path, + BenchmarkManifest, + ) + self._validate_manifest(config, manifest) + return config, base_run, manifest + except FileNotFoundError as error: + raise BenchmarkNotFoundError(f"benchmark {self.benchmark_id!r} is not initialized") from error + except BenchmarkConflictError: + raise + except (OSError, ValidationError, ValueError) as error: + raise BenchmarkConflictError(f"benchmark {self.benchmark_id!r} contains invalid metadata") from error + + def publish_report(self, report: BenchmarkReport) -> Path: + """Persist one immutable point-in-time report.""" + if report.benchmark_id != self.benchmark_id: + raise BenchmarkConflictError("benchmark report identity does not match its store") + _, _, manifest = self.load() + expected_reference = ArtifactReference(path=self.manifest_path.as_posix(), sha256=manifest.compute_sha256()) + if report.benchmark_manifest != expected_reference: + raise BenchmarkConflictError("benchmark report does not bind the persisted manifest") + report_root = self.reports_root / report.analysis_id + report_path = report_root / "report.json" + try: + ensure_private_directory(self.reports_root, parents=False) + ensure_private_directory(report_root, parents=False) + with open_verified_directory(report_root, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + "report.json", + report.serialize_json(), + report_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + return report_path + except FileExistsError as error: + raise BenchmarkConflictError("benchmark report identity already contains different bytes") from error + except OSError as error: + raise BenchmarkStoreError("benchmark report cannot be persisted") from error + + def manifest_reference(self, manifest: BenchmarkManifest) -> ArtifactReference: + """Return the immutable reference for a verified manifest.""" + return ArtifactReference(path=self.manifest_path.as_posix(), sha256=manifest.compute_sha256()) + + def _validate_manifest( + self, + config: DataDesignerSlurmBenchmarkConfig, + manifest: BenchmarkManifest, + ) -> None: + expected = ArtifactReference(path=self.config_path.as_posix(), sha256=config.compute_sha256()) + if manifest.benchmark_id != self.benchmark_id or manifest.benchmark_config != expected: + raise BenchmarkConflictError("benchmark manifest does not match its immutable config") + for child in manifest.children: + expected_path = self.workspace_root / "runs" / child.child_run_id / "authored-config.json" + if child.child_authored_config.path != expected_path.as_posix(): + raise BenchmarkConflictError("benchmark child reference is outside its ordinary run state") + + @staticmethod + def _read_record( + directory_descriptor: int, + name: str, + path: Path, + record_type: type[_RecordT], + ) -> _RecordT: + content = read_regular_text( + directory_descriptor, + name, + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + record = record_type.model_validate_json(content) + if record.serialize_json() != content: + raise BenchmarkConflictError("benchmark metadata is not canonical") + return record + + +__all__ = [ + "BenchmarkConflictError", + "BenchmarkNotFoundError", + "BenchmarkStore", + "BenchmarkStoreError", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 730ce513c..18d8a4c2f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -14,6 +14,7 @@ import typer from pydantic import BaseModel, ValidationError +from data_designer.slurm.cli_benchmark import create_benchmark_app from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json from data_designer.slurm.images.records import validate_oci_source_for_lifecycle @@ -364,6 +365,9 @@ def _bounded_error_message(message: str, *, fallback: str) -> str: return sanitized if len(sanitized) <= 512 else f"{sanitized[:509]}..." +app.add_typer(create_benchmark_app(_invoke, _emit_result), name="benchmark") + + def create_cli() -> click.Command: """Create the Slurm CLI group.""" return typer.main.get_command(app) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py new file mode 100644 index 000000000..83c3f56b4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Thin benchmark commands for the Slurm CLI.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import partial +from pathlib import Path + +import typer +from pydantic import BaseModel + +from data_designer.slurm.config import load_benchmark_config +from data_designer.slurm.services import create_slurm_benchmark_service +from data_designer.slurm.services.errors import SlurmServiceOperation + +_Invoke = Callable[[SlurmServiceOperation, Callable[[], BaseModel]], BaseModel] +_Emit = Callable[[BaseModel], None] + + +def _run_benchmark( + benchmark_file: Path, + *, + profile_file: Path | None, + cluster: str | None, + force: bool, +) -> BaseModel: + config = load_benchmark_config(benchmark_file) + service = create_slurm_benchmark_service(profile_file=profile_file, cluster=cluster) + return service.run(config, source_root=benchmark_file.resolve().parent, force=force) + + +def _analyze_benchmark( + benchmark: str, + *, + profile_file: Path | None, + cluster: str | None, + refresh: bool, + fail_if_incomplete: bool, +) -> BaseModel: + benchmark_id = Path(benchmark.rstrip("/")).name + service = create_slurm_benchmark_service(profile_file=profile_file, cluster=cluster) + return service.analyze( + benchmark_id, + refresh_state=refresh, + fail_if_incomplete=fail_if_incomplete, + ) + + +def create_benchmark_app(invoke: _Invoke, emit: _Emit) -> typer.Typer: + """Create benchmark commands using the root CLI error and output policy.""" + benchmark_app = typer.Typer(help="Run and analyze Slurm benchmarks", no_args_is_help=True) + + @benchmark_app.command("run") + def run_command( + benchmark_file: Path = typer.Argument(..., exists=True, dir_okay=False, readable=True), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), + force: bool = typer.Option(False, "--force"), + ) -> None: + """Expand and submit ordinary child runs.""" + emit( + invoke( + SlurmServiceOperation.RUN_BENCHMARK, + partial( + _run_benchmark, + benchmark_file, + profile_file=profile_file, + cluster=cluster, + force=force, + ), + ) + ) + + @benchmark_app.command("analyze") + def analyze_command( + benchmark: str = typer.Argument(..., help="Managed benchmark ID or benchmark directory"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), + refresh: bool = typer.Option(False, "--refresh-state", "--refresh"), + fail_if_incomplete: bool = typer.Option(False, "--fail-if-incomplete"), + ) -> None: + """Write one point-in-time report from ordinary child state.""" + emit( + invoke( + SlurmServiceOperation.ANALYZE_BENCHMARK, + partial( + _analyze_benchmark, + benchmark, + profile_file=profile_file, + cluster=cluster, + refresh=refresh, + fail_if_incomplete=fail_if_incomplete, + ), + ) + ) + + return benchmark_app + + +__all__ = ["create_benchmark_app"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 00b7f9bf3..3fca0ea0f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -33,6 +33,7 @@ from data_designer.slurm.config.loading import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + load_benchmark_config, load_builder_payload, load_profile_catalog, load_run_config, @@ -120,6 +121,7 @@ "VllmServerConfig", "injected_profile", "load_builder_payload", + "load_benchmark_config", "load_profile_catalog", "load_run_config", "resolve_profile", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 5c43b9f64..b4cad4be5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -17,6 +17,7 @@ from yaml.nodes import MappingNode from data_designer.slurm._errors import format_parse_error, format_validation_error +from data_designer.slurm.config.benchmark import DataDesignerSlurmBenchmarkConfig from data_designer.slurm.config.errors import SlurmConfigLoadError from data_designer.slurm.config.profiles import ( SelectedSlurmProfile, @@ -30,7 +31,12 @@ PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" -_ConfigT = TypeVar("_ConfigT", DataDesignerSlurmConfig, SlurmProfileCatalog) +_ConfigT = TypeVar( + "_ConfigT", + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + SlurmProfileCatalog, +) _HostnameResolver = Callable[[], tuple[str, ...]] @@ -69,6 +75,11 @@ def load_run_config(path: str | Path) -> DataDesignerSlurmConfig: return _load_config(path, DataDesignerSlurmConfig) +def load_benchmark_config(path: str | Path) -> DataDesignerSlurmBenchmarkConfig: + """Load one strict local YAML or JSON benchmark declaration.""" + return _load_config(path, DataDesignerSlurmBenchmarkConfig) + + def load_profile_catalog(path: str | Path) -> SlurmProfileCatalog: """Load one strict local YAML or JSON cluster-profile catalog.""" return _load_config(path, SlurmProfileCatalog) @@ -161,6 +172,8 @@ def _load_config(path: str | Path, config_type: type[_ConfigT]) -> _ConfigT: payload = _parse_mapping(contents, suffix=resolved_path.suffix) if config_type is DataDesignerSlurmConfig: _reject_run_environment_interpolation(payload) + elif config_type is DataDesignerSlurmBenchmarkConfig: + _reject_benchmark_environment_interpolation(payload) else: _reject_environment_interpolation(payload) return config_type.model_validate(payload) @@ -223,6 +236,20 @@ def _reject_run_environment_interpolation(payload: Mapping[str, object]) -> None _reject_environment_interpolation(builder_value) +def _reject_benchmark_environment_interpolation(payload: Mapping[str, object]) -> None: + for key, value in payload.items(): + _reject_environment_interpolation(key) + if key != "base_run" or not isinstance(value, Mapping): + _reject_environment_interpolation(value) + continue + for base_key, base_value in value.items(): + _reject_environment_interpolation(base_key) + if base_key == "inline" and isinstance(base_value, Mapping): + _reject_run_environment_interpolation(base_value) + else: + _reject_environment_interpolation(base_value) + + def _resolve_profile_path( explicit_path: str | Path | None, *, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index c8abeda28..9f83c0c92 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -531,13 +531,15 @@ def _publish_readiness(self, state: ReadinessState) -> None: if self._readiness is not None and not _can_advance_readiness(self._readiness.state, state): return revision = 1 if self._readiness is None else self._readiness.revision + 1 + timestamp = self._now() readiness = AttemptReadiness( schema_version=1, run_id=self._attempt.run_id, shard_id=self._attempt.shard_id, attempt_id=self._attempt.attempt_id, revision=revision, - updated_at=self._now(), + updated_at=timestamp, + started_at=self._attempt.updated_at if self._readiness is None else self._readiness.started_at, state=state, deployments=tuple( DeploymentReadiness( @@ -578,12 +580,6 @@ def _publish_stopped_readiness(self) -> None: for status in self._statuses: status.state = ReadinessState.STOPPED status.ready_backends = 0 - status.last_probe = _probe_evidence( - self._now(), - ProbeOutcome.SUCCESS, - "runtime_stopped", - "allocation processes stopped", - ) self._publish_readiness(ReadinessState.STOPPED) def _publish_terminal_attempt( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 28a0ea265..d7f125a36 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -136,6 +136,7 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non attempt_id=context.attempt.attempt_id, revision=previous.revision + 1, updated_at=timestamp, + started_at=previous.started_at, state=ReadinessState.READY, deployments=tuple( DeploymentReadiness( @@ -299,6 +300,7 @@ def _readiness( attempt_id=context.attempt.attempt_id, revision=1 if previous is None else previous.revision + 1, updated_at=timestamp, + started_at=timestamp if previous is None else previous.started_at, state=state, deployments=tuple( DeploymentReadiness( @@ -327,6 +329,7 @@ def _write_failed_and_stopped_readiness(context: AllocationContext, writer: Slur attempt_id=previous.attempt_id, revision=previous.revision + 1, updated_at=failed_at, + started_at=previous.started_at, state=ReadinessState.FAILED, deployments=tuple( deployment.model_copy( @@ -369,18 +372,13 @@ def _write_stopped_readiness( attempt_id=previous.attempt_id, revision=previous.revision + 1, updated_at=timestamp, + started_at=previous.started_at, state=ReadinessState.STOPPED, deployments=tuple( deployment.model_copy( update={ "state": ReadinessState.STOPPED, "ready_backends": 0, - "last_probe": _probe( - timestamp, - ProbeOutcome.SUCCESS, - "runtime_stopped", - "allocation processes stopped", - ), } ) for deployment in previous.deployments diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 79722df62..e6a3107f1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -36,12 +36,17 @@ if TYPE_CHECKING: from data_designer.slurm.services.wiring import ( # noqa: F401 SlurmRunArtifactPublisher, + create_slurm_benchmark_service, create_slurm_image_service, create_slurm_run_service, ) _LAZY_IMPORTS = { "SlurmRunArtifactPublisher": ("data_designer.slurm.services.wiring", "SlurmRunArtifactPublisher"), + "create_slurm_benchmark_service": ( + "data_designer.slurm.services.wiring", + "create_slurm_benchmark_service", + ), "create_slurm_image_service": ("data_designer.slurm.services.wiring", "create_slurm_image_service"), "create_slurm_run_service": ("data_designer.slurm.services.wiring", "create_slurm_run_service"), } @@ -71,6 +76,7 @@ "SlurmServiceError", "SlurmServiceErrorCode", "SlurmServiceOperation", + "create_slurm_benchmark_service", "create_slurm_image_service", "create_slurm_profile_service", "create_slurm_run_service", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py index 5ae28389c..6e01c2d46 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -38,6 +38,7 @@ publish_immutable_text, sync_directory, ) +from data_designer.slurm.state.storage import StateStorage _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _TEMPORARY_PREFIX = ".artifact." @@ -62,8 +63,16 @@ def initialize( ) -> None: if force: raise StateConflictError("force cannot replace durable run state") - created_at = self._clock() writer = SlurmStateWriter(self._workspace_root, plan.run_id) + try: + created_at = writer.load_run().created_at + except StateNotFoundError: + try: + created_at = ( + StateStorage(self._workspace_root, plan.run_id).read_shard(plan.shards[0].shard_id).created_at + ) + except FileNotFoundError: + created_at = self._clock() plan_reference = ArtifactReference( path=(writer.run_root / "resolved-plan.json").as_posix(), sha256=plan.compute_sha256(), @@ -119,14 +128,22 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte ) ) - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - """Mark every initial attempt failed after its held job is cancelled.""" + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + """Fail initial attempts owned by the cancelled held job.""" writer = SlurmStateWriter(self._workspace_root, plan.run_id) for shard in plan.shards: try: attempt = writer.load_attempt(shard.shard_id, "attempt-0001") except StateNotFoundError: continue + if attempt.scheduler is None or attempt.scheduler.array_job_id != job_id: + continue writer.update_attempt( attempt.model_copy( update={ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py index 4c4fd2192..ceb04ec03 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py @@ -5,6 +5,7 @@ from __future__ import annotations +from pathlib import Path from typing import Protocol from pydantic import TypeAdapter, ValidationError @@ -24,13 +25,25 @@ class SlurmBenchmarkBackend(Protocol): """Process benchmarks through a supported service dependency.""" - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: """Persist and start the ordinary child runs for one benchmark. Any non-``INTERNAL`` service error must contain a caller-safe message. """ - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: """Return one point-in-time benchmark report. Any non-``INTERNAL`` service error must contain a caller-safe message. @@ -49,7 +62,13 @@ class SlurmBenchmarkService: def __init__(self, backend: SlurmBenchmarkBackend) -> None: self._backend = backend - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: str | Path = ".", + force: bool = False, + ) -> BenchmarkManifest: """Start all ordinary child runs and return their immutable mapping. The returned manifest must reference the exact serialized benchmark config. @@ -60,9 +79,17 @@ def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: operation = SlurmServiceOperation.RUN_BENCHMARK if not isinstance(config, DataDesignerSlurmBenchmarkConfig): raise _make_invalid_request_error(operation, "config must be a DataDesignerSlurmBenchmarkConfig") + if not isinstance(source_root, str | Path): + raise _make_invalid_request_error(operation, "source_root must be a path") + if type(force) is not bool: + raise _make_invalid_request_error(operation, "force must be a boolean") def run_benchmark() -> BenchmarkManifest: - manifest = self._backend.run(config) + manifest = self._backend.run( + config, + source_root=Path(source_root).expanduser().resolve(), + force=force, + ) if not isinstance(manifest, BenchmarkManifest): raise TypeError("benchmark backend returned an invalid manifest") if manifest.benchmark_config.sha256 != config.compute_sha256(): @@ -71,12 +98,19 @@ def run_benchmark() -> BenchmarkManifest: return _invoke_service_backend(operation, run_benchmark) - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool = False, + fail_if_incomplete: bool = False, + ) -> BenchmarkReport: """Analyze one persisted benchmark without a resident monitor. Args: benchmark_id: Persisted benchmark identity to analyze. refresh_state: Request one point-in-time state refresh before analysis. + fail_if_incomplete: Fail after persisting a report containing non-success cases. Raises: SlurmServiceError: If the request is invalid or analysis fails. @@ -86,11 +120,15 @@ def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> B validated_id = _IDENTIFIER_ADAPTER.validate_python(benchmark_id, strict=True) except ValidationError: raise _make_invalid_request_error(operation, "benchmark_id must be a valid identifier") from None - if type(refresh_state) is not bool: - raise _make_invalid_request_error(operation, "refresh_state must be a boolean") + if type(refresh_state) is not bool or type(fail_if_incomplete) is not bool: + raise _make_invalid_request_error(operation, "analysis actions must be booleans") def analyze_benchmark() -> BenchmarkReport: - report = self._backend.analyze(validated_id, refresh_state=refresh_state) + report = self._backend.analyze( + validated_id, + refresh_state=refresh_state, + fail_if_incomplete=fail_if_incomplete, + ) if not isinstance(report, BenchmarkReport): raise TypeError("benchmark backend returned an invalid report") if report.benchmark_id != validated_id: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 691849bc4..6055e6604 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -17,6 +17,8 @@ from pydantic import JsonValue +from data_designer.slurm.benchmark.execution import SystemBenchmarkBackend +from data_designer.slurm.benchmark.observer import PersistedBenchmarkRunObserver from data_designer.slurm.client.dependencies import ( ClientDependencyResolutionError, ClientDependencyResolver, @@ -51,6 +53,7 @@ from data_designer.slurm.runtime.bundle import stage_runtime_bundle from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.services.artifacts import StateRunArtifactPublisher +from data_designer.slurm.services.benchmark import SlurmBenchmarkService from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.image_lifecycle import SlurmImageLifecycleManager from data_designer.slurm.services.images import SlurmImageService @@ -115,8 +118,14 @@ def initialize( def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: """Persist the submitted scheduler identity for every initial attempt.""" - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - """Mark initial attempts failed after a held submission is cancelled.""" + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + """Fail initial attempts owned by the cancelled held job.""" @dataclass(frozen=True, slots=True) @@ -338,7 +347,7 @@ def execute( f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", ) from error try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -353,7 +362,7 @@ def execute( pass else: try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except BaseException: pass raise @@ -369,7 +378,7 @@ def execute( f"held Slurm job {receipt.job_id} could not be released or cancelled", ) from error try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -498,9 +507,14 @@ def _record_submission( "submission state cannot be recorded", ) from None - def _record_submission_failure(self, publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan) -> None: + def _record_submission_failure( + self, + publisher: SlurmRunArtifactPublisher, + plan: ResolvedSlurmRunPlan, + job_id: int, + ) -> None: try: - publisher.record_submission_failure(plan, failed_at=self._clock()) + publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) except (StateNotFoundError, StateConflictError, SlurmStateError): raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -724,6 +738,29 @@ def create_slurm_run_service( ) -> SlurmRunService: """Create the production run service for one selected cluster profile.""" selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + return _create_slurm_run_service( + selected, + artifact_publisher=artifact_publisher, + dependency_resolver=dependency_resolver, + launcher=launcher, + run_id_factory=run_id_factory, + clock=clock, + package_version=package_version, + source_environment=source_environment, + ) + + +def _create_slurm_run_service( + selected: SelectedSlurmProfile, + *, + artifact_publisher: SlurmRunArtifactPublisher | None = None, + dependency_resolver: ClientDependencyResolver | None = None, + launcher: SlurmCommandClient | None = None, + run_id_factory: RunIdFactory | None = None, + clock: Clock | None = None, + package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, +) -> SlurmRunService: command_client = launcher or SlurmCommandClient() selected_clock = clock or _utc_now preparer = _RunPreparer( @@ -770,6 +807,54 @@ def create_slurm_image_service( return SlurmImageService(backend, backend) +def create_slurm_benchmark_service( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + artifact_publisher: SlurmRunArtifactPublisher | None = None, + dependency_resolver: ClientDependencyResolver | None = None, + launcher: SlurmCommandClient | None = None, + clock: Clock | None = None, + package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, +) -> SlurmBenchmarkService: + """Create the production benchmark service for one selected cluster profile.""" + selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + command_client = launcher or SlurmCommandClient() + selected_clock = clock or _utc_now + publisher = artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock) + dependencies = dependency_resolver or ClientDependencyResolver() + environment = dict(os.environ if source_environment is None else source_environment) + version = package_version or importlib.metadata.version("data-designer-slurm") + + def create_child_service(run_id: Identifier) -> SlurmRunService: + return _create_slurm_run_service( + selected, + artifact_publisher=publisher, + dependency_resolver=dependencies, + launcher=command_client, + run_id_factory=lambda: run_id, + clock=selected_clock, + package_version=version, + source_environment=environment, + ) + + observer = PersistedBenchmarkRunObserver( + selected.profile.workspace_root, + command_client, + selected_clock, + ) + backend = SystemBenchmarkBackend( + selected.profile.workspace_root, + create_child_service, + observer, + selected_clock, + ) + return SlurmBenchmarkService(backend) + + def _resolve_builder_payload( authored: DataDesignerSlurmConfig, source_root: Path, @@ -802,6 +887,7 @@ def _format_job_ids(job_ids: list[int]) -> str: __all__ = [ "SlurmRunArtifactPublisher", + "create_slurm_benchmark_service", "create_slurm_image_service", "create_slurm_run_service", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index 1fc65422c..15985b2a3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -33,6 +33,7 @@ CollectionStatus, ) from data_designer.slurm.state.errors import ( + SchedulerStateConflictError, SlurmStateError, StateConflictError, StateCorruptionError, @@ -163,6 +164,7 @@ "SchedulerQueueRecord", "SchedulerObservation", "SchedulerState", + "SchedulerStateConflictError", "Sha256Digest", "ShardManifest", "ShardStatus", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py index e43897dcc..c3d5bb528 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py @@ -20,3 +20,7 @@ class StateConflictError(SlurmStateError): class StateCorruptionError(SlurmStateError): """Raised when persisted state cannot be safely read or validated.""" + + +class SchedulerStateConflictError(StateCorruptionError): + """Raised when persisted winners conflict with scheduler evidence.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index 1c49d2952..aa459a6de 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -16,6 +16,7 @@ from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.base import SchedulerIdentity, SchedulerJobIdentity from data_designer.slurm.state.errors import ( + SchedulerStateConflictError, SlurmStateError, StateConflictError, StateCorruptionError, @@ -302,7 +303,7 @@ def _validate_winner_scheduler_consistency( return winning = next(status for status in statuses if status.attempt.attempt_id == winner.attempt_id) if winning.effective_state is not EffectiveAttemptState.SUCCEEDED: - raise StateCorruptionError("persisted winner conflicts with terminal scheduler evidence") + raise SchedulerStateConflictError("persisted winner conflicts with terminal scheduler evidence") def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py index b5957da9b..27fc325a2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -86,13 +86,21 @@ class AttemptReadiness(StateRecord): attempt_id: AttemptId revision: PositiveInt updated_at: datetime + started_at: datetime | None = Field(default=None, exclude_if=lambda value: value is None) state: ReadinessState deployments: tuple[DeploymentReadiness, ...] = Field(min_length=1) _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + @field_validator("started_at") + @classmethod + def validate_started_at(cls, value: datetime | None) -> datetime | None: + return None if value is None else validate_utc_timestamp(value) + @model_validator(mode="after") def validate_deployments(self) -> AttemptReadiness: + if self.started_at is not None and self.started_at > self.updated_at: + raise ValueError("readiness started_at cannot follow updated_at") _validate_deployment_uniqueness(self.deployments) _validate_probe_chronology(self.deployments, self.updated_at) _validate_attempt_state(self.state, self.deployments) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 473150f43..3dd9ab644 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -70,6 +70,7 @@ def validate_readiness_transition( _require(previous.shard_id == current.shard_id, "readiness shard_id cannot change") _require(previous.attempt_id == current.attempt_id, "readiness attempt_id cannot change") _require(current.revision == previous.revision + 1, "readiness revision must increase by exactly one") + _require(current.started_at == previous.started_at, "readiness started_at cannot change") _require(current.updated_at >= previous.updated_at, "readiness updated_at cannot move backward") _require( current.state in _ALLOWED_READINESS_TRANSITIONS[previous.state], diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 4f6f60694..39e2e6f71 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -53,6 +53,7 @@ _RUNTIME_DIRECTORY_NAME = "runtime" _RESUME_LOCK_FILENAME = "resume.lock" _LOCK_DIRECTORY_NAME = ".locks" +_SUBMISSION_LOCK_PREFIX = "submission" _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _ATTEMPT_NAME_PATTERN = re.compile(r"^attempt-[0-9]{4,}$") _RecordT = TypeVar("_RecordT", bound=ContractRecord) @@ -144,6 +145,20 @@ def acquire_run_lock(self) -> Iterator[None]: except FileNotFoundError as error: raise StateNotFoundError(f"run {self.run_id!r} is not initialized") from error + @contextmanager + def acquire_submission_lock(self) -> Iterator[None]: + """Serialize the initial submission decision for one run identity.""" + self.ensure_storage() + with open_verified_directory(self.runs_root, require_private=True) as runs_descriptor: + with open_verified_child_directory( + runs_descriptor, + _LOCK_DIRECTORY_NAME, + self.locks_root, + ) as locks_descriptor: + lock_name = f"{_SUBMISSION_LOCK_PREFIX}-{self.run_id}.lock" + with acquire_file_lock(locks_descriptor, lock_name, self.locks_root / lock_name): + yield + @contextmanager def acquire_shard_lock(self, shard_id: ShardId) -> Iterator[None]: try: diff --git a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py new file mode 100644 index 000000000..907acd5ea --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from data_designer.slurm.benchmark import BenchmarkCompiler +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig + + +def test_compiler_expands_stable_ordered_ordinary_runs( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + compiled = BenchmarkCompiler.compile(benchmark_config, authored_run) + equivalent = DataDesignerSlurmBenchmarkConfig.model_validate(benchmark_config.model_dump(mode="json")) + + assert compiled == BenchmarkCompiler.compile(equivalent, authored_run) + assert tuple(case.case_id for case in compiled.cases) == ( + "two-independent-replicas-c32", + "one-two-node-replica-c32", + "two-independent-replicas-c64", + "one-two-node-replica-c64", + "two-independent-replicas-c128", + "one-two-node-replica-c128", + ) + assert len({case.child_run_id for case in compiled.cases}) == 6 + assert all(case.child_run_config.invocation.num_records == 1000 for case in compiled.cases) + assert [deployment.resources.nodes for deployment in compiled.cases[0].child_run_config.deployments] == [2, 1] + assert [deployment.topology.nodes_per_replica for deployment in compiled.cases[1].child_run_config.deployments] == [ + 2, + 1, + ] + assert compiled.cases[0].child_run_config.invocation.model_concurrency == {"generator": 32, "judge": 32} + + +def test_compiled_child_maps_are_immutable( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + compiled = BenchmarkCompiler.compile(benchmark_config, authored_run) + + with pytest.raises(TypeError, match="frozen dictionary"): + compiled.cases[0].child_run_config.invocation.model_concurrency["generator"] = 1 + + +def test_compiler_rejects_unknown_alias_at_the_boundary( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + invalid = benchmark_config.model_copy(update={"model_aliases": ["missing"]}) + + with pytest.raises(ValueError, match="unknown model aliases"): + BenchmarkCompiler.compile(invalid, authored_run) + + +@pytest.mark.parametrize( + ("concurrency", "expected_records"), + [(1, 1000), (2000, 2000), (6000, 5000)], +) +def test_adaptive_record_policy_is_bounded( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + concurrency: int, + expected_records: int, +) -> None: + config = benchmark_config.model_copy( + update={ + "concurrency_values": [concurrency], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + + compiled = BenchmarkCompiler.compile(config, authored_run) + + assert compiled.cases[0].child_run_config.invocation.num_records == expected_records diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py new file mode 100644 index 000000000..99e23af5f --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from threading import Event, Lock + +import pytest + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkAnalyzer, + BenchmarkMeasurements, + BenchmarkObservationFailure, + BenchmarkRunObservation, + BenchmarkShardMeasurements, +) +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.execution import SystemBenchmarkBackend +from data_designer.slurm.benchmark.records import ( + BenchmarkChildRun, + BenchmarkOutcome, + BenchmarkRecommendationKind, +) +from data_designer.slurm.benchmark.store import BenchmarkConflictError, BenchmarkStore +from data_designer.slurm.config import ( + BenchmarkBaseRun, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) + +NOW = datetime(2026, 8, 19, 12, tzinfo=timezone.utc) + + +class _Observer: + def __init__(self, responses=None) -> None: + self.responses = {} if responses is None else responses + self.calls = [] + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + self.calls.append((run_id, refresh_state)) + response = self.responses[run_id] + if isinstance(response, Exception): + raise response + return response + + +class _RunService: + def __init__(self, run_id, configs, submissions, calls, benchmark_root, failures) -> None: + self.run_id = run_id + self.configs = configs + self.submissions = submissions + self.calls = calls + self.benchmark_root = benchmark_root + self.failures = failures + + def execute(self, config, *, source_root, dry_run, force): + assert (self.benchmark_root / "benchmark.json").is_file() + self.calls.append((self.run_id, config, source_root, dry_run, force)) + self.configs[self.run_id] = config + if self.run_id in self.failures: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "submission unavailable", + ) + self.submissions.add(self.run_id) + return SlurmRunExecution( + run_id=self.run_id, + state="submitted", + plan_sha256="1" * 64, + shard_count=1, + job_id=4101, + ) + + +class _ConcurrentRunService(_RunService): + def __init__(self, *args, second_check: Event) -> None: + super().__init__(*args) + self.second_check = second_check + + def execute(self, config, *, source_root, dry_run, force): + self.second_check.wait(timeout=1) + return super().execute(config, source_root=source_root, dry_run=dry_run, force=force) + + +class _ConcurrentSubmissionLoader: + def __init__(self, submissions: set[str], second_check: Event) -> None: + self._submissions = submissions + self._second_check = second_check + self._lock = Lock() + self._checks = 0 + + def __call__(self, run_id: str) -> bool: + with self._lock: + self._checks += 1 + if self._checks == 2: + self._second_check.set() + return run_id in self._submissions + + +def _inline_config( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> DataDesignerSlurmBenchmarkConfig: + return benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run), + "concurrency_values": [32], + } + ) + + +def _manifest(workspace: Path, config, compiled): + store = BenchmarkStore(workspace, compiled.benchmark_id) + children = tuple( + BenchmarkChildRun( + case_id=case.case_id, + child_run_id=case.child_run_id, + child_authored_config=ArtifactReference( + path=(workspace / "runs" / case.child_run_id / "authored-config.json").as_posix(), + sha256=case.child_run_config.compute_sha256(), + ), + ) + for case in compiled.cases + ) + return store, store.build_manifest(config, children) + + +def _measurements( + actual_records: int, + boot_seconds: float, + generation_seconds: float, + wall_seconds: float, + shard_count: int, +) -> BenchmarkMeasurements: + records_per_shard = actual_records // shard_count + return BenchmarkMeasurements( + shards=tuple( + BenchmarkShardMeasurements( + actual_records=( + actual_records - records_per_shard * (shard_count - 1) + if index == shard_count - 1 + else records_per_shard + ), + boot_seconds=boot_seconds, + generation_seconds=generation_seconds, + wall_seconds=wall_seconds, + ) + for index in range(shard_count) + ) + ) + + +def test_store_is_convergent_and_rejects_tampered_metadata( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + + assert store.publish(config, authored_run, manifest) == manifest + assert store.publish(config, authored_run, manifest) == manifest + assert store.load() == (config, authored_run, manifest) + + store.config_path.write_text("{}\n") + with pytest.raises(BenchmarkConflictError): + store.load() + + +def test_run_persists_manifest_before_submission_and_is_idempotent( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + configs = {} + submissions = set() + calls = [] + failures = set() + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + first = backend.run(config, source_root=tmp_path, force=False) + second = backend.run(config, source_root=tmp_path, force=False) + + assert first == second + assert len(calls) == len(compiled.cases) + assert tuple(configs) == tuple(case.child_run_id for case in compiled.cases) + assert all(force is False for *_, force in calls) + + +def test_concurrent_runs_submit_each_deterministic_child_once( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + configs = {} + submissions = set() + calls = [] + failures = set() + second_check = Event() + submission_loader = _ConcurrentSubmissionLoader(submissions, second_check) + store, manifest = _manifest(tmp_path, config, compiled) + store.publish(config, authored_run, manifest) + + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _ConcurrentRunService( + run_id, + configs, + submissions, + calls, + store.benchmark_root, + failures, + second_check=second_check, + ), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submission_loader, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(backend.run, config, source_root=tmp_path, force=False) for _ in range(2)) + assert tuple(future.result() for future in futures) == (manifest, manifest) + + assert len(calls) == len(compiled.cases) + + +def test_initialized_child_without_submission_requires_force_to_resume( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + missing_id = compiled.cases[0].child_run_id + configs = {} + submissions = set() + calls = [] + failures = {missing_id} + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + with pytest.raises(SlurmServiceError, match="1 of 2"): + backend.run(config, source_root=tmp_path, force=False) + failures.clear() + with pytest.raises(SlurmServiceError, match="1 of 2") as conflict: + backend.run(config, source_root=tmp_path, force=False) + assert conflict.value.code is SlurmServiceErrorCode.CONFLICT + + backend.run(config, source_root=tmp_path, force=True) + + assert [run_id for run_id, *_ in calls].count(missing_id) == 2 + + +def test_partial_submission_attempts_all_children_and_resumes_only_missing( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + missing_id = compiled.cases[0].child_run_id + configs = {} + submissions = set() + calls = [] + failures = {missing_id} + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + with pytest.raises(SlurmServiceError, match="1 of 2"): + backend.run(config, source_root=tmp_path, force=True) + assert len(calls) == 2 + + failures.clear() + backend.run(config, source_root=tmp_path, force=True) + + assert [run_id for run_id, *_ in calls].count(missing_id) == 2 + assert len(calls) == 3 + assert all(force is False for *_, force in calls) + + +def test_analysis_preserves_mixed_outcomes_and_stable_order( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + successful = compiled.cases[0] + plan = multi_node_plan.model_copy(update={"run_id": successful.child_run_id}) + observer = _Observer( + { + successful.child_run_id: BenchmarkRunObservation( + authored_config=successful.child_run_config, + resolved_plan=plan, + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements(1000, 60, 120, 180, len(plan.shards)), + ), + compiled.cases[1].child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING), + } + ) + + report = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=True, + ) + repeated = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=True, + ) + + assert tuple(case.outcome for case in report.cases) == ( + BenchmarkOutcome.SUCCEEDED, + BenchmarkOutcome.MISSING, + ) + assert report.cases[0].rows_per_second == 1000 / 120 + assert report.cases[0].generation_seconds == 14340 + assert report.cases[0].target_jobs == 17 + assert report.cases[0].total_gpu_hours == report.cases[0].gpu_hours_per_job * 17 + assert report.serialize_json() == repeated.serialize_json() + assert observer.calls == [(case.child_run_id, True) for case in compiled.cases] * 2 + + +def test_analysis_keeps_measured_infeasible_case_successful( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run).model_copy( + update={"deployment_cases": [benchmark_config.deployment_cases[0]]} + ) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + case = compiled.cases[0] + observation = BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements( + case.child_run_config.invocation.num_records, + 15000, + 120, + 15120, + len(multi_node_plan.shards), + ), + ) + + report = BenchmarkAnalyzer(_Observer({case.child_run_id: observation}), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert report.cases[0].outcome is BenchmarkOutcome.SUCCEEDED + assert report.cases[0].feasible is False + assert report.cases[0].generation_seconds == 0 + assert report.cases[0].target_jobs is None + assert report.cases[0].total_gpu_hours is None + + +def test_analysis_rejects_manifest_rewritten_away_from_compiled_children( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + first = manifest.children[0] + rewritten_id = f"{first.child_run_id}-other" + rewritten = BenchmarkChildRun( + case_id=first.case_id, + child_run_id=rewritten_id, + child_authored_config=ArtifactReference( + path=(tmp_path / "runs" / rewritten_id / "authored-config.json").as_posix(), + sha256="f" * 64, + ), + ) + tampered = manifest.model_copy(update={"children": (rewritten, *manifest.children[1:])}) + + with pytest.raises(ValueError, match="deterministic child expansion"): + BenchmarkAnalyzer(_Observer(), lambda: NOW).analyze( + config, + authored_run, + tampered, + store.manifest_reference(tampered), + refresh_state=False, + ) + + +def test_analysis_recommendations_are_stable_and_select_dominating_case( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + responses = {} + for index, case in enumerate(compiled.cases, start=1): + responses[case.child_run_id] = BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements( + case.child_run_config.invocation.num_records, + 60, + 120 * index, + 180 * index, + len(multi_node_plan.shards), + ), + ) + + report = BenchmarkAnalyzer(_Observer(responses), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert tuple(item.kind for item in report.recommendations) == ( + BenchmarkRecommendationKind.PARETO, + BenchmarkRecommendationKind.MINIMUM_JOBS, + BenchmarkRecommendationKind.MINIMUM_GPU_HOURS, + ) + assert {item.case_id for item in report.recommendations} == {compiled.cases[0].case_id} + + +@pytest.mark.parametrize( + "outcome", + [ + BenchmarkOutcome.PENDING, + BenchmarkOutcome.ACCOUNTING_LAG, + BenchmarkOutcome.FAILED, + BenchmarkOutcome.INCOMPLETE, + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + ], +) +def test_analysis_keeps_each_non_success_outcome_explicit( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + outcome: BenchmarkOutcome, +) -> None: + config = _inline_config(benchmark_config, authored_run).model_copy( + update={"deployment_cases": [benchmark_config.deployment_cases[0]]} + ) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + case = compiled.cases[0] + response = ( + BenchmarkObservationFailure(outcome) + if outcome + in { + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + } + else BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=outcome, + ) + ) + + report = BenchmarkAnalyzer(_Observer({case.child_run_id: response}), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert report.cases[0].outcome is outcome + + +def test_fail_if_incomplete_persists_the_point_in_time_report_first( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + store.publish(config, authored_run, manifest) + observer = _Observer( + {case.child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING) for case in compiled.cases} + ) + backend = SystemBenchmarkBackend( + tmp_path, + lambda _: pytest.fail("analysis must not create a run service"), + observer, + lambda: NOW, + ) + + with pytest.raises(SlurmServiceError, match="2 incomplete cases"): + backend.analyze( + manifest.benchmark_id, + refresh_state=False, + fail_if_incomplete=True, + ) + + assert len(tuple(store.reports_root.iterdir())) == 1 diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 7abca1b66..c2ff3087f 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -13,12 +13,14 @@ from data_designer.slurm.config import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, DataDesignerSlurmConfigBuilder, ProfileSelectionSource, SlurmConfigBuilderError, SlurmConfigLoadError, SlurmProfileCatalog, + load_benchmark_config, load_builder_payload, load_profile_catalog, load_run_config, @@ -79,6 +81,16 @@ def test_builder_write_config_round_trips_supported_formats(tmp_path: Path, suff assert load_run_config(path) == builder.build() +def test_benchmark_loader_round_trips_strict_config( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, +) -> None: + path = tmp_path / "benchmark.json" + path.write_text(benchmark_config.serialize_json()) + + assert load_benchmark_config(path) == benchmark_config + + def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: with pytest.raises(SlurmConfigBuilderError, match="must end"): _config_builder().write_config(tmp_path / "run.txt") diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index 90009a9b4..e35ba4f98 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -262,6 +262,22 @@ def test_successful_benchmark_case_requires_complete_positive_output(field: str, BenchmarkReport.model_validate_json(json.dumps(payload)) +def test_successful_infeasible_benchmark_case_allows_zero_budgeted_generation() -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + case = payload["cases"][0] + case.update( + feasible=False, + generation_seconds=0, + target_jobs=None, + total_gpu_hours=None, + ) + payload["recommendations"] = [] + + report = BenchmarkReport.model_validate_json(json.dumps(payload)) + + assert report.cases[0].feasible is False + + def test_benchmark_report_allows_pareto_frontier_but_singleton_minima() -> None: payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) second_case = deepcopy(payload["cases"][0]) diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index df3439944..b8bf75f5f 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -193,6 +193,10 @@ def write_result_under_dataset_lease() -> None: assert ReadinessState.READY in [item.state for item in state.readiness] assert state.readiness[-1].state is ReadinessState.STOPPED assert state.readiness[-1].deployments[0].endpoint_publication is EndpointPublicationState.PUBLISHED + assert state.readiness[-1].deployments[0].last_probe is not None + assert state.readiness[-1].deployments[0].last_probe.reason_code == "endpoint_ready" + assert state.readiness[0].started_at is not None + assert {item.started_at for item in state.readiness} == {state.readiness[0].started_at} assert all(process.poll() is not None for process in runner.processes) assert {step.stdout_path.parent.name for step in runner.steps} == {"execution-00000001"} diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index b3df00eb7..df0b0e47e 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -558,23 +558,28 @@ def test_benchmark_service_delegates_run_and_analysis( benchmark_report: BenchmarkReport, ) -> None: benchmark_manifest = manifest_with_matching_config_digest + source_root = Path.cwd().resolve() backend = FakeBenchmarkBackend( - run_responses=((benchmark_config, benchmark_manifest),), - analysis_responses=(((benchmark_manifest.benchmark_id, True), benchmark_report),), + run_responses=(((benchmark_config, source_root, True), benchmark_manifest),), + analysis_responses=(((benchmark_manifest.benchmark_id, True, True), benchmark_report),), ) service = SlurmBenchmarkService(backend) - assert service.run(benchmark_config) is benchmark_manifest - assert service.analyze(benchmark_manifest.benchmark_id, refresh_state=True) is benchmark_report - assert backend.run_calls == [benchmark_config] - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, True)] + assert service.run(benchmark_config, force=True) is benchmark_manifest + assert ( + service.analyze(benchmark_manifest.benchmark_id, refresh_state=True, fail_if_incomplete=True) + is benchmark_report + ) + assert backend.run_calls == [(benchmark_config, source_root, True)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, True, True)] backend.assert_complete() def test_benchmark_service_rejects_invalid_manifest( benchmark_config: DataDesignerSlurmBenchmarkConfig, ) -> None: - backend = FakeBenchmarkBackend(run_responses=((benchmark_config, object()),)) # type: ignore[arg-type] + request = (benchmark_config, Path.cwd().resolve(), False) + backend = FakeBenchmarkBackend(run_responses=((request, object()),)) # type: ignore[arg-type] with pytest.raises(SlurmServiceError) as caught: SlurmBenchmarkService(backend).run(benchmark_config) @@ -582,13 +587,13 @@ def test_benchmark_service_rejects_invalid_manifest( assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK assert str(caught.value) == "run benchmark failed" - assert backend.run_calls == [benchmark_config] + assert backend.run_calls == [request] backend.assert_complete() def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkManifest) -> None: backend = FakeBenchmarkBackend( - analysis_responses=(((benchmark_manifest.benchmark_id, False), object()),), # type: ignore[arg-type] + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), object()),), # type: ignore[arg-type] ) with pytest.raises(SlurmServiceError) as caught: @@ -597,7 +602,7 @@ def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkM assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK assert str(caught.value) == "analyze benchmark failed" - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False, False)] backend.assert_complete() @@ -607,14 +612,15 @@ def test_benchmark_service_rejects_manifest_for_another_config( ) -> None: reference = manifest_with_matching_config_digest.benchmark_config.model_copy(update={"sha256": "0" * 64}) manifest = manifest_with_matching_config_digest.model_copy(update={"benchmark_config": reference}) - backend = FakeBenchmarkBackend(run_responses=((benchmark_config, manifest),)) + request = (benchmark_config, Path.cwd().resolve(), False) + backend = FakeBenchmarkBackend(run_responses=((request, manifest),)) with pytest.raises(SlurmServiceError) as caught: SlurmBenchmarkService(backend).run(benchmark_config) assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK - assert backend.run_calls == [benchmark_config] + assert backend.run_calls == [request] backend.assert_complete() @@ -629,17 +635,22 @@ def test_benchmark_service_rejects_untyped_config() -> None: @pytest.mark.parametrize( - ("benchmark_id", "refresh_state"), - [("", False), ("invalid/id", False), ("benchmark-001", 1)], + ("benchmark_id", "refresh_state", "fail_if_incomplete"), + [("", False, False), ("invalid/id", False, False), ("benchmark-001", 1, False), ("benchmark-001", False, 1)], ) def test_benchmark_service_validates_analysis_actions( benchmark_id: object, refresh_state: object, + fail_if_incomplete: object, ) -> None: service = SlurmBenchmarkService(FakeBenchmarkBackend()) with pytest.raises(SlurmServiceError) as caught: - service.analyze(benchmark_id, refresh_state=refresh_state) # type: ignore[arg-type] + service.analyze( # type: ignore[arg-type] + benchmark_id, + refresh_state=refresh_state, + fail_if_incomplete=fail_if_incomplete, + ) assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK @@ -651,7 +662,7 @@ def test_benchmark_service_rejects_uncorrelated_report( ) -> None: report = benchmark_report.model_copy(update={"benchmark_id": "other-benchmark"}) backend = FakeBenchmarkBackend( - analysis_responses=(((benchmark_manifest.benchmark_id, False), report),), + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), report),), ) with pytest.raises(SlurmServiceError) as caught: @@ -659,7 +670,7 @@ def test_benchmark_service_rejects_uncorrelated_report( assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False, False)] backend.assert_complete() diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index cd453f598..02f072ebe 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -12,13 +12,17 @@ import pytest import data_designer.slurm.services.retry_collection as retry_collection_module +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler from data_designer.slurm.client.dependencies import ResolvedClientDependencies from data_designer.slurm.config import ( + BenchmarkBaseRun, BuilderInput, + DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, SecretRef, SlurmProfile, SlurmProfileCatalog, + select_profile, ) from data_designer.slurm.contracts import ArtifactReference, canonical_json from data_designer.slurm.images.records import RegisteredImage @@ -36,6 +40,7 @@ SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, + create_slurm_benchmark_service, create_slurm_image_service, create_slurm_run_service, ) @@ -59,9 +64,11 @@ def __init__( self, gpu_counts: tuple[int, ...] = (), *, - submission_job_ids: tuple[int, ...] = (42,), + submission_job_ids: tuple[int, ...] | None = None, cancel_error: Exception | None = None, + job_id: int = 42, release_error: Exception | None = None, + submission_error: Exception | None = None, ) -> None: self.submissions: list[str] = [] self.cancellations: list[int] = [] @@ -71,10 +78,11 @@ def __init__( self.queue_entries: tuple[SlurmQueueEntry, ...] = () self.accounting_entries: tuple[SlurmAccountingEntry, ...] = () self.submission_matches: tuple[SlurmSubmissionMatch, ...] = () - self.submission_job_ids = submission_job_ids + self.submission_job_ids = (job_id,) if submission_job_ids is None else submission_job_ids self.gpu_counts = gpu_counts self.cancel_error = cancel_error self.release_error = release_error + self.submission_error = submission_error def submit_script( self, @@ -86,6 +94,8 @@ def submit_script( self.submissions.append(script) self.held_submissions.append(hold) self.exported_environments.append(dict(export_environment or {})) + if self.submission_error is not None: + raise self.submission_error index = min(len(self.submissions) - 1, len(self.submission_job_ids) - 1) return SlurmJobSubmissionReceipt(job_id=self.submission_job_ids[index]) @@ -129,8 +139,9 @@ def __init__( submission_error: BaseException | None = None, ) -> None: self.initializations: list[tuple[str, bool]] = [] + self.plans: list[ResolvedSlurmRunPlan] = [] self.submissions: list[tuple[str, int, datetime]] = [] - self.submission_failures: list[tuple[str, datetime]] = [] + self.submission_failures: list[tuple[str, int, datetime]] = [] self.initialization_error = initialization_error self.submission_error = submission_error @@ -148,6 +159,7 @@ def initialize( assert plan.authored_config.sha256 == authored.compute_sha256() assert builder_payload is None assert all(path.is_file() for path in dependencies.wheel_sources) + self.plans.append(plan) self.initializations.append((plan.run_id, force)) def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: @@ -155,8 +167,14 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte raise self.submission_error self.submissions.append((plan.run_id, job_id, submitted_at)) - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - self.submission_failures.append((plan.run_id, failed_at)) + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + self.submission_failures.append((plan.run_id, job_id, failed_at)) def _profile(tmp_path: Path, profile_catalog: SlurmProfileCatalog) -> SlurmProfile: @@ -213,6 +231,116 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( assert launcher.submissions == [] +def test_production_benchmark_wiring_submits_each_case_as_an_ordinary_run( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + publisher = _Publisher() + service = create_slurm_benchmark_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + artifact_publisher=publisher, + package_version="0.9.2", + ) + + manifest = service.run(config, source_root=tmp_path, force=True) + + assert len(manifest.children) == 1 + assert publisher.initializations == [(manifest.children[0].child_run_id, False)] + assert len(launcher.submissions) == 1 + assert (tmp_path / "benchmarks" / manifest.benchmark_id / "benchmark.json").is_file() + + +def test_production_benchmark_children_preserve_catalog_profile_provenance( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + catalog = profile_catalog.model_copy( + update={ + "clusters": { + name: profile.model_copy(update={"workspace_root": tmp_path.as_posix()}) + for name, profile in profile_catalog.clusters.items() + } + } + ) + _register_images(tmp_path, authored_run_single, single_node_plan) + publisher = _Publisher() + service = create_slurm_benchmark_service( + catalog=catalog, + cluster="primary", + launcher=_Launcher(), # type: ignore[arg-type] + artifact_publisher=publisher, + package_version="0.9.2", + ) + + service.run(config, source_root=tmp_path) + + assert publisher.plans[0].selected_profile == select_profile(catalog, cluster="primary") + + +@pytest.mark.parametrize("drop_run_manifest", [False, True]) +def test_production_benchmark_force_resumes_initialized_child_without_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + drop_run_manifest: bool, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + child_run_id = BenchmarkCompiler.compile(config, authored_run_single).cases[0].child_run_id + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher(submission_error=SlurmLauncherError("unavailable")) + service = create_slurm_benchmark_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="1 of 1") as unavailable: + service.run(config, source_root=tmp_path) + assert unavailable.value.code is SlurmServiceErrorCode.UNAVAILABLE + run_root = tmp_path / "runs" / child_run_id + if drop_run_manifest: + (run_root / "run.json").unlink() + launcher.submission_error = None + + manifest = service.run(config, source_root=tmp_path, force=True) + + writer = SlurmStateWriter(tmp_path, manifest.children[0].child_run_id) + shard = writer.load_shards()[0] + assert len(launcher.submissions) == 2 + assert writer.load_attempts(shard.shard_id)[0].scheduler is not None + + def test_public_plan_resolves_builder_relative_to_explicit_source_root( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -882,6 +1010,39 @@ def test_recording_conflict_cancels_the_accepted_job( assert len(publisher.submission_failures) == 1 +def test_recording_conflict_does_not_fail_another_submission_attempt( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + profile = _profile(tmp_path, profile_catalog) + first = create_slurm_run_service( + profile=profile, + launcher=_Launcher(job_id=41), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + first.execute(authored_run_single, source_root=tmp_path) + launcher = _Launcher(job_id=42) + second = create_slurm_run_service( + profile=profile, + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + second.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.cancellations == [42] + attempt = SlurmStateWriter(tmp_path, "run-wired").load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.SUBMITTED + assert attempt.scheduler == SchedulerIdentity(array_job_id=41, array_task_id=0) + + def test_partial_submission_recording_failure_cancels_job_and_fails_created_attempts( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py index de32b94d1..9db2d7af0 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py @@ -119,21 +119,38 @@ class FakeBenchmarkBackend(SlurmBenchmarkBackend): def __init__( self, *, - run_responses: Iterable[tuple[DataDesignerSlurmBenchmarkConfig, BenchmarkManifest | BaseException]] = (), - analysis_responses: Iterable[tuple[tuple[Identifier, bool], BenchmarkReport | BaseException]] = (), + run_responses: Iterable[ + tuple[tuple[DataDesignerSlurmBenchmarkConfig, Path, bool], BenchmarkManifest | BaseException] + ] = (), + analysis_responses: Iterable[tuple[tuple[Identifier, bool, bool], BenchmarkReport | BaseException]] = (), ) -> None: self._run_script = _ScriptedResponses(run_responses) self._analysis_script = _ScriptedResponses(analysis_responses) self.run_calls = self._run_script.calls self.analysis_calls = self._analysis_script.calls - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: """Return the manifest scripted for one exact benchmark config.""" - return self._run_script.next(config, operation="benchmark run") + return self._run_script.next((config, source_root, force), operation="benchmark run") - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: """Return the report scripted for one benchmark and refresh action.""" - return self._analysis_script.next((benchmark_id, refresh_state), operation="benchmark analysis") + return self._analysis_script.next( + (benchmark_id, refresh_state, fail_if_incomplete), + operation="benchmark analysis", + ) def assert_complete(self) -> None: """Assert that every scripted benchmark response was consumed.""" diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 41bc6dc73..4781fdccf 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -12,8 +12,11 @@ from typing import cast import pytest -from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask +from data_designer.slurm.benchmark.analysis import BenchmarkObservationFailure +from data_designer.slurm.benchmark.observer import PersistedBenchmarkRunObserver +from data_designer.slurm.benchmark.records import BenchmarkOutcome from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 @@ -23,13 +26,19 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptReadiness, AttemptTerminalClassification, CandidateOutcome, CandidateOutputFile, CandidateOutputManifest, + DeploymentReadiness, EffectiveAttemptState, EffectiveRunState, + EndpointPublicationState, GenerationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, RunManifest, SchedulerIdentity, SchedulerJobIdentity, @@ -193,6 +202,290 @@ def test_fresh_process_refresh_persists_one_fixed_accounting_lag_deadline( assert stat.S_IMODE(scheduler_path.stat().st_mode) == 0o600 +def test_benchmark_observer_refreshes_from_a_fresh_process( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state=None) + observed_at = case.created_at + timedelta(minutes=3) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: observed_at, + ).observe(case.plan.run_id, refresh_state=True) + + assert observation.outcome is BenchmarkOutcome.ACCOUNTING_LAG + assert observation.authored_config == authored_run_single + assert observation.resolved_plan == case.plan + + +def test_benchmark_observer_preserves_missing_and_tampered_child_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan, submitted=False) + observer = PersistedBenchmarkRunObserver(case.workspace, _StaticSchedulerClient((), ()), lambda: case.created_at) + + with pytest.raises(BenchmarkObservationFailure) as missing: + observer.observe("run-missing", refresh_state=False) + assert missing.value.outcome is BenchmarkOutcome.MISSING + + authored_path = case.writer.run_root / "authored-config.json" + authored_path.write_text("{}\n") + with pytest.raises(BenchmarkObservationFailure) as stale: + observer.observe(case.plan.run_id, refresh_state=False) + assert stale.value.outcome is BenchmarkOutcome.STALE + + +@pytest.mark.parametrize( + ("attempt_state", "terminal_classification", "expected"), + [ + (AttemptLifecycleState.SUBMITTED, None, BenchmarkOutcome.PENDING), + ( + AttemptLifecycleState.FAILED, + AttemptTerminalClassification.FAILED, + BenchmarkOutcome.FAILED, + ), + ], +) +def test_benchmark_observer_reconstructs_persisted_non_success_outcomes( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + attempt_state: AttemptLifecycleState, + terminal_classification: AttemptTerminalClassification | None, + expected: BenchmarkOutcome, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + if attempt_state is not AttemptLifecycleState.SUBMITTED: + case.writer.update_attempt( + _copy_attempt( + case.attempt, + state=attempt_state, + terminal_classification=terminal_classification, + ) + ) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + _StaticSchedulerClient((), ()), + lambda: case.created_at, + ).observe(case.plan.run_id, refresh_state=False) + + assert observation.outcome is expected + assert observation.measurements is None + + +def test_benchmark_observer_requires_measurements_for_persisted_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + _publish_winner_state(case) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + _StaticSchedulerClient((), ()), + lambda: case.created_at, + ).observe(case.plan.run_id, refresh_state=False) + + assert observation.outcome is BenchmarkOutcome.SUCCEEDED + assert observation.measurements is None + + +def test_benchmark_observer_reads_complete_measurements_from_ordinary_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + ready_at = case.created_at + timedelta(minutes=2) + deployment = case.plan.deployments[0] + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=1, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.PENDING, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.PENDING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ), + ), + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=2, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.STARTING, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.STARTING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ), + ), + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=3, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.READY, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.READY, + expected_backends=deployment.topology.replica_count, + ready_backends=deployment.topology.replica_count, + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=ProbeEvidence( + observed_at=ready_at, + outcome=ProbeOutcome.SUCCESS, + reason_code="backend_ready", + redacted_message="Backend is ready", + ), + ), + ), + ) + ) + _publish_winner_state(case) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state="COMPLETED") + + observation = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: case.created_at + timedelta(minutes=8), + ).observe(case.plan.run_id, refresh_state=True) + + assert observation.outcome is BenchmarkOutcome.SUCCEEDED + assert observation.measurements is not None + assert observation.measurements.actual_records == authored_run_single.invocation.num_records + assert observation.measurements.boot_seconds == 60 + assert observation.measurements.generation_seconds == 180 + assert observation.measurements.wall_seconds == 300 + + +def test_benchmark_observer_excludes_queue_stagger_from_cross_shard_timings( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + cases = _initialized_cases(tmp_path, authored_run, multi_node_plan) + for index, case in enumerate(cases): + started_at = case.created_at + timedelta(minutes=2, hours=index) + _publish_ready_state(case, case.attempt, started_at=started_at) + _publish_winner_state( + case, + completed_at=started_at + timedelta(minutes=5), + stopped_at=started_at + timedelta(minutes=6), + ) + + observation = PersistedBenchmarkRunObserver( + cases[0].workspace, + _StaticSchedulerClient((), ()), + lambda: cases[-1].created_at + timedelta(hours=2), + ).observe(cases[0].plan.run_id, refresh_state=False) + + assert observation.measurements is not None + assert observation.measurements.actual_records == authored_run.invocation.num_records + assert observation.measurements.boot_seconds == 60 + assert observation.measurements.generation_seconds == 240 + assert observation.measurements.wall_seconds == 360 + + +def test_benchmark_observer_ignores_superseded_attempt_evidence_after_retry_wins( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + case.writer.update_attempt( + _copy_attempt( + case.attempt, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.PREEMPTED, + updated_at=case.attempt.created_at + timedelta(minutes=2), + ) + ) + retry = AttemptManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id="attempt-0002", + attempt_ordinal=2, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4102, array_task_id=0), + created_at=case.attempt.created_at + timedelta(minutes=3), + updated_at=case.attempt.created_at + timedelta(minutes=3), + ) + case.writer.create_attempt(retry) + _publish_ready_state(case, retry) + _, winner = _publish_winner_state(case, retry) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4102, array_task_id=0)),)), + ) + ) + observer = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(runner), + lambda: winner.published_at + timedelta(minutes=1), + ) + + persisted = observer.observe(case.plan.run_id, refresh_state=False) + first_scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + retry_scheduler = cast(SchedulerIdentity, retry.scheduler) + SlurmCommandClient(runner).submit("first.sbatch") + SlurmCommandClient(runner).submit("retry.sbatch") + runner.set_task_state(first_scheduler, queue_state=None, accounting_state="PREEMPTED") + runner.set_task_state(retry_scheduler, queue_state=None, accounting_state="COMPLETED") + refreshed = observer.observe(case.plan.run_id, refresh_state=True) + + assert persisted.outcome is BenchmarkOutcome.SUCCEEDED + assert persisted.measurements is not None + assert persisted.measurements.boot_seconds == 60 + assert refreshed.outcome is BenchmarkOutcome.SUCCEEDED + assert refreshed.measurements == persisted.measurements + + def test_refresh_uses_terminal_accounting_over_stale_active_queue_state( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -352,6 +645,15 @@ def test_refresh_rejects_winner_that_conflicts_with_terminal_scheduler_evidence( SlurmCommandClient(fake_slurm_runner), ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + observer = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: winner.published_at + timedelta(minutes=1), + ) + with pytest.raises(BenchmarkObservationFailure) as conflict: + observer.observe(case.plan.run_id, refresh_state=True) + assert conflict.value.outcome is BenchmarkOutcome.SCHEDULER_INCONSISTENT + def test_refresh_reports_a_validated_winner_as_succeeded( tmp_path: Path, @@ -589,6 +891,18 @@ def _initialized_case( *, submitted: bool = True, ) -> _ReconciliationCase: + cases = _initialized_cases(tmp_path, authored_config, plan, submitted=submitted) + assert len(cases) == 1 + return cases[0] + + +def _initialized_cases( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + *, + submitted: bool = True, +) -> tuple[_ReconciliationCase, ...]: workspace = tmp_path / "workspace" workspace.mkdir() relocated_plan = _relocate_plan(plan, workspace) @@ -603,35 +917,44 @@ def _initialized_case( path=(run_root / "resolved-plan.json").as_posix(), sha256=relocated_plan.compute_sha256(), ), - shard_count=1, + shard_count=len(relocated_plan.shards), ) - planned_shard = relocated_plan.shards[0] - shard = ShardManifest( - schema_version=1, - run_id=relocated_plan.run_id, - shard_id=planned_shard.shard_id, - shard_index=planned_shard.shard_index, - record_range=planned_shard.record_range, - input_partition=planned_shard.input_partition, - resume_workspace=planned_shard.resume_workspace, - created_at=created_at, + shards = tuple( + ShardManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=planned_shard.shard_id, + shard_index=planned_shard.shard_index, + record_range=planned_shard.record_range, + input_partition=planned_shard.input_partition, + resume_workspace=planned_shard.resume_workspace, + created_at=created_at, + ) + for planned_shard in relocated_plan.shards ) writer = SlurmStateWriter(workspace, relocated_plan.run_id) - writer.initialize_run(authored_config, relocated_plan, run, (shard,)) - attempt = AttemptManifest( - schema_version=1, - run_id=relocated_plan.run_id, - shard_id=shard.shard_id, - attempt_id="attempt-0001", - attempt_ordinal=1, - resolved_plan=run.resolved_plan, - state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, - scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0) if submitted else None, - created_at=created_at + timedelta(minutes=1), - updated_at=created_at + timedelta(minutes=2), + writer.initialize_run(authored_config, relocated_plan, run, shards) + attempts = tuple( + AttemptManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, + scheduler=(SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index) if submitted else None), + created_at=created_at + timedelta(minutes=1), + updated_at=created_at + timedelta(minutes=2), + ) + for shard in shards + ) + for attempt in attempts: + writer.create_attempt(attempt) + return tuple( + _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, writer, created_at) + for shard, attempt in zip(shards, attempts, strict=True) ) - writer.create_attempt(attempt) - return _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, writer, created_at) def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: @@ -646,23 +969,36 @@ def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurm return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) -def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, ShardWinner]: +def _publish_winner_state( + case: _ReconciliationCase, + attempt: AttemptManifest | None = None, + *, + completed_at: datetime | None = None, + stopped_at: datetime | None = None, +) -> tuple[AttemptManifest, ShardWinner]: + attempt = case.attempt if attempt is None else attempt + result_completed_at = completed_at or attempt.created_at + timedelta(minutes=4) + attempt_stopped_at = stopped_at or attempt.created_at + timedelta(minutes=5) running = _copy_attempt( - case.attempt, + attempt, state=AttemptLifecycleState.RUNNING, - updated_at=case.created_at + timedelta(minutes=3), + updated_at=attempt.created_at + timedelta(minutes=2), ) case.writer.update_attempt(running) - candidate_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/output-manifest.json" + candidate_path = ( + case.writer.run_root / f"shards/{attempt.shard_id}/attempts/{attempt.attempt_id}/output-manifest.json" + ) dataset_path = candidate_path.parent / "dataset" - requested = case.plan.shards[0].requested_records + requested = next(shard.requested_records for shard in case.plan.shards if shard.shard_id == case.shard.shard_id) candidate = CandidateOutputManifest( schema_version=1, run_id=case.plan.run_id, shard_id=running.shard_id, attempt_id=running.attempt_id, attempt_ordinal=running.attempt_ordinal, - created_at=case.created_at + timedelta(minutes=4), + created_at=( + result_completed_at - timedelta(minutes=1) if completed_at else attempt.created_at + timedelta(minutes=3) + ), dataset_path=dataset_path.as_posix(), requested_records=requested, actual_records=requested, @@ -684,7 +1020,7 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S run_id=case.plan.run_id, shard_id=running.shard_id, attempt_id=running.attempt_id, - completed_at=case.created_at + timedelta(minutes=5), + completed_at=result_completed_at, requested_records=requested, actual_records=requested, outcome=ClientOutcome.COMPLETE, @@ -700,7 +1036,7 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S state=AttemptLifecycleState.SUCCEEDED, terminal_classification=AttemptTerminalClassification.SUCCEEDED, candidate_output=candidate_reference, - updated_at=case.created_at + timedelta(minutes=6), + updated_at=attempt_stopped_at, ) case.writer.update_attempt(completed) winner = ShardWinner( @@ -710,14 +1046,87 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S attempt_id=completed.attempt_id, attempt_ordinal=completed.attempt_ordinal, candidate_manifest=candidate_reference, - published_at=case.created_at + timedelta(minutes=7), + published_at=attempt_stopped_at + timedelta(minutes=1), ) - winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + winner_path = case.writer.run_root / f"shards/{attempt.shard_id}/winner.json" winner_path.write_text(winner.serialize_json()) winner_path.chmod(0o600) return completed, winner +def _publish_ready_state( + case: _ReconciliationCase, + attempt: AttemptManifest, + *, + started_at: datetime | None = None, +) -> None: + started_at = started_at or attempt.created_at + timedelta(minutes=2) + ready_at = started_at + timedelta(minutes=1) + pending = AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=started_at, + started_at=started_at, + state=ReadinessState.PENDING, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.PENDING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + for deployment in case.plan.deployments + ), + ) + case.writer.write_readiness(pending) + case.writer.write_readiness( + pending.model_copy( + update={ + "revision": 2, + "state": ReadinessState.STARTING, + "deployments": tuple( + deployment.model_copy(update={"state": ReadinessState.STARTING}) + for deployment in pending.deployments + ), + } + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=3, + updated_at=ready_at, + started_at=started_at, + state=ReadinessState.READY, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.READY, + expected_backends=deployment.topology.replica_count, + ready_backends=deployment.topology.replica_count, + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=ProbeEvidence( + observed_at=ready_at, + outcome=ProbeOutcome.SUCCESS, + reason_code="backend_ready", + redacted_message="Backend is ready", + ), + ) + for deployment in case.plan.deployments + ), + ) + ) + + def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: payload = attempt.model_dump(mode="json") payload.update(updates) diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index aad4367ee..dfa23f957 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -10,7 +10,9 @@ from click.testing import CliRunner import data_designer.slurm.cli as cli_module -from data_designer.slurm.config import DataDesignerSlurmConfig +import data_designer.slurm.cli_benchmark as benchmark_cli_module +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig from data_designer.slurm.services import ( SlurmCollectionExecution, SlurmRetryExecution, @@ -83,6 +85,22 @@ def collect( ) +class _BenchmarkService: + def __init__(self, manifest: BenchmarkManifest, report: BenchmarkReport) -> None: + self.manifest = manifest + self.report = report + self.run_calls = [] + self.analysis_calls = [] + + def run(self, config, *, source_root, force): + self.run_calls.append((config, source_root, force)) + return self.manifest + + def analyze(self, benchmark_id, *, refresh_state, fail_if_incomplete): + self.analysis_calls.append((benchmark_id, refresh_state, fail_if_incomplete)) + return self.report + + def test_execute_emits_deterministic_json_and_forwards_actions( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -107,6 +125,35 @@ def test_execute_emits_deterministic_json_and_forwards_actions( assert service.calls == [(authored_run_single, tmp_path, True, True)] +def test_benchmark_cli_forwards_run_and_analysis_actions( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, + monkeypatch, +) -> None: + benchmark_file = tmp_path / "benchmark.json" + benchmark_file.write_text(benchmark_config.serialize_json()) + service = _BenchmarkService(benchmark_manifest, benchmark_report) + monkeypatch.setattr(benchmark_cli_module, "create_slurm_benchmark_service", lambda **_: service) + + run_result = CliRunner().invoke( + cli_module.create_cli(), + ["benchmark", "run", str(benchmark_file), "--force"], + ) + analyze_result = CliRunner().invoke( + cli_module.create_cli(), + ["benchmark", "analyze", "/workspace/benchmarks/benchmark-001", "--refresh", "--fail-if-incomplete"], + ) + + assert run_result.exit_code == 0 + assert json.loads(run_result.stdout)["benchmark_id"] == benchmark_manifest.benchmark_id + assert analyze_result.exit_code == 0 + assert json.loads(analyze_result.stdout)["analysis_id"] == benchmark_report.analysis_id + assert service.run_calls == [(benchmark_config, tmp_path, True)] + assert service.analysis_calls == [("benchmark-001", True, True)] + + def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: service = _RunService() monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) @@ -383,11 +430,11 @@ def test_image_add_rejects_credential_bearing_oci_source() -> None: assert "secret" not in result.stderr -def test_cli_exposes_m3c_run_commands() -> None: +def test_cli_exposes_m3_public_commands() -> None: result = CliRunner().invoke(cli_module.create_cli(), ["--help"]) assert result.exit_code == 0 assert all( - command in result.stdout for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile") + command in result.stdout + for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile", "benchmark") ) - assert "benchmark" not in result.stdout diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index b5f54d241..82a7cc4b5 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -135,10 +135,13 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non statement += f""" slurm_help_result = CliRunner().invoke(app, ["slurm", "--help"]) assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) -assert all(command in slurm_help_result.output for command in ("execute", "status", "cancel", "retry", "merge", "image")) -assert "benchmark" not in slurm_help_result.output +assert all( + command in slurm_help_result.output + for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile", "benchmark") +) assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} +from data_designer.slurm.benchmark import BenchmarkCompiler from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference from data_designer.slurm.contracts import RecordRange as ContractRecordRange from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace @@ -160,6 +163,9 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non SlurmRetryCoordinator, SlurmStateReconciler, ) +from data_designer.slurm.services import create_slurm_benchmark_service +assert BenchmarkCompiler.__name__ == "BenchmarkCompiler" +assert callable(create_slurm_benchmark_service) assert CollectionResult.__name__ == "CollectionResult" assert RetryPlan.__name__ == "RetryPlan" assert RunManifest.__name__ == "RunManifest" @@ -252,7 +258,7 @@ def main() -> None: assert str(leaf_base_requirement.specifier) == f"=={version}" assert leaf_packaging_requirement.specifier == Requirement("packaging>=25,<27").specifier assert leaf_pip_requirement.specifier == Requirement("pip>=25,<27").specifier - assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier + assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.12,<3").specifier assert leaf_pyyaml_requirement.specifier == Requirement("pyyaml>=6.0.1,<7").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) diff --git a/uv.lock b/uv.lock index 822556cd3..197492074 100644 --- a/uv.lock +++ b/uv.lock @@ -967,7 +967,7 @@ requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, { name = "packaging", specifier = ">=25,<27" }, { name = "pip", specifier = ">=25,<27" }, - { name = "pydantic", specifier = ">=2.9.2,<3" }, + { name = "pydantic", specifier = ">=2.12,<3" }, { name = "pyyaml", specifier = ">=6.0.1,<7" }, ] @@ -1198,7 +1198,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1559,17 +1559,17 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version < '3.11'" }, - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "jedi", marker = "python_full_version < '3.11'" }, - { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, - { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, - { name = "pygments", marker = "python_full_version < '3.11'" }, - { name = "stack-data", marker = "python_full_version < '3.11'" }, - { name = "traitlets", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e5/61/1810830e8b93c72dcd3c0f150c80a00c3deb229562d9423807ec92c3a539/ipython-8.38.0.tar.gz", hash = "sha256:9cfea8c903ce0867cc2f23199ed8545eb741f3a69420bfcf3743ad1cec856d39", size = 5513996, upload-time = "2026-01-05T10:59:06.901Z" } wheels = [ @@ -1584,17 +1584,17 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version == '3.11.*'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version == '3.11.*'" }, - { name = "jedi", marker = "python_full_version == '3.11.*'" }, - { name = "matplotlib-inline", marker = "python_full_version == '3.11.*'" }, - { name = "pexpect", marker = "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version == '3.11.*'" }, - { name = "pygments", marker = "python_full_version == '3.11.*'" }, - { name = "stack-data", marker = "python_full_version == '3.11.*'" }, - { name = "traitlets", marker = "python_full_version == '3.11.*'" }, - { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a6/60/2111715ea11f39b1535bed6024b7dec7918b71e5e5d30855a5b503056b50/ipython-9.10.0.tar.gz", hash = "sha256:cd9e656be97618a0676d058134cd44e6dc7012c0e5cb36a9ce96a8c904adaf77", size = 4426526, upload-time = "2026-02-02T10:00:33.594Z" } wheels = [ @@ -1610,16 +1610,16 @@ resolution-markers = [ "python_full_version >= '3.12' and python_full_version < '3.14'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version >= '3.12'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.12'" }, - { name = "jedi", marker = "python_full_version >= '3.12'" }, - { name = "matplotlib-inline", marker = "python_full_version >= '3.12'" }, - { name = "pexpect", marker = "python_full_version >= '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "stack-data", marker = "python_full_version >= '3.12'" }, - { name = "traitlets", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, ] sdist = { url = "https://files.pythonhosted.org/packages/86/28/a4698eda5a8928a45d6b693578b135b753e14fa1c2b36ee9441e69a45576/ipython-9.11.0.tar.gz", hash = "sha256:2a94bc4406b22ecc7e4cb95b98450f3ea493a76bec8896cda11b78d7752a6667", size = 4427354, upload-time = "2026-03-05T08:57:30.549Z" } wheels = [ @@ -1631,7 +1631,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments", marker = "python_full_version >= '3.11'" }, + { name = "pygments" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -4288,7 +4288,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -4349,7 +4349,7 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [