diff --git a/.agents/recipes/dependencies/recipe.md b/.agents/recipes/dependencies/recipe.md index 9e0aebf10..f96bae48e 100644 --- a/.agents/recipes/dependencies/recipe.md +++ b/.agents/recipes/dependencies/recipe.md @@ -11,7 +11,7 @@ permissions: # Dependency Audit -Audit the dependency graph across all three packages. Write findings to +Audit the dependency graph across all four packages. Write findings to `/tmp/audit-{{suite}}.md`. Dependabot handles version bump PRs. This recipe focuses on what Dependabot @@ -39,11 +39,12 @@ Read the `pyproject.toml` for each package: cat packages/data-designer-config/pyproject.toml cat packages/data-designer-engine/pyproject.toml cat packages/data-designer/pyproject.toml +cat packages/data-designer-slurm/pyproject.toml ``` -Note: engine and interface packages use `uv-dynamic-versioning` to inject -dependencies. Check both static declarations and the dynamic versioning -config. +Note: engine, interface, and Slurm packages use `uv-dynamic-versioning` to +inject dependencies. Check both static declarations and the dynamic +versioning config, including selected optional dependencies. ### 2. Direct dependency declaration gaps @@ -79,7 +80,7 @@ intentionally deferred but still need to be declared as dependencies. Check that shared dependencies use consistent version constraints: ```bash -# Extract dependency specs from all three pyproject.toml files +# Extract dependency specs from all four pyproject.toml files grep -E "^\s+\"[a-zA-Z]" packages/*/pyproject.toml ``` diff --git a/.agents/recipes/structure/recipe.md b/.agents/recipes/structure/recipe.md index 26a435b30..694ee44fb 100644 --- a/.agents/recipes/structure/recipe.md +++ b/.agents/recipes/structure/recipe.md @@ -19,21 +19,24 @@ to `/tmp/audit-{{suite}}.md`. Before starting, read the authoritative sources for the structural rules: 1. **`AGENTS.md`** (repo root) - "The Layering Is Structural" section defines - the three packages, their ownership, and the dependency direction rule. + the four packages, their ownership, and the dependency direction rule. 2. **`architecture/overview.md`** - system architecture diagram, package layout, and the "no reverse imports" rule. The canonical rules: ``` -data-designer (interface) -> data-designer-engine -> data-designer-config +data-designer-slurm -> data-designer (interface) -> data-designer-engine -> data-designer-config ``` - `packages/data-designer-config/` must NOT import from `data_designer.engine`, - `data_designer.interface`, or `data_designer.cli` + `data_designer.interface`, `data_designer.cli`, or `data_designer.slurm` - `packages/data-designer-engine/` must NOT import from - `data_designer.interface` or `data_designer.cli` -- `packages/data-designer/` CAN import from both engine and config + `data_designer.interface`, `data_designer.cli`, or `data_designer.slurm` +- `packages/data-designer/` CAN import from engine and config but must NOT + import from `data_designer.slurm` +- `packages/data-designer-slurm/` uses public interface/config APIs and must + NOT import engine internals **What CI already enforces**: ruff rule `TID` catches relative imports. The CI test matrix runs config, engine, and interface tests in isolation (separate diff --git a/.agents/recipes/test-health/recipe.md b/.agents/recipes/test-health/recipe.md index b9253033a..cac26fddb 100644 --- a/.agents/recipes/test-health/recipe.md +++ b/.agents/recipes/test-health/recipe.md @@ -306,7 +306,7 @@ Write the report to `/tmp/audit-{{suite}}.md`: | Check | Status | Detail | |-------|--------|--------| -| Package imports | OK/FAIL | All three packages import cleanly | +| Package imports | OK/FAIL | All four packages import cleanly | | Import timing | OK/FAIL | X.XXs (budget: 3s) | | Registry completeness | OK/WARN | Column types resolve to config classes | diff --git a/.github/workflows/agentic-ci-daily.yml b/.github/workflows/agentic-ci-daily.yml index 5e7637f2a..1348c905f 100644 --- a/.github/workflows/agentic-ci-daily.yml +++ b/.github/workflows/agentic-ci-daily.yml @@ -144,6 +144,7 @@ jobs: if: matrix.suite == 'dependencies' run: | .venv/bin/python scripts/audit_package_dependencies.py \ + --extra data-designer:slurm \ --output /tmp/dependency-inventory.json jq '.packages[] | {package, missing, unresolved_modules}' \ /tmp/dependency-inventory.json diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7dd5b584d..350b25a8b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -129,6 +129,42 @@ jobs: uv run --with pytest --with pytest-asyncio --with pytest-httpx --with pytest-env \ pytest packages/data-designer/tests + test-slurm-package: + name: ${{ matrix.name }} + needs: validate-dispatch + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - python-version: "3.10" + test-path: packages/data-designer-slurm/tests/runtime/test_proxy.py + name: Test Slurm proxy (Python 3.10) + - python-version: "3.11" + test-path: packages/data-designer-slurm/tests + name: Test Slurm package wheels + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Install uv + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "latest" + python-version: ${{ matrix.python-version }} + enable-cache: true + + - name: Install development dependencies + run: make install-dev + + - name: Run package tests + run: .venv/bin/pytest ${{ matrix.test-path }} + + - name: Run built-wheel installation tests + if: matrix.python-version == '3.11' + run: make test-slurm-wheel-install + # =========================================================================== # Combined Coverage Check # Runs all tests together to verify overall coverage threshold @@ -284,7 +320,7 @@ jobs: test-summary: name: Test (Python ${{ matrix.python-version }} on ${{ matrix.os }}) runs-on: ubuntu-latest - needs: [validate-dispatch, test-config, test-engine, test-interface] + needs: [validate-dispatch, test-config, test-engine, test-interface, test-slurm-package] if: always() strategy: matrix: @@ -297,12 +333,14 @@ jobs: if [[ "${{ needs.validate-dispatch.result }}" != "success" ]] || \ [[ "${{ needs.test-config.result }}" != "success" ]] || \ [[ "${{ needs.test-engine.result }}" != "success" ]] || \ - [[ "${{ needs.test-interface.result }}" != "success" ]]; then + [[ "${{ needs.test-interface.result }}" != "success" ]] || \ + [[ "${{ needs.test-slurm-package.result }}" != "success" ]]; then echo "One or more test jobs failed" echo "validate-dispatch: ${{ needs.validate-dispatch.result }}" echo "test-config: ${{ needs.test-config.result }}" echo "test-engine: ${{ needs.test-engine.result }}" echo "test-interface: ${{ needs.test-interface.result }}" + echo "test-slurm-package: ${{ needs.test-slurm-package.result }}" exit 1 fi echo "All test jobs passed successfully" diff --git a/AGENTS.md b/AGENTS.md index 434a68e5f..fc02c55de 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,15 +7,16 @@ If you are an agent helping a user **build a dataset**, use the [`data-designer` ## The Layering Is Structural -The `data_designer` namespace is split across three installable packages that merge at runtime via PEP 420 implicit namespace packages (no top-level `__init__.py`). +The `data_designer` namespace is split across four installable packages that merge at runtime via PEP 420 implicit namespace packages (no top-level `__init__.py`). | Package | Path | Owns | |---------|------|------| | `data-designer-config` | `packages/data-designer-config/` | `data_designer.config` — column configs, model configs, sampler params, builder API, plugin system, lazy imports | | `data-designer-engine` | `packages/data-designer-engine/` | `data_designer.engine` — column generators, dataset builders, DAG execution, model facade, validators, sampling | | `data-designer` | `packages/data-designer/` | `data_designer.interface` — public `DataDesigner` class, results, errors; `data_designer.cli` — CLI entry point; `data_designer.integrations` | +| `data-designer-slurm` | `packages/data-designer-slurm/` | `data_designer.slurm` — optional Slurm execution | -**Dependency direction (left depends on right):** interface → engine → config. Never import against this flow. +**Import direction (left imports right):** Slurm → interface → engine → config. The `data-designer[slurm]` extra creates a packaging-only reverse edge; no code may import against this flow. ## Core Concepts @@ -34,7 +35,7 @@ The `data_designer` namespace is split across three installable packages that me ## Structural Invariants -- **Import direction** — interface → engine → config (left depends on right). No reverse imports. +- **Import direction** — Slurm → interface → engine → config (left imports right). No reverse imports. - **Fast imports** — heavy third-party libraries are lazy-loaded via `data_designer.lazy_heavy_imports`. See [STYLEGUIDE.md](STYLEGUIDE.md) for the pattern. - **No relative imports** — absolute imports only, enforced by ruff rule `TID`. - **Typed code** — all functions, methods, and class attributes require type annotations. Modern syntax: `list[str]`, `str | None`. diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 0b5396cc3..1e6869807 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -110,16 +110,17 @@ uv run ruff format --check # Check formatting ### Running Tests -`make test` runs all three package test suites in sequence (config, engine, interface). When iterating on a single package, run its tests directly: +`make test` runs all four package test suites in sequence (config, engine, interface, Slurm). When iterating on a single package, run its tests directly: ```bash -# Run all tests (config + engine + interface) +# Run all tests (config + engine + interface + Slurm) make test # Run a single package's tests make test-config # data-designer-config make test-engine # data-designer-engine make test-interface # data-designer (interface) +make test-slurm # data-designer-slurm # Run a specific test file uv run pytest tests/config/test_sampler_constraints.py diff --git a/Makefile b/Makefile index ff3e6e8b4..1540c8955 100644 --- a/Makefile +++ b/Makefile @@ -12,17 +12,20 @@ GRAPHIFY_REQUIREMENTS ?= $(REPO_PATH)/.agents/requirements-graphify.txt # Package directories CONFIG_PKG := packages/data-designer-config ENGINE_PKG := packages/data-designer-engine +SLURM_PKG := packages/data-designer-slurm INTERFACE_PKG := packages/data-designer # Package source and test paths CONFIG_PATHS := $(CONFIG_PKG)/src $(CONFIG_PKG)/tests ENGINE_PATHS := $(ENGINE_PKG)/src $(ENGINE_PKG)/tests +SLURM_PATHS := $(SLURM_PKG)/src $(SLURM_PKG)/tests INTERFACE_PATHS := $(INTERFACE_PKG)/src $(INTERFACE_PKG)/tests $(INTERFACE_PKG)/dev-tools ALL_PKG_PATHS := packages/ scripts/ tests_e2e/ # Test directories CONFIG_TESTS := $(CONFIG_PKG)/tests ENGINE_TESTS := $(ENGINE_PKG)/tests +SLURM_TESTS := $(SLURM_PKG)/tests INTERFACE_TESTS := $(INTERFACE_PKG)/tests define install-pre-commit-hooks @@ -44,7 +47,7 @@ help: @echo "═════════════════════════════════════════════════════════════" @echo "" @echo "📦 Installation (uv workspace - all packages in editable mode):" - @echo " install - Install all packages (config → engine → interface)" + @echo " install - Install all packages (config → engine → interface → slurm)" @echo " install-dev - Install all packages + dev tools (pytest, etc.)" @echo " install-graphify - Install isolated Graphify CI dependencies" @echo " install-dev-notebooks - Install all packages + dev + docs + notebook tools" @@ -55,13 +58,14 @@ help: @echo " test-graphify - Run Graphify CLI compatibility tests" @echo " coverage - Run tests with coverage report" @echo " test-e2e - Run e2e plugin tests" + @echo " test-slurm-wheel-install - Test optional Slurm package wheel installation" @echo " health-checks - Run provider health checks" @echo " test-run-tutorials - Run tutorial notebooks as e2e tests" @echo " test-run-recipes - Run recipe scripts as e2e tests" @echo " test-run-all-examples - Run all tutorials and recipes as e2e tests" @echo "" @echo "🔬 Isolated Testing (mirrors CI - uses temp venv):" - @echo " test-isolated - Run all isolated tests (config → engine → interface)" + @echo " test-isolated - Run core package isolated tests (config → engine → interface)" @echo " test-config-isolated - Test config with ONLY config installed" @echo " test-engine-isolated - Test engine with ONLY engine+config installed" @echo " test-interface-isolated - Test interface with full package installed" @@ -115,7 +119,7 @@ help: @echo " publish VERSION=X.Y.Z ALLOW_BRANCH=1 - Publish from non-main branch" @echo " publish VERSION=X.Y.Z FORCE_TAG=1 - Overwrite existing git tag" @echo "" - @echo "📦 Per-Package Commands (use suffix: -config, -engine, -interface):" + @echo "📦 Per-Package Commands (use suffix: -config, -engine, -slurm, -interface):" @echo " test- - Run tests for a specific package" @echo " lint- - Lint a specific package" @echo " lint-fix- - Fix lint issues in a specific package" @@ -135,7 +139,7 @@ help: install: @echo "📦 Installing DataDesigner workspace (all packages in editable mode)..." - @echo " Packages: data-designer-config → data-designer-engine → data-designer" + @echo " Packages: data-designer-config → data-designer-engine → data-designer → data-designer-slurm" uv sync --all-packages @echo "✅ Installation complete!" @echo "" @@ -143,7 +147,7 @@ install: install-dev: @echo "📦 Installing DataDesigner workspace in development mode..." - @echo " Packages: data-designer-config → data-designer-engine → data-designer" + @echo " Packages: data-designer-config → data-designer-engine → data-designer → data-designer-slurm" @echo " Groups: dev (pytest, coverage, etc.)" uv sync --all-packages --group dev $(call install-pre-commit-hooks) @@ -154,11 +158,12 @@ install-dev: @echo " packages/data-designer-config/ - Configuration layer (lightweight)" @echo " packages/data-designer-engine/ - Generation engine (heavy deps)" @echo " packages/data-designer/ - Full package with CLI" + @echo " packages/data-designer-slurm/ - Optional Slurm execution" @echo "" @echo "💡 Next steps:" @echo " make verify-imports - Verify all packages are working" @echo " make test - Run all tests across packages" - @echo " make test- - Run tests for specific package (config, engine, interface)" + @echo " make test- - Run tests for specific package (config, engine, slurm, interface)" @echo " make lint - Lint all code" @echo " make build - Build all package wheels" @@ -175,7 +180,7 @@ test-graphify: install-graphify install-dev-notebooks: @echo "📦 Installing DataDesigner workspace with notebook dependencies..." - @echo " Packages: data-designer-config → data-designer-engine → data-designer" + @echo " Packages: data-designer-config → data-designer-engine → data-designer → data-designer-slurm" @echo " Groups: dev + docs + notebooks (Jupyter, jupytext, etc.)" uv sync --all-packages --group dev --group docs --group notebooks @$(MAKE) install-notebook-css @@ -186,7 +191,7 @@ install-dev-notebooks: install-dev-recipes: @echo "📦 Installing DataDesigner workspace with recipe dependencies..." - @echo " Packages: data-designer-config → data-designer-engine → data-designer" + @echo " Packages: data-designer-config → data-designer-engine → data-designer → data-designer-slurm" @echo " Groups: dev + recipes (bm25s, pymupdf, etc.)" uv sync --all-packages --group dev --group recipes $(call install-pre-commit-hooks) @@ -198,7 +203,7 @@ install-dev-recipes: # TESTING # ============================================================================== -test: test-config test-engine test-interface +test: test-config test-engine test-interface test-slurm @echo "✅ All package tests complete!" test-config: @@ -213,6 +218,14 @@ test-interface: @echo "🧪 Testing data-designer (interface)..." uv run --group dev pytest $(INTERFACE_TESTS) +test-slurm: + @echo "🧪 Testing data-designer-slurm..." + .venv/bin/pytest $(SLURM_TESTS) + +test-slurm-wheel-install: + @echo "🧪 Testing data-designer-slurm wheel installation..." + .venv/bin/python scripts/test_slurm_package_install.py + # ------------------------------------------------------------------------------ # Isolated Testing (mirrors CI behavior) # Each package is installed independently to verify dependency boundaries @@ -343,12 +356,12 @@ test-run-all-examples: test-run-tutorials test-run-recipes # CODE QUALITY - FORMATTING # ============================================================================== -format: format-config format-engine format-interface +format: format-config format-engine format-interface format-slurm @echo "📐 Formatting scripts and tests_e2e..." uv run ruff format scripts/ tests_e2e/ @echo "✅ Formatting complete!" -format-check: format-check-config format-check-engine format-check-interface +format-check: format-check-config format-check-engine format-check-interface format-check-slurm @echo "📐 Checking scripts and tests_e2e formatting..." uv run ruff format --check scripts/ tests_e2e/ @echo "✅ Formatting check complete! Run 'make format' to auto-fix issues." @@ -365,6 +378,10 @@ format-interface: @echo "📐 Formatting data-designer (interface)..." uv run ruff format $(INTERFACE_PATHS) --exclude '**/_version.py' +format-slurm: + @echo "📐 Formatting data-designer-slurm..." + .venv/bin/ruff format $(SLURM_PATHS) + format-check-config: @echo "📐 Checking data-designer-config formatting..." uv run ruff format --check $(CONFIG_PATHS) --exclude '**/_version.py' @@ -377,16 +394,20 @@ format-check-interface: @echo "📐 Checking data-designer (interface) formatting..." uv run ruff format --check $(INTERFACE_PATHS) --exclude '**/_version.py' +format-check-slurm: + @echo "📐 Checking data-designer-slurm formatting..." + .venv/bin/ruff format --check $(SLURM_PATHS) + # ============================================================================== # CODE QUALITY - LINTING # ============================================================================== -lint: lint-config lint-engine lint-interface +lint: lint-config lint-engine lint-interface lint-slurm @echo "🔍 Linting scripts and tests_e2e..." uv run ruff check --output-format=full scripts/ tests_e2e/ @echo "✅ Linting complete! Run 'make lint-fix' to auto-fix issues." -lint-fix: lint-fix-config lint-fix-engine lint-fix-interface +lint-fix: lint-fix-config lint-fix-engine lint-fix-interface lint-fix-slurm @echo "🔍 Fixing lint issues in scripts and tests_e2e..." uv run ruff check --fix scripts/ tests_e2e/ @echo "✅ Linting with autofix complete!" @@ -403,6 +424,10 @@ lint-interface: @echo "🔍 Linting data-designer (interface)..." uv run ruff check --output-format=full $(INTERFACE_PATHS) --exclude '**/_version.py' +lint-slurm: + @echo "🔍 Linting data-designer-slurm..." + .venv/bin/ruff check --output-format=full $(SLURM_PATHS) + lint-fix-config: @echo "🔍 Fixing lint issues in data-designer-config..." uv run ruff check --fix $(CONFIG_PATHS) --exclude '**/_version.py' @@ -415,6 +440,10 @@ lint-fix-interface: @echo "🔍 Fixing lint issues in data-designer (interface)..." uv run ruff check --fix $(INTERFACE_PATHS) --exclude '**/_version.py' +lint-fix-slurm: + @echo "🔍 Fixing lint issues in data-designer-slurm..." + .venv/bin/ruff check --fix $(SLURM_PATHS) + # ============================================================================== # CODE QUALITY - COMBINED CHECKS # ============================================================================== @@ -434,11 +463,14 @@ check-engine: format-check-engine lint-engine check-interface: format-check-interface lint-interface @echo "✅ Checks complete for data-designer (interface)!" +check-slurm: format-check-slurm lint-slurm + @echo "✅ Checks complete for data-designer-slurm!" + # ============================================================================== # BUILD # ============================================================================== -build: build-config build-engine build-interface +build: build-config build-engine build-interface build-slurm @echo "✅ All packages built!" build-config: @@ -453,6 +485,10 @@ build-interface: @echo "🏗️ Building data-designer (interface)..." cd $(INTERFACE_PKG) && uv build -o dist +build-slurm: + @echo "🏗️ Building data-designer-slurm..." + cd $(SLURM_PKG) && uv build -o dist + # ============================================================================== # UTILITIES # ============================================================================== @@ -462,6 +498,7 @@ verify-imports: uv run python -c "from data_designer.config.config_builder import DataDesignerConfigBuilder; print(' ✓ config')" uv run python -c "from data_designer.engine.compiler import compile_data_designer_config; print(' ✓ engine')" uv run python -c "from data_designer.interface.data_designer import DataDesigner; print(' ✓ interface')" + .venv/bin/python -c "import data_designer.slurm; print(' ✓ slurm')" @echo "✅ All imports verified!" show-versions: @@ -469,6 +506,7 @@ show-versions: @uv run python -c "from data_designer.config._version import __version__; print(f' data-designer-config: {__version__}')" 2>/dev/null || echo " data-designer-config: (not installed)" @uv run python -c "from data_designer.engine._version import __version__; print(f' data-designer-engine: {__version__}')" 2>/dev/null || echo " data-designer-engine: (not installed)" @uv run python -c "from data_designer.interface._version import __version__; print(f' data-designer: {__version__}')" 2>/dev/null || echo " data-designer: (not installed)" + @.venv/bin/python -c 'from importlib.metadata import version; print(" data-designer-slurm: " + version("data-designer-slurm"))' 2>/dev/null || echo " data-designer-slurm: (not installed)" # ============================================================================== # LICENSE CHECKS @@ -777,6 +815,7 @@ clean-dist: rm -rf $(CONFIG_PKG)/dist rm -rf $(ENGINE_PKG)/dist rm -rf $(INTERFACE_PKG)/dist + rm -rf $(SLURM_PKG)/dist rm -f packages/*/src/data_designer/*/_version.py @echo "✅ Dist directories cleaned!" @@ -795,20 +834,20 @@ clean-test-coverage: # ============================================================================== .PHONY: bench-cli-startup bench-cli-startup-verbose \ - build build-config build-engine build-interface \ - check-all check-all-fix check-config check-engine check-interface \ + build build-config build-engine build-interface build-slurm \ + check-all check-all-fix check-config check-engine check-interface check-slurm \ check-dependency-licenses check-fern-docs check-fern-docs-locally check-fern-links check-fern-published-docs check-fern-release-version check-fern-theme-access check-license-headers \ clean clean-dist clean-notebooks clean-pycache clean-test-coverage \ convert-execute-notebooks \ coverage coverage-config coverage-engine coverage-interface \ - format format-check format-check-config format-check-engine format-check-interface \ - format-config format-engine format-interface \ + format format-check format-check-config format-check-engine format-check-interface format-check-slurm \ + format-config format-engine format-interface format-slurm \ generate-colab-notebooks generate-fern-notebooks generate-fern-notebooks-with-outputs help \ install install-dev install-dev-notebooks install-dev-recipes install-docs-deps install-graphify install-notebook-css \ - lint lint-config lint-engine lint-fix lint-fix-config lint-fix-engine lint-fix-interface lint-interface lock-graphify \ + lint lint-config lint-engine lint-fix lint-fix-config lint-fix-engine lint-fix-interface lint-fix-slurm lint-interface lint-slurm lock-graphify \ perf-import perf-import-runtime prepare-fern-docs prepare-fern-release publish serve-fern-docs-dev serve-fern-docs-local-theme serve-fern-docs-locally show-versions \ health-checks \ test test-config test-config-isolated test-e2e test-engine test-engine-isolated \ - test-graphify test-interface test-interface-isolated test-isolated \ + test-graphify test-interface test-interface-isolated test-isolated test-slurm test-slurm-wheel-install \ test-run-all-examples test-run-recipes test-run-tutorials \ update-license-headers verify-imports diff --git a/architecture/cli.md b/architecture/cli.md index 1c9e9cf4a..ecdcece1e 100644 --- a/architecture/cli.md +++ b/architecture/cli.md @@ -20,6 +20,33 @@ The CLI is built on Typer with lazy command loading to keep startup fast. Config `create_lazy_typer_group` and `_LazyCommand` stubs defer importing command modules until a command is actually invoked. This keeps `data-designer --help` fast — only the command names and descriptions are loaded eagerly; the full module (and its dependencies) loads on first use. +The root group also discovers optional command groups from the `data_designer.cli` +entry-point group. The entry-point name is the top-level command name. Its target +must be a zero-argument callable that returns a `click.Command`, normally a Click +group produced from a Typer application. Root help uses the providing +distribution's `Summary` metadata and does not load the target. + +```toml +[project.entry-points."data_designer.cli"] +slurm = "data_designer.slurm.cli:create_cli" +``` + +Extension distributions must declare a compatible dependency on +`data-designer`. Compatibility and the loaded object are validated only when the +extension command is selected. The entry-point name is authoritative and is +assigned to the returned command. Built-in command names are reserved. A +built-in collision makes that command unavailable with a targeted error rather +than allowing the extension to replace it. A duplicate extension name remains +visible as unavailable in root help and fails deterministically without loading +any conflicting target. Invalid extension names and discovery failures are +reported as warnings without disabling built-in commands. + +The installed-wheel smoke test measures nine warm root-help subprocesses for a +base-only environment and an environment with the Slurm extension. The accepted +base median is at most one second, and the median overhead from installing the +extension is at most 100 ms. It also verifies that root help does not import the +extension or compatibility-checking modules. + ### Layering Pattern (Setup Workflows) Config management commands (models, providers, MCP providers, tools) follow a consistent four-layer pattern: diff --git a/architecture/overview.md b/architecture/overview.md index 35c53e591..946e5b162 100644 --- a/architecture/overview.md +++ b/architecture/overview.md @@ -1,11 +1,13 @@ # System Architecture -DataDesigner is split across three installable packages that merge at runtime into a single `data_designer` namespace via PEP 420 implicit namespace packages (no top-level `__init__.py`). +DataDesigner is split across four installable packages that merge at runtime into a single `data_designer` namespace via PEP 420 implicit namespace packages (no top-level `__init__.py`). ## Overview ``` ┌─────────────────────────────────────────────────────────┐ +│ data-designer-slurm (optional Slurm execution) │ +├─────────────────────────────────────────────────────────┤ │ data-designer (interface + CLI + integrations) │ │ DataDesigner class, CLI commands, HuggingFace Hub │ ├─────────────────────────────────────────────────────────┤ @@ -19,7 +21,7 @@ DataDesigner is split across three installable packages that merge at runtime in └─────────────────────────────────────────────────────────┘ ``` -**Dependency direction:** interface → engine → config. No reverse imports. +**Import direction:** Slurm → interface → engine → config. The `data-designer[slurm]` extra creates a packaging-only reverse edge from the interface distribution to the optional leaf; no base package imports `data_designer.slurm`. Users declare what their data should look like through config objects (columns, types, relationships, validation rules). The engine compiles those configs into an execution plan and generates the dataset. The interface package provides the public `DataDesigner` class and CLI that wire everything together. @@ -50,7 +52,7 @@ Users declare what their data should look like through config objects (columns, ## Design Decisions -- **PEP 420 namespace packages** allow the three packages to be installed independently while sharing the `data_designer` namespace. This enables lighter installs (e.g., config-only for validation tooling) without import conflicts. +- **PEP 420 namespace packages** allow the four packages to be installed independently while sharing the `data_designer` namespace. This enables lighter installs (e.g., config-only for validation tooling) and an optional Slurm leaf without import conflicts. - **Lazy imports throughout** — `__getattr__`-based lazy loading in `data_designer.config` and `data_designer.interface`, plus `lazy_heavy_imports` for numpy/pandas, keep startup fast. - **Async-only execution** gives `DatasetBuilder` one scheduling path with row-group parallelism and DAG-aware dispatch behind the public interface. - **`TaskRegistry` subclasses: one instance per class** — `TaskRegistry.__new__` (`registry/base.py`) ensures a single instance of each concrete registry (column generators, profilers, processors). **`ModelRegistry`** and **`MCPRegistry`** are ordinary classes, constructed per run with injected dependencies. **`PluginRegistry`** (`plugins/registry.py`) uses `__new__` so entry points are discovered once per process. diff --git a/dependency-license-policy.toml b/dependency-license-policy.toml index 4b8642bb8..a858b2b70 100644 --- a/dependency-license-policy.toml +++ b/dependency-license-policy.toml @@ -16,6 +16,13 @@ allowed_licenses = [ # Existing dependencies outside the global whitelist are reviewed by package, # version, and exact reported license. A version or license metadata change must # be reviewed explicitly. +[exceptions.aiohappyeyeballs] +reason = "Unmodified async connection helper installed as a separate transitive dependency under the PSF license." +reports = [ + { version = "2.6.1", license = "PSF-2.0" }, + { version = "2.6.1", license = "Python Software Foundation License" }, +] + [exceptions.certifi] reason = "Unmodified TLS certificate bundle installed as a separate transitive dependency." reports = [{ version = "2026.10.19", license = "Mozilla Public License 2.0 (MPL 2.0)" }] diff --git a/fern/versions/latest.yml b/fern/versions/latest.yml index 0d44b79b0..b903bad0c 100644 --- a/fern/versions/latest.yml +++ b/fern/versions/latest.yml @@ -79,6 +79,30 @@ 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: Overview + path: ./latest/pages/slurm/overview.mdx + - page: Getting Started + path: ./latest/pages/slurm/getting-started.mdx + - page: Profiles and Clusters + path: ./latest/pages/slurm/profiles.mdx + - page: Images + path: ./latest/pages/slurm/images.mdx + - page: Run Configuration + path: ./latest/pages/slurm/run-configuration.mdx + - page: Run Operations + path: ./latest/pages/slurm/operations.mdx + - page: Retry and Collection + path: ./latest/pages/slurm/retry-and-collection.mdx + - page: Benchmarks + path: ./latest/pages/slurm/benchmarks.mdx + - page: Dependencies and Tools + path: ./latest/pages/slurm/dependencies-and-tools.mdx + - page: Troubleshooting + path: ./latest/pages/slurm/troubleshooting.mdx + - page: CLI Reference + path: ./latest/pages/slurm/cli-reference.mdx - section: Recipes contents: - page: Recipe Cards diff --git a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx index 5c4c20209..5e2d23023 100644 --- a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx +++ b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx @@ -222,6 +222,9 @@ results = designer.create( dataset_name="training_data", resume=ResumeMode.IF_POSSIBLE, ) + +print(results.dataset_path) +print(results.requested_resume_mode, results.effective_resume_mode) ``` Resume modes: @@ -241,7 +244,7 @@ Resume has a few important invariants: - Once `process_after_generation()` has run, the dataset is considered terminal for resume. Re-running with the same target returns the existing dataset; extending requires a fresh run. - If a run crashed after every row group was written but before `process_after_generation()` could start, resume runs after-generation on the existing on-disk dataset (the parquet files are still clean) and marks it terminal afterwards. A crash _during_ `process_after_generation()` still raises — the parquet files may have been partially rewritten and starting fresh is the only safe option. -The `DatasetCreationResults` returned by a resume invocation reflects the full dataset on disk for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`). Per-run observability — `task_traces`, model-usage logs, and telemetry events emitted during the call — is scoped to the resume invocation only; the original run's in-memory traces are not persisted across process boundaries. +The `DatasetCreationResults` returned by a resume invocation reflects the full dataset on disk for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`). `dataset_path` is the resolved output directory, which may differ from the requested name when Data Designer starts fresh. Compare `requested_resume_mode` with `effective_resume_mode` to determine whether the invocation resumed or started fresh. Per-run observability — `task_traces`, model-usage logs, and telemetry events emitted during the call — is scoped to the resume invocation only; the original run's in-memory traces are not persisted across process boundaries. Only resume datasets from trusted artifact directories. Resume reads local `metadata.json`, `builder_config.json`, and parquet files to determine checkpoint state. @@ -385,18 +388,23 @@ config_builder.add_model_config( ### Run outcomes -A run can finish with fewer records than requested when non-retryable errors drop rows. Inspect `len(result.load_dataset())` to detect. +A run can finish with fewer records than requested when non-retryable errors drop rows. Compare `actual_num_records` with `requested_num_records`, or inspect `is_partial`, without loading the dataset. -If the rate of non-retryable errors crosses `RunConfig.shutdown_error_rate`, generation stops early and raises `DataDesignerEarlyShutdownError` (a subclass of `DataDesignerGenerationError`). Catch it separately when a typed retry path is appropriate: +If the rate of non-retryable errors crosses `RunConfig.shutdown_error_rate`, generation stops early. A run with no surviving records raises `DataDesignerEarlyShutdownError` (a subclass of `DataDesignerGenerationError`). If some records survive, `create()` returns a partial result with `early_shutdown=True`. ```python -from data_designer.interface.errors import DataDesignerEarlyShutdownError +from data_designer.interface import DataDesignerEarlyShutdownError try: result = dd_instance.create(config_builder, num_records=1000) except DataDesignerEarlyShutdownError: # e.g. retry against a different model alias ... +else: + if result.early_shutdown: + print(f"Early shutdown preserved {result.actual_num_records} records") + if result.is_partial: + print(f"Generated {result.actual_num_records} of {result.requested_num_records} requested records") ``` ## Local OpenTelemetry Metrics diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx new file mode 100644 index 000000000..b5db16406 --- /dev/null +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -0,0 +1,87 @@ +# 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. + +Complete [Getting Started](/slurm/getting-started) first so the profile, images, and base run are ready. + +## Define cases + +Create `benchmark.yaml` next to an existing `run.yaml`: + +```yaml +schema_version: 1 +name: generator-scaling +base_run: run.yaml +model_aliases: all +concurrency_values: + - 32 + - 64 +deployment_cases: + - name: single-node-groups + deployments: + generator: + nodes: 2 + nodes_per_replica: 1 + - name: two-node-group + deployments: + generator: + nodes: 2 + nodes_per_replica: 2 +analysis: + target_total_records: 1000000 + target_runtime: 4h +``` + +`model_aliases` is required. Set it to `all` to apply each concurrency value to every model in the base run. For a two-model run with `generator` and `reviewer`, values `[32, 64]` produce concurrency pairs `(32, 32)` and `(64, 64)` at each deployment case. To vary only one model, set `model_aliases: [generator]`; the reviewer then keeps its base-run concurrency. Every selected model receives the same value, so this sweep does not produce mixed pairs such as `(32, 64)`. + +The compiler expands the authored order deterministically. For this example, the first cases are `single-node-groups-c32` and `two-node-group-c32`. By default, each case requests `min(5000, max(1000, ceil(concurrency * 1.0)))` records. Thus concurrency 512 requests 1,000 records and concurrency 1,024 requests 1,024. For a longer sample at those concurrency values, set `record_policy: {type: adaptive, base_records: 1000, max_records: 5000, records_per_concurrency: 4.0}` to request 2,048 and 4,096 records. Raise `max_records` when testing larger concurrency values. A fixed policy with `type: fixed` and `records` is also available. + +Each benchmark case runs as one Slurm array task, even when the base run uses multiple tasks. The record policy sets the sample size for that task. This measures the throughput of one allocation, which analysis uses to estimate how many such tasks are needed for `target_total_records` within `target_runtime`. + +`nodes_per_replica` controls the number of nodes in each pipeline-parallel group. The resolved replica count is `(nodes / nodes_per_replica) * (gpus_per_node / tensor_parallel)`. Benchmark deployment cases change `nodes` and `nodes_per_replica`; they keep each deployment's `tensor_parallel` from the base run. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count) for concrete eight-GPU examples. + +## Run + +```bash +data-designer slurm benchmark run benchmark.yaml +``` + +This uses `DATA_DESIGNER_SLURM_PROFILE_FILE` and the default cluster from Getting Started. 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 with the same installed revision. 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 the allocation, so queue wait does not affect topology comparisons. For new one-task cases, `rows_per_second` measures that task's generation rate. Reports for saved multi-task benchmarks retain their aggregate shard rate. `generation_seconds` is the target runtime remaining after startup, and GPU-hour estimates assume each projected task uses the full target runtime. 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() + +manifest = service.run(config, source_root=benchmark_file.parent) +report = service.analyze(manifest.benchmark_id, refresh_state=True) +``` + +For child failure recovery and output collection, use [Retry and collection](/slurm/retry-and-collection). diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx new file mode 100644 index 000000000..ba2583ea6 --- /dev/null +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -0,0 +1,47 @@ +# Slurm CLI reference + +Install the extension with `python -m pip install "data-designer[slurm]"`, verify it with `python -m pip show data-designer-slurm`, then use `data-designer slurm COMMAND --help` for the authoritative option list. + +All commands accept `--profile-file PATH` and `--cluster NAME` where cluster selection is required. + +| Command | Purpose | +| --- | --- | +| `profile init` | Create a restrictive starter profile catalog without overwriting. | +| `profile validate` | Validate catalog selection, shared workspace access, and scheduler facts. | +| `image add` | Resolve an official versioned vLLM tag, or import/inspect a pinned OCI or SQSH image through a CPU Slurm job, then publish a verified alias. | +| `image ls` | List aliases in the selected profile workspace. | +| `image info NAME` | Show one registered alias and immutable image facts. | +| `image rm NAME` | Remove an alias and retain its SQSH artifact. Use `--delete-artifact` to also remove an unshared, package-owned OCI artifact. | +| `execute RUN_FILE` | Resolve and submit a run. Add `--dry-run` to render without submission. | +| `status RUN_ID` | Reconcile scheduler observations with durable managed state. | +| `cancel RUN_ID` | Cancel active jobs owned by a managed run. | +| `retry RUN_OR_JOB_ID` | Preview or submit a sparse retry. Repeat `--task-id` to select shards. | +| `merge RUN_ID` | Submit a zero-GPU winner collection job for a managed run. | +| `benchmark run BENCHMARK_FILE` | Expand benchmark cases into ordinary child runs and submit them. | +| `benchmark analyze BENCHMARK` | Write a report from persisted child state. | + +## Common options + +| Option | Commands | Meaning | +| --- | --- | --- | +| `--output FORMAT` | `slurm` group, before a subcommand | Use `human` (default) or `json` for automation. | +| `--follow-logs` / `--no-follow-logs` | `image add` | Follow bounded image-job logs by default in an interactive terminal; override either way. | +| `--profile-file PATH` | Most commands | Use a catalog other than the default or environment-selected file. | +| `--cluster NAME` | Most commands | Select one catalog entry explicitly. | +| `--dry-run` | `execute`, `retry` | Perform resolution and render a batch script without submission. | +| `--input-path PATH` | `merge` | Select a managed run directory instead of passing `RUN_ID`. | +| `--output-path PATH` | `merge` | Select an authorized compute-visible collection directory instead of `/runs//merge-output`. | +| `--num-partitions N` | `merge` | Write N Parquet files; defaults to 1 for a new collection. | +| `--force` | `retry`, `benchmark run` | Skip retry confirmation or resume a safely recoverable benchmark submission. It does not overwrite immutable evidence. | +| `--refresh-state` | `benchmark analyze` | Reconcile each child with the scheduler before analysis. | +| `--fail-if-incomplete` | `benchmark analyze` | Persist the report, then return a conflict if any case is incomplete. | + +## Output and errors + +Success is readable text on stdout by default; failure is readable text on stderr. For automation, use `data-designer slurm --output json COMMAND` or set `DATA_DESIGNER_SLURM_OUTPUT=json`. JSON mode keeps the existing canonical success value on stdout and error object on stderr: + +```json +{"error":{"code":"invalid_request","message":"...","operation":"execute_run"}} +``` + +Exit codes are `1` for internal errors, `2` for invalid requests, `3` for missing managed resources, `4` for conflicts, and `5` for unavailable external systems. Error messages are bounded and sanitized; inspect durable logs for operational detail. diff --git a/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx b/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx new file mode 100644 index 000000000..157d9a196 --- /dev/null +++ b/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx @@ -0,0 +1,111 @@ +# Dependencies, plugins, MCP, and secrets + +The shared Python environment used to invoke the CLI provides Data Designer and its base client environment. A run may add immutable pure Python wheels as an allocation-local overlay. The overlay is reconstructed for each allocation and does not modify the shared environment. + +## Add package requirements + +```yaml +client: + dependencies: + requirements: + - example-plugin==1.2.3 +``` + +The submit-side resolver accepts package requirements or direct HTTPS wheel URLs with a SHA-256 fragment. It rejects editable installs, source directories, VCS references, markers, duplicate normalized names, and non-wheel URLs. The resolved lock and wheels are persisted with the run. + +The overlay cannot replace any distribution already installed in the shared client environment. Create a new shared environment when Data Designer or another base package needs to change, and use that environment to submit new runs. + +## Reuse a dependency lock + +Use exactly one of `requirements` or `lock_file`: + +```yaml +client: + dependencies: + requirements: null + lock_file: inputs/dependency-lock.json +``` + +The lock must target the recorded shared client runtime inventory and reference available wheel artifacts. + +`client.dependencies.index_credentials` reserves named secret references for the allocation runtime. Submit-side inline requirement resolution does not use them and always targets the public index. + +Inline requirements resolve from the public Python package index with ambient pip configuration disabled. For private or unpublished plugins, install the package into a new shared client environment, use an immutable HTTPS wheel URL that does not embed credentials, or supply a previously resolved compatible lock and wheel artifacts. + +## Reference secrets + +Secrets are environment references, never literal values in persisted YAML: + +```yaml +deployments: + - model_alias: generator + model: example/generator + server: + type: vllm + image: + name: vllm + environment: + MODEL_ACCESS_TOKEN: + type: secret + environment: MODEL_ACCESS_TOKEN +``` + +Literal non-secret settings use an explicit literal binding: + +```yaml + environment: + VLLM_LOGGING_LEVEL: + type: literal + value: INFO + MODEL_ACCESS_TOKEN: + type: secret + environment: MODEL_ACCESS_TOKEN +``` + +Export referenced variables in the trusted submit environment before running `execute`. Values are forwarded only at the required runtime boundary and are not serialized into run contracts or logs. Secret-shaped environment names cannot use literal bindings. + +## Use plugins + +Install plugin packages into the shared client environment or declare a compatible pure-wheel dependency. The serialized builder config can then reference the plugin column, seed reader, or processor. See the [plugin overview](/plugins/overview) for Data Designer plugin contracts. + +## Configure MCP providers + +Remote providers support SSE and streamable HTTP: + +```yaml +invocation: + num_records: 100 + dataset_name: tool-use + mcp_providers: + - provider_type: streamable_http + name: search + endpoint: https://tools.example.test/mcp + api_key: + type: secret + environment: MCP_API_KEY +``` + +Endpoints cannot contain credentials, query parameters, or fragments. + +Local stdio providers run inside the client allocation: + +```yaml +invocation: + num_records: 100 + dataset_name: tool-use + mcp_providers: + - provider_type: stdio + name: local-tool + command: example-mcp-server + args: + - --mode + - read-only + environment: + TOOL_API_KEY: + type: secret + environment: TOOL_API_KEY +``` + +The command is one executable token and arguments are passed without a shell. The shared client environment or dependency overlay must provide the executable. See [Tool Use and MCP](/concepts/tool-use-mcp/overview) for builder-side tool configuration. + +Code-sandbox sidecars are not supported by the v1 Slurm extension. Use an MCP service with an appropriate isolation boundary or run the workload through a separately managed sandbox system. diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx new file mode 100644 index 000000000..b4ec3815f --- /dev/null +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -0,0 +1,188 @@ +# Slurm Getting Started + +This flow starts from a clean Python environment in a site-approved Slurm submit environment and submits one small Data Designer run. Replace the example account, partitions, image sources, and model with values available on your cluster. + +Before starting, confirm the [Slurm prerequisites and host responsibilities](/slurm/overview). + +Obtain these site-specific values from your Slurm or platform administrator before continuing: + +- A Python 3.10 or later executable +- An account, GPU partition, CPU image-build partition, and shared workspace path +- A supported official vLLM release tag, a digest-qualified custom serving OCI image, or an absolute SQSH path containing a supported vLLM installation +- A model available to the serving image and compute nodes + +The site values and `example/generator` model below are placeholders. The official vLLM tag is a real image reference; select a supported release your site can access. + +## 1. Create and activate an environment + +Create the environment on a shared filesystem at the same absolute path seen by submit and compute nodes. If software installation is not allowed on login nodes, run these steps from a dedicated editor node or a short CPU allocation approved by your administrator. + +```bash +export DD_SLURM_WORKSPACE=/workspace/alice/data-designer-slurm +export DATA_DESIGNER_SLURM_PROFILE_FILE="$DD_SLURM_WORKSPACE/config/slurm-profiles.yml" +export DD_PYTHON="${DD_PYTHON:-python3}" + +mkdir -p "$DD_SLURM_WORKSPACE/config" "$DD_SLURM_WORKSPACE/envs" +"$DD_PYTHON" -c 'import sys; print(sys.version); raise SystemExit(sys.version_info < (3, 10))' +"$DD_PYTHON" -m venv "$DD_SLURM_WORKSPACE/envs/control" +source "$DD_SLURM_WORKSPACE/envs/control/bin/activate" +python -m pip install --upgrade pip +``` + +If the version check fails, load a site-provided Python module or set `DD_PYTHON` to a supported executable such as `python3.10`, then rerun the block. + +## 2. Install Data Designer with Slurm support + +```bash +python -m pip install "data-designer[slurm]" +python -m pip show data-designer-slurm +``` + +The base `data-designer` wheel does not install or import Slurm support. The extra installs the same-version `data-designer-slurm` wheel and exposes the lazy `data-designer slurm` command group. Do not continue if `pip show` cannot find the extension. + +The active environment is also the client runtime for submitted jobs. Planning records its exact Python executable, ABI, package inventory, and available install provenance; compute-node preflight verifies those facts before generation starts. Keep the environment at the same absolute shared path for the lifetime of submitted runs, and do not modify it while runs are active. + +## 3. Discover the CLI + +```bash +data-designer --help +data-designer slurm --help +``` + +If the Slurm command group is absent, see [Troubleshooting](/slurm/troubleshooting#the-slurm-command-group-is-missing). + +## 4. Create a profile + +Create the durable shared workspace first, then initialize an explicit profile: + +```bash +data-designer slurm profile init \ + --profile-file "$DATA_DESIGNER_SLURM_PROFILE_FILE" \ + --cluster primary \ + --workspace-root "$DD_SLURM_WORKSPACE" \ + --image-build-partition cpu \ + --account research \ + --partition gpu +``` + +The command writes the profile with restrictive permissions and never overwrites an existing file. It leaves `host_patterns` empty, so this single-cluster catalog uses `default_cluster: primary` unless `--cluster` is specified. Before submitting work, review every generated site value. In particular, confirm the account, partitions, workspace, GPU request mode, `gpus_per_node`, and the image-build CPU, memory, and time limit. The starter sets `gpus_per_node: auto`, which requests the GPU count reported by eligible nodes. The run below uses one GPU **per vLLM replica**, not necessarily one GPU for the allocation: on an eight-GPU node, it starts eight replicas. If your site permits partial-node allocations and you want to request only one GPU, set `gpus_per_node: 1` in the profile. Generation discovers the Slurm executable directory from the submit host; normally you do not need `scheduler.bin_path`. The starter image-build time limit is four hours and may exceed site policy. Edit it before `image add` when necessary. + +## 5. Validate the profile + +```bash +data-designer slurm profile validate +``` + +Validation checks strict profile loading, cluster selection, workspace access, and scheduler facts. See [Profiles and clusters](/slurm/profiles-and-clusters) for non-default files and multiple clusters. + +## 6. Register a serving image + +```bash +data-designer slurm image add \ + vllm/vllm-openai:v0.22.0 \ + --kind serving \ + --name vllm +``` + +Replace the tag with a supported release available to your site. The CLI pins the tag to a registry digest before submission, then submits a CPU Slurm job, waits for it to finish, checks the installed vLLM version, and only then publishes the alias. Keep the terminal open. You can instead import a custom digest-pinned OCI image or register an existing absolute `.sqsh` path. No Data Designer client image is needed. See [Images](/slurm/images). + +## 7. Create a Data Designer builder config + +Create `builder.yaml`: + +```yaml +data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + columns: + - name: tone + column_type: sampler + sampler_type: category + params: + values: [friendly, formal, playful] + - name: greeting + column_type: llm-text + prompt: Write one short {{ tone }} greeting. + model_alias: generator +``` + +The sampler chooses a tone for each row, and the model uses that value when writing the greeting. The model alias in this file must match a deployment in the Slurm run config. + +## 8. Create a Slurm run config + +Create `run.yaml` next to `builder.yaml`: + +```yaml +schema_version: 1 +name: greeting-run +builder: + source: builder.yaml +invocation: + num_records: 8 + dataset_name: greetings + model_concurrency: + generator: 8 +deployments: + - model_alias: generator + served_model_name: example/generator + model: example/generator + server: + type: vllm + image: + name: vllm + startup_timeout: 15m + distributed_init_timeout: 10m + topology: + tensor_parallel: 1 +submission: + time_limit: 01:00:00 +``` + +The repository also contains checked examples under `packages/data-designer-slurm/examples`. `served_model_name` is the name sent by the client to the vLLM-compatible endpoint. It may differ from `model`, which can be an absolute model path visible only to the serving container. Slurm writes Parquet shard and collection files; the run config needs no output section. See [GPU topology](/slurm/run-configuration#gpu-topology-and-replica-count) before changing `resources.nodes` or `tensor_parallel`. + +## 9. Dry-run the submission + +```bash +data-designer slurm execute run.yaml --dry-run +``` + +Dry-run resolves the selected profile, records the active shared Python runtime, verifies registered serving-image contents, loads the builder envelope, resolves client dependencies, and renders the batch script without calling `sbatch`. Full builder and plugin validation runs with that Python environment on the compute node. Preview staging is temporary: dry-run does not initialize durable managed-run state and cannot be resumed or passed to `status`. + +## 10. Submit the run + +```bash +data-designer slurm execute run.yaml +``` + +The command prints a readable result, a status command, and a first-shard log-tail command in an interactive terminal. Save its run ID; all lifecycle commands use the managed run ID rather than a raw Slurm job ID. To capture a JSON object instead, use `data-designer slurm --output json execute run.yaml`. + +## 11. Inspect or cancel the run + +```bash +data-designer slurm status +data-designer slurm cancel +``` + +Status reconciles active queue and accounting evidence with durable run state. Cancellation targets only jobs recorded for the managed run. + +## 12. Inspect results + +Each successful shard publishes a canonical winner at: + +```text +/workspace/alice/data-designer-slurm/runs//shards//winner.json +``` + +The winner's `candidate_manifest.path` points to the winning attempt's `output-manifest.json`. That manifest records `dataset_path` and the relative path, size, digest, and record count of every canonical output file. Follow this chain instead of recursively scanning for Parquet files, because failed or superseded attempts can retain non-canonical files. + +Use `status` to confirm every shard succeeded. To create one collected dataset at the default destination under the run directory, submit a managed collection job: + +```bash +data-designer slurm merge +``` + +The command writes to `/workspace/alice/data-designer-slurm/runs//merge-output`. `--output-path` selects another authorized destination. `--num-partitions` controls the collected Parquet file count and defaults to 1 for a new collection. + +Continue with [Run operations](/slurm/run-operations), [Retry and collection](/slurm/retry-and-collection), [Scaling runs](/slurm/run-configuration#scale-runs), [Dependencies and tools](/slurm/dependencies-and-tools), or the [CLI reference](/slurm/cli-reference). diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx new file mode 100644 index 000000000..1bab843a0 --- /dev/null +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -0,0 +1,96 @@ +# Slurm images + +Data Designer Slurm images run vLLM model servers. The Data Designer client runs directly from the shared Python environment used to invoke the CLI; it does not require an image. + +## Choose a serving image + +For vLLM, the simplest path is to select a supported release of the official `vllm/vllm-openai` image. You do not need to build a Data Designer client image or write a Dockerfile for this path: + +```bash +data-designer slurm image add vllm/vllm-openai:v0.22.0 \ + --kind serving \ + --name vllm +``` + +Replace `v0.22.0` with a supported release that is available to your site. The CLI resolves that tag to an immutable registry digest before submitting the import job, then imports by digest on Enroot 3.5 or later. The imported image is inspected, and registration fails if its installed vLLM version does not match the requested release. The alias records the inspected version and source digest. A later change to the tag cannot silently change the submitted import. + +If your site requires a customized serving environment, obtain a digest-qualified OCI reference or absolute compute-visible SQSH path from your platform administrator, or publish one through your site's container build process. `image add` imports or inspects that image; it does not build a serving environment from a Dockerfile. The image must provide a supported vLLM executable and all runtime dependencies for the selected model. + +Before registration, confirm that the image source is accessible from the configured CPU image-build partition and that the model is accessible from the GPU partition. Keep registry credentials in the compute environment or Enroot configuration. + +## Import a custom digest-pinned OCI image + +The digest below demonstrates the required syntax and does not identify a real image. Replace it before running the command. + +```bash +data-designer slurm image add \ + registry.example.test/inference/vllm@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ + --kind serving \ + --name vllm +``` + +Custom OCI sources must use `repository@sha256:<64 lowercase hex characters>`. The official versioned `vllm/vllm-openai:vX.Y.Z` form above is the only supported tag-only exception. Other tags without a digest, schemes, embedded credentials, query strings, and fragments are rejected. Registry credentials belong in the compute environment or Enroot configuration, not in the command or profile. + +The command stages package-owned inspection code, submits a CPU job using `image_build`, waits for terminal Slurm state, verifies the SQSH and inspection record, then atomically publishes the artifact and alias. The terminal stays attached while the job runs. + +In an interactive terminal, `image add` reports the job ID, scheduler transitions, import and inspection phases, and elapsed time. It checks for new logs every 30 seconds and prints a status heartbeat after two minutes without new activity. It follows the image job's standard output and error by default: + +```bash +data-designer slurm image add vllm/vllm-openai:v0.22.0 \ + --kind serving --name vllm +``` + +The progress and log stream is written to standard error with UTC timestamps; followed logs are bounded to 1 MiB and blank lines are omitted. The final readable result remains on standard output. Logs start only after the CPU job receives an allocation and may contain details emitted by the image or container runtime. The final available log output is shown before temporary job files are cleaned up, including when the job fails. Use `--no-follow-logs` to suppress streaming, or `--follow-logs` to enable it explicitly in a non-interactive session. For a machine-readable result, put `--output json` immediately after `slurm`. If you need durable logs, capture standard error in your own private file. + +## Register an existing SQSH + +```bash +data-designer slurm image add \ + /workspace/alice/images/vllm.sqsh \ + --kind serving \ + --name vllm +``` + +The absolute SQSH path must be visible on the image-build node and all later compute nodes. Existing images are inspected by a CPU job in place. + +## Inspect aliases + +```bash +data-designer slurm image ls +data-designer slurm image info vllm +``` + +`image ls` shows each alias, image kind, detected runtime version, and SQSH path. Use `image info` for the full image details and hashes. `--output json` retains the complete registration record for scripts. + +Aliases are cluster-profile scoped because each profile has its own workspace. Run configuration can reference an alias with `name` or a registered absolute path with `path`. + +## Replace or remove an alias + +An alias is immutable unless replacement is explicit: + +```bash +data-designer slurm image add /workspace/alice/images/vllm-v2.sqsh \ + --kind serving \ + --name vllm \ + --replace +``` + +Removing an alias does not delete its SQSH artifact: + +```bash +data-designer slurm image rm vllm +``` + +To delete the SQSH too, use `--delete-artifact`. This works only for an OCI artifact imported into the selected workspace by `image add`, when no other registered alias or persisted run plan references the file. Existing SQSH files registered from outside the workspace are never deleted by this command. + +```bash +data-designer slurm image rm vllm --delete-artifact +``` + +The command confirms the alias removal and whether the SQSH was retained or deleted. + +Published OCI imports are content-addressed. Identical imports can share an artifact, while conflicting bytes or inspection facts fail closed. + +## Compatibility rules + +The shared client environment must remain available at the exact recorded Python path with the same Python ABI and package inventory on every compute node. It must include the Slurm extra's compatible `aiohttp` dependency for the native inference proxy. Dependency overlays may add pure Python wheels, but cannot replace distributions already installed in that environment. Serving images must expose a supported vLLM executable and version. diff --git a/fern/versions/latest/pages/slurm/operations.mdx b/fern/versions/latest/pages/slurm/operations.mdx new file mode 100644 index 000000000..a965b3131 --- /dev/null +++ b/fern/versions/latest/pages/slurm/operations.mdx @@ -0,0 +1,82 @@ +# Slurm run operations + +The CLI prints readable labels and values by default. To get canonical JSON for scripts or agents, put `--output json` immediately after `slurm`, before the command. Errors are readable on stderr by default and retain structured JSON and the same exit codes in JSON mode. + +```bash +data-designer slurm --output json execute run.yaml --dry-run > dry-run.json +data-designer slurm --output json status > status.json +``` + +You can also set `DATA_DESIGNER_SLURM_OUTPUT=json` for a script that calls several Slurm commands. The default remains human-readable when that variable is unset. In either mode, progress and optional image-job logs go to stderr with UTC display timestamps; do not parse them as the command result. + +## Resolve without submitting + +```bash +data-designer slurm execute run.yaml --dry-run +``` + +Dry-run performs profile selection, image verification, builder-envelope loading, dependency resolution, plan compilation, runtime staging, and batch rendering. It does not call `sbatch`. Full builder and plugin validation runs later in the client environment. Because resolution is real, the selected workspace and registered images must already be available. Preview staging is temporary: dry-run does not create durable managed-run state, so its result cannot be resumed, inspected with `status`, or used as collection input. + +## Submit + +```bash +data-designer slurm execute run.yaml +``` + +A successful response includes `run_id`, `state`, `plan_sha256`, `shard_count`, and `job_id`. The package persists immutable inputs before submission and records scheduler identity after acceptance. Conflicting durable state is never overwritten. + +In an interactive terminal, `execute` reports when it invokes `sbatch` with the rendered script on standard input. After Slurm accepts and releases the job, it prints a `status` command and a `tail -F` example for the first shard's stdout and stderr. Logs appear only after the job receives an allocation. The invocation display omits exported environment details and is not a replay command; use `--dry-run` to inspect the rendered script. Do not submit again just because a job is pending. + +Use `--profile-file` and `--cluster` on any command when the defaults are not appropriate. + +For GPU allocation, tensor parallelism, replica counts, and array task concurrency, see [Scale runs](/slurm/run-configuration#scale-runs) in Run Configuration. `status` reports the resulting run and shard state; topology is chosen in the run file before submission. + +## Inspect status + +```bash +data-designer slurm status +``` + +Status loads the managed run, queries active queue and accounting records for its known jobs, and reconciles those observations with persisted attempts, readiness, and shard winners. Accounting can lag after a job leaves the active queue; rerun status after the cluster accounting service catches up. + +The response includes every shard in planned order. Use the attempt lifecycle and readiness details to distinguish queue wait, model startup, active generation, publication, terminal failure, and success. + +## Cancel + +```bash +data-designer slurm cancel +``` + +Cancellation selects only active job IDs recorded for the managed run. It is safe to call again when no managed jobs remain active. + +## Durable artifacts + +Managed files are below `/runs/`: + +- Authored and resolved configuration, dependency lock, runtime bundle reference, and run manifest +- Slurm stdout and stderr +- Attempt and readiness records +- Candidate results, shard winners, collection state, and output + +Allocation-local scratch is deliberately absent after cleanup. The stable container path `/run/data-designer-slurm` is not a durable troubleshooting location. + +For each shard, `shards//winner.json` is the canonical pointer. Its `candidate_manifest.path` identifies the winning attempt's `output-manifest.json`, which records the canonical `dataset_path` and output files. Do not identify results by recursively scanning attempt directories: failed and superseded attempts may retain valid-looking data files. The `merge` command validates these winner chains before creating a run-level dataset at the requested destination. + +## Use the Python service API + +```python +from pathlib import Path + +from data_designer.slurm.config import load_run_config +from data_designer.slurm.services import create_slurm_run_service + +run_file = Path("run.yaml").resolve() +config = load_run_config(run_file) +service = create_slurm_run_service(cluster="primary") + +preview = service.execute(config, source_root=run_file.parent, dry_run=True) +submission = service.execute(config, source_root=run_file.parent) +status = service.status(submission.run_id) +``` + +Public service methods apply the same validation, persistence, error model, and output contracts as the CLI. diff --git a/fern/versions/latest/pages/slurm/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx new file mode 100644 index 000000000..e6f64deac --- /dev/null +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -0,0 +1,39 @@ +# Data Designer on Slurm + +Data Designer Slurm runs Data Designer clients and vLLM model servers in batch allocations. The optional extension owns profile selection, image inspection, submission, status, retry, result collection, and benchmark workflows. + +Start with [Getting Started](/slurm/getting-started). Use [Dependencies and tools](/slurm/dependencies-and-tools) for plugins, MCP providers, and secrets, and keep the [CLI reference](/slurm/cli-reference) nearby for automation. + +## Prerequisites + +The submit environment needs: + +- Python 3.10 or later +- `sbatch`, `squeue`, `sacct`, `scancel`, `scontrol`, and `sinfo` +- Access to a shared workspace that is visible from the submit host and every eligible compute node + +Compute nodes need access to the same shared Python environment at the exact path used by the submit command. Enroot 3.5 or later is required for both digest-pinned OCI imports and existing SQSH inspection. Serving images must contain vLLM. + +Your Slurm administrator can confirm the account, GPU partition, CPU image-build partition, GPU request mode, and shared filesystem paths for your cluster. + +## Where each task runs + +| Persona | Responsibilities | +| --- | --- | +| Workstation | Author and version dataset and Slurm configuration. Copy files to a cluster-visible location before submission. | +| Site-approved submit environment | Install the CLI, select a profile, register images, submit work, inspect status, retry, cancel, and collect output. This may be a login host, editor node, or CPU allocation, depending on site policy. | +| Compute node | Import or inspect serving images, start vLLM containers, run the native Data Designer client from the shared Python environment, generate records, and write managed state and results. | + +Run Slurm commands from a site-approved submit environment with the cluster CLI and shared filesystem. Do not install packages or run substantive work on a login node when site policy prohibits it. A laptop without the cluster CLI and shared filesystem can author configuration, but cannot validate or submit it. + +## Storage model + +The configured `workspace_root` is durable and shared. It contains image metadata, immutable run inputs, Slurm logs, attempt state, winning results, and collection state. `slurm merge RUN_ID` collects winners under `/runs//merge-output` by default. Use `--output-path` to select another authorized destination. + +Runtime bundle extraction, the client dependency overlay, Data Designer home and cache directories, and Enroot state use allocation-local scratch. Serving containers mount that scratch at `/run/data-designer-slurm`; native client steps use the corresponding host path. The runtime removes it on success, failure, cancellation, and retry. Do not configure a mount that overlaps this path or expect transient scratch files to remain after a job. + +## Supported in v1 + +The public v1 surface supports named cluster profiles, serving image aliases, one or more vLLM deployments, independent replicas, multi-node replicas, Slurm arrays, bounded array concurrency, durable status, cancellation, sparse retry, winner-driven collection, dependencies, plugins, MCP providers, and benchmarks. + +Interactive allocations, code-sandbox sidecars, NeMo Run, SGLang, Dynamo, Kubernetes, Ray, and platform-specific orchestration are not part of the v1 Slurm extension. diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx new file mode 100644 index 000000000..30ad8d69e --- /dev/null +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -0,0 +1,119 @@ +# Slurm profiles and clusters + +A profile catalog stores site facts separately from portable run intent. It selects one cluster and supplies scheduler defaults, workspace storage, image-build resources, GPU facts, and container mounts. + +## Locate the catalog + +Commands resolve the catalog in this order: + +1. `--profile-file PATH` +2. `DATA_DESIGNER_SLURM_PROFILE_FILE` +3. `~/.data-designer-slurm-profile.yml` + +Pass `--cluster NAME` to select an entry explicitly. Without it, hostname patterns are considered first, followed by `default_cluster` when there is no match. Duplicate patterns and ambiguous hostname matches are rejected. + +## Initialize a catalog + +```bash +data-designer slurm profile init \ + --profile-file /workspace/alice/config/slurm-profiles.yml \ + --workspace-root /workspace/alice/data-designer-slurm \ + --image-build-partition cpu \ + --cluster primary \ + --account research \ + --partition gpu +``` + +The starter sets `host_patterns: []`, `gpus_per_node: auto`, `gpu_request_mode: gres`, and CPU image-build resources of 2 CPUs, 8 GiB, and four hours. Initialization creates missing parent directories of a custom `--profile-file` with private permissions. Review the generated account, partitions, workspace, GPU settings, and image-build resources against site policy before validation or image registration. For hostname-based selection in a multi-cluster catalog, pass one or more explicit `--host-pattern` values during initialization or edit `host_patterns` in the catalog. + +## Configure multiple clusters + +```yaml +schema_version: 1 +default_cluster: primary +clusters: + primary: + schema_version: 1 + host_patterns: + - login*.example.test + scheduler: + account: research + partition: gpu + mem_per_gpu: 80G + gpus_per_node: 8 + workspace_root: /workspace/alice/data-designer-slurm + image_build: + partition: cpu + cpus_per_task: 2 + memory: 8G + time_limit: "04:00:00" + gpu_request_mode: gres + lab: + schema_version: 1 + host_patterns: + - lab-login*.example.test + scheduler: + partition: accelerated + gpus_per_node: 4 + workspace_root: /workspace/alice/data-designer-slurm-lab + image_build: + partition: general + cpus_per_task: 2 + memory: 8G + time_limit: "04:00:00" + gpu_request_mode: visible +``` + +Run a command against a non-default cluster: + +```bash +data-designer slurm profile validate --cluster lab +data-designer slurm image ls --cluster lab +data-designer slurm execute run.yaml --cluster lab +``` + +`visible` GPU mode asks Slurm for visible GPUs without GRES memory directives. It cannot be combined with `mem_per_gpu`. + +To let Slurm choose among GPU partitions, set `scheduler.partition` to a comma-separated value or a YAML list: + +```yaml +scheduler: + partition: + - gpu-a + - gpu-b +``` + +The CLI also accepts `--partition gpu-a,gpu-b` during profile initialization. Slurm chooses a partition when scheduling the job; the CPU `image_build.partition` remains a single partition. + +`gpus_per_node` is both the GPU count used for deployment topology and, in `gres` mode, the count requested from Slurm on each job node. `auto` queries eligible nodes and uses their reported count, so an eight-GPU partition produces `--gres=gpu:8`. Set an explicit smaller count, such as `1`, for partial-node allocations when the cluster permits them. Use `auto` only when requesting every GPU on each selected node is intended, and use an explicit count when eligible nodes are heterogeneous. + +Multi-node generation jobs request exclusive nodes, including in `gres` mode. This keeps serving ports isolated because Slurm reports batch-job GPU IDs for the first node only. A partial-node `gpus_per_node` setting still limits the GPUs used per node, but other jobs cannot share those nodes during the run. + +The profile decides how many GPUs the job requests **per node**. Each deployment's `tensor_parallel` and `nodes_per_replica` decide how those GPUs form vLLM replicas. For example, one node with `gpus_per_node: 8` and `tensor_parallel: 1` runs eight one-GPU replicas; it does not request only one GPU. See [GPU topology and replica count](/slurm/run-configuration#gpu-topology-and-replica-count). + +When selecting multiple partitions with `gpus_per_node: auto`, every selected partition must report GPU resources and all eligible nodes must report the same GPU count. Choose partitions that can also satisfy the run's GPU topology, model, and image requirements. + +Generation automatically discovers the directory containing `srun` on the submit host and records it in the immutable run plan, then prepends it to the batch job's controlled `PATH`. Load your site's Slurm environment before submitting; no profile setting is normally needed. If the submit-host path is not available on compute nodes, set `scheduler.bin_path` as an override. It must be one absolute directory containing runnable `srun` and `scontrol`, visible from compute nodes; do not provide a colon-separated list. + +## Add mounts + +Use `container_mounts` for absolute compute-visible paths that the client or model needs: + +```yaml +container_mounts: + - source: /workspace/alice/models + target: /models + read_only: true +``` + +Targets must be unique and must not overlap `/run/data-designer-slurm`, which is reserved for allocation-local scratch. + +## Validate before submission + +```bash +data-designer slurm profile validate \ + --profile-file /workspace/alice/config/slurm-profiles.yml \ + --cluster primary +``` + +Validation creates and removes a probe in `workspace_root`. With `gpus_per_node: auto`, it also queries eligible nodes and requires one consistent positive GPU count. diff --git a/fern/versions/latest/pages/slurm/retry-and-collection.mdx b/fern/versions/latest/pages/slurm/retry-and-collection.mdx new file mode 100644 index 000000000..c3ba3dc15 --- /dev/null +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -0,0 +1,67 @@ +# Retry and collection + +Retry and collection operate on persisted run state. They do not infer ownership from arbitrary Slurm jobs or directories. + +## Preview a retry + +```bash +data-designer slurm retry --dry-run +``` + +By default, retry selects failed or incomplete shards and resolves `--resume if_possible` to either `always` or `never` from persisted compatibility evidence. The response includes selected shard IDs, new attempt IDs, effective resume mode, and a batch script. + +Select specific zero-based array tasks when needed: + +```bash +data-designer slurm retry \ + --task-id 2 \ + --task-id 7 \ + --resume never \ + --dry-run +``` + +## Submit a retry + +```bash +data-designer slurm retry +``` + +Interactive use asks for confirmation. Add `--force` to skip that prompt in automation. Force does not replace immutable run metadata or discard scheduler evidence. + +You may pass a managed Slurm job ID instead of a run ID when the package can resolve it unambiguously. Prefer the run ID in scripts. + +Retries create new attempts. The original attempt records remain durable, and each shard publishes at most one winner. Status shows the complete attempt history. + +## Collect winners + +Shards with at least one generated record can publish Parquet winners, even when they produce fewer records than requested. Use `merge` to submit a zero-GPU CPU job that collects winners under `/runs//merge-output`. The merged dataset may contain fewer records than the run requested: + +```bash +data-designer slurm merge +``` + +The merge reads the immutable winning shard outputs and verifies their schemas. If nested struct fields appear in different orders across shards, it writes them in the first shard's field order. Missing fields and changed field types are rejected. The collection job uses the installed Slurm package version and runtime bundle from the merge command, so a newer installation can collect an earlier run without changing its shard records. + +For an explicit destination, use `--output-path`: + +```bash +data-designer slurm merge --output-path /workspace/alice/exports/final-dataset +``` + +You can also select the run with `--input-path` instead of the positional run ID: + +```bash +data-designer slurm merge \ + --input-path /workspace/alice/data-designer-slurm/runs/ \ + --output-path /workspace/alice/exports/final-dataset +``` + +`--input-path` must be exactly a managed run directory under the selected profile workspace. An explicit output path must be absolute, writable, compute-visible, and authorized by the profile mounts. + +`--num-partitions` selects the number of collected Parquet files. It defaults to 1 for a new collection and cannot exceed the total number of records in the winning shards: + +```bash +data-designer slurm merge --num-partitions 20 +``` + +Repeating the same merge command returns or reuses managed collection state when safe. If you omit `--output-path` or `--num-partitions` on a repeat call, the existing collection keeps its destination and partition count. Explicitly changing either for an active collection is a conflict. diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx new file mode 100644 index 000000000..0b8186340 --- /dev/null +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -0,0 +1,329 @@ +# Slurm run configuration + +A run file is portable intent. Cluster-specific accounts, partitions, workspace paths, GPU facts, and image-build resources belong in a [profile](/slurm/profiles-and-clusters). + +## Required structure + +```yaml +schema_version: 1 +name: product-descriptions +builder: + source: builder.yaml +invocation: + num_records: 1000 + dataset_name: product-descriptions + model_concurrency: + generator: 64 +deployments: + - model_alias: generator + served_model_name: example/generator + model: example/generator + server: + type: vllm + image: + name: vllm + resources: + nodes: 1 + topology: + tensor_parallel: 8 + nodes_per_replica: 1 +array_tasks: + count: 1 +submission: + job_name: product-descriptions + time_limit: 03:55:00 +``` + +`builder` accepts exactly one complete serialized Data Designer config through `source` or `inline`. Relative source paths resolve from the run file directory. `invocation.run_config` accepts public `RunConfig` fields except `display_tui`, which is disabled for Slurm jobs. `input_bindings.seed_path` and `managed_assets_path` accept absolute compute-visible paths. + +Every `model_alias` in `model_concurrency` and the builder must have one deployment. Deployment aliases must be unique. `model` selects the Hugging Face model ID or absolute model path loaded by the serving container. `served_model_name` is the name exposed to the client and defaults to `model`; set it explicitly when `model` is a container-visible path but the builder uses a stable model name. Shard and collected output files are Parquet. Choose the destination and partition count with `slurm merge` after the run completes. + +The client runs from the same shared Python environment as the CLI on the first allocated GPU node. Its step requests no GPUs; the serving steps use the GPUs and their registered images. The `client` block is optional; `client.cpus` defaults to 16. It sets the client step's CPU request and contributes to the batch allocation's CPU request. + +Use `submission.partition` to override the profile's GPU partition selection for one run. It accepts a comma-separated value or a YAML list, such as: + +```yaml +submission: + partition: [gpu-a, gpu-b] +``` + +Slurm schedules the job in one of the selected partitions. With `gpus_per_node: auto`, every selected partition must report GPU resources with one consistent count. + +## Configure vLLM + +Each server supports readiness and launch controls: + +```yaml +server: + type: vllm + image: + name: vllm + startup_timeout: 20m + distributed_init_timeout: 10m + lead_boot_standoff: 60s + rank_launch_stagger: 5s + readiness_path: /health + queue_backpressure: + max_waiting_requests: 128 + retry_after_seconds: 1 + extra_args: + - --dtype + - bfloat16 +``` + +`distributed_init_timeout` must not exceed `startup_timeout`. Increase both when distributed initialization legitimately needs more time. + +Set `server.io_processor_plugin` to a registered vLLM IO processor name, such as `deberta_gliner_io` for GLiNER. The serving image must include the plugin and its dependencies, and the model must be prepared for that plugin. Pooling options such as `--runner pooling` belong in `extra_args`. + +`extra_args` is a list of individual arguments, not a shell command. Options owned by the orchestrator and unsafe shell forms are rejected. Configure model identity, topology, networking, and lifecycle through their typed fields instead of passing reserved options such as `--model`, `--served-model-name`, `--host`, `--port`, `--tensor-parallel-size`, `--pipeline-parallel-size`, or distributed-launch flags. In particular, `served_model_name` belongs on the deployment, alongside `model`, not under `server.extra_args`. + +## Scale runs + +Each array task gets a separate allocation containing all deployments. A deployment's `resources.nodes` is its node count **within one array task**; different deployments use different nodes by default. For example, a run with a two-node writer and a one-node reviewer requests three nodes per array task. `array_tasks.max_concurrent` limits how many such allocations may run at once. + +### Arrays + +Split the requested total into independent shards and cap the number of allocations that can run at once: + +```yaml +invocation: + num_records: 1000 + dataset_name: sharded-run +array_tasks: + count: 5 + max_concurrent: 2 +``` + +This submits one Slurm job array with tasks `0` through `4` (`--array=0-4%2`). Each task runs the **entire deployment set** and plans 200 records: task `0` owns indices `[0, 200)`, task `1` owns `[200, 400)`, and so on through task `4` with `[800, 1000)`. Task `2` corresponds to `shard-00002`. At most two tasks can run concurrently, subject to available cluster resources. Omit `max_concurrent` to let Slurm schedule all five as resources permit. When the total does not divide evenly, the last task gets the remainder. `count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. + +`num_records` is the planned total, not a guaranteed output count. A nonempty shard can succeed with fewer records; an empty shard needs a retry. Status, retry, and collection retain each task's shard identity and winning attempt. Array task count does not set the number of collected Parquet files: `slurm merge RUN_ID --num-partitions 5` requests five files from the winning records. See [Retry and collection](/slurm/retry-and-collection) for partial results. + +After submission, inspect the managed run ID. If task `2` fails, preview a retry of only that shard: + +```bash +data-designer slurm status +data-designer slurm retry --task-id 2 --dry-run +``` + +Runs with more than one array task do not support global processors or profilers; image, plugin, custom, or unknown column types; local-callable validators; shuffled seed input; or authored seed-selection strategies. Seeded array runs must bind the dataset through `invocation.input_bindings.seed_path`. + +### Independent replicas + +Allocate multiple nodes and use all eight GPUs per node for one vLLM replica on each node: + +```yaml +resources: + nodes: 4 +topology: + tensor_parallel: 8 + nodes_per_replica: 1 +``` + +In an eight-GPU profile, this deployment starts four independent eight-GPU servers. To start one server **per GPU** instead, use `tensor_parallel: 1`; with the same four nodes, that starts 32 servers. Both layouts reserve four eight-GPU nodes in a GRES profile. + +```yaml +resources: + nodes: 4 +topology: + tensor_parallel: 1 + nodes_per_replica: 1 +``` + +### Multi-node replicas + +Use more than one node for each replica: + +```yaml +resources: + nodes: 4 +topology: + tensor_parallel: 8 + nodes_per_replica: 2 +``` + +In an eight-GPU profile, this starts two replicas, each spanning two nodes with 16 GPUs total. `nodes_per_replica` maps to pipeline parallelism; it does not increase the number of independent replicas. + +### GPU topology and replica count + +`tensor_parallel` is the number of GPUs each replica uses **on each node**. `nodes_per_replica` is the number of nodes in its pipeline-parallel group. The planner starts `gpus_per_node / tensor_parallel` replica lanes per node group, so the total for one deployment is: + +```text +replicas = (resources.nodes / nodes_per_replica) * (gpus_per_node / tensor_parallel) +GPUs per replica = tensor_parallel * nodes_per_replica +``` + +For a profile that requests eight GPUs per node, the following are separate choices for **one deployment in one array task**: + +| `resources.nodes` | `tensor_parallel` | `nodes_per_replica` | vLLM replicas | GPUs per replica | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 1 | 1 | 8 | 1 | +| 4 | 1 | 1 | 32 | 1 | +| 4 | 2 | 1 | 16 | 2 | +| 4 | 4 | 1 | 8 | 4 | +| 4 | 8 | 1 | 4 | 8 | +| 4 | 8 | 2 | 2 | 16 | + +Increasing `resources.nodes` at TP=1 creates more independent one-GPU servers; it does not combine their GPUs into one larger server. The client uses one logical endpoint per model alias, backed by a local proxy that routes requests among that alias's replicas. `invocation.model_concurrency` limits in-flight model calls **per alias within each array task**; it does not set the replica count or the number of array tasks. + +`tensor_parallel` must divide the GPUs requested per node in the selected profile. `nodes_per_replica` must divide `resources.nodes`. With `gpus_per_node: auto`, the eligible nodes must agree on their GPU count. Multi-node expert parallelism is not supported in v1. + +### Multiple models and images + +Add one deployment per builder model alias. Each deployment may use a different registered serving image: + +```yaml +deployments: + - model_alias: writer + model: example/writer + server: + type: vllm + image: + name: vllm-writer + - model_alias: judge + model: example/judge + server: + type: vllm + image: + name: vllm-judge +invocation: + num_records: 1000 + dataset_name: judged-output + model_concurrency: + writer: 64 + judge: 32 +``` + +### Different models on one node + +By default, deployments use disjoint nodes and each receives the profile's GPU budget. To place different models on the same physical node, give their resources the same `shared_node` name and explicit GPU quotas: + +```yaml +deployments: + - model_alias: writer + model: Qwen/Qwen2.5-0.5B-Instruct + resources: + shared_node: models + gpus_per_node: 1 + topology: + tensor_parallel: 1 + server: + type: vllm + image: {name: vllm} + - model_alias: judge + model: microsoft/Phi-3-mini-4k-instruct + resources: + shared_node: models + gpus_per_node: 1 + topology: + tensor_parallel: 1 + server: + type: vllm + image: {name: vllm} +``` + +Declare both aliases in the Data Designer builder. For example, the writer can generate an `answer` column and the judge can evaluate `{{ answer }}` in a dependent column. A complete four-record example is included at `packages/data-designer-slurm/examples/colocated-models.yaml`. + +The profile's `gpus_per_node` is the Slurm allocation budget. Set it to `2` for the example above; the quota sum must fit that budget. Each quota is a positive number of GPUs reserved for that deployment. `tensor_parallel` must divide its quota: a quota of four with TP two starts two replicas of that model. The compiler assigns immutable disjoint GPU slices and separate backend and logical endpoint ports. GPU identifiers are resolved from the serving step's actual CUDA visibility, including Slurm device remapping. Each alias retains its own readiness, backend order, concurrency, and status; a model failure stops the shared worker and fails the attempt. + +Shared placement requires one group covering every deployment in the run, at least two deployments, `resources.nodes: 1`, `nodes_per_replica: 1`, GRES GPU request mode, the same resolved serving image, and identical `server.environment` and `queue_backpressure` settings. Model identities, served names, extra arguments, and launch/readiness timings may differ. Multiple groups, mixtures of shared and disjoint placements, multi-node shared groups, and different images within one group are rejected before submission. Memory fit is not inferred. An allocation can use fewer GPUs than the physical node has; sharing does not add node exclusivity. + +To share a GPU explicitly, set `resources.gpu_start_index: 0` and `gpus_per_node: 1` on both deployments, with a profile budget of one GPU. For a TP-two model sharing its second GPU with a detector, use offsets `0` and `1` with quotas `2` and `1` and a profile budget of two. Each slice must fit the budget; automatic placement follows the highest assigned index. Overlapping models start in launch-delay order, waiting for a model's readiness before starting another model on its GPUs. Configure per-model `--gpu-memory-utilization`, context length, and concurrency through `server.extra_args` to leave sufficient memory for both models and startup allocations. + +Python callers use the same resource fields with `with_deployment`: + +```python +slurm.with_deployment({ + "model_alias": "writer", + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "resources": {"shared_node": "models", "gpus_per_node": 1}, + "topology": {"tensor_parallel": 1}, + "server": {"type": "vllm", "image": {"name": "vllm"}}, +}) +``` + +Add the judge deployment with the same group name before calling `build()` or `write_config()`. + +### Put topologies and shards together + +This example plans 1,000 records across five array tasks. Each task allocates two nodes for one-GPU `writer` replicas and one node for a single eight-GPU `reviewer` replica. The two aliases may load different model families and serving images: + +```yaml +schema_version: 1 +name: sharded-two-model-run +builder: + source: builder.yaml +invocation: + num_records: 1000 + dataset_name: reviewed-records + model_concurrency: + writer: 16 + reviewer: 16 +deployments: + - model_alias: writer + model: example/writer-model + server: + type: vllm + image: + name: vllm-writer + resources: + nodes: 2 + topology: + tensor_parallel: 1 + nodes_per_replica: 1 + - model_alias: reviewer + model: example/reviewer-model + server: + type: vllm + image: + name: vllm-reviewer + resources: + nodes: 1 + topology: + tensor_parallel: 8 + nodes_per_replica: 1 +array_tasks: + count: 5 + max_concurrent: 2 +``` + +With eight GPUs per node, **each task** has 16 writer replicas and one reviewer replica on three nodes. At most two tasks can run together, so the run can occupy up to six nodes and has up to 32 writer replicas and two reviewer replicas while both tasks are active. The client can send up to 16 in-flight calls per alias in each task, or up to 32 per alias across two active tasks. Each task has its own client, proxies, and 200 planned records; the tasks do not share a model endpoint. `builder.yaml` must declare `writer` and `reviewer` model aliases, and both image aliases and model IDs above must be replaced with registered images and models available to your cluster. + +## Build a run in Python + +```python +from data_designer.config import DataDesignerConfigBuilder, LLMTextColumnConfig, ModelConfig +from data_designer.slurm.config import DataDesignerSlurmConfigBuilder + +builder = DataDesignerConfigBuilder( + model_configs=[ + ModelConfig(alias="generator", model="example/generator", provider="openai") + ] +) +builder.add_column( + LLMTextColumnConfig( + name="greeting", + prompt="Write one short greeting.", + model_alias="generator", + ) +) + +slurm = ( + DataDesignerSlurmConfigBuilder.from_config_builder(builder, name="generated-run") + .with_invocation( + num_records=100, + dataset_name="generated", + model_concurrency={"generator": 16}, + ) + .with_client(cpus=16) + .with_deployment( + { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + } + ) +) +slurm.write_config("run.yaml") +``` + +The Slurm builder validates each orchestration step and writes strict YAML or JSON. Full Data Designer builder and plugin validation runs in the prepared client environment. The CLI rejects duplicate keys, YAML anchors, merge keys, unknown fields, and environment interpolation outside the serialized Data Designer builder payload. diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx new file mode 100644 index 000000000..e1a7476ca --- /dev/null +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -0,0 +1,98 @@ +# Slurm troubleshooting + +Start with the error on stderr, then inspect the managed run directory and Slurm state. Use `data-designer slurm --output json COMMAND` for a structured error. The CLI uses these exit codes: + +| Exit code | Error code | Meaning | +| --- | --- | --- | +| 1 | `internal` | Unexpected package failure | +| 2 | `invalid_request` | Invalid command input or configuration | +| 3 | `not_found` | Managed run, job, image, or benchmark not found | +| 4 | `conflict` | Existing immutable state conflicts with the request | +| 5 | `unavailable` | Slurm, storage, image, or runtime dependency unavailable | + +## The Slurm command group is missing + +Confirm the optional extra is installed in the active environment: + +```bash +python -m pip show data-designer +python -m pip show data-designer-slurm +data-designer slurm --help +``` + +Install `data-designer[slurm]` rather than only `data-designer`. The extension wheel must have the same version as the base wheel. If pip warns that the selected `data-designer` release does not provide the `slurm` extra, that package index or version does not contain the extension; select a release that publishes both packages instead of continuing with the base CLI. + +## Profile selection fails + +Pass both selectors explicitly while debugging: + +```bash +data-designer slurm profile validate \ + --profile-file /workspace/alice/config/slurm-profiles.yml \ + --cluster primary +``` + +Check that the profile file and workspace are accessible from the submit host, hostname patterns do not overlap, and `default_cluster` names an existing entry. Set `gpus_per_node` explicitly if an eligible partition contains heterogeneous nodes. + +## Slurm commands are unavailable in a batch job + +Generation discovers the `srun` directory from the submit host and checks that `scontrol` is there before submitting. If discovery fails, load your site's Slurm environment and retry preparation; the error names the missing tool. If the discovered directory does not exist on compute nodes, set `scheduler.bin_path` in the selected profile to a shared absolute directory containing runnable `srun` and `scontrol`. The generated batch script prepends that directory to its controlled `PATH`. + +## A small run requests every GPU on each node + +`gpus_per_node: auto` uses the GPU count reported by eligible nodes and requests that count in GRES mode. Set an explicit count such as `gpus_per_node: 1` when the cluster permits partial-node allocations, then keep deployment topology compatible with that count. + +## Slurm rejects the submission + +If `sbatch` reports an invalid account/partition combination, check `submission.account` and `submission.partition` in the run file as well as the selected profile defaults. An explicit run-file value takes precedence. A dry-run renders the request but does not verify account authorization with Slurm. On the submit host, inspect associations and partition availability, then validate a candidate without creating a job: + +```bash +sacctmgr -nP show assoc where user="$USER" format=Account,Partition,QOS +sinfo --all --noheader --format='%P|%a|%G' +sbatch --test-only --account="$ACCOUNT" --partition="$PARTITION" \ + --nodes=1 --gres=gpu:"$GPU_COUNT" --time="$TIME_LIMIT" --wrap=true +``` + +An association with an empty partition field does not guarantee access to every partition. Match the test-only request to the run's resources; if access remains unclear, ask the cluster administrator. If the CLI reports that the submission outcome is uncertain, inspect queue and accounting state before retrying. A rejected or uncertain submission can leave initialized run state without a recorded attempt. + +## Image registration fails + +- Confirm the image-build partition accepts CPU jobs. +- Confirm Enroot is available on the allocated node. +- Use Enroot 3.5 or later for OCI import and existing SQSH inspection. The CLI passes the resolved digest as the registry manifest reference; it does not require Enroot 4.0's `tag@digest` URI syntax. +- Use a supported official `vllm/vllm-openai:vX.Y.Z` tag, a digest-qualified credential-free custom OCI reference, or an existing absolute SQSH path. The official tag path also needs Docker Hub access from the submit host. +- If version verification fails, check that the imported image contains the requested vLLM release; no alias is published on a mismatch. +- Confirm the workspace path contains no whitespace, colon, comma, or backslash that Enroot mount syntax cannot represent. +- While a lifecycle job is active, inspect it under `/images/.tmp/jobs` and read its Slurm log. + +An interrupted or unknown submission can retain a temporary lifecycle job for investigation, but does not publish an unverified alias. Normal success and known failure paths clean temporary lifecycle state. + +## Dry-run cannot find an image + +Run `data-designer slurm image ls` against the same profile and cluster. Aliases are profile-workspace scoped. If the file changed after registration, add and inspect it again; planning rehashes registered SQSH content. + +## A run is absent from `squeue` + +Run `data-designer slurm status `. Completed jobs leave the active queue and may appear in `sacct` after a delay. Status reports accounting lag without inventing a terminal result. + +## Model readiness times out + +Inspect the attempt readiness record and Slurm stderr under the managed run directory. Confirm model files and mounts are visible, the serving image has the expected vLLM runtime, tensor parallelism matches available GPUs, and multi-node traffic is permitted. Increase `startup_timeout` only after identifying a normal slow startup rather than a failed launch. If increasing `distributed_init_timeout`, keep it less than or equal to `startup_timeout`. + +## A shard failed + +Use `status` to find the failed task and attempt. Batch stderr identifies the allocation runtime phase and exit status; synchronous step failures also report their durable stdout and stderr paths. Preview a sparse retry before submission: + +```bash +data-designer slurm retry --task-id --dry-run +``` + +Do not edit persisted manifests. Conflicts preserve scheduler and provenance evidence for operator review. + +## Output is incomplete + +Check that every shard has a `winner.json`, then follow its `candidate_manifest.path` to the canonical `output-manifest.json`. Do not count files by recursively scanning attempts. A shard with at least one generated record can become a winner even if it produced fewer records than requested; a shard with no records cannot. Retry shards without winners, then use `merge` with the original run directory. The merged dataset contains the actual winning records and may be smaller than the requested record count. + +## Scratch files are gone + +This is expected. Runtime extraction, client overlays, caches, and Enroot state live in allocation-local scratch and are cleaned on every terminal path. Use durable Slurm logs, attempt state, readiness records, and candidate results beneath `workspace_root` for diagnosis. diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py index c3fc1195f..09c908474 100644 --- a/packages/data-designer-config/src/data_designer/config/__init__.py +++ b/packages/data-designer-config/src/data_designer/config/__init__.py @@ -32,6 +32,11 @@ from data_designer.config.config_builder import DataDesignerConfigBuilder # noqa: F401 from data_designer.config.custom_column import custom_column_generator # noqa: F401 from data_designer.config.data_designer_config import DataDesignerConfig # noqa: F401 + from data_designer.config.errors import ( # noqa: F401 + InvalidConfigError, + InvalidFileFormatError, + InvalidFilePathError, + ) from data_designer.config.mcp import ( # noqa: F401 LocalStdioMCPProvider, MCPProvider, @@ -160,6 +165,10 @@ "custom_column_generator": (f"{_MOD_BASE}.custom_column", "custom_column_generator"), # data_designer_config "DataDesignerConfig": (f"{_MOD_BASE}.data_designer_config", "DataDesignerConfig"), + # errors + "InvalidConfigError": (f"{_MOD_BASE}.errors", "InvalidConfigError"), + "InvalidFileFormatError": (f"{_MOD_BASE}.errors", "InvalidFileFormatError"), + "InvalidFilePathError": (f"{_MOD_BASE}.errors", "InvalidFilePathError"), # mcp "LocalStdioMCPProvider": (_MOD_MCP, "LocalStdioMCPProvider"), "MCPProvider": (_MOD_MCP, "MCPProvider"), diff --git a/packages/data-designer-config/src/data_designer/config/config_builder.py b/packages/data-designer-config/src/data_designer/config/config_builder.py index f69e63706..08a1e1b17 100644 --- a/packages/data-designer-config/src/data_designer/config/config_builder.py +++ b/packages/data-designer-config/src/data_designer/config/config_builder.py @@ -8,6 +8,7 @@ import logging from pathlib import Path +import yaml from pydantic import model_validator from pygments import highlight from pygments.formatters import HtmlFormatter @@ -24,7 +25,13 @@ ) from data_designer.config.data_designer_config import DataDesignerConfig from data_designer.config.default_model_settings import get_default_model_configs -from data_designer.config.errors import BuilderConfigurationError, BuilderSerializationError, InvalidColumnTypeError +from data_designer.config.errors import ( + BuilderConfigurationError, + BuilderSerializationError, + InvalidColumnTypeError, + InvalidFileFormatError, + InvalidFilePathError, +) from data_designer.config.exportable_config import ExportableConfigBase from data_designer.config.mcp import ToolConfig from data_designer.config.models import ModelConfig, load_model_configs @@ -111,11 +118,19 @@ def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self: Raises: ValueError: If the config format is invalid. ValidationError: If the builder config loaded from the config is invalid. + InvalidFilePathError: If a local config path cannot be read. + InvalidFileFormatError: If a local or inline config contains malformed YAML or JSON. """ if isinstance(config, BuilderConfig): builder_config = config else: - json_config = json.loads(serialize_data(smart_load_yaml(config))) + try: + loaded_config = smart_load_yaml(config) + except OSError as e: + raise InvalidFilePathError(f"Failed to load builder config: {e}") from e + except (UnicodeError, yaml.YAMLError) as e: + raise InvalidFileFormatError(f"Failed to parse builder config: {e}") from e + json_config = json.loads(serialize_data(loaded_config)) # Normalize shorthand DataDesignerConfig into full BuilderConfig if "columns" in json_config and "data_designer" not in json_config: json_config = {"data_designer": json_config} diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index 8ba5b301e..587084673 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -196,10 +196,15 @@ def _smart_load_yaml_internal(yaml_in: str | Path | dict, *, from_url: bool) -> with open(yaml_in) as file: yaml_out = yaml.safe_load(file) elif isinstance(yaml_in, str): - if not from_url and yaml_in.endswith((".yaml", ".yml")) and not os.path.isfile(yaml_in): + is_local_config_path = not from_url and yaml_in.lower().endswith((".yaml", ".yml", ".json")) + has_explicit_path_evidence = Path(yaml_in).is_absolute() or ( + ("/" in yaml_in or "\\" in yaml_in) and ": " not in yaml_in and "\n" not in yaml_in and "\r" not in yaml_in + ) + if is_local_config_path and has_explicit_path_evidence: + raise FileNotFoundError(f"File not found: {yaml_in}") + yaml_out = yaml.safe_load(yaml_in) + if is_local_config_path and not isinstance(yaml_out, dict): raise FileNotFoundError(f"File not found: {yaml_in}") - else: - yaml_out = yaml.safe_load(yaml_in) else: raise ValueError( f"'{yaml_in}' is an invalid yaml config format. Valid options are: dict, yaml string, or yaml file path." diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py new file mode 100644 index 000000000..c11efe29a --- /dev/null +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import data_designer.config.errors as config_errors +from data_designer.config import ( + DataDesignerConfigBuilder, + InvalidConfigError, + InvalidFileFormatError, + InvalidFilePathError, +) + + +def test_invalid_config_error_is_publicly_exported() -> None: + assert InvalidConfigError is config_errors.InvalidConfigError + + +def test_builder_file_errors_are_publicly_exported() -> None: + assert InvalidFileFormatError is config_errors.InvalidFileFormatError + assert InvalidFilePathError is config_errors.InvalidFilePathError + + +@pytest.mark.parametrize( + "filename", ["missing.yaml", "missing.json", "missing.YAML", "missing.JsOn", "[broken.JSON", "missing: config.JSON"] +) +def test_from_config_normalizes_missing_file_error(tmp_path: Path, filename: str) -> None: + with pytest.raises(InvalidFilePathError) as exc_info: + DataDesignerConfigBuilder.from_config(str(tmp_path / filename)) + + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + +@pytest.mark.parametrize("extension", ["json", "yaml", "yml"]) +def test_from_config_accepts_inline_config_ending_in_supported_extension(extension: str) -> None: + builder = DataDesignerConfigBuilder.from_config( + f""" +model_configs: + - alias: stub-model + model: stub-model + provider: provider-1 +columns: + - name: category + column_type: sampler + sampler_type: category + params: + values: + - value.{extension}""" + ) + + assert builder.get_column_config("category").params.values == [f"value.{extension}"] + + +def test_from_config_normalizes_undecodable_file_error(tmp_path: Path) -> None: + config_path = tmp_path / "invalid.yaml" + config_path.write_bytes(b"\xff") + + with pytest.raises(InvalidFileFormatError) as exc_info: + DataDesignerConfigBuilder.from_config(config_path) + + assert isinstance(exc_info.value.__cause__, UnicodeDecodeError) + + +@pytest.mark.parametrize("extension", ["json", "yaml", "yml"]) +def test_from_config_normalizes_malformed_config_error(extension: str) -> None: + with pytest.raises(InvalidFileFormatError): + DataDesignerConfigBuilder.from_config(f"data_designer: [/tmp/value.{extension}") diff --git a/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py b/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py index b2cf04a45..0ba28fcb6 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py @@ -52,6 +52,8 @@ class DeploymentTypeEnum(str, Enum): LIBRARY = "library" API = "api" NVIDIA_INTERNAL = "nvidia-internal" + SLURM = "slurm" + SLURM_NVIDIA_INTERNAL = "slurm-nvidia-internal" UNDEFINED = "undefined" diff --git a/packages/data-designer-engine/tests/engine/models/test_telemetry.py b/packages/data-designer-engine/tests/engine/models/test_telemetry.py index dc9b52070..b31180796 100644 --- a/packages/data-designer-engine/tests/engine/models/test_telemetry.py +++ b/packages/data-designer-engine/tests/engine/models/test_telemetry.py @@ -18,12 +18,16 @@ ) -def test_nvidia_internal_deployment_type_uses_schema_version_1_9() -> None: +@pytest.mark.parametrize( + "deployment_type", + (DeploymentTypeEnum.NVIDIA_INTERNAL, DeploymentTypeEnum.SLURM, DeploymentTypeEnum.SLURM_NVIDIA_INTERNAL), +) +def test_supported_deployment_type_uses_schema_version_1_9(deployment_type: DeploymentTypeEnum) -> None: event = InferenceEvent( nemo_source=NemoSourceEnum.DATADESIGNER, task="batch", task_status=TaskStatusEnum.SUCCESS, - deployment_type=DeploymentTypeEnum("nvidia-internal"), + deployment_type=deployment_type, model="test-model", ) payload = build_payload( @@ -32,7 +36,19 @@ def test_nvidia_internal_deployment_type_uses_schema_version_1_9() -> None: ) assert payload["eventSchemaVer"] == "1.9" - assert payload["events"][0]["parameters"]["deploymentType"] == "nvidia-internal" + assert payload["events"][0]["parameters"]["deploymentType"] == deployment_type.value + + +@pytest.mark.parametrize( + ("value", "expected"), + (("slurm", DeploymentTypeEnum.SLURM), ("Slurm-nvidia-internal", DeploymentTypeEnum.SLURM_NVIDIA_INTERNAL)), +) +def test_slurm_environment_deployment_type_is_recognized( + monkeypatch: pytest.MonkeyPatch, value: str, expected: DeploymentTypeEnum +) -> None: + monkeypatch.setenv("NEMO_DEPLOYMENT_TYPE", value) + + assert get_nemo_deployment_type() is expected def test_unrecognized_env_deployment_type_defaults_to_undefined(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/packages/data-designer-slurm/LICENSE b/packages/data-designer-slurm/LICENSE new file mode 100644 index 000000000..d67d6cffa --- /dev/null +++ b/packages/data-designer-slurm/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2025 NVIDIA + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/packages/data-designer-slurm/README.md b/packages/data-designer-slurm/README.md new file mode 100644 index 000000000..1c000bd69 --- /dev/null +++ b/packages/data-designer-slurm/README.md @@ -0,0 +1,11 @@ +# data-designer-slurm + +Optional Slurm execution support for Data Designer. + +Install it through the Data Designer extra: + +```bash +pip install "data-designer[slurm]" +``` + +Follow the [Slurm Getting Started guide](https://docs.nvidia.com/nemo/datadesigner/slurm/getting-started) for profile setup, image registration, dry-run validation, submission, status, and result collection. Sanitized configuration examples are available in the [Data Designer repository](https://github.com/NVIDIA-NeMo/DataDesigner/tree/main/packages/data-designer-slurm/examples). diff --git a/packages/data-designer-slurm/examples/benchmark.yaml b/packages/data-designer-slurm/examples/benchmark.yaml new file mode 100644 index 000000000..f3608b15e --- /dev/null +++ b/packages/data-designer-slurm/examples/benchmark.yaml @@ -0,0 +1,21 @@ +schema_version: 1 +name: generator-scaling +base_run: run.yaml +model_aliases: all +concurrency_values: + - 32 + - 64 +deployment_cases: + - name: single-node-groups + deployments: + generator: + nodes: 2 + nodes_per_replica: 1 + - name: two-node-group + deployments: + generator: + nodes: 2 + nodes_per_replica: 2 +analysis: + target_total_records: 1000000 + target_runtime: 4h diff --git a/packages/data-designer-slurm/examples/builder.yaml b/packages/data-designer-slurm/examples/builder.yaml new file mode 100644 index 000000000..f2cf333b0 --- /dev/null +++ b/packages/data-designer-slurm/examples/builder.yaml @@ -0,0 +1,15 @@ +data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + columns: + - name: tone + column_type: sampler + sampler_type: category + params: + values: [friendly, formal, playful] + - name: greeting + column_type: llm-text + prompt: Write one short {{ tone }} greeting. + model_alias: generator diff --git a/packages/data-designer-slurm/examples/colocated-models.yaml b/packages/data-designer-slurm/examples/colocated-models.yaml new file mode 100644 index 000000000..12e1bc433 --- /dev/null +++ b/packages/data-designer-slurm/examples/colocated-models.yaml @@ -0,0 +1,42 @@ +schema_version: 1 +name: colocated-models +builder: + inline: + data_designer: + model_configs: + - alias: writer + model: Qwen/Qwen2.5-0.5B-Instruct + provider: openai + - alias: judge + model: microsoft/Phi-3-mini-4k-instruct + provider: openai + columns: + - name: answer + column_type: llm-text + model_alias: writer + prompt: Write one short greeting. + - name: review + column_type: llm-text + model_alias: judge + prompt: 'Briefly evaluate this greeting: {{ answer }}' +invocation: + num_records: 4 + dataset_name: judged-output + model_concurrency: {writer: 1, judge: 1} +client: + cpus: 4 +deployments: + - model_alias: writer + model: Qwen/Qwen2.5-0.5B-Instruct + resources: {shared_node: models, gpus_per_node: 1} + topology: {tensor_parallel: 1} + server: + type: vllm + image: {name: vllm} + - model_alias: judge + model: microsoft/Phi-3-mini-4k-instruct + resources: {shared_node: models, gpus_per_node: 1} + topology: {tensor_parallel: 1} + server: + type: vllm + image: {name: vllm} diff --git a/packages/data-designer-slurm/examples/profile-catalog.yaml b/packages/data-designer-slurm/examples/profile-catalog.yaml new file mode 100644 index 000000000..add258b4a --- /dev/null +++ b/packages/data-designer-slurm/examples/profile-catalog.yaml @@ -0,0 +1,33 @@ +schema_version: 1 +default_cluster: primary +clusters: + primary: + schema_version: 1 + host_patterns: + - login*.example.test + scheduler: + account: research + partition: gpu + mem_per_gpu: 80G + gpus_per_node: 8 + workspace_root: /workspace/alice/data-designer-slurm + image_build: + partition: cpu + cpus_per_task: 2 + memory: 8G + time_limit: "04:00:00" + gpu_request_mode: gres + lab: + schema_version: 1 + host_patterns: + - lab-login*.example.test + scheduler: + partition: accelerated + gpus_per_node: 4 + workspace_root: /workspace/alice/data-designer-slurm-lab + image_build: + partition: general + cpus_per_task: 2 + memory: 8G + time_limit: "04:00:00" + gpu_request_mode: visible diff --git a/packages/data-designer-slurm/examples/run.yaml b/packages/data-designer-slurm/examples/run.yaml new file mode 100644 index 000000000..71abe7c55 --- /dev/null +++ b/packages/data-designer-slurm/examples/run.yaml @@ -0,0 +1,23 @@ +schema_version: 1 +name: greeting-run +builder: + source: builder.yaml +invocation: + num_records: 8 + dataset_name: greetings + model_concurrency: + generator: 8 +deployments: + - model_alias: generator + served_model_name: example/generator + model: example/generator + server: + type: vllm + image: + name: vllm + startup_timeout: 15m + distributed_init_timeout: 10m + topology: + tensor_parallel: 1 +submission: + time_limit: 01:00:00 diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml new file mode 100644 index 000000000..5d4a08779 --- /dev/null +++ b/packages/data-designer-slurm/pyproject.toml @@ -0,0 +1,61 @@ +[project] +name = "data-designer-slurm" +dynamic = ["version", "dependencies"] +description = "Slurm execution for Data Designer" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +license-files = ["LICENSE"] + +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", +] + +[project.entry-points."data_designer.cli"] +slurm = "data_designer.slurm.cli:create_cli" + +[build-system] +requires = ["hatchling", "uv-dynamic-versioning>=0.7.0"] +build-backend = "hatchling.build" + +[tool.hatch.version] +source = "uv-dynamic-versioning" + +[tool.uv-dynamic-versioning] +vcs = "git" +style = "pep440" +bump = true + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +dependencies = [ + "aiohttp>=3.14.3,<4", + "data-designer=={{ version }}", + "packaging>=25,<27", + "pip>=25,<27", + "pydantic>=2.12,<3", + "pyyaml>=6.0.1,<7", +] + +[tool.hatch.build.targets.wheel] +packages = ["src/data_designer"] + +[tool.ruff] +extend = "../../pyproject.toml" + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.uv] +package = true + +[tool.uv.sources] +data-designer = { workspace = true } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/__init__.py new file mode 100644 index 000000000..0c471e0c7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Slurm execution for Data Designer.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py new file mode 100644 index 000000000..ef4cdbbea --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Safe formatting for normalized Slurm boundary errors.""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping, Sequence +from types import UnionType +from typing import Annotated, Any, Literal, Union, get_args, get_origin + +import yaml +from pydantic import BaseModel, ValidationError + +_ERROR_DESCRIPTIONS = { + "extra_forbidden": "field is not permitted", + "greater_than": "must be greater than the allowed minimum", + "greater_than_equal": "must be at least the allowed minimum", + "less_than": "must be less than the allowed maximum", + "less_than_equal": "must not exceed the allowed maximum", + "literal_error": "must use an allowed value", + "missing": "field is required", + "string_pattern_mismatch": "does not match the required pattern", + "string_too_long": "is longer than allowed", + "string_too_short": "is shorter than allowed", + "value_error": "value is invalid", +} +_SAFE_VALUE_ERROR_MESSAGES = { + **{ + f"Value error, {message}": message + for message in ( + "shared_node and gpus_per_node must be provided together", + "shared_node requires exactly one deployment node", + "gpu_start_index requires shared_node", + "shared_node supports only single-node replicas", + "tensor_parallel must divide the deployment GPU quota", + "shared_node requires at least two deployments in each group", + "shared_node requires one group covering all deployments", + "shared_node requires matching server environment and queue backpressure", + "shared_node GPU quotas exceed the profile GPU budget", + "shared_node requires GRES GPU request mode", + "shared_node requires the same resolved serving image", + "deployment placement and GPU ownership must match authored resources", + ) + }, + "Value error, builder content digest does not match the resolved input": ( + "resolved builder digest does not match its input" + ), + "Value error, image inspection digest does not match the resolved SQSH": ( + "resolved image digest does not match its inspection record" + ), + "Value error, mem_per_gpu requires GRES GPU request mode": "mem_per_gpu requires GRES GPU request mode", + "Value error, resolved model aliases do not match the inline builder": ( + "resolved model aliases do not match the inline builder" + ), + "Value error, resolved deployment aliases must exactly cover Data Designer model aliases": ( + "resolved deployment aliases must exactly cover Data Designer model aliases" + ), +} + + +def format_validation_error( + error: ValidationError, + *, + subject: str, + models: type[BaseModel] | tuple[type[BaseModel], ...], +) -> str: + """Summarize validation without rendering user-controlled values.""" + candidates = models if isinstance(models, tuple) else (models,) + model_type = next((model for model in candidates if model.__name__ == error.title), None) + details = error.errors(include_url=False, include_context=False, include_input=False) + summaries = sorted({_format_error_detail(detail, model_type=model_type) for detail in details}) + count = error.error_count() + noun = "error" if count == 1 else "errors" + summary = f": {'; '.join(summaries)}" if summaries else "" + return f"{subject} failed validation ({count} {noun}{summary})" + + +def _format_error_detail(detail: dict[str, Any], *, model_type: type[BaseModel] | None) -> str: + error_type = str(detail["type"]) + message = _SAFE_VALUE_ERROR_MESSAGES.get(str(detail.get("msg"))) + description = message or _ERROR_DESCRIPTIONS.get(error_type, error_type.replace("_", " ")) + location = _format_location(detail.get("loc", ()), model_type=model_type) + return f"{location}: {description}" if location else description + + +def _format_location(location: Iterable[object], *, model_type: type[BaseModel] | None) -> str: + if model_type is None: + return "" + parts: list[str] = [] + schemas: tuple[object, ...] = (model_type,) + for segment in location: + schemas = tuple(candidate for schema in schemas for candidate in _expand_schema(schema)) + if any(_is_mapping_schema(schema) for schema in schemas): + break + if isinstance(segment, int): + item_schemas = tuple(item for schema in schemas for item in _sequence_item_schemas(schema, index=segment)) + if not parts or not item_schemas: + break + parts[-1] = f"{parts[-1]}[{segment}]" + schemas = item_schemas + continue + if not isinstance(segment, str): + break + field_schemas = tuple( + field.annotation + for schema in schemas + if isinstance(schema, type) + and issubclass(schema, BaseModel) + and (field := schema.model_fields.get(segment)) is not None + ) + if not field_schemas: + branch_schemas = _tagged_union_schemas(schemas, tag=segment) + if not branch_schemas: + break + schemas = branch_schemas + continue + parts.append(segment) + schemas = field_schemas + return ".".join(parts) + + +def _expand_schema(schema: object) -> tuple[object, ...]: + origin = get_origin(schema) + if origin is Annotated: + return _expand_schema(get_args(schema)[0]) + if origin in (Union, UnionType): + return tuple(candidate for item in get_args(schema) for candidate in _expand_schema(item)) + return (schema,) + + +def _is_mapping_schema(schema: object) -> bool: + origin = get_origin(schema) + candidate = origin or schema + return isinstance(candidate, type) and issubclass(candidate, Mapping) + + +def _tagged_union_schemas(schemas: tuple[object, ...], *, tag: str) -> tuple[object, ...]: + if len(schemas) < 2: + return () + return tuple( + schema + for schema in schemas + if isinstance(schema, type) + and issubclass(schema, BaseModel) + and ( + schema.__name__ == tag + or any(_annotation_contains_literal(field.annotation, value=tag) for field in schema.model_fields.values()) + ) + ) + + +def _annotation_contains_literal(annotation: object, *, value: str) -> bool: + return any( + get_origin(candidate) is Literal and value in get_args(candidate) for candidate in _expand_schema(annotation) + ) + + +def _sequence_item_schemas(schema: object, *, index: int) -> tuple[object, ...]: + origin = get_origin(schema) + candidate = origin or schema + if not isinstance(candidate, type) or not issubclass(candidate, Sequence): + return () + arguments = get_args(schema) + if not arguments: + return () + if origin is tuple and arguments[-1] is not Ellipsis: + return (arguments[index],) if index < len(arguments) else () + return (arguments[0],) + + +def format_parse_error(error: json.JSONDecodeError | yaml.YAMLError) -> str: + """Summarize a parse failure without rendering source text.""" + if isinstance(error, json.JSONDecodeError): + return f"invalid JSON at line {error.lineno}, column {error.colno}" + mark = getattr(error, "problem_mark", None) + if mark is None: + return "invalid YAML" + return f"invalid YAML at line {mark.line + 1}, column {mark.column + 1}" 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 new file mode 100644 index 000000000..52a7ed22f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable benchmark records for Data Designer Slurm.""" + +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, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkRecommendation, + BenchmarkRecommendationKind, + BenchmarkReport, +) + +__all__ = [ + "BenchmarkCompiler", + "BenchmarkCaseResult", + "BenchmarkChildRun", + "BenchmarkManifest", + "BenchmarkOutcome", + "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..9a7241ffc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -0,0 +1,345 @@ +# 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 + 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 = len({index for deployment in plan.deployments for index in deployment.node_indices}) + gpus = nodes * plan.resolved_gpus_per_node + 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"), + **( + {"gpu_start_index": deployment.gpu_start_index} + if deployment.authored.resources.gpu_start_index is not None + else {} + ), + } + 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, ...]: + actual = tuple( + (child.case_id, child.child_run_id, child.child_authored_config.sha256) for child in manifest.children + ) + for legacy in (False, True): + try: + compiled = BenchmarkCompiler._compile(config, base_run, legacy=legacy) + except ValueError: + continue + expected = tuple( + (case.case_id, case.child_run_id, case.child_run_config.compute_sha256()) for case in compiled.cases + ) + if manifest.benchmark_id == compiled.benchmark_id and actual == expected: + return tuple(case.child_run_config.invocation.num_records for case in compiled.cases) + raise BenchmarkManifestMismatchError("benchmark manifest does not match deterministic child expansion") + + +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..4f6f7404c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py @@ -0,0 +1,183 @@ +# 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 + +from pydantic import Field + +from data_designer.slurm.config import ( + AdaptiveRecordPolicy, + ArrayTasksConfig, + BenchmarkDeploymentCase, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, +) +from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 + +_CASE_SEMANTICS_VERSION = 2 # Keep single-task cases distinct from earlier multi-task benchmark records. + + +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: + return BenchmarkCompiler._compile(config, base_run, legacy=False) + + @staticmethod + def _compile( + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + *, + legacy: bool, + ) -> 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 = { + "benchmark": config.model_dump(mode="json"), + "base_run": base_run.model_dump(mode="json"), + } + if not legacy: + identity["case_semantics_version"] = _CASE_SEMANTICS_VERSION + identity_digest = compute_canonical_json_sha256(identity) + 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, + case_id=case_id, + concurrency=concurrency, + selected_aliases=selected_aliases, + deployment_case=deployment_case, + requested_records=resolve_requested_records(config, concurrency), + single_task=not legacy, + ), + ) + ) + + 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, + *, + case_id: str, + concurrency: int, + selected_aliases: tuple[str, ...], + deployment_case: BenchmarkDeploymentCase, + requested_records: int, + single_task: bool, +) -> 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} + ), + } + ) + ) + + updates = { + "name": _bounded_identifier(f"{base_run.name}-{case_id}"), + "invocation": invocation, + "deployments": deployments, + } + if single_task: + updates["array_tasks"] = ArrayTasksConfig() + return DataDesignerSlurmConfig.model_validate(base_run.model_copy(update=updates).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..7e85ea97b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -0,0 +1,330 @@ +# 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.security import redact_sensitive_diagnostic +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] +ProgressReporter = Callable[[str], None] + + +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, + progress: ProgressReporter | 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 + self._progress = progress + + 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] = [] + first_failure: str | None = None + submitted = 0 + reused = 0 + total = len(compiled.cases) + self._report_progress(f"Benchmark manifest saved; checking {total} child runs.") + for index, case in enumerate(compiled.cases, start=1): + 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): + reused += 1 + self._report_progress(f"Benchmark case {index}/{total}: already submitted.") + continue + execution = self._run_service_factory(case.child_run_id).execute( + case.child_run_config, + source_root=child_source_root, + dry_run=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", + ) + submitted += 1 + self._report_progress(f"Benchmark case {index}/{total}: submitted.") + except SlurmServiceError as error: + failures.append(error.code) + reason = ( + "internal child submission error" + if error.code is SlurmServiceErrorCode.INTERNAL + else " ".join(redact_sensitive_diagnostic(str(error)).split())[:256] + ) + detail = f"{case.case_id} (run {case.child_run_id}): {reason}" + first_failure = first_failure or detail + self._report_progress(f"Benchmark case {index}/{total} submission failed: {detail}; continuing.") + except Exception: + failures.append(SlurmServiceErrorCode.INTERNAL) + detail = f"{case.case_id} (run {case.child_run_id}): internal child submission error" + first_failure = first_failure or detail + self._report_progress(f"Benchmark case {index}/{total} submission failed: {detail}; continuing.") + self._report_progress( + f"Benchmark launch: {submitted} submitted, {reused} already submitted, {len(failures)} failed." + ) + if failures: + code = _partial_failure_code(failures) + message = ( + f"Benchmark {manifest.benchmark_id}: {len(failures)} of {total} child runs could not be submitted; " + f"first failure: {first_failure}" + ) + raise SlurmServiceError( + code, + SlurmServiceOperation.RUN_BENCHMARK, + message[:509] + "..." if len(message) > 512 else message, + ) + return manifest + + def _report_progress(self, message: str) -> None: + if self._progress is not None: + try: + self._progress(message) + except Exception: + # A closed progress stream must not change persisted submission state. + pass + + 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..0ed6b35db --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py @@ -0,0 +1,223 @@ +# 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 not in {ClientOutcome.COMPLETE, ClientOutcome.PARTIAL} + or not client_result.actual_records + 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 new file mode 100644 index 000000000..c626266ca --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta +from enum import Enum +from typing import Annotated + +from pydantic import ( + Field, + NonNegativeFloat, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.slurm.contracts import ArtifactReference, ContractRecord, ContractValue, Identifier + + +class BenchmarkChildRun(ContractValue): + case_id: Identifier + child_run_id: Identifier + child_authored_config: ArtifactReference + + @model_validator(mode="after") + def validate_authored_config(self) -> BenchmarkChildRun: + expected_suffix = f"/runs/{self.child_run_id}/authored-config.json" + if not self.child_authored_config.path.endswith(expected_suffix): + raise ValueError("child authored config path must match the child run identity") + return self + + +class BenchmarkManifest(ContractRecord): + """Stable mapping from benchmark cases to ordinary child runs.""" + + benchmark_id: Identifier + benchmark_config: ArtifactReference + children: tuple[BenchmarkChildRun, ...] = Field(min_length=1) + + @model_validator(mode="after") + def validate_children(self) -> BenchmarkManifest: + case_ids = tuple(child.case_id for child in self.children) + run_ids = tuple(child.child_run_id for child in self.children) + if len(case_ids) != len(set(case_ids)): + raise ValueError("benchmark case IDs must be unique") + if len(run_ids) != len(set(run_ids)): + raise ValueError("benchmark child run IDs must be unique") + return self + + +class BenchmarkOutcome(str, Enum): + PENDING = "pending" + ACCOUNTING_LAG = "accounting_lag" + 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}$")] | None = None + requested_records: PositiveInt + actual_records: NonNegativeInt | None = None + boot_seconds: NonNegativeFloat | None = None + generation_seconds: NonNegativeFloat | None = None + wall_seconds: NonNegativeFloat | None = None + rows_per_second: NonNegativeFloat | None = None + request_count: NonNegativeInt | None = None + token_count: NonNegativeInt | None = None + 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 + feasible: bool | None = None + + @model_validator(mode="after") + def validate_metrics(self) -> BenchmarkCaseResult: + if self.actual_records is not None and self.actual_records > self.requested_records: + raise ValueError("benchmark actual_records must not exceed requested_records") + required = ( + self.actual_records, + self.boot_seconds, + self.generation_seconds, + self.wall_seconds, + self.rows_per_second, + self.gpus_per_job, + self.nodes_per_job, + self.gpu_hours_per_job, + self.feasible, + ) + if self.outcome is BenchmarkOutcome.SUCCEEDED and any(value is None for value in required): + raise ValueError("successful benchmark cases require complete timing and feasibility metrics") + if self.outcome is BenchmarkOutcome.SUCCEEDED: + if self.actual_records == 0: + raise ValueError("successful benchmark cases require at least one record") + 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 + + +class BenchmarkRecommendationKind(str, Enum): + PARETO = "pareto" + MINIMUM_JOBS = "minimum_jobs" + MINIMUM_GPU_HOURS = "minimum_gpu_hours" + + +class BenchmarkRecommendation(ContractValue): + kind: BenchmarkRecommendationKind + case_id: Identifier + + +class BenchmarkReport(ContractRecord): + """Atomic point-in-time benchmark analysis output.""" + + benchmark_id: Identifier + analysis_id: Identifier + benchmark_manifest: ArtifactReference + created_at: datetime + cases: tuple[BenchmarkCaseResult, ...] = Field(min_length=1) + recommendations: tuple[BenchmarkRecommendation, ...] = () + + @field_validator("created_at") + @classmethod + def validate_created_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("created_at must be timezone-aware UTC") + return value + + @model_validator(mode="after") + def validate_report(self) -> BenchmarkReport: + case_ids = tuple(case.case_id for case in self.cases) + child_run_ids = tuple(case.child_run_id for case in self.cases) + if len(case_ids) != len(set(case_ids)): + raise ValueError("benchmark report case IDs must be unique") + if len(child_run_ids) != len(set(child_run_ids)): + raise ValueError("benchmark report child run IDs must be unique") + unknown = {recommendation.case_id for recommendation in self.recommendations}.difference(case_ids) + if unknown: + raise ValueError(f"recommendations reference unknown cases: {', '.join(sorted(unknown))}") + recommendable = { + case.case_id for case in self.cases if case.outcome is BenchmarkOutcome.SUCCEEDED and case.feasible is True + } + identities: set[tuple[BenchmarkRecommendationKind, str]] = set() + singleton_kinds: set[BenchmarkRecommendationKind] = set() + for recommendation in self.recommendations: + if recommendation.case_id not in recommendable: + raise ValueError("benchmark recommendations must reference successful feasible cases") + identity = (recommendation.kind, recommendation.case_id) + if identity in identities: + raise ValueError("benchmark recommendations must be unique") + identities.add(identity) + if recommendation.kind is not BenchmarkRecommendationKind.PARETO: + if recommendation.kind in singleton_kinds: + raise ValueError("minimum benchmark recommendation kinds must be unique") + singleton_kinds.add(recommendation.kind) + 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 new file mode 100644 index 000000000..60546f5ae --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -0,0 +1,593 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import re +import sys +from collections.abc import Callable, Mapping +from contextvars import ContextVar +from datetime import UTC, datetime +from enum import Enum +from functools import partial +from pathlib import Path +from typing import NoReturn, TypeVar + +import click +import typer +from pydantic import BaseModel, ValidationError + +from data_designer.slurm.cli_benchmark import create_benchmark_app +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildRequest, + ServingImageInspection, + SlurmConfigLoadError, + load_run_config, +) +from data_designer.slurm.contracts import canonical_json +from data_designer.slurm.images.records import RegisteredImage, validate_oci_source_for_lifecycle +from data_designer.slurm.images.vllm_source import ( + VllmSourceResolutionError, + VllmTagNotFoundError, + is_versioned_vllm_tag, + resolve_versioned_vllm_source, +) +from data_designer.slurm.security import redact_sensitive_diagnostic +from data_designer.slurm.services import ( + SlurmRetryExecution, + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, + create_slurm_image_service, + create_slurm_profile_service, + create_slurm_run_service, +) + +_ResultT = TypeVar("_ResultT") +_EXIT_CODES = { + SlurmServiceErrorCode.INVALID_REQUEST: 2, + SlurmServiceErrorCode.NOT_FOUND: 3, + SlurmServiceErrorCode.CONFLICT: 4, + SlurmServiceErrorCode.UNAVAILABLE: 5, + SlurmServiceErrorCode.INTERNAL: 1, +} + + +class _RetryResumeMode(str, Enum): + NEVER = "never" + ALWAYS = "always" + IF_POSSIBLE = "if_possible" + + +class _OutputFormat(str, Enum): + HUMAN = "human" + JSON = "json" + + +_ACTIVE_OUTPUT_FORMAT: ContextVar[_OutputFormat] = ContextVar("slurm_output_format", default=_OutputFormat.HUMAN) + + +app = typer.Typer( + name="slurm", + help="Run Data Designer workloads on Slurm", + no_args_is_help=True, +) +image_app = typer.Typer(help="Manage verified Slurm images", no_args_is_help=True) +profile_app = typer.Typer(help="Initialize and validate Slurm profiles", no_args_is_help=True) +app.add_typer(image_app, name="image") +app.add_typer(profile_app, name="profile") + + +@app.callback() +def slurm_callback( + ctx: typer.Context, + output: _OutputFormat = typer.Option( + _OutputFormat.HUMAN, + "--output", + envvar="DATA_DESIGNER_SLURM_OUTPUT", + help="Output format: human or json", + ), +) -> None: + """Choose readable output by default or stable JSON for automation.""" + ctx.obj = output + token = _ACTIVE_OUTPUT_FORMAT.set(_OutputFormat(output)) + ctx.call_on_close(lambda: _ACTIVE_OUTPUT_FORMAT.reset(token)) + + +@app.command("execute") +def execute_command( + run_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"), + dry_run: bool = typer.Option(False, "--dry-run"), +) -> None: + """Prepare and submit one authored run.""" + operation = SlurmServiceOperation.EXECUTE_RUN + + def execute() -> BaseModel: + _emit_progress("Validating and rendering run..." if dry_run else "Preparing and submitting run...") + config = load_run_config(run_file) + service = create_slurm_run_service(profile_file=profile_file, cluster=cluster, progress=_emit_progress) + return service.execute(config, source_root=run_file.resolve().parent, dry_run=dry_run) + + _emit_result(_invoke(operation, execute)) + + +@app.command("status") +def status_command( + run_id: str = typer.Argument(..., help="Managed Data Designer run ID"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Reconcile scheduler observations and show persisted run status.""" + operation = SlurmServiceOperation.STATUS_RUN + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).status(run_id), + ) + _emit_result(result) + + +@app.command("cancel") +def cancel_command( + run_id: str = typer.Argument(..., help="Managed Data Designer run ID"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Request cancellation of active jobs.""" + operation = SlurmServiceOperation.CANCEL_RUN + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).cancel(run_id), + ) + _emit_result(result) + + +@app.command("retry") +def retry_command( + run_or_job_id: str = typer.Argument(..., help="Managed run ID or Slurm array job ID"), + task_ids: list[int] | None = typer.Option(None, "--task-id", min=0, help="Array task ID to retry; repeatable"), + resume: _RetryResumeMode = typer.Option(_RetryResumeMode.IF_POSSIBLE, "--resume"), + dry_run: bool = typer.Option(False, "--dry-run"), + force: bool = typer.Option(False, "--force", help="Submit without confirmation"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Retry failed shards from immutable persisted run state.""" + operation = SlurmServiceOperation.RETRY_RUN + shard_ids = None if task_ids is None else tuple(f"shard-{task_id:05d}" for task_id in task_ids) + service = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster), + ) + + if not dry_run and not force: + planned = _invoke( + operation, + partial( + service.retry, + run_or_job_id, + shard_ids=shard_ids, + resume=resume.value, + dry_run=True, + ), + ) + typer.echo( + f"Retry {', '.join(planned.shard_ids)} with resume={planned.effective_resume_mode}", + err=True, + ) + try: + confirmed = click.confirm("Submit this retry?", default=False, err=True) + except click.Abort: + typer.echo(err=True) + _fail( + SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "interactive confirmation is unavailable; pass --force or --dry-run", + ) + ) + if not confirmed: + _emit_value({"operation": operation.value, "state": "declined"}) + return + shard_ids = planned.shard_ids + resume = _RetryResumeMode(planned.effective_resume_mode) + if not dry_run: + _emit_progress("Preparing and submitting retry...") + result = _invoke( + operation, + partial( + service.retry, + run_or_job_id, + shard_ids=shard_ids, + resume=resume.value, + dry_run=dry_run, + ), + ) + _emit_result(result) + + +@app.command("merge") +def merge_command( + run_id: str | None = typer.Argument(None, help="Managed run ID to collect"), + input_path: Path | None = typer.Option(None, "--input-path", file_okay=False), + output_path: Path | None = typer.Option(None, "--output-path", file_okay=False), + num_partitions: int | None = typer.Option( + None, "--num-partitions", min=1, help="Parquet partitions (default: 1 for a new collection)" + ), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Submit winner-driven collection as a zero-GPU Slurm job.""" + operation = SlurmServiceOperation.COLLECT_RUN + if (run_id is None) == (input_path is None): + _fail( + SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "specify a run ID or --input-path", + ) + ) + _emit_progress("Preparing and submitting collection job...") + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).collect( + run_id if run_id is not None else input_path, + destination=output_path, + num_partitions=num_partitions, + ), + ) + _emit_result(result) + + +@profile_app.command("init") +def profile_init_command( + workspace_root: Path = typer.Option(..., "--workspace-root", file_okay=False), + image_build_partition: str = typer.Option(..., "--image-build-partition"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str = typer.Option("default", "--cluster"), + account: str | None = typer.Option(None, "--account"), + partition: str | None = typer.Option(None, "--partition", help="GPU partition or comma-separated partitions"), + host_pattern: list[str] | None = typer.Option(None, "--host-pattern"), +) -> None: + """Create a safe portable starter profile without overwriting.""" + operation = SlurmServiceOperation.INIT_PROFILE + result = _invoke( + operation, + lambda: create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=workspace_root, + image_build_partition=image_build_partition, + cluster=cluster, + account=account, + partition=partition, + host_patterns=tuple(host_pattern or ()), + ), + ) + _emit_result(result) + + +@profile_app.command("validate") +def profile_validate_command( + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Validate strict loading, cluster selection, workspace, and Slurm facts.""" + operation = SlurmServiceOperation.VALIDATE_PROFILE + result = _invoke( + operation, + lambda: create_slurm_profile_service(profile_file=profile_file, cluster=cluster).validate(), + ) + _emit_result(result) + + +@image_app.command("add") +def image_add_command( + source: str = typer.Argument(..., help="Digest-qualified OCI, official versioned vLLM tag, or absolute SQSH path"), + kind: str = typer.Option("serving", "--kind", help="Image role: serving or client"), + name: str | None = typer.Option(None, "--name", help="Alias to register (derived from the source by default)"), + replace: bool = typer.Option(False, "--replace", help="Allow replacement of an existing alias"), + follow_logs: bool | None = typer.Option( + None, + "--follow-logs/--no-follow-logs", + help="Show bounded image-job logs on stderr (default in an interactive terminal)", + ), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Import or inspect an image and register its alias, reporting progress in terminals.""" + operation = SlurmServiceOperation.ADD_IMAGE + + def add() -> BaseModel: + resolved_source = source + if is_versioned_vllm_tag(source): + if kind != "serving": + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "versioned vLLM images require --kind serving", + ) + _emit_progress("Resolving vLLM image tag...") + try: + resolved_source = resolve_versioned_vllm_source(source) + except VllmTagNotFoundError as error: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, str(error)) from None + except VllmSourceResolutionError as error: + raise SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, operation, str(error)) from None + if ( + not resolved_source.endswith(".sqsh") + and re.fullmatch(r"[^\s]+@sha256:[0-9a-f]{64}", resolved_source) is None + ): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "OCI image source must be digest-qualified as name@sha256:", + ) + try: + validate_oci_source_for_lifecycle(resolved_source) + except ValueError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "OCI image source must be a credential-free registry reference without a scheme", + ) from None + request = ImageBuildRequest(name=name or _derive_image_name(resolved_source), kind=kind, source=resolved_source) + _emit_progress("Preparing image registration...") + return create_slurm_image_service( + profile_file=profile_file, + cluster=cluster, + progress=_emit_progress, + logs=_emit_image_log if (follow_logs if follow_logs is not None else _progress_enabled()) else None, + ).add(request, replace=replace) + + _emit_registered_image(_invoke(operation, add), registered=True) + + +@image_app.command("ls") +def image_list_command( + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """List registered image aliases.""" + operation = SlurmServiceOperation.LIST_IMAGES + images = _invoke( + operation, + lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).list(), + ) + if _output_format() is _OutputFormat.JSON: + _emit_value([image.model_dump(mode="json") for image in images]) + elif not images: + typer.echo("No registered images.") + else: + typer.echo("\n\n".join("\n".join(_image_list_lines(image)) for image in images)) + + +@image_app.command("info") +def image_info_command( + name: str = typer.Argument(..., help="Registered image alias"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Show one registered image alias.""" + operation = SlurmServiceOperation.GET_IMAGE + result = _invoke( + operation, + lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).get(name), + ) + _emit_registered_image(result) + + +@image_app.command("rm") +def image_remove_command( + name: str = typer.Argument(..., help="Registered image alias to remove"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), + delete_artifact: bool = typer.Option( + False, "--delete-artifact", help="Also delete an unshared SQSH imported by this workspace" + ), +) -> None: + """Unregister one alias; optionally delete its package-owned SQSH artifact.""" + operation = SlurmServiceOperation.REMOVE_IMAGE + result = _invoke( + operation, + lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).remove( + name, delete_artifact=delete_artifact + ), + ) + if _output_format() is _OutputFormat.JSON: + _emit_result(result) + else: + typer.echo(f"Removed image alias: {_human_scalar(result.name)}") + outcome = "deleted" if delete_artifact else "retained" + typer.echo(f"SQSH artifact {outcome}: {_human_scalar(result.path)}") + + +def _invoke(operation: SlurmServiceOperation, call: Callable[[], _ResultT]) -> _ResultT: + try: + return call() + except SlurmServiceError as error: + _fail(error) + except SlurmConfigLoadError as error: + message = _bounded_error_message(str(error), fallback="invalid command input") + _fail(SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, message)) + except ValidationError: + _fail(SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, "invalid command input")) + except Exception: + _fail( + SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + f"{operation.value.replace('_', ' ')} failed", + ) + ) + + +def _fail(error: SlurmServiceError) -> NoReturn: + if _output_format() is _OutputFormat.JSON: + _emit_json( + { + "error": { + "code": error.code.value, + "message": str(error), + "operation": error.operation.value, + } + }, + err=True, + ) + else: + typer.echo(f"Error: {error}", err=True) + typer.echo(f"Operation: {error.operation.value.replace('_', ' ')}", err=True) + raise typer.Exit(_EXIT_CODES[error.code]) + + +def _emit_result(result: BaseModel) -> None: + if _output_format() is _OutputFormat.HUMAN and isinstance(result, SlurmRunExecution | SlurmRetryExecution): + _emit_value(result.model_dump(mode="json", exclude_none=True)) + else: + _emit_value(result.model_dump(mode="json")) + + +def _emit_registered_image(result: BaseModel, *, registered: bool = False) -> None: + if _output_format() is _OutputFormat.JSON or not isinstance(result, RegisteredImage): + _emit_result(result) + return + typer.echo("\n".join(_registered_image_lines(result, registered=registered))) + + +def _registered_image_lines(result: RegisteredImage, *, registered: bool = False) -> list[str]: + alias_label = "Registered image alias" if registered else "Image alias" + lines = [ + f"{alias_label}: {_human_scalar(result.name)}", + f"SQSH path: {_human_scalar(result.path)}", + f"SQSH sha256: {_human_scalar(result.sqsh_sha256)}", + ] + if result.source_oci_digest is not None: + lines.append(f"Source OCI digest: {_human_scalar(result.source_oci_digest)}") + lines.extend( + ( + "Image details:", + *_human_lines(result.inspection.inspection.model_dump(mode="json"), indent=2), + ) + ) + return lines + + +def _image_list_lines(image: RegisteredImage) -> list[str]: + details = image.inspection.inspection + if isinstance(details, ServingImageInspection): + runtime = f"vLLM {details.runtime_version}" + else: + assert isinstance(details, ClientImageInspection) + runtime = f"Python {details.python_version}" + summary = f"{image.name} ({details.kind.value}, {runtime})" + return [ + " ".join(_human_scalar(summary).splitlines()), + f" SQSH: {_human_scalar(image.path)}", + ] + + +def _emit_value(value: object, *, err: bool = False) -> None: + if _output_format() is _OutputFormat.JSON: + _emit_json(value, err=err) + else: + typer.echo("\n".join(_human_lines(value)), err=err) + + +def _output_format() -> _OutputFormat: + context = click.get_current_context(silent=True) + while context is not None: + if isinstance(context.obj, _OutputFormat): + return context.obj + context = context.parent + return _ACTIVE_OUTPUT_FORMAT.get() + + +def _human_lines(value: object, *, indent: int = 0) -> list[str]: + prefix = " " * indent + if isinstance(value, Mapping): + if not value: + return [f"{prefix}(none)"] + lines: list[str] = [] + for key, item in value.items(): + label = str(key).replace("_", " ").capitalize() + if isinstance(item, Mapping | list | tuple) or isinstance(item, str) and "\n" in item: + lines.append(f"{prefix}{label}:") + lines.extend(_human_lines(item, indent=indent + 2)) + else: + lines.append(f"{prefix}{label}: {_human_scalar(item)}") + return lines + if isinstance(value, list | tuple): + if not value: + return [f"{prefix}(none)"] + lines = [] + for index, item in enumerate(value, start=1): + if isinstance(item, Mapping | list | tuple): + lines.append(f"{prefix}{index}.") + lines.extend(_human_lines(item, indent=indent + 2)) + else: + lines.append(f"{prefix}- {_human_scalar(item)}") + return lines + if isinstance(value, str) and "\n" in value: + return [f"{prefix}{_human_scalar(line)}" for line in value.rstrip("\n").splitlines()] + return [f"{prefix}{_human_scalar(value)}"] + + +def _human_scalar(value: object) -> str: + if value is None: + return "not set" + if isinstance(value, bool): + return "yes" if value else "no" + return redact_sensitive_diagnostic(str(value)) + + +def _emit_json(value: object, *, err: bool = False) -> None: + typer.echo(canonical_json(value).decode("utf-8"), err=err) + + +def _emit_progress(message: str) -> None: + try: + if _progress_enabled(): + typer.echo(f"[{_utc_timestamp()}] {message}", err=True) + except Exception: + # Progress is optional and must not change the submission outcome. + pass + + +def _emit_image_log(message: str) -> None: + try: + typer.echo(f"[{_utc_timestamp()}] {message}", err=True) + except Exception: + # Log display is optional and must not change the submission outcome. + pass + + +def _progress_enabled() -> bool: + return sys.stderr.isatty() + + +def _utc_timestamp() -> str: + return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _derive_image_name(source: str) -> str: + if source.endswith(".sqsh"): + candidate = Path(source).stem + else: + candidate = source.rpartition("@sha256:")[0].rsplit("/", maxsplit=1)[-1].replace(":", "-") + return re.sub(r"[^A-Za-z0-9._-]+", "-", candidate).strip("-._") or "image" + + +def _bounded_error_message(message: str, *, fallback: str) -> str: + sanitized = "".join(" " if ord(character) < 32 or ord(character) == 127 else character for character in message) + if not sanitized: + return fallback + return sanitized if len(sanitized) <= 512 else f"{sanitized[:509]}..." + + +app.add_typer(create_benchmark_app(_invoke, _emit_result, _emit_progress), 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..5e8d84c21 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py @@ -0,0 +1,118 @@ +# 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, resolve_profile +from data_designer.slurm.services import create_slurm_benchmark_service +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation + +_Invoke = Callable[[SlurmServiceOperation, Callable[[], BaseModel]], BaseModel] +_Emit = Callable[[BaseModel], None] +_Progress = Callable[[str], None] + + +def _run_benchmark( + benchmark_file: Path, + *, + profile_file: Path | None, + cluster: str | None, + force: bool, + progress: _Progress, +) -> BaseModel: + progress("Preparing benchmark launch...") + config = load_benchmark_config(benchmark_file) + service = create_slurm_benchmark_service(profile_file=profile_file, cluster=cluster, progress=progress) + 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: + normalized = benchmark.rstrip("/") + benchmark_id = Path(normalized).name + if normalized != benchmark_id: + selected = resolve_profile(profile_file=profile_file, cluster=cluster) + requested = Path(normalized).expanduser().resolve() + expected = (Path(selected.profile.workspace_root) / "benchmarks" / benchmark_id).resolve() + if requested != expected: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark directory is outside the selected profile's benchmarks directory", + ) + 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, progress: _Progress) -> 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, + progress=progress, + ), + ) + ) + + @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/client/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py new file mode 100644 index 000000000..da4cf64fa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Semantic client records shared with Slurm state consumers.""" + +from __future__ import annotations + +from data_designer.slurm.client.records import ( + ClientEnvironmentManifest, + ClientEnvironmentOutcome, + ClientErrorCode, + ClientInstallerOutcome, + ClientOutcome, + ClientPluginEntryPoint, + ClientProgress, + ClientProgressPhase, + ClientResult, +) + +__all__ = [ + "ClientEnvironmentManifest", + "ClientEnvironmentOutcome", + "ClientErrorCode", + "ClientInstallerOutcome", + "ClientOutcome", + "ClientPluginEntryPoint", + "ClientProgress", + "ClientProgressPhase", + "ClientResult", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py b/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py new file mode 100644 index 000000000..aee3843d9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py @@ -0,0 +1,339 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve immutable client dependency locks before Slurm submission.""" + +from __future__ import annotations + +import hashlib +import os +import stat +import subprocess +import sys +import tempfile +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path + +from packaging.utils import canonicalize_name, parse_wheel_filename + +from data_designer.slurm.client.filesystem import compute_file_sha256, read_regular_bytes +from data_designer.slurm.client.records import ClientErrorCode +from data_designer.slurm.config import ClientDependencies +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution, validate_absolute_path +from data_designer.slurm.planning import LockedPackage, ResolvedClientRuntime, ResolvedDependencyLock + +_RESOLVER_VERSION = "pip-pure-wheel-1" +_DEFAULT_INDEX_URL = "https://pypi.org/simple" +_LOCK_SOURCE_DIRECTORY = "inputs" +_DEPENDENCY_DIRECTORY = "dependencies" + +DependencyCommandRunner = Callable[[tuple[str, ...], Path, Mapping[str, str]], None] + + +class ClientDependencyResolutionError(RuntimeError): + """Raised when authored client dependencies cannot produce a safe lock.""" + + +@dataclass(frozen=True, slots=True) +class ResolvedClientDependencies: + """Resolved lock plus source artifacts for atomic run initialization.""" + + lock: ResolvedDependencyLock + wheel_sources: tuple[Path, ...] + lock_source: Path | None = None + + +class ClientDependencyResolver: + """Resolve pure wheels against the immutable native client inventory.""" + + def __init__( + self, + *, + command_runner: DependencyCommandRunner | None = None, + environ: Mapping[str, str] | None = None, + python_executable: str | None = None, + ) -> None: + self._command_runner = command_runner or _run_pip_download + self._environ = os.environ if environ is None else environ + self._python_executable = sys.executable if python_executable is None else python_executable + + @contextmanager + def resolve( + self, + dependencies: ClientDependencies, + client_runtime: ResolvedClientRuntime, + *, + run_root: str | Path, + source_root: str | Path, + ) -> Iterator[ResolvedClientDependencies]: + """Yield one verified resolution while temporary wheel sources remain live.""" + if not isinstance(dependencies, ClientDependencies): + raise ClientDependencyResolutionError("client dependencies are invalid") + try: + normalized_run_root = Path(validate_absolute_path(Path(run_root).as_posix())) + normalized_source_root = Path(validate_absolute_path(Path(source_root).as_posix())) + except ValueError as error: + raise ClientDependencyResolutionError("dependency resolution paths are invalid") from error + + with tempfile.TemporaryDirectory(prefix="data-designer-slurm-dependencies-") as temporary: + temporary_root = Path(temporary).resolve() + try: + if dependencies.requirements is None: + result = self._load_lock( + dependencies, + client_runtime, + normalized_run_root, + normalized_source_root, + ) + else: + result = self._resolve_requirements( + dependencies, + client_runtime, + normalized_run_root, + temporary_root, + ) + except ClientDependencyResolutionError: + raise + except OSError as error: + raise ClientDependencyResolutionError("client dependency artifacts are unavailable") from error + yield result + + def _resolve_requirements( + self, + dependencies: ClientDependencies, + client_runtime: ResolvedClientRuntime, + run_root: Path, + temporary_root: Path, + ) -> ResolvedClientDependencies: + requirements = dependencies.requirements + if requirements is None: + raise ClientDependencyResolutionError("client dependency requirements are unavailable") + base_distributions = tuple(sorted(client_runtime.distributions, key=lambda item: item.name)) + if not requirements: + return ResolvedClientDependencies( + lock=_build_lock(client_runtime, base_distributions=base_distributions, authored_requirements=()), + wheel_sources=(), + ) + + requirements_path = temporary_root / "requirements.in" + constraints_path = temporary_root / "constraints.txt" + wheelhouse = temporary_root / "wheels" + wheelhouse.mkdir(mode=0o700) + requirements_path.write_text("".join(f"{item}\n" for item in requirements), encoding="utf-8") + constraints_path.write_text( + "".join(f"{item.name}=={item.version}\n" for item in base_distributions), + encoding="utf-8", + ) + command = ( + self._python_executable, + "-m", + "pip", + "download", + "--disable-pip-version-check", + "--no-input", + f"--index-url={_DEFAULT_INDEX_URL}", + "--only-binary=:all:", + "--platform=any", + "--implementation=py", + "--abi=none", + f"--python-version={client_runtime.python_version.rsplit('.', maxsplit=1)[0]}", + f"--dest={wheelhouse.as_posix()}", + f"--requirement={requirements_path.as_posix()}", + f"--constraint={constraints_path.as_posix()}", + ) + environment = self._resolution_environment(dependencies) + try: + self._command_runner(command, temporary_root, environment) + except ClientDependencyResolutionError: + raise + except Exception as error: + raise ClientDependencyResolutionError("client dependency resolution failed") from error + + packages, sources = _load_downloaded_wheels(wheelhouse, run_root, base_distributions) + try: + lock = _build_lock( + client_runtime, + base_distributions=base_distributions, + authored_requirements=tuple(requirements), + overlay_packages=packages, + ) + except ValueError as error: + raise ClientDependencyResolutionError("resolved client dependencies are incompatible") from error + return ResolvedClientDependencies(lock=lock, wheel_sources=sources) + + def _load_lock( + self, + dependencies: ClientDependencies, + client_runtime: ResolvedClientRuntime, + run_root: Path, + source_root: Path, + ) -> ResolvedClientDependencies: + lock_file = dependencies.lock_file + if lock_file is None: + raise ClientDependencyResolutionError("client dependency lock is unavailable") + source = source_root / lock_file + try: + content = read_regular_bytes(source, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) + supplied = ResolvedDependencyLock.model_validate_json(content, strict=True) + except Exception as error: + raise ClientDependencyResolutionError("client dependency lock is invalid") from error + if ( + supplied.client_runtime_sha256 != client_runtime.runtime_sha256 + or supplied.python_abi != client_runtime.python_abi + or supplied.base_distributions != tuple(sorted(client_runtime.distributions, key=lambda item: item.name)) + ): + raise ClientDependencyResolutionError("client dependency lock targets another runtime") + + packages: list[LockedPackage] = [] + sources: list[Path] = [] + for package in supplied.overlay_packages: + wheel = Path(package.artifact.path) + _verify_locked_wheel(wheel, package) + packages.append( + LockedPackage( + name=package.name, + version=package.version, + artifact=ArtifactReference( + path=(run_root / _DEPENDENCY_DIRECTORY / wheel.name).as_posix(), + sha256=package.artifact.sha256, + ), + ) + ) + sources.append(wheel) + source_reference = ArtifactReference( + path=(run_root / _LOCK_SOURCE_DIRECTORY / source.name).as_posix(), + sha256=hashlib.sha256(content).hexdigest(), + ) + try: + lock = ResolvedDependencyLock( + schema_version=1, + resolver_version=supplied.resolver_version, + python_abi=supplied.python_abi, + client_runtime_sha256=supplied.client_runtime_sha256, + authored_requirements=supplied.authored_requirements, + authored_source=lock_file, + source=source_reference, + base_distributions=supplied.base_distributions, + overlay_packages=tuple(packages), + ) + except ValueError as error: + raise ClientDependencyResolutionError("client dependency lock is incompatible") from error + return ResolvedClientDependencies(lock=lock, wheel_sources=tuple(sources), lock_source=source) + + def _resolution_environment(self, dependencies: ClientDependencies) -> dict[str, str]: + credential_names = {reference.environment for reference in dependencies.index_credentials.values()} + for reference in dependencies.index_credentials.values(): + value = self._environ.get(reference.environment) + if value is None or not value: + raise ClientDependencyResolutionError("client dependency index credential is unavailable") + # These credentials belong to the allocation runtime, not submit-side pip. + environment = { + name: value + for name, value in self._environ.items() + if not name.startswith("PIP_") and name not in credential_names + } + environment["PIP_CONFIG_FILE"] = os.devnull + return environment + + +def _build_lock( + client_runtime: ResolvedClientRuntime, + *, + base_distributions: tuple[InstalledDistribution, ...], + authored_requirements: tuple[str, ...], + overlay_packages: tuple[LockedPackage, ...] = (), +) -> ResolvedDependencyLock: + return ResolvedDependencyLock( + schema_version=1, + resolver_version=_RESOLVER_VERSION, + python_abi=client_runtime.python_abi, + client_runtime_sha256=client_runtime.runtime_sha256, + authored_requirements=authored_requirements, + base_distributions=base_distributions, + overlay_packages=overlay_packages, + ) + + +def _load_downloaded_wheels( + wheelhouse: Path, + run_root: Path, + base_distributions: tuple[InstalledDistribution, ...], +) -> tuple[tuple[LockedPackage, ...], tuple[Path, ...]]: + base = {item.name: item.version for item in base_distributions} + packages: list[tuple[LockedPackage, Path]] = [] + names: set[str] = set() + for wheel in sorted(wheelhouse.iterdir(), key=lambda path: path.name): + status = wheel.lstat() + if not stat.S_ISREG(status.st_mode) or wheel.suffix != ".whl": + raise ClientDependencyResolutionError("dependency resolver produced an invalid artifact") + try: + name, version, _, tags = parse_wheel_filename(wheel.name) + except ValueError as error: + raise ClientDependencyResolutionError("dependency resolver produced an invalid wheel") from error + normalized_name = canonicalize_name(name) + normalized_version = str(version) + if not tags or any(tag.abi != "none" or tag.platform != "any" for tag in tags): + raise ClientDependencyResolutionError("client dependencies must resolve to platform-independent wheels") + if normalized_name in base: + if base[normalized_name] != normalized_version: + raise ClientDependencyResolutionError("client dependency conflicts with the base runtime") + continue + if normalized_name in names: + raise ClientDependencyResolutionError("dependency resolver produced duplicate packages") + names.add(normalized_name) + digest = compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) + packages.append( + ( + LockedPackage( + name=normalized_name, + version=normalized_version, + artifact=ArtifactReference( + path=(run_root / _DEPENDENCY_DIRECTORY / wheel.name).as_posix(), + sha256=digest, + ), + ), + wheel, + ) + ) + packages.sort(key=lambda item: item[0].name) + return tuple(item[0] for item in packages), tuple(item[1] for item in packages) + + +def _verify_locked_wheel(wheel: Path, package: LockedPackage) -> None: + try: + name, version, _, tags = parse_wheel_filename(wheel.name) + digest = compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) + except Exception as error: + raise ClientDependencyResolutionError("client dependency artifact is invalid") from error + if ( + not tags + or any(tag.abi != "none" or tag.platform != "any" for tag in tags) + or canonicalize_name(name) != package.name + or str(version) != package.version + or digest != package.artifact.sha256 + ): + raise ClientDependencyResolutionError("client dependency artifact differs from its lock") + + +def _run_pip_download(command: tuple[str, ...], cwd: Path, environment: Mapping[str, str]) -> None: + try: + subprocess.run( + command, + cwd=cwd, + env=environment, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=True, + ) + except (OSError, subprocess.SubprocessError) as error: + raise ClientDependencyResolutionError("client dependency resolution failed") from error + + +__all__ = [ + "ClientDependencyResolutionError", + "ClientDependencyResolver", + "ResolvedClientDependencies", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py new file mode 100644 index 000000000..e588cd11c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -0,0 +1,597 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +import re +import site +import stat +import subprocess +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import cast + +from packaging.tags import interpreter_name, interpreter_version +from packaging.utils import canonicalize_name, parse_wheel_filename + +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ( + compute_file_sha256, + ensure_private_directory, + read_regular_bytes, + replace_private_text, +) +from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution, canonical_json + +DistributionInventory = Callable[[Path | None], tuple[InstalledDistribution, ...]] +CommandRunner = Callable[[tuple[str, ...]], None] +_OVERLAY_INVENTORY_FILENAME = ".data-designer-overlay-files.json" + + +@dataclass(frozen=True) +class PreparedClientEnvironment: + run_id: str + shard_id: str + attempt_id: str + attempt_dir: Path + scratch_root: Path + overlay_path: Path + dependency_lock: ArtifactReference + client_runtime_sha256: str + python_abi: str + installer_outcome: ClientInstallerOutcome + installed_distributions: tuple[InstalledDistribution, ...] + + +@dataclass(frozen=True) +class _BootstrapInputs: + run_id: str + logical_run_root: Path + shard_id: str + attempt_id: str + attempt_dir: Path + scratch_root: Path + client_runtime_sha256: str + python_abi: str + python_executable: Path + runtime: dict[str, object] + dependency_lock: ArtifactReference + + +@dataclass(frozen=True) +class _VerifiedDependencies: + base_distributions: tuple[InstalledDistribution, ...] + overlay_packages: tuple[dict[str, object], ...] + + +class ClientEnvironmentBuilder: + """Prepare one immutable allocation-local package environment.""" + + def __init__( + self, + *, + inventory: DistributionInventory | None = None, + command_runner: CommandRunner | None = None, + ) -> None: + self._inventory = inventory or inspect_distributions + self._command_runner = command_runner or _run_installer + + def prepare( + self, + plan_path: Path, + *, + shard_id: str, + attempt_id: str, + attempt_dir: Path, + scratch_root: Path, + ) -> PreparedClientEnvironment: + """Verify the plan and lock subset needed before plugin-aware imports.""" + inputs = self._load_bootstrap_inputs( + plan_path, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_dir=attempt_dir, + scratch_root=scratch_root, + ) + dependencies = self._verify_dependency_lock(inputs) + overlay_path, installer_outcome, installed = self._prepare_overlay(inputs, dependencies) + return PreparedClientEnvironment( + run_id=inputs.run_id, + shard_id=inputs.shard_id, + attempt_id=inputs.attempt_id, + attempt_dir=inputs.attempt_dir, + scratch_root=inputs.scratch_root, + overlay_path=overlay_path, + dependency_lock=inputs.dependency_lock, + client_runtime_sha256=inputs.client_runtime_sha256, + python_abi=inputs.python_abi, + installer_outcome=installer_outcome, + installed_distributions=installed, + ) + + @staticmethod + def _load_bootstrap_inputs( + plan_path: Path, + *, + shard_id: str, + attempt_id: str, + attempt_dir: Path, + scratch_root: Path, + ) -> _BootstrapInputs: + _validate_input_path(plan_path, "resolved plan") + _validate_input_path(attempt_dir, "attempt directory") + _validate_scratch_root(scratch_root, attempt_dir) + if not re.fullmatch(r"shard-[0-9]{5,}", shard_id): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "shard identifier is invalid") + if not re.fullmatch(r"attempt-[0-9]{4,}", attempt_id) or int(attempt_id.removeprefix("attempt-")) < 1: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "attempt identifier is invalid") + plan_payload = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) + run_id = _require_string(plan_payload, "run_id") + authored_reference = _artifact_reference(_require_object(plan_payload, "authored_config")) + logical_run_root = Path(authored_reference.path).parent + if ( + plan_path.name != "resolved-plan.json" + or logical_run_root.name != run_id + or logical_run_root.parent.name != "runs" + or (logical_run_root / plan_path.name).as_posix() != plan_path.as_posix() + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved plan path is not canonical") + logical_attempt = logical_run_root / "shards" / shard_id / "attempts" / attempt_id + expected_attempt = logical_attempt + if attempt_dir.as_posix() != expected_attempt.as_posix(): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "attempt directory does not match the plan") + ensure_private_directory(attempt_dir) + + client = _require_object(plan_payload, "client") + runtime = _require_object(client, "runtime") + runtime_sha256 = _require_digest(runtime, "runtime_sha256") + runtime_identity = dict(runtime) + runtime_identity.pop("runtime_sha256") + if hashlib.sha256(canonical_json(runtime_identity)).hexdigest() != runtime_sha256: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client runtime fingerprint is invalid") + python_abi = _require_string(runtime, "python_abi") + if python_abi != f"{interpreter_name()}{interpreter_version()}": + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client Python ABI differs from the plan") + python_executable = Path(_require_string(runtime, "python_executable")) + if python_executable.as_posix() != sys.executable: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "client Python executable differs from the plan" + ) + + lock_reference = _artifact_reference(_require_object(client, "dependency_lock")) + if Path(lock_reference.path).as_posix() != (logical_run_root / "dependency-lock.json").as_posix(): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency lock path is not canonical") + return _BootstrapInputs( + run_id=run_id, + logical_run_root=logical_run_root, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_dir=attempt_dir, + scratch_root=scratch_root, + client_runtime_sha256=runtime_sha256, + python_abi=python_abi, + python_executable=python_executable, + runtime=runtime, + dependency_lock=lock_reference, + ) + + def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependencies: + lock_bytes = read_regular_bytes( + Path(inputs.dependency_lock.path), + missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, + ) + if _sha256_bytes(lock_bytes) != inputs.dependency_lock.sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency lock digest differs") + lock = _parse_json_bytes(lock_bytes, ClientErrorCode.DEPENDENCY_CONFLICT) + if _require_digest(lock, "client_runtime_sha256") != inputs.client_runtime_sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency lock targets another runtime") + if _require_string(lock, "python_abi") != inputs.python_abi: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency lock targets another Python ABI") + + expected_base = _parse_distributions(lock.get("base_distributions")) + inspected_base = _parse_distributions(inputs.runtime.get("distributions")) + try: + actual_base = self._inventory(None) + except ClientWorkerError: + raise + except Exception as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "client runtime inventory cannot be verified" + ) from error + if expected_base != inspected_base or actual_base != expected_base: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "client runtime inventory differs from the lock" + ) + + source = lock.get("source") + if source is not None: + source_reference = _artifact_reference(_as_object(source)) + _verify_input_artifact( + source_reference, + inputs.logical_run_root / "inputs", + ) + return _VerifiedDependencies( + base_distributions=expected_base, + overlay_packages=tuple(_as_object_list(lock.get("overlay_packages"))), + ) + + def _prepare_overlay( + self, + inputs: _BootstrapInputs, + dependencies: _VerifiedDependencies, + ) -> tuple[Path, ClientInstallerOutcome, tuple[InstalledDistribution, ...]]: + expected_overlay, wheels = _verify_wheels( + dependencies.overlay_packages, + inputs.logical_run_root / "dependencies", + dependencies.base_distributions, + ) + overlay_path = inputs.scratch_root / "client-env" / "site-packages" + outcome = self._install_overlay(inputs.python_executable, wheels, expected_overlay, overlay_path) + installed = tuple(sorted((*dependencies.base_distributions, *expected_overlay), key=lambda item: item.name)) + return overlay_path, outcome, installed + + def _install_overlay( + self, + python_executable: Path, + wheels: tuple[Path, ...], + expected: tuple[InstalledDistribution, ...], + target: Path, + ) -> ClientInstallerOutcome: + ensure_private_directory(target) + if not wheels: + if any(target.iterdir()): + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay contains unexpected files" + ) + return ClientInstallerOutcome.NOT_REQUIRED + existing = self._inventory(target) + if existing == expected and any(target.iterdir()): + _verify_overlay_files(target) + return ClientInstallerOutcome.REUSED + if any(target.iterdir()): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay contains unexpected files") + try: + self._command_runner( + ( + python_executable.as_posix(), + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-deps", + "--no-index", + "--target", + target.as_posix(), + *(wheel.as_posix() for wheel in wheels), + ) + ) + except (OSError, subprocess.SubprocessError) as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_INSTALL_FAILED, "client dependency installation failed" + ) from error + if self._inventory(target) != expected: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_INSTALL_FAILED, "installed client dependencies differ from the lock" + ) + replace_private_text( + target / _OVERLAY_INVENTORY_FILENAME, + json.dumps(_overlay_file_hashes(target), sort_keys=True, separators=(",", ":")), + ) + return ClientInstallerOutcome.INSTALLED + + +def _overlay_file_hashes(target: Path, *, recorded: dict[str, str] | None = None) -> dict[str, str]: + """Fingerprint installed files and discard bytecode generated after installation.""" + hashes: dict[str, str] = {} + try: + for path in target.rglob("*"): + relative = path.relative_to(target) + if path.name == _OVERLAY_INVENTORY_FILENAME and path.parent == target: + continue + mode = path.lstat().st_mode + if stat.S_ISDIR(mode): + continue + if not stat.S_ISREG(mode): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay contains unsafe files") + if ( + recorded is not None + and "__pycache__" in relative.parts + and path.suffix == ".pyc" + and relative.as_posix() not in recorded + ): + path.unlink() + continue + hashes[relative.as_posix()] = compute_file_sha256(path, missing_code=ClientErrorCode.DEPENDENCY_CONFLICT) + except OSError as error: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay cannot be verified") from error + return hashes + + +def _verify_overlay_files(target: Path) -> None: + marker = target / _OVERLAY_INVENTORY_FILENAME + try: + recorded = json.loads(read_regular_bytes(marker, missing_code=ClientErrorCode.DEPENDENCY_CONFLICT)) + if not isinstance(recorded, dict) or not all( + isinstance(path, str) and isinstance(digest, str) for path, digest in recorded.items() + ): + raise ValueError("invalid overlay inventory") + except (ClientWorkerError, ValueError, UnicodeError) as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay has no valid file inventory" + ) from error + if _overlay_file_hashes(target, recorded=recorded) != recorded: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay files differ from the lock") + + +def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ...]: + """Return an exact immutable distribution inventory.""" + if path is not None: + importlib.invalidate_caches() + search_paths = _active_distribution_search_paths() if path is None else [path.as_posix()] + distributions = importlib.metadata.distributions(path=search_paths) + installed: list[InstalledDistribution] = [] + names: set[str] = set() + for distribution in distributions: + name = canonicalize_name(distribution.metadata["Name"] or "") + version = distribution.version + if not name or not version or name in names: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "installed distribution inventory is invalid") + direct_url = distribution.read_text("direct_url.json") + if direct_url is not None and _is_mutable_direct_url(direct_url): + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, + f"mutable installed distribution {name!r} is forbidden", + ) + names.add(name) + installed.append( + InstalledDistribution( + name=name, + version=version, + provenance_sha256=(None if path is not None or direct_url is None else _direct_url_sha256(direct_url)), + ) + ) + return tuple(sorted(installed, key=lambda item: item.name)) + + +def _active_distribution_search_paths() -> list[str]: + """Exclude user-site and submit-only ``PYTHONPATH`` entries from the inventory.""" + search_paths = sys.path.copy() + if site.ENABLE_USER_SITE: + try: + search_paths = [ + entry + for entry in json.loads( + subprocess.check_output( + (sys.executable, "-s", "-c", "import json,sys; print(json.dumps(sys.path))"), + text=True, + timeout=10, + ) + ) + if entry + ] + except (OSError, subprocess.SubprocessError, ValueError) as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, + "client Python import paths cannot be inspected with user-site disabled", + ) from error + configured_paths = { + _normalize_import_path(entry) for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep) if entry + } + user_site = site.getusersitepackages() + user_site_paths = { + _normalize_import_path(entry) for entry in ([user_site] if isinstance(user_site, str) else user_site) + } + standard_site_paths = {_normalize_import_path(entry) for entry in site.getsitepackages()} + active_paths: list[str] = [] + for entry in search_paths: + normalized_entry = _normalize_import_path(entry) + if normalized_entry in user_site_paths: + continue + if normalized_entry in configured_paths and normalized_entry not in standard_site_paths: + continue + if normalized_entry not in active_paths: + active_paths.append(normalized_entry) + return active_paths + + +def _normalize_import_path(path: str) -> str: + return os.path.normcase(os.path.realpath(path or os.curdir)) + + +def activate_environment(prepared: PreparedClientEnvironment) -> None: + """Activate one verified overlay and isolate Data Designer attempt state.""" + home = prepared.scratch_root / "data-designer-home" + cache = prepared.scratch_root / "cache" + ensure_private_directory(home) + ensure_private_directory(cache) + os.environ["DATA_DESIGNER_HOME"] = home.as_posix() + os.environ["XDG_CACHE_HOME"] = cache.as_posix() + os.environ["PYTHONNOUSERSITE"] = "1" + os.environ["DISABLE_DATA_DESIGNER_PLUGINS"] = "false" + if prepared.overlay_path.as_posix() not in sys.path: + sys.path.append(prepared.overlay_path.as_posix()) + importlib.invalidate_caches() + + +def _validate_input_path(path: Path, label: str) -> None: + if not path.is_absolute() or path != Path(os.path.normpath(path.as_posix())): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, f"{label} path is not canonical") + + +def _validate_scratch_root(scratch_root: Path, attempt_dir: Path) -> None: + _validate_input_path(scratch_root, "allocation scratch") + try: + status = scratch_root.lstat() + except OSError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "allocation scratch is unavailable") from error + if ( + not stat.S_ISDIR(status.st_mode) + or stat.S_ISLNK(status.st_mode) + or status.st_uid != os.geteuid() + or scratch_root == attempt_dir + or scratch_root.is_relative_to(attempt_dir) + or attempt_dir.is_relative_to(scratch_root) + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "allocation scratch is invalid") + ensure_private_directory(scratch_root) + + +def _run_installer(command: tuple[str, ...]) -> None: + subprocess.run(command, check=True, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + +def _verify_wheels( + packages: tuple[dict[str, object], ...], + logical_dependencies_root: Path, + base_distributions: tuple[InstalledDistribution, ...], +) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: + expected: list[InstalledDistribution] = [] + wheels: list[Path] = [] + base_names = {item.name for item in base_distributions} + for package in packages: + name = canonicalize_name(_require_string(package, "name")) + version = _require_string(package, "version") + if name in base_names or name in {item.name for item in expected}: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency distributions overlap") + artifact = _artifact_reference(_require_object(package, "artifact")) + logical_wheel = Path(artifact.path) + if logical_wheel.parent != logical_dependencies_root or logical_wheel.suffix != ".whl": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") + wheel = logical_wheel + if compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != artifact.sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency wheel digest differs") + try: + wheel_name, wheel_version, _, _ = parse_wheel_filename(wheel.name) + except ValueError as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "dependency artifact is not a valid wheel" + ) from error + if canonicalize_name(wheel_name) != name or str(wheel_version) != version: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "dependency wheel identity differs from the lock" + ) + expected.append(InstalledDistribution(name=name, version=version)) + wheels.append(wheel) + sorted_pairs = sorted(zip(expected, wheels, strict=True), key=lambda pair: pair[0].name) + if tuple(item.name for item in expected) != tuple(pair[0].name for pair in sorted_pairs): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency overlay is not sorted") + return tuple(pair[0] for pair in sorted_pairs), tuple(pair[1] for pair in sorted_pairs) + + +def _verify_input_artifact(reference: ArtifactReference, logical_root: Path) -> None: + logical_path = Path(reference.path) + if logical_path.parent != logical_root or logical_path.suffix != ".json": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") + path = logical_path + if compute_file_sha256(path, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != reference.sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency source digest differs") + + +def _parse_distributions(value: object) -> tuple[InstalledDistribution, ...]: + parsed = tuple( + InstalledDistribution( + name=canonicalize_name(_require_string(item, "name")), + version=_require_string(item, "version"), + provenance_sha256=( + None if item.get("provenance_sha256") is None else _require_digest(item, "provenance_sha256") + ), + ) + for item in _as_object_list(value) + ) + names = tuple(item.name for item in parsed) + if names != tuple(sorted(names)) or len(names) != len(set(names)): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency inventory is not sorted and unique") + return parsed + + +def _load_json_object(path: Path, code: ClientErrorCode) -> dict[str, object]: + return _parse_json_bytes(read_regular_bytes(path, missing_code=code), code) + + +def _parse_json_bytes(value: bytes, code: ClientErrorCode) -> dict[str, object]: + try: + return _as_object(json.loads(value)) + except (UnicodeDecodeError, json.JSONDecodeError, TypeError) as error: + raise ClientWorkerError(code, "client JSON artifact is invalid") from error + + +def _as_object(value: object) -> dict[str, object]: + if not isinstance(value, dict) or not all(isinstance(key, str) for key in value): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client JSON object is invalid") + return cast(dict[str, object], value) + + +def _as_object_list(value: object) -> list[dict[str, object]]: + if not isinstance(value, list): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency package list is invalid") + return [_as_object(item) for item in value] + + +def _require_object(value: dict[str, object], key: str) -> dict[str, object]: + try: + return _as_object(value[key]) + except KeyError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "required client plan field is missing") from error + + +def _require_string(value: dict[str, object], key: str) -> str: + item = value.get(key) + if not isinstance(item, str) or not item or any(ord(character) < 32 for character in item): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "required client string field is invalid") + return item + + +def _require_digest(value: dict[str, object], key: str) -> str: + digest = _require_string(value, key) + if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "required client digest is invalid") + return digest + + +def _artifact_reference(value: dict[str, object]) -> ArtifactReference: + try: + return ArtifactReference(path=_require_string(value, "path"), sha256=_require_digest(value, "sha256")) + except ValueError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client artifact reference is invalid") from error + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _is_mutable_direct_url(value: str) -> bool: + try: + direct_url = _as_object(json.loads(value)) + except (json.JSONDecodeError, TypeError): + return True + if "dir_info" in direct_url: + return True + vcs_info = direct_url.get("vcs_info") + if vcs_info is not None: + if not isinstance(vcs_info, dict) or vcs_info.get("vcs") != "git": + return True + commit_id = vcs_info.get("commit_id") + return not isinstance(commit_id, str) or re.fullmatch(r"[0-9a-f]{40}|[0-9a-f]{64}", commit_id) is None + archive_info = direct_url.get("archive_info") + if archive_info is None: + return True + if not isinstance(archive_info, dict): + return True + hashes = archive_info.get("hashes") + digest = None if not isinstance(hashes, dict) else hashes.get("sha256") + return not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None + + +def _direct_url_sha256(value: str) -> str: + try: + payload = _as_object(json.loads(value)) + except (json.JSONDecodeError, TypeError) as error: # pragma: no cover - guarded by mutability validation + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "installed provenance is invalid") from error + return hashlib.sha256(canonical_json(payload)).hexdigest() diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/client/errors.py new file mode 100644 index 000000000..0a90165b1 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/errors.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.client.records import ClientErrorCode + + +class ClientWorkerError(RuntimeError): + """Failure safe to classify at the allocation boundary.""" + + def __init__(self, code: ClientErrorCode, message: str) -> None: + super().__init__(message) + self.code = code + self.redacted_message = message diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py new file mode 100644 index 000000000..ff9b4edf5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -0,0 +1,836 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import logging +import os +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Protocol, cast +from urllib.parse import urlsplit + +from pydantic import ValidationError + +import data_designer.lazy_heavy_imports as lazy +from data_designer.config import ( + DataDesignerConfigBuilder, + LocalStdioMCPProvider, + MCPProvider, + ModelProvider, + PartitionBlock, + ResumeMode, + RunConfig, +) +from data_designer.interface import DataDesigner +from data_designer.slurm.client.environment import PreparedClientEnvironment +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ( + compute_file_sha256, + ensure_private_directory, + publish_private_text, + read_regular_bytes, + replace_private_text, +) +from data_designer.slurm.client.records import ( + ClientEnvironmentManifest, + ClientEnvironmentOutcome, + ClientErrorCode, + ClientOutcome, + ClientPluginEntryPoint, + ClientProgress, + ClientProgressPhase, + ClientResult, +) +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.images import InstalledDistribution +from data_designer.slurm.config.run import LocalStdioMCPProviderConfig, RemoteMCPProviderConfig +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.ports import resolve_allocation_plan +from data_designer.slurm.state import ( + MAXIMUM_CANDIDATE_OUTPUT_FILES, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + compute_candidate_schema_digest, +) + +Clock = Callable[[], datetime] +DataDesignerFactory = Callable[..., DataDesigner] +logger = logging.getLogger(__name__) + + +class _ResultArtifactStorage(Protocol): + @property + def final_dataset_path(self) -> Path: + """Return the directory containing Data Designer's final batch files.""" + + +class CreationResults(Protocol): + dataset_path: Path + requested_num_records: int + actual_num_records: int + early_shutdown: bool | None + requested_resume_mode: ResumeMode | None + effective_resume_mode: ResumeMode | None + + @property + def artifact_storage(self) -> _ResultArtifactStorage: + """Return the storage used by this generation result.""" + + def export(self, path: Path, *, format: str) -> Path: + """Export the generated dataset.""" + + +@dataclass(frozen=True) +class _ExecutionContext: + plan: ResolvedSlurmRunPlan + shard: PlannedShard + builder: DataDesignerConfigBuilder + designer: DataDesigner + requested_resume: ResumeMode + retry_resume: ResumeMode | None + dataset_path: Path + + +@dataclass(frozen=True) +class _ValidatedCreation: + requested_records: int + actual_records: int + dataset_path: Path + expected_path: Path + effective_resume: ResumeMode + early_shutdown: bool + + +class _ProgressWriter: + def __init__( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + clock: Clock, + *, + revision: int = 0, + ) -> None: + self._context = context + self._prepared = prepared + self._clock = clock + self._revision = revision + self._seen_batches: set[Path] = set() + self._pending_batches: set[Path] = set() + self._completed_records = 0 + self._batches_initialized = False + + def required( + self, + phase: ClientProgressPhase, + *, + completed_records: int | None = None, + error_code: ClientErrorCode | None = None, + ) -> None: + self._revision += 1 + progress = ClientProgress( + schema_version=1, + run_id=self._context.plan.run_id, + shard_id=self._context.shard.shard_id, + attempt_id=self._prepared.attempt_id, + revision=self._revision, + updated_at=self._clock(), + phase=phase, + requested_records=self._context.shard.requested_records, + completed_records=completed_records, + error_code=error_code, + ) + replace_private_text(self._prepared.attempt_dir / "client-progress.json", progress.serialize_json()) + + def best_effort( + self, + phase: ClientProgressPhase, + *, + completed_records: int | None = None, + error_code: ClientErrorCode | None = None, + ) -> None: + try: + self.required(phase, completed_records=completed_records, error_code=error_code) + except Exception: + logger.warning("Could not persist client %s progress", phase.value) + + def on_batch_complete(self, path: Path) -> None: + try: + if not self._batches_initialized: + batches = tuple(path.parent.glob("*.parquet")) + completed = sum(lazy.pq.read_metadata(item).num_rows for item in batches) + self._seen_batches = set(batches) + self._completed_records = completed + self._batches_initialized = True + except Exception: + logger.warning("Could not read completed batch metadata for progress") + return + if path not in self._seen_batches: + self._pending_batches.add(path) + for pending_path in tuple(sorted(self._pending_batches)): + try: + count = lazy.pq.read_metadata(pending_path).num_rows + except Exception: + logger.warning("Could not read completed batch metadata for progress") + continue + self._completed_records += count + self._seen_batches.add(pending_path) + self._pending_batches.remove(pending_path) + self.best_effort( + ClientProgressPhase.GENERATING, + completed_records=min(self._completed_records, self._context.shard.requested_records), + ) + + +class ClientWorker: + """Validate and execute one planned Data Designer shard.""" + + def __init__( + self, + *, + data_designer_factory: DataDesignerFactory = DataDesigner, + clock: Clock | None = None, + environment: Mapping[str, str] | None = None, + ) -> None: + self._data_designer_factory = data_designer_factory + self._clock = clock or (lambda: datetime.now(timezone.utc)) + self._environment = os.environ if environment is None else environment + + def preflight( + self, + plan_path: Path, + *, + prepared: PreparedClientEnvironment, + endpoints: Mapping[str, str], + plugins: tuple[ClientPluginEntryPoint, ...], + retry_resume: ResumeMode | None = None, + ) -> ClientEnvironmentManifest: + """Validate packages, plugins, assets, and config without generation.""" + try: + context = self._build_context( + plan_path, + prepared=prepared, + endpoints=endpoints, + retry_resume=retry_resume, + ) + progress = _ProgressWriter(context, prepared, self._clock) + progress.required(ClientProgressPhase.VALIDATING_PLUGINS) + progress.required(ClientProgressPhase.VALIDATING_CONFIG) + context.designer.validate(context.builder) + manifest = ClientEnvironmentManifest.from_prepared( + prepared, + created_at=self._clock(), + outcome=ClientEnvironmentOutcome.READY, + plugins=plugins, + ) + replace_private_text(prepared.attempt_dir / "client-environment.json", manifest.serialize_json()) + return manifest + except ClientWorkerError as error: + self._write_preflight_failure(plan_path, prepared, endpoints, plugins, error) + raise + except Exception as error: + failure = ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "Data Designer configuration is invalid") + self._write_preflight_failure(plan_path, prepared, endpoints, plugins, failure) + raise failure from error + + def run( + self, + plan_path: Path, + *, + prepared: PreparedClientEnvironment, + endpoints: Mapping[str, str], + plugins: tuple[ClientPluginEntryPoint, ...], + retry_resume: ResumeMode | None = None, + ) -> ClientResult: + """Invoke the public Data Designer generation contract and persist its result.""" + context: _ExecutionContext | None = None + progress: _ProgressWriter | None = None + try: + context = self._build_context( + plan_path, + prepared=prepared, + endpoints=endpoints, + retry_resume=retry_resume, + ) + progress = _ProgressWriter(context, prepared, self._clock, revision=2) + self._validate_environment_manifest(context, prepared, plugins) + progress.required(ClientProgressPhase.GENERATING, completed_records=0) + self._prepare_dataset_workspace(context, prepared) + results = self._generate_dataset(context, progress) + try: + creation = self._validate_creation_result(context, prepared, results) + progress.required(ClientProgressPhase.FINALIZING, completed_records=creation.actual_records) + result = self._finalize(context, prepared, results, creation) + except ClientWorkerError: + raise + except Exception as error: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer output is invalid") from error + progress.best_effort( + ClientProgressPhase.COMPLETE, + completed_records=cast(int, result.actual_records), + ) + return result + except ClientWorkerError as error: + if context is not None: + self._write_failure(context, prepared, error, progress=progress) + raise + except (KeyboardInterrupt, SystemExit) as error: + failure = ClientWorkerError(ClientErrorCode.INTERRUPTED, "Data Designer generation was interrupted") + if context is not None: + self._write_failure(context, prepared, failure, progress=progress) + raise failure from error + except Exception as error: + failure = ClientWorkerError(ClientErrorCode.GENERATION_FAILED, "Data Designer generation failed") + if context is not None: + self._write_failure(context, prepared, failure, progress=progress) + raise failure from error + + def _build_context( + self, + plan_path: Path, + *, + prepared: PreparedClientEnvironment, + endpoints: Mapping[str, str], + retry_resume: ResumeMode | None = None, + ) -> _ExecutionContext: + try: + plan = ResolvedSlurmRunPlan.model_validate_json( + read_regular_bytes(plan_path, missing_code=ClientErrorCode.INVALID_INPUT) + ) + shard = next(item for item in plan.shards if item.shard_id == prepared.shard_id) + self._validate_prepared(plan, shard, prepared) + lock = ResolvedDependencyLock.model_validate_json( + read_regular_bytes( + Path(plan.client.dependency_lock.path), + missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, + ) + ) + if lock.compute_sha256() != plan.client.dependency_lock.sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency lock digest differs") + expected_distributions = tuple( + sorted( + ( + *lock.base_distributions, + *( + InstalledDistribution(name=package.name, version=package.version) + for package in lock.overlay_packages + ), + ), + key=lambda item: item.name, + ) + ) + if prepared.installed_distributions != expected_distributions: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client environment differs from the lock") + builder_payload = self._load_builder(plan) + builder = DataDesignerConfigBuilder.from_config(builder_payload) + allocation_plan = resolve_allocation_plan(plan, self._environment) + providers = self._materialize_model_endpoints(allocation_plan, builder, endpoints) + self._validate_model_references(builder) + self._materialize_seed(plan, shard, builder) + mcp_providers = self._materialize_mcp_providers(plan) + managed_assets_path = self._validate_assets(plan) + requested_resume = ResumeMode(plan.invocation.authored.resume) + execution_resume = requested_resume if retry_resume is None else retry_resume + if ( + retry_resume is not None + and requested_resume is not ResumeMode.IF_POSSIBLE + and retry_resume is not requested_resume + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "retry resume mode differs from the plan") + dataset_path = self._dataset_path(plan, shard, prepared, execution_resume) + designer = self._data_designer_factory( + artifact_path=dataset_path.parent, + model_providers=providers, + managed_assets_path=managed_assets_path.as_posix(), + mcp_providers=mcp_providers, + auto_configure_logging=False, + ) + designer.set_run_config(RunConfig.model_validate(plan.invocation.effective_run_config)) + return _ExecutionContext(plan, shard, builder, designer, requested_resume, retry_resume, dataset_path) + except ClientWorkerError: + raise + except Exception as error: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "Data Designer configuration is invalid") from error + + @staticmethod + def _validate_prepared( + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + prepared: PreparedClientEnvironment, + ) -> None: + logical_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + expected_attempt = logical_attempt + if ( + plan.run_id != prepared.run_id + or prepared.attempt_dir != expected_attempt + or plan.client.dependency_lock != prepared.dependency_lock + or plan.client.runtime.runtime_sha256 != prepared.client_runtime_sha256 + or plan.client.runtime.python_abi != prepared.python_abi + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "prepared client environment differs from the plan") + + @staticmethod + def _load_builder(plan: ResolvedSlurmRunPlan) -> dict[str, object]: + if plan.builder.inline is not None: + return cast(dict[str, object], plan.builder.inline) + assert plan.builder.source is not None + payload = read_regular_bytes( + Path(plan.builder.source.path), + missing_code=ClientErrorCode.INVALID_INPUT, + ) + if hashlib.sha256(payload).hexdigest() != plan.builder.source.sha256: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "builder artifact digest differs") + try: + loaded = json.loads(payload) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder artifact is invalid") from error + if not isinstance(loaded, dict): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder artifact is invalid") + return cast(dict[str, object], loaded) + + @staticmethod + def _materialize_model_endpoints( + plan: ResolvedSlurmRunPlan, + builder: DataDesignerConfigBuilder, + endpoints: Mapping[str, str], + ) -> list[ModelProvider]: + aliases = tuple(deployment.authored.model_alias for deployment in plan.deployments) + if set(endpoints) != set(aliases): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "runtime endpoints do not match model aliases") + configs = {config.alias: config for config in builder.model_configs} + if set(configs) != set(aliases): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder model aliases do not match the plan") + updated = [] + providers = [] + for deployment, port in zip(plan.deployments, plan.client.ports, strict=True): + alias = deployment.authored.model_alias + endpoint = endpoints[alias] + expected_endpoint = f"http://127.0.0.1:{port.port}/v1" + parsed = urlsplit(endpoint) + if endpoint != expected_endpoint or parsed.username is not None or parsed.password is not None: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "runtime model endpoint is invalid") + provider_name = f"slurm-{alias}" + model_config = configs[alias] + inference = model_config.inference_parameters + concurrency = plan.invocation.authored.model_concurrency.get(alias) + if concurrency is not None: + inference = inference.model_copy(update={"max_parallel_requests": concurrency}) + updated.append( + model_config.model_copy( + update={ + "model": deployment.served_model_name, + "provider": provider_name, + "inference_parameters": inference, + } + ) + ) + providers.append(ModelProvider(name=provider_name, endpoint=endpoint, provider_type="openai")) + for model_config in tuple(builder.model_configs): + builder.delete_model_config(model_config.alias) + for model_config in updated: + builder.add_model_config(model_config) + return providers + + @staticmethod + def _validate_model_references(builder: DataDesignerConfigBuilder) -> None: + available = {config.alias for config in builder.model_configs} + referenced = {alias for column in builder.get_column_configs() for alias in column.get_model_aliases()} + referenced.update( + profiler.model_alias for profiler in builder.get_profilers() if hasattr(profiler, "model_alias") + ) + if not referenced.issubset(available): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder references unavailable model aliases") + + @staticmethod + def _materialize_seed( + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + builder: DataDesignerConfigBuilder, + ) -> None: + seed_path = plan.invocation.effective_input_bindings.seed_path + if seed_path is None: + return + seed = builder.get_seed_config() + if seed is None or "path" not in type(seed.source).model_fields: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed binding does not match the builder") + path = Path(seed_path) + if not path.exists() or not os.access(path, os.R_OK): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed input is unavailable") + if shard.input_partition is None: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "seed partition artifact is missing") + payload = read_regular_bytes( + Path(shard.input_partition.path), + missing_code=ClientErrorCode.INVALID_INPUT, + ) + if hashlib.sha256(payload).hexdigest() != shard.input_partition.sha256: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "seed partition artifact digest differs") + source = seed.source.model_copy(update={"path": path.as_posix()}) + builder.with_seed_dataset( + source, + sampling_strategy=seed.sampling_strategy, + selection_strategy=PartitionBlock(index=shard.shard_index, num_partitions=len(plan.shards)), + ) + + @staticmethod + def _materialize_mcp_providers( + plan: ResolvedSlurmRunPlan, + ) -> list[MCPProvider | LocalStdioMCPProvider]: + providers: list[MCPProvider | LocalStdioMCPProvider] = [] + for provider in plan.invocation.authored.mcp_providers: + if isinstance(provider, RemoteMCPProviderConfig): + api_key = _resolve_secret(provider.api_key) if provider.api_key is not None else None + providers.append( + MCPProvider( + provider_type=provider.provider_type, + name=provider.name, + endpoint=provider.endpoint, + api_key=api_key, + ) + ) + elif isinstance(provider, LocalStdioMCPProviderConfig): + environment = { + name: binding.value if isinstance(binding, LiteralEnvironmentBinding) else _resolve_secret(binding) + for name, binding in provider.environment.items() + } + providers.append( + LocalStdioMCPProvider( + name=provider.name, + command=provider.command, + args=list(provider.args), + env=environment, + ) + ) + return providers + + @staticmethod + def _validate_assets(plan: ResolvedSlurmRunPlan) -> Path: + value = plan.invocation.effective_input_bindings.managed_assets_path + assert value is not None + path = Path(value) + if not path.is_dir() or not os.access(path, os.R_OK | os.X_OK): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "managed assets are unavailable") + return path + + @staticmethod + def _dataset_path( + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + prepared: PreparedClientEnvironment, + resume: ResumeMode, + ) -> Path: + resume_path = Path(shard.resume_workspace.path) + if resume_path.is_symlink(): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "resume workspace is invalid") + if resume is ResumeMode.ALWAYS and (not resume_path.is_dir() or not any(resume_path.iterdir())): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "required resume workspace is unavailable") + if resume is ResumeMode.NEVER: + return prepared.attempt_dir / "dataset" + if resume_path.exists() and not resume_path.is_dir(): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "resume workspace is invalid") + return resume_path + + @staticmethod + def _prepare_dataset_workspace( + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + ) -> None: + if context.dataset_path.parent == prepared.attempt_dir: + ensure_private_directory(context.dataset_path.parent) + elif not context.dataset_path.parent.is_dir(): + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "shard dataset workspace is unavailable") + execution_resume = context.requested_resume if context.retry_resume is None else context.retry_resume + if execution_resume is not ResumeMode.NEVER: + ensure_private_directory(context.dataset_path) + if execution_resume is ResumeMode.NEVER and context.dataset_path.exists(): + if not context.dataset_path.is_dir() or any(context.dataset_path.iterdir()): + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "attempt dataset workspace is not empty") + + @staticmethod + def _generate_dataset( + context: _ExecutionContext, + progress: _ProgressWriter, + ) -> CreationResults: + try: + return cast( + CreationResults, + context.designer.create( + context.builder, + num_records=context.shard.requested_records, + dataset_name=context.dataset_path.name, + resume=context.requested_resume if context.retry_resume is None else context.retry_resume, + artifact_path=context.dataset_path.parent, + on_batch_complete=progress.on_batch_complete, + ), + ) + except ClientWorkerError: + raise + except (KeyboardInterrupt, SystemExit) as error: + raise ClientWorkerError(ClientErrorCode.INTERRUPTED, "Data Designer generation was interrupted") from error + except Exception as error: + raise ClientWorkerError(ClientErrorCode.GENERATION_FAILED, "Data Designer generation failed") from error + + def _finalize( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + results: CreationResults, + creation: _ValidatedCreation, + ) -> ClientResult: + dataset_path, output_paths = self._place_dataset(creation, results) + candidate = self._build_candidate_manifest(context, prepared, creation, dataset_path, output_paths) + return self._publish_success(context, prepared, creation, candidate) + + @staticmethod + def _validate_creation_result( + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + results: CreationResults, + ) -> _ValidatedCreation: + requested = context.shard.requested_records + actual = results.actual_num_records + if results.requested_num_records != requested or not 0 <= actual <= requested: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer result counts are invalid") + if results.early_shutdown is None or results.effective_resume_mode is None: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer result metadata is incomplete") + execution_resume = context.requested_resume if context.retry_resume is None else context.retry_resume + if results.requested_resume_mode is not execution_resume: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer resume metadata differs") + + dataset_path = Path(results.dataset_path) + effective_resume = results.effective_resume_mode + if context.retry_resume is not None and effective_resume is not context.retry_resume: + raise ClientWorkerError( + ClientErrorCode.OUTPUT_INVALID, "Data Designer effective resume mode differs from retry intent" + ) + shared_path = Path(context.shard.resume_workspace.path) + expected_path = shared_path if effective_resume is ResumeMode.ALWAYS else prepared.attempt_dir / "dataset" + if ( + not dataset_path.is_absolute() + or dataset_path != Path(os.path.normpath(dataset_path.as_posix())) + or not dataset_path.is_dir() + or dataset_path.is_symlink() + ): + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer dataset path is invalid") + if dataset_path != expected_path and ( + effective_resume is not ResumeMode.NEVER + or expected_path.exists() + or not ( + dataset_path.parent == prepared.attempt_dir + or dataset_path == shared_path + or (dataset_path.parent == shared_path.parent and dataset_path.name.startswith(f"{shared_path.name}_")) + ) + ): + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer dataset path is invalid") + return _ValidatedCreation( + requested_records=requested, + actual_records=actual, + dataset_path=dataset_path, + expected_path=expected_path, + effective_resume=effective_resume, + early_shutdown=results.early_shutdown, + ) + + @staticmethod + def _place_dataset( + creation: _ValidatedCreation, + results: CreationResults, + ) -> tuple[Path, tuple[Path, ...]]: + dataset_path = creation.dataset_path + output_paths = ClientWorker._candidate_output_paths(creation, results) + if dataset_path != creation.expected_path: + ensure_private_directory(creation.expected_path.parent) + dataset_path.rename(creation.expected_path) + output_paths = tuple(creation.expected_path / path.relative_to(dataset_path) for path in output_paths) + dataset_path = creation.expected_path + if dataset_path != creation.expected_path: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer dataset path is invalid") + return dataset_path, output_paths + + @staticmethod + def _candidate_output_paths(creation: _ValidatedCreation, results: CreationResults) -> tuple[Path, ...]: + if not creation.actual_records: + return () + batches = tuple(sorted(results.artifact_storage.final_dataset_path.glob("batch_*.parquet"))) + if not batches: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer produced no final batch files") + if len(batches) <= MAXIMUM_CANDIDATE_OUTPUT_FILES: + schema_digests = {compute_candidate_schema_digest(lazy.pq.read_schema(path)) for path in batches} + if len(schema_digests) == 1: + return batches + # The export keeps the candidate bounded when there are many batches and + # unifies schemas when batch column types differ. + return (results.export(creation.dataset_path / "part-00000.parquet", format="parquet"),) + + def _build_candidate_manifest( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + creation: _ValidatedCreation, + dataset_path: Path, + output_paths: tuple[Path, ...], + ) -> CandidateOutputManifest: + files: tuple[CandidateOutputFile, ...] = () + schema_digest = context.plan.builder.content_sha256 + if output_paths: + record_counts = tuple(lazy.pq.read_metadata(path).num_rows for path in output_paths) + if sum(record_counts) != creation.actual_records: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "candidate record count differs") + schema_digests = {compute_candidate_schema_digest(lazy.pq.read_schema(path)) for path in output_paths} + if len(schema_digests) != 1: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "candidate batch schemas differ") + schema_digest = schema_digests.pop() + files = tuple( + CandidateOutputFile( + relative_path=path.relative_to(dataset_path).as_posix(), + sha256=compute_file_sha256(path, missing_code=ClientErrorCode.OUTPUT_INVALID), + byte_size=path.stat().st_size, + record_count=record_count, + ) + for path, record_count in zip(output_paths, record_counts, strict=True) + ) + created_at = self._clock() + return CandidateOutputManifest( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=prepared.attempt_id, + attempt_ordinal=int(prepared.attempt_id.removeprefix("attempt-")), + created_at=created_at, + dataset_path=dataset_path.as_posix(), + requested_records=creation.requested_records, + actual_records=creation.actual_records, + outcome=( + CandidateOutcome.EMPTY + if creation.actual_records == 0 + else CandidateOutcome.COMPLETE + if creation.actual_records == creation.requested_records + else CandidateOutcome.PARTIAL + ), + files=files, + dataset_schema_digest=schema_digest, + provenance_digest=context.plan.compute_sha256(), + ) + + def _publish_success( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + creation: _ValidatedCreation, + candidate: CandidateOutputManifest, + ) -> ClientResult: + candidate_path = prepared.attempt_dir / "output-manifest.json" + publish_private_text(candidate_path, candidate.serialize_json()) + result = ClientResult( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=prepared.attempt_id, + completed_at=self._clock(), + requested_records=creation.requested_records, + actual_records=creation.actual_records, + outcome=( + ClientOutcome.COMPLETE + if creation.actual_records == creation.requested_records + else ClientOutcome.PARTIAL + ), + dataset_path=candidate.dataset_path, + early_shutdown=creation.early_shutdown, + requested_resume_mode=context.requested_resume.value, + effective_resume_mode=creation.effective_resume.value, + candidate_output_manifest=ArtifactReference( + path=candidate_path.as_posix(), + sha256=candidate.compute_sha256(), + ), + ) + publish_private_text(prepared.attempt_dir / "client-result.json", result.serialize_json()) + return result + + def _validate_environment_manifest( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + plugins: tuple[ClientPluginEntryPoint, ...], + ) -> None: + try: + manifest = ClientEnvironmentManifest.model_validate_json( + read_regular_bytes( + prepared.attempt_dir / "client-environment.json", + missing_code=ClientErrorCode.CONFIG_INVALID, + ) + ) + except ValidationError as error: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "client environment manifest is invalid") from error + if ( + manifest.outcome is not ClientEnvironmentOutcome.READY + or manifest.run_id != context.plan.run_id + or manifest.shard_id != context.shard.shard_id + or manifest.attempt_id != prepared.attempt_id + or manifest.dependency_lock != prepared.dependency_lock + or manifest.client_runtime_sha256 != prepared.client_runtime_sha256 + or manifest.python_abi != prepared.python_abi + or manifest.overlay_path != prepared.overlay_path.as_posix() + or manifest.installed_distributions != prepared.installed_distributions + or manifest.plugins != plugins + ): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "client environment manifest differs") + + def _write_preflight_failure( + self, + plan_path: Path, + prepared: PreparedClientEnvironment, + endpoints: Mapping[str, str], + plugins: tuple[ClientPluginEntryPoint, ...], + error: ClientWorkerError, + ) -> None: + manifest = ClientEnvironmentManifest.from_prepared( + prepared, + created_at=self._clock(), + outcome=ClientEnvironmentOutcome.FAILED, + plugins=plugins, + error_code=error.code, + redacted_message=error.redacted_message, + ) + replace_private_text(prepared.attempt_dir / "client-environment.json", manifest.serialize_json()) + try: + context = self._build_context(plan_path, prepared=prepared, endpoints=endpoints) + except ClientWorkerError: + return + self._write_failure(context, prepared, error, progress=None) + + def _write_failure( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + error: ClientWorkerError, + *, + progress: _ProgressWriter | None, + ) -> None: + result = ClientResult( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=prepared.attempt_id, + completed_at=self._clock(), + requested_records=context.shard.requested_records, + actual_records=None, + outcome=ClientOutcome.FAILED, + requested_resume_mode=context.requested_resume.value, + error_code=error.code.value, + redacted_message=error.redacted_message, + ) + publish_private_text(prepared.attempt_dir / "client-result.json", result.serialize_json()) + progress = progress or _ProgressWriter(context, prepared, self._clock, revision=2) + progress.best_effort(ClientProgressPhase.FAILED, error_code=error.code) + + +def _resolve_secret(reference: SecretRef) -> str: + try: + return os.environ[reference.environment] + except KeyError as error: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "required client secret is unavailable") from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/client/filesystem.py new file mode 100644 index 000000000..8c7364ebb --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/filesystem.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +import stat +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.records import ClientErrorCode + +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_DIRECTORY_MODE = 0o700 +_FILE_MODE = 0o600 +_DIRECTORY_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW +_FILE_READ_FLAGS = os.O_RDONLY | os.O_NOFOLLOW | getattr(os, "O_NONBLOCK", 0) + + +def ensure_private_directory(path: Path) -> None: + """Create one private directory tree and reject symlink components.""" + try: + with _open_directory(path, create=True) as descriptor: + os.fchmod(descriptor, _DIRECTORY_MODE) + except OSError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client workspace is not a regular directory") from error + + +def read_regular_bytes(path: Path, *, missing_code: ClientErrorCode) -> bytes: + """Read one bounded regular file without following a final symlink.""" + try: + with _open_directory(path.parent, create=False) as directory_descriptor: + with _open_regular_file(directory_descriptor, path.name) as (descriptor, info): + if not stat.S_ISREG(info.st_mode): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client input is not a regular file") + return _read_bounded_bytes(descriptor, info) + except ClientWorkerError: + raise + except FileNotFoundError as error: + raise ClientWorkerError(missing_code, "required client artifact is missing") from error + except OSError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "required client artifact is unreadable") from error + + +def compute_file_sha256(path: Path, *, missing_code: ClientErrorCode) -> str: + """Hash one regular file without retaining its contents.""" + try: + with _open_directory(path.parent, create=False) as directory_descriptor: + with _open_regular_file(directory_descriptor, path.name) as (descriptor, info): + if not stat.S_ISREG(info.st_mode): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client artifact is not a regular file") + digest = hashlib.sha256() + while chunk := os.read(descriptor, 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + except ClientWorkerError: + raise + except FileNotFoundError as error: + raise ClientWorkerError(missing_code, "required client artifact is missing") from error + except OSError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "required client artifact is unreadable") from error + + +def replace_private_text(path: Path, text: str) -> None: + """Atomically replace one private attempt-local record.""" + with _write_private_temporary(path, text) as (_, directory_descriptor, temporary_name): + os.replace( + temporary_name, + path.name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + ) + + +def publish_private_text(path: Path, text: str) -> None: + """Convergently publish one immutable private attempt-local record.""" + with _write_private_temporary(path, text) as (payload, directory_descriptor, temporary_name): + try: + os.link( + temporary_name, + path.name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + try: + with _open_regular_file(directory_descriptor, path.name) as (descriptor, info): + if not stat.S_ISREG(info.st_mode): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client input is not a regular file") + if _read_bounded_bytes(descriptor, info) != payload: + raise ClientWorkerError( + ClientErrorCode.OUTPUT_INVALID, "immutable client record already differs" + ) + os.fchmod(descriptor, _FILE_MODE) + except ClientWorkerError: + raise + except FileNotFoundError as error: + raise ClientWorkerError( + ClientErrorCode.OUTPUT_INVALID, "required client artifact is missing" + ) from error + except OSError as error: + raise ClientWorkerError( + ClientErrorCode.INVALID_INPUT, "required client artifact is unreadable" + ) from error + + +@contextmanager +def _write_private_temporary(path: Path, text: str) -> Iterator[tuple[bytes, int, str]]: + payload = text.encode("utf-8") + if len(payload) > _MAXIMUM_RECORD_SIZE: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "client record exceeds the size limit") + with _open_directory(path.parent, create=True) as directory_descriptor: + os.fchmod(directory_descriptor, _DIRECTORY_MODE) + temporary_name = f".client.{uuid.uuid4().hex}.tmp" + descriptor = os.open( + temporary_name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, + _FILE_MODE, + dir_fd=directory_descriptor, + ) + try: + try: + os.fchmod(descriptor, _FILE_MODE) + remaining = memoryview(payload) + while remaining: + written = os.write(descriptor, remaining) + if written == 0: + raise OSError("client temporary file write made no progress") + remaining = remaining[written:] + os.fsync(descriptor) + finally: + os.close(descriptor) + yield payload, directory_descriptor, temporary_name + finally: + try: + os.unlink(temporary_name, dir_fd=directory_descriptor) + except FileNotFoundError: + pass + + +@contextmanager +def _open_directory(path: Path, *, create: bool) -> Iterator[int]: + start = path.anchor or "." + parts = path.parts[1:] if path.anchor else path.parts + descriptor = os.open(start, _DIRECTORY_FLAGS) + try: + for part in parts: + if part in ("", ".", ".."): + raise OSError("client directory path is not canonical") + try: + child_descriptor = os.open(part, _DIRECTORY_FLAGS, dir_fd=descriptor) + except FileNotFoundError: + if not create: + raise + try: + os.mkdir(part, _DIRECTORY_MODE, dir_fd=descriptor) + except FileExistsError: + pass + child_descriptor = os.open(part, _DIRECTORY_FLAGS, dir_fd=descriptor) + try: + if not stat.S_ISDIR(os.fstat(child_descriptor).st_mode): + raise OSError("client workspace component is not a directory") + except BaseException: + os.close(child_descriptor) + raise + previous_descriptor = descriptor + descriptor = child_descriptor + os.close(previous_descriptor) + yield descriptor + finally: + os.close(descriptor) + + +@contextmanager +def _open_regular_file(directory_descriptor: int, name: str) -> Iterator[tuple[int, os.stat_result]]: + descriptor = os.open(name, _FILE_READ_FLAGS, dir_fd=directory_descriptor) + try: + yield descriptor, os.fstat(descriptor) + finally: + os.close(descriptor) + + +def _read_bounded_bytes(descriptor: int, info: os.stat_result) -> bytes: + if info.st_size > _MAXIMUM_RECORD_SIZE: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client input exceeds the size limit") + payload = bytearray() + while len(payload) <= _MAXIMUM_RECORD_SIZE: + chunk = os.read(descriptor, min(1024 * 1024, _MAXIMUM_RECORD_SIZE + 1 - len(payload))) + if not chunk: + return bytes(payload) + payload.extend(chunk) + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client input exceeds the size limit") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/plugins.py b/packages/data-designer-slurm/src/data_designer/slurm/client/plugins.py new file mode 100644 index 000000000..5da494c30 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/plugins.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from importlib.metadata import entry_points + +from packaging.utils import canonicalize_name + +from data_designer.plugins import Plugin +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.records import ClientErrorCode, ClientPluginEntryPoint +from data_designer.slurm.contracts import InstalledDistribution + + +def discover_plugins( + installed_distributions: tuple[InstalledDistribution, ...], +) -> tuple[ClientPluginEntryPoint, ...]: + """Load and validate every installed Data Designer plugin entry point.""" + installed = {distribution.name: distribution.version for distribution in installed_distributions} + discovered: list[ClientPluginEntryPoint] = [] + plugin_names: set[str] = set() + try: + for entry_point in entry_points(group="data_designer.plugins"): + if entry_point.dist is None: + raise ValueError("plugin entry point has no owning distribution") + distribution = canonicalize_name(entry_point.dist.name) + version = entry_point.dist.version + if installed.get(distribution) != version: + raise ValueError("plugin entry point is outside the verified environment") + plugin = entry_point.load() + if not isinstance(plugin, Plugin): + raise TypeError("plugin entry point did not load a Plugin") + plugin.config_cls + plugin.impl_cls + if plugin.name in plugin_names: + raise ValueError("plugin names must be unique") + plugin_names.add(plugin.name) + discovered.append( + ClientPluginEntryPoint( + entry_point=entry_point.name, + value=entry_point.value, + distribution=distribution, + distribution_version=version, + plugin_name=plugin.name, + plugin_type=plugin.plugin_type.value, + ) + ) + except ClientWorkerError: + raise + except Exception as error: + raise ClientWorkerError(ClientErrorCode.PLUGIN_LOAD_FAILED, "Data Designer plugin validation failed") from error + return tuple(sorted(discovered, key=lambda plugin: (plugin.distribution, plugin.entry_point))) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/process.py b/packages/data-designer-slurm/src/data_designer/slurm/client/process.py new file mode 100644 index 000000000..f66895e42 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/process.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bootstrap-safe process boundary for the plugin-aware client worker. + +This module must remain free of Data Designer configuration, interface, plugin, +and worker imports. The child process activates its verified dependency overlay +before importing any of those modules. +""" + +from __future__ import annotations + +import subprocess +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field + +ProcessExecutor = Callable[[tuple[str, ...]], int] + + +def execute_process(command: tuple[str, ...]) -> int: + """Execute a child process and return its exit status.""" + return subprocess.run(command, check=False).returncode + + +@dataclass(frozen=True) +class ClientWorkerProcess: + """Launch the plugin-aware client worker in a fresh Python interpreter.""" + + executable: str = field(default_factory=lambda: sys.executable) + executor: ProcessExecutor = execute_process + + def run(self, arguments: Sequence[str]) -> int: + """Run one client-worker operation without inheriting imported modules.""" + command = (self.executable, "-m", "data_designer.slurm.client.worker", *arguments) + return self.executor(command) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py new file mode 100644 index 000000000..ce8cf1cd7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta +from enum import Enum +from typing import TYPE_CHECKING, Annotated, Literal + +from pydantic import NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + ContractValue, + Identifier, + InstalledDistribution, + Sha256Digest, + ShardId, + validate_absolute_path, + validate_plain_text, +) + +if TYPE_CHECKING: + from data_designer.slurm.client.environment import PreparedClientEnvironment + + +class ClientErrorCode(str, Enum): + INVALID_INPUT = "invalid_input" + DEPENDENCY_ARTIFACT_MISSING = "dependency_artifact_missing" + DEPENDENCY_DIGEST_MISMATCH = "dependency_digest_mismatch" + DEPENDENCY_CONFLICT = "dependency_conflict" + DEPENDENCY_INSTALL_FAILED = "dependency_install_failed" + PLUGIN_LOAD_FAILED = "plugin_load_failed" + CONFIG_INVALID = "config_invalid" + GENERATION_FAILED = "generation_failed" + INTERRUPTED = "interrupted" + OUTPUT_INVALID = "output_invalid" + + +class ClientEnvironmentOutcome(str, Enum): + READY = "ready" + FAILED = "failed" + + +class ClientInstallerOutcome(str, Enum): + NOT_REQUIRED = "not_required" + INSTALLED = "installed" + REUSED = "reused" + + +class ClientProgressPhase(str, Enum): + PREPARING_ENVIRONMENT = "preparing_environment" + VALIDATING_PLUGINS = "validating_plugins" + VALIDATING_CONFIG = "validating_config" + GENERATING = "generating" + FINALIZING = "finalizing" + COMPLETE = "complete" + FAILED = "failed" + + +class ClientPluginEntryPoint(ContractValue): + entry_point: Annotated[str, StringConstraints(min_length=1, max_length=256)] + value: Annotated[str, StringConstraints(min_length=1, max_length=512)] + distribution: Annotated[str, StringConstraints(min_length=1, max_length=128)] + distribution_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + plugin_name: Identifier + plugin_type: Literal["column-generator", "seed-reader", "processor"] + + @field_validator("entry_point", "value", "distribution", "distribution_version") + @classmethod + def validate_text(cls, value: str) -> str: + return validate_plain_text(value, field_name="plugin entry point") + + +class ClientEnvironmentManifest(ContractRecord): + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + created_at: datetime + outcome: ClientEnvironmentOutcome + dependency_lock: ArtifactReference + client_runtime_sha256: Sha256Digest + python_abi: Identifier + overlay_path: str + installer_outcome: ClientInstallerOutcome + installed_distributions: tuple[InstalledDistribution, ...] + plugins: tuple[ClientPluginEntryPoint, ...] + error_code: ClientErrorCode | None = None + redacted_message: Annotated[str, StringConstraints(max_length=512)] | None = None + + _overlay_path_is_absolute = field_validator("overlay_path")(validate_absolute_path) + + @classmethod + def from_prepared( + cls, + prepared: PreparedClientEnvironment, + *, + created_at: datetime, + outcome: ClientEnvironmentOutcome, + plugins: tuple[ClientPluginEntryPoint, ...], + error_code: ClientErrorCode | None = None, + redacted_message: str | None = None, + ) -> ClientEnvironmentManifest: + return cls( + schema_version=1, + run_id=prepared.run_id, + shard_id=prepared.shard_id, + attempt_id=prepared.attempt_id, + created_at=created_at, + outcome=outcome, + dependency_lock=prepared.dependency_lock, + client_runtime_sha256=prepared.client_runtime_sha256, + python_abi=prepared.python_abi, + overlay_path=prepared.overlay_path.as_posix(), + installer_outcome=prepared.installer_outcome, + installed_distributions=prepared.installed_distributions, + plugins=plugins, + error_code=error_code, + redacted_message=redacted_message, + ) + + @field_validator("created_at") + @classmethod + def validate_created_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("created_at must be timezone-aware UTC") + return value + + @field_validator("redacted_message") + @classmethod + def validate_redacted_message(cls, value: str | None) -> str | None: + return None if value is None else validate_plain_text(value, field_name="redacted message") + + @model_validator(mode="after") + def validate_environment(self) -> ClientEnvironmentManifest: + distribution_names = tuple(distribution.name for distribution in self.installed_distributions) + if distribution_names != tuple(sorted(distribution_names)) or len(distribution_names) != len( + set(distribution_names) + ): + raise ValueError("installed distributions must be sorted and unique") + plugin_keys = tuple((plugin.distribution, plugin.entry_point) for plugin in self.plugins) + if plugin_keys != tuple(sorted(plugin_keys)) or len(plugin_keys) != len(set(plugin_keys)): + raise ValueError("plugin entry points must be sorted and unique") + if self.outcome is ClientEnvironmentOutcome.READY: + if self.error_code is not None or self.redacted_message is not None: + raise ValueError("ready client environments cannot contain failure details") + elif self.error_code is None: + raise ValueError("failed client environments require error_code") + return self + + +class ClientProgress(ContractRecord): + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + revision: PositiveInt + updated_at: datetime + phase: ClientProgressPhase + requested_records: PositiveInt + completed_records: NonNegativeInt | None = None + error_code: ClientErrorCode | None = None + + @field_validator("updated_at") + @classmethod + def validate_updated_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("updated_at must be timezone-aware UTC") + return value + + @model_validator(mode="after") + def validate_progress(self) -> ClientProgress: + if self.completed_records is not None and self.completed_records > self.requested_records: + raise ValueError("completed_records must not exceed requested_records") + if (self.phase is ClientProgressPhase.FAILED) != (self.error_code is not None): + raise ValueError("only failed progress requires an error code") + return self + + +class ClientOutcome(str, Enum): + COMPLETE = "complete" + PARTIAL = "partial" + FAILED = "failed" + + +class ClientResult(ContractRecord): + """Semantic Data Designer outcome independent of engine-internal result types.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + completed_at: datetime + requested_records: PositiveInt + actual_records: NonNegativeInt | None + outcome: ClientOutcome + dataset_path: str | None = None + early_shutdown: bool | None = None + requested_resume_mode: Literal["never", "always", "if_possible"] + effective_resume_mode: Literal["never", "always"] | None = None + candidate_output_manifest: ArtifactReference | None = None + error_code: Identifier | None = None + redacted_message: Annotated[str, StringConstraints(max_length=512)] | None = None + + @field_validator("completed_at") + @classmethod + def validate_completed_at(cls, value: datetime) -> datetime: + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise ValueError("completed_at must be timezone-aware UTC") + return value + + @field_validator("dataset_path") + @classmethod + def validate_dataset_path(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + @field_validator("redacted_message") + @classmethod + def validate_message(cls, value: str | None) -> str | None: + if value is not None and any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("redacted_message must not contain control characters") + return value + + @model_validator(mode="after") + def validate_outcome(self) -> ClientResult: + if self.actual_records is not None and self.actual_records > self.requested_records: + raise ValueError("actual_records must not exceed requested_records") + if self.requested_resume_mode == "never" and self.effective_resume_mode not in {None, "never"}: + raise ValueError("resume mode never cannot become effective resume") + if self.outcome is not ClientOutcome.FAILED: + if self.early_shutdown is None or self.effective_resume_mode is None: + raise ValueError("non-failed client results require resume and early-shutdown facts") + if self.outcome is ClientOutcome.COMPLETE: + if self.actual_records != self.requested_records: + raise ValueError("complete client results require the requested record count") + if self.early_shutdown: + raise ValueError("complete client results cannot report early shutdown") + self._require_success_artifacts() + elif self.outcome is ClientOutcome.PARTIAL: + if self.actual_records is None or self.actual_records >= self.requested_records: + raise ValueError("partial client results require fewer than the requested record count") + self._require_success_artifacts() + else: + if self.candidate_output_manifest is not None: + raise ValueError("failed client results cannot reference a candidate output manifest") + if self.error_code is None: + raise ValueError("failed client results require error_code") + return self + + def _require_success_artifacts(self) -> None: + if self.dataset_path is None or self.candidate_output_manifest is None: + raise ValueError("successful client results require dataset and candidate manifest paths") + if self.error_code is not None or self.redacted_message is not None: + raise ValueError("successful client results cannot contain failure details") + shard_root = f"/runs/{self.run_id}/shards/{self.shard_id}" + if self.effective_resume_mode == "never": + expected_dataset = f"{shard_root}/attempts/{self.attempt_id}/dataset" + else: + expected_dataset = f"{shard_root}/dataset" + if not self.dataset_path.endswith(expected_dataset): + raise ValueError("dataset path must match the run, shard, attempt, and resume policy") + expected_manifest = f"{shard_root}/attempts/{self.attempt_id}/output-manifest.json" + if not self.candidate_output_manifest.path.endswith(expected_manifest): + raise ValueError("candidate output reference must match the run, shard, and attempt") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py b/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py new file mode 100644 index 000000000..b6c4f6a63 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inspect and bind the native Python runtime used by Slurm client steps.""" + +from __future__ import annotations + +import os +import platform +import sys +from pathlib import Path + +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion, Version + +from data_designer.slurm.client.environment import inspect_distributions +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.contracts import compute_canonical_json_sha256, validate_absolute_path +from data_designer.slurm.planning import ResolvedClientRuntime + +_REQUIRED_DISTRIBUTIONS = frozenset( + { + "aiohttp", + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + } +) +_PROXY_AIOHTTP_VERSIONS = SpecifierSet(">=3.14.3,<4") +_DATA_DESIGNER_DISTRIBUTIONS = frozenset( + {"data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm"} +) + + +class ClientRuntimeInspectionError(RuntimeError): + """Raised when the active CLI interpreter cannot be used on compute nodes.""" + + +class ClientRuntimeInspector: + """Resolve the active CLI interpreter into an immutable planning contract.""" + + def inspect(self) -> ResolvedClientRuntime: + """Return verified facts for the Python interpreter running the CLI.""" + try: + python_executable = validate_absolute_path(Path(sys.executable).absolute().as_posix()) + if not os.access(python_executable, os.X_OK): + raise ClientRuntimeInspectionError("client Python executable is unavailable") + cache_tag = sys.implementation.cache_tag + if cache_tag is None: + raise ClientRuntimeInspectionError("client Python does not expose an ABI tag") + python_abi = f"cp{cache_tag.removeprefix('cpython-')}" if cache_tag.startswith("cpython-") else cache_tag + distributions = inspect_distributions(None) + names = {distribution.name for distribution in distributions} + missing = _REQUIRED_DISTRIBUTIONS.difference(names) + if missing: + raise ClientRuntimeInspectionError( + f"required client distributions are not installed with user-site disabled: {', '.join(sorted(missing))}; " + "install them in the shared client Python environment, for example a virtual environment" + ) + package_versions = { + distribution.version + for distribution in distributions + if distribution.name in _DATA_DESIGNER_DISTRIBUTIONS + } + if len(package_versions) != 1: + raise ClientRuntimeInspectionError("Data Designer packages in the client Python must share one version") + aiohttp_version = next( + distribution.version for distribution in distributions if distribution.name == "aiohttp" + ) + try: + compatible_aiohttp = Version(aiohttp_version) in _PROXY_AIOHTTP_VERSIONS + except InvalidVersion: + compatible_aiohttp = False + if not compatible_aiohttp: + raise ClientRuntimeInspectionError("client Python requires aiohttp>=3.14.3,<4 for the inference proxy") + identity = { + "python_executable": python_executable, + "python_implementation": platform.python_implementation().casefold(), + "python_version": platform.python_version(), + "python_abi": python_abi, + "distributions": distributions, + } + fingerprint_payload = identity | { + "distributions": [distribution.model_dump(mode="json") for distribution in distributions] + } + return ResolvedClientRuntime( + **identity, + runtime_sha256=compute_canonical_json_sha256(fingerprint_payload), + ) + except ClientRuntimeInspectionError: + raise + except ClientWorkerError as error: + raise ClientRuntimeInspectionError(error.redacted_message) from None + except (OSError, ValueError): + raise ClientRuntimeInspectionError("client Python runtime cannot be inspected") from None + + +__all__ = ["ClientRuntimeInspectionError", "ClientRuntimeInspector"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py new file mode 100644 index 000000000..87a8f9ca7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import importlib +import os +import signal +import sys +from collections.abc import Sequence +from datetime import datetime, timezone +from pathlib import Path + +from data_designer.config import ResumeMode +from data_designer.slurm.client.environment import ( + ClientEnvironmentBuilder, + PreparedClientEnvironment, + activate_environment, +) +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import replace_private_text +from data_designer.slurm.client.records import ( + ClientEnvironmentManifest, + ClientEnvironmentOutcome, + ClientErrorCode, + ClientPluginEntryPoint, +) + + +def main(argv: Sequence[str] | None = None) -> int: + """Run one allocation-local client preflight or generation operation.""" + arguments = _parse_arguments(argv) + previous_sigterm_handler = signal.signal(signal.SIGTERM, _handle_sigterm) + prepared: PreparedClientEnvironment | None = None + plugins: tuple[ClientPluginEntryPoint, ...] = () + try: + endpoints = _parse_endpoints(arguments.endpoint) + scratch_root = os.environ.get("DATA_DESIGNER_SLURM_SCRATCH_ROOT") + if scratch_root is None: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "allocation scratch is unavailable") + prepared = ClientEnvironmentBuilder().prepare( + arguments.plan, + shard_id=arguments.shard_id, + attempt_id=arguments.attempt_id, + attempt_dir=arguments.attempt_dir, + scratch_root=Path(scratch_root), + ) + activate_environment(prepared) + plugins_module = importlib.import_module("data_designer.slurm.client.plugins") + discover_plugins = getattr(plugins_module, "discover_plugins") + plugins = discover_plugins(prepared.installed_distributions) + execution_module = importlib.import_module("data_designer.slurm.client.execution") + ClientWorker = getattr(execution_module, "ClientWorker") + worker = ClientWorker() + retry_resume = None if arguments.resume_mode is None else ResumeMode(arguments.resume_mode) + if arguments.operation == "preflight": + worker.preflight( + arguments.plan, + prepared=prepared, + endpoints=endpoints, + plugins=plugins, + retry_resume=retry_resume, + ) + else: + worker.run( + arguments.plan, + prepared=prepared, + endpoints=endpoints, + plugins=plugins, + retry_resume=retry_resume, + ) + return 0 + except ClientWorkerError as error: + if prepared is not None and error.code is ClientErrorCode.PLUGIN_LOAD_FAILED: + _write_plugin_failure(prepared, plugins, error) + print(f"client worker failed: {error.code.value}", file=sys.stderr) + return 2 + except (KeyboardInterrupt, SystemExit): + print(f"client worker failed: {ClientErrorCode.INTERRUPTED.value}", file=sys.stderr) + return 2 + except Exception: + print(f"client worker failed: {ClientErrorCode.INVALID_INPUT.value}", file=sys.stderr) + return 2 + finally: + signal.signal(signal.SIGTERM, previous_sigterm_handler) + + +def _handle_sigterm(_signum: int, _frame: object) -> None: + raise KeyboardInterrupt + + +def _parse_arguments(argv: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(prog="python -m data_designer.slurm.client.worker") + parser.add_argument("operation", choices=("preflight", "run")) + parser.add_argument("--plan", required=True, type=Path) + parser.add_argument("--shard-id", required=True) + parser.add_argument("--attempt-id", required=True) + parser.add_argument("--attempt-dir", required=True, type=Path) + parser.add_argument("--resume-mode", choices=("never", "always")) + parser.add_argument("--endpoint", action="append", default=[]) + return parser.parse_args(argv) + + +def _parse_endpoints(values: list[str]) -> dict[str, str]: + endpoints: dict[str, str] = {} + for value in values: + alias, separator, endpoint = value.rpartition("=") + if not separator or not alias or not endpoint or alias in endpoints: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "runtime endpoint argument is invalid") + endpoints[alias] = endpoint + return endpoints + + +def _write_plugin_failure( + prepared: PreparedClientEnvironment, + plugins: tuple[ClientPluginEntryPoint, ...], + error: ClientWorkerError, +) -> None: + manifest = ClientEnvironmentManifest.from_prepared( + prepared, + created_at=datetime.now(timezone.utc), + outcome=ClientEnvironmentOutcome.FAILED, + plugins=plugins, + error_code=error.code, + redacted_message=error.redacted_message, + ) + replace_private_text(prepared.attempt_dir / "client-environment.json", manifest.serialize_json()) + + +if __name__ == "__main__": + raise SystemExit(main()) 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 new file mode 100644 index 000000000..9842577f0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public authored configuration contracts for Data Designer Slurm.""" + +from __future__ import annotations + +from data_designer.slurm.config.benchmark import ( + AdaptiveRecordPolicy, + BenchmarkAnalysisTargets, + BenchmarkBaseRun, + BenchmarkDeploymentCase, + BenchmarkDeploymentOverride, + DataDesignerSlurmBenchmarkConfig, + FixedRecordPolicy, +) +from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) +from data_designer.slurm.config.errors import SlurmConfigBuilderError, SlurmConfigLoadError +from data_designer.slurm.config.images import ( + ClientImageInspection, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + ServingImageInspection, +) +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, + resolve_profile, +) +from data_designer.slurm.config.profiles import ( + ContainerMount, + GpuRequestMode, + ImageBuildProfile, + ProfileSelectionSource, + SchedulerProfile, + SelectedSlurmProfile, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, + validate_selected_profile, +) +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + DeploymentResources, + DeploymentTopology, + InputBindings, + InvocationConfig, + LocalStdioMCPProviderConfig, + RemoteMCPProviderConfig, + ServerDeploymentConfig, + SubmissionConfig, +) +from data_designer.slurm.config.vllm import QueueBackpressureConfig, VllmServerConfig + +__all__ = [ + "AdaptiveRecordPolicy", + "ArrayTasksConfig", + "BenchmarkAnalysisTargets", + "BenchmarkBaseRun", + "BenchmarkDeploymentCase", + "BenchmarkDeploymentOverride", + "BuilderInput", + "ClientConfig", + "ClientDependencies", + "ClientImageInspection", + "ContainerMount", + "DataDesignerSlurmBenchmarkConfig", + "DataDesignerSlurmConfig", + "DataDesignerSlurmConfigBuilder", + "DEFAULT_PROFILE_FILE_NAME", + "DeploymentResources", + "DeploymentTopology", + "FixedRecordPolicy", + "GpuRequestMode", + "ImageBuildProfile", + "ImageBuildRequest", + "ImageInspectionRecord", + "ImageKind", + "ImageRef", + "InputBindings", + "InstalledDistribution", + "InvocationConfig", + "LiteralEnvironmentBinding", + "LocalStdioMCPProviderConfig", + "ProfileSelectionSource", + "PROFILE_FILE_ENVIRONMENT", + "QueueBackpressureConfig", + "RemoteMCPProviderConfig", + "SchedulerProfile", + "SecretRef", + "collect_secret_environment_names", + "SelectedSlurmProfile", + "ServerDeploymentConfig", + "ServingImageInspection", + "SlurmConfigBuilderError", + "SlurmConfigLoadError", + "SlurmProfile", + "SlurmProfileCatalog", + "SubmissionConfig", + "VllmServerConfig", + "injected_profile", + "load_builder_payload", + "load_benchmark_config", + "load_profile_catalog", + "load_run_config", + "resolve_profile", + "select_profile", + "validate_selected_profile", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py new file mode 100644 index 000000000..10a2386eb --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import Field, PositiveFloat, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.contracts import ( + AuthoredConfig, + Duration, + Identifier, + SchemaVersion, + validate_local_config_path, +) + + +class BenchmarkBaseRun(AuthoredConfig): + source: str | None = None + inline: DataDesignerSlurmConfig | None = None + + @model_validator(mode="before") + @classmethod + def normalize_source(cls, value: object) -> object: + if isinstance(value, str): + return {"source": value} + return value + + @field_validator("source") + @classmethod + def validate_source(cls, value: str | None) -> str | None: + return None if value is None else validate_local_config_path(value) + + @model_validator(mode="after") + def validate_base_run(self) -> BenchmarkBaseRun: + if (self.source is None) == (self.inline is None): + raise ValueError("base_run requires exactly one of source or inline") + return self + + +class BenchmarkDeploymentOverride(AuthoredConfig): + nodes: PositiveInt + nodes_per_replica: PositiveInt + + @model_validator(mode="after") + def validate_topology(self) -> BenchmarkDeploymentOverride: + if self.nodes % self.nodes_per_replica: + raise ValueError("nodes_per_replica must divide benchmark deployment nodes") + return self + + +class BenchmarkDeploymentCase(AuthoredConfig): + name: Identifier + deployments: dict[str, BenchmarkDeploymentOverride] = Field(min_length=1) + + +class FixedRecordPolicy(AuthoredConfig): + type: Literal["fixed"] + records: PositiveInt + + +class AdaptiveRecordPolicy(AuthoredConfig): + type: Literal["adaptive"] = "adaptive" + base_records: PositiveInt = 1000 + max_records: PositiveInt = 5000 + records_per_concurrency: PositiveFloat = 1.0 + + @model_validator(mode="after") + def validate_bounds(self) -> AdaptiveRecordPolicy: + if self.max_records < self.base_records: + raise ValueError("adaptive max_records must not be less than base_records") + return self + + +BenchmarkRecordPolicy = Annotated[FixedRecordPolicy | AdaptiveRecordPolicy, Field(discriminator="type")] + + +class BenchmarkAnalysisTargets(AuthoredConfig): + target_total_records: PositiveInt + target_runtime: Duration + + +class DataDesignerSlurmBenchmarkConfig(AuthoredConfig): + """Authored benchmark intent expanded into ordinary Slurm run configs.""" + + schema_version: SchemaVersion + name: Identifier + base_run: BenchmarkBaseRun + model_aliases: Literal["all"] | list[str] + concurrency_values: list[PositiveInt] = Field(min_length=1) + deployment_cases: list[BenchmarkDeploymentCase] = Field(min_length=1) + record_policy: BenchmarkRecordPolicy = Field(default_factory=AdaptiveRecordPolicy) + analysis: BenchmarkAnalysisTargets + + @model_validator(mode="after") + def validate_benchmark(self) -> DataDesignerSlurmBenchmarkConfig: + if isinstance(self.model_aliases, list): + if not self.model_aliases: + raise ValueError("model_aliases must not be empty") + if len(self.model_aliases) != len(set(self.model_aliases)): + raise ValueError("benchmark model aliases must be unique") + if len(self.concurrency_values) != len(set(self.concurrency_values)): + raise ValueError("benchmark concurrency values must be unique") + case_names = [case.name for case in self.deployment_cases] + if len(case_names) != len(set(case_names)): + raise ValueError("benchmark deployment case names must be unique") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py new file mode 100644 index 000000000..21670bb5e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure convenience builder for authored Slurm run declarations.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TypeVar + +import yaml +from pydantic import BaseModel, ValidationError + +from data_designer.config import DataDesignerConfigBuilder +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.config.errors import SlurmConfigBuilderError +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + InputBindings, + InvocationConfig, + MCPProviderConfig, + ServerDeploymentConfig, + SubmissionConfig, +) + +_ConfigValueT = TypeVar("_ConfigValueT", bound=BaseModel) + + +class DataDesignerSlurmConfigBuilder: + """Build one strict authored Slurm run without resolving or submitting it.""" + + def __init__(self, builder: BuilderInput, *, name: str = "data-designer") -> None: + self._name = name + self._builder = builder + self._invocation: InvocationConfig | None = None + self._client: ClientConfig | None = None + self._deployments: list[ServerDeploymentConfig] = [] + self._array_tasks = ArrayTasksConfig() + self._submission = SubmissionConfig() + + @classmethod + def from_config_builder( + cls, + builder: DataDesignerConfigBuilder, + *, + name: str = "data-designer", + ) -> DataDesignerSlurmConfigBuilder: + """Start from one public Data Designer configuration builder.""" + return cls(_validate_model(BuilderInput, {"inline": builder.get_builder_config().to_dict()}), name=name) + + @classmethod + def from_builder_source( + cls, + source: str, + *, + name: str = "data-designer", + ) -> DataDesignerSlurmConfigBuilder: + """Start from one local serialized Data Designer builder path.""" + return cls(_validate_model(BuilderInput, {"source": source}), name=name) + + def with_invocation( + self, + *, + num_records: int, + dataset_name: str, + resume: str = "never", + run_config: Mapping[str, object] | None = None, + input_bindings: InputBindings | Mapping[str, object] | None = None, + mcp_providers: Sequence[MCPProviderConfig | Mapping[str, object]] = (), + model_concurrency: Mapping[str, int] | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set typed Data Designer invocation intent.""" + self._invocation = _validate_model( + InvocationConfig, + { + "num_records": num_records, + "dataset_name": dataset_name, + "resume": resume, + "run_config": {} if run_config is None else dict(run_config), + "input_bindings": {} if input_bindings is None else input_bindings, + "mcp_providers": list(mcp_providers), + "model_concurrency": {} if model_concurrency is None else dict(model_concurrency), + }, + ) + return self + + def with_client( + self, + *, + cpus: int = 16, + dependencies: ClientDependencies | Mapping[str, object] | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set the separate zero-GPU Data Designer client declaration.""" + self._client = _validate_model( + ClientConfig, + { + "cpus": cpus, + "dependencies": {} if dependencies is None else dependencies, + }, + ) + return self + + def with_deployment( + self, + deployment: ServerDeploymentConfig | Mapping[str, object], + ) -> DataDesignerSlurmConfigBuilder: + """Append one deployment while preserving authored order.""" + self._deployments.append(_validate_model(ServerDeploymentConfig, deployment)) + return self + + def with_array_tasks( + self, + *, + count: int, + max_concurrent: int | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set deterministic horizontal sharding.""" + self._array_tasks = _validate_model( + ArrayTasksConfig, + {"count": count, "max_concurrent": max_concurrent}, + ) + return self + + def with_submission(self, **values: object) -> DataDesignerSlurmConfigBuilder: + """Set typed Slurm submission intent.""" + self._submission = _validate_model(SubmissionConfig, values) + return self + + def build(self) -> DataDesignerSlurmConfig: + """Return the complete authored declaration without resolving ambient state.""" + missing = [] + if self._invocation is None: + missing.append("invocation") + if not self._deployments: + missing.append("deployment") + if missing: + raise SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") + assert self._invocation is not None + return _validate_model( + DataDesignerSlurmConfig, + { + "schema_version": 1, + "name": self._name, + "builder": self._builder, + "invocation": self._invocation, + "client": self._client or ClientConfig(), + "deployments": self._deployments, + "array_tasks": self._array_tasks, + "submission": self._submission, + }, + ) + + def write_config(self, path: str | Path) -> None: + """Serialize the authored declaration as deterministic JSON or YAML.""" + output_path = Path(path) + config = self.build() + if output_path.suffix == ".json": + contents = config.serialize_json() + elif output_path.suffix in {".yaml", ".yml"}: + contents = yaml.safe_dump( + config.model_dump(mode="json"), + default_flow_style=False, + sort_keys=True, + ) + else: + raise SlurmConfigBuilderError("config path must end in .json, .yaml, or .yml") + try: + output_path.write_text(contents, encoding="utf-8") + except OSError: + raise SlurmConfigBuilderError(f"cannot write Slurm config {output_path}") from None + + +def _validate_model(config_type: type[_ConfigValueT], value: object) -> _ConfigValueT: + try: + return config_type.model_validate(value) + except ValidationError as error: + message = format_validation_error(error, subject=config_type.__name__, models=config_type) + raise SlurmConfigBuilderError(message) from None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py new file mode 100644 index 000000000..006effec8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed environment bindings and secret-shape validation.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping +from typing import Annotated, Literal + +from pydantic import BaseModel, Field, JsonValue, StringConstraints, field_validator + +from data_designer.slurm.contracts import AuthoredConfig, validate_plain_text +from data_designer.slurm.types import EnvironmentName + +__all__ = [ + "EnvironmentBinding", + "LiteralEnvironmentBinding", + "SecretRef", + "collect_secret_environment_names", +] + + +class LiteralEnvironmentBinding(AuthoredConfig): + """Literal, non-secret environment value materialized for a Slurm process. + + Attributes: + type: Discriminator for serialized environment bindings. + value: Plain-text value. Secret-shaped environment names cannot use this binding. + """ + + type: Literal["literal"] + value: Annotated[str, StringConstraints(max_length=4096)] + + @field_validator("value") + @classmethod + def validate_value(cls, value: str) -> str: + return validate_plain_text(value, field_name="environment value") + + +class SecretRef(AuthoredConfig): + """Reference to a login-environment variable containing secret material. + + The referenced value is materialized only at the runtime boundary and is never + persisted into authored or resolved contracts. + + Attributes: + type: Discriminator for serialized environment bindings. + environment: Name of the source variable in the trusted login environment. + """ + + type: Literal["secret"] + environment: EnvironmentName + + +EnvironmentBinding = Annotated[ + LiteralEnvironmentBinding | SecretRef, + Field(discriminator="type"), +] + + +def collect_secret_environment_names(value: object) -> tuple[str, ...]: + """Return every external environment variable referenced below a config value.""" + names: set[str] = set() + _collect_secret_environment_names(value, names) + return tuple(sorted(names)) + + +def is_secret_bearing_name(value: str) -> bool: + """Return whether a Slurm-owned name conventionally carries secret material.""" + segments = _secret_name_segments(value) + return bool( + _SECRET_NAME_PARTS.intersection(segments) + or {"access", "key"}.issubset(segments) + or segments[-1] in {"auth", "key"} + ) + + +def validate_environment_bindings(values: Mapping[EnvironmentName, EnvironmentBinding]) -> None: + """Reject literal values for environment variables that look secret-bearing.""" + if any( + is_secret_bearing_name(name) and isinstance(binding, LiteralEnvironmentBinding) + for name, binding in values.items() + ): + raise ValueError("secret-shaped environment names require external secret references") + + +def validate_no_plaintext_secrets(value: JsonValue, *, field_name: str) -> None: + """Reject non-null values stored under recursively secret-bearing JSON keys.""" + if isinstance(value, dict): + for key, child in value.items(): + if is_secret_bearing_json_key(key) and child is not None: + raise ValueError(f"{field_name} must not contain plaintext values under secret-bearing keys") + validate_no_plaintext_secrets(child, field_name=field_name) + elif isinstance(value, list): + for child in value: + validate_no_plaintext_secrets(child, field_name=field_name) + + +def is_secret_bearing_json_key(value: str) -> bool: + """Return whether a persisted JSON key conventionally carries secret material.""" + segments = _secret_name_segments(value) + return tuple(segments[-2:]) not in _NON_SECRET_KEY_SUFFIXES and is_secret_bearing_name(value) + + +_SECRET_NAME_PARTS = frozenset({"credential", "credentials", "password", "secret", "token"}) +_NON_SECRET_KEY_SUFFIXES = frozenset( + { + ("foreign", "key"), + ("idempotency", "key"), + ("partition", "key"), + ("primary", "key"), + ("sort", "key"), + } +) + + +def _secret_name_segments(value: str) -> list[str]: + snake_case = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", value) + normalized = re.sub(r"[^a-z0-9]+", "_", snake_case.casefold()).strip("_") + return normalized.split("_") + + +def _collect_secret_environment_names(value: object, names: set[str]) -> None: + if isinstance(value, SecretRef): + names.add(value.environment) + elif isinstance(value, BaseModel): + for field_name in type(value).model_fields: + _collect_secret_environment_names(getattr(value, field_name), names) + elif isinstance(value, Mapping): + for child in value.values(): + _collect_secret_environment_names(child, names) + elif isinstance(value, (tuple, list)): + for child in value: + _collect_secret_environment_names(child, names) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py new file mode 100644 index 000000000..e45a33b33 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/errors.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalized authored Slurm configuration errors.""" + +from __future__ import annotations + + +class SlurmConfigError(ValueError): + """Base error for authored Slurm configuration boundaries.""" + + +class SlurmConfigBuilderError(SlurmConfigError): + """Raised when the Slurm config builder is incomplete or cannot serialize.""" + + +class SlurmConfigLoadError(SlurmConfigError): + """Raised when a local Slurm configuration file is not strict and valid.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/images.py b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py new file mode 100644 index 000000000..4c520394f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import re +from enum import Enum +from typing import Annotated, Literal + +from pydantic import Field, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AuthoredConfig, + ContractRecord, + ContractValue, + Identifier, + InstalledDistribution, + Sha256Digest, + validate_absolute_path, + validate_plain_text, +) + + +class ImageKind(str, Enum): + CLIENT = "client" + SERVING = "serving" + + +class ImageRef(AuthoredConfig): + """Authored reference to one registered image alias or SQSH path.""" + + name: Identifier | None = None + path: str | None = None + + @field_validator("path") + @classmethod + def validate_path(cls, value: str | None) -> str | None: + if value is None: + return None + validate_absolute_path(value) + if not value.endswith(".sqsh"): + raise ValueError("image path must end in .sqsh") + return value + + @model_validator(mode="after") + def validate_reference(self) -> ImageRef: + if (self.name is None) == (self.path is None): + raise ValueError("image reference requires exactly one of name or path") + return self + + +class ImageBuildRequest(AuthoredConfig): + """Typed input for one image import or existing-SQSH registration.""" + + name: Identifier + kind: Literal["client", "serving"] + source: str + + @field_validator("source") + @classmethod + def validate_source(cls, value: str) -> str: + validate_plain_text(value, field_name="source") + if value.endswith(".sqsh"): + return validate_absolute_path(value) + if not re.fullmatch(r"[^\s]+@sha256:[0-9a-f]{64}", value): + raise ValueError("OCI image source must be digest-qualified") + return value + + +class ClientImageInspection(ContractValue): + kind: Literal[ImageKind.CLIENT] + python_implementation: Identifier + python_version: Annotated[str, StringConstraints(pattern=r"^[0-9]+\.[0-9]+\.[0-9]+$")] + python_abi: Identifier + distributions: tuple[InstalledDistribution, ...] + installer_path: str + installer_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + + _installer_path_is_absolute = field_validator("installer_path")(validate_absolute_path) + + @model_validator(mode="after") + def validate_distributions(self) -> ClientImageInspection: + names = tuple(distribution.name for distribution in self.distributions) + if len(names) != len(set(names)): + raise ValueError("installed distribution names must be unique") + return self + + +class ServingImageInspection(ContractValue): + kind: Literal[ImageKind.SERVING] + server_type: Literal["vllm"] + runtime_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + executable_path: str + + _executable_path_is_absolute = field_validator("executable_path")(validate_absolute_path) + + +ImageInspection = Annotated[ClientImageInspection | ServingImageInspection, Field(discriminator="kind")] + + +class ImageInspectionRecord(ContractRecord): + """Digest-bound factual inspection output produced inside an SQSH.""" + + inspector_version: Identifier + sqsh_sha256: Sha256Digest + inspection: ImageInspection 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 new file mode 100644 index 000000000..e95bab69c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict local loading and cluster-profile selection.""" + +from __future__ import annotations + +import json +import os +import socket +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import TypeVar, cast + +import yaml +from pydantic import JsonValue, ValidationError +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, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, +) +from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig + +PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" +DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" + +_ConfigT = TypeVar( + "_ConfigT", + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + SlurmProfileCatalog, +) +_HostnameResolver = Callable[[], tuple[str, ...]] + + +class _StrictYamlLoader(yaml.SafeLoader): + pass + + +def _construct_unique_mapping( + loader: _StrictYamlLoader, + node: MappingNode, + deep: bool = False, +) -> dict[object, object]: + mapping: dict[object, object] = {} + for key_node, value_node in node.value: + if key_node.tag == "tag:yaml.org,2002:merge": + raise SlurmConfigLoadError("YAML merge keys are not supported") + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in mapping + except TypeError: + raise SlurmConfigLoadError("configuration mapping keys must be scalar values") from None + if duplicate: + raise SlurmConfigLoadError("duplicate configuration key") + mapping[key] = loader.construct_object(value_node, deep=deep) + return mapping + + +_StrictYamlLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_mapping, +) + + +def load_run_config(path: str | Path) -> DataDesignerSlurmConfig: + """Load one strict local YAML or JSON run declaration.""" + 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) + + +def load_builder_payload(path: str | Path) -> dict[str, JsonValue]: + """Load one complete serialized Data Designer builder config.""" + resolved_path = _normalize_file_path(path) + try: + contents = resolved_path.read_text(encoding="utf-8") + except FileNotFoundError: + raise SlurmConfigLoadError(f"builder source file not found: {resolved_path}") from None + except OSError: + raise SlurmConfigLoadError(f"cannot read configuration file {resolved_path}") from None + try: + payload = _parse_mapping(contents, suffix=resolved_path.suffix) + validated = BuilderInput(inline=payload).inline + assert validated is not None + return validated + except SlurmConfigLoadError: + raise + except ValidationError as error: + message = format_validation_error( + error, + subject=f"builder configuration file {resolved_path}", + models=BuilderInput, + ) + raise SlurmConfigLoadError(message) from None + except (json.JSONDecodeError, yaml.YAMLError) as error: + raise SlurmConfigLoadError(f"invalid configuration file {resolved_path}: {format_parse_error(error)}") from None + + +def resolve_profile( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + hostnames: tuple[str, ...] | None = None, + hostname_resolver: _HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, +) -> SelectedSlurmProfile: + """Resolve an injected profile or select one catalog entry.""" + try: + sources = sum(source is not None for source in (profile, catalog, profile_file)) + if sources > 1: + raise SlurmConfigLoadError("profile, catalog, and profile_file are mutually exclusive") + if profile is not None: + if cluster is not None: + raise SlurmConfigLoadError("an injected profile cannot be combined with cluster selection") + return injected_profile(profile) + + catalog_path: str | None = None + if catalog is None: + path = _resolve_profile_path( + profile_file, + environ=os.environ if environ is None else environ, + home_directory=home_directory, + ) + catalog = load_profile_catalog(path) + catalog_path = path.as_posix() + + if cluster is None and hostnames is None: + resolver = hostname_resolver or _local_hostnames + hostnames = resolver() + normalized_hostnames = tuple( + dict.fromkeys(hostname.strip().casefold() for hostname in (hostnames or ()) if hostname) + ) + return select_profile( + catalog, + cluster=cluster, + hostnames=normalized_hostnames, + catalog_path=catalog_path, + ) + except SlurmConfigLoadError: + raise + except ValidationError as error: + message = format_validation_error(error, subject="profile selection", models=SelectedSlurmProfile) + raise SlurmConfigLoadError(message) from None + except ValueError as error: + raise SlurmConfigLoadError(str(error)) from None + + +def _load_config(path: str | Path, config_type: type[_ConfigT]) -> _ConfigT: + resolved_path = _normalize_file_path(path) + try: + contents = resolved_path.read_text(encoding="utf-8") + except OSError: + raise SlurmConfigLoadError(f"cannot read configuration file {resolved_path}") from None + try: + 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) + except SlurmConfigLoadError: + raise + except ValidationError as error: + message = format_validation_error( + error, + subject=f"configuration file {resolved_path}", + models=config_type, + ) + raise SlurmConfigLoadError(message) from None + except (json.JSONDecodeError, yaml.YAMLError) as error: + raise SlurmConfigLoadError(f"invalid configuration file {resolved_path}: {format_parse_error(error)}") from None + + +def _parse_mapping(contents: str, *, suffix: str) -> dict[str, object]: + if suffix == ".json": + payload = json.loads(contents, object_pairs_hook=_unique_json_object) + else: + events = yaml.parse(contents, Loader=yaml.SafeLoader) + if any(getattr(event, "anchor", None) is not None for event in events): + raise SlurmConfigLoadError("YAML anchors and aliases are not supported") + payload = yaml.load(contents, Loader=_StrictYamlLoader) + if not isinstance(payload, dict) or any(not isinstance(key, str) for key in payload): + raise SlurmConfigLoadError("configuration root must be an object with string keys") + return cast(dict[str, object], payload) + + +def _unique_json_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise SlurmConfigLoadError("duplicate configuration key") + result[key] = value + return result + + +def _reject_environment_interpolation(value: object) -> None: + if isinstance(value, str) and "${" in value: + raise SlurmConfigLoadError("environment interpolation is not supported") + if isinstance(value, Mapping): + for key, item in value.items(): + _reject_environment_interpolation(key) + _reject_environment_interpolation(item) + elif isinstance(value, list | tuple): + for item in value: + _reject_environment_interpolation(item) + + +def _reject_run_environment_interpolation(payload: Mapping[str, object]) -> None: + for key, value in payload.items(): + _reject_environment_interpolation(key) + if key != "builder" or not isinstance(value, Mapping): + _reject_environment_interpolation(value) + continue + for builder_key, builder_value in value.items(): + _reject_environment_interpolation(builder_key) + if builder_key != "inline": + _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, + *, + environ: Mapping[str, str], + home_directory: str | Path | None, +) -> Path: + source = explicit_path + if source is None: + environment_path = environ.get(PROFILE_FILE_ENVIRONMENT) + if environment_path is not None and not environment_path: + raise SlurmConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") + source = environment_path + if source is None: + home = Path.home() if home_directory is None else Path(home_directory) + source = home / DEFAULT_PROFILE_FILE_NAME + return _normalize_file_path(source) + + +def _normalize_file_path(path: str | Path) -> Path: + resolved = Path(path).expanduser().resolve() + if resolved.suffix not in {".json", ".yaml", ".yml"}: + raise SlurmConfigLoadError("configuration path must end in .json, .yaml, or .yml") + return resolved + + +def _local_hostnames() -> tuple[str, ...]: + return socket.gethostname(), socket.getfqdn() diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py new file mode 100644 index 000000000..4c16b4610 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from enum import Enum +from fnmatch import fnmatchcase +from typing import Annotated, Literal + +from pydantic import Field, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AuthoredConfig, + ContractRecord, + Identifier, + SchemaVersion, + Sha256Digest, + compute_canonical_json_sha256, + validate_absolute_path, + validate_plain_text, +) +from data_designer.slurm.types import PartitionSelection + +_ALLOCATION_SCRATCH_CONTAINER_ROOT = "/run/data-designer-slurm" + + +class GpuRequestMode(str, Enum): + GRES = "gres" + VISIBLE = "visible" + + +class SchedulerProfile(AuthoredConfig): + account: Identifier | None = None + partition: PartitionSelection | None = None + mem_per_gpu: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] | None = None + bin_path: str | None = None + + @field_validator("bin_path") + @classmethod + def validate_bin_path(cls, value: str | None) -> str | None: + if value is None: + return None + validated = validate_absolute_path(value) + if ":" in validated: + raise ValueError("scheduler bin path must name one directory") + return validated + + +class ImageBuildProfile(AuthoredConfig): + partition: Identifier + cpus_per_task: PositiveInt + memory: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] + time_limit: Annotated[str, StringConstraints(pattern=r"^[0-9]+:[0-5][0-9]:[0-5][0-9]$")] + + @field_validator("time_limit") + @classmethod + def validate_time_limit(cls, value: str) -> str: + hours, minutes, seconds = (int(component) for component in value.split(":")) + if hours == minutes == seconds == 0: + raise ValueError("image-build time limit must be positive") + return value + + +class ContainerMount(AuthoredConfig): + source: str + target: str + read_only: bool = False + + _paths_are_absolute = field_validator("source", "target")(validate_absolute_path) + + +class SlurmProfile(AuthoredConfig): + """Strict facts for one Slurm cluster.""" + + schema_version: SchemaVersion + host_patterns: list[str] = Field(default_factory=list) + scheduler: SchedulerProfile = Field(default_factory=SchedulerProfile) + gpus_per_node: PositiveInt | Literal["auto"] + workspace_root: str + image_build: ImageBuildProfile + gpu_request_mode: Literal["gres", "visible"] = "gres" + container_mounts: list[ContainerMount] = Field(default_factory=list) + + _workspace_root_is_absolute = field_validator("workspace_root")(validate_absolute_path) + + @field_validator("host_patterns") + @classmethod + def validate_host_patterns(cls, values: list[str]) -> list[str]: + normalized: set[str] = set() + for value in values: + validate_plain_text(value, field_name="host pattern") + _validate_hostname_glob(value) + pattern = value.casefold() + if pattern in normalized: + raise ValueError(f"duplicate hostname glob: {value!r}") + normalized.add(pattern) + return values + + @model_validator(mode="after") + def validate_mounts(self) -> SlurmProfile: + targets = [mount.target for mount in self.container_mounts] + if len(targets) != len(set(targets)): + raise ValueError("container mount targets must be unique") + if any( + target == _ALLOCATION_SCRATCH_CONTAINER_ROOT + or target.startswith(f"{_ALLOCATION_SCRATCH_CONTAINER_ROOT}/") + or _ALLOCATION_SCRATCH_CONTAINER_ROOT.startswith(f"{target}/") + for target in targets + ): + raise ValueError("container mount targets must not overlap allocation scratch") + if self.gpu_request_mode == "visible" and self.scheduler.mem_per_gpu is not None: + raise ValueError("mem_per_gpu requires GRES GPU request mode") + return self + + +class SlurmProfileCatalog(AuthoredConfig): + """Versioned catalog of independently complete cluster profiles.""" + + schema_version: SchemaVersion + default_cluster: Identifier + clusters: dict[Identifier, SlurmProfile] = Field(min_length=1) + + @model_validator(mode="after") + def validate_catalog(self) -> SlurmProfileCatalog: + if self.default_cluster not in self.clusters: + raise ValueError("default_cluster must name a configured cluster") + + patterns: dict[str, str] = {} + for cluster_name, profile in self.clusters.items(): + for pattern in profile.host_patterns: + normalized = pattern.casefold() + if normalized in patterns: + raise ValueError( + f"hostname glob {pattern!r} is duplicated by clusters " + f"{patterns[normalized]!r} and {cluster_name!r}" + ) + patterns[normalized] = cluster_name + return self + + +class ProfileSelectionSource(str, Enum): + EXPLICIT = "explicit" + HOSTNAME = "hostname" + DEFAULT = "default" + INJECTED = "injected" + + +class SelectedSlurmProfile(ContractRecord): + """Selected profile and provenance persisted in a resolved run plan.""" + + cluster_name: Identifier | None = None + selection_source: ProfileSelectionSource + matched_pattern: str | None = None + catalog_path: str | None = None + catalog_sha256: Sha256Digest | None = None + profile_sha256: Sha256Digest + profile: SlurmProfile + + @field_validator("catalog_path") + @classmethod + def validate_catalog_path(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + @model_validator(mode="after") + def validate_selection(self) -> SelectedSlurmProfile: + if self.profile_sha256 != _profile_digest(self.profile): + raise ValueError("profile_sha256 does not match the selected profile") + + catalog_fields = (self.cluster_name, self.catalog_sha256) + if self.selection_source is ProfileSelectionSource.INJECTED: + if any(value is not None for value in (*catalog_fields, self.catalog_path, self.matched_pattern)): + raise ValueError("injected profiles must not contain catalog selection fields") + else: + if any(value is None for value in catalog_fields): + raise ValueError("catalog selections require cluster_name and catalog_sha256") + if self.selection_source is ProfileSelectionSource.HOSTNAME: + if self.matched_pattern is None: + raise ValueError("hostname selection requires matched_pattern") + elif self.matched_pattern is not None: + raise ValueError("only hostname selection may contain matched_pattern") + return self + + +def select_profile( + catalog: SlurmProfileCatalog, + *, + cluster: str | None = None, + hostnames: tuple[str, ...] = (), + catalog_path: str | None = None, +) -> SelectedSlurmProfile: + """Select a profile with explicit, hostname, then default precedence.""" + catalog_sha256 = compute_canonical_json_sha256(catalog.model_dump(mode="json")) + if cluster is not None: + if cluster not in catalog.clusters: + raise ValueError(f"unknown cluster {cluster!r}") + return _selection( + catalog, + cluster, + ProfileSelectionSource.EXPLICIT, + catalog_sha256, + catalog_path=catalog_path, + ) + + normalized_hosts = {hostname.casefold() for hostname in hostnames if hostname} + matches: dict[str, list[str]] = {} + for cluster_name, profile in catalog.clusters.items(): + matching_patterns = sorted( + pattern + for pattern in profile.host_patterns + if any(fnmatchcase(hostname, pattern.casefold()) for hostname in normalized_hosts) + ) + if matching_patterns: + matches[cluster_name] = matching_patterns + + if len(matches) > 1: + raise ValueError(f"hostname matches multiple clusters: {', '.join(sorted(matches))}") + if matches: + selected_name = next(iter(matches)) + return _selection( + catalog, + selected_name, + ProfileSelectionSource.HOSTNAME, + catalog_sha256, + catalog_path=catalog_path, + matched_pattern=matches[selected_name][0], + ) + return _selection( + catalog, + catalog.default_cluster, + ProfileSelectionSource.DEFAULT, + catalog_sha256, + catalog_path=catalog_path, + ) + + +def injected_profile(profile: SlurmProfile) -> SelectedSlurmProfile: + """Record a directly injected effective profile.""" + return SelectedSlurmProfile( + schema_version=1, + selection_source=ProfileSelectionSource.INJECTED, + profile_sha256=_profile_digest(profile), + profile=profile, + ) + + +def validate_selected_profile( + catalog: SlurmProfileCatalog, + selected: SelectedSlurmProfile, +) -> SelectedSlurmProfile: + """Validate a persisted catalog selection against its source catalog.""" + if selected.selection_source is ProfileSelectionSource.INJECTED: + raise ValueError("injected profile selection has no source catalog") + if selected.catalog_sha256 != compute_canonical_json_sha256(catalog.model_dump(mode="json")): + raise ValueError("selected profile catalog digest does not match the catalog") + if selected.cluster_name not in catalog.clusters: + raise ValueError("selected cluster is absent from the catalog") + if selected.profile != catalog.clusters[selected.cluster_name]: + raise ValueError("selected profile does not match its catalog entry") + if selected.selection_source is ProfileSelectionSource.DEFAULT and selected.cluster_name != catalog.default_cluster: + raise ValueError("default profile selection does not match the catalog default") + if ( + selected.selection_source is ProfileSelectionSource.HOSTNAME + and selected.matched_pattern not in selected.profile.host_patterns + ): + raise ValueError("hostname selection pattern is absent from the selected profile") + return selected + + +def _selection( + catalog: SlurmProfileCatalog, + cluster_name: str, + source: ProfileSelectionSource, + catalog_sha256: Sha256Digest, + *, + catalog_path: str | None, + matched_pattern: str | None = None, +) -> SelectedSlurmProfile: + profile = catalog.clusters[cluster_name] + return SelectedSlurmProfile( + schema_version=1, + cluster_name=cluster_name, + selection_source=source, + matched_pattern=matched_pattern, + catalog_path=catalog_path, + catalog_sha256=catalog_sha256, + profile_sha256=_profile_digest(profile), + profile=profile, + ) + + +def _profile_digest(profile: SlurmProfile) -> Sha256Digest: + return compute_canonical_json_sha256(profile.model_dump(mode="json")) + + +def _validate_hostname_glob(value: str) -> None: + if "/" in value or any(character.isspace() for character in value): + raise ValueError(f"invalid hostname glob: {value!r}") + open_class: int | None = None + for index, character in enumerate(value): + if character == "[": + if open_class is not None: + raise ValueError(f"invalid hostname glob: {value!r}") + open_class = index + elif character == "]": + if open_class is None or index == open_class + 1: + raise ValueError(f"invalid hostname glob: {value!r}") + open_class = None + if open_class is not None: + raise ValueError(f"invalid hostname glob: {value!r}") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py new file mode 100644 index 000000000..dbc4c0fde --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Authored Data Designer run, client, deployment, and submission configuration.""" + +from __future__ import annotations + +import posixpath +import re +from typing import Annotated, Literal +from urllib.parse import urlsplit + +from packaging.requirements import InvalidRequirement, Requirement +from packaging.utils import canonicalize_name +from pydantic import ( + Field, + JsonValue, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.config import RunConfig +from data_designer.slurm.config.environment import ( + EnvironmentBinding, + LiteralEnvironmentBinding, + SecretRef, + is_secret_bearing_name, + validate_environment_bindings, + validate_no_plaintext_secrets, +) +from data_designer.slurm.config.vllm import QueueBackpressureConfig, VllmServerConfig +from data_designer.slurm.contracts import ( + AuthoredConfig, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, + validate_url, +) +from data_designer.slurm.types import EnvironmentName, Identifier, PartitionSelection, SchemaVersion + +__all__ = [ + "ArrayTasksConfig", + "BuilderInput", + "ClientConfig", + "ClientDependencies", + "DataDesignerSlurmConfig", + "DeploymentResources", + "DeploymentTopology", + "EnvironmentBinding", + "InputBindings", + "InvocationConfig", + "LiteralEnvironmentBinding", + "LocalStdioMCPProviderConfig", + "MCPProviderConfig", + "QueueBackpressureConfig", + "RemoteMCPProviderConfig", + "SecretRef", + "ServerDeploymentConfig", + "SubmissionConfig", + "VllmServerConfig", +] + + +class BuilderInput(AuthoredConfig): + source: str | None = None + inline: dict[str, JsonValue] | None = None + + @field_validator("source") + @classmethod + def validate_source(cls, value: str | None) -> str | None: + return None if value is None else validate_local_config_path(value) + + @model_validator(mode="after") + def validate_input(self) -> BuilderInput: + if (self.source is None) == (self.inline is None): + raise ValueError("builder requires exactly one of source or inline") + if self.inline is not None: + if not self.inline: + raise ValueError("inline builder input must not be empty") + validate_no_plaintext_secrets(self.inline, field_name="inline builder input") + retired = {"dependencies", "sandbox_config", "server_configs"}.intersection(self.inline) + if retired: + raise ValueError(f"builder input contains retired fields: {', '.join(sorted(retired))}") + if "data_designer" in self.inline: + unknown = set(self.inline).difference({"data_designer", "library_version"}) + library_version = self.inline.get("library_version") + valid = not unknown and isinstance(self.inline["data_designer"], dict) + valid = valid and (library_version is None or isinstance(library_version, str)) + else: + valid = isinstance(self.inline.get("columns"), list) + if not valid: + raise ValueError("inline builder input must be one complete serialized Data Designer config") + return self + + +class InputBindings(AuthoredConfig): + seed_path: str | None = None + managed_assets_path: str | None = None + + @field_validator("seed_path", "managed_assets_path") + @classmethod + def validate_paths(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + +class RemoteMCPProviderConfig(AuthoredConfig): + provider_type: Literal["sse", "streamable_http"] + name: Identifier + endpoint: str + api_key: SecretRef | None = None + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + validate_url(value, field_name="MCP endpoint") + parsed = urlsplit(value) + if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment: + raise ValueError("MCP endpoint must not embed credentials, query parameters, or fragments") + return value + + +class LocalStdioMCPProviderConfig(AuthoredConfig): + provider_type: Literal["stdio"] + name: Identifier + command: str + args: list[str] = Field(default_factory=list) + environment: dict[EnvironmentName, EnvironmentBinding] = Field(default_factory=dict) + + @field_validator("command") + @classmethod + def validate_command(cls, value: str) -> str: + validate_plain_text(value, field_name="MCP command") + if any(character.isspace() for character in value): + raise ValueError("MCP command must be one executable token") + return value + + @field_validator("args") + @classmethod + def validate_args(cls, values: list[str]) -> list[str]: + for value in values: + validate_plain_text(value, field_name="MCP argument") + option = re.split(r"[=\s]", value.lstrip(), maxsplit=1)[0].lstrip("-") + if is_secret_bearing_name(option): + raise ValueError("secret-shaped MCP arguments must use an environment secret reference") + return values + + @field_validator("environment") + @classmethod + def validate_environment( + cls, values: dict[EnvironmentName, EnvironmentBinding] + ) -> dict[EnvironmentName, EnvironmentBinding]: + validate_environment_bindings(values) + return values + + +MCPProviderConfig = Annotated[ + RemoteMCPProviderConfig | LocalStdioMCPProviderConfig, + Field(discriminator="provider_type"), +] + + +class InvocationConfig(AuthoredConfig): + num_records: PositiveInt + dataset_name: Identifier + resume: Literal["never", "always", "if_possible"] = "never" + run_config: dict[str, JsonValue] = Field(default_factory=dict) + input_bindings: InputBindings = Field(default_factory=InputBindings) + mcp_providers: list[MCPProviderConfig] = Field(default_factory=list) + model_concurrency: dict[str, PositiveInt] = Field(default_factory=dict) + + @field_validator("run_config") + @classmethod + def validate_run_config_keys(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + unknown = set(value).difference(RunConfig.model_fields.keys() - {"display_tui"}) + if unknown: + raise ValueError(f"unknown Data Designer RunConfig fields: {', '.join(sorted(unknown))}") + RunConfig.model_validate(value) + return value + + @model_validator(mode="after") + def validate_mcp_providers(self) -> InvocationConfig: + names = [provider.name for provider in self.mcp_providers] + if len(names) != len(set(names)): + raise ValueError("MCP provider names must be unique") + return self + + +class ClientDependencies(AuthoredConfig): + requirements: list[str] | None = Field(default_factory=list) + lock_file: str | None = None + index_credentials: dict[str, SecretRef] = Field(default_factory=dict) + + @field_validator("requirements") + @classmethod + def validate_requirements(cls, values: list[str] | None) -> list[str] | None: + if values is None: + return None + names: list[str] = [] + for value in values: + validate_plain_text(value, field_name="dependency requirement") + if value != value.strip() or value.startswith(("-e ", "/", "./", "../")) or "git+" in value: + raise ValueError("dependency requirement must identify a package or immutable wheel") + try: + requirement = Requirement(value) + except InvalidRequirement as error: + raise ValueError("invalid dependency requirement") from error + if requirement.marker is not None: + raise ValueError("dependency requirement environment markers are not supported") + if requirement.url is not None: + validate_url(requirement.url, field_name="direct dependency URL") + parsed = urlsplit(requirement.url) + valid_wheel = ( + parsed.scheme == "https" + and parsed.hostname is not None + and parsed.username is None + and parsed.password is None + and not parsed.query + and parsed.path.endswith(".whl") + and re.fullmatch(r"sha256=[0-9a-f]{64}", parsed.fragment) is not None + ) + if not valid_wheel: + raise ValueError("direct dependency URLs must be HTTPS wheels with a SHA-256 fragment") + names.append(canonicalize_name(requirement.name)) + if len(names) != len(set(names)): + raise ValueError("dependency requirements must have unique normalized names") + return values + + @field_validator("lock_file") + @classmethod + def validate_lock_file(cls, value: str | None) -> str | None: + if value is None: + return None + validate_plain_text(value, field_name="dependency lock path") + if "://" in value or ".." in value.split("/"): + raise ValueError("dependency lock must be a normalized local path") + normalized = posixpath.normpath(value) + if not normalized.endswith(".json"): + raise ValueError("dependency lock path must end in .json") + return normalized + + @model_validator(mode="after") + def validate_source(self) -> ClientDependencies: + if self.lock_file is None and self.requirements is None: + raise ValueError("client dependencies require requirements or lock_file") + if self.lock_file is not None and self.requirements is not None: + raise ValueError("client dependencies cannot contain both requirements and lock_file") + return self + + +class ClientConfig(AuthoredConfig): + cpus: PositiveInt = 16 + dependencies: ClientDependencies = Field(default_factory=ClientDependencies) + + +class DeploymentResources(AuthoredConfig): + nodes: PositiveInt = 1 + shared_node: Identifier | None = Field(default=None, exclude_if=lambda value: value is None) + gpus_per_node: PositiveInt | None = Field(default=None, exclude_if=lambda value: value is None) + gpu_start_index: NonNegativeInt | None = Field(default=None, exclude_if=lambda value: value is None) + + @model_validator(mode="after") + def validate_shared_node(self) -> DeploymentResources: + if (self.shared_node is None) != (self.gpus_per_node is None): + raise ValueError("shared_node and gpus_per_node must be provided together") + if self.shared_node is not None and self.nodes != 1: + raise ValueError("shared_node requires exactly one deployment node") + if self.gpu_start_index is not None and self.shared_node is None: + raise ValueError("gpu_start_index requires shared_node") + return self + + +class DeploymentTopology(AuthoredConfig): + tensor_parallel: PositiveInt = 1 + nodes_per_replica: PositiveInt = 1 + + +class ServerDeploymentConfig(AuthoredConfig): + model_alias: str + served_model_name: str | None = None + model: str + server: VllmServerConfig + resources: DeploymentResources = Field(default_factory=DeploymentResources) + topology: DeploymentTopology = Field(default_factory=DeploymentTopology) + + @field_validator("model") + @classmethod + def validate_model(cls, value: str) -> str: + validate_plain_text(value, field_name="model") + if value.startswith("/"): + return validate_absolute_path(value) + if any(character.isspace() for character in value): + raise ValueError("Hugging Face model identifiers must not contain whitespace") + return value + + @field_validator("served_model_name") + @classmethod + def validate_served_model_name(cls, value: str | None) -> str | None: + return None if value is None else validate_plain_text(value, field_name="served model name") + + @model_validator(mode="after") + def validate_topology(self) -> ServerDeploymentConfig: + if self.resources.shared_node is not None: + if self.topology.nodes_per_replica != 1: + raise ValueError("shared_node supports only single-node replicas") + quota = self.resources.gpus_per_node + assert quota is not None + if quota % self.topology.tensor_parallel: + raise ValueError("tensor_parallel must divide the deployment GPU quota") + if self.resources.nodes % self.topology.nodes_per_replica: + raise ValueError("nodes_per_replica must divide deployment nodes") + if self.server.enable_expert_parallel and self.topology.nodes_per_replica > 1: + raise ValueError("multi-node expert parallel is not supported in v1") + return self + + +class ArrayTasksConfig(AuthoredConfig): + count: PositiveInt = 1 + max_concurrent: PositiveInt | None = None + + @model_validator(mode="after") + def validate_concurrency(self) -> ArrayTasksConfig: + if self.max_concurrent is not None and self.max_concurrent > self.count: + raise ValueError("array task concurrency must not exceed task count") + return self + + +class SubmissionConfig(AuthoredConfig): + account: Identifier | None = None + partition: PartitionSelection | None = None + job_name: Identifier = "data-designer" + time_limit: Annotated[str, StringConstraints(pattern=r"^(?:[0-9]+-)?[0-9]{2}:[0-9]{2}:[0-9]{2}$")] = "03:55:00" + comment: Annotated[str, StringConstraints(max_length=256)] | None = None + + @field_validator("time_limit") + @classmethod + def validate_time_limit(cls, value: str) -> str: + clock = value.rsplit("-", maxsplit=1)[-1] + _, minutes, seconds = (int(part) for part in clock.split(":")) + if minutes >= 60 or seconds >= 60: + raise ValueError("time_limit minutes and seconds must be below 60") + return value + + @field_validator("comment") + @classmethod + def validate_comment(cls, value: str | None) -> str | None: + return None if value is None else validate_plain_text(value, field_name="submission comment") + + +class DataDesignerSlurmConfig(AuthoredConfig): + """Complete portable intent for one Data Designer Slurm run.""" + + schema_version: SchemaVersion + name: Identifier + builder: BuilderInput + invocation: InvocationConfig + client: ClientConfig = Field(default_factory=ClientConfig) + deployments: list[ServerDeploymentConfig] = Field(min_length=1) + array_tasks: ArrayTasksConfig = Field(default_factory=ArrayTasksConfig) + submission: SubmissionConfig = Field(default_factory=SubmissionConfig) + + @model_validator(mode="after") + def validate_run(self) -> DataDesignerSlurmConfig: + if self.array_tasks.count > self.invocation.num_records: + raise ValueError("array task count must not exceed requested records") + aliases = [deployment.model_alias for deployment in self.deployments] + if len(aliases) != len(set(aliases)): + raise ValueError("deployment model aliases must be unique") + groups: dict[str, list[ServerDeploymentConfig]] = {} + for deployment in self.deployments: + if deployment.resources.shared_node is not None: + groups.setdefault(deployment.resources.shared_node, []).append(deployment) + if groups and ( + len(groups) != 1 or any(deployment.resources.shared_node is None for deployment in self.deployments) + ): + raise ValueError("shared_node requires one group covering all deployments") + for members in groups.values(): + if len(members) < 2: + raise ValueError("shared_node requires at least two deployments in each group") + first = members[0].server + if any( + member.server.environment != first.environment + or member.server.queue_backpressure != first.queue_backpressure + for member in members[1:] + ): + raise ValueError("shared_node requires matching server environment and queue backpressure") + unknown_concurrency = set(self.invocation.model_concurrency).difference(aliases) + if unknown_concurrency: + raise ValueError( + f"model concurrency references undeclared aliases: {', '.join(sorted(unknown_concurrency))}" + ) + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/utils.py b/packages/data-designer-slurm/src/data_designer/slurm/config/utils.py new file mode 100644 index 000000000..ebfba1845 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/utils.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small shared utilities for authored Slurm configuration.""" + +from __future__ import annotations + + +def convert_duration_to_seconds(value: str) -> int: + """Convert a validated Slurm duration to whole seconds.""" + factor = {"s": 1, "m": 60, "h": 3600, "d": 86400}[value[-1]] + return int(value[:-1]) * factor diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py new file mode 100644 index 000000000..2db28d086 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Authored vLLM server configuration.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.environment import EnvironmentBinding +from data_designer.slurm.config.images import ImageRef +from data_designer.slurm.config.utils import convert_duration_to_seconds +from data_designer.slurm.config.vllm_validation import ( + validate_vllm_environment_bindings, + validate_vllm_extra_args, + validate_vllm_readiness_path, +) +from data_designer.slurm.contracts import AuthoredConfig +from data_designer.slurm.types import Duration, EnvironmentName, Identifier, NonNegativeDuration + +__all__ = [ + "QueueBackpressureConfig", + "VllmServerConfig", +] + + +class QueueBackpressureConfig(AuthoredConfig): + """Admission policy applied by the logical endpoint under queue pressure. + + Attributes: + max_waiting_requests: Maximum queued requests before admission is rejected. + retry_after_seconds: Retry delay returned to callers, or ``None`` to omit it. + """ + + max_waiting_requests: NonNegativeInt = 128 + retry_after_seconds: PositiveInt | None = 1 + + +class VllmServerConfig(AuthoredConfig): + """Authored vLLM image, lifecycle, and safe launch customization. + + Duration fields use Slurm duration strings; resolution converts them to integer + seconds. ``extra_args`` contains individual argv items, never shell fragments. + + Attributes: + type: Serving-backend discriminator. + image: Authored image reference to inspect and resolve. + startup_timeout: Total deadline for backend readiness. + distributed_init_timeout: Deadline for distributed initialization. + lead_boot_standoff: Delay before non-leading replicas launch. + rank_launch_stagger: Additional delay per deployment-wide replica index. + readiness_path: Absolute HTTP path used for readiness probes. + enable_expert_parallel: Whether vLLM expert parallelism is requested. + io_processor_plugin: Registered vLLM IO processor plugin name. + queue_backpressure: Logical-endpoint queue admission policy. + extra_args: Additional shell-free vLLM argv items not reserved by orchestration. + environment: Explicit literal or secret-reference environment bindings. + """ + + type: Literal["vllm"] + image: ImageRef + startup_timeout: Duration = "15m" + distributed_init_timeout: Duration = "10m" + lead_boot_standoff: NonNegativeDuration = "60s" + rank_launch_stagger: NonNegativeDuration = "5s" + readiness_path: str = "/health" + enable_expert_parallel: bool = False + io_processor_plugin: Identifier | None = Field(default=None, exclude_if=lambda value: value is None) + queue_backpressure: QueueBackpressureConfig = Field(default_factory=QueueBackpressureConfig) + extra_args: list[str] = Field(default_factory=list) + environment: dict[EnvironmentName, EnvironmentBinding] = Field(default_factory=dict) + + @field_validator("readiness_path") + @classmethod + def validate_readiness_path(cls, value: str) -> str: + return validate_vllm_readiness_path(value) + + @field_validator("extra_args") + @classmethod + def validate_extra_args(cls, values: list[str]) -> list[str]: + validate_vllm_extra_args(values) + return values + + @field_validator("environment") + @classmethod + def validate_environment( + cls, values: dict[EnvironmentName, EnvironmentBinding] + ) -> dict[EnvironmentName, EnvironmentBinding]: + validate_vllm_environment_bindings(values) + return values + + @model_validator(mode="after") + def validate_timeouts(self) -> VllmServerConfig: + if convert_duration_to_seconds(self.distributed_init_timeout) > convert_duration_to_seconds( + self.startup_timeout + ): + raise ValueError("distributed_init_timeout must not exceed startup_timeout") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/vllm_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm_validation.py new file mode 100644 index 000000000..4ba5bcd6a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm_validation.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validation and normalization owned by the vLLM configuration boundary.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Mapping, Sequence + +from data_designer.slurm.config.environment import ( + EnvironmentBinding, + is_secret_bearing_json_key, + is_secret_bearing_name, + validate_environment_bindings, + validate_no_plaintext_secrets, +) +from data_designer.slurm.contracts import validate_plain_text +from data_designer.slurm.types import EnvironmentName + + +def validate_vllm_environment_bindings(values: Mapping[EnvironmentName, EnvironmentBinding]) -> None: + """Reject environment names derived by the serving compiler or runtime.""" + validate_environment_bindings(values) + conflicts = sorted( + name + for name in values + if name in _RESERVED_VLLM_ENVIRONMENT_NAMES or name.startswith(_RESERVED_VLLM_ENVIRONMENT_PREFIXES) + ) + if conflicts: + raise ValueError(f"vLLM environment names are owned by the compiler or runtime: {', '.join(conflicts)}") + + +def validate_vllm_readiness_path(value: str) -> str: + """Require an HTTP path that cannot be interpreted as another authority.""" + validate_plain_text(value, field_name="readiness path") + if ( + not value.startswith("/") + or value.startswith("//") + or any(character.isspace() for character in value) + or "?" in value + or "#" in value + ): + raise ValueError("readiness path must be an absolute URL path without whitespace, query, or fragment") + return value + + +def validate_vllm_extra_args(values: Sequence[str]) -> None: + """Validate shell-free vLLM argv without interpreting plugin-owned values.""" + seen_flags: set[str] = set() + for value in values: + validate_plain_text(value, field_name="vLLM argument") + if value != value.strip(): + raise ValueError("vLLM arguments must not contain leading or trailing whitespace") + flag = value.partition("=")[0].replace("_", "-") + if not flag: + raise ValueError("vLLM argument option name must not be empty") + candidate_flag = flag.split(maxsplit=1)[0] + if flag == "--": + raise ValueError("vLLM argument option terminators are not supported") + if _is_reserved_vllm_flag(candidate_flag): + raise ValueError(f"vLLM argument {candidate_flag!r} is owned by the compiler or runtime") + if candidate_flag.startswith("--") and is_secret_bearing_name(candidate_flag.removeprefix("--")): + raise ValueError("secret-shaped vLLM arguments must use an environment secret reference") + _validate_vllm_argument_has_no_plaintext_secrets(value) + if flag.startswith("-") and re.fullmatch(r"--[A-Za-z][A-Za-z0-9.-]*", flag) is None: + raise ValueError( + "short or combined vLLM options are owned by the compiler or runtime; " + "each extra option must use one complete long flag in one token" + ) + if flag.startswith("--"): + canonical_flag = f"--{flag.removeprefix('--no-')}" if flag.startswith("--no-") else flag + if canonical_flag in seen_flags: + raise ValueError(f"duplicate or conflicting vLLM argument {flag!r}") + seen_flags.add(canonical_flag) + + +def _validate_vllm_argument_has_no_plaintext_secrets(value: str) -> None: + payload = value.partition("=")[2] if value.startswith("--") and "=" in value else value + try: + decoded = json.loads(payload) + except (json.JSONDecodeError, RecursionError): + pass + else: + validate_no_plaintext_secrets(decoded, field_name="vLLM argument payload") + if any(is_secret_bearing_json_key(match.group("name")) for match in _ARGUMENT_ASSIGNMENT_PATTERN.finditer(payload)): + raise ValueError("vLLM argument payloads must not contain plaintext values under secret-bearing keys") + + +_RESERVED_VLLM_FLAGS = frozenset( + { + "--api-key", + "--api-server-count", + "--config", + "--cpu-distributed-timeout-seconds", + "--cp-kv-cache-interleave-size", + "--cpunodebind", + "--dcp-comm-backend", + "--dcp-kv-cache-interleave-size", + "--default-mm-loras", + "--distributed-executor-backend", + "--distributed-init-address", + "--ec-transfer-config", + "--enable-elastic-ep", + "--enable-expert-parallel", + "--enable-lora", + "--enable-ssl-refresh", + "--grpc", + "--headless", + "--host", + "--io-processor-plugin", + "--kv-events-config", + "--kv-transfer-config", + "--logits-processors", + "--lora-modules", + "--master-addr", + "--master-port", + "--max-cpu-loras", + "--max-loras", + "--middleware", + "--model", + "--nnodes", + "--node-rank", + "--no-data-parallel-external-lb", + "--no-data-parallel-hybrid-lb", + "--no-enable-elastic-ep", + "--no-enable-expert-parallel", + "--no-enable-lora", + "--no-enable-ssl-refresh", + "--no-numa-bind", + "--physcpubind", + "--pipeline-parallel-size", + "--port", + "--reasoning-parser-plugin", + "--root-path", + "--served-model-name", + "--spec-model", + "--speculative-config", + "--tensor-parallel-size", + "--tool-parser-plugin", + "--uds", + "--weight-transfer-config", + "--worker-cls", + "--worker-extension-cls", + } +) +_VLLM_EXACT_PASSTHROUGH_FLAGS = frozenset({"--reasoning-parser"}) +_ARGUMENT_ASSIGNMENT_PATTERN = re.compile(r"""(?:^|[,{;?&\s])["']?(?P[A-Za-z][A-Za-z0-9_.-]*)["']?\s*[:=]""") +_RESERVED_VLLM_FLAG_PREFIXES = ( + "--api-server-", + "--data-parallel-", + "--decode-context-parallel-", + "--distributed-", + "--dp-supervisor-", + "--master-", + "--numa-", + "--pipeline-parallel-", + "--prefill-context-parallel-", + "--ssl-", + "--tensor-parallel-", +) +_RESERVED_VLLM_ENVIRONMENT_NAMES = frozenset( + { + "CACHE_ROOT", + "CUDA_VISIBLE_DEVICES", + "DP_GLOBAL", + "DP_LOCAL", + "ENABLE_EP", + "GEN_WORKFLOW", + "GLOBAL_REPLICA_COUNT", + "GROUP_RANK", + "HEAD_IP", + "HEAD_NODE_NAME", + "HF_HOME", + "LOCAL_RANK", + "LOCAL_WORLD_SIZE", + "LOG_DIR", + "MASTER_ADDR", + "MASTER_PORT", + "NODE_RANK", + "NODE_GROUP_COUNT", + "NVIDIA_VISIBLE_DEVICES", + "PIPELINE_PARALLEL_SIZE", + "PYTHON_EXEC", + "RANK", + "REPLICAS_PER_NODE_GROUP", + "ROLE_RANK", + "ROLE_WORLD_SIZE", + "SCRATCH_ROOT", + "SERVER_HTTP_PORT", + "SERVICE_DP_RPC_PORT", + "SERVICE_GPU_IDS", + "SERVICE_HTTP_PORT_BASE", + "SERVICE_MASTER_PORT_BASE", + "TENSOR_PARALLEL_SIZE", + "VLLM_ALLOW_RUNTIME_LORA_UPDATING", + "VLLM_API_KEY", + "VLLM_CACHE_ROOT", + "VLLM_DISTRIBUTED_TIMEOUT_SECONDS", + "VLLM_HOST_IP", + "VLLM_LEAD_BOOT_STANDOFF_SECONDS", + "VLLM_MASTER_PORT_BASE", + "VLLM_MODEL_REDIRECT_PATH", + "VLLM_PLUGINS", + "VLLM_PORT", + "VLLM_RANK_LAUNCH_STAGGER_SECONDS", + "VLLM_RPC_BASE_PATH", + "WORLD_SIZE", + } +) +_RESERVED_VLLM_ENVIRONMENT_PREFIXES = ( + "SLURM_", + "TORCHELASTIC_", + "VLLM_DP_", + "VLLM_LORA_", + "VLLM_MOONCAKE_", + "VLLM_NIXL_", + "VLLM_RAY_", +) + + +def _is_reserved_vllm_flag(flag: str) -> bool: + if flag in _VLLM_EXACT_PASSTHROUGH_FLAGS: + return False + return ( + flag in _RESERVED_VLLM_FLAGS + or flag.startswith(_RESERVED_VLLM_FLAG_PREFIXES) + or (flag.startswith("--") and any(reserved_flag.startswith(flag) for reserved_flag in _RESERVED_VLLM_FLAGS)) + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py new file mode 100644 index 000000000..4e511d6ca --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared immutable contracts for Slurm planning and runtime state.""" + +from __future__ import annotations + +import hashlib +import json +import posixpath +from collections.abc import Mapping +from typing import Annotated, TypeVar +from urllib.parse import urlsplit + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +from data_designer.slurm.types import ( + AttemptId, + Duration, + EnvironmentName, + Identifier, + NetworkPort, + NonNegativeDuration, + PartitionSelection, + SchemaVersion, + Sha256Digest, + ShardId, +) + +# TODO: Remove these compatibility exports after Stage 2 branches import shared scalars from slurm.types. +ModelAlias = str + +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") + + +class _FrozenList(list[_ValueT]): + """List that retains JSON compatibility without exposing mutation.""" + + def _immutable(self, *args: object, **kwargs: object) -> None: + del args, kwargs + raise TypeError("frozen list cannot be modified") + + __delitem__ = _immutable + __iadd__ = _immutable + __imul__ = _immutable + __setitem__ = _immutable + append = _immutable + clear = _immutable + extend = _immutable + insert = _immutable + pop = _immutable + remove = _immutable + reverse = _immutable + sort = _immutable + + +class _FrozenDict(dict[_KeyT, _ValueT]): + """Dictionary that retains JSON compatibility without exposing mutation.""" + + def _immutable(self, *args: object, **kwargs: object) -> None: + del args, kwargs + raise TypeError("frozen dictionary cannot be modified") + + __delitem__ = _immutable + __ior__ = _immutable + __setitem__ = _immutable + clear = _immutable + pop = _immutable + popitem = _immutable + setdefault = _immutable + update = _immutable + + +def _freeze_collections(value: object) -> object: + if isinstance(value, Mapping): + return _FrozenDict({key: _freeze_collections(item) for key, item in value.items()}) + if isinstance(value, list): + return _FrozenList(_freeze_collections(item) for item in value) + if isinstance(value, tuple): + return tuple(_freeze_collections(item) for item in value) + return value + + +class ContractValue(BaseModel): + """Base for strict immutable values shared across Slurm boundaries.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + hide_input_in_errors=True, + allow_inf_nan=False, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + @field_validator("*", mode="after") + @classmethod + def freeze_collections(cls, value: object) -> object: + return _freeze_collections(value) + + +DistributionName = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[a-z0-9]+(?:-[a-z0-9]+)*$", + ), +] + + +class InstalledDistribution(ContractValue): + name: DistributionName + version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + provenance_sha256: Sha256Digest | None = Field(default=None, exclude_if=lambda value: value is None) + + +class AuthoredConfig(ContractValue): + """Base for strict authored configuration values.""" + + def serialize_canonical_json(self) -> bytes: + return canonical_json(self.model_dump(mode="json")) + + def serialize_json(self) -> str: + return pretty_json(self.model_dump(mode="json")) + + def compute_sha256(self) -> Sha256Digest: + """Compute the digest of the exact bytes written by ``serialize_json``.""" + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() + + +class ContractRecord(ContractValue): + """Base for immutable, explicitly versioned Slurm records.""" + + schema_version: SchemaVersion + + def serialize_canonical_json(self) -> bytes: + """Serialize the record to stable compact UTF-8 bytes.""" + return canonical_json(self.model_dump(mode="json")) + + def serialize_json(self) -> str: + """Serialize the record to deterministic persisted text.""" + return pretty_json(self.model_dump(mode="json")) + + def compute_sha256(self) -> Sha256Digest: + """Compute the digest of the exact bytes written by ``serialize_json``.""" + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() + + +def canonical_json(value: object) -> bytes: + """Serialize a JSON-compatible value to stable UTF-8 bytes.""" + return json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + +def pretty_json(value: object) -> str: + """Serialize a JSON-compatible value to deterministic persisted text.""" + return ( + json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + +def compute_canonical_json_sha256(value: object) -> Sha256Digest: + """Compute the canonical JSON digest of a JSON-compatible value.""" + return hashlib.sha256(canonical_json(value)).hexdigest() + + +def compute_serialized_json_sha256(value: object) -> Sha256Digest: + """Compute the digest of deterministic persisted JSON bytes.""" + return hashlib.sha256(pretty_json(value).encode("utf-8")).hexdigest() + + +def derive_managed_assets_path(workspace_root: str) -> str: + """Derive the default managed-assets path from a workspace root.""" + return posixpath.join(workspace_root, "managed-assets") + + +def is_path_below(path: str, root: str) -> bool: + """Return whether a path is strictly below a root.""" + return path != root and posixpath.commonpath((path, root)) == root + + +def paths_overlap(left: str, right: str) -> bool: + """Return whether either path contains the other.""" + return left == right or is_path_below(left, right) or is_path_below(right, left) + + +# TODO: Remove after in-flight Stage 2 branches adopt the explicit canonical-JSON digest name. +compute_sha256 = compute_canonical_json_sha256 + + +def validate_absolute_path(value: str) -> str: + """Validate a normalized, absolute POSIX path below the filesystem root.""" + if not value.startswith("/"): + raise ValueError("path must be absolute") + if value.startswith("//"): + raise ValueError("path must have exactly one leading slash") + if value == "/": + raise ValueError("path must not be the filesystem root") + validate_plain_text(value, field_name="path") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + if posixpath.normpath(value) != value: + raise ValueError("path must be normalized") + return value + + +def validate_relative_path(value: str) -> str: + """Validate a normalized relative POSIX path without parent traversal.""" + if not value or value.startswith("/"): + raise ValueError("path must be a non-empty relative path") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("path must not contain control characters") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + if posixpath.normpath(value) != value or value == ".": + raise ValueError("path must be normalized") + return value + + +def validate_local_config_path(value: str) -> str: + """Validate a local JSON or YAML configuration path.""" + validate_plain_text(value, field_name="path") + if "://" in value: + raise ValueError("builder and config sources must be local paths") + if ".." in value.split("/"): + raise ValueError("path must not contain parent-directory components") + normalized = posixpath.normpath(value) + if posixpath.splitext(normalized)[1] not in {".json", ".yaml", ".yml"}: + raise ValueError("config path must end in .json, .yaml, or .yml") + return normalized + + +def validate_plain_text(value: str, *, field_name: str) -> str: + """Reject empty text and control characters at persisted boundaries.""" + if not value: + raise ValueError(f"{field_name} must not be empty") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError(f"{field_name} must not contain control characters") + return value + + +def validate_url(value: str, *, field_name: str) -> str: + """Validate an HTTP(S) URL with a valid host and port.""" + validate_plain_text(value, field_name=field_name) + try: + parsed = urlsplit(value) + parsed.port + except ValueError as error: + raise ValueError(f"{field_name} must be an HTTP(S) URL with a valid host and port") from error + if ( + parsed.scheme not in {"http", "https"} + or parsed.hostname is None + or any(character.isspace() for character in value) + ): + raise ValueError(f"{field_name} must be an HTTP(S) URL") + return value + + +class ArtifactReference(ContractValue): + """Immutable reference to persisted file bytes and their digest.""" + + path: str + sha256: Sha256Digest + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +class RecordRange(ContractValue): + """Half-open global record range assigned to one shard.""" + + start_index: NonNegativeInt + end_index_exclusive: PositiveInt + + @property + def record_count(self) -> int: + return self.end_index_exclusive - self.start_index + + @model_validator(mode="after") + def validate_bounds(self) -> RecordRange: + if self.end_index_exclusive <= self.start_index: + raise ValueError("end_index_exclusive must be greater than start_index") + return self + + +class ResumeWorkspace(ContractValue): + """Canonical shard-owned dataset workspace.""" + + path: str + + _path_is_absolute = field_validator("path")(validate_absolute_path) + + +__all__ = [ + "ArtifactReference", + "AttemptId", + "AuthoredConfig", + "ContractRecord", + "ContractValue", + "Duration", + "EnvironmentName", + "Identifier", + "ModelAlias", + "NetworkPort", + "NonNegativeDuration", + "PartitionSelection", + "RecordRange", + "ResumeWorkspace", + "SchemaVersion", + "Sha256Digest", + "ShardId", + "canonical_json", + "compute_canonical_json_sha256", + "compute_serialized_json_sha256", + "compute_sha256", + "derive_managed_assets_path", + "is_path_below", + "paths_overlap", + "pretty_json", + "validate_absolute_path", + "validate_local_config_path", + "validate_plain_text", + "validate_relative_path", + "validate_url", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py new file mode 100644 index 000000000..8dfddb188 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared descriptor-bound filesystem primitives for package-owned state.""" + +from __future__ import annotations + +import fcntl +import os +import secrets +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +PRIVATE_DIRECTORY_MODE = 0o700 +PRIVATE_FILE_MODE = 0o600 + + +@contextmanager +def acquire_restrictive_file_lock( + directory_descriptor: int, + name: str, + display_path: Path, + *, + resource_name: str, + shared: bool = False, +) -> Iterator[None]: + """Acquire one private single-link lock beneath a verified directory.""" + descriptor: int | None = None + try: + descriptor = open_lock_file(directory_descriptor, name, display_path, resource_name=resource_name) + fcntl.flock(descriptor, fcntl.LOCK_SH if shared else fcntl.LOCK_EX) + locked_status = os.fstat(descriptor) + current_status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + get_file_facts(locked_status)[:2] != get_file_facts(current_status)[:2] + or not stat.S_ISREG(current_status.st_mode) + or current_status.st_nlink != 1 + ): + raise OSError(f"{resource_name} lock {display_path} changed while it was being acquired") + yield + finally: + if descriptor is not None: + try: + fcntl.flock(descriptor, fcntl.LOCK_UN) + finally: + os.close(descriptor) + + +def open_lock_file( + directory_descriptor: int, + name: str, + display_path: Path, + *, + resource_name: str, +) -> int: + """Open and restrict one package-owned lock file.""" + for _ in range(10): + try: + descriptor = os.open( + name, + os.O_CREAT | os.O_RDWR | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + PRIVATE_FILE_MODE, + dir_fd=directory_descriptor, + ) + except FileNotFoundError: + continue + try: + status = os.fstat(descriptor) + if not stat.S_ISREG(status.st_mode) or status.st_nlink != 1: + raise OSError(f"{resource_name} lock {display_path} is not a regular file") + os.fchmod(descriptor, PRIVATE_FILE_MODE) + if os.fstat(descriptor).st_mode & 0o077: + raise OSError(f"{resource_name} lock {display_path} does not have restrictive permissions") + except BaseException: + os.close(descriptor) + raise + return descriptor + raise OSError(f"{resource_name} lock {display_path} could not be opened") + + +@contextmanager +def open_verified_directory( + path: Path, + *, + resource_name: str, + require_private: bool = False, +) -> Iterator[int]: + """Open one real directory and reject replacement during the open.""" + descriptor: int | None = None + try: + before_open = path.lstat() + if not stat.S_ISDIR(before_open.st_mode): + raise OSError(f"{resource_name} directory {path} is not a directory") + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(path, flags) + after_open = os.fstat(descriptor) + validate_open_directory(before_open, after_open, path, resource_name=resource_name) + if require_private and after_open.st_mode & 0o077: + raise OSError(f"{resource_name} directory {path} does not have restrictive permissions") + yield descriptor + finally: + if descriptor is not None: + os.close(descriptor) + + +@contextmanager +def open_verified_child_directory( + parent_descriptor: int, + name: str, + display_path: Path, + *, + resource_name: str, + require_private: bool = False, +) -> Iterator[int]: + """Open one real child directory relative to a verified parent.""" + descriptor: int | None = None + try: + before_open = os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + if not stat.S_ISDIR(before_open.st_mode): + raise OSError(f"{resource_name} directory {display_path} is not a directory") + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(name, flags, dir_fd=parent_descriptor) + after_open = os.fstat(descriptor) + validate_open_directory(before_open, after_open, display_path, resource_name=resource_name) + if require_private and after_open.st_mode & 0o077: + raise OSError(f"{resource_name} directory {display_path} does not have restrictive permissions") + yield descriptor + finally: + if descriptor is not None: + os.close(descriptor) + + +def validate_open_directory( + before_open: os.stat_result, + after_open: os.stat_result, + display_path: Path, + *, + resource_name: str, +) -> None: + """Require an opened descriptor to identify the directory that was inspected.""" + if get_file_facts(before_open)[:2] != get_file_facts(after_open)[:2]: + raise OSError(f"{resource_name} directory {display_path} changed while it was being opened") + if not stat.S_ISDIR(after_open.st_mode): + raise OSError(f"{resource_name} directory {display_path} is not a directory") + + +def create_restrictive_temporary_file( + directory_descriptor: int, + *, + prefix: str, + suffix: str, +) -> tuple[int, str]: + """Create one private temporary file beneath an open directory.""" + for _ in range(100): + name = f"{prefix}{secrets.token_hex(8)}{suffix}" + try: + descriptor = os.open( + name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + PRIVATE_FILE_MODE, + dir_fd=directory_descriptor, + ) + except FileExistsError: + continue + try: + os.fchmod(descriptor, PRIVATE_FILE_MODE) + except OSError: + os.close(descriptor) + try: + os.unlink(name, dir_fd=directory_descriptor) + except OSError: + pass + raise + return descriptor, name + raise OSError("cannot allocate a unique temporary file name") + + +def is_managed_temporary_name(name: str, *, prefix: str, suffix: str) -> bool: + """Return whether a name has the package's random temporary-file shape.""" + if not name.startswith(prefix) or not name.endswith(suffix): + return False + token = name[len(prefix) : len(name) - len(suffix)] + return len(token) == 16 and all(character in "0123456789abcdef" for character in token) + + +def get_file_facts(status: os.stat_result) -> tuple[int, int, int, int, int]: + """Return stable identity and content facts used for replacement checks.""" + return (status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns, status.st_ctime_ns) + + +__all__ = [ + "PRIVATE_DIRECTORY_MODE", + "PRIVATE_FILE_MODE", + "acquire_restrictive_file_lock", + "create_restrictive_temporary_file", + "get_file_facts", + "is_managed_temporary_name", + "open_lock_file", + "open_verified_child_directory", + "open_verified_directory", + "validate_open_directory", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/images/__init__.py new file mode 100644 index 000000000..ff8cb1be0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal SQSH inspection and registry primitives for Data Designer Slurm.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py new file mode 100644 index 000000000..156ffe1cd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical errors for Slurm image inspection and registry operations.""" + +from __future__ import annotations + + +class SlurmImageError(RuntimeError): + """Base error for package-owned Slurm image operations.""" + + +class ImageInspectionError(SlurmImageError): + """Raised when an image environment cannot produce valid factual metadata.""" + + +class ImageRegistryError(SlurmImageError): + """Raised when persisted image registry state cannot be read or written.""" + + +class ImageArtifactRemovalError(ImageRegistryError): + """Raised when an alias was removed but its artifact could not be deleted.""" + + +class ImageNotFoundError(SlurmImageError): + """Raised when an image alias or path is not registered.""" + + +class ImageConflictError(SlurmImageError): + """Raised when an image registration conflicts with immutable registry state.""" + + +class ImageVerificationError(SlurmImageError): + """Raised when an image file or inspection record fails verification.""" + + +class ImageLifecycleError(SlurmImageError): + """Raised when a CPU Slurm image lifecycle job cannot be prepared or submitted.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py new file mode 100644 index 000000000..c57b285f3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound filesystem operations for package-owned Slurm image state.""" + +from __future__ import annotations + +import os +import stat +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from pathlib import Path + +from data_designer.slurm.filesystem import ( + PRIVATE_DIRECTORY_MODE, + acquire_restrictive_file_lock, + create_restrictive_temporary_file, + get_file_facts, +) +from data_designer.slurm.filesystem import ( + open_verified_child_directory as open_shared_child_directory, +) +from data_designer.slurm.filesystem import ( + open_verified_directory as open_shared_directory, +) +from data_designer.slurm.images.errors import ImageRegistryError + + +@contextmanager +def acquire_file_lock( + directory_descriptor: int, name: str, display_path: Path, *, shared: bool = False +) -> Iterator[None]: + """Acquire one advisory lock without reopening its parent path.""" + with ExitStack() as resources: + try: + resources.enter_context( + acquire_restrictive_file_lock( + directory_descriptor, + name, + display_path, + resource_name="registry", + shared=shared, + ) + ) + except OSError as error: + raise ImageRegistryError(f"cannot lock image registry target {display_path}") from error + yield + + +def ensure_private_directory(path: Path, *, parents: bool) -> None: + """Create or restrict one real directory to its owning user.""" + path.mkdir(mode=PRIVATE_DIRECTORY_MODE, parents=parents, exist_ok=True) + with open_verified_directory(path) as descriptor: + os.fchmod(descriptor, PRIVATE_DIRECTORY_MODE) + + +@contextmanager +def open_verified_directory(path: Path) -> Iterator[int]: + """Open one real directory and reject path replacement during the open.""" + with open_shared_directory(path, resource_name="registry") as descriptor: + yield descriptor + + +@contextmanager +def open_verified_child_directory( + parent_descriptor: int, + name: str, + display_path: Path, +) -> Iterator[int]: + """Open one real child directory relative to an already verified parent.""" + with open_shared_child_directory( + parent_descriptor, + name, + display_path, + resource_name="registry", + ) as descriptor: + yield descriptor + + +def create_temporary_file( + directory_descriptor: int, + *, + prefix: str, + suffix: str, +) -> tuple[int, str]: + """Create one restrictive temporary file beneath an open directory.""" + return create_restrictive_temporary_file(directory_descriptor, prefix=prefix, suffix=suffix) + + +def read_regular_text(directory_descriptor: int, name: str, display_path: Path) -> str: + """Read one non-symlink regular text file beneath an open directory.""" + descriptor: int | None = None + try: + before_open = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(before_open.st_mode): + raise OSError(f"registry path {display_path} is not a regular file") + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + after_open = os.fstat(descriptor) + if (before_open.st_dev, before_open.st_ino) != (after_open.st_dev, after_open.st_ino): + raise OSError(f"registry path {display_path} changed while it was being opened") + registry_file = os.fdopen(descriptor, "r", encoding="utf-8") + descriptor = None + with registry_file: + content = registry_file.read() + after_read = os.fstat(registry_file.fileno()) + after_path = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + not stat.S_ISREG(after_read.st_mode) + or not stat.S_ISREG(after_path.st_mode) + or get_file_facts(after_open) != get_file_facts(after_read) + or get_file_facts(after_read) != get_file_facts(after_path) + ): + raise OSError(f"registry path {display_path} changed while it was being read") + return content + finally: + if descriptor is not None: + os.close(descriptor) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py new file mode 100644 index 000000000..d0bc57abf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package-owned factual inspectors executed inside target image environments.""" + +from __future__ import annotations + +import importlib.metadata +import platform +import sys +from typing import Protocol + +from packaging.utils import canonicalize_name +from pydantic import ValidationError + +from data_designer.slurm.config import ( + ClientImageInspection, + ImageInspectionRecord, + InstalledDistribution, + ServingImageInspection, +) +from data_designer.slurm.contracts import Identifier, Sha256Digest +from data_designer.slurm.images.errors import ImageInspectionError +from data_designer.slurm.images.resources.inspect_image import ( + find_distribution_console_script, + find_unique_distribution, +) + +INSPECTOR_VERSION: Identifier = "inspector-1" +_REQUIRED_CLIENT_DISTRIBUTIONS = ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", +) + + +class InspectionEnvironment(Protocol): + """Minimal environment facts required by the package image inspectors.""" + + def get_python_implementation(self) -> str: + """Return the active Python implementation name.""" + + def get_python_version(self) -> str: + """Return the active Python semantic version.""" + + def get_python_abi(self) -> str: + """Return the active Python wheel ABI tag.""" + + def list_distributions(self) -> tuple[InstalledDistribution, ...]: + """Return the installed Python distribution inventory.""" + + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" + + +class SystemInspectionEnvironment: + """Read factual metadata from the environment in which the inspector runs.""" + + def get_python_implementation(self) -> str: + """Return the normalized active Python implementation.""" + return platform.python_implementation().casefold() + + def get_python_version(self) -> str: + """Return the active Python version.""" + return platform.python_version() + + def get_python_abi(self) -> str: + """Return the active Python wheel ABI tag.""" + cache_tag = sys.implementation.cache_tag + if cache_tag is None: + raise ImageInspectionError("active Python does not expose an ABI cache tag") + if cache_tag.startswith("cpython-"): + return f"cp{cache_tag.removeprefix('cpython-')}" + return cache_tag + + def list_distributions(self) -> tuple[InstalledDistribution, ...]: + """Return a normalized, deterministic distribution inventory.""" + versions_by_name: dict[str, str] = {} + for distribution in importlib.metadata.distributions(): + raw_name = distribution.metadata.get("Name") + if not raw_name: + raise ImageInspectionError("installed distribution is missing its canonical name") + name = canonicalize_name(raw_name) + existing_version = versions_by_name.setdefault(name, distribution.version) + if existing_version != distribution.version: + raise ImageInspectionError(f"installed distribution {name!r} has conflicting versions") + try: + return tuple( + InstalledDistribution(name=name, version=version) for name, version in sorted(versions_by_name.items()) + ) + except ValidationError as error: + raise ImageInspectionError("installed distribution inventory is invalid") from error + + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" + try: + distribution = find_unique_distribution(importlib.metadata.distributions(), name) + executable = find_distribution_console_script(distribution, name) + except (OSError, RuntimeError) as error: + raise ImageInspectionError(str(error)) from error + return (distribution.version, executable) + + +class ClientImageInspector: + """Inspect the Python and dependency-installation facts of a client image.""" + + def __init__(self, environment: InspectionEnvironment | None = None) -> None: + self._environment = environment if environment is not None else SystemInspectionEnvironment() + + def inspect(self, sqsh_sha256: Sha256Digest) -> ImageInspectionRecord: + """Return a digest-bound client inspection for the active image environment.""" + environment = self._environment + try: + distributions = tuple(sorted(environment.list_distributions(), key=lambda item: item.name)) + versions_by_name = {distribution.name: distribution.version for distribution in distributions} + missing_distributions = tuple( + name for name in _REQUIRED_CLIENT_DISTRIBUTIONS if name not in versions_by_name + ) + if missing_distributions: + missing = ", ".join(repr(name) for name in missing_distributions) + raise ImageInspectionError(f"required client distributions are not installed: {missing}") + installer_version, installer_path = environment.get_distribution_console_script("pip") + if installer_version != versions_by_name["pip"]: + raise ImageInspectionError("pip console script does not match the distribution inventory") + inspection = ClientImageInspection( + kind="client", + python_implementation=environment.get_python_implementation(), + python_version=environment.get_python_version(), + python_abi=environment.get_python_abi(), + distributions=distributions, + installer_path=installer_path, + installer_version=installer_version, + ) + return ImageInspectionRecord( + schema_version=1, + inspector_version=INSPECTOR_VERSION, + sqsh_sha256=sqsh_sha256, + inspection=inspection, + ) + except ImageInspectionError: + raise + except (OSError, ValueError) as error: + raise ImageInspectionError("client image inspection produced invalid facts") from error + + +class ServingImageInspector: + """Inspect package-owned vLLM runtime facts in a serving image.""" + + def __init__(self, environment: InspectionEnvironment | None = None) -> None: + self._environment = environment if environment is not None else SystemInspectionEnvironment() + + def inspect(self, sqsh_sha256: Sha256Digest) -> ImageInspectionRecord: + """Return a digest-bound serving inspection for the active image environment.""" + environment = self._environment + try: + runtime_version, executable_path = environment.get_distribution_console_script("vllm") + inspection = ServingImageInspection( + kind="serving", + server_type="vllm", + runtime_version=runtime_version, + executable_path=executable_path, + ) + return ImageInspectionRecord( + schema_version=1, + inspector_version=INSPECTOR_VERSION, + sqsh_sha256=sqsh_sha256, + inspection=inspection, + ) + except ImageInspectionError: + raise + except (OSError, ValueError) as error: + raise ImageInspectionError("serving image inspection produced invalid facts") from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py new file mode 100644 index 000000000..2f8add476 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -0,0 +1,718 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Structured CPU Slurm jobs for OCI import and existing-SQSH inspection.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from importlib import resources +from pathlib import Path +from typing import Literal + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, SelectedSlurmProfile +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.images.errors import ImageLifecycleError +from data_designer.slurm.images.filesystem import ( + ensure_private_directory, + open_verified_child_directory, + open_verified_directory, +) +from data_designer.slurm.images.records import ( + ImageLifecycleOperation, + ImageLifecyclePlan, + RegisteredImage, + validate_enroot_mount_path, + validate_oci_source_for_lifecycle, +) +from data_designer.slurm.images.service import VerifiedImageRegistry +from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt + +_INSPECTOR_FILENAME = "inspect_image.py" +_ENROOT_RC_FILENAME = "enroot.rc" +_PLAN_FILENAME = "image-lifecycle-plan.json" +_SCRIPT_FILENAME = "image-lifecycle.sbatch" +_RESOURCE_PACKAGE = "data_designer.slurm.images.resources" +_SCRATCH_RESOURCE_PACKAGE = "data_designer.slurm.runtime" +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_MINIMUM_ENROOT_OCI_VERSION = (3, 5) +_MINIMUM_ENROOT_SQSH_VERSION = (3, 5) +_MAXIMUM_INSPECTION_SIZE = 1024 * 1024 +_PHASE_FILENAME = "phase" + + +@dataclass(frozen=True, slots=True) +class PreparedImageLifecycleJob: + """Persisted, checksum-bound files ready for one Slurm submission. + + Attributes: + plan: Structured lifecycle intent and package-owned paths. + plan_file: Checksum-bound persisted plan. + script_file: Checksum-bound batch script. + """ + + plan: ImageLifecyclePlan + plan_file: ArtifactReference + script_file: ArtifactReference + _job_directory_identity: tuple[int, int] = field(repr=False, compare=False) + + +def prepare_image_lifecycle_job( + request: ImageBuildRequest, + selected_profile: SelectedSlurmProfile, + *, + lifecycle_id: Identifier, +) -> PreparedImageLifecycleJob: + """Stage one deterministic image plan, runtime, and batch script beneath the selected workspace.""" + try: + lifecycle_id = _IDENTIFIER_ADAPTER.validate_python(lifecycle_id, strict=True) + except ValidationError as error: + raise ImageLifecycleError("image lifecycle ID is invalid") from error + try: + validate_oci_source_for_lifecycle(request.source) + except ValueError as error: + raise ImageLifecycleError( + "OCI image source must be a credential-free registry reference without a scheme" + ) from error + workspace_root = Path(selected_profile.profile.workspace_root) + try: + validate_enroot_mount_path(workspace_root.as_posix()) + except ValueError as error: + raise ImageLifecycleError("selected workspace cannot be represented as an Enroot mount") from error + image_root = workspace_root / "images" + temporary_root = image_root / ".tmp" + job_root = temporary_root / "jobs" + job_directory = job_root / lifecycle_id + job_directory_created = False + job_directory_identity: tuple[int, int] | None = None + try: + ensure_private_directory(image_root, parents=True) + ensure_private_directory(temporary_root, parents=False) + ensure_private_directory(job_root, parents=False) + job_directory.mkdir(mode=0o700) + job_directory_created = True + job_directory_status = job_directory.lstat() + job_directory_identity = (job_directory_status.st_dev, job_directory_status.st_ino) + inspection_directory = job_directory / "output" + ensure_private_directory(inspection_directory, parents=False) + inspector_script = _stage_resource(job_directory, _INSPECTOR_FILENAME) + enroot_rc = _stage_resource(job_directory, _ENROOT_RC_FILENAME) + is_existing_sqsh = request.source.endswith(".sqsh") + plan = ImageLifecyclePlan( + schema_version=1, + lifecycle_id=lifecycle_id, + request=request, + selected_profile=selected_profile, + operation=( + ImageLifecycleOperation.INSPECT_SQSH if is_existing_sqsh else ImageLifecycleOperation.IMPORT_OCI + ), + job_directory=job_directory.as_posix(), + sqsh_path=(request.source if is_existing_sqsh else (job_directory / "candidate.sqsh").as_posix()), + inspection_output_path=(inspection_directory / "inspection.json").as_posix(), + inspector_script=inspector_script, + enroot_rc=enroot_rc, + source_oci_digest=None if is_existing_sqsh else request.source.rpartition("@sha256:")[2], + ) + plan_file = _write_file(job_directory / _PLAN_FILENAME, plan.serialize_json().encode(), mode=0o600) + script_file = _write_file( + job_directory / _SCRIPT_FILENAME, + render_image_lifecycle_script(plan).encode(), + mode=0o500, + ) + except (OSError, ValueError) as error: + if job_directory_created and job_directory_identity is not None: + _try_remove_job_directory(job_directory, expected_identity=job_directory_identity) + raise ImageLifecycleError(f"cannot prepare image lifecycle job {lifecycle_id!r}") from error + assert job_directory_identity is not None + return PreparedImageLifecycleJob( + plan=plan, + plan_file=plan_file, + script_file=script_file, + _job_directory_identity=job_directory_identity, + ) + + +def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: + """Render one thin CPU batch entrypoint from structured image lifecycle intent.""" + try: + plan = ImageLifecyclePlan.model_validate(plan.model_dump(mode="python"), strict=True) + except ValueError as error: + raise ImageLifecycleError("cannot render an invalid image lifecycle plan") from error + + profile = plan.selected_profile.profile + image_build = profile.image_build + scratch_source = resources.files(_SCRATCH_RESOURCE_PACKAGE).joinpath("scratch.sh").read_text() + directives = render_batch_directives( + ( + ("job-name", f"dd-image-{plan.request.kind}"), + ("account", profile.scheduler.account), + ("partition", profile.image_build.partition), + ("nodes", "1"), + ("ntasks", "1"), + ("cpus-per-task", str(image_build.cpus_per_task)), + ("mem", image_build.memory), + ("time", image_build.time_limit), + ("chdir", plan.job_directory), + ("output", f"{plan.job_directory}/slurm-%j.out"), + ("error", f"{plan.job_directory}/slurm-%j.err"), + ) + ) + source_block = "" + existing_sqsh_preflight = "" + if plan.operation is ImageLifecycleOperation.IMPORT_OCI: + minimum_major, minimum_minor = _MINIMUM_ENROOT_OCI_VERSION + source_block = f"""readonly DD_OCI_SOURCE={quote_shell_value(_format_enroot_oci_uri(plan.request.source))} +if [[ -e "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then + printf '%s\\n' 'candidate SQSH path already exists' >&2 + exit 73 +fi +verify_enroot_compatibility {minimum_major} {minimum_minor} "digest-pinned OCI imports" +DD_REMOVE_SQSH_ON_FAILURE=1 +report_phase importing_oci +enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" +report_phase verifying_sqsh +""" + else: + minimum_major, minimum_minor = _MINIMUM_ENROOT_SQSH_VERSION + existing_sqsh_preflight = ( + f'verify_enroot_compatibility {minimum_major} {minimum_minor} "existing SQSH inspection"\n' + ) + inspection_source_block = ( + f"""DD_INSPECTION_SQSH={quote_shell_value(f"{plan.job_directory}/verified.sqsh")} +if [[ -e "${{DD_INSPECTION_SQSH}}" || -L "${{DD_INSPECTION_SQSH}}" ]]; then + printf '%s\\n' 'verified SQSH snapshot path already exists' >&2 + exit 73 +fi +DD_REMOVE_INSPECTION_SQSH_ON_EXIT=1 +if cp --help 2>&1 | grep -q -- '--reflink'; then + cp --reflink=auto -- "${{DD_VERIFIED_SQSH}}" "${{DD_INSPECTION_SQSH}}" +else + cp -- "${{DD_VERIFIED_SQSH}}" "${{DD_INSPECTION_SQSH}}" +fi +chmod 0400 "${{DD_INSPECTION_SQSH}}" +readonly DD_INSPECTION_SQSH_SHA256="$(compute_file_sha256 "${{DD_INSPECTION_SQSH}}")" +if [[ "${{DD_INSPECTION_SQSH_SHA256}}" != "${{DD_SQSH_SHA256}}" ]]; then + printf '%s\\n' 'SQSH bytes changed while creating the inspection snapshot' >&2 + exit 74 +fi +""" + if plan.operation is ImageLifecycleOperation.INSPECT_SQSH + else 'DD_INSPECTION_SQSH="${DD_IMAGE_SQSH}"\n' + ) + + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +umask 077 +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +{scratch_source} +readonly DD_JOB_DIR={quote_shell_value(plan.job_directory)} +readonly DD_IMAGE_KIND={quote_shell_value(plan.request.kind)} +readonly DD_IMAGE_SQSH={quote_shell_value(plan.sqsh_path)} +readonly DD_INSPECTION_DIRECTORY={quote_shell_value(Path(plan.inspection_output_path).parent.as_posix())} +readonly DD_INSPECTION_OUTPUT={quote_shell_value(plan.inspection_output_path)} +readonly DD_INSPECTOR={quote_shell_value(plan.inspector_script.path)} +readonly DD_INSPECTOR_SHA256={quote_shell_value(plan.inspector_script.sha256)} +readonly DD_ENROOT_RC={quote_shell_value(plan.enroot_rc.path)} +readonly DD_ENROOT_RC_SHA256={quote_shell_value(plan.enroot_rc.sha256)} +DD_CONTAINER_NAME="" +DD_REMOVE_SQSH_ON_FAILURE=0 +DD_INSPECTION_SQSH="" +DD_REMOVE_INSPECTION_SQSH_ON_EXIT=0 + +cleanup() {{ + local status="$?" + if [[ -n "${{DD_CONTAINER_NAME}}" ]]; then + enroot remove -f "${{DD_CONTAINER_NAME}}" >/dev/null 2>&1 || true + fi + exec 9<&- 2>/dev/null || true + if (( DD_REMOVE_INSPECTION_SQSH_ON_EXIT == 1 )); then + rm -f -- "${{DD_INSPECTION_SQSH}}" >/dev/null 2>&1 || true + fi + if (( status != 0 )); then + rm -f -- "${{DD_INSPECTION_OUTPUT}}" >/dev/null 2>&1 || true + if (( DD_REMOVE_SQSH_ON_FAILURE == 1 )); then + rm -f -- "${{DD_IMAGE_SQSH}}" >/dev/null 2>&1 || true + fi + fi + if [[ -n ${{DD_SCRATCH_ROOT:-}} ]]; then + dd_remove_local_scratch "${{DD_SCRATCH_ROOT}}" >/dev/null 2>&1 || true + fi + return "${{status}}" +}} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +compute_file_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$1")" + printf '%s\\n' "${{actual_sha256%% *}}" +}} + +verify_sha256() {{ + [[ "$(compute_file_sha256 "$2")" == "$1" ]] +}} + +report_phase() {{ + printf '%s\\n' "$1" > "${{DD_JOB_DIR}}/.phase.next" && + mv -f -- "${{DD_JOB_DIR}}/.phase.next" "${{DD_JOB_DIR}}/{_PHASE_FILENAME}" || true +}} + +verify_enroot_compatibility() {{ + local required_major="$1" required_minor="$2" purpose="$3" + local version major minor + version="$(enroot version)" + if [[ ! ${{version}} =~ ^([0-9]+)\\.([0-9]+)\\.([0-9]+)([-+][0-9A-Za-z.-]+)?$ ]]; then + printf '%s\\n' 'Enroot version output is invalid' >&2 + exit 78 + fi + major="${{BASH_REMATCH[1]}}" + minor="${{BASH_REMATCH[2]}}" + if (( 10#${{major}} < 10#${{required_major}} )) || + (( 10#${{major}} == 10#${{required_major}} && 10#${{minor}} < 10#${{required_minor}} )); then + printf 'Enroot %s.%s or newer is required for %s\\n' \\ + "${{required_major}}" "${{required_minor}}" "${{purpose}}" >&2 + exit 78 + fi +}} + +verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" +verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" +dd_initialize_local_scratch +install -d -m 0700 "${{DD_INSPECTION_DIRECTORY}}" +export HOME="${{DD_SCRATCH_ROOT}}/home" +if [[ ! ${{SLURM_CPUS_PER_TASK:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_CPUS_PER_TASK must be a positive integer' >&2 + exit 64 +fi +export ENROOT_MAX_PROCESSORS="${{SLURM_CPUS_PER_TASK}}" + +{source_block}{"report_phase preparing_sqsh" if plan.operation is ImageLifecycleOperation.INSPECT_SQSH else ""} +if [[ ! -f "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then + printf '%s\\n' 'SQSH path must be a regular non-symlink file' >&2 + exit 66 +fi +{existing_sqsh_preflight}readonly DD_SQSH_SHA256="$(compute_file_sha256 "${{DD_IMAGE_SQSH}}")" +if [[ ! ${{DD_SQSH_SHA256}} =~ ^[0-9a-f]{{64}}$ ]]; then + printf '%s\\n' 'SQSH checksum output is invalid' >&2 + exit 65 +fi +if [[ ! ${{SLURM_JOB_ID:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_JOB_ID must be a positive integer' >&2 + exit 64 +fi +DD_CONTAINER_NAME="dd-image-${{SLURM_JOB_ID}}" +exec 9<"${{DD_IMAGE_SQSH}}" +if [[ -e /proc/self/fd/9 ]]; then + readonly DD_VERIFIED_SQSH="/proc/self/fd/9" +else + readonly DD_VERIFIED_SQSH="${{DD_IMAGE_SQSH}}" +fi +if [[ ! -f "${{DD_VERIFIED_SQSH}}" ]]; then + printf '%s\\n' 'SQSH descriptor is not a regular file' >&2 + exit 74 +fi +readonly DD_VERIFIED_SQSH_SHA256="$(compute_file_sha256 "${{DD_VERIFIED_SQSH}}")" +if [[ "${{DD_VERIFIED_SQSH_SHA256}}" != "${{DD_SQSH_SHA256}}" ]]; then + printf '%s\\n' 'SQSH bytes changed before inspection' >&2 + exit 74 +fi +{inspection_source_block} + +report_phase inspecting_image +enroot create -f --name "${{DD_CONTAINER_NAME}}" "${{DD_INSPECTION_SQSH}}" +ENROOT_LOGIN_SHELL=no ENROOT_MOUNT_HOME=no enroot start --root \ + --rc "${{DD_ENROOT_RC}}" \ + --mount "${{DD_INSPECTOR}}:/opt/data-designer-slurm/inspect_image.py:x-create=file,bind,ro" \ + --mount "${{DD_INSPECTION_DIRECTORY}}:/opt/data-designer-slurm/output:x-create=dir,bind" \ + "${{DD_CONTAINER_NAME}}" -- /bin/sh -c ' +python_path="$(command -v python3 || command -v python || true)" +if [ -z "${{python_path}}" ]; then + printf "%s\\n" "target image does not contain Python" >&2 + exit 69 +fi +exec "${{python_path}}" "$@" +' dd-image-inspector \ + /opt/data-designer-slurm/inspect_image.py \ + "${{DD_IMAGE_KIND}}" \ + "${{DD_SQSH_SHA256}}" \ + /opt/data-designer-slurm/output/inspection.json + +[[ -s "${{DD_INSPECTION_OUTPUT}}" ]] +readonly DD_FINAL_INSPECTION_SQSH_SHA256="$(compute_file_sha256 "${{DD_INSPECTION_SQSH}}")" +if [[ "${{DD_FINAL_INSPECTION_SQSH_SHA256}}" != "${{DD_SQSH_SHA256}}" ]]; then + printf '%s\n' 'SQSH inspection snapshot changed during inspection' >&2 + exit 74 +fi +readonly DD_FINAL_SQSH_SHA256="$(compute_file_sha256 "${{DD_VERIFIED_SQSH}}")" +if [[ "${{DD_FINAL_SQSH_SHA256}}" != "${{DD_SQSH_SHA256}}" ]]; then + printf '%s\\n' 'SQSH bytes changed during inspection' >&2 + exit 74 +fi +if [[ ! -f "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then + printf '%s\\n' 'SQSH path changed during inspection' >&2 + exit 74 +fi +readonly DD_FINAL_PATH_SHA256="$(compute_file_sha256 "${{DD_IMAGE_SQSH}}")" +if [[ "${{DD_FINAL_PATH_SHA256}}" != "${{DD_SQSH_SHA256}}" ]]; then + printf '%s\\n' 'SQSH path changed during inspection' >&2 + exit 74 +fi +report_phase inspection_complete +true +""" + + +def submit_prepared_image_lifecycle( + prepared: PreparedImageLifecycleJob, + client: SlurmCommandClient, +) -> SlurmJobSubmissionReceipt: + """Verify and submit one prepared image lifecycle script.""" + with _open_prepared_job_directory(prepared) as job_directory_descriptor: + verified_script = _verify_prepared_lifecycle(prepared, job_directory_descriptor) + try: + script_text = verified_script.decode("utf-8", errors="strict") + except UnicodeError as error: + raise ImageLifecycleError("prepared image lifecycle script is not valid UTF-8") from error + return client.submit_script(script_text) + + +def publish_completed_image_lifecycle( + prepared: PreparedImageLifecycleJob, + *, + replace: bool = False, +) -> RegisteredImage: + """Validate, atomically publish, and register one completed lifecycle result. + + Args: + prepared: The exact checksum-bound job whose inspection has completed. + replace: Whether an existing alias may be replaced explicitly. + + Returns: + The durable alias binding for the verified SQSH. + + Raises: + SlurmImageError: If validation, publication, registration, or cleanup fails. + Cleanup failure after a committed publication does not invalidate the alias. + """ + try: + with _open_prepared_job_directory(prepared) as job_directory_descriptor: + _verify_prepared_lifecycle(prepared, job_directory_descriptor) + inspection_path = Path(prepared.plan.inspection_output_path) + try: + with open_verified_child_directory( + job_directory_descriptor, + inspection_path.parent.name, + inspection_path.parent, + ) as inspection_directory_descriptor: + inspection_content = _read_regular_file( + inspection_directory_descriptor, + inspection_path.name, + inspection_path, + maximum_size=_MAXIMUM_INSPECTION_SIZE, + ) + except OSError as error: + raise ImageLifecycleError("cannot read prepared image lifecycle inspection output") from error + try: + inspection = ImageInspectionRecord.model_validate_json(inspection_content, strict=True) + except (UnicodeError, ValueError, ValidationError) as error: + raise ImageLifecycleError("image lifecycle inspection output is invalid") from error + registry = VerifiedImageRegistry(prepared.plan.selected_profile.profile.workspace_root) + if prepared.plan.operation is ImageLifecycleOperation.IMPORT_OCI: + source_oci_digest = prepared.plan.source_oci_digest + if source_oci_digest is None: + raise ImageLifecycleError("OCI image lifecycle result is missing its source digest") + registered = registry.publish_imported( + prepared.plan.request, + inspection, + Path(prepared.plan.sqsh_path), + source_oci_digest=source_oci_digest, + candidate_directory_descriptor=job_directory_descriptor, + replace=replace, + ) + else: + registered = registry.register_existing(prepared.plan.request, inspection, replace=replace) + except BaseException: + _try_remove_job_directory( + Path(prepared.plan.job_directory), + expected_identity=prepared._job_directory_identity, + ) + raise + cleanup_prepared_image_lifecycle(prepared) + return registered + + +def cleanup_prepared_image_lifecycle(prepared: PreparedImageLifecycleJob) -> None: + """Remove package-owned temporary state for one completed or failed lifecycle job. + + Raises: + ImageLifecycleError: If the package-owned job directory cannot be removed. + """ + job_directory = Path(prepared.plan.job_directory) + try: + _delete_job_directory(job_directory, expected_identity=prepared._job_directory_identity) + except FileNotFoundError: + return + except OSError as error: + raise ImageLifecycleError(f"cannot clean image lifecycle job {prepared.plan.lifecycle_id!r}") from error + + +def read_image_lifecycle_log( + prepared: PreparedImageLifecycleJob, + job_id: int, + stream: Literal["stdout", "stderr"], + *, + offset: int, + maximum_bytes: int, +) -> tuple[bytes, int]: + """Read a bounded chunk from this job's regular, non-symlink Slurm log.""" + if job_id <= 0 or offset < 0 or maximum_bytes <= 0 or stream not in ("stdout", "stderr"): + raise ValueError("invalid image lifecycle log read") + extension = "out" if stream == "stdout" else "err" + filename = f"slurm-{job_id}.{extension}" + with _open_prepared_job_directory(prepared) as directory_descriptor: + try: + descriptor = os.open( + filename, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + except FileNotFoundError: + # The files are only created once Slurm starts the batch job. + return b"", offset + try: + status = os.fstat(descriptor) + if not stat.S_ISREG(status.st_mode): + raise ImageLifecycleError("image lifecycle log is not a regular file") + if status.st_size < offset: + offset = 0 + content = os.pread(descriptor, maximum_bytes, offset) + return content, offset + len(content) + finally: + os.close(descriptor) + + +def read_image_lifecycle_phase(prepared: PreparedImageLifecycleJob) -> str | None: + """Read the job's bounded phase checkpoint without following symlinks.""" + with _open_prepared_job_directory(prepared) as directory_descriptor: + try: + content = _read_regular_file( + directory_descriptor, + _PHASE_FILENAME, + Path(prepared.plan.job_directory) / _PHASE_FILENAME, + maximum_size=64, + ) + except ImageLifecycleError as error: + if isinstance(error.__cause__, FileNotFoundError): + return None + raise + phase = content.decode("ascii", errors="ignore").strip() + valid_phases = {"importing_oci", "preparing_sqsh", "verifying_sqsh", "inspecting_image", "inspection_complete"} + return phase if phase in valid_phases else None + + +def _verify_prepared_lifecycle( + prepared: PreparedImageLifecycleJob, + job_directory_descriptor: int, +) -> bytes: + expected_artifacts = ( + ( + "plan", + prepared.plan_file, + Path(prepared.plan.job_directory) / _PLAN_FILENAME, + prepared.plan.serialize_json().encode(), + ), + ( + "script", + prepared.script_file, + Path(prepared.plan.job_directory) / _SCRIPT_FILENAME, + render_image_lifecycle_script(prepared.plan).encode(), + ), + ) + verified_script: bytes | None = None + for label, artifact, expected_path, expected_content in expected_artifacts: + expected_sha256 = hashlib.sha256(expected_content).hexdigest() + actual_content = _read_regular_file( + job_directory_descriptor, + expected_path.name, + expected_path, + ) + if ( + artifact.path != expected_path.as_posix() + or artifact.sha256 != expected_sha256 + or actual_content != expected_content + ): + raise ImageLifecycleError(f"prepared image lifecycle {label} no longer matches its digest") + if label == "script": + verified_script = actual_content + if verified_script is None: + raise ImageLifecycleError("prepared image lifecycle script was not verified") + return verified_script + + +@contextmanager +def _open_prepared_job_directory(prepared: PreparedImageLifecycleJob) -> Iterator[int]: + job_directory = Path(prepared.plan.job_directory) + with ExitStack() as stack: + try: + parent_descriptor = stack.enter_context(open_verified_directory(job_directory.parent)) + job_directory_descriptor = stack.enter_context( + open_verified_child_directory( + parent_descriptor, + job_directory.name, + job_directory, + ) + ) + status = os.fstat(job_directory_descriptor) + except OSError as error: + raise ImageLifecycleError("cannot access prepared image lifecycle job directory") from error + if (status.st_dev, status.st_ino) != prepared._job_directory_identity: + raise ImageLifecycleError("prepared image lifecycle job directory no longer matches its identity") + yield job_directory_descriptor + + +def _stage_resource(job_directory: Path, filename: str) -> ArtifactReference: + content = resources.files(_RESOURCE_PACKAGE).joinpath(filename).read_bytes() + return _write_file(job_directory / filename, content, mode=0o500) + + +def _format_enroot_oci_uri(source: str) -> str: + repository, _, digest = source.rpartition("@sha256:") + # Enroot 3.5 passes the value after the last colon to the registry's + # /manifests/ endpoint. A digest there pins the import without + # using the tag@digest URI syntax introduced in Enroot 4.0. + tag_separator = repository.rfind(":") + if tag_separator > repository.rfind("/"): + repository = repository[:tag_separator] + registry_or_namespace, separator, remainder = repository.partition("/") + if not separator: + return f"docker://registry-1.docker.io#library/{repository}:sha256:{digest}" + if "." in registry_or_namespace or ":" in registry_or_namespace or registry_or_namespace == "localhost": + registry = "registry-1.docker.io" if registry_or_namespace == "docker.io" else registry_or_namespace + return f"docker://{registry}#{remainder}:sha256:{digest}" + return f"docker://registry-1.docker.io#{repository}:sha256:{digest}" + + +def _write_file(path: Path, content: bytes, *, mode: int) -> ArtifactReference: + descriptor: int | None = None + try: + descriptor = os.open( + path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + mode, + ) + os.fchmod(descriptor, mode) + with os.fdopen(descriptor, "wb") as output: + descriptor = None + output.write(content) + output.flush() + os.fsync(output.fileno()) + finally: + if descriptor is not None: + os.close(descriptor) + return ArtifactReference(path=path.as_posix(), sha256=hashlib.sha256(content).hexdigest()) + + +def _read_regular_file( + directory_descriptor: int, + name: str, + display_path: Path, + *, + maximum_size: int | None = None, +) -> bytes: + descriptor: int | None = None + try: + before_open = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(before_open.st_mode): + raise ImageLifecycleError("prepared image lifecycle file is not regular") + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + path_status = os.fstat(descriptor) + if not stat.S_ISREG(path_status.st_mode) or (before_open.st_dev, before_open.st_ino) != ( + path_status.st_dev, + path_status.st_ino, + ): + raise ImageLifecycleError("prepared image lifecycle file is not regular") + with os.fdopen(descriptor, "rb") as source: + descriptor = None + content = source.read() if maximum_size is None else source.read(maximum_size + 1) + after_read = os.fstat(source.fileno()) + after_path = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + not stat.S_ISREG(after_read.st_mode) + or not stat.S_ISREG(after_path.st_mode) + or _get_file_facts(path_status) != _get_file_facts(after_read) + or _get_file_facts(after_read) != _get_file_facts(after_path) + ): + raise ImageLifecycleError("prepared image lifecycle file changed while it was being read") + if maximum_size is not None and len(content) > maximum_size: + raise ImageLifecycleError("image lifecycle inspection output is too large") + return content + except OSError as error: + raise ImageLifecycleError(f"cannot read prepared image lifecycle file {display_path!s}") from error + finally: + if descriptor is not None: + os.close(descriptor) + + +def _get_file_facts(status: os.stat_result) -> tuple[int, int, int, int, int]: + return (status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns, status.st_ctime_ns) + + +def _delete_job_directory(job_directory: Path, *, expected_identity: tuple[int, int]) -> None: + with open_verified_directory(job_directory.parent) as parent_descriptor: + with open_verified_child_directory( + parent_descriptor, + job_directory.name, + job_directory, + ) as job_directory_descriptor: + status = os.fstat(job_directory_descriptor) + if (status.st_dev, status.st_ino) != expected_identity: + raise OSError(f"image lifecycle job directory {job_directory} changed before cleanup") + _clear_directory(job_directory_descriptor, job_directory) + current_status = os.stat(job_directory.name, dir_fd=parent_descriptor, follow_symlinks=False) + if (current_status.st_dev, current_status.st_ino) != expected_identity: + raise OSError(f"image lifecycle job directory {job_directory} changed during cleanup") + os.rmdir(job_directory.name, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + + +def _clear_directory(directory_descriptor: int, display_path: Path) -> None: + for name in os.listdir(directory_descriptor): + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + child_path = display_path / name + if stat.S_ISDIR(status.st_mode): + with open_verified_child_directory( + directory_descriptor, + name, + child_path, + ) as child_descriptor: + _clear_directory(child_descriptor, child_path) + current_status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if (current_status.st_dev, current_status.st_ino) != (status.st_dev, status.st_ino): + raise OSError(f"image lifecycle directory {child_path} changed during cleanup") + os.rmdir(name, dir_fd=directory_descriptor) + else: + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +def _try_remove_job_directory(job_directory: Path, *, expected_identity: tuple[int, int]) -> None: + try: + _delete_job_directory(job_directory, expected_identity=expected_identity) + except OSError: + pass diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py new file mode 100644 index 000000000..a44d24bbf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable records for Slurm image lifecycle and registry operations.""" + +from __future__ import annotations + +import posixpath +from enum import Enum + +from pydantic import field_validator, model_validator + +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, SelectedSlurmProfile +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + SchemaVersion, + Sha256Digest, + validate_absolute_path, +) + + +class ImageLifecycleOperation(str, Enum): + """Compute-node operation required by one image request.""" + + IMPORT_OCI = "import_oci" + INSPECT_SQSH = "inspect_sqsh" + + +class ImageLifecyclePlan(ContractRecord): + """Immutable inputs for one CPU Slurm image lifecycle job.""" + + lifecycle_id: Identifier + request: ImageBuildRequest + selected_profile: SelectedSlurmProfile + operation: ImageLifecycleOperation + job_directory: str + sqsh_path: str + inspection_output_path: str + inspector_script: ArtifactReference + enroot_rc: ArtifactReference + source_oci_digest: Sha256Digest | None = None + + _paths_are_absolute = field_validator("job_directory", "sqsh_path", "inspection_output_path")( + validate_absolute_path + ) + + @model_validator(mode="after") + def validate_lifecycle(self) -> ImageLifecyclePlan: + validate_oci_source_for_lifecycle(self.request.source) + workspace_root = self.selected_profile.profile.workspace_root + expected_job_directory = posixpath.join( + workspace_root, + "images", + ".tmp", + "jobs", + self.lifecycle_id, + ) + if self.job_directory != expected_job_directory: + raise ValueError("image lifecycle job directory must derive from the selected workspace") + validate_enroot_mount_path(self.job_directory) + if self.inspection_output_path != posixpath.join(self.job_directory, "output", "inspection.json"): + raise ValueError("image lifecycle inspection output must belong to its dedicated output directory") + expected_runtime_artifacts = ( + (self.inspector_script, "inspect_image.py"), + (self.enroot_rc, "enroot.rc"), + ) + for artifact, filename in expected_runtime_artifacts: + if artifact.path != posixpath.join(self.job_directory, filename): + raise ValueError("image lifecycle runtime artifacts must use their package-owned job paths") + + if self.request.source.endswith(".sqsh"): + if self.operation is not ImageLifecycleOperation.INSPECT_SQSH: + raise ValueError("existing SQSH requests require inspection without import") + if self.sqsh_path != self.request.source: + raise ValueError("existing SQSH plans must inspect the authored source path in place") + if self.source_oci_digest is not None: + raise ValueError("existing SQSH plans must not contain an OCI source digest") + else: + if self.operation is not ImageLifecycleOperation.IMPORT_OCI: + raise ValueError("OCI requests require an import operation") + if self.sqsh_path != posixpath.join(self.job_directory, "candidate.sqsh"): + raise ValueError("OCI import output must remain attempt-local until publication") + expected_source_digest = self.request.source.rpartition("@sha256:")[2] + if self.source_oci_digest != expected_source_digest: + raise ValueError("OCI source digest does not match the authored source") + return self + + +def validate_oci_source_for_lifecycle(source: str) -> str: + """Reject OCI source forms that could persist credentials or ambiguous schemes.""" + if source.endswith(".sqsh"): + return source + repository, separator, _digest = source.rpartition("@sha256:") + if ( + separator != "@sha256:" + or "@" in repository + or "://" in repository + or any(delimiter in repository for delimiter in ("?", "#")) + ): + raise ValueError("OCI image source must be a credential-free registry reference without a scheme") + return source + + +def validate_enroot_mount_path(path: str) -> str: + """Reject host paths that Enroot's fstab-style mount parser cannot represent safely.""" + validate_absolute_path(path) + if any(character.isspace() or character in {":", ",", "\\"} for character in path): + raise ValueError("image lifecycle workspace path cannot be represented as an Enroot mount") + return path + + +class RegisteredImage(ContractRecord): + """One immutable alias binding for a verified SQSH artifact.""" + + name: Identifier + path: str + sqsh_sha256: Sha256Digest + source_oci_digest: Sha256Digest | None = None + inspection: ImageInspectionRecord + + @property + def kind(self) -> ImageKind: + """Return the factual role discovered by the package inspector.""" + return self.inspection.inspection.kind + + @property + def immutable_facts(self) -> tuple[Sha256Digest, Sha256Digest | None, ImageInspectionRecord]: + """Return the facts that every alias for this SQSH path must share.""" + return (self.sqsh_sha256, self.source_oci_digest, self.inspection) + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + validate_absolute_path(value) + if not value.endswith(".sqsh"): + raise ValueError("registered image path must end in .sqsh") + return value + + @model_validator(mode="after") + def validate_inspection_digest(self) -> RegisteredImage: + if self.inspection.sqsh_sha256 != self.sqsh_sha256: + raise ValueError("image inspection digest does not match the registered SQSH") + return self + + +class ImageRegistryDocument(ContractValue): + """Versioned contents of the package-owned YAML image registry.""" + + schema_version: SchemaVersion + images: tuple[RegisteredImage, ...] = () + + @model_validator(mode="after") + def validate_images(self) -> ImageRegistryDocument: + names = tuple(image.name for image in self.images) + if names != tuple(sorted(names)): + raise ValueError("registered images must be sorted by alias") + if len(names) != len(set(names)): + raise ValueError("registered image aliases must be unique") + + facts_by_path: dict[str, tuple[Sha256Digest, Sha256Digest | None, ImageInspectionRecord]] = {} + for image in self.images: + facts = image.immutable_facts + existing = facts_by_path.setdefault(image.path, facts) + if existing != facts: + raise ValueError("aliases for one SQSH path must share identical immutable facts") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py new file mode 100644 index 000000000..841e9fa9e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py @@ -0,0 +1,693 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Atomic workspace-derived persistence for immutable Slurm image facts.""" + +from __future__ import annotations + +import hashlib +import json +import os +import stat +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path + +import yaml +from pydantic import TypeAdapter, ValidationError +from yaml.constructor import ConstructorError +from yaml.nodes import MappingNode +from yaml.resolver import BaseResolver + +from data_designer.slurm.config import ImageRef +from data_designer.slurm.contracts import Identifier, validate_absolute_path +from data_designer.slurm.filesystem import get_file_facts +from data_designer.slurm.images.errors import ( + ImageArtifactRemovalError, + ImageConflictError, + ImageNotFoundError, + ImageRegistryError, +) +from data_designer.slurm.images.filesystem import ( + acquire_file_lock, + create_temporary_file, + ensure_private_directory, + open_verified_child_directory, + open_verified_directory, + read_regular_text, +) +from data_designer.slurm.images.records import ImageRegistryDocument, RegisteredImage + +_LOCK_DIRECTORY_NAME = ".locks" +_REGISTRY_FILENAME = "registry.yaml" +_REGISTRY_LOCK_FILENAME = "registry.lock" +_RUN_REFERENCE_LOCK_FILENAME = "run-references.lock" +_REGISTRY_ROLLBACK_FILENAME = ".registry.rollback.yaml" +_REGISTRY_COMMITTED_FILENAME = ".registry.committed.yaml" +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class ImageRegistryStore: + """Persist image aliases beneath one explicitly selected shared workspace. + + Writers acquire alias, target when present, then registry locks; readers + acquire only the registry lock. Keeping this order fixed prevents deadlocks. + """ + + def __init__(self, workspace_root: str | Path) -> None: + try: + normalized_root = validate_absolute_path(Path(workspace_root).as_posix()) + except ValueError as error: + raise ImageRegistryError(f"invalid image workspace root {workspace_root!s}") from error + self._image_root = Path(normalized_root) / "images" + self._registry_path = self._image_root / _REGISTRY_FILENAME + self._lock_directory = self._image_root / _LOCK_DIRECTORY_NAME + + @property + def image_root(self) -> Path: + """Return the workspace-derived image root.""" + return self._image_root + + @property + def registry_path(self) -> Path: + """Return the workspace-derived registry record path.""" + return self._registry_path + + @contextmanager + def acquire_run_reference_lock(self, *, exclusive: bool = False) -> Iterator[None]: + """Allow concurrent run publication while excluding artifact removal.""" + self._ensure_storage() + with self._open_storage() as (_, lock_directory_descriptor): + with acquire_file_lock( + lock_directory_descriptor, + _RUN_REFERENCE_LOCK_FILENAME, + self._lock_directory / _RUN_REFERENCE_LOCK_FILENAME, + shared=not exclusive, + ): + yield + + def list_images(self) -> tuple[RegisteredImage, ...]: + """Return all registered aliases in deterministic order.""" + return self._load().images + + def get_by_name(self, name: Identifier) -> RegisteredImage: + """Return one registered alias.""" + name = _validate_image_alias(name) + for image in self._load().images: + if image.name == name: + return image + raise ImageNotFoundError(f"image alias {name!r} is not registered") + + def get_by_path(self, path: str | Path) -> RegisteredImage: + """Return immutable facts for one registered absolute SQSH path.""" + try: + normalized_path = ImageRef(path=Path(path).as_posix()).path + except ValidationError as error: + raise ImageRegistryError(f"invalid registered image path {path!s}") from error + assert normalized_path is not None + matches = tuple(image for image in self._load().images if image.path == normalized_path) + if not matches: + raise ImageNotFoundError(f"image path {normalized_path!r} is not registered") + return matches[0] + + def register( + self, + image: RegisteredImage, + *, + verify_before_publish: Callable[[RegisteredImage], None], + prepare_before_publish: Callable[[RegisteredImage], None] | None = None, + verify_after_publish: Callable[[RegisteredImage], None] | None = None, + rollback_after_failure: Callable[[RegisteredImage], None] | None = None, + replace: bool = False, + ) -> RegisteredImage: + """Atomically add or explicitly replace one verified image alias. + + Args: + image: Immutable alias and artifact facts to persist. + verify_before_publish: Final verification or artifact-publication callback run under all writer locks. + prepare_before_publish: Optional expensive preparation run under the alias and target locks, but outside + the global registry lock. The registry is revalidated after this callback returns. + verify_after_publish: Optional verification run after registry replacement but before commit. + rollback_after_failure: Optional no-fail cleanup run under the alias and target locks after registry + rollback. + replace: Whether an existing alias may be replaced explicitly. + + Returns: + The registered alias facts. + + Raises: + ImageConflictError: If the alias or target facts conflict with existing state. + ImageRegistryError: If registry storage cannot be locked, read, recovered, or persisted. + """ + self._ensure_storage() + alias_lock_name = f"alias-{image.name}.lock" + target_lock_name = _get_target_lock_name(image.path) + with self._open_storage() as (image_root_descriptor, lock_directory_descriptor): + with ( + acquire_file_lock( + lock_directory_descriptor, + alias_lock_name, + self._lock_directory / alias_lock_name, + ), + acquire_file_lock( + lock_directory_descriptor, + target_lock_name, + self._lock_directory / target_lock_name, + ), + ): + try: + if prepare_before_publish is not None: + with acquire_file_lock( + lock_directory_descriptor, + _REGISTRY_LOCK_FILENAME, + self._lock_directory / _REGISTRY_LOCK_FILENAME, + ): + self._recover_transaction(image_root_descriptor) + self._build_registration( + self._load_from_directory(image_root_descriptor), + image, + replace=replace, + ) + prepare_before_publish(image) + with acquire_file_lock( + lock_directory_descriptor, + _REGISTRY_LOCK_FILENAME, + self._lock_directory / _REGISTRY_LOCK_FILENAME, + ): + self._recover_transaction(image_root_descriptor) + snapshot = self._load_from_directory(image_root_descriptor) + updated = self._build_registration(snapshot, image, replace=replace) + self._save( + image_root_descriptor, + snapshot, + updated, + verify_before_publish=lambda: verify_before_publish(image), + verify_after_publish=( + None if verify_after_publish is None else lambda: verify_after_publish(image) + ), + ) + except _RegistryCommitUncertainError: + raise + except BaseException: + if rollback_after_failure is not None: + rollback_after_failure(image) + raise + return image + + def unregister(self, name: Identifier, *, delete_artifact: bool = False) -> RegisteredImage: + """Remove an alias and optionally its unshared, package-owned SQSH artifact.""" + name = _validate_image_alias(name) + if delete_artifact: + return self._unregister_with_artifact(name) + self._ensure_storage() + alias_lock_name = f"alias-{name}.lock" + with self._open_storage() as (image_root_descriptor, lock_directory_descriptor): + with ( + acquire_file_lock( + lock_directory_descriptor, + alias_lock_name, + self._lock_directory / alias_lock_name, + ), + acquire_file_lock( + lock_directory_descriptor, + _REGISTRY_LOCK_FILENAME, + self._lock_directory / _REGISTRY_LOCK_FILENAME, + ), + ): + self._recover_transaction(image_root_descriptor) + snapshot = self._load_from_directory(image_root_descriptor) + try: + removed = next(image for image in snapshot.images if image.name == name) + except StopIteration: + raise ImageNotFoundError(f"image alias {name!r} is not registered") from None + updated = ImageRegistryDocument( + schema_version=1, + images=tuple(image for image in snapshot.images if image.name != name), + ) + self._save(image_root_descriptor, snapshot, updated) + return removed + + def _unregister_with_artifact(self, name: Identifier) -> RegisteredImage: + with self.acquire_run_reference_lock(exclusive=True): + return self._unregister_with_artifact_locked(name) + + def _unregister_with_artifact_locked(self, name: Identifier) -> RegisteredImage: + selected = self.get_by_name(name) + artifact_directory = self._image_root / "artifacts" + expected_path = artifact_directory / f"{name}-{selected.sqsh_sha256}.sqsh" + if selected.source_oci_digest is None or selected.path != expected_path.as_posix(): + raise ImageConflictError("only package-owned OCI artifacts can be deleted") + + alias_lock_name = f"alias-{name}.lock" + target_lock_name = _get_target_lock_name(selected.path) + with self._open_storage() as (image_root_descriptor, lock_directory_descriptor): + with ( + acquire_file_lock( + lock_directory_descriptor, + alias_lock_name, + self._lock_directory / alias_lock_name, + ), + acquire_file_lock( + lock_directory_descriptor, + target_lock_name, + self._lock_directory / target_lock_name, + ), + ): + # The caller holds the exclusive run-reference lock through the final unlink. + self._require_no_persisted_run_reference(selected.path) + with acquire_file_lock( + lock_directory_descriptor, + _REGISTRY_LOCK_FILENAME, + self._lock_directory / _REGISTRY_LOCK_FILENAME, + ): + self._recover_transaction(image_root_descriptor) + snapshot = self._load_from_directory(image_root_descriptor) + current = next((image for image in snapshot.images if image.name == name), None) + if current != selected: + raise ImageConflictError("image alias changed during artifact removal") + if any(image.name != name and image.path == selected.path for image in snapshot.images): + raise ImageConflictError("SQSH artifact is still used by another image alias") + with open_verified_child_directory( + image_root_descriptor, + "artifacts", + artifact_directory, + ) as artifact_directory_descriptor: + try: + artifact_descriptor = os.open( + expected_path.name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=artifact_directory_descriptor, + ) + except OSError as error: + raise ImageConflictError("SQSH artifact is unavailable for deletion") from error + try: + artifact_status = os.fstat(artifact_descriptor) + path_status = os.stat( + expected_path.name, + dir_fd=artifact_directory_descriptor, + follow_symlinks=False, + ) + if not stat.S_ISREG(artifact_status.st_mode) or ( + artifact_status.st_dev, + artifact_status.st_ino, + ) != (path_status.st_dev, path_status.st_ino): + raise ImageConflictError("SQSH artifact changed before deletion") + updated = ImageRegistryDocument( + schema_version=1, + images=tuple(image for image in snapshot.images if image.name != name), + ) + self._save(image_root_descriptor, snapshot, updated) + try: + path_status = os.stat( + expected_path.name, + dir_fd=artifact_directory_descriptor, + follow_symlinks=False, + ) + if ( + artifact_status.st_dev, + artifact_status.st_ino, + ) != (path_status.st_dev, path_status.st_ino): + raise OSError("SQSH artifact changed after alias removal") + os.unlink(expected_path.name, dir_fd=artifact_directory_descriptor) + os.fsync(artifact_directory_descriptor) + except OSError as error: + raise ImageArtifactRemovalError( + "image alias was removed, but SQSH artifact deletion could not be confirmed" + ) from error + finally: + os.close(artifact_descriptor) + return selected + + def _require_no_persisted_run_reference(self, artifact_path: str) -> None: + """Keep artifacts needed by submitted runs and their possible retries.""" + try: + workspace_root = self._image_root.parent + runs_root = workspace_root / "runs" + with open_verified_directory(workspace_root) as workspace_descriptor: + try: + os.stat("runs", dir_fd=workspace_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + with open_verified_child_directory(workspace_descriptor, "runs", runs_root) as runs_descriptor: + for run_name in os.listdir(runs_descriptor): + run_status = os.stat(run_name, dir_fd=runs_descriptor, follow_symlinks=False) + if stat.S_ISLNK(run_status.st_mode): + raise ImageConflictError("run state cannot be checked before SQSH artifact deletion") + if not stat.S_ISDIR(run_status.st_mode): + continue + with open_verified_child_directory( + runs_descriptor, run_name, runs_root / run_name + ) as run_descriptor: + plan = self._read_persisted_run_plan(run_descriptor) + if plan is None: + continue + deployments = plan.get("deployments") + if not isinstance(deployments, list) or not deployments: + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + for deployment in deployments: + if not isinstance(deployment, dict) or not isinstance(deployment.get("image"), dict): + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + image_path = deployment["image"].get("path") + if not isinstance(image_path, str) or not image_path.endswith(".sqsh"): + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + try: + validate_absolute_path(image_path) + except ValueError: + raise ImageConflictError( + "run plan cannot be checked before SQSH artifact deletion" + ) from None + if image_path == artifact_path: + raise ImageConflictError(f"SQSH artifact is still referenced by run {run_name!r}") + except (OSError, ValueError, TypeError, UnicodeError) as error: + raise ImageConflictError("run state cannot be checked before SQSH artifact deletion") from error + + @staticmethod + def _read_persisted_run_plan(run_descriptor: int) -> dict[str, object] | None: + """Read one run plan through its directory descriptor with a hard size limit.""" + plan_name = "resolved-plan.json" + try: + before_open = os.stat(plan_name, dir_fd=run_descriptor, follow_symlinks=False) + except FileNotFoundError: + return None + if not stat.S_ISREG(before_open.st_mode) or before_open.st_size > 16 * 1024 * 1024: + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + plan_descriptor = os.open( + plan_name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=run_descriptor, + ) + try: + opened = os.fstat(plan_descriptor) + if not stat.S_ISREG(opened.st_mode) or get_file_facts(before_open) != get_file_facts(opened): + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + remaining = 16 * 1024 * 1024 + 1 + chunks: list[bytes] = [] + while remaining: + chunk = os.read(plan_descriptor, min(remaining, 1024 * 1024)) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + if remaining == 0: + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + after_read = os.fstat(plan_descriptor) + after_path = os.stat(plan_name, dir_fd=run_descriptor, follow_symlinks=False) + if get_file_facts(opened) != get_file_facts(after_read) or get_file_facts(opened) != get_file_facts( + after_path + ): + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + plan = json.loads(b"".join(chunks)) + if not isinstance(plan, dict): + raise ImageConflictError("run plan cannot be checked before SQSH artifact deletion") + return plan + finally: + os.close(plan_descriptor) + + def _load(self) -> ImageRegistryDocument: + try: + self._image_root.lstat() + except FileNotFoundError: + return ImageRegistryDocument(schema_version=1) + except OSError as error: + raise ImageRegistryError(f"cannot load image registry {self._registry_path}") from error + try: + self._ensure_storage() + except ImageRegistryError as error: + raise ImageRegistryError(f"cannot load image registry {self._registry_path}") from error + try: + with self._open_storage() as (image_root_descriptor, lock_directory_descriptor): + with acquire_file_lock( + lock_directory_descriptor, + _REGISTRY_LOCK_FILENAME, + self._lock_directory / _REGISTRY_LOCK_FILENAME, + ): + self._recover_transaction(image_root_descriptor) + return self._load_from_directory(image_root_descriptor) + except OSError as error: + raise ImageRegistryError(f"cannot load image registry {self._registry_path}") from error + + def _load_from_directory(self, image_root_descriptor: int) -> ImageRegistryDocument: + try: + payload = yaml.load( + read_regular_text(image_root_descriptor, _REGISTRY_FILENAME, self._registry_path), + Loader=_UniqueKeySafeLoader, + ) + return ImageRegistryDocument.model_validate_json(json.dumps(payload)) + except FileNotFoundError: + return ImageRegistryDocument(schema_version=1) + except (OSError, RecursionError, TypeError, UnicodeError, ValueError, ValidationError, yaml.YAMLError) as error: + raise ImageRegistryError(f"cannot load image registry {self._registry_path}") from error + + def _save( + self, + image_root_descriptor: int, + previous: ImageRegistryDocument, + snapshot: ImageRegistryDocument, + *, + verify_before_publish: Callable[[], None] | None = None, + verify_after_publish: Callable[[], None] | None = None, + ) -> None: + temporary_name: str | None = None + rollback_temporary_name: str | None = None + rollback_installed = False + try: + rollback_temporary_name = self._write_temporary_snapshot( + image_root_descriptor, + previous, + prefix=".registry.rollback.", + ) + os.replace( + rollback_temporary_name, + _REGISTRY_ROLLBACK_FILENAME, + src_dir_fd=image_root_descriptor, + dst_dir_fd=image_root_descriptor, + ) + rollback_temporary_name = None + rollback_installed = True + os.fsync(image_root_descriptor) + temporary_name = self._write_temporary_snapshot( + image_root_descriptor, + snapshot, + prefix=".registry.", + ) + if verify_before_publish is not None: + verify_before_publish() + os.replace( + temporary_name, + _REGISTRY_FILENAME, + src_dir_fd=image_root_descriptor, + dst_dir_fd=image_root_descriptor, + ) + temporary_name = None + if verify_after_publish is not None: + verify_after_publish() + os.fsync(image_root_descriptor) + os.replace( + _REGISTRY_ROLLBACK_FILENAME, + _REGISTRY_COMMITTED_FILENAME, + src_dir_fd=image_root_descriptor, + dst_dir_fd=image_root_descriptor, + ) + rollback_installed = False + try: + os.fsync(image_root_descriptor) + except OSError as error: + raise _RegistryCommitUncertainError( + f"cannot persist image registry {self._registry_path}; commit state requires recovery" + ) from error + try: + os.unlink(_REGISTRY_COMMITTED_FILENAME, dir_fd=image_root_descriptor) + os.fsync(image_root_descriptor) + except OSError: + pass + except BaseException as error: + if rollback_installed: + try: + self._recover_transaction(image_root_descriptor) + rollback_installed = False + except ImageRegistryError: + pass + if isinstance(error, (OSError, TypeError, yaml.YAMLError)): + raise ImageRegistryError(f"cannot persist image registry {self._registry_path}") from error + raise + finally: + for name in (temporary_name, rollback_temporary_name): + if name is None: + continue + try: + os.unlink(name, dir_fd=image_root_descriptor) + except OSError: + pass + + def _write_temporary_snapshot( + self, + image_root_descriptor: int, + snapshot: ImageRegistryDocument, + *, + prefix: str, + ) -> str: + descriptor: int | None = None + temporary_name: str | None = None + try: + descriptor, temporary_name = create_temporary_file( + image_root_descriptor, + prefix=prefix, + suffix=".tmp", + ) + output = os.fdopen(descriptor, "w", encoding="utf-8") + descriptor = None + with output: + yaml.safe_dump( + snapshot.model_dump(mode="json"), + output, + default_flow_style=False, + sort_keys=True, + ) + output.flush() + os.fsync(output.fileno()) + return temporary_name + except BaseException: + if descriptor is not None: + os.close(descriptor) + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=image_root_descriptor) + except OSError: + pass + raise + + def _recover_transaction(self, image_root_descriptor: int) -> None: + try: + rollback_exists = _has_regular_marker(image_root_descriptor, _REGISTRY_ROLLBACK_FILENAME) + committed_exists = _has_regular_marker(image_root_descriptor, _REGISTRY_COMMITTED_FILENAME) + if rollback_exists and committed_exists: + raise OSError("image registry contains conflicting transaction markers") + if committed_exists: + os.unlink(_REGISTRY_COMMITTED_FILENAME, dir_fd=image_root_descriptor) + os.fsync(image_root_descriptor) + return + if not rollback_exists: + return + os.replace( + _REGISTRY_ROLLBACK_FILENAME, + _REGISTRY_FILENAME, + src_dir_fd=image_root_descriptor, + dst_dir_fd=image_root_descriptor, + ) + os.fsync(image_root_descriptor) + except OSError as error: + raise ImageRegistryError(f"cannot recover image registry {self._registry_path}") from error + + def _ensure_storage(self) -> None: + try: + ensure_private_directory(self._image_root, parents=True) + ensure_private_directory(self._lock_directory, parents=False) + except OSError as error: + raise ImageRegistryError(f"cannot initialize image registry beneath {self._image_root}") from error + + @contextmanager + def _open_storage(self) -> Iterator[tuple[int, int]]: + try: + with open_verified_directory(self._image_root) as image_root_descriptor: + with open_verified_child_directory( + image_root_descriptor, + _LOCK_DIRECTORY_NAME, + self._lock_directory, + ) as lock_directory_descriptor: + yield image_root_descriptor, lock_directory_descriptor + except OSError as error: + raise ImageRegistryError(f"cannot access image registry beneath {self._image_root}") from error + + @staticmethod + def _build_registration( + snapshot: ImageRegistryDocument, + image: RegisteredImage, + *, + replace: bool, + ) -> ImageRegistryDocument: + existing = next((entry for entry in snapshot.images if entry.name == image.name), None) + if existing is not None and not replace: + raise ImageConflictError(f"image alias {image.name!r} is already registered") + ImageRegistryStore._validate_path_facts(snapshot, image, replacing=existing) + images = tuple(entry for entry in snapshot.images if entry.name != image.name) + (image,) + return ImageRegistryDocument( + schema_version=1, + images=tuple(sorted(images, key=lambda entry: entry.name)), + ) + + @staticmethod + def _validate_path_facts( + snapshot: ImageRegistryDocument, + image: RegisteredImage, + *, + replacing: RegisteredImage | None, + ) -> None: + for existing in snapshot.images: + if replacing is not None and existing.name == replacing.name: + continue + if existing.path != image.path: + continue + if existing.immutable_facts != image.immutable_facts: + raise ImageConflictError(f"image path {image.path!r} already has different immutable facts") + + +class _RegistryCommitUncertainError(ImageRegistryError): + """Report a failed commit sync while preserving its recovery marker.""" + + +class _UniqueKeySafeLoader(yaml.SafeLoader): + """Safe YAML loader that rejects ambiguous duplicate mapping keys.""" + + +def _construct_unique_mapping( + loader: _UniqueKeySafeLoader, + node: MappingNode, + deep: bool = False, +) -> dict[object, object]: + mapping: dict[object, object] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in mapping + except TypeError as error: + raise ConstructorError( + "while constructing a mapping", + node.start_mark, + "found an unhashable mapping key", + key_node.start_mark, + ) from error + if duplicate: + raise ConstructorError( + "while constructing a mapping", + node.start_mark, + f"found duplicate key {key!r}", + key_node.start_mark, + ) + mapping[key] = loader.construct_object(value_node, deep=deep) + return mapping + + +_UniqueKeySafeLoader.add_constructor(BaseResolver.DEFAULT_MAPPING_TAG, _construct_unique_mapping) + + +def _validate_image_alias(name: str) -> Identifier: + """Validate an image alias before deriving any filesystem path from it.""" + try: + return _IDENTIFIER_ADAPTER.validate_python(name, strict=True) + except ValidationError as error: + raise ImageRegistryError(f"invalid image alias {name!r}") from error + + +def _get_target_lock_name(path: str) -> str: + digest = hashlib.sha256(path.encode("utf-8")).hexdigest() + return f"target-{digest}.lock" + + +def _has_regular_marker(directory_descriptor: int, name: str) -> bool: + try: + marker_status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return False + if not stat.S_ISREG(marker_status.st_mode): + raise OSError(f"image registry transaction marker {name} is not a regular file") + return True diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py new file mode 100644 index 000000000..d3be375ef --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone resources staged for image inspection jobs.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc new file mode 100644 index 000000000..9e31e54c5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc @@ -0,0 +1,6 @@ +#!/bin/sh +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[ "$1" = "--" ] && shift +exec "$@" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py new file mode 100644 index 000000000..01dc0f6b6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Emit factual image metadata using only the target image's Python standard library.""" + +from __future__ import annotations + +import importlib.metadata +import json +import os +import platform +import re +import sys +from collections.abc import Iterable +from pathlib import Path + +INSPECTOR_VERSION = "inspector-1" +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_REQUIRED_CLIENT_DISTRIBUTIONS = ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", +) + + +def inspect_image(kind: str, sqsh_sha256: str) -> dict[str, object]: + """Return one JSON-compatible digest-bound inspection record.""" + if _SHA256_PATTERN.fullmatch(sqsh_sha256) is None: + raise ValueError("SQSH digest must be lowercase SHA-256 text") + if kind == "client": + inspection = _inspect_client() + elif kind == "serving": + inspection = _inspect_serving() + else: + raise ValueError("image kind must be 'client' or 'serving'") + return { + "schema_version": 1, + "inspector_version": INSPECTOR_VERSION, + "sqsh_sha256": sqsh_sha256, + "inspection": inspection, + } + + +def write_inspection(path: Path, record: dict[str, object]) -> None: + """Atomically write one restrictive inspection record.""" + path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.{os.getpid()}.tmp") + descriptor = os.open(temporary_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as output: + json.dump(record, output, allow_nan=False, ensure_ascii=False, indent=2, sort_keys=True) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + os.replace(temporary_path, path) + directory_descriptor = os.open(path.parent, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + finally: + try: + temporary_path.unlink() + except FileNotFoundError: + pass + + +def main(arguments: list[str] | None = None) -> int: + """Run the standalone inspector entrypoint.""" + values = sys.argv[1:] if arguments is None else arguments + if len(values) != 3: + raise ValueError("expected KIND SQSH_SHA256 OUTPUT_PATH") + kind, sqsh_sha256, output_path = values + output = Path(output_path) + if not output.is_absolute(): + raise ValueError("inspection output path must be absolute") + write_inspection(output, inspect_image(kind, sqsh_sha256)) + return 0 + + +def _inspect_client() -> dict[str, object]: + distributions = tuple(importlib.metadata.distributions()) + versions = _list_distribution_versions(distributions) + missing = tuple(name for name in _REQUIRED_CLIENT_DISTRIBUTIONS if name not in versions) + if missing: + raise RuntimeError(f"required client distributions are not installed: {', '.join(missing)}") + installer_distribution = find_unique_distribution(distributions, "pip") + installer_path = find_distribution_console_script(installer_distribution, "pip") + cache_tag = sys.implementation.cache_tag + if cache_tag is None: + raise RuntimeError("active Python does not expose an ABI cache tag") + python_abi = f"cp{cache_tag.removeprefix('cpython-')}" if cache_tag.startswith("cpython-") else cache_tag + return { + "kind": "client", + "python_implementation": platform.python_implementation().casefold(), + "python_version": platform.python_version(), + "python_abi": python_abi, + "distributions": [{"name": name, "version": version} for name, version in sorted(versions.items())], + "installer_path": installer_path, + "installer_version": installer_distribution.version, + } + + +def _inspect_serving() -> dict[str, object]: + distribution = find_unique_distribution(importlib.metadata.distributions(), "vllm") + executable_path = find_distribution_console_script(distribution, "vllm") + return { + "kind": "serving", + "server_type": "vllm", + "runtime_version": distribution.version, + "executable_path": executable_path, + } + + +def _list_distribution_versions(distributions: Iterable[importlib.metadata.Distribution]) -> dict[str, str]: + versions: dict[str, str] = {} + for distribution in distributions: + name = _distribution_name(distribution) + existing = versions.setdefault(name, distribution.version) + if existing != distribution.version: + raise RuntimeError(f"installed distribution {name!r} has conflicting versions") + return versions + + +def _distribution_name(distribution: importlib.metadata.Distribution) -> str: + raw_name = distribution.metadata.get("Name") + if not raw_name: + raise RuntimeError("installed distribution is missing its canonical name") + return re.sub(r"[-_.]+", "-", raw_name).casefold() + + +def find_unique_distribution( + distributions: Iterable[importlib.metadata.Distribution], + name: str, +) -> importlib.metadata.Distribution: + """Select exactly one installed distribution by canonical name.""" + canonical_name = re.sub(r"[-_.]+", "-", name).casefold() + matches = tuple( + distribution for distribution in distributions if _distribution_name(distribution) == canonical_name + ) + if not matches: + raise RuntimeError(f"required distribution {canonical_name!r} is not installed") + if len(matches) != 1: + raise RuntimeError(f"required distribution {canonical_name!r} is not installed exactly once") + return matches[0] + + +def find_distribution_console_script(distribution: importlib.metadata.Distribution, name: str) -> str: + """Resolve exactly one executable console script owned by one distribution.""" + entry_points = tuple( + entry_point + for entry_point in distribution.entry_points + if entry_point.group == "console_scripts" and entry_point.name == name + ) + if len(entry_points) != 1: + raise RuntimeError(f"required distribution {name!r} does not expose one console script") + files = distribution.files + if files is None: + raise RuntimeError(f"required distribution {name!r} does not expose an installed-file inventory") + executable_paths: set[Path] = set() + for installed_file in files: + if Path(str(installed_file)).name != name: + continue + candidate = Path(distribution.locate_file(installed_file)) + if not candidate.is_absolute(): + continue + try: + resolved = candidate.resolve(strict=True) + except OSError: + continue + if resolved.is_file() and os.access(resolved, os.X_OK): + executable_paths.add(resolved) + if len(executable_paths) != 1: + raise RuntimeError(f"required distribution {name!r} does not own one executable console script") + return next(iter(executable_paths)).as_posix() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py new file mode 100644 index 000000000..f1004f68a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -0,0 +1,523 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Verified existing-image registry operations for Slurm planning.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path + +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion, Version + +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + ServingImageInspection, +) +from data_designer.slurm.contracts import Sha256Digest +from data_designer.slurm.images.errors import ImageConflictError, ImageVerificationError +from data_designer.slurm.images.filesystem import ensure_private_directory, open_verified_directory +from data_designer.slurm.images.inspection import INSPECTOR_VERSION +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.images.vllm_source import pinned_vllm_version +from data_designer.slurm.planning import ResolvedImage + +_HASH_BLOCK_SIZE = 1024 * 1024 +_ARTIFACT_DIRECTORY_NAME = "artifacts" +_TEMPORARY_DIRECTORY_NAME = ".tmp" +_PROXY_HTTP_REQUIREMENT = SpecifierSet(">=3.14.3,<4") + + +class VerifiedImageRegistry: + """Provide verified existing-SQSH operations for the public image facade.""" + + def __init__(self, workspace_root: str | Path) -> None: + self._registry = ImageRegistryStore(workspace_root) + + def register_existing( + self, + request: ImageBuildRequest, + inspection: ImageInspectionRecord, + *, + replace: bool = False, + ) -> RegisteredImage: + """Verify and register an existing compute-visible SQSH in place.""" + if not request.source.endswith(".sqsh"): + raise ImageVerificationError("OCI image sources require the CPU Slurm image lifecycle") + path = Path(request.source) + with _open_verified_sqsh(path) as snapshot: + _validate_inspection( + inspection, + expected_kind=ImageKind(request.kind), + expected_sha256=snapshot.sha256, + ) + image = RegisteredImage( + schema_version=1, + name=request.name, + path=path.as_posix(), + sqsh_sha256=inspection.sqsh_sha256, + inspection=inspection, + ) + verify_identity = lambda _image: _verify_snapshot_identity(image, snapshot) + return self._registry.register( + image, + verify_before_publish=verify_identity, + verify_after_publish=verify_identity, + replace=replace, + ) + + def publish_imported( + self, + request: ImageBuildRequest, + inspection: ImageInspectionRecord, + candidate_path: Path, + *, + source_oci_digest: Sha256Digest, + candidate_directory_descriptor: int | None = None, + replace: bool = False, + ) -> RegisteredImage: + """Publish one verified lifecycle-owned OCI import and register its alias. + + The candidate must live in this workspace's lifecycle temporary tree. + Publication keeps a prior alias or artifact intact unless replacement is + explicit and the new artifact is fully verified. + + Args: + request: Authored digest-pinned OCI image request. + inspection: Package-owned facts bound to the imported SQSH digest. + candidate_path: Lifecycle-owned SQSH candidate beneath workspace temporary storage. + source_oci_digest: Resolved source digest carried by the lifecycle plan. + candidate_directory_descriptor: Optional verified lifecycle directory used to access the candidate + descriptor-relative. + replace: Whether an existing alias may be replaced explicitly. + + Returns: + The durable alias binding for the published artifact. + + Raises: + ImageConflictError: If an alias or artifact collision is not identical. + ImageVerificationError: If source, candidate, inspection, or publication verification fails. + ImageRegistryError: If registry publication or recovery fails. + """ + if request.source.endswith(".sqsh"): + raise ImageVerificationError("existing SQSH sources must be registered in place") + expected_source_digest = request.source.rpartition("@sha256:")[2] + if source_oci_digest != expected_source_digest: + raise ImageVerificationError("OCI source digest does not match the authored source") + temporary_root = self._registry.image_root / _TEMPORARY_DIRECTORY_NAME + if not candidate_path.is_absolute() or ".." in candidate_path.parts: + raise ImageVerificationError("imported SQSH candidate must belong to lifecycle temporary storage") + try: + candidate_path.relative_to(temporary_root) + except ValueError as error: + raise ImageVerificationError( + "imported SQSH candidate must belong to lifecycle temporary storage" + ) from error + _validate_inspection(inspection, expected_kind=ImageKind(request.kind)) + requested_vllm_version = pinned_vllm_version(request.source) + if requested_vllm_version is not None: + inspected = inspection.inspection + if not isinstance(inspected, ServingImageInspection): + raise ImageVerificationError("versioned vLLM source requires a serving image inspection") + try: + version_matches = Version(inspected.runtime_version) == Version(requested_vllm_version) + except InvalidVersion: + version_matches = False + if not version_matches: + raise ImageVerificationError("inspected vLLM version does not match the requested image version") + artifact_directory = self._registry.image_root / _ARTIFACT_DIRECTORY_NAME + artifact_path = artifact_directory / f"{request.name}-{inspection.sqsh_sha256}.sqsh" + image = RegisteredImage( + schema_version=1, + name=request.name, + path=artifact_path.as_posix(), + sqsh_sha256=inspection.sqsh_sha256, + source_oci_digest=source_oci_digest, + inspection=inspection, + ) + created_artifact_identity: tuple[int, int] | None = None + with ( + _open_verified_sqsh( + candidate_path, + directory_descriptor=candidate_directory_descriptor, + ) as candidate_snapshot, + ExitStack() as final_stack, + ): + if candidate_snapshot.sha256 != inspection.sqsh_sha256: + raise ImageVerificationError("image inspection does not match the imported SQSH digest") + final_snapshot: _VerifiedSqshSnapshot | None = None + + def prepare_before_registry(_image: RegisteredImage) -> None: + nonlocal created_artifact_identity, final_snapshot + _verify_snapshot(image, candidate_snapshot) + ensure_private_directory(artifact_directory, parents=False) + created_artifact_identity = _publish_candidate( + candidate_path, + artifact_path, + expected_identity=candidate_snapshot.facts[:2], + candidate_directory_descriptor=candidate_directory_descriptor, + ) + try: + final_snapshot = final_stack.enter_context(_open_verified_sqsh(artifact_path)) + except ImageVerificationError as error: + if created_artifact_identity is None: + raise ImageConflictError( + f"image artifact path {artifact_path!s} cannot be reused as a verified SQSH artifact" + ) from error + raise + if final_snapshot.sha256 != inspection.sqsh_sha256: + if created_artifact_identity is not None: + raise ImageVerificationError("published SQSH artifact does not match its inspection digest") + raise ImageConflictError(f"image artifact path {artifact_path!s} contains conflicting bytes") + _sync_snapshot(final_snapshot) + _verify_snapshot(image, final_snapshot) + + def verify_before_registry(_image: RegisteredImage) -> None: + if final_snapshot is None: + raise ImageVerificationError("published SQSH artifact was not opened for verification") + _verify_snapshot_identity(image, final_snapshot) + + def verify_after_registry(_image: RegisteredImage) -> None: + if final_snapshot is None: + raise ImageVerificationError("published SQSH artifact was not opened for verification") + _verify_snapshot_identity(image, final_snapshot) + + def rollback_after_failure(_image: RegisteredImage) -> None: + if created_artifact_identity is not None: + _remove_artifact(artifact_path, expected_identity=created_artifact_identity) + + return self._registry.register( + image, + prepare_before_publish=prepare_before_registry, + verify_before_publish=verify_before_registry, + verify_after_publish=verify_after_registry, + rollback_after_failure=rollback_after_failure, + replace=replace, + ) + + def list_images(self) -> tuple[RegisteredImage, ...]: + """Return all registered aliases in deterministic order.""" + return self._registry.list_images() + + def resolve_for_planning(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: + """Resolve and reverify one registered alias or path for plan compilation.""" + if reference.name is not None: + image = self._registry.get_by_name(reference.name) + else: + assert reference.path is not None + image = self._registry.get_by_path(reference.path) + _verify_registered_image(image) + _validate_inspection( + image.inspection, + expected_kind=expected_kind, + expected_sha256=image.sqsh_sha256, + ) + if expected_kind is ImageKind.CLIENT: + _validate_client_proxy_dependency(image.inspection) + return ResolvedImage( + authored_ref=reference, + path=image.path, + sha256=image.sqsh_sha256, + inspection=image.inspection, + ) + + def unregister(self, name: str, *, delete_artifact: bool = False) -> RegisteredImage: + """Remove an alias and optionally its unshared, package-owned artifact.""" + return self._registry.unregister(name, delete_artifact=delete_artifact) + + +def compute_sqsh_file_sha256(path: Path) -> Sha256Digest: + """Compute the SHA-256 of one regular, non-symlink SQSH file.""" + with _open_verified_sqsh(path) as snapshot: + return snapshot.sha256 + + +def _validate_inspection( + inspection: ImageInspectionRecord, + *, + expected_kind: ImageKind, + expected_sha256: Sha256Digest | None = None, +) -> None: + if inspection.inspector_version != INSPECTOR_VERSION: + raise ImageVerificationError(f"unsupported image inspector version {inspection.inspector_version!r}") + if expected_sha256 is not None and inspection.sqsh_sha256 != expected_sha256: + raise ImageVerificationError("image inspection does not match the SQSH digest") + if inspection.inspection.kind is not expected_kind: + raise ImageVerificationError( + f"image inspection kind {inspection.inspection.kind.value!r} does not match {expected_kind.value!r}" + ) + + +def _validate_client_proxy_dependency(inspection: ImageInspectionRecord) -> None: + client = inspection.inspection + assert isinstance(client, ClientImageInspection) + version = next((item.version for item in client.distributions if item.name == "aiohttp"), None) + try: + compatible = version is not None and version in _PROXY_HTTP_REQUIREMENT + except InvalidVersion: + compatible = False + if not compatible: + raise ImageVerificationError( + "client image must install aiohttp>=3.14.3,<4 for the inference proxy; rebuild and replace the image" + ) + + +def _verify_registered_image(image: RegisteredImage) -> None: + with _open_verified_sqsh(Path(image.path)) as snapshot: + _verify_snapshot(image, snapshot) + + +@dataclass(frozen=True, slots=True) +class _VerifiedSqshSnapshot: + path: Path + descriptor: int + facts: tuple[int, int, int, int] + sha256: Sha256Digest + directory_descriptor: int | None = None + + def verify_unchanged(self) -> None: + """Verify that the open bytes and their public path still match the snapshot.""" + descriptor_status = os.fstat(self.descriptor) + if not stat.S_ISREG(descriptor_status.st_mode) or _get_file_facts(descriptor_status) != self.facts: + raise ImageVerificationError(f"image path {self.path} changed while it was being verified") + actual_sha256 = _hash_descriptor(self.descriptor) + descriptor_status = os.fstat(self.descriptor) + path_status = self._stat_path() + if ( + not stat.S_ISREG(descriptor_status.st_mode) + or not stat.S_ISREG(path_status.st_mode) + or _get_file_facts(descriptor_status) != self.facts + or _get_file_facts(path_status) != self.facts + or actual_sha256 != self.sha256 + ): + raise ImageVerificationError(f"image path {self.path} changed while it was being verified") + + def verify_identity_unchanged(self) -> None: + """Verify descriptor and path identity without rereading the SQSH bytes.""" + descriptor_status = os.fstat(self.descriptor) + path_status = self._stat_path() + if ( + not stat.S_ISREG(descriptor_status.st_mode) + or not stat.S_ISREG(path_status.st_mode) + or _get_file_facts(descriptor_status) != self.facts + or _get_file_facts(path_status) != self.facts + ): + raise ImageVerificationError(f"image path {self.path} changed while it was being verified") + + def _stat_path(self) -> os.stat_result: + if self.directory_descriptor is None: + return self.path.lstat() + return os.stat( + self.path.name, + dir_fd=self.directory_descriptor, + follow_symlinks=False, + ) + + +@contextmanager +def _open_verified_sqsh( + path: Path, + *, + directory_descriptor: int | None = None, +) -> Iterator[_VerifiedSqshSnapshot]: + descriptor: int | None = None + try: + before_open = ( + path.lstat() + if directory_descriptor is None + else os.stat(path.name, dir_fd=directory_descriptor, follow_symlinks=False) + ) + if not stat.S_ISREG(before_open.st_mode): + raise ImageVerificationError(f"image path {path} must be a regular non-symlink file") + open_path = path if directory_descriptor is None else Path(path.name) + descriptor = os.open( + open_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + after_open = os.fstat(descriptor) + if not stat.S_ISREG(after_open.st_mode) or (before_open.st_dev, before_open.st_ino) != ( + after_open.st_dev, + after_open.st_ino, + ): + raise ImageVerificationError(f"image path {path} changed while it was being opened") + facts = _get_file_facts(after_open) + sha256 = _hash_descriptor(descriptor) + snapshot = _VerifiedSqshSnapshot( + path=path, + descriptor=descriptor, + facts=facts, + sha256=sha256, + directory_descriptor=directory_descriptor, + ) + snapshot.verify_identity_unchanged() + yield snapshot + except FileNotFoundError as error: + raise ImageVerificationError(f"image path {path} does not exist") from error + except ImageVerificationError: + raise + except OSError as error: + raise ImageVerificationError(f"cannot read image path {path}") from error + finally: + if descriptor is not None: + os.close(descriptor) + + +def _hash_descriptor(descriptor: int) -> Sha256Digest: + digest = hashlib.sha256() + offset = 0 + while block := os.pread(descriptor, _HASH_BLOCK_SIZE, offset): + digest.update(block) + offset += len(block) + return digest.hexdigest() + + +def _get_file_facts(status: os.stat_result) -> tuple[int, int, int, int]: + # Publishing through a hard link changes ctime without changing the verified + # inode or its bytes, so ctime cannot be part of the stable identity tuple. + return (status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns) + + +def _verify_snapshot(image: RegisteredImage, snapshot: _VerifiedSqshSnapshot) -> None: + try: + snapshot.verify_unchanged() + except ImageVerificationError as error: + raise ImageVerificationError(f"registered image {image.name!r} no longer matches its SQSH digest") from error + if snapshot.sha256 != image.sqsh_sha256: + raise ImageVerificationError(f"registered image {image.name!r} no longer matches its SQSH digest") + + +def _verify_snapshot_identity(image: RegisteredImage, snapshot: _VerifiedSqshSnapshot) -> None: + try: + snapshot.verify_identity_unchanged() + except ImageVerificationError as error: + raise ImageVerificationError(f"registered image {image.name!r} no longer matches its SQSH digest") from error + if snapshot.sha256 != image.sqsh_sha256: + raise ImageVerificationError(f"registered image {image.name!r} no longer matches its SQSH digest") + + +def _publish_candidate( + candidate_path: Path, + artifact_path: Path, + *, + expected_identity: tuple[int, int], + candidate_directory_descriptor: int | None = None, +) -> tuple[int, int] | None: + try: + with ExitStack() as stack: + if candidate_directory_descriptor is None: + candidate_directory_descriptor = stack.enter_context(open_verified_directory(candidate_path.parent)) + artifact_directory_descriptor = stack.enter_context(open_verified_directory(artifact_path.parent)) + candidate_status = os.stat( + candidate_path.name, + dir_fd=candidate_directory_descriptor, + follow_symlinks=False, + ) + if ( + not stat.S_ISREG(candidate_status.st_mode) + or ( + candidate_status.st_dev, + candidate_status.st_ino, + ) + != expected_identity + ): + raise OSError(f"imported SQSH candidate {candidate_path!s} changed before publication") + try: + os.link( + candidate_path.name, + artifact_path.name, + src_dir_fd=candidate_directory_descriptor, + dst_dir_fd=artifact_directory_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + return None + try: + published_status = os.stat( + artifact_path.name, + dir_fd=artifact_directory_descriptor, + follow_symlinks=False, + ) + published_identity = (published_status.st_dev, published_status.st_ino) + if published_identity != expected_identity: + raise OSError(f"published SQSH artifact {artifact_path!s} has an unexpected identity") + os.fsync(artifact_directory_descriptor) + candidate_status = os.stat( + candidate_path.name, + dir_fd=candidate_directory_descriptor, + follow_symlinks=False, + ) + if (candidate_status.st_dev, candidate_status.st_ino) != published_identity: + raise OSError(f"imported SQSH candidate {candidate_path!s} changed during publication") + os.unlink(candidate_path.name, dir_fd=candidate_directory_descriptor) + os.fsync(candidate_directory_descriptor) + except OSError: + _remove_created_artifact_from_directory( + artifact_directory_descriptor, + artifact_path.name, + ) + raise + return published_identity + except OSError as error: + raise ImageVerificationError(f"cannot publish image artifact {artifact_path!s}") from error + + +def _sync_snapshot(snapshot: _VerifiedSqshSnapshot) -> None: + try: + os.fsync(snapshot.descriptor) + except OSError as error: + raise ImageVerificationError(f"cannot synchronize image artifact {snapshot.path!s}") from error + + +def _remove_artifact(artifact_path: Path, *, expected_identity: tuple[int, int]) -> None: + try: + with open_verified_directory(artifact_path.parent) as artifact_directory_descriptor: + _remove_artifact_from_directory( + artifact_directory_descriptor, + artifact_path.name, + expected_identity=expected_identity, + ) + except OSError: + pass + + +def _remove_artifact_from_directory( + artifact_directory_descriptor: int, + artifact_name: str, + *, + expected_identity: tuple[int, int], +) -> None: + try: + artifact_status = os.stat( + artifact_name, + dir_fd=artifact_directory_descriptor, + follow_symlinks=False, + ) + if (artifact_status.st_dev, artifact_status.st_ino) != expected_identity: + return + os.unlink(artifact_name, dir_fd=artifact_directory_descriptor) + os.fsync(artifact_directory_descriptor) + except OSError: + pass + + +def _remove_created_artifact_from_directory( + artifact_directory_descriptor: int, + artifact_name: str, +) -> None: + try: + os.unlink(artifact_name, dir_fd=artifact_directory_descriptor) + os.fsync(artifact_directory_descriptor) + except OSError: + pass diff --git a/packages/data-designer-slurm/src/data_designer/slurm/images/vllm_source.py b/packages/data-designer-slurm/src/data_designer/slurm/images/vllm_source.py new file mode 100644 index 000000000..e9252e034 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/vllm_source.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve versioned public vLLM images before the immutable OCI import path.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Mapping +from http.client import HTTPException +from urllib.error import HTTPError, URLError +from urllib.parse import urlencode +from urllib.request import HTTPRedirectHandler, Request, build_opener + +_VLLM_REPOSITORY = "vllm/vllm-openai" +_VERSION = r"v[0-9]+\.[0-9]+\.[0-9]+" +_VERSIONED_TAG = re.compile(rf"^(?:docker\.io/)?{_VLLM_REPOSITORY}:(?P{_VERSION})$") +_PINNED_VERSION = re.compile(rf"^(?:docker\.io/)?{_VLLM_REPOSITORY}:(?P{_VERSION})@sha256:[0-9a-f]{{64}}$") +_SHA256 = re.compile(r"^sha256:[0-9a-f]{64}$") +_TOKEN_URL = "https://auth.docker.io/token" +_MANIFEST_URL = f"https://registry-1.docker.io/v2/{_VLLM_REPOSITORY}/manifests/" +_MANIFEST_ACCEPT = ", ".join( + ( + "application/vnd.oci.image.index.v1+json", + "application/vnd.docker.distribution.manifest.list.v2+json", + "application/vnd.oci.image.manifest.v1+json", + "application/vnd.docker.distribution.manifest.v2+json", + ) +) +_MAX_TOKEN_RESPONSE_BYTES = 65536 +_REQUEST_TIMEOUT_SECONDS = 10 + + +class VllmSourceResolutionError(Exception): + """A versioned vLLM source could not be pinned to an OCI digest.""" + + +class VllmTagNotFoundError(VllmSourceResolutionError): + """The requested official vLLM release tag does not exist.""" + + +class _RejectRedirects(HTTPRedirectHandler): + def redirect_request( + self, + request: Request, + response: object, + code: int, + message: str, + headers: object, + new_url: str, + ) -> None: + return None + + +def is_versioned_vllm_tag(source: str) -> bool: + """Return whether a source is an unpinned official vLLM version tag.""" + return _VERSIONED_TAG.fullmatch(source) is not None + + +def pinned_vllm_version(source: str) -> str | None: + """Extract the release version recorded alongside an immutable OCI digest.""" + match = _PINNED_VERSION.fullmatch(source) + return None if match is None else match.group("tag").removeprefix("v") + + +def resolve_versioned_vllm_source(source: str) -> str: + """Resolve one official vLLM version tag and return a digest-pinned source. + + Only fixed Docker Hub hosts and the official vLLM repository are contacted. + Bearer credentials never enter persisted plans, output, or exception text. + """ + match = _VERSIONED_TAG.fullmatch(source) + if match is None: + raise ValueError("source is not an official versioned vLLM image") + tag = match.group("tag") + try: + query = urlencode({"service": "registry.docker.io", "scope": f"repository:{_VLLM_REPOSITORY}:pull"}) + token_body, _ = _request(Request(f"{_TOKEN_URL}?{query}", method="GET")) + token_payload = json.loads(token_body) + token = token_payload.get("token") if isinstance(token_payload, dict) else None + if not isinstance(token, str) or not token or len(token) > 16384: + raise VllmSourceResolutionError("Docker Hub did not return a valid vLLM pull token") + try: + _, headers = _request( + Request( + f"{_MANIFEST_URL}{tag}", + headers={"Authorization": f"Bearer {token}", "Accept": _MANIFEST_ACCEPT}, + method="HEAD", + ) + ) + except HTTPError as error: + if error.code == 404: + raise VllmTagNotFoundError("the requested vLLM image version was not found") from None + raise + digest = headers.get("Docker-Content-Digest", "") + if _SHA256.fullmatch(digest) is None: + raise VllmSourceResolutionError("Docker Hub did not return a valid vLLM manifest digest") + except HTTPError: + raise VllmSourceResolutionError("the vLLM image version could not be resolved") from None + except (URLError, TimeoutError, OSError, HTTPException, UnicodeError, json.JSONDecodeError): + raise VllmSourceResolutionError("the vLLM image version could not be resolved") from None + return f"{_VLLM_REPOSITORY}:{tag}@{digest}" + + +def _request(request: Request) -> tuple[bytes, Mapping[str, str]]: + opener = build_opener(_RejectRedirects) + with opener.open(request, timeout=_REQUEST_TIMEOUT_SECONDS) as response: + body = response.read(_MAX_TOKEN_RESPONSE_BYTES + 1) + if len(body) > _MAX_TOKEN_RESPONSE_BYTES: + raise VllmSourceResolutionError("Docker Hub returned an oversized response") + return body, response.headers diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py new file mode 100644 index 000000000..f19f4a04e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public integration validation across Slurm plan and runtime-state contracts.""" + +from __future__ import annotations + +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError + +IntegrationContractError = PlanStateContractError + + +class PlanStateValidator(PersistedPlanStateValidator): + """Public compatibility facade for persisted-plan validation.""" + + +__all__ = [ + "IntegrationContractError", + "PlanStateValidator", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py new file mode 100644 index 000000000..61da221a6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal Slurm submission, observation, and batch-rendering helpers.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py new file mode 100644 index 000000000..7342d3922 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared safe rendering primitives for package-owned Slurm batch scripts.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +from data_designer.slurm.launcher.errors import SlurmBatchRenderError + +_DIRECTIVE_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9-]*$") +_DIRECTIVE_TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:/,%+-]*$") + + +@dataclass(frozen=True) +class BatchDirective: + """One validated ``#SBATCH`` option.""" + + name: str + value: str + + def render(self) -> str: + """Render the directive as one non-executable scheduler line.""" + if type(self.name) is not str or _DIRECTIVE_NAME_PATTERN.fullmatch(self.name) is None: + raise SlurmBatchRenderError("batch directive name is invalid") + if type(self.value) is not str: + raise SlurmBatchRenderError("batch directive value must be text") + reject_control_characters(self.value, field_name=f"--{self.name} value") + value = self.value if _DIRECTIVE_TOKEN_PATTERN.fullmatch(self.value) else _quote_sbatch_option_value(self.value) + return f"#SBATCH --{self.name}={value}" + + +def render_batch_directives(values: tuple[tuple[str, str | None], ...]) -> str: + """Render ordered optional directive values as safe scheduler lines.""" + return "\n".join(BatchDirective(name=name, value=value).render() for name, value in values if value is not None) + + +def quote_shell_value(value: str) -> str: + """Quote one validated value for literal use in package-owned Bash.""" + if type(value) is not str: + raise SlurmBatchRenderError("shell value must be text") + reject_control_characters(value, field_name="shell value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") + return f'"{escaped}"' + + +def reject_control_characters(value: str, *, field_name: str) -> None: + """Reject characters that can split a directive or shell assignment.""" + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SlurmBatchRenderError(f"{field_name} must not contain control characters") + + +def _quote_sbatch_option_value(value: str) -> str: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py new file mode 100644 index 000000000..a872daa19 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal typed argument-vector client for Slurm command-line tools.""" + +from __future__ import annotations + +import os +import re +import subprocess +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta +from pathlib import Path + +from data_designer.slurm.contracts import Identifier, PartitionSelection +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmNamedJobEntry, + SlurmQueueEntry, + SlurmSubmissionMatch, +) +from data_designer.slurm.launcher.parsing import ( + parse_accounting, + parse_default_partition_gpu_counts, + parse_gpu_counts, + parse_named_jobs, + parse_queue, + parse_selected_partition_gpu_counts, + parse_submission, +) +from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner +from data_designer.slurm.security import redact_sensitive_diagnostic +from data_designer.slurm.state import SchedulerIdentity, SchedulerJobIdentity + +_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_PARTITION_SELECTION_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}(?:,[A-Za-z0-9][A-Za-z0-9._-]{0,127})*$") +_ENVIRONMENT_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 +_UNKNOWN_QUEUE_JOB_DETAILS = frozenset( + { + "Invalid job id specified", + "slurm_load_jobs error: Invalid job id specified", + } +) + + +@dataclass(frozen=True) +class SlurmExecutables: + """Executable paths used for bounded Slurm operations.""" + + sbatch: str = "sbatch" + squeue: str = "squeue" + sacct: str = "sacct" + scancel: str = "scancel" + scontrol: str = "scontrol" + sinfo: str = "sinfo" + + def __post_init__(self) -> None: + for executable in (self.sbatch, self.squeue, self.sacct, self.scancel, self.scontrol, self.sinfo): + _validate_argument(executable, field_name="Slurm executable") + if any(character.isspace() for character in executable): + raise ValueError("Slurm executable must be one argument-vector token") + + +class SlurmCommandClient: + """Submit, observe, and cancel Slurm jobs through structured commands.""" + + _executables: SlurmExecutables + _runner: CommandRunner + + def __init__( + self, + runner: CommandRunner | None = None, + *, + executables: SlurmExecutables | None = None, + ) -> None: + self._runner = runner if runner is not None else SubprocessRunner() + self._executables = executables if executables is not None else SlurmExecutables() + + def submit( + self, + script_path: str | Path, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: + """Submit one rendered batch script and return its assigned job ID.""" + path = str(script_path) + _validate_argument(path, field_name="batch script path") + if path.startswith("-"): + raise ValueError("batch script path must not begin with '-'; prefix relative paths with './'") + hold_arguments = ("--hold",) if hold else () + export_argument, environment = _format_export_environment(export_environment) + output = self._run( + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument, path), + environment=environment, + ) + return parse_submission(output) + + def submit_script( + self, + script: str, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: + """Submit verified batch-script text through standard input.""" + if type(script) is not str or not script or "\0" in script: + raise ValueError("batch script text must be non-empty UTF-8 text without NUL") + hold_arguments = ("--hold",) if hold else () + export_argument, environment = _format_export_environment(export_environment) + try: + output = self._run( + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), + input_text=script, + environment=environment, + ) + except SlurmCommandError as error: + raise SlurmSubmissionError( + str(error), + may_have_succeeded=error.command_may_have_completed, + ) from error + try: + return parse_submission(output) + except SlurmCommandOutputError as error: + raise SlurmSubmissionError(str(error), may_have_succeeded=True) from error + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + """Return normalized active-queue rows for explicit managed jobs.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + try: + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) + ) + except SlurmCommandError as error: + if _is_unknown_queue_job_error(error): + return () + raise + entries = parse_queue(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), + requested, + command="squeue", + ) + return tuple(entry for entry in entries if entry.job_identity not in ignored) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + """Return normalized accounting rows for explicit managed jobs.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + output = self._run( + ( + self._executables.sacct, + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", + f"--jobs={jobs}", + ) + ) + entries = parse_accounting(output) + ignored = _validate_observed_job_identities( + tuple(entry.job_identity for entry in entries), + requested, + command="sacct", + ) + return tuple(entry for entry in entries if entry.job_identity not in ignored) + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return current-user allocations matching one exact recovery name.""" + if type(job_name) is not str or _IDENTIFIER_PATTERN.fullmatch(job_name) is None: + raise ValueError("Slurm job name must be a valid identifier") + accounting_start = _format_accounting_start(submitted_after) + queue_output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%.128j", + "--me", + f"--name={job_name}", + ) + ) + accounting_output = self._run( + ( + self._executables.sacct, + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,JobName%128", + f"--uid={os.getuid()}", + f"--starttime={accounting_start}", + f"--name={job_name}", + ) + ) + entries = ( + *parse_named_jobs(queue_output, command="squeue"), + *parse_named_jobs(accounting_output, command="sacct"), + ) + if any(entry.job_name != job_name for entry in entries): + raise SlurmCommandOutputError("scheduler returned a job outside the requested exact name") + return _merge_submission_matches(entries) + + def cancel(self, selector: SchedulerJobIdentity) -> None: + """Cancel one managed Slurm job, array, or array task.""" + self._run((self._executables.scancel, _format_selector(selector))) + + def release(self, job_id: int) -> None: + """Release one held managed Slurm job or array.""" + self._run((self._executables.scontrol, "release", _format_job_id(job_id))) + + def query_gpu_counts(self, *, partition: PartitionSelection | None = None) -> tuple[int, ...]: + """Return configured GPU counts for requested partitions or the default.""" + if partition is None: + command = (self._executables.sinfo, "--noheader", "--format=%P|%G") + return parse_default_partition_gpu_counts(self._run(command)) + if ( + type(partition) is not str + or len(partition) > 1024 + or _PARTITION_SELECTION_PATTERN.fullmatch(partition) is None + or len(set(partition.split(","))) != len(partition.split(",")) + ): + raise ValueError("Slurm partitions must be a comma-separated list of unique valid identifiers") + partitions = tuple(partition.split(",")) + if len(partitions) == 1: + command = (self._executables.sinfo, "--noheader", "--format=%G", f"--partition={partition}") + return parse_gpu_counts(self._run(command)) + command = (self._executables.sinfo, "--noheader", "--format=%P|%G", f"--partition={partition}") + return parse_selected_partition_gpu_counts(self._run(command), partitions=partitions) + + def _run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> str: + command_name = Path(command[0]).name + try: + if environment is None: + completed = ( + self._runner.run(command) + if input_text is None + else self._runner.run(command, input_text=input_text) + ) + else: + completed = self._runner.run(command, input_text=input_text, environment=environment) + except subprocess.TimeoutExpired as error: + raise SlurmCommandError( + f"{command_name} could not be executed: {_format_error_detail(error)}", + command_may_have_completed=True, + ) from error + except OSError as error: + raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error + except subprocess.SubprocessError as error: + raise SlurmCommandError( + f"{command_name} could not be executed: {_format_error_detail(error)}", + command_may_have_completed=True, + ) from error + returncode = getattr(completed, "returncode", None) + stdout = getattr(completed, "stdout", None) + stderr = getattr(completed, "stderr", None) + if type(returncode) is not int or not isinstance(stdout, str) or not isinstance(stderr, str): + raise SlurmCommandError( + f"{command_name} returned a malformed process result", + command_may_have_completed=True, + ) + if returncode: + detail = _normalize_bounded_text(stderr) or "no diagnostic output" + raise SlurmCommandError(f"{command_name} failed with exit code {returncode}: {detail}") + return stdout + + +def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[str, Mapping[str, str]]: + """Export selected runner values without loading a node login environment.""" + exported = dict(environment or {}) + exported.setdefault("NEMO_DEPLOYMENT_TYPE", os.environ.get("NEMO_DEPLOYMENT_TYPE", "slurm")) + names = tuple(sorted(exported)) + for name in names: + if _ENVIRONMENT_NAME_PATTERN.fullmatch(name) is None: + raise ValueError("exported environment names must be valid identifiers") + value = exported[name] + if type(value) is not str or "\0" in value: + raise ValueError("exported environment values must be strings without NUL") + return "--export=ALL", exported + + +def _format_selectors(selectors: Sequence[SchedulerJobIdentity]) -> str: + if not selectors: + raise ValueError("at least one managed Slurm job selector is required") + return ",".join(dict.fromkeys(_format_selector(selector) for selector in selectors)) + + +def _format_selector(selector: SchedulerJobIdentity) -> str: + if isinstance(selector, SchedulerIdentity): + job_id = _format_job_id(selector.array_job_id) + if selector.array_task_id > _MAX_SLURM_INTEGER: + raise ValueError("Slurm array-task IDs must be non-negative 32-bit integers") + return f"{job_id}_{selector.array_task_id}" + return _format_job_id(selector) + + +def _format_job_id(value: object) -> str: + if type(value) is not int or not 0 < value <= _MAX_SLURM_INTEGER: + raise ValueError("Slurm job IDs must be positive 32-bit integers") + return str(value) + + +def _validate_observed_job_identities( + job_identities: Sequence[SchedulerJobIdentity], + selectors: Sequence[SchedulerJobIdentity], + *, + command: str, +) -> frozenset[SchedulerJobIdentity]: + """Validate result correlation and return unselected aggregate rows.""" + selected_job_ids = {selector for selector in selectors if type(selector) is int} + selected_array_tasks = {selector for selector in selectors if isinstance(selector, SchedulerIdentity)} + selected_array_job_ids = {selector.array_job_id for selector in selected_array_tasks} + ignored: set[SchedulerJobIdentity] = set() + for job_identity in job_identities: + if type(job_identity) is int and job_identity in selected_job_ids: + continue + if type(job_identity) is int and job_identity in selected_array_job_ids: + ignored.add(job_identity) + continue + if isinstance(job_identity, SchedulerIdentity) and ( + job_identity in selected_array_tasks or job_identity.array_job_id in selected_job_ids + ): + continue + raise SlurmCommandOutputError(f"{command} returned an unrequested job or array-task ID") + return frozenset(ignored) + + +def _validate_argument(value: str, *, field_name: str) -> None: + if type(value) is not str or not value: + raise ValueError(f"{field_name} must not be empty") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError(f"{field_name} must not contain control characters") + + +def _normalize_bounded_text(value: str, *, limit: int = 512) -> str: + redacted = redact_sensitive_diagnostic(value) + normalized = " ".join(redacted.split()) + return normalized if len(normalized) <= limit else f"{normalized[: limit - 3]}..." + + +def _format_error_detail(error: BaseException) -> str: + if isinstance(error, subprocess.TimeoutExpired): + return "command timed out" + return _normalize_bounded_text(str(error)) or error.__class__.__name__ + + +def _is_unknown_queue_job_error(error: SlurmCommandError) -> bool: + message = str(error) + return any(message.endswith(f": {detail}") for detail in _UNKNOWN_QUEUE_JOB_DETAILS) + + +def _format_accounting_start(value: datetime) -> str: + if not isinstance(value, datetime) or value.tzinfo is None or value.utcoffset() is None: + raise ValueError("submission lookup timestamp must be timezone-aware") + return (value - timedelta(minutes=1)).astimezone().strftime("%Y-%m-%dT%H:%M:%S") + + +def _merge_submission_matches(entries: Sequence[SlurmNamedJobEntry]) -> tuple[SlurmSubmissionMatch, ...]: + grouped: dict[tuple[int, str], set[int]] = {} + ordinary: set[tuple[int, str]] = set() + for entry in entries: + key = (entry.job_id, entry.job_name) + if entry.array_task_id is None: + ordinary.add(key) + else: + grouped.setdefault(key, set()).add(entry.array_task_id) + matches: list[SlurmSubmissionMatch] = [] + for job_id, job_name in sorted(ordinary | set(grouped)): + key = (job_id, job_name) + task_ids = tuple(sorted(grouped[key])) if key in grouped else None + matches.append(SlurmSubmissionMatch(job_id=job_id, job_name=job_name, array_task_ids=task_ids)) + return tuple(matches) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py new file mode 100644 index 000000000..79cac21d3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Safe rendering for zero-GPU CPU collection jobs.""" + +from __future__ import annotations + +import posixpath +from importlib import resources + +from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.destinations import CollectionDestination +from data_designer.slurm.state.outputs import CollectionPlan + + +def render_collection_script( + resolved_plan: ResolvedSlurmRunPlan, + collection_plan: CollectionPlan, + destination: CollectionDestination, +) -> str: + """Render one CPU-only job that invokes the allocation-gated collection worker.""" + if collection_plan.run_id != resolved_plan.run_id: + raise SlurmBatchRenderError("collection run identity does not match the resolved plan") + if collection_plan.host_destination != destination.host_path: + raise SlurmBatchRenderError("collection host destination does not match its resolved mount") + + collection_root = posixpath.join( + posixpath.dirname(resolved_plan.authored_config.path), + "collections", + collection_plan.collection_id, + ) + collection_plan_path = posixpath.join(collection_root, "plan.json") + directives = render_batch_directives( + ( + ("job-name", collection_plan.submission_job_name), + ("account", resolved_plan.submission.account), + ("partition", resolved_plan.selected_profile.profile.image_build.partition), + ("nodes", "1"), + ("ntasks", "1"), + ("cpus-per-task", str(resolved_plan.client.authored.cpus)), + ("time", resolved_plan.submission.time_limit), + ("chdir", collection_root), + ("output", f"{collection_root}/slurm-%j.out"), + ("error", f"{collection_root}/slurm-%j.err"), + ) + ) + scratch_source = resources.files("data_designer.slurm.runtime").joinpath("scratch.sh").read_text() + package_version = collection_plan.package_version or resolved_plan.package_version + runtime_bundle = collection_plan.runtime_bundle or resolved_plan.runtime_bundle + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +export PYTHONNOUSERSITE=1 + +{scratch_source} +readonly DD_COLLECTION_PLAN={quote_shell_value(collection_plan_path)} +readonly DD_COLLECTION_PLAN_SHA256={quote_shell_value(collection_plan.compute_sha256())} +readonly DD_COLLECTION_PYTHON={quote_shell_value(collection_plan.python_executable)} +readonly DD_PACKAGE_VERSION={quote_shell_value(package_version)} +readonly DD_RUNTIME_ARCHIVE={quote_shell_value(runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={quote_shell_value(runtime_bundle.sha256)} +readonly DD_WORKSPACE_ROOT={quote_shell_value(resolved_plan.selected_profile.profile.workspace_root)} +readonly DD_RUN_ID={quote_shell_value(resolved_plan.run_id)} +readonly DD_COLLECTION_ID={quote_shell_value(collection_plan.collection_id)} + +verify_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] +}} + +verify_sha256 "${{DD_COLLECTION_PLAN_SHA256}}" "${{DD_COLLECTION_PLAN}}" +[[ -x "${{DD_COLLECTION_PYTHON}}" ]] +actual_package_version="$("${{DD_COLLECTION_PYTHON}}" -c \ + 'from importlib.metadata import version; print(version("data-designer-slurm"))')" +[[ "${{actual_package_version}}" == "${{DD_PACKAGE_VERSION}}" ]] +trap 'status=$?; dd_remove_local_scratch "${{DD_SCRATCH_ROOT}}" || true; exit "${{status}}"' EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +dd_initialize_local_scratch +dd_extract_runtime_bundle "${{DD_RUNTIME_ARCHIVE}}" "${{DD_RUNTIME_SHA256}}" +export PYTHONPATH="${{DD_SCRATCH_ROOT}}/runtime" +"${{DD_COLLECTION_PYTHON}}" -m data_designer.slurm.state.collection_worker \ + --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" --collection-id "${{DD_COLLECTION_ID}}" +""" + + +__all__ = ["render_collection_script"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py new file mode 100644 index 000000000..6e483d86e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal normalized errors for the Slurm launcher boundary.""" + +from __future__ import annotations + + +class SlurmLauncherError(RuntimeError): + """Base error for structured Slurm launcher operations.""" + + +class SlurmCommandError(SlurmLauncherError): + """A Slurm command could not be executed successfully.""" + + def __init__(self, message: str, *, command_may_have_completed: bool = False) -> None: + super().__init__(message) + self.command_may_have_completed = command_may_have_completed + + +class SlurmSubmissionError(SlurmLauncherError): + """An sbatch submission failed with an explicit ambiguity classification.""" + + def __init__(self, message: str, *, may_have_succeeded: bool) -> None: + super().__init__(message) + self.may_have_succeeded = may_have_succeeded + + +class SlurmCommandOutputError(SlurmLauncherError, ValueError): + """A Slurm command returned output that violates its requested format.""" + + +class SlurmBatchRenderError(SlurmLauncherError, ValueError): + """A resolved plan cannot be rendered as a safe batch script.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py new file mode 100644 index 000000000..64111099e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transient typed values returned by Slurm commands.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from data_designer.slurm.state import SchedulerJobIdentity, SchedulerState + + +@dataclass(frozen=True) +class SlurmJobSubmissionReceipt: + """Job identity returned for one accepted non-federated submission.""" + + job_id: int + + +@dataclass(frozen=True) +class SlurmProcessExitCode: + """Slurm's process status and terminating signal pair.""" + + exit_status: int + termination_signal: int + + +@dataclass(frozen=True) +class SlurmQueueEntry: + """One transient normalized active-queue entry.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + + +@dataclass(frozen=True) +class SlurmAccountingEntry: + """One transient normalized accounting entry.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + process_exit_code: SlurmProcessExitCode + + +@dataclass(frozen=True) +class SlurmNamedJobEntry: + """One scheduler allocation found through an exact job-name lookup.""" + + job_id: int + array_task_id: int | None + job_name: str + + +@dataclass(frozen=True) +class SlurmSubmissionMatch: + """One named scheduler allocation with its ordinary or array shape.""" + + job_id: int + job_name: str + array_task_ids: tuple[int, ...] | None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py new file mode 100644 index 000000000..41a9f865b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal strict parsers for bounded, machine-readable Slurm output.""" + +from __future__ import annotations + +import re +from typing import Literal + +from data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmNamedJobEntry, + SlurmProcessExitCode, + SlurmQueueEntry, +) +from data_designer.slurm.state import SchedulerIdentity, SchedulerJobIdentity, SchedulerState + +_ARRAY_ID_PATTERN = re.compile(r"^(?P[1-9][0-9]*)_(?P[0-9]+)$") +_JOB_ID_PATTERN = re.compile(r"^[1-9][0-9]*$") +_CLUSTER_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_EXIT_CODE_PATTERN = re.compile(r"^(?P[0-9]+):(?P[0-9]+)$") +_GRES_GPU_PATTERN = re.compile(r"^gpu:(?:(?:[^:,()]+):)*(?P[1-9][0-9]*)(?:\([^\r\n]*\))?$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 + +_STATE_MAP = { + "BOOT_FAIL": SchedulerState.FAILED, + "CANCELLED": SchedulerState.CANCELLED, + "COMPLETED": SchedulerState.COMPLETED, + "COMPLETING": SchedulerState.RUNNING, + "CONFIGURING": SchedulerState.PENDING, + "DEADLINE": SchedulerState.FAILED, + "FAILED": SchedulerState.FAILED, + "NODE_FAIL": SchedulerState.NODE_FAILED, + "OUT_OF_MEMORY": SchedulerState.OUT_OF_MEMORY, + "PENDING": SchedulerState.PENDING, + "PREEMPTED": SchedulerState.PREEMPTED, + "REQUEUED": SchedulerState.REQUEUED, + "REQUEUE_FED": SchedulerState.PENDING, + "REQUEUE_HOLD": SchedulerState.PENDING, + "RESV_DEL_HOLD": SchedulerState.PENDING, + "RESIZING": SchedulerState.RUNNING, + "REVOKED": SchedulerState.FAILED, + "RUNNING": SchedulerState.RUNNING, + "SIGNALING": SchedulerState.RUNNING, + "SPECIAL_EXIT": SchedulerState.PENDING, + "STAGE_OUT": SchedulerState.RUNNING, + "STOPPED": SchedulerState.RUNNING, + "SUSPENDED": SchedulerState.RUNNING, + "TIMEOUT": SchedulerState.TIMED_OUT, +} + + +def parse_submission(output: str) -> SlurmJobSubmissionReceipt: + """Parse non-federated ``sbatch --parsable`` output.""" + value = output.strip() + job_id, separator, cluster_name = value.partition(";") + if not job_id.isascii() or not job_id.isdecimal(): + raise SlurmCommandOutputError("sbatch returned an invalid job ID") + parsed_job_id = _parse_decimal(job_id, message="sbatch returned an invalid job ID") + if parsed_job_id <= 0: + raise SlurmCommandOutputError("sbatch returned an invalid job ID") + if separator: + if _CLUSTER_NAME_PATTERN.fullmatch(cluster_name) is None: + raise SlurmCommandOutputError("sbatch returned an invalid cluster name") + raise SlurmCommandOutputError("federated Slurm submissions are not supported") + return SlurmJobSubmissionReceipt(job_id=parsed_job_id) + + +def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: + """Parse ``squeue --format=%i|%T`` rows.""" + entries: list[SlurmQueueEntry] = [] + identities: set[SchedulerJobIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 2: + raise SlurmCommandOutputError(f"squeue line {line_number} must contain two fields") + job_identity = _parse_job_identity(fields[0], command="squeue", line_number=line_number) + _reject_duplicate(job_identity, identities, command="squeue", line_number=line_number) + entries.append(SlurmQueueEntry(job_identity=job_identity, state=parse_state(fields[1]))) + return tuple(entries) + + +def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: + """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" + entries: list[SlurmAccountingEntry] = [] + identities: set[SchedulerJobIdentity] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = line.split("|") + if len(fields) != 3: + raise SlurmCommandOutputError(f"sacct line {line_number} must contain three fields") + job_identity = _parse_job_identity(fields[0], command="sacct", line_number=line_number) + _reject_duplicate(job_identity, identities, command="sacct", line_number=line_number) + entries.append( + SlurmAccountingEntry( + job_identity=job_identity, + state=parse_state(fields[1]), + process_exit_code=_parse_exit_code(fields[2], line_number=line_number), + ) + ) + return tuple(entries) + + +def parse_named_jobs(output: str, *, command: Literal["sacct", "squeue"]) -> tuple[SlurmNamedJobEntry, ...]: + """Parse scheduler allocations returned for an exact job-name lookup.""" + entries: list[SlurmNamedJobEntry] = [] + identities: set[tuple[int, int | None, str]] = set() + for line_number, line in _collect_nonempty_lines(output): + fields = tuple(field.strip() for field in line.split("|")) + if len(fields) != 2 or not fields[1]: + raise SlurmCommandOutputError(f"{command} line {line_number} must contain a job ID and name") + identity = _parse_job_identity(fields[0], command=command, line_number=line_number) + job_id = identity.array_job_id if isinstance(identity, SchedulerIdentity) else identity + array_task_id = identity.array_task_id if isinstance(identity, SchedulerIdentity) else None + key = (job_id, array_task_id, fields[1]) + if key in identities: + continue + identities.add(key) + entries.append(SlurmNamedJobEntry(job_id=job_id, array_task_id=array_task_id, job_name=fields[1])) + return tuple(entries) + + +def parse_gpu_counts(output: str) -> tuple[int, ...]: + """Parse configured per-node GPU counts from ``sinfo --format=%G`` rows.""" + counts: list[int] = [] + for line_number, line in _collect_nonempty_lines(output): + if line in {"(null)", "N/A"}: + continue + line_counts: list[int] = [] + for gres in _split_gres_fields(line, line_number=line_number): + if not gres.startswith("gpu:"): + continue + match = _GRES_GPU_PATTERN.fullmatch(gres) + if match is None: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + line_counts.append( + _parse_decimal( + match.group("count"), + message=f"sinfo line {line_number} contains an invalid GPU resource", + ) + ) + if line_counts: + counts.append(sum(line_counts)) + return tuple(counts) + + +def parse_selected_partition_gpu_counts(output: str, *, partitions: tuple[str, ...]) -> tuple[int, ...]: + """Require a GPU count from every requested partition in ``%P|%G`` rows.""" + requested = set(partitions) + found: set[str] = set() + counts: list[int] = [] + for line_number, line in _collect_nonempty_lines(output): + fields = tuple(field.strip() for field in line.split("|")) + if len(fields) != 2 or not all(fields): + raise SlurmCommandOutputError(f"sinfo line {line_number} must contain a partition and resources") + partition, gres = fields + partition = partition.removesuffix("*") + if partition not in requested: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an unexpected partition") + row_counts = parse_gpu_counts(gres) + if len(row_counts) != 1: + raise SlurmCommandOutputError(f"sinfo partition {partition!r} has no GPU resources") + found.add(partition) + counts.extend(row_counts) + if found != requested: + raise SlurmCommandOutputError("sinfo omitted a requested partition") + return tuple(counts) + + +def parse_default_partition_gpu_counts(output: str) -> tuple[int, ...]: + """Parse GPU counts from ``sinfo --format=%P|%G`` default-partition rows.""" + default_partition: str | None = None + resources: list[str] = [] + for line_number, line in _collect_nonempty_lines(output): + fields = tuple(field.strip() for field in line.split("|")) + if len(fields) != 2 or not all(fields): + raise SlurmCommandOutputError(f"sinfo line {line_number} must contain a partition and resources") + partition, gres = fields + if not partition.endswith("*"): + continue + partition = partition.removesuffix("*") + if _CLUSTER_NAME_PATTERN.fullmatch(partition) is None: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid default partition") + if default_partition is not None and partition != default_partition: + raise SlurmCommandOutputError("sinfo returned multiple default partitions") + default_partition = partition + resources.append(gres) + return parse_gpu_counts("\n".join(resources)) + + +def parse_state(value: str) -> SchedulerState: + """Normalize one Slurm long state spelling without guessing unknown states.""" + normalized = value.strip().upper().removesuffix("+") + if not normalized: + raise SlurmCommandOutputError("scheduler state must not be empty") + if normalized.startswith("CANCELLED BY "): + canceller = normalized.removeprefix("CANCELLED BY ") + if not canceller.isascii() or not canceller.isdecimal(): + raise SlurmCommandOutputError("cancelled scheduler state has an invalid owner") + normalized = "CANCELLED" + elif any(character.isspace() for character in normalized): + raise SlurmCommandOutputError("scheduler state contains unexpected whitespace") + return _STATE_MAP.get(normalized, SchedulerState.UNKNOWN) + + +def _split_gres_fields(value: str, *, line_number: int) -> tuple[str, ...]: + fields: list[str] = [] + start = 0 + annotation_depth = 0 + for index, character in enumerate(value): + if character == "(": + annotation_depth += 1 + if annotation_depth > 1: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + elif character == ")": + annotation_depth -= 1 + if annotation_depth < 0: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + elif character == "," and annotation_depth == 0: + fields.append(value[start:index]) + start = index + 1 + if annotation_depth: + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + fields.append(value[start:]) + if any(not field for field in fields): + raise SlurmCommandOutputError(f"sinfo line {line_number} contains an invalid GPU resource") + return tuple(fields) + + +def _collect_nonempty_lines(output: str) -> tuple[tuple[int, str], ...]: + return tuple( + (line_number, line) + for line_number, raw_line in enumerate(output.splitlines(), start=1) + if (line := raw_line.strip()) + ) + + +def _parse_array_identity(value: str, *, command: str, line_number: int) -> SchedulerIdentity: + match = _ARRAY_ID_PATTERN.fullmatch(value) + if match is None: + raise SlurmCommandOutputError(f"{command} line {line_number} contains an invalid array-task ID") + message = f"{command} line {line_number} contains an invalid array-task ID" + return SchedulerIdentity( + array_job_id=_parse_decimal(match.group("job"), message=message), + array_task_id=_parse_decimal(match.group("task"), message=message), + ) + + +def _parse_job_identity(value: str, *, command: str, line_number: int) -> SchedulerJobIdentity: + message = f"{command} line {line_number} contains an invalid job or array-task ID" + if _JOB_ID_PATTERN.fullmatch(value) is not None: + return _parse_decimal(value, message=message) + try: + return _parse_array_identity(value, command=command, line_number=line_number) + except SlurmCommandOutputError as error: + raise SlurmCommandOutputError(message) from error + + +def _parse_exit_code(value: str, *, line_number: int) -> SlurmProcessExitCode: + match = _EXIT_CODE_PATTERN.fullmatch(value) + if match is None: + raise SlurmCommandOutputError(f"sacct line {line_number} contains an invalid exit code") + message = f"sacct line {line_number} contains an invalid exit code" + return SlurmProcessExitCode( + exit_status=_parse_decimal(match.group("status"), message=message), + termination_signal=_parse_decimal(match.group("signal"), message=message), + ) + + +def _parse_decimal(value: str, *, message: str) -> int: + if len(value) > 10: + raise SlurmCommandOutputError(message) + try: + parsed = int(value) + except ValueError as error: + raise SlurmCommandOutputError(message) from error + if parsed > _MAX_SLURM_INTEGER: + raise SlurmCommandOutputError(message) + return parsed + + +def _reject_duplicate( + job_identity: SchedulerJobIdentity, + identities: set[SchedulerJobIdentity], + *, + command: str, + line_number: int, +) -> None: + if job_identity in identities: + raise SlurmCommandOutputError(f"{command} line {line_number} duplicates a job or array-task ID") + identities.add(job_identity) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py new file mode 100644 index 000000000..8b57f5149 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal safe rendering for thin deterministic Slurm batch entrypoints.""" + +from __future__ import annotations + +import posixpath + +from data_designer.slurm.images.records import validate_enroot_mount_path +from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.outputs import RetryPlan + +_SYSTEM_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + + +def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: + """Render a resolved generation plan as one thin deterministic entrypoint.""" + if type(attempt_ordinal) is not int or attempt_ordinal <= 0: + raise SlurmBatchRenderError("attempt_ordinal must be a positive integer") + + run_root = posixpath.dirname(plan.authored_config.path) + plan_path = posixpath.join(run_root, "resolved-plan.json") + attempt = f"{attempt_ordinal:04d}" + directive_text = render_batch_directives( + _build_generation_directives( + plan, + output_path=posixpath.join(run_root, f"slurm-attempt-{attempt}-%A_%a.out"), + error_path=posixpath.join(run_root, f"slurm-attempt-{attempt}-%A_%a.err"), + ) + ) + if _requires_exclusive_nodes(plan): + directive_text = f"{directive_text}\n#SBATCH --exclusive" + return f"""#!/usr/bin/env bash +{directive_text} +set -Eeuo pipefail +export PATH={quote_shell_value(_get_command_path(plan))} +export PYTHONNOUSERSITE=1 + +readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={quote_shell_value(plan_path)} +readonly DD_PLAN_SHA256={quote_shell_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={quote_shell_value(run_root)} +readonly DD_ATTEMPT_ORDINAL={quote_shell_value(attempt)} + +verify_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] +}} + +verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" +verify_sha256 "${{DD_PLAN_SHA256}}" "${{DD_PLAN}}" +if [[ ! ${{SLURM_ARRAY_TASK_ID:-}} =~ ^[0-9]+$ ]]; then + printf '%s\\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +readonly DD_ARRAY_TASK_ID="${{SLURM_ARRAY_TASK_ID}}" +printf -v DD_SHARD_ID 'shard-%05d' "${{DD_ARRAY_TASK_ID}}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${{DD_RUN_ROOT}}/shards/${{DD_SHARD_ID}}/attempts/attempt-${{DD_ATTEMPT_ORDINAL}}" +source <(tar -xOf "${{DD_RUNTIME_ARCHIVE}}" scratch.sh) +trap 'dd_cleanup_allocation_scratch || true' EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +dd_initialize_local_scratch +dd_stage_allocation_runtime +readonly DD_RUNTIME_ROOT="${{DD_SCRATCH_ROOT}}/runtime" + +source "${{DD_RUNTIME_ROOT}}/entrypoint.sh" +dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" +""" + + +def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) -> str: + """Render selected failed shards as one immutable retry array submission.""" + run_root = posixpath.dirname(plan.authored_config.path) + plan_path = posixpath.join(run_root, "resolved-plan.json") + if retry.run_id != plan.run_id: + raise SlurmBatchRenderError("retry run identity does not match the resolved plan") + if retry.resolved_plan.path != plan_path or retry.resolved_plan.sha256 != plan.compute_sha256(): + raise SlurmBatchRenderError("retry does not bind the persisted resolved plan") + try: + validate_enroot_mount_path(plan.selected_profile.profile.workspace_root) + except ValueError as error: + raise SlurmBatchRenderError("retry workspace cannot be represented as a safe Enroot mount") from error + planned_by_id = {shard.shard_id: shard for shard in plan.shards} + for retry_shard in retry.planned_shards: + planned = planned_by_id.get(retry_shard.shard_id) + if planned is None or planned.array_task_index != retry_shard.array_task_index: + raise SlurmBatchRenderError("retry shard does not match the resolved plan") + + array_tasks = ",".join(str(shard.array_task_index) for shard in retry.planned_shards) + if plan.array_tasks.max_concurrent is not None: + array_tasks = f"{array_tasks}%{plan.array_tasks.max_concurrent}" + retry_root = posixpath.join(run_root, "retries", retry.retry_id) + directives = render_batch_directives( + _build_generation_directives( + plan, + array=array_tasks, + job_name=retry.submission_job_name, + output_path=posixpath.join(retry_root, "slurm-%A_%a.out"), + error_path=posixpath.join(retry_root, "slurm-%A_%a.err"), + ) + ) + if _requires_exclusive_nodes(plan): + directives = f"{directives}\n#SBATCH --exclusive" + attempt_cases = "\n".join( + f" {shard.array_task_index}) DD_ATTEMPT_ORDINAL={quote_shell_value(f'{shard.attempt_ordinal:04d}')} ;;" + for shard in retry.planned_shards + ) + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH={quote_shell_value(_get_command_path(plan))} +export PYTHONNOUSERSITE=1 + +readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} +readonly DD_PLAN={quote_shell_value(plan_path)} +readonly DD_PLAN_SHA256={quote_shell_value(plan.compute_sha256())} +readonly DD_RUN_ROOT={quote_shell_value(run_root)} +readonly DD_WORKSPACE_ROOT={quote_shell_value(plan.selected_profile.profile.workspace_root)} +readonly DD_RUN_ID={quote_shell_value(plan.run_id)} +DD_CLIENT_PYTHON={quote_shell_value(plan.client.runtime.python_executable)} +readonly DD_RETRY_ID={quote_shell_value(retry.retry_id)} +readonly DD_RETRY_PLAN_SHA256={quote_shell_value(retry.compute_sha256())} +readonly DD_EFFECTIVE_RESUME_MODE={quote_shell_value(retry.effective_resume_mode)} + +verify_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] +}} + +verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" +verify_sha256 "${{DD_PLAN_SHA256}}" "${{DD_PLAN}}" +if [[ ! ${{SLURM_ARRAY_TASK_ID:-}} =~ ^[0-9]+$ ]]; then + printf '%s\\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +readonly DD_ARRAY_TASK_ID="${{SLURM_ARRAY_TASK_ID}}" +case "${{DD_ARRAY_TASK_ID}}" in +{attempt_cases} + *) printf '%s\\n' 'array task is absent from retry plan' >&2; exit 64 ;; +esac +readonly DD_ATTEMPT_ORDINAL +printf -v DD_SHARD_ID 'shard-%05d' "${{DD_ARRAY_TASK_ID}}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${{DD_RUN_ROOT}}/shards/${{DD_SHARD_ID}}/attempts/attempt-${{DD_ATTEMPT_ORDINAL}}" +readonly DD_ATTEMPT_MANIFEST="${{DD_ATTEMPT_DIR}}/attempt.json" +for ((DD_WAIT_COUNT = 0; DD_WAIT_COUNT < 300; DD_WAIT_COUNT++)); do + [[ -f "${{DD_ATTEMPT_MANIFEST}}" ]] && break + sleep 1 +done +if [[ ! -f "${{DD_ATTEMPT_MANIFEST}}" ]]; then + printf '%s\\n' 'retry attempt state was not published before allocation startup' >&2 + exit 70 +fi +if [[ ! ${{SLURM_ARRAY_JOB_ID:-}} =~ ^[1-9][0-9]*$ ]]; then + printf '%s\\n' 'SLURM_ARRAY_JOB_ID must be a positive integer' >&2 + exit 64 +fi +readonly DD_ARRAY_JOB_ID="${{SLURM_ARRAY_JOB_ID}}" +readonly DD_ATTEMPT_ID="attempt-${{DD_ATTEMPT_ORDINAL}}" +source <(tar -xOf "${{DD_RUNTIME_ARCHIVE}}" scratch.sh) +trap 'dd_cleanup_allocation_scratch || true' EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +dd_initialize_local_scratch +dd_stage_allocation_runtime +readonly DD_RUNTIME_ROOT="${{DD_SCRATCH_ROOT}}/runtime" +export PYTHONPATH="${{DD_RUNTIME_ROOT}}" +srun \\ + --nodes=1 \\ + --ntasks=1 \\ + --cpus-per-task=1 \\ + --exact \\ + --overlap \\ + --unbuffered \\ + --gres=none \\ + --export=ALL \\ + -- \\ + "${{DD_CLIENT_PYTHON}}" -m data_designer.slurm.state.attempt_identity \\ + --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" \\ + --shard-id "${{DD_SHARD_ID}}" --attempt-id "${{DD_ATTEMPT_ID}}" \\ + --array-job-id "${{DD_ARRAY_JOB_ID}}" --array-task-id "${{DD_ARRAY_TASK_ID}}" + +source "${{DD_RUNTIME_ROOT}}/entrypoint.sh" +dd_slurm_run_allocation \ + "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" "${{DD_RETRY_ID}}" \ + "${{DD_RETRY_PLAN_SHA256}}" "${{DD_EFFECTIVE_RESUME_MODE}}" +""" + + +def _get_command_path(plan: ResolvedSlurmRunPlan) -> str: + scheduler_bin_path = plan.scheduler_bin_path or plan.selected_profile.profile.scheduler.bin_path + return _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" + + +def _build_generation_directives( + plan: ResolvedSlurmRunPlan, + *, + array: str | None = None, + job_name: str | None = None, + output_path: str, + error_path: str, +) -> tuple[tuple[str, str | None], ...]: + node_indices = ( + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + ) + node_count = max(node_indices) + 1 + resolved_array = array or "0" + if array is None and plan.array_tasks.count > 1: + resolved_array = f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.max_concurrent is not None: + resolved_array = f"{resolved_array}%{plan.array_tasks.max_concurrent}" + + profile = plan.selected_profile.profile + if profile.gpu_request_mode == "gres": + # Reserve CPUs with every GPU so socket-local CPU binding cannot hide GPUs from a serving step. + cpus_per_gpu = max( + 1, (plan.client.authored.cpus + plan.resolved_gpus_per_node - 1) // plan.resolved_gpus_per_node + ) + cpu_directive = ("cpus-per-gpu", str(cpus_per_gpu)) + else: + if profile.scheduler.mem_per_gpu is not None: + raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") + cpu_directive = ("cpus-per-task", str(plan.client.authored.cpus)) + + values: list[tuple[str, str | None]] = [ + ("job-name", plan.submission.job_name if job_name is None else job_name), + ("account", plan.submission.account), + ("partition", plan.submission.partition), + ("nodes", str(node_count)), + cpu_directive, + ("time", plan.submission.time_limit), + ("array", resolved_array), + ("chdir", posixpath.dirname(plan.authored_config.path)), + ("output", output_path), + ("error", error_path), + ] + if profile.gpu_request_mode == "gres": + values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) + if profile.scheduler.mem_per_gpu is not None: + values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) + if plan.submission.comment is not None: + values.append(("comment", plan.submission.comment)) + return tuple(values) + + +def _requires_exclusive_nodes(plan: ResolvedSlurmRunPlan) -> bool: + """Avoid remote-node port collisions when batch GPU IDs describe node zero only.""" + if plan.selected_profile.profile.gpu_request_mode == "visible": + return True + return any(index > 0 for deployment in plan.deployments for index in deployment.node_indices) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py new file mode 100644 index 000000000..de0ceb1d5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal injectable process execution for Slurm command-line tools.""" + +from __future__ import annotations + +import math +import os +import subprocess +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Protocol + + +class CommandRunner(Protocol): + """Minimal command boundary implemented by production and fake runners.""" + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + """Execute one argument-vector command.""" + ... + + +class SubprocessRunner: + """Run commands without a shell or unrestricted ambient environment.""" + + _environment: Mapping[str, str] + _timeout_seconds: float + + def __init__( + self, + *, + environment: Mapping[str, str] | None = None, + timeout_seconds: float = 30.0, + ) -> None: + if type(timeout_seconds) not in {int, float} or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("timeout_seconds must be a finite positive number") + if environment is None: + explicit_environment = {"PATH": os.environ.get("PATH") or os.defpath} + if "SLURM_CONF" in os.environ: + explicit_environment["SLURM_CONF"] = os.environ["SLURM_CONF"] + else: + explicit_environment = dict(environment) + for name, value in explicit_environment.items(): + if type(name) is not str or not name or "=" in name or "\0" in name: + raise ValueError("environment names must be non-empty and must not contain '=' or NUL") + if type(value) is not str or "\0" in value: + raise ValueError("environment values must not contain NUL") + self._environment = MappingProxyType({**explicit_environment, "LC_ALL": "C"}) + self._timeout_seconds = float(timeout_seconds) + + @property + def environment(self) -> Mapping[str, str]: + """Return the allowlisted environment forwarded to child processes.""" + return self._environment + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + """Execute an argument vector with captured text output.""" + process_environment = dict(self._environment) + if environment is not None: + _validate_environment(environment) + process_environment.update(environment) + if input_text is not None: + return subprocess.run( + tuple(command), + check=False, + input=input_text, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + env=process_environment, + timeout=self._timeout_seconds, + ) + return subprocess.run( + tuple(command), + check=False, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + env=process_environment, + timeout=self._timeout_seconds, + ) + + +def _validate_environment(environment: Mapping[str, str]) -> None: + for name, value in environment.items(): + if type(name) is not str or not name or "=" in name or "\0" in name: + raise ValueError("environment names must be non-empty and must not contain '=' or NUL") + if type(value) is not str or "\0" in value: + raise ValueError("environment values must not contain NUL") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/scheduler_path.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/scheduler_path.py new file mode 100644 index 000000000..a365f0118 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/scheduler_path.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve the Slurm tools needed inside a generated allocation.""" + +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +from data_designer.slurm.contracts import validate_absolute_path + + +class SchedulerPathError(ValueError): + """The submit environment cannot provide the required allocation tools.""" + + +def resolve_scheduler_bin_path(configured: str | None) -> str: + """Find a single executable directory for the batch job's controlled PATH.""" + if configured is None: + srun = shutil.which("srun") + if srun is None: + raise SchedulerPathError( + "srun was not found on the submit host PATH; load your site's Slurm environment before submitting" + ) + directory = Path(os.path.abspath(srun)).parent + else: + directory = Path(configured) + try: + validate_absolute_path(str(directory)) + except ValueError: + raise SchedulerPathError("scheduler.bin_path must be one absolute directory") from None + if ":" in str(directory): + raise SchedulerPathError("scheduler.bin_path must be one absolute directory") + if not directory.is_dir() or any( + not (directory / executable).is_file() or not os.access(directory / executable, os.X_OK) + for executable in ("srun", "scontrol") + ): + raise SchedulerPathError( + "the Slurm executable directory must contain runnable srun and scontrol; " + "load your site's Slurm environment or set scheduler.bin_path to their shared directory" + ) + return str(directory) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py new file mode 100644 index 000000000..d0f837c41 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolved execution-plan contracts for Data Designer Slurm.""" + +from __future__ import annotations + +from data_designer.slurm.planning.models import ( + ArtifactReference, + LockedPackage, + PlannedShard, + PortClaim, + RecordRange, + ResolvedBuilderInput, + ResolvedClient, + ResolvedClientRuntime, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedImage, + ResolvedInvocation, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, + ResumeWorkspace, +) + +__all__ = [ + "ArtifactReference", + "LockedPackage", + "PlannedShard", + "PortClaim", + "RecordRange", + "ResolvedBuilderInput", + "ResolvedClient", + "ResolvedClientRuntime", + "ResolvedDependencyLock", + "ResolvedDeployment", + "ResolvedImage", + "ResolvedInvocation", + "ResolvedSlurmRunPlan", + "ResolvedSubmission", + "ResolvedTopology", + "ResumeWorkspace", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py new file mode 100644 index 000000000..21f59e381 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Submission-safe identity extraction for Data Designer builder payloads.""" + +from __future__ import annotations + +from pydantic import JsonValue + +from data_designer.slurm.contracts import ModelAlias, Sha256Digest, compute_serialized_json_sha256 + + +def get_declared_model_aliases(builder: dict[str, JsonValue]) -> tuple[ModelAlias, ...]: + """Return aliases declared by the known ``model_configs`` envelope.""" + data_designer = builder.get("data_designer", builder) + if not isinstance(data_designer, dict): + raise ValueError("builder data_designer value must be an object") + model_configs = data_designer.get("model_configs") or [] + if not isinstance(model_configs, list): + raise ValueError("builder model_configs must be a list") + + model_aliases: list[ModelAlias] = [] + for model_config in model_configs: + if not isinstance(model_config, dict) or not isinstance(model_config.get("alias"), str): + raise ValueError("each builder model config must contain a string alias") + model_aliases.append(model_config["alias"]) + return tuple(model_aliases) + + +def get_persisted_builder_identity( + builder: dict[str, JsonValue], +) -> tuple[tuple[ModelAlias, ...], Sha256Digest]: + """Return declared aliases and the digest of persisted builder JSON.""" + return get_declared_model_aliases(builder), compute_serialized_json_sha256(builder) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py new file mode 100644 index 000000000..f107a9e2e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal pure deterministic Slurm plan compilation.""" + +from __future__ import annotations + +import posixpath + +from pydantic import ValidationError + +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.contracts import ( + ArtifactReference, + RecordRange, + ResumeWorkspace, + compute_serialized_json_sha256, +) +from data_designer.slurm.planning.errors import SlurmPlanCompilationError, SlurmPlanContractError +from data_designer.slurm.planning.models import ( + PlannedShard, + PortClaim, + ResolvedClient, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedTopology, + derive_deployment_placement, +) +from data_designer.slurm.planning.resolution import ( + EffectiveDataDesignerSlurmConfig, + validate_effective_slurm_config, +) +from data_designer.slurm.planning.validation import validate_resolved_plan + +__all__: list[str] = [] + +_LOGICAL_ENDPOINT_PORT = 17000 +_HTTP_PORT = 18000 +_RENDEZVOUS_PORT = 19000 +_PORT_RANGE_SIZE = 1000 +_COMPILER_VALIDATION_MODELS = ( + ArtifactReference, + EffectiveDataDesignerSlurmConfig, + PlannedShard, + PortClaim, + RecordRange, + ResolvedClient, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedTopology, + ResumeWorkspace, +) + + +class SlurmRunCompiler: + """Compile one fully resolved configuration without ambient state or I/O.""" + + @staticmethod + def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: + """Return one immutable deterministic execution plan.""" + try: + effective = EffectiveDataDesignerSlurmConfig.model_validate(effective.model_dump(mode="python")) + effective = validate_effective_slurm_config(effective) + deployments = _compile_deployments(effective) + client = _compile_client(effective, deployments) + _validate_port_claims(effective, client, deployments) + plan = ResolvedSlurmRunPlan( + schema_version=1, + run_id=effective.run_id, + package_version=effective.package_version, + authored_config=ArtifactReference( + path=posixpath.join(_run_root(effective), "authored-config.json"), + sha256=effective.authored.compute_sha256(), + ), + selected_profile=effective.selected_profile, + scheduler_bin_path=effective.scheduler_bin_path, + resolved_gpus_per_node=effective.resolved_gpus_per_node, + builder=effective.builder, + invocation=effective.invocation, + client=client, + deployments=deployments, + array_tasks=effective.authored.array_tasks, + shards=_compile_shards(effective), + submission=effective.submission, + container_mounts=tuple(effective.selected_profile.profile.container_mounts), + runtime_bundle=effective.runtime_bundle, + ) + return validate_resolved_plan( + effective.authored, + effective.dependency_lock, + plan, + builder_payload=effective.builder_payload, + ) + except (SlurmPlanCompilationError, SlurmPlanContractError): + raise + except ValidationError as error: + message = format_validation_error( + error, + subject="Slurm plan compilation", + models=_COMPILER_VALIDATION_MODELS, + ) + raise SlurmPlanCompilationError(message) from None + except ValueError as error: + raise SlurmPlanCompilationError(str(error)) from None + + +def _compile_deployments( + effective: EffectiveDataDesignerSlurmConfig, +) -> tuple[ResolvedDeployment, ...]: + if len(effective.authored.deployments) > _PORT_RANGE_SIZE: + raise SlurmPlanCompilationError("deployment count exceeds the compiler-owned logical endpoint port range") + resolved: list[ResolvedDeployment] = [] + placements = derive_deployment_placement(effective.authored.deployments, effective.resolved_gpus_per_node) + http_offsets: dict[int, int] = {} + for index, (authored, image) in enumerate( + zip(effective.authored.deployments, effective.deployment_images, strict=True) + ): + node_indices, gpu_start_index, quota = placements[index] + topology = ResolvedTopology.derive( + node_count=authored.resources.nodes, + gpus_per_node=quota, + tensor_parallel=authored.topology.tensor_parallel, + nodes_per_replica=authored.topology.nodes_per_replica, + ) + if topology.replicas_per_node_group > _PORT_RANGE_SIZE: + raise SlurmPlanCompilationError("replica lanes exceed the compiler-owned deployment port range") + deployment_id = f"deployment-{index:05d}" + http_offset = http_offsets.get(node_indices[0], 0) + if http_offset + topology.replicas_per_node_group > _PORT_RANGE_SIZE: + raise SlurmPlanCompilationError("replica lanes exceed the compiler-owned node port range") + ports = _compile_deployment_ports(deployment_id, node_indices, topology, http_offset=http_offset) + for node in node_indices: + http_offsets[node] = http_offset + topology.replicas_per_node_group + resolved.append( + ResolvedDeployment( + deployment_id=deployment_id, + authored=authored, + served_model_name=authored.served_model_name or authored.model, + image=image, + node_indices=node_indices, + gpus_per_node=quota, + gpu_start_index=gpu_start_index, + topology=topology, + ports=ports, + ) + ) + return tuple(resolved) + + +def _compile_deployment_ports( + deployment_id: str, + node_indices: tuple[int, ...], + topology: ResolvedTopology, + *, + http_offset: int = 0, +) -> tuple[PortClaim, ...]: + http: list[PortClaim] = [] + rendezvous: list[PortClaim] = [] + for group_index in range(topology.node_group_count): + head = node_indices[group_index * topology.nodes_per_replica] + for lane_index in range(topology.replicas_per_node_group): + replica_index = group_index * topology.replicas_per_node_group + lane_index + http.append( + PortClaim( + name=f"{deployment_id}-http-{replica_index:05d}", + role="http", + node_index=head, + port=_HTTP_PORT + http_offset + lane_index, + ) + ) + if topology.nodes_per_replica > 1: + rendezvous.append( + PortClaim( + name=f"{deployment_id}-rendezvous-{replica_index:05d}", + role="rendezvous", + node_index=head, + port=_RENDEZVOUS_PORT + lane_index, + ) + ) + return tuple(http + rendezvous) + + +def _compile_client( + effective: EffectiveDataDesignerSlurmConfig, + deployments: tuple[ResolvedDeployment, ...], +) -> ResolvedClient: + host_node_index = deployments[0].node_indices[0] + ports = tuple( + PortClaim( + name=f"{deployment.deployment_id}-logical-endpoint", + role="logical_endpoint", + node_index=host_node_index, + port=_LOGICAL_ENDPOINT_PORT + index, + ) + for index, deployment in enumerate(deployments) + ) + return ResolvedClient( + authored=effective.authored.client, + runtime=effective.client_runtime, + dependency_lock=ArtifactReference( + path=posixpath.join(_run_root(effective), "dependency-lock.json"), + sha256=effective.dependency_lock.compute_sha256(), + ), + host_node_index=host_node_index, + gpu_count=0, + ports=ports, + ) + + +def _validate_port_claims( + effective: EffectiveDataDesignerSlurmConfig, + client: ResolvedClient, + deployments: tuple[ResolvedDeployment, ...], +) -> None: + ports = client.ports + tuple(port for deployment in deployments for port in deployment.ports) + addresses = tuple((port.node_index, port.port) for port in ports) + if len(addresses) != len(set(addresses)): + raise SlurmPlanCompilationError("compiler-owned port claims collide on one node") + otel_port = effective.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and (client.host_node_index, otel_port) in addresses: + raise SlurmPlanCompilationError("client OTEL metrics port collides with a compiler-owned port") + + +def _compile_shards(effective: EffectiveDataDesignerSlurmConfig) -> tuple[PlannedShard, ...]: + count = effective.authored.array_tasks.count + requested = effective.authored.invocation.num_records + floor_count = requested // count + start = 0 + shards: list[PlannedShard] = [] + for index in range(count): + record_count = requested - floor_count * (count - 1) if index == count - 1 else floor_count + end = start + record_count + shard_id = f"shard-{index:05d}" + shard_root = posixpath.join(_run_root(effective), "shards", shard_id) + record_range = RecordRange(start_index=start, end_index_exclusive=end) + partition = None + seed_path = effective.invocation.effective_input_bindings.seed_path + if seed_path is not None: + partition = ArtifactReference( + path=posixpath.join(shard_root, "input-partition.json"), + sha256=compute_serialized_json_sha256( + { + "record_range": record_range.model_dump(mode="json"), + "seed_path": seed_path, + } + ), + ) + shards.append( + PlannedShard( + shard_id=shard_id, + shard_index=index, + array_task_index=index, + record_range=record_range, + input_partition=partition, + resume_workspace=ResumeWorkspace(path=posixpath.join(shard_root, "dataset")), + ) + ) + start = end + return tuple(shards) + + +def _run_root(effective: EffectiveDataDesignerSlurmConfig) -> str: + return posixpath.join(effective.selected_profile.profile.workspace_root, "runs", effective.run_id) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py new file mode 100644 index 000000000..63a6471a3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalized Slurm planning errors.""" + +from __future__ import annotations + + +class SlurmPlanningError(ValueError): + """Base error for Slurm configuration resolution and plan validation.""" + + +class SlurmConfigResolutionError(SlurmPlanningError): + """Raised when resolved inputs do not match one authored declaration.""" + + +class SlurmPlanCompilationError(SlurmPlanningError): + """Raised when one effective configuration cannot produce a valid plan.""" + + +class SlurmPlanContractError(SlurmPlanningError): + """Raised when a resolved plan does not match its authored inputs.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py new file mode 100644 index 000000000..10ee11ee4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -0,0 +1,602 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath +from collections.abc import Sequence +from typing import Annotated, Literal +from urllib.parse import urlsplit + +from packaging.requirements import Requirement +from packaging.utils import canonicalize_name +from packaging.version import InvalidVersion +from pydantic import Field, JsonValue, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.config import RunConfig +from data_designer.slurm.config.environment import validate_no_plaintext_secrets +from data_designer.slurm.config.images import ImageInspection, ImageInspectionRecord, ImageKind, ImageRef +from data_designer.slurm.config.profiles import ContainerMount, SelectedSlurmProfile +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + ClientConfig, + ClientDependencies, + InputBindings, + InvocationConfig, + ServerDeploymentConfig, + SubmissionConfig, +) +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + DistributionName, + Identifier, + InstalledDistribution, + ModelAlias, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, + compute_canonical_json_sha256, + compute_serialized_json_sha256, + derive_managed_assets_path, + is_path_below, + paths_overlap, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, +) +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases +from data_designer.slurm.types import NetworkPort, PartitionSelection + + +class ResolvedImage(ContractValue): + """Immutable SQSH path and the digest-bound inspection that approved it.""" + + authored_ref: ImageRef + path: str + sha256: Sha256Digest + inspection: ImageInspectionRecord + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + validate_absolute_path(value) + if not value.endswith(".sqsh"): + raise ValueError("resolved image path must end in .sqsh") + return value + + @model_validator(mode="after") + def validate_image(self) -> ResolvedImage: + if self.inspection.sqsh_sha256 != self.sha256: + raise ValueError("image inspection digest does not match the resolved SQSH") + if self.authored_ref.path is not None and self.authored_ref.path != self.path: + raise ValueError("resolved image path does not match the authored path") + return self + + @property + def kind(self) -> ImageKind: + return self.inspection.inspection.kind + + @property + def inspection_facts(self) -> ImageInspection: + """Return factual image inspection data without exposing record nesting to consumers.""" + return self.inspection.inspection + + +class ResolvedClientRuntime(ContractValue): + """Exact native Python runtime used by allocation client processes.""" + + python_executable: str + python_implementation: Identifier + python_version: Annotated[str, StringConstraints(pattern=r"^[0-9]+\.[0-9]+\.[0-9]+$")] + python_abi: Identifier + distributions: tuple[InstalledDistribution, ...] + runtime_sha256: Sha256Digest + + _python_executable_is_absolute = field_validator("python_executable")(validate_absolute_path) + + @model_validator(mode="after") + def validate_runtime(self) -> ResolvedClientRuntime: + names = tuple(distribution.name for distribution in self.distributions) + if names != tuple(sorted(names)) or len(names) != len(set(names)): + raise ValueError("client runtime distributions must be sorted and unique") + identity = self.model_dump(mode="json", exclude={"runtime_sha256"}) + if self.runtime_sha256 != compute_canonical_json_sha256(identity): + raise ValueError("client runtime fingerprint does not match its facts") + return self + + +class LockedPackage(ContractValue): + name: DistributionName + version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + artifact: ArtifactReference + + @model_validator(mode="after") + def validate_artifact(self) -> LockedPackage: + if not self.artifact.path.endswith(".whl"): + raise ValueError("locked overlay artifacts must be wheels") + return self + + +class ResolvedDependencyLock(ContractRecord): + """Immutable client dependency resolution against one fixed Python runtime.""" + + resolver_version: Identifier + python_abi: Identifier + client_runtime_sha256: Sha256Digest + authored_requirements: tuple[str, ...] + authored_source: str | None = None + source: ArtifactReference | None = None + base_distributions: tuple[InstalledDistribution, ...] + overlay_packages: tuple[LockedPackage, ...] + + @field_validator("authored_source") + @classmethod + def validate_authored_source(cls, value: str | None) -> str | None: + if value is None: + return None + normalized = validate_local_config_path(value) + if not normalized.endswith(".json"): + raise ValueError("dependency lock source must end in .json") + return normalized + + @model_validator(mode="after") + def validate_packages(self) -> ResolvedDependencyLock: + if (self.authored_source is None) != (self.source is None): + raise ValueError("dependency lock authored and resolved sources must be provided together") + ClientDependencies(requirements=list(self.authored_requirements)) + base_names = tuple(distribution.name for distribution in self.base_distributions) + overlay_names = tuple(package.name for package in self.overlay_packages) + if base_names != tuple(sorted(base_names)) or overlay_names != tuple(sorted(overlay_names)): + raise ValueError("dependency lock distributions must be sorted by normalized name") + if len(base_names) != len(set(base_names)) or len(overlay_names) != len(set(overlay_names)): + raise ValueError("dependency lock distribution names must be unique") + overlap = set(base_names).intersection(overlay_names) + if overlap: + raise ValueError(f"overlay packages overlap base runtime distributions: {', '.join(sorted(overlap))}") + base_packages = {distribution.name: distribution for distribution in self.base_distributions} + overlay_packages = {package.name: package for package in self.overlay_packages} + for value in self.authored_requirements: + requirement = Requirement(value) + name = canonicalize_name(requirement.name) + if requirement.url is not None: + package = overlay_packages.get(name) + digest = urlsplit(requirement.url).fragment.removeprefix("sha256=") + if package is None or package.artifact.sha256 != digest: + raise ValueError(f"direct requirement {name!r} must match one locked overlay artifact") + continue + package = overlay_packages.get(name) or base_packages.get(name) + if package is None: + raise ValueError(f"authored requirement {name!r} is missing from the dependency lock") + try: + satisfied = requirement.specifier.contains(package.version, prereleases=True) + except InvalidVersion as error: + raise ValueError(f"locked package {name!r} has an invalid version") from error + if not satisfied: + raise ValueError(f"locked package {name!r} does not satisfy its authored requirement") + return self + + +class ResolvedBuilderInput(ContractValue): + authored_source: str | None = None + source: ArtifactReference | None = None + inline: dict[str, JsonValue] | None = None + # Both forms use deterministic persisted builder JSON bytes. + content_sha256: Sha256Digest + model_aliases: tuple[ModelAlias, ...] + + @model_validator(mode="after") + def validate_input(self) -> ResolvedBuilderInput: + if (self.source is None) == (self.inline is None): + raise ValueError("resolved builder requires exactly one of source or inline") + if self.source is None: + if self.authored_source is not None: + raise ValueError("inline builder input cannot contain authored_source") + validate_no_plaintext_secrets(self.inline, field_name="resolved inline builder input") + expected_digest = compute_serialized_json_sha256(self.inline) + if self.model_aliases != get_declared_model_aliases(self.inline): + raise ValueError("resolved model aliases do not match the inline builder") + else: + if self.authored_source is None: + raise ValueError("resolved builder source requires authored_source") + expected_digest = self.source.sha256 + if len(self.model_aliases) != len(set(self.model_aliases)): + raise ValueError("resolved builder model aliases must be unique") + if self.content_sha256 != expected_digest: + raise ValueError("builder content digest does not match the resolved input") + return self + + +class ResolvedInvocation(ContractValue): + authored: InvocationConfig + effective_input_bindings: InputBindings + effective_run_config: dict[str, JsonValue] + + @field_validator("effective_run_config") + @classmethod + def validate_run_config_is_materialized(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + materialized = RunConfig.model_validate(value).model_dump(mode="json") + if value != materialized: + raise ValueError("effective_run_config must contain the fully materialized Data Designer RunConfig") + return value + + +class PortClaim(ContractValue): + name: Identifier + role: Literal["http", "rendezvous", "logical_endpoint"] + # TODO: Rename node_index to allocation_node_index after downstream Stage 2 branches converge on this contract. + node_index: NonNegativeInt + port: NetworkPort + + +class ResolvedTopology(ContractValue): + """Resource-derived serving topology shared by planning and serving.""" + + tensor_parallel: PositiveInt + nodes_per_replica: PositiveInt + pipeline_parallel: PositiveInt + node_group_count: PositiveInt + replicas_per_node_group: PositiveInt + replica_count: PositiveInt + gpus_per_replica: PositiveInt + + @classmethod + def derive( + cls, + *, + node_count: int, + gpus_per_node: int, + tensor_parallel: int, + nodes_per_replica: int, + ) -> ResolvedTopology: + """Derive the canonical v1 topology from placement and GPU resources.""" + if min(node_count, gpus_per_node, tensor_parallel, nodes_per_replica) <= 0: + raise ValueError("topology derivation inputs must be positive") + if node_count % nodes_per_replica: + raise ValueError("nodes_per_replica must divide the deployment node count") + if gpus_per_node % tensor_parallel: + raise ValueError("tensor_parallel must divide resolved GPUs per node") + node_group_count = node_count // nodes_per_replica + replicas_per_node_group = gpus_per_node // tensor_parallel + return cls( + tensor_parallel=tensor_parallel, + nodes_per_replica=nodes_per_replica, + pipeline_parallel=nodes_per_replica, + node_group_count=node_group_count, + replicas_per_node_group=replicas_per_node_group, + replica_count=node_group_count * replicas_per_node_group, + gpus_per_replica=tensor_parallel * nodes_per_replica, + ) + + +class ResolvedDeployment(ContractValue): + deployment_id: Identifier + authored: ServerDeploymentConfig + served_model_name: str + image: ResolvedImage + node_indices: tuple[NonNegativeInt, ...] = Field(min_length=1) + gpus_per_node: PositiveInt + gpu_start_index: NonNegativeInt = Field(default=0, exclude_if=lambda value: value == 0) + topology: ResolvedTopology + ports: tuple[PortClaim, ...] = () + + @model_validator(mode="after") + def validate_deployment(self) -> ResolvedDeployment: + validate_plain_text(self.served_model_name, field_name="served model name") + if self.served_model_name != (self.authored.served_model_name or self.authored.model): + raise ValueError("resolved served model name does not match the authored deployment") + if self.image.kind is not ImageKind.SERVING: + raise ValueError("server deployments require serving images") + if self.image.authored_ref != self.authored.server.image: + raise ValueError("resolved serving image does not match the authored image reference") + if len(self.node_indices) != self.authored.resources.nodes: + raise ValueError("deployment placement must contain exactly the requested node count") + if self.node_indices != tuple(sorted(set(self.node_indices))): + raise ValueError("deployment node indices must be sorted and unique") + expected = ResolvedTopology.derive( + node_count=len(self.node_indices), + gpus_per_node=self.gpus_per_node, + tensor_parallel=self.authored.topology.tensor_parallel, + nodes_per_replica=self.authored.topology.nodes_per_replica, + ) + if self.topology != expected: + raise ValueError("resolved topology does not match deployment resources") + if any(port.node_index not in self.node_indices for port in self.ports): + raise ValueError("deployment port claims must use deployment nodes") + names = tuple(port.name for port in self.ports) + if len(names) != len(set(names)): + raise ValueError("deployment port claim names must be unique") + if any(not name.startswith(f"{self.deployment_id}-") for name in names): + raise ValueError("deployment port claim names must use the deployment ID") + if any(port.role == "logical_endpoint" for port in self.ports): + raise ValueError("logical endpoint ports belong to the resolved client") + + group_heads = self.node_indices[:: self.topology.nodes_per_replica] + expected_http_nodes = tuple(head for head in group_heads for _ in range(self.topology.replicas_per_node_group)) + http_ports = tuple(port for port in self.ports if port.role == "http") + http_nodes = tuple(port.node_index for port in http_ports) + if http_nodes != expected_http_nodes: + raise ValueError("deployment requires one ordered HTTP port claim per replica lane") + expected_http_names = tuple(f"{self.deployment_id}-http-{index:05d}" for index in range(len(http_ports))) + if tuple(port.name for port in http_ports) != expected_http_names: + raise ValueError("deployment HTTP port names must match their ordered replica lane") + + expected_rendezvous_nodes = ( + tuple(head for head in group_heads for _ in range(self.topology.replicas_per_node_group)) + if self.topology.nodes_per_replica > 1 + else () + ) + rendezvous_ports = tuple(port for port in self.ports if port.role == "rendezvous") + rendezvous_nodes = tuple(port.node_index for port in rendezvous_ports) + if rendezvous_nodes != expected_rendezvous_nodes: + raise ValueError("deployment requires one ordered rendezvous port claim per multi-node replica lane") + expected_rendezvous_names = tuple( + f"{self.deployment_id}-rendezvous-{index:05d}" for index in range(len(rendezvous_ports)) + ) + if tuple(port.name for port in rendezvous_ports) != expected_rendezvous_names: + raise ValueError("deployment rendezvous port names must match their ordered replica lane") + return self + + +class ResolvedClient(ContractValue): + authored: ClientConfig + runtime: ResolvedClientRuntime + dependency_lock: ArtifactReference + host_node_index: NonNegativeInt + gpu_count: Literal[0] + ports: tuple[PortClaim, ...] = () + + @model_validator(mode="after") + def validate_client(self) -> ResolvedClient: + if any(port.role != "logical_endpoint" for port in self.ports): + raise ValueError("resolved client ports must be logical endpoints") + if any(port.node_index != self.host_node_index for port in self.ports): + raise ValueError("logical endpoint ports must use the client host") + names = tuple(port.name for port in self.ports) + if len(names) != len(set(names)): + raise ValueError("logical endpoint port claim names must be unique") + return self + + +class PlannedShard(ContractValue): + shard_id: ShardId + shard_index: NonNegativeInt + array_task_index: NonNegativeInt + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace: ResumeWorkspace + + @property + def requested_records(self) -> int: + return self.record_range.record_count + + +class ResolvedSubmission(ContractValue): + account: Identifier | None = None + partition: PartitionSelection | None = None + job_name: Identifier + time_limit: str + comment: str | None = None + + @model_validator(mode="after") + def validate_submission(self) -> ResolvedSubmission: + SubmissionConfig.model_validate(self.model_dump(mode="python")) + return self + + +class ResolvedSlurmRunPlan(ContractRecord): + """Immutable allocation input consumed without ambient configuration.""" + + run_id: Identifier + package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + authored_config: ArtifactReference + selected_profile: SelectedSlurmProfile + scheduler_bin_path: str | None = Field(default=None, exclude_if=lambda value: value is None) + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + invocation: ResolvedInvocation + client: ResolvedClient + deployments: tuple[ResolvedDeployment, ...] = Field(min_length=1) + array_tasks: ArrayTasksConfig + shards: tuple[PlannedShard, ...] = Field(min_length=1) + submission: ResolvedSubmission + container_mounts: tuple[ContainerMount, ...] = () + runtime_bundle: ArtifactReference + + @field_validator("scheduler_bin_path") + @classmethod + def validate_scheduler_bin_path(cls, value: str | None) -> str | None: + if value is None: + return None + validate_absolute_path(value) + if ":" in value: + raise ValueError("scheduler bin path must name one directory") + return value + + @model_validator(mode="after") + def validate_plan(self) -> ResolvedSlurmRunPlan: + profile = self.selected_profile.profile + if profile.scheduler.bin_path is not None and self.scheduler_bin_path not in (None, profile.scheduler.bin_path): + raise ValueError("resolved scheduler bin path does not match the selected profile override") + if profile.gpus_per_node != "auto" and profile.gpus_per_node != self.resolved_gpus_per_node: + raise ValueError("resolved GPU count does not match the selected profile") + if profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: + raise ValueError("mem_per_gpu requires GRES GPU request mode") + placements = derive_deployment_placement( + tuple(deployment.authored for deployment in self.deployments), self.resolved_gpus_per_node + ) + for deployment, (nodes, gpu_start, quota) in zip(self.deployments, placements, strict=True): + if (deployment.node_indices, deployment.gpu_start_index, deployment.gpus_per_node) != ( + nodes, + gpu_start, + quota, + ): + raise ValueError("deployment placement and GPU ownership must match authored resources") + shared: dict[str, ResolvedDeployment] = {} + for deployment in self.deployments: + group = deployment.authored.resources.shared_node + if group is None: + continue + if profile.gpu_request_mode != "gres": + raise ValueError("shared_node requires GRES GPU request mode") + first = shared.setdefault(group, deployment) + if (deployment.image.path, deployment.image.sha256, deployment.image.inspection_facts) != ( + first.image.path, + first.image.sha256, + first.image.inspection_facts, + ): + raise ValueError("shared_node requires the same resolved serving image") + if ( + deployment.authored.server.environment != first.authored.server.environment + or deployment.authored.server.queue_backpressure != first.authored.server.queue_backpressure + ): + raise ValueError("shared_node requires matching server environment and queue backpressure") + if tuple(profile.container_mounts) != self.container_mounts: + raise ValueError("plan mount mappings must match the selected profile") + + deployment_ids = tuple(deployment.deployment_id for deployment in self.deployments) + aliases = tuple(deployment.authored.model_alias for deployment in self.deployments) + expected_deployment_ids = tuple(f"deployment-{index:05d}" for index in range(len(self.deployments))) + if deployment_ids != expected_deployment_ids: + raise ValueError("resolved deployment IDs must use complete ordered zero-based identities") + if len(aliases) != len(set(aliases)): + raise ValueError("resolved deployment aliases must be unique") + if set(aliases) != set(self.builder.model_aliases): + raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") + + if self.client.host_node_index != self.deployments[0].node_indices[0]: + raise ValueError("client must be colocated on the first node of the first deployment") + authored_bindings = self.invocation.authored.input_bindings + expected_bindings = InputBindings( + seed_path=authored_bindings.seed_path, + managed_assets_path=authored_bindings.managed_assets_path + or derive_managed_assets_path(profile.workspace_root), + ) + if self.invocation.effective_input_bindings != expected_bindings: + raise ValueError("effective input bindings must match the authored and selected profile input") + managed_assets_path = self.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + workspace_state = tuple( + posixpath.join(profile.workspace_root, name) for name in ("images", "runtime", "benchmarks", "runs") + ) + if any(paths_overlap(managed_assets_path, path) for path in workspace_state): + raise ValueError("managed_assets_path must not overlap package-managed workspace state") + if "non_inference_max_parallel_workers" not in self.invocation.authored.run_config: + workers = self.invocation.effective_run_config["non_inference_max_parallel_workers"] + if workers != 4: + raise ValueError("default non-inference worker count must match the Data Designer default") + + expected_logical_names = tuple( + f"{deployment.deployment_id}-logical-endpoint" for deployment in self.deployments + ) + logical_names = tuple(port.name for port in self.client.ports) + if logical_names != expected_logical_names: + raise ValueError("client requires one ordered logical endpoint port per deployment") + + ports = self.client.ports + tuple(port for deployment in self.deployments for port in deployment.ports) + port_keys = tuple((port.node_index, port.port) for port in ports) + if len(port_keys) != len(set(port_keys)): + raise ValueError("plan port claims must be unique per node") + port_names = tuple(port.name for port in ports) + if len(port_names) != len(set(port_names)): + raise ValueError("plan port claim names must be unique") + otel_port = self.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and any( + port.node_index == self.client.host_node_index and port.port == otel_port for port in ports + ): + raise ValueError("client OTEL metrics port collides with a plan port claim") + + run_root = posixpath.join(profile.workspace_root, "runs", self.run_id) + if self.authored_config.path != posixpath.join(run_root, "authored-config.json"): + raise ValueError("authored config reference must use the plan run root") + if self.client.dependency_lock.path != posixpath.join(run_root, "dependency-lock.json"): + raise ValueError("dependency lock reference must use the plan run root") + runtime_root = posixpath.join(profile.workspace_root, "runtime") + runtime_name = posixpath.basename(self.runtime_bundle.path) + if ( + not is_path_below(self.runtime_bundle.path, runtime_root) + or runtime_name != f"{self.runtime_bundle.sha256}.tar.gz" + ): + raise ValueError("runtime bundle must be a content-addressed tar archive below the workspace runtime root") + self._validate_shards(run_root) + return self + + def _validate_shards(self, run_root: str) -> None: + if len(self.shards) != self.array_tasks.count: + raise ValueError("plan must contain exactly one shard per array task") + requested_records = self.invocation.authored.num_records + floor_count = requested_records // self.array_tasks.count + expected_start = 0 + shard_ids: list[ShardId] = [] + workspace_paths: list[str] = [] + partition_paths: list[str] = [] + requires_partition = self.invocation.effective_input_bindings.seed_path is not None + for index, shard in enumerate(self.shards): + if shard.shard_index != index or shard.array_task_index != index: + raise ValueError("shards must use complete ordered zero-based identities") + if shard.shard_id != f"shard-{index:05d}": + raise ValueError("shard IDs must match their zero-based shard index") + if shard.record_range.start_index != expected_start: + raise ValueError("shard record ranges must be contiguous") + expected_count = ( + requested_records - floor_count * (self.array_tasks.count - 1) + if index == self.array_tasks.count - 1 + else floor_count + ) + if shard.requested_records != expected_count: + raise ValueError("shards must use deterministic floor/remainder record counts") + expected_workspace = posixpath.join(run_root, "shards", shard.shard_id, "dataset") + if shard.resume_workspace.path != expected_workspace: + raise ValueError("shard resume workspace must match the run and shard identity") + if (shard.input_partition is not None) != requires_partition: + raise ValueError("shard input partition presence must match the authored seed input") + if shard.input_partition is not None: + expected_partition = posixpath.join(run_root, "shards", shard.shard_id, "input-partition.json") + if shard.input_partition.path != expected_partition: + raise ValueError("shard input partition must match the run and shard identity") + partition_paths.append(shard.input_partition.path) + shard_ids.append(shard.shard_id) + workspace_paths.append(shard.resume_workspace.path) + expected_start = shard.record_range.end_index_exclusive + if expected_start != requested_records: + raise ValueError("shard record ranges must cover the requested records") + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("shard IDs must be unique") + if len(workspace_paths) != len(set(workspace_paths)): + raise ValueError("shard resume workspaces must be unique") + if len(partition_paths) != len(set(partition_paths)): + raise ValueError("shard input partitions must be unique") + + +def derive_deployment_placement( + deployments: Sequence[ServerDeploymentConfig], gpus_per_node: int +) -> tuple[tuple[tuple[int, ...], int, int], ...]: + """Assign nodes and GPU slices, allowing explicit shared offsets.""" + groups: dict[str, tuple[int, int, int]] = {} + placements: list[tuple[tuple[int, ...], int, int]] = [] + next_node = 0 + shared_nodes = {deployment.resources.shared_node for deployment in deployments} + if shared_nodes != {None} and (len(shared_nodes) != 1 or None in shared_nodes): + raise ValueError("shared_node requires one group covering all deployments") + for deployment in deployments: + resources = deployment.resources + if resources.shared_node is None: + placements.append((tuple(range(next_node, next_node + resources.nodes)), 0, gpus_per_node)) + next_node += resources.nodes + continue + quota = resources.gpus_per_node + assert quota is not None + if resources.shared_node not in groups: + groups[resources.shared_node] = (next_node, 0, 0) + next_node += 1 + node, next_gpu, members = groups[resources.shared_node] + gpu_start = next_gpu if resources.gpu_start_index is None else resources.gpu_start_index + if gpu_start + quota > gpus_per_node: + raise ValueError("shared_node GPU quotas exceed the profile GPU budget") + placements.append(((node,), gpu_start, quota)) + groups[resources.shared_node] = (node, max(next_gpu, gpu_start + quota), members + 1) + if any(members < 2 for _, _, members in groups.values()): + raise ValueError("shared_node requires at least two deployments in each group") + return tuple(placements) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py new file mode 100644 index 000000000..5f735142d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -0,0 +1,385 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal pure authored-configuration resolution for Slurm planning.""" + +from __future__ import annotations + +import posixpath +from typing import Annotated + +from pydantic import JsonValue, PositiveInt, StringConstraints, TypeAdapter, ValidationError + +from data_designer.config import RunConfig +from data_designer.slurm._errors import format_validation_error +from data_designer.slurm.config.images import ImageKind +from data_designer.slurm.config.profiles import SelectedSlurmProfile +from data_designer.slurm.config.run import BuilderInput, DataDesignerSlurmConfig, InputBindings +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractValue, + Identifier, + compute_serialized_json_sha256, + derive_managed_assets_path, + is_path_below, + paths_overlap, +) +from data_designer.slurm.planning.builder_identity import ( + get_declared_model_aliases, + get_persisted_builder_identity, +) +from data_designer.slurm.planning.errors import SlurmConfigResolutionError +from data_designer.slurm.planning.models import ( + ResolvedBuilderInput, + ResolvedClientRuntime, + ResolvedDependencyLock, + ResolvedImage, + ResolvedInvocation, + ResolvedSubmission, + derive_deployment_placement, +) + +__all__: list[str] = [] + +_SHARDABLE_COLUMN_TYPES = frozenset( + { + "embedding", + "expression", + "llm-code", + "llm-judge", + "llm-structured", + "llm-text", + "sampler", + "seed-dataset", + "validation", + } +) +_COMPATIBILITY_RUN_DEFAULTS: dict[str, JsonValue] = { + "buffer_size": 16384, + "disable_early_shutdown": True, + "display_tui": False, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "non_inference_max_parallel_workers": 4, + "otel_metrics_port": None, + "shutdown_error_rate": 1.0, +} +_RUN_ID_ADAPTER = TypeAdapter(Identifier) +_PACKAGE_VERSION_ADAPTER = TypeAdapter(Annotated[str, StringConstraints(min_length=1, max_length=128)]) +_RESOLUTION_VALIDATION_MODELS = ( + RunConfig, + ArtifactReference, + BuilderInput, + InputBindings, + ResolvedBuilderInput, + ResolvedInvocation, + ResolvedSubmission, +) + + +class EffectiveDataDesignerSlurmConfig(ContractValue): + """Fully materialized, side-effect-free input to the plan compiler.""" + + run_id: Identifier + package_version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + authored: DataDesignerSlurmConfig + selected_profile: SelectedSlurmProfile + scheduler_bin_path: str | None = None + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + builder_payload: dict[str, JsonValue] | None = None + invocation: ResolvedInvocation + client_runtime: ResolvedClientRuntime + deployment_images: tuple[ResolvedImage, ...] + dependency_lock: ResolvedDependencyLock + submission: ResolvedSubmission + runtime_bundle: ArtifactReference + + +def resolve_slurm_config( + authored: DataDesignerSlurmConfig, + *, + selected_profile: SelectedSlurmProfile, + client_runtime: ResolvedClientRuntime, + deployment_images: tuple[ResolvedImage, ...], + dependency_lock: ResolvedDependencyLock, + runtime_bundle: ArtifactReference, + run_id: str, + package_version: str, + scheduler_bin_path: str | None = None, + resolved_gpus_per_node: int | None = None, + builder_payload: dict[str, JsonValue] | None = None, +) -> EffectiveDataDesignerSlurmConfig: + """Materialize all non-secret defaults from explicitly supplied resolved inputs.""" + try: + run_id = _RUN_ID_ADAPTER.validate_python(run_id, strict=True) + package_version = _PACKAGE_VERSION_ADAPTER.validate_python(package_version, strict=True) + gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) + workspace_root = selected_profile.profile.workspace_root + run_root = posixpath.join(workspace_root, "runs", run_id) + builder, resolved_builder_payload = _resolve_builder( + authored, + run_root=run_root, + builder_payload=builder_payload, + ) + effective = EffectiveDataDesignerSlurmConfig( + run_id=run_id, + package_version=package_version, + authored=authored, + selected_profile=selected_profile, + scheduler_bin_path=scheduler_bin_path, + resolved_gpus_per_node=gpus_per_node, + builder=builder, + builder_payload=resolved_builder_payload, + invocation=_materialize_invocation(authored, workspace_root), + client_runtime=client_runtime, + deployment_images=deployment_images, + dependency_lock=dependency_lock, + submission=_materialize_submission(authored, selected_profile), + runtime_bundle=runtime_bundle, + ) + return validate_effective_slurm_config(effective) + except SlurmConfigResolutionError: + raise + except ValidationError as error: + message = format_validation_error( + error, + subject="Slurm configuration resolution", + models=(*_RESOLUTION_VALIDATION_MODELS, EffectiveDataDesignerSlurmConfig), + ) + raise SlurmConfigResolutionError(message) from None + except ValueError as error: + raise SlurmConfigResolutionError(str(error)) from None + + +def validate_effective_slurm_config( + effective: EffectiveDataDesignerSlurmConfig, +) -> EffectiveDataDesignerSlurmConfig: + """Validate one fully materialized compiler input.""" + authored = effective.authored + profile = effective.selected_profile.profile + workspace_root = profile.workspace_root + if profile.gpus_per_node != "auto" and profile.gpus_per_node != effective.resolved_gpus_per_node: + raise SlurmConfigResolutionError("resolved GPU count does not match the selected profile") + if profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: + raise SlurmConfigResolutionError("mem_per_gpu requires GRES GPU request mode") + + if authored.builder.inline is not None: + if effective.builder_payload is not None: + raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") + else: + if effective.builder_payload is None: + raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") + BuilderInput(inline=effective.builder_payload) + expected_invocation = _materialize_invocation(authored, workspace_root) + if effective.invocation != expected_invocation: + raise SlurmConfigResolutionError("resolved invocation does not match the authored invocation") + expected_submission = _materialize_submission(authored, effective.selected_profile) + if effective.submission != expected_submission: + raise SlurmConfigResolutionError("resolved submission does not match the authored and profile input") + + _validate_dependency_resolution(authored, effective.client_runtime, effective.dependency_lock) + _validate_resolved_images(authored, effective.deployment_images) + derive_deployment_placement(authored.deployments, effective.resolved_gpus_per_node) + shared_images: dict[str, ResolvedImage] = {} + for deployment, image in zip(authored.deployments, effective.deployment_images, strict=True): + group = deployment.resources.shared_node + if group is None: + continue + if profile.gpu_request_mode != "gres": + raise SlurmConfigResolutionError("shared_node requires GRES GPU request mode") + first = shared_images.setdefault(group, image) + if (image.path, image.sha256, image.inspection_facts) != (first.path, first.sha256, first.inspection_facts): + raise SlurmConfigResolutionError("shared_node requires the same resolved serving image") + _validate_sharding_constraints(authored, builder_payload=effective.builder_payload) + managed_assets_path = effective.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + _validate_managed_assets_path(managed_assets_path, workspace_root) + runtime_root = posixpath.join(workspace_root, "runtime") + runtime_name = posixpath.basename(effective.runtime_bundle.path) + if ( + not is_path_below(effective.runtime_bundle.path, runtime_root) + or runtime_name != f"{effective.runtime_bundle.sha256}.tar.gz" + ): + raise SlurmConfigResolutionError( + "runtime bundle must be a content-addressed tar archive below the selected workspace runtime root" + ) + return effective + + +def _resolve_gpu_count(selected: SelectedSlurmProfile, resolved: int | None) -> int: + configured = selected.profile.gpus_per_node + if configured == "auto": + if type(resolved) is not int or resolved <= 0: + raise SlurmConfigResolutionError("auto gpus_per_node requires one resolved positive integer") + return resolved + if resolved is not None and resolved != configured: + raise SlurmConfigResolutionError("resolved GPU count conflicts with the selected profile") + return configured + + +def _resolve_builder( + authored: DataDesignerSlurmConfig, + *, + run_root: str, + builder_payload: dict[str, JsonValue] | None, +) -> tuple[ResolvedBuilderInput, dict[str, JsonValue] | None]: + if authored.builder.inline is not None: + if builder_payload is not None: + raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") + return ( + ResolvedBuilderInput( + inline=authored.builder.inline, + content_sha256=compute_serialized_json_sha256(authored.builder.inline), + model_aliases=get_declared_model_aliases(authored.builder.inline), + ), + None, + ) + if builder_payload is None: + raise SlurmConfigResolutionError("sourced builder input requires its resolved payload") + validated_payload = BuilderInput(inline=builder_payload).inline + assert validated_payload is not None + aliases, digest = get_persisted_builder_identity(validated_payload) + source = ArtifactReference( + path=posixpath.join(run_root, "builder-config.json"), + sha256=digest, + ) + return ( + ResolvedBuilderInput( + authored_source=authored.builder.source, + source=source, + content_sha256=source.sha256, + model_aliases=aliases, + ), + validated_payload, + ) + + +def _materialize_invocation(authored: DataDesignerSlurmConfig, workspace_root: str) -> ResolvedInvocation: + input_bindings = authored.invocation.input_bindings + return ResolvedInvocation( + authored=authored.invocation, + effective_input_bindings=InputBindings( + seed_path=input_bindings.seed_path, + managed_assets_path=input_bindings.managed_assets_path or derive_managed_assets_path(workspace_root), + ), + effective_run_config=_materialize_run_config(authored), + ) + + +def _materialize_submission( + authored: DataDesignerSlurmConfig, + selected_profile: SelectedSlurmProfile, +) -> ResolvedSubmission: + return ResolvedSubmission( + account=authored.submission.account or selected_profile.profile.scheduler.account, + partition=authored.submission.partition or selected_profile.profile.scheduler.partition, + job_name=authored.submission.job_name, + time_limit=authored.submission.time_limit, + comment=authored.submission.comment, + ) + + +def _materialize_run_config(authored: DataDesignerSlurmConfig) -> dict[str, JsonValue]: + values = dict(authored.invocation.run_config) + authored_early_shutdown = {"disable_early_shutdown", "shutdown_error_rate", "shutdown_error_window"}.intersection( + values + ) + for name, value in _COMPATIBILITY_RUN_DEFAULTS.items(): + if authored_early_shutdown and name in {"disable_early_shutdown", "shutdown_error_rate"}: + continue + values.setdefault(name, value) + return RunConfig.model_validate(values).model_dump(mode="json") + + +def _validate_sharding_constraints( + authored: DataDesignerSlurmConfig, + *, + builder_payload: dict[str, JsonValue] | None, +) -> None: + if authored.array_tasks.count == 1: + return + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + data_designer = payload.get("data_designer", payload) + if not isinstance(data_designer, dict): + raise SlurmConfigResolutionError("builder data_designer value must be an object") + if data_designer.get("processors"): + raise SlurmConfigResolutionError("multi-shard runs do not support global processors") + if data_designer.get("profilers"): + raise SlurmConfigResolutionError("multi-shard runs do not support global profilers") + + seed_config = data_designer.get("seed_config") + if isinstance(seed_config, dict): + if seed_config.get("sampling_strategy") == "shuffle": + raise SlurmConfigResolutionError("multi-shard runs do not support shuffled seed input") + if seed_config.get("selection_strategy") is not None: + raise SlurmConfigResolutionError("multi-shard runs do not support authored seed selection strategies") + if authored.invocation.input_bindings.seed_path is None: + raise SlurmConfigResolutionError("multi-shard seed input requires a typed seed_path binding") + + columns = data_designer.get("columns", []) + if not isinstance(columns, list): + raise SlurmConfigResolutionError("builder columns must be a list") + for column in columns: + if not isinstance(column, dict) or not isinstance(column.get("column_type"), str): + raise SlurmConfigResolutionError("multi-shard runs require known column semantics") + column_type = column["column_type"] + if column_type == "image": + raise SlurmConfigResolutionError("multi-shard runs do not support media output columns") + if column_type not in _SHARDABLE_COLUMN_TYPES: + raise SlurmConfigResolutionError( + "multi-shard runs do not support custom, plugin, or unknown column semantics" + ) + if column_type == "validation" and column.get("validator_type") == "local_callable": + raise SlurmConfigResolutionError("multi-shard runs do not support local callable validators") + + +def _validate_dependency_resolution( + authored: DataDesignerSlurmConfig, + client_runtime: ResolvedClientRuntime, + dependency_lock: ResolvedDependencyLock, +) -> None: + if dependency_lock.client_runtime_sha256 != client_runtime.runtime_sha256: + raise SlurmConfigResolutionError("dependency lock does not match the resolved client runtime") + if dependency_lock.python_abi != client_runtime.python_abi: + raise SlurmConfigResolutionError("dependency lock Python ABI does not match the client runtime") + if dependency_lock.base_distributions != client_runtime.distributions: + raise SlurmConfigResolutionError("dependency lock inventory does not match the client runtime") + requirements = authored.client.dependencies.requirements + if requirements is not None: + if dependency_lock.authored_source is not None or dependency_lock.source is not None: + raise SlurmConfigResolutionError("inline requirements cannot resolve from an authored lock file") + if dependency_lock.authored_requirements != tuple(requirements): + raise SlurmConfigResolutionError("dependency lock requirements do not match authored requirements") + elif dependency_lock.authored_source != authored.client.dependencies.lock_file or dependency_lock.source is None: + raise SlurmConfigResolutionError("dependency lock source does not match the authored lock file") + + +def _validate_resolved_images( + authored: DataDesignerSlurmConfig, + deployment_images: tuple[ResolvedImage, ...], +) -> None: + if len(deployment_images) != len(authored.deployments): + raise SlurmConfigResolutionError("resolved serving images must match the authored deployment count") + for deployment, image in zip(authored.deployments, deployment_images, strict=True): + _validate_resolved_image_identity(image) + if image.kind is not ImageKind.SERVING: + raise SlurmConfigResolutionError("resolved deployment image must contain serving inspection facts") + if image.authored_ref != deployment.server.image: + raise SlurmConfigResolutionError("resolved deployment image does not match the authored reference") + + +def _validate_resolved_image_identity(image: ResolvedImage) -> None: + if image.sha256 != image.inspection.sqsh_sha256: + raise SlurmConfigResolutionError("resolved image digest does not match its inspection record") + if image.authored_ref.path is not None and image.path != image.authored_ref.path: + raise SlurmConfigResolutionError("resolved image path does not match the authored path") + + +def _validate_managed_assets_path(managed_assets_path: str, workspace_root: str) -> None: + workspace_state = tuple( + posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks", "runs") + ) + if any(paths_overlap(managed_assets_path, path) for path in workspace_state): + raise SlurmConfigResolutionError("managed_assets_path must not overlap package-managed workspace state") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py new file mode 100644 index 000000000..8a35033c2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath + +from pydantic import JsonValue + +from data_designer.config import RunConfig +from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.contracts import derive_managed_assets_path +from data_designer.slurm.planning.builder_identity import get_persisted_builder_identity +from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.models import ( + ResolvedDependencyLock, + ResolvedSlurmRunPlan, +) + + +def validate_resolved_plan( + authored: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + plan: ResolvedSlurmRunPlan, + *, + builder_payload: dict[str, JsonValue] | None = None, +) -> ResolvedSlurmRunPlan: + """Validate cross-record identities and digests for one resolved plan.""" + _require( + plan.authored_config.sha256 == authored.compute_sha256(), + "authored config digest does not match the resolved plan", + ) + _require(plan.invocation.authored == authored.invocation, "resolved invocation does not match authored input") + expected_managed_assets = authored.invocation.input_bindings.managed_assets_path or derive_managed_assets_path( + plan.selected_profile.profile.workspace_root + ) + _require( + plan.invocation.effective_input_bindings.seed_path == authored.invocation.input_bindings.seed_path + and plan.invocation.effective_input_bindings.managed_assets_path == expected_managed_assets, + "resolved input bindings do not match authored/profile input", + ) + explicit_run_config = RunConfig.model_validate(authored.invocation.run_config).model_dump( + mode="json", + exclude_unset=True, + ) + _require_json_subset(plan.invocation.effective_run_config, explicit_run_config, path="run_config") + _require(plan.client.authored == authored.client, "resolved client does not match authored input") + _require( + tuple(deployment.authored for deployment in plan.deployments) == tuple(authored.deployments), + "resolved deployments do not match authored order and values", + ) + _require(plan.array_tasks == authored.array_tasks, "resolved array task policy does not match authored input") + + if authored.builder.inline is not None: + _require(builder_payload is None, "inline builder input must not retain a separate payload") + _require( + plan.builder.inline == authored.builder.inline, "resolved inline builder does not match authored input" + ) + model_aliases, digest = get_persisted_builder_identity(authored.builder.inline) + _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match inline builder") + _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match inline builder") + else: + _require( + plan.builder.authored_source == authored.builder.source, + "resolved builder source does not match authored input", + ) + if builder_payload is None: + raise SlurmPlanContractError("sourced builder validation requires its resolved payload") + model_aliases, digest = get_persisted_builder_identity(builder_payload) + _require(plan.builder.source is not None, "resolved builder source artifact is missing") + assert plan.builder.source is not None + expected_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "builder-config.json") + _require(plan.builder.source.path == expected_path, "resolved builder artifact path does not match the run") + _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match builder source") + _require(plan.builder.source.sha256 == digest, "resolved builder digest does not match builder source") + _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match builder source") + + expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account + expected_partition = authored.submission.partition or plan.selected_profile.profile.scheduler.partition + _require(plan.submission.account == expected_account, "resolved account does not match authored/profile input") + _require( + plan.submission.partition == expected_partition, "resolved partition does not match authored/profile input" + ) + _require( + plan.submission.job_name == authored.submission.job_name, "resolved job name does not match authored input" + ) + _require( + plan.submission.time_limit == authored.submission.time_limit, + "resolved time limit does not match authored input", + ) + _require(plan.submission.comment == authored.submission.comment, "resolved comment does not match authored input") + + _require( + plan.client.dependency_lock.sha256 == dependency_lock.compute_sha256(), + "dependency lock digest does not match the resolved plan", + ) + _require( + dependency_lock.client_runtime_sha256 == plan.client.runtime.runtime_sha256, + "dependency lock runtime fingerprint does not match the resolved client runtime", + ) + _require( + dependency_lock.python_abi == plan.client.runtime.python_abi, + "dependency lock Python ABI does not match client runtime", + ) + _require( + dependency_lock.base_distributions == plan.client.runtime.distributions, + "dependency lock inventory does not match client runtime inspection", + ) + authored_requirements = authored.client.dependencies.requirements + if authored_requirements is not None: + _require( + dependency_lock.authored_source is None and dependency_lock.source is None, + "dependency lock source is present for authored requirements", + ) + _require( + dependency_lock.authored_requirements == tuple(authored_requirements), + "dependency lock requirements do not match authored requirements", + ) + else: + _require( + dependency_lock.authored_source == authored.client.dependencies.lock_file + and dependency_lock.source is not None, + "dependency lock source does not match the authored lock file", + ) + return plan + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise SlurmPlanContractError(message) + + +def _require_json_subset(actual: JsonValue, expected: JsonValue, *, path: str) -> None: + if isinstance(expected, dict): + _require(isinstance(actual, dict), f"explicit {path} value does not match the resolved plan") + for key, value in expected.items(): + _require(key in actual, f"explicit {path}.{key} value is missing from the resolved plan") + _require_json_subset(actual[key], value, path=f"{path}.{key}") + return + _require(actual == expected, f"explicit {path} value does not match the resolved plan") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py new file mode 100644 index 000000000..02864bfce --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local runtime primitives for resolved Slurm plans.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from data_designer.slurm.runtime.bundle import stage_runtime_bundle # noqa: F401 + from data_designer.slurm.runtime.controller import OneNodeAllocationController # noqa: F401 + from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode # noqa: F401 + +_LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "OneNodeAllocationController": ( + "data_designer.slurm.runtime.controller", + "OneNodeAllocationController", + ), + "SlurmRuntimeError": ("data_designer.slurm.runtime.errors", "SlurmRuntimeError"), + "SlurmRuntimeErrorCode": ("data_designer.slurm.runtime.errors", "SlurmRuntimeErrorCode"), + "stage_runtime_bundle": ("data_designer.slurm.runtime.bundle", "stage_runtime_bundle"), +} + +__all__ = [ + "OneNodeAllocationController", + "SlurmRuntimeError", + "SlurmRuntimeErrorCode", + "stage_runtime_bundle", +] + + +def __getattr__(name: str) -> object: + """Lazily load orchestration dependencies only when requested.""" + if name in _LAZY_IMPORTS: + module_name, attribute_name = _LAZY_IMPORTS[name] + attribute = getattr(importlib.import_module(module_name), attribute_name) + globals()[name] = attribute + return attribute + raise AttributeError(f"module 'data_designer.slurm.runtime' has no attribute {name!r}") + + +def __dir__() -> list[str]: + """Return public runtime exports for interactive discovery.""" + return __all__ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py new file mode 100644 index 000000000..95ff51045 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM middleware that exposes bounded queue admission to client AIMD.""" + +from __future__ import annotations + +import importlib +import json +import math +import os +import threading +import time +from collections.abc import Awaitable, Callable, Iterable, Mapping +from dataclasses import dataclass +from http import HTTPStatus +from typing import Any, Protocol + +MAX_WAITING_REQUESTS_ENVIRONMENT = "DD_VLLM_MAX_WAITING_REQUESTS" +RETRY_AFTER_SECONDS_ENVIRONMENT = "DD_VLLM_RETRY_AFTER_SECONDS" +_METRIC_NAME = "vllm:num_requests_waiting" +_EXEMPT_PATHS = ("/health", "/ready", "/metrics", "/version", "/v1/models", "/ping") + +AsgiMessage = dict[str, Any] +AsgiScope = dict[str, Any] +AsgiReceive = Callable[[], Awaitable[AsgiMessage]] +AsgiSend = Callable[[AsgiMessage], Awaitable[None]] +AsgiApp = Callable[[AsgiScope, AsgiReceive, AsgiSend], Awaitable[None]] + + +class QueueDepthReader(Protocol): + """Read the latest aggregate waiting-request count.""" + + def __call__(self) -> int | None: + """Return queue depth, or ``None`` when metrics are unavailable.""" + ... + + +@dataclass(frozen=True, slots=True) +class QueueBackpressureSettings: + """Validated worker-owned queue admission settings.""" + + max_waiting_requests: int + retry_after_seconds: int | None + poll_interval_seconds: float = 0.1 + stale_after_seconds: float = 1.0 + + def __post_init__(self) -> None: + if type(self.max_waiting_requests) is not int or self.max_waiting_requests < 0: + raise ValueError("maximum waiting requests must be non-negative") + if self.retry_after_seconds is not None and ( + type(self.retry_after_seconds) is not int or self.retry_after_seconds <= 0 + ): + raise ValueError("retry-after seconds must be positive or absent") + if self.poll_interval_seconds <= 0 or self.stale_after_seconds <= 0: + raise ValueError("queue sampler intervals must be positive") + + @classmethod + def from_environment(cls, environment: Mapping[str, str] | None = None) -> QueueBackpressureSettings: + """Load the policy transported by the structured runtime step.""" + source = os.environ if environment is None else environment + maximum = _parse_non_negative_integer(source.get(MAX_WAITING_REQUESTS_ENVIRONMENT), default=128) + retry_value = source.get(RETRY_AFTER_SECONDS_ENVIRONMENT) + retry_after = None if retry_value == "" else _parse_positive_integer(retry_value, default=1) + return cls(maximum, retry_after) + + +@dataclass(frozen=True, slots=True) +class QueueSnapshot: + """One sampled queue depth and its monotonic observation time.""" + + depth: int | None + observed_at: float + + +class QueueBackpressureController: + """Cache queue metrics away from the request path and decide admission.""" + + def __init__( + self, + settings: QueueBackpressureSettings | None = None, + reader: QueueDepthReader | None = None, + *, + start_sampler: bool = True, + ) -> None: + self.settings = settings or QueueBackpressureSettings.from_environment() + self._reader = reader or read_vllm_queue_depth + self._start_sampler = start_sampler + self._snapshot = QueueSnapshot(None, 0.0) + self._lock = threading.Lock() + self._thread: threading.Thread | None = None + + def sample_once(self) -> QueueSnapshot: + """Refresh the cached queue depth, failing open on reader errors.""" + try: + depth = self._reader() + except Exception: + depth = None + if depth is not None: + depth = max(0, depth) + snapshot = QueueSnapshot(depth, time.monotonic()) + with self._lock: + self._snapshot = snapshot + return snapshot + + def should_reject(self) -> tuple[bool, QueueSnapshot]: + """Return the fail-open admission decision and supporting snapshot.""" + self._ensure_sampler() + with self._lock: + snapshot = self._snapshot + stale = time.monotonic() - snapshot.observed_at > self.settings.stale_after_seconds + reject = snapshot.depth is not None and not stale and snapshot.depth > self.settings.max_waiting_requests + return reject, snapshot + + def _ensure_sampler(self) -> None: + if not self._start_sampler or self._thread is not None: + return + with self._lock: + if self._thread is None: + self._thread = threading.Thread(target=self._sample_forever, name="dd-vllm-queue-depth", daemon=True) + self._thread.start() + + def _sample_forever(self) -> None: + while True: + self.sample_once() + time.sleep(self.settings.poll_interval_seconds) + + +class QueueDepthBackpressureMiddleware: + """Reject non-health HTTP requests with 429 above the resolved queue threshold.""" + + def __init__(self, app: AsgiApp, controller: QueueBackpressureController | None = None) -> None: + self.app = app + self.controller = controller or QueueBackpressureController() + + async def __call__(self, scope: AsgiScope, receive: AsgiReceive, send: AsgiSend) -> None: + """Apply queue admission without changing exempt or accepted requests.""" + path = str(scope.get("path", "")) + if scope.get("type") != "http" or path in _EXEMPT_PATHS: + await self.app(scope, receive, send) + return + reject, snapshot = self.controller.should_reject() + if not reject: + await self.app(scope, receive, send) + return + await _send_overload(send, self.controller.settings, snapshot) + + +def read_vllm_queue_depth() -> int | None: + """Return aggregate vLLM queue depth from supported metrics registries.""" + for source in (_read_vllm_metrics, _read_prometheus_metrics): + values = _collect_metric_values(source()) + if values: + return max(0, int(sum(values))) + return None + + +async def _send_overload( + send: AsgiSend, + settings: QueueBackpressureSettings, + snapshot: QueueSnapshot, +) -> None: + body = json.dumps( + { + "error": { + "message": "serving queue admission threshold exceeded", + "type": "rate_limit_exceeded", + "code": HTTPStatus.TOO_MANY_REQUESTS.value, + "queue_depth": snapshot.depth, + "max_waiting_requests": settings.max_waiting_requests, + } + }, + separators=(",", ":"), + ).encode() + headers = [(b"content-type", b"application/json")] + if settings.retry_after_seconds is not None: + headers.append((b"retry-after", str(settings.retry_after_seconds).encode())) + await send({"type": "http.response.start", "status": HTTPStatus.TOO_MANY_REQUESTS.value, "headers": headers}) + await send({"type": "http.response.body", "body": body}) + + +def _read_vllm_metrics() -> Iterable[object]: + try: + module = importlib.import_module("vllm.v1.metrics.reader") + return module.get_metrics_snapshot() + except (ImportError, AttributeError, RuntimeError): + return () + + +def _read_prometheus_metrics() -> Iterable[object]: + try: + module = importlib.import_module("prometheus_client") + return tuple(sample for family in module.REGISTRY.collect() for sample in getattr(family, "samples", ())) + except (ImportError, AttributeError, RuntimeError): + return () + + +def _collect_metric_values(metrics: Iterable[object]) -> tuple[float, ...]: + names = {_METRIC_NAME, _METRIC_NAME.replace(":", "_")} + values: list[float] = [] + for metric in metrics: + if str(getattr(metric, "name", "")) not in names: + continue + try: + value = float(getattr(metric, "value")) + if math.isfinite(value) and value >= 0: + values.append(value) + except (AttributeError, TypeError, ValueError): + continue + return tuple(values) + + +def _parse_non_negative_integer(value: str | None, *, default: int) -> int: + if value is None: + return default + if not value.isascii() or not value.isdigit(): + raise ValueError("queue threshold is invalid") + return int(value) + + +def _parse_positive_integer(value: str | None, *, default: int) -> int: + parsed = _parse_non_negative_integer(value, default=default) + if parsed <= 0: + raise ValueError("retry-after seconds must be positive") + return parsed + + +__all__ = [ + "MAX_WAITING_REQUESTS_ENVIRONMENT", + "QueueBackpressureController", + "QueueBackpressureSettings", + "QueueDepthBackpressureMiddleware", + "QueueDepthReader", + "QueueSnapshot", + "RETRY_AFTER_SECONDS_ENVIRONMENT", + "read_vllm_queue_depth", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py new file mode 100644 index 000000000..214711b7f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed allocation step manifest produced by the native client runtime.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping +from pathlib import Path +from typing import Literal + +from data_designer.slurm.config.environment import ( + SecretRef, + collect_secret_environment_names, +) +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.manifest import RuntimeBootstrapManifest, RuntimeProbeSpec, RuntimeStepSpec +from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.paths import ALLOCATION_SCRATCH_CONTAINER_ROOT +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.preflight import AllocationLayout, validate_allocation_layout +from data_designer.slurm.runtime.server_manifest import build_server_steps +from data_designer.slurm.runtime.steps import build_client_command, build_endpoint_command +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment + + +def build_runtime_manifest( + context: AllocationContext, + environment: Mapping[str, str], + *, + runtime_root: Path, + log_directory: Path, + layout: AllocationLayout, +) -> RuntimeBootstrapManifest: + """Build the secret-free command handoff for the Bash controller.""" + plan = context.plan + validate_allocation_layout(plan, layout) + if ( + not runtime_root.is_absolute() + or runtime_root.name != "runtime" + or re.fullmatch(r"data-designer-slurm-[1-9][0-9]*-[0-9]+", runtime_root.parent.name) is None + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime root is outside allocation-local scratch", + ) + runtime_container_root = Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime" + deployments = resolve_allocation_deployments(context, environment) + endpoints = tuple( + RuntimeEndpoint( + model_alias=deployment.model_alias, + served_model_name=deployment.served_model_name, + host="127.0.0.1", + port=deployment.logical_endpoint.port, + ) + for deployment in deployments + ) + steps: list[RuntimeStepSpec] = [ + _build_client_step( + RuntimeStepRole.CLIENT_PREFLIGHT, + "client-preflight", + "preflight", + context, + environment, + endpoints, + runtime_root, + log_directory, + layout, + ) + ] + steps.extend(build_server_steps(deployments, context, runtime_container_root.as_posix(), log_directory, layout)) + steps.extend( + _build_endpoint_step(deployment, context, runtime_root, log_directory, layout) for deployment in deployments + ) + steps.append( + _build_client_step( + RuntimeStepRole.CLIENT, + "client-generation", + "client", + context, + environment, + endpoints, + runtime_root, + log_directory, + layout, + ) + ) + secret_names = set(collect_secret_environment_names(plan)) + secret_names.update( + name + for deployment in deployments + for name, binding in deployment.launch_policy.environment.items() + if isinstance(binding, SecretRef) + ) + return RuntimeBootstrapManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + plan_sha256=plan.compute_sha256(), + all_secret_environment_names=tuple(sorted(secret_names)), + steps=tuple(steps), + ) + + +def _build_client_step( + role: RuntimeStepRole, + step_id: str, + operation: str, + context: AllocationContext, + environment: Mapping[str, str], + endpoints: tuple[RuntimeEndpoint, ...], + runtime_host_root: Path, + log_directory: Path, + layout: AllocationLayout, +) -> RuntimeStepSpec: + plan = context.plan + retry_resume_mode = None if context.retry_plan is None else context.retry_plan.effective_resume_mode + command = build_client_command( + "preflight" if operation == "preflight" else "run", + plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + retry_resume_mode, + ) + if operation == "client": + endpoint_arguments = tuple( + argument + for endpoint in endpoints + for argument in ("--endpoint", f"{endpoint.model_alias}=http://{endpoint.host}:{endpoint.port}/v1") + ) + retry_binding = ( + () + if context.retry_plan is None + else ( + "--retry-id", + context.retry_plan.retry_id, + "--retry-plan-sha256", + context.retry_plan.compute_sha256(), + "--effective-resume-mode", + context.retry_plan.effective_resume_mode, + ) + ) + command = ( + plan.client.runtime.python_executable, + "-m", + "data_designer.slurm.runtime.entrypoint", + "client", + *command[4:6], + *command[10:12], + *retry_binding, + *endpoint_arguments, + ) + secret_names = collect_secret_environment_names( + (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ) + allocation_environment = ( + {"SLURM_JOB_GPUS": environment["SLURM_JOB_GPUS"]} + if plan.selected_profile.profile.gpu_request_mode == "gres" + else {} + ) + return _step( + step_id=step_id, + role=role, + execution="native", + image_path=None, + command=command, + cpus=plan.client.authored.cpus, + gpu_indices=(), + literal_environment={ + "DATA_DESIGNER_SLURM_SCRATCH_ROOT": runtime_host_root.parent.as_posix(), + "LC_ALL": "C", + "PYTHONPATH": runtime_host_root.as_posix(), + **allocation_environment, + }, + secret_environment={name: name for name in secret_names}, + environment_prefixes={}, + container_environment=(), + log_directory=log_directory, + node_hosts=(layout.get_host(plan.client.host_node_index),), + ) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_root: Path, + log_directory: Path, + layout: AllocationLayout, +) -> RuntimeStepSpec: + proxy_path = runtime_root / "data_designer/slurm/runtime/proxy.py" + backend_hosts = ( + tuple(layout.get_host(backend.node_index) for backend in deployment.backend_endpoints) + if len(layout.node_hosts) > 1 + else None + ) + command = build_endpoint_command( + deployment, + context.plan, + proxy_path, + deployment.logical_endpoint.port, + backend_hosts=backend_hosts, + ) + return _step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + execution="native", + image_path=None, + command=command, + cpus=context.plan.client.authored.cpus, + gpu_indices=(), + literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_root.as_posix()}, + secret_environment={}, + environment_prefixes={}, + container_environment=(), + log_directory=log_directory, + node_hosts=(layout.get_host(context.plan.client.host_node_index),), + readiness=( + RuntimeProbeSpec( + host="127.0.0.1", + port=deployment.logical_endpoint.port, + path="/health", + deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + ), + ), + ) + + +def _step( + *, + step_id: str, + role: RuntimeStepRole, + execution: Literal["native", "container"], + image_path: str | None, + command: tuple[str, ...], + cpus: int, + gpu_indices: tuple[int, ...], + node_hosts: tuple[str, ...], + literal_environment: dict[str, str], + secret_environment: dict[str, str], + environment_prefixes: dict[str, str], + container_environment: tuple[str, ...], + log_directory: Path, + kill_on_bad_exit: bool = False, + launch_delay_seconds: int = 0, + readiness: tuple[RuntimeProbeSpec, ...] = (), +) -> RuntimeStepSpec: + return RuntimeStepSpec( + step_id=step_id, + role=role, + execution=execution, + image_path=image_path, + command=command, + cpus=cpus, + gpu_indices=gpu_indices, + node_hosts=node_hosts, + kill_on_bad_exit=kill_on_bad_exit, + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=container_environment, + stdout_path=(log_directory / f"{step_id}.out").as_posix(), + stderr_path=(log_directory / f"{step_id}.err").as_posix(), + launch_delay_seconds=launch_delay_seconds, + readiness=readiness, + ) + + +__all__ = ["RuntimeBootstrapManifest", "RuntimeProbeSpec", "RuntimeStepSpec", "build_runtime_manifest"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py new file mode 100644 index 000000000..8e70ffc14 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic content-addressed staging for the allocation runtime bundle.""" + +from __future__ import annotations + +import gzip +import hashlib +import io +import os +import re +import secrets +import stat +import tarfile +from contextlib import suppress +from pathlib import Path + +from data_designer.slurm.contracts import ArtifactReference, validate_absolute_path +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode + +_DIRECTORY_MODE = 0o700 +_FILE_MODE = 0o600 +_ENTRYPOINT_MODE = 0o500 +_SOURCE_MODE = 0o400 +_MAXIMUM_SOURCE_SIZE = 16 * 1024 * 1024 +_TEMPORARY_NAME_PATTERN = re.compile(r"^\.runtime\.[0-9a-f]{16}\.tmp$") +_SHELL_SOURCE_NAMES = ("scratch.sh", "entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh") +_SLURM_PACKAGE_ROOT = "data_designer/slurm" +_SOURCE_MANIFEST_NAME = f"{_SLURM_PACKAGE_ROOT}/runtime/slurm-sources.txt" + + +def stage_runtime_bundle(workspace_root: str | Path) -> ArtifactReference: + """Convergently stage the versioned allocation entrypoint by content digest. + + Raises: + SlurmRuntimeError: If the workspace or runtime artifact is unsafe or cannot be staged. + """ + try: + normalized_workspace = Path(validate_absolute_path(Path(workspace_root).as_posix())) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime workspace is invalid") from error + try: + archive = _build_runtime_archive() + digest = hashlib.sha256(archive).hexdigest() + runtime_root = normalized_workspace / "runtime" + archive_path = runtime_root / f"{digest}.tar.gz" + _ensure_private_runtime_directory(normalized_workspace, runtime_root) + _ArchivePublisher(runtime_root).publish(archive_path.name, archive) + except OSError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "cannot stage allocation runtime bundle" + ) from error + return ArtifactReference(path=archive_path.as_posix(), sha256=digest) + + +def _build_runtime_archive() -> bytes: + source_root = Path(__file__).parents[1] + sources = _collect_slurm_sources(source_root) + output = io.BytesIO() + with ( + gzip.GzipFile(fileobj=output, mode="wb", filename="", mtime=0) as compressed, + tarfile.open(fileobj=compressed, mode="w", format=tarfile.PAX_FORMAT) as archive, + ): + for name in _SHELL_SOURCE_NAMES: + mode = _ENTRYPOINT_MODE if name == "entrypoint.sh" else _SOURCE_MODE + _add_archive_file(archive, name, _read_runtime_source(source_root / "runtime" / name), mode=mode) + manifest = "".join(f"{archive_name}\n" for archive_name, _ in sources).encode() + _add_archive_file(archive, _SOURCE_MANIFEST_NAME, manifest, mode=_SOURCE_MODE) + for archive_name, source_path in sources: + _add_archive_file(archive, archive_name, _read_runtime_source(source_path), mode=_SOURCE_MODE) + return output.getvalue() + + +def _collect_slurm_sources(source_root: Path) -> tuple[tuple[str, Path], ...]: + relative_sources = sorted( + (path.relative_to(source_root) for path in source_root.rglob("*.py")), + key=lambda path: path.as_posix(), + ) + return tuple( + (f"{_SLURM_PACKAGE_ROOT}/{relative.as_posix()}", source_root / relative) for relative in relative_sources + ) + + +def _add_archive_file(archive: tarfile.TarFile, name: str, content: bytes, *, mode: int) -> None: + entry = tarfile.TarInfo(name) + entry.size = len(content) + entry.mode = mode + entry.uid = 0 + entry.gid = 0 + entry.uname = "" + entry.gname = "" + entry.mtime = 0 + archive.addfile(entry, io.BytesIO(content)) + + +def _read_runtime_source(path: Path) -> bytes: + before = path.lstat() + if not stat.S_ISREG(before.st_mode) or before.st_size > _MAXIMUM_SOURCE_SIZE: + raise OSError(f"allocation runtime source {path} is not a bounded regular file") + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)) + try: + opened = os.fstat(descriptor) + if _file_identity(before) != _file_identity(opened): + raise OSError(f"allocation runtime source {path} changed while it was opened") + chunks: list[bytes] = [] + content_size = 0 + while chunk := os.read(descriptor, min(1024 * 1024, _MAXIMUM_SOURCE_SIZE + 1 - content_size)): + chunks.append(chunk) + content_size += len(chunk) + if content_size > _MAXIMUM_SOURCE_SIZE: + raise OSError(f"allocation runtime source {path} exceeds its size limit") + after = os.fstat(descriptor) + if _file_identity(opened) != _file_identity(after): + raise OSError(f"allocation runtime source {path} changed while it was read") + return b"".join(chunks) + finally: + os.close(descriptor) + + +def _ensure_private_runtime_directory(workspace_root: Path, runtime_root: Path) -> None: + workspace_status = workspace_root.lstat() + if not stat.S_ISDIR(workspace_status.st_mode) or workspace_status.st_mode & 0o022: + raise OSError(f"runtime workspace {workspace_root} is not a private directory") + runtime_root.mkdir(mode=_DIRECTORY_MODE, exist_ok=True) + runtime_status = runtime_root.lstat() + if not stat.S_ISDIR(runtime_status.st_mode): + raise OSError(f"runtime root {runtime_root} is not a directory") + descriptor = os.open( + runtime_root, + os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0), + ) + try: + opened_status = os.fstat(descriptor) + if (runtime_status.st_dev, runtime_status.st_ino) != (opened_status.st_dev, opened_status.st_ino): + raise OSError(f"runtime root {runtime_root} changed while it was opened") + os.fchmod(descriptor, _DIRECTORY_MODE) + os.fsync(descriptor) + finally: + os.close(descriptor) + + +class _ArchivePublisher: + """Own descriptor-bound publication of immutable runtime archives.""" + + def __init__(self, runtime_root: Path) -> None: + self._runtime_root = runtime_root + + def publish(self, name: str, content: bytes) -> None: + """Publish one archive or confirm an identical durable publication.""" + directory_descriptor = os.open( + self._runtime_root, + os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0), + ) + try: + self._publish_in_directory(directory_descriptor, name, content) + finally: + os.close(directory_descriptor) + + def _publish_in_directory(self, directory_descriptor: int, name: str, content: bytes) -> None: + existing = _read_existing_archive(directory_descriptor, name, expected_size=len(content)) + if existing is not None: + self._accept_existing(directory_descriptor, existing, content) + return + temporary_name = self._write_temporary(directory_descriptor, content) + try: + self._link_or_accept_existing(directory_descriptor, temporary_name, name, content) + finally: + with suppress(FileNotFoundError): + os.unlink(temporary_name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + @staticmethod + def _accept_existing(directory_descriptor: int, existing: bytes, content: bytes) -> None: + if existing != content: + raise OSError("content-addressed runtime bundle contains different bytes") + os.fsync(directory_descriptor) + + @staticmethod + def _write_temporary(directory_descriptor: int, content: bytes) -> str: + temporary_name, descriptor = _create_temporary_archive(directory_descriptor) + try: + os.fchmod(descriptor, _FILE_MODE) + with os.fdopen(descriptor, "wb") as output: + descriptor = -1 + output.write(content) + output.flush() + os.fsync(output.fileno()) + except BaseException: + with suppress(OSError): + os.unlink(temporary_name, dir_fd=directory_descriptor) + raise + finally: + if descriptor >= 0: + os.close(descriptor) + return temporary_name + + @staticmethod + def _link_or_accept_existing( + directory_descriptor: int, + temporary_name: str, + name: str, + content: bytes, + ) -> None: + try: + os.link( + temporary_name, + name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + existing = _read_existing_archive(directory_descriptor, name, expected_size=len(content)) + if existing != content: + raise OSError("content-addressed runtime bundle contains different bytes") from None + + +def _create_temporary_archive(directory_descriptor: int) -> tuple[str, int]: + for _ in range(100): + temporary_name = f".runtime.{secrets.token_hex(8)}.tmp" + try: + descriptor = os.open( + temporary_name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + _FILE_MODE, + dir_fd=directory_descriptor, + ) + except FileExistsError: + continue + return temporary_name, descriptor + raise OSError("cannot allocate a unique runtime bundle temporary name") + + +def _read_existing_archive(directory_descriptor: int, name: str, *, expected_size: int) -> bytes | None: + try: + before = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return None + before = _repair_interrupted_archive_publication(directory_descriptor, name, before) + if not _is_valid_archive_status(before, expected_size): + raise OSError("runtime bundle is not a restrictive single-link regular file") + return _read_bound_archive(directory_descriptor, name, before, expected_size) + + +def _read_bound_archive( + directory_descriptor: int, + name: str, + before: os.stat_result, + expected_size: int, +) -> bytes: + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + try: + after_open = os.fstat(descriptor) + if _file_identity(before) != _file_identity(after_open): + raise OSError("runtime bundle changed while it was opened") + with os.fdopen(descriptor, "rb") as existing: + descriptor = -1 + content = existing.read(expected_size + 1) + after_read = os.fstat(existing.fileno()) + if _file_identity(after_open) != _file_identity(after_read): + raise OSError("runtime bundle changed while it was read") + return content + finally: + if descriptor >= 0: + os.close(descriptor) + + +def _is_valid_archive_status(status: os.stat_result, expected_size: int) -> bool: + return ( + stat.S_ISREG(status.st_mode) + and status.st_nlink == 1 + and not status.st_mode & 0o077 + and status.st_size == expected_size + ) + + +def _file_identity(status: os.stat_result) -> tuple[int, int, int, int, int]: + return status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns, status.st_ctime_ns + + +def _repair_interrupted_archive_publication( + directory_descriptor: int, + name: str, + status: os.stat_result, +) -> os.stat_result: + if status.st_nlink == 1: + return status + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077 or status.st_nlink != 2: + return status + matching_names: list[str] = [] + for candidate in os.listdir(directory_descriptor): + if _TEMPORARY_NAME_PATTERN.fullmatch(candidate) is None: + continue + try: + candidate_status = os.stat(candidate, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + continue + if (candidate_status.st_dev, candidate_status.st_ino) == (status.st_dev, status.st_ino): + matching_names.append(candidate) + if len(matching_names) != 1: + return os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + try: + os.unlink(matching_names[0], dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + except FileNotFoundError: + pass + repaired = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(repaired.st_mode) or repaired.st_nlink != 1 or repaired.st_mode & 0o077: + raise OSError("runtime bundle changed while recovering interrupted publication") + return repaired diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh new file mode 100644 index 000000000..c5e43336a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +DD_MANAGED_PIDS=() +DD_REQUIRED_PIDS=() +DD_CLEANUP_COMPLETE=0 + +dd_start_runtime_timer() { + coproc DD_RUNTIME_TIMER { while IFS= read -r _; do :; done; } +} + +dd_sleep() { + if [[ ${DD_RUNTIME_TIMER[0]:-} =~ ^[0-9]+$ ]]; then + read -r -t "$1" -u "${DD_RUNTIME_TIMER[0]}" _ || true + else + sleep "$1" + fi +} + +dd_register_required_pid() { + DD_MANAGED_PIDS+=("$1") + DD_REQUIRED_PIDS+=("$1") +} + +dd_require_running() { + local pid + for pid in "${DD_REQUIRED_PIDS[@]+"${DD_REQUIRED_PIDS[@]}"}"; do + kill -0 "${pid}" 2>/dev/null || return 1 + done +} + +dd_wait_for_client() { + local client_pid=$1 + local status=0 + while kill -0 "${client_pid}" 2>/dev/null; do + wait -n || status=$? + if ! kill -0 "${client_pid}" 2>/dev/null; then + return "${status}" + fi + dd_require_running || return 70 + done + wait "${client_pid}" +} + +dd_cleanup_steps() { + ((DD_CLEANUP_COMPLETE == 0)) || return 0 + local index pid position + local -a indices=() + ((${#DD_MANAGED_PIDS[@]} == 0)) || indices=("${!DD_MANAGED_PIDS[@]}") + for ((position = ${#indices[@]} - 1; position >= 0; position--)); do + index=${indices[position]} + pid=${DD_MANAGED_PIDS[index]} + kill -0 "${pid}" 2>/dev/null && kill -TERM "${pid}" 2>/dev/null || true + done + dd_sleep 0.2 + for ((position = ${#indices[@]} - 1; position >= 0; position--)); do + index=${indices[position]} + pid=${DD_MANAGED_PIDS[index]} + kill -0 "${pid}" 2>/dev/null && kill -KILL "${pid}" 2>/dev/null || true + done + for ((position = ${#indices[@]} - 1; position >= 0; position--)); do + index=${indices[position]} + wait "${DD_MANAGED_PIDS[index]}" 2>/dev/null || true + done + DD_CLEANUP_COMPLETE=1 +} + +dd_stop_runtime_timer() { + if [[ ${DD_RUNTIME_TIMER_PID:-} =~ ^[0-9]+$ ]]; then + kill -TERM "${DD_RUNTIME_TIMER_PID}" 2>/dev/null || true + wait "${DD_RUNTIME_TIMER_PID}" 2>/dev/null || true + fi +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py new file mode 100644 index 000000000..0d54156fa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load allocation identity from native compute-visible persisted state.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from pydantic import ValidationError + +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.state import RetryPlan, SlurmStateError, SlurmStateWriter +from data_designer.slurm.state.filesystem import open_verified_directory, read_regular_text + +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +def load_allocation_context( + plan_path: Path, + attempt_directory: Path, + environment: Mapping[str, str], + *, + retry_id: str | None = None, + retry_plan_sha256: str | None = None, + effective_resume_mode: str | None = None, +) -> tuple[AllocationContext, SlurmStateWriter]: + """Load one scheduler-selected shard attempt through native host paths.""" + writer = _load_state_writer(plan_path, attempt_directory) + plan = writer.load_resolved_plan() + expected_plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + if plan_path != expected_plan_path: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime plan path does not match persisted run intent", + ) + shard = _select_shard(plan.shards, _scheduler_task_id(environment.get("SLURM_ARRAY_TASK_ID"))) + host_attempt_directory = expected_plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt_directory.name + if attempt_directory != host_attempt_directory: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "attempt directory does not match the scheduler array task", + ) + attempt = writer.load_attempt(shard.shard_id, attempt_directory.name) + array_job_id = _scheduler_task_id(environment.get("SLURM_ARRAY_JOB_ID")) + if attempt.scheduler is None or attempt.scheduler.array_job_id != array_job_id: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "scheduler array job does not match the persisted attempt", + ) + retry_plan = _load_retry_plan( + writer, + retry_id=retry_id, + retry_plan_sha256=retry_plan_sha256, + effective_resume_mode=effective_resume_mode, + ) + return AllocationContext(plan, shard, attempt, host_attempt_directory, retry_plan), writer + + +def _load_retry_plan( + writer: SlurmStateWriter, + *, + retry_id: str | None, + retry_plan_sha256: str | None, + effective_resume_mode: str | None, +) -> RetryPlan | None: + arguments = (retry_id, retry_plan_sha256, effective_resume_mode) + if all(argument is None for argument in arguments): + return None + if any(argument is None for argument in arguments): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry binding is incomplete") + assert retry_id is not None and retry_plan_sha256 is not None and effective_resume_mode is not None + if effective_resume_mode not in {"never", "always"}: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry resume mode is invalid") + try: + retry_plan = writer.load_retry_plan(retry_id) + except SlurmStateError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime retry plan is unavailable") from error + if retry_plan.compute_sha256() != retry_plan_sha256 or retry_plan.effective_resume_mode != effective_resume_mode: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime retry binding differs from persisted state", + ) + return retry_plan + + +def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWriter: + if not plan_path.is_absolute() or not attempt_directory.is_absolute(): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime paths must be absolute") + if plan_path.name != "resolved-plan.json": + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + try: + with open_verified_directory(plan_path.parent, require_private=True) as descriptor: + content = read_regular_text( + descriptor, + plan_path.name, + plan_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + plan = ResolvedSlurmRunPlan.model_validate_json(content) + except (OSError, UnicodeError, ValueError, ValidationError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan is unavailable or invalid" + ) from error + run_id = plan.run_id + logical_workspace_root = plan.selected_profile.profile.workspace_root + logical_plan_path = Path(logical_workspace_root) / "runs" / run_id / plan_path.name + if logical_plan_path != plan_path: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + return SlurmStateWriter(logical_workspace_root, run_id) + + +def _select_shard(shards: tuple[PlannedShard, ...], task_id: int) -> PlannedShard: + selected = tuple(shard for shard in shards if shard.array_task_index == task_id) + if len(selected) != 1: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "scheduler array task does not identify exactly one planned shard", + ) + return selected[0] + + +def _scheduler_task_id(value: str | None) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "SLURM_ARRAY_TASK_ID must be a non-negative integer", + ) + return int(value) + + +__all__ = ["load_allocation_context"] 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 new file mode 100644 index 000000000..91330e39e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -0,0 +1,693 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One-node allocation orchestration with one process and cleanup owner.""" + +from __future__ import annotations + +import subprocess +from collections.abc import Mapping +from contextlib import AbstractContextManager +from dataclasses import dataclass +from datetime import datetime, timedelta +from pathlib import Path +from typing import Literal, Protocol + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.logs import bind_execution_logs, execution_log_directory +from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStep +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.preflight import AllocationPreflight +from data_designer.slurm.runtime.probes import ReadinessProber +from data_designer.slurm.runtime.records import load_nonempty_client_candidate +from data_designer.slurm.runtime.steps import ( + ClientStepBuilder, + build_endpoint_steps, + build_vllm_steps, +) +from data_designer.slurm.runtime.supervisor import ManagedStep, RuntimeClock, StepSupervisor +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + ShardWinner, + StateNotFoundError, +) + +_READINESS_PROGRESS = { + ReadinessState.PENDING: 0, + ReadinessState.RESTARTING: 0, + ReadinessState.STARTING: 1, + ReadinessState.READY: 2, +} + + +class RuntimeStateStore(Protocol): + """State-writer operations required by one allocation.""" + + def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: + """Persist a monotonic attempt update.""" + ... + + def load_attempt(self, shard_id: str, attempt_id: str) -> AttemptManifest: + """Load one persisted attempt.""" + ... + + def publish_attempt_result( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + """Publish one complete result pair and bind it to its running attempt.""" + ... + + def acquire_dataset_workspace( + self, + shard_id: str, + attempt_id: str, + resume_mode: Literal["never", "always", "if_possible"], + ) -> AbstractContextManager[Path]: + """Hold exclusive ownership of the dataset workspace for generation.""" + ... + + def finalize_winner( + self, + shard_id: str, + attempt_id: str, + *, + published_at: datetime, + completed_at: datetime | None = None, + ) -> ShardWinner: + """Commit attempt success and publish its immutable winning candidate.""" + ... + + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: + """Persist the exact next readiness revision.""" + ... + + def load_readiness(self, shard_id: str, attempt_id: str) -> AttemptReadiness: + """Load the latest readiness revision for restart recovery.""" + ... + + +@dataclass(slots=True) +class _DeploymentStatus: + deployment: ResolvedVllmServerDeployment + state: ReadinessState = ReadinessState.PENDING + ready_backends: int = 0 + endpoint_publication: EndpointPublicationState = EndpointPublicationState.PENDING + last_probe: ProbeEvidence | None = None + + +@dataclass(slots=True) +class _RunOutcome: + candidate_reference: ArtifactReference | None = None + client_completed_at: datetime | None = None + failure: SlurmRuntimeError | None = None + failure_cause: BaseException | None = None + + +@dataclass(frozen=True, slots=True) +class _RuntimeTopology: + deployments: tuple[ResolvedVllmServerDeployment, ...] + endpoint_steps: tuple[tuple[RuntimeStep, RuntimeEndpoint], ...] + endpoints: tuple[RuntimeEndpoint, ...] + + +class OneNodeAllocationController: + """Execute one planned shard attempt and finalize nonempty candidates.""" + + def __init__( + self, + context: AllocationContext, + *, + runtime_proxy_path: Path, + runtime_root: Path | None = None, + state: RuntimeStateStore, + supervisor: StepSupervisor, + preflight: AllocationPreflight, + client_steps: ClientStepBuilder, + prober: ReadinessProber, + clock: RuntimeClock, + environment: Mapping[str, str], + poll_interval_seconds: float = 0.5, + probe_timeout_seconds: float = 1.0, + ) -> None: + if poll_interval_seconds <= 0 or probe_timeout_seconds <= 0: + raise ValueError("runtime probe intervals must be positive") + if any(deployment.authored.resources.shared_node is not None for deployment in context.plan.deployments): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "shared_node requires the coordinated allocation runtime", + ) + self._context = context + self._runtime_proxy_path = runtime_proxy_path + self._runtime_root = runtime_root if runtime_root is not None else runtime_proxy_path.parent + self._state = state + self._supervisor = supervisor + self._preflight = preflight + self._client_steps = client_steps + self._prober = prober + self._clock = clock + self._environment = environment + self._poll_interval_seconds = poll_interval_seconds + self._probe_timeout_seconds = probe_timeout_seconds + self._attempt = context.attempt + self._readiness: AttemptReadiness | None = None + self._statuses: list[_DeploymentStatus] = [] + self._log_directory: Path | None = None + + def run(self) -> AttemptManifest: + """Run the allocation and persist one terminal attempt on every owned path.""" + outcome = self._capture_execution() + self._record_outcome_failure(outcome) + self._cleanup_runtime(outcome) + self._record_stopped_readiness(outcome) + if outcome.failure is None: + try: + terminal = self._finalize_success(outcome) + except BaseException as error: + outcome.failure = _normalize_failure(error) + outcome.failure_cause = error + terminal = self._persist_terminal_outcome(outcome) + else: + terminal = self._persist_terminal_outcome(outcome) + if outcome.failure is not None: + if outcome.failure_cause is outcome.failure: + raise outcome.failure + raise outcome.failure from outcome.failure_cause + return terminal + + def _finalize_success(self, outcome: _RunOutcome) -> AttemptManifest: + if outcome.candidate_reference is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", + ) + completed_at = self._now() + if outcome.client_completed_at is not None and outcome.client_completed_at > completed_at: + completed_at = outcome.client_completed_at + self._state.finalize_winner( + self._attempt.shard_id, + self._attempt.attempt_id, + completed_at=completed_at, + published_at=max(self._now(), completed_at), + ) + self._attempt = self._state.load_attempt(self._attempt.shard_id, self._attempt.attempt_id) + return self._attempt + + def _capture_execution(self) -> _RunOutcome: + try: + candidate_reference, completed_at = self._execute() + return _RunOutcome(candidate_reference=candidate_reference, client_completed_at=completed_at) + except BaseException as error: + return _RunOutcome(failure=_normalize_failure(error), failure_cause=error) + + def _record_outcome_failure(self, outcome: _RunOutcome) -> None: + if outcome.failure is not None: + outcome.failure = self._record_failed_readiness(outcome.failure) + + def _cleanup_runtime(self, outcome: _RunOutcome) -> None: + try: + self._supervisor.cleanup() + except BaseException as error: + if outcome.failure is None: + outcome.failure = _normalize_failure(error) + outcome.failure_cause = error + outcome.failure = self._record_failed_readiness(outcome.failure) + else: + outcome.failure.add_note(f"cleanup also failed: {type(error).__name__}") + + def _record_stopped_readiness(self, outcome: _RunOutcome) -> None: + if self._readiness is None or not self._supervisor.cleanup_complete: + return + try: + self._publish_stopped_readiness() + except BaseException as error: + if outcome.failure is None: + outcome.failure = _normalize_failure(error) + outcome.failure_cause = error + else: + outcome.failure.add_note(f"stopped readiness could not be persisted: {type(error).__name__}") + + def _persist_terminal_outcome(self, outcome: _RunOutcome) -> AttemptManifest: + try: + return self._publish_terminal_attempt( + failure=outcome.failure, + candidate_reference=outcome.candidate_reference, + client_completed_at=outcome.client_completed_at, + ) + except BaseException as error: + terminal_failure = _normalize_failure(error) + if outcome.failure is None: + outcome.failure = terminal_failure + outcome.failure_cause = error + try: + return self._publish_terminal_attempt( + failure=outcome.failure, + candidate_reference=None, + client_completed_at=None, + ) + except BaseException as retry_error: + outcome.failure.add_note( + f"failed terminal attempt could not be persisted: {type(retry_error).__name__}" + ) + return self._attempt + outcome.failure.add_note(f"terminal attempt could not be persisted: {type(error).__name__}") + return self._attempt + + def _execute(self) -> tuple[ArtifactReference, datetime]: + topology = self._prepare_runtime() + required_processes = self._start_runtime(topology) + return self._run_client(topology.endpoints, required_processes) + + def _prepare_runtime(self) -> _RuntimeTopology: + self._validate_attempt_state() + deployments = resolve_allocation_deployments(self._context, self._environment) + restarting = self._readiness is not None + if restarting: + self._begin_execution(deployments, ReadinessState.RESTARTING) + self._preflight.verify(self._context, self._environment) + self._mark_attempt_running() + if not restarting: + self._begin_execution(deployments, ReadinessState.PENDING) + endpoint_steps = tuple( + (self._bind_logs(step), endpoint) + for step, endpoint in build_endpoint_steps( + deployments, + self._context.plan, + self._context.attempt_directory, + self._environment, + self._runtime_proxy_path, + self._runtime_root, + ) + ) + endpoints = tuple(endpoint for _, endpoint in endpoint_steps) + client_preflight = self._bind_logs( + self._client_steps.build_preflight_step( + self._context.plan, + self._context.shard, + self._attempt, + self._context.attempt_directory, + endpoints, + self._environment, + retry_resume_mode=self._retry_resume_mode, + ) + ) + self._supervisor.wait(self._supervisor.start(client_preflight)) + return _RuntimeTopology(deployments, endpoint_steps, endpoints) + + def _begin_execution( + self, + deployments: tuple[ResolvedVllmServerDeployment, ...], + state: ReadinessState, + ) -> None: + self._statuses = [_DeploymentStatus(deployment, state=state) for deployment in deployments] + self._publish_readiness(state) + if self._readiness is None: # pragma: no cover - persistence returns the published snapshot + raise AssertionError("runtime readiness was not published") + self._log_directory = execution_log_directory( + self._context.attempt_directory, + self._readiness.revision, + ) + + def _start_runtime(self, topology: _RuntimeTopology) -> tuple[ManagedStep, ...]: + for status in self._statuses: + if status.state in {ReadinessState.PENDING, ReadinessState.RESTARTING}: + status.state = ReadinessState.STARTING + self._publish_readiness(ReadinessState.STARTING) + server_processes = self._start_servers(topology.deployments) + for status, required in zip(self._statuses, server_processes, strict=True): + self._wait_for_backends(status, required) + self._publish_readiness(ReadinessState.STARTING) + endpoint_processes = tuple(self._supervisor.start(step) for step, _ in topology.endpoint_steps) + required_processes = tuple(process for group in server_processes for process in group) + endpoint_processes + self._wait_for_endpoints(topology.endpoints, required_processes) + for status in self._statuses: + status.state = ReadinessState.READY + status.endpoint_publication = EndpointPublicationState.PUBLISHED + status.last_probe = _probe_evidence(self._now(), ProbeOutcome.SUCCESS, "endpoint_ready", "endpoint ready") + self._publish_readiness(ReadinessState.READY) + return required_processes + + def _run_client( + self, + endpoints: tuple[RuntimeEndpoint, ...], + required_processes: tuple[ManagedStep, ...], + ) -> tuple[ArtifactReference, datetime]: + generation = self._bind_logs( + self._client_steps.build_generation_step( + self._context.plan, + self._context.shard, + self._attempt, + self._context.attempt_directory, + endpoints, + self._environment, + retry_resume_mode=self._retry_resume_mode, + ) + ) + workspace_mode = self._retry_resume_mode + if workspace_mode is None: + workspace_mode = "never" if self._context.plan.invocation.authored.resume == "never" else "always" + with self._state.acquire_dataset_workspace( + self._attempt.shard_id, + self._attempt.attempt_id, + workspace_mode, + ): + generation_started_at = self._now() + self._supervisor.wait(self._supervisor.start(generation), required=required_processes) + self._supervisor.require_running(required_processes) + client_result, candidate = load_nonempty_client_candidate(self._context, self._attempt) + if self._retry_resume_mode is not None and client_result.effective_resume_mode != self._retry_resume_mode: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client effective resume mode differs from the persisted retry plan", + ) + self._validate_client_timestamps(candidate.created_at, client_result.completed_at, generation_started_at) + candidate_reference = client_result.candidate_output_manifest + if candidate_reference is None: # pragma: no cover - nonempty results require this reference + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "nonempty client result has no candidate reference", + ) + self._state.publish_attempt_result(client_result, candidate) + self._attempt = _copy_attempt(self._attempt, candidate_output=candidate_reference) + return candidate_reference, client_result.completed_at + + @property + def _retry_resume_mode(self) -> Literal["never", "always"] | None: + retry_plan = self._context.retry_plan + return None if retry_plan is None else retry_plan.effective_resume_mode + + def _validate_client_timestamps( + self, + candidate_created_at: datetime, + client_completed_at: datetime, + generation_started_at: datetime, + ) -> None: + if candidate_created_at < generation_started_at or client_completed_at < generation_started_at: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client result predates the current generation step", + ) + if client_completed_at > self._now(): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client completion timestamp is later than the allocation clock", + ) + + def _validate_attempt_state(self) -> None: + if self._attempt.state not in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + }: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt is not executable", + ) + if self._attempt.state is AttemptLifecycleState.RUNNING: + self._load_restart_readiness() + + def _load_restart_readiness(self) -> None: + try: + readiness = self._state.load_readiness(self._attempt.shard_id, self._attempt.attempt_id) + except StateNotFoundError: + return + if readiness.state in {ReadinessState.FAILED, ReadinessState.STOPPED}: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt has terminal readiness and cannot be restarted", + ) + self._readiness = readiness + + def _mark_attempt_running(self) -> None: + if self._attempt.state is AttemptLifecycleState.RUNNING: + return + self._attempt = _copy_attempt( + self._attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=self._now(), + ) + self._attempt = self._state.update_attempt(self._attempt) + + def _start_servers( + self, + deployments: tuple[ResolvedVllmServerDeployment, ...], + ) -> tuple[tuple[ManagedStep, ...], ...]: + all_managed: list[tuple[ManagedStep, ...]] = [] + for deployment in deployments: + started_at = self._clock.monotonic() + managed: list[ManagedStep] = [] + steps = tuple( + self._bind_logs(step) + for step in build_vllm_steps( + deployment, + self._context.plan, + self._context.attempt_directory, + self._environment, + self._runtime_root, + ) + ) + for process, step in zip(deployment.processes, steps, strict=True): + self._supervisor.require_running(tuple(managed)) + remaining_delay = process.launch_delay_seconds - (self._clock.monotonic() - started_at) + if remaining_delay > 0: + self._clock.sleep(remaining_delay) + managed.append(self._supervisor.start(step)) + all_managed.append(tuple(managed)) + return tuple(all_managed) + + def _wait_for_backends(self, status: _DeploymentStatus, required: tuple[ManagedStep, ...]) -> None: + probes = status.deployment.readiness_probes + deadline = self._clock.monotonic() + max(probe.deadline_seconds for probe in probes) + previous_ready = -1 + while True: + self._supervisor.require_running(required) + ready = sum( + self._prober.is_ready( + "127.0.0.1", + probe.port, + probe.path, + timeout_seconds=self._probe_timeout_seconds, + ) + for probe in probes + ) + if ready != previous_ready: + self._record_backend_probe(status, ready, len(probes)) + previous_ready = ready + if ready == len(probes): + return + if self._clock.monotonic() >= deadline: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.READINESS_TIMEOUT, + f"deployment {status.deployment.deployment_id!r} readiness timed out", + ) + self._clock.sleep(self._poll_interval_seconds) + + def _record_backend_probe(self, status: _DeploymentStatus, ready: int, expected: int) -> None: + if status.state is ReadinessState.READY and ready != expected: + return + status.ready_backends = ready + status.last_probe = _probe_evidence( + self._now(), + ProbeOutcome.SUCCESS if ready == expected else ProbeOutcome.FAILURE, + "backend_ready" if ready == expected else "backend_starting", + f"{ready} of {expected} backends ready", + ) + self._publish_readiness(ReadinessState.STARTING) + + def _wait_for_endpoints( + self, + endpoints: tuple[RuntimeEndpoint, ...], + required: tuple[ManagedStep, ...], + ) -> None: + timeout = max(status.deployment.launch_policy.startup_timeout_seconds for status in self._statuses) + deadline = self._clock.monotonic() + timeout + while True: + self._supervisor.require_running(required) + if all( + self._prober.is_ready( + endpoint.host, + endpoint.port, + "/health", + timeout_seconds=self._probe_timeout_seconds, + ) + for endpoint in endpoints + ): + return + if self._clock.monotonic() >= deadline: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.READINESS_TIMEOUT, + "logical endpoint readiness timed out", + ) + self._clock.sleep(self._poll_interval_seconds) + + 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=timestamp, + started_at=self._attempt.updated_at if self._readiness is None else self._readiness.started_at, + state=state, + deployments=tuple( + DeploymentReadiness( + deployment_id=status.deployment.deployment_id, + model_alias=status.deployment.model_alias, + state=status.state, + expected_backends=len(status.deployment.backend_endpoints), + ready_backends=status.ready_backends, + endpoint_publication=status.endpoint_publication, + last_probe=status.last_probe, + ) + for status in self._statuses + ), + ) + self._readiness = self._state.write_readiness(readiness) + + def _record_failed_readiness(self, failure: SlurmRuntimeError) -> SlurmRuntimeError: + if self._readiness is None: + return failure + try: + for status in self._statuses: + status.state = ReadinessState.FAILED + status.ready_backends = 0 + if status.endpoint_publication is EndpointPublicationState.PENDING: + status.endpoint_publication = EndpointPublicationState.FAILED + status.last_probe = _probe_evidence( + self._now(), + ProbeOutcome.FAILURE, + failure.code.value, + "allocation runtime failed", + ) + self._publish_readiness(ReadinessState.FAILED) + except BaseException as error: + failure.add_note(f"failed readiness could not be persisted: {type(error).__name__}") + return failure + + def _publish_stopped_readiness(self) -> None: + for status in self._statuses: + status.state = ReadinessState.STOPPED + status.ready_backends = 0 + self._publish_readiness(ReadinessState.STOPPED) + + def _publish_terminal_attempt( + self, + *, + failure: SlurmRuntimeError | None, + candidate_reference: ArtifactReference | None, + client_completed_at: datetime | None, + ) -> AttemptManifest: + timestamp = self._now() + if client_completed_at is not None and client_completed_at > timestamp: + timestamp = client_completed_at + terminal = self._build_terminal_attempt(failure, candidate_reference, timestamp) + self._attempt = self._state.update_attempt(terminal) + return self._attempt + + def _build_terminal_attempt( + self, + failure: SlurmRuntimeError | None, + candidate_reference: ArtifactReference | None, + timestamp: datetime, + ) -> AttemptManifest: + if failure is None: + if candidate_reference is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", + ) + terminal = _copy_attempt( + self._attempt, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + candidate_output=candidate_reference, + updated_at=timestamp, + ) + else: + terminal = _copy_attempt( + self._attempt, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.FAILED, + updated_at=timestamp, + ) + return terminal + + def _now(self) -> datetime: + value = self._clock.now() + if value.tzinfo is None or value.utcoffset() != timedelta(0): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock is not UTC") + minimum = self._attempt.updated_at + if self._readiness is not None and self._readiness.updated_at > minimum: + minimum = self._readiness.updated_at + if value < minimum: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock moved backward") + return value + + def _bind_logs(self, step: RuntimeStep) -> RuntimeStep: + if self._log_directory is None: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime log directory is unavailable") + return bind_execution_logs(step, self._log_directory) + + +def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: + payload = attempt.model_dump(mode="python") + payload.update(updates) + return AttemptManifest.model_validate(payload) + + +def _probe_evidence( + observed_at: datetime, + outcome: ProbeOutcome, + reason_code: str, + message: str, +) -> ProbeEvidence: + return ProbeEvidence( + observed_at=observed_at, + outcome=outcome, + reason_code=reason_code, + redacted_message=message, + ) + + +def _can_advance_readiness(previous: ReadinessState, current: ReadinessState) -> bool: + if current is ReadinessState.RESTARTING: + return previous not in {ReadinessState.FAILED, ReadinessState.STOPPED} + previous_progress = _READINESS_PROGRESS.get(previous) + current_progress = _READINESS_PROGRESS.get(current) + if previous_progress is None or current_progress is None: + return True + return current_progress >= previous_progress + + +def _normalize_failure(error: BaseException) -> SlurmRuntimeError: + if isinstance(error, SlurmRuntimeError): + return error + if isinstance(error, (OSError, subprocess.SubprocessError)): + return SlurmRuntimeError(SlurmRuntimeErrorCode.STEP_FAILED, "allocation process operation failed") + if isinstance(error, KeyboardInterrupt): + return SlurmRuntimeError(SlurmRuntimeErrorCode.CLIENT_FAILED, "allocation execution was interrupted") + return SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "allocation runtime failed at an internal boundary", + ) + + +__all__ = ["OneNodeAllocationController", "RuntimeStateStore"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py new file mode 100644 index 000000000..6d55c5893 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose resolved deployments into node-worker specifications.""" + +from __future__ import annotations + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.node_spec import NodeProcessSpec, NodeReadinessSpec, NodeSpec, NodeWorkerSpec +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.steps import build_vllm_command +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess + + +def build_node_worker_spec( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + layout: AllocationLayout, + *, + colocated_deployments: tuple[ResolvedVllmServerDeployment, ...] = (), +) -> NodeWorkerSpec: + """Build the validated work assigned to each node in one deployment.""" + model = _resolve_model(deployment, plan) + deployments = (deployment, *colocated_deployments) + assigned_gpus = [ + (process.node_index, gpu) + for member in deployments + for process in member.processes + for gpu in process.gpu_indices + ] + sharing = len(set(assigned_gpus)) != len(assigned_gpus) + nodes = tuple( + NodeSpec( + node_index=node_index, + host=layout.get_host(node_index), + ports=tuple(sorted(port for member in deployments for port in _get_node_ports(member, node_index))), + processes=tuple( + sorted( + ( + NodeProcessSpec( + process_id=process.process_id, + command=( + build_vllm_command(member, process, plan) + if len(layout.node_hosts) == 1 + else build_vllm_process_command(member, process, plan, layout) + ), + gpu_indices=tuple(process.gpu_indices), + launch_delay_seconds=process.launch_delay_seconds, + readiness=NodeReadinessSpec( + port=member.backend_endpoints[process.deployment_replica_index].port, + path=member.launch_policy.readiness_path, + deadline_seconds=member.launch_policy.startup_timeout_seconds, + ) + if sharing + else None, + ) + for member in deployments + for process in member.processes + if process.node_index == node_index + ), + key=lambda process: process.launch_delay_seconds, + ) + ), + ) + for node_index in deployment.node_indices + ) + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=plan.resolved_gpus_per_node if colocated_deployments else deployment.gpus_per_node, + required_model_path=model if deployment.model.startswith("/") and not colocated_deployments else None, + nodes=nodes, + required_model_paths=tuple( + _resolve_model(member, plan) for member in deployments if member.model.startswith("/") + ) + if colocated_deployments + else (), + allow_gpu_sharing=sharing, + ) + + +def build_vllm_process_command( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + plan: ResolvedSlurmRunPlan, + layout: AllocationLayout, +) -> tuple[str, ...]: + """Build one shell-free vLLM lane command at its resolved host placement.""" + backend = deployment.backend_endpoints[process.deployment_replica_index] + model = _resolve_model(deployment, plan) + command: tuple[str, ...] = ( + deployment.executable_path, + "serve", + model, + "--served-model-name", + deployment.served_model_name, + "--host", + "0.0.0.0", + "--port", + str(backend.port), + "--tensor-parallel-size", + str(process.tensor_parallel), + "--distributed-executor-backend", + "uni" if process.tensor_parallel * process.pipeline_parallel == 1 else "mp", + "--data-parallel-backend", + "mp", + "--middleware", + "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware", + *_distributed_arguments(process, layout), + ) + if deployment.launch_policy.enable_expert_parallel: + command += ("--enable-expert-parallel",) + if deployment.launch_policy.io_processor_plugin is not None: + command += ("--io-processor-plugin", deployment.launch_policy.io_processor_plugin) + return command + deployment.launch_policy.extra_args + + +def _resolve_model(deployment: ResolvedVllmServerDeployment, plan: ResolvedSlurmRunPlan) -> str: + return get_container_path(plan, deployment.model) if deployment.model.startswith("/") else deployment.model + + +def _distributed_arguments( + process: ResolvedVllmProcess, + layout: AllocationLayout, +) -> tuple[str, ...]: + if process.pipeline_parallel == 1: + return () + rendezvous = process.rendezvous + if rendezvous is None: # pragma: no cover - resolved contracts enforce this + raise AssertionError("distributed process has no rendezvous") + arguments = ( + "--pipeline-parallel-size", + str(process.pipeline_parallel), + "--nnodes", + str(process.pipeline_parallel), + "--node-rank", + str(process.pipeline_rank), + "--master-addr", + layout.get_host(rendezvous.master_node_index), + "--master-port", + str(rendezvous.port), + "--distributed-timeout-seconds", + str(rendezvous.timeout_seconds), + ) + return arguments + (("--headless",) if process.pipeline_rank > 0 else ()) + + +def _get_node_ports(deployment: ResolvedVllmServerDeployment, node_index: int) -> tuple[int, ...]: + http_ports = tuple(endpoint.port for endpoint in deployment.backend_endpoints if endpoint.node_index == node_index) + rendezvous_ports = tuple( + process.rendezvous.port + for process in deployment.processes + if process.node_index == node_index and process.pipeline_rank == 0 and process.rendezvous is not None + ) + return tuple(sorted((*http_ports, *rendezvous_ports))) + + +__all__ = ["build_node_worker_spec", "build_vllm_process_command"] 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 new file mode 100644 index 000000000..ce4897d65 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed control phases executed with the resolved native client runtime.""" + +from __future__ import annotations + +import argparse +import os +import sys +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone +from pathlib import Path + +from data_designer.slurm.client.filesystem import ensure_private_directory, replace_private_text +from data_designer.slurm.client.process import ClientWorkerProcess +from data_designer.slurm.runtime.bootstrap import build_runtime_manifest +from data_designer.slurm.runtime.context import load_allocation_context +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.logs import execution_log_directory +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.preflight import ( + AllocationLayout, + SystemAllocationPreflight, + validate_allocation_layout, +) +from data_designer.slurm.runtime.records import load_nonempty_client_candidate +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + SlurmStateWriter, + StateNotFoundError, +) + + +def main(arguments: Sequence[str] | None = None) -> int: + """Execute one native allocation control phase.""" + parsed = _parse_arguments(arguments) + try: + if parsed.operation == "prepare": + _prepare(parsed, os.environ) + elif parsed.operation == "ready": + _ready(parsed, os.environ) + elif parsed.operation == "client": + _client(parsed, os.environ) + elif parsed.operation == "succeed": + _succeed(parsed, os.environ) + else: + _fail(parsed, os.environ) + except SlurmRuntimeError as error: + print(f"allocation runtime failed ({error.code.value}): {error}", file=sys.stderr) + return ( + 64 if error.code in {SlurmRuntimeErrorCode.INVALID_CONTEXT, SlurmRuntimeErrorCode.PREFLIGHT_FAILED} else 70 + ) + except KeyboardInterrupt: + print("allocation runtime interrupted", file=sys.stderr) + return 130 + except Exception: + print("allocation runtime failed at an internal boundary", file=sys.stderr) + return 70 + return 0 + + +def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(prog="data-designer-slurm-runtime") + subparsers = parser.add_subparsers(dest="operation", required=True) + for operation in ("prepare", "ready", "succeed", "fail"): + subparser = subparsers.add_parser(operation) + _add_context_arguments(subparser) + prepare = subparsers.choices["prepare"] + prepare.add_argument("--runtime-root", required=True, type=Path) + prepare.add_argument("--manifest", required=True, type=Path) + prepare.add_argument("--node-host", action="append", required=True) + client = subparsers.add_parser("client") + _add_context_arguments(client) + client.add_argument("--endpoint", action="append", default=[]) + return parser.parse_args(arguments) + + +def _add_context_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--plan", required=True, type=Path) + parser.add_argument("--attempt-dir", required=True, type=Path) + parser.add_argument("--retry-id") + parser.add_argument("--retry-plan-sha256") + parser.add_argument("--effective-resume-mode", choices=("never", "always")) + + +def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = _load_context(arguments, environment) + _validate_attempt_is_executable(context.attempt) + layout = AllocationLayout(tuple(arguments.node_host)) + validate_allocation_layout(context.plan, layout) + SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) + SystemAllocationPreflight.verify_client_runtime(context) + SystemAllocationPreflight.verify_artifacts(context) + SystemAllocationPreflight.verify_ports(context, environment) + readiness = _begin_attempt(context, writer, environment) + log_directory = execution_log_directory(context.attempt_directory, readiness.revision) + ensure_private_directory(log_directory) + manifest = build_runtime_manifest( + context, + environment, + runtime_root=arguments.runtime_root, + log_directory=log_directory, + layout=layout, + ) + expected_manifest = context.attempt_directory / "runtime-manifest.json" + if arguments.manifest != expected_manifest: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime manifest path is invalid") + replace_private_text(arguments.manifest, manifest.serialize_json()) + + +def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = _load_context(arguments, environment) + previous = writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + timestamp = _now(context.attempt, previous) + deployments = _resolve_deployments(context, environment) + writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + revision=previous.revision + 1, + updated_at=timestamp, + started_at=previous.started_at, + state=ReadinessState.READY, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.model_alias, + state=ReadinessState.READY, + expected_backends=len(deployment.backend_endpoints), + ready_backends=len(deployment.backend_endpoints), + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=_probe(timestamp, ProbeOutcome.SUCCESS, "endpoint_ready", "endpoint ready"), + ) + for deployment in deployments + ), + ) + ) + + +def _client( + arguments: argparse.Namespace, + environment: Mapping[str, str], + *, + client_worker: ClientWorkerProcess | None = None, +) -> None: + context, writer = _load_context(arguments, environment) + worker_process = client_worker or ClientWorkerProcess() + generation_started_at = _now(context.attempt, _load_optional_readiness(context, writer)) + resume_mode = ( + context.plan.invocation.authored.resume + if context.retry_plan is None + else context.retry_plan.effective_resume_mode + ) + with writer.acquire_dataset_workspace( + context.shard.shard_id, + context.attempt.attempt_id, + resume_mode, + ): + return_code = worker_process.run( + ( + "run", + "--plan", + arguments.plan.as_posix(), + "--shard-id", + context.shard.shard_id, + "--attempt-id", + context.attempt.attempt_id, + "--attempt-dir", + arguments.attempt_dir.as_posix(), + *(() if context.retry_plan is None else ("--resume-mode", context.retry_plan.effective_resume_mode)), + *(argument for endpoint in arguments.endpoint for argument in ("--endpoint", endpoint)), + ) + ) + if return_code != 0: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.CLIENT_FAILED, "client generation failed") + client_result, candidate = load_nonempty_client_candidate( + context, + context.attempt, + attempt_directory=arguments.attempt_dir, + ) + if ( + context.retry_plan is not None + and client_result.effective_resume_mode != context.retry_plan.effective_resume_mode + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client effective resume mode differs from the persisted retry plan", + ) + completed_at = client_result.completed_at + if candidate.created_at < generation_started_at or completed_at < generation_started_at: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client result predates the current generation step", + ) + if completed_at > datetime.now(timezone.utc): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client completion timestamp is later than the allocation clock", + ) + writer.publish_attempt_result(client_result, candidate) + + +def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = _load_context(arguments, environment) + stopped_at = _write_stopped_readiness(context, writer) + attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) + if attempt.candidate_output is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", + ) + completed_at = max(stopped_at, attempt.updated_at) + writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=max(datetime.now(timezone.utc), completed_at), + ) + + +def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = _load_context(arguments, environment) + attempt = writer.load_attempt(context.shard.shard_id, context.attempt.attempt_id) + if attempt.state in {AttemptLifecycleState.SUCCEEDED, AttemptLifecycleState.FAILED}: + return + timestamp = _write_failed_and_stopped_readiness(context, writer) + writer.update_attempt( + attempt.model_copy( + update={ + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.FAILED, + "updated_at": max(timestamp, attempt.updated_at), + } + ) + ) + + +def _load_context( + arguments: argparse.Namespace, + environment: Mapping[str, str], +) -> tuple[AllocationContext, SlurmStateWriter]: + return load_allocation_context( + arguments.plan, + arguments.attempt_dir, + environment, + retry_id=arguments.retry_id, + retry_plan_sha256=arguments.retry_plan_sha256, + effective_resume_mode=arguments.effective_resume_mode, + ) + + +def _begin_attempt( + context: AllocationContext, + writer: SlurmStateWriter, + environment: Mapping[str, str], +) -> AttemptReadiness: + attempt = context.attempt + previous = _load_optional_readiness(context, writer) + timestamp = _now(attempt, previous) + if attempt.state is not AttemptLifecycleState.RUNNING: + attempt = writer.update_attempt( + attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING, "updated_at": timestamp}) + ) + deployments = _resolve_deployments(context, environment) + initial_state = ReadinessState.RESTARTING if previous is not None else ReadinessState.PENDING + initial = writer.write_readiness(_readiness(context, deployments, previous, initial_state, timestamp)) + return writer.write_readiness( + _readiness(context, deployments, initial, ReadinessState.STARTING, _now(attempt, initial)) + ) + + +def _readiness( + context: AllocationContext, + deployments: tuple[ResolvedVllmServerDeployment, ...], + previous: AttemptReadiness | None, + state: ReadinessState, + timestamp: datetime, +) -> AttemptReadiness: + return AttemptReadiness( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + 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( + deployment_id=deployment.deployment_id, + model_alias=deployment.model_alias, + state=state, + expected_backends=len(deployment.backend_endpoints), + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + for deployment in deployments + ), + ) + + +def _write_failed_and_stopped_readiness(context: AllocationContext, writer: SlurmStateWriter) -> datetime: + previous = _load_optional_readiness(context, writer) + if previous is None: + return max(datetime.now(timezone.utc), context.attempt.updated_at) + failed_at = _now(context.attempt, previous) + failed = writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=previous.run_id, + shard_id=previous.shard_id, + 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( + update={ + "state": ReadinessState.FAILED, + "ready_backends": 0, + "endpoint_publication": ( + EndpointPublicationState.FAILED + if deployment.endpoint_publication is EndpointPublicationState.PENDING + else deployment.endpoint_publication + ), + "last_probe": _probe( + failed_at, + ProbeOutcome.FAILURE, + "allocation_failed", + "allocation runtime failed", + ), + } + ) + for deployment in previous.deployments + ), + ) + ) + return _write_stopped_readiness(context, writer, previous=failed) + + +def _write_stopped_readiness( + context: AllocationContext, + writer: SlurmStateWriter, + *, + previous: AttemptReadiness | None = None, +) -> datetime: + previous = previous or writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + timestamp = _now(context.attempt, previous) + writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=previous.run_id, + shard_id=previous.shard_id, + 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, + } + ) + for deployment in previous.deployments + ), + ) + ) + return timestamp + + +def _load_optional_readiness(context: AllocationContext, writer: SlurmStateWriter) -> AttemptReadiness | None: + try: + return writer.load_readiness(context.shard.shard_id, context.attempt.attempt_id) + except StateNotFoundError: + return None + + +def _resolve_deployments( + context: AllocationContext, + environment: Mapping[str, str], +) -> tuple[ResolvedVllmServerDeployment, ...]: + return resolve_allocation_deployments(context, environment) + + +def _validate_attempt_is_executable(attempt: AttemptManifest) -> None: + if attempt.state not in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + }: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation attempt is not executable") + + +def _probe(observed_at: datetime, outcome: ProbeOutcome, reason_code: str, message: str) -> ProbeEvidence: + return ProbeEvidence( + observed_at=observed_at, + outcome=outcome, + reason_code=reason_code, + redacted_message=message, + ) + + +def _now(attempt: AttemptManifest, readiness: AttemptReadiness | None) -> datetime: + value = datetime.now(timezone.utc) + minimum = attempt.updated_at + if readiness is not None and readiness.updated_at > minimum: + minimum = readiness.updated_at + if value < minimum: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock moved backward") + if value.tzinfo is None or value.utcoffset() != timedelta(0): # pragma: no cover - system clock is UTC + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime clock is not UTC") + return value + + +if __name__ == "__main__": # pragma: no cover - exercised through the installed module + raise SystemExit(main()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh new file mode 100644 index 000000000..edac5241e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -Eeuo pipefail + +readonly DD_RUNTIME_DIR=$(cd -- "${BASH_SOURCE[0]%/*}" && pwd -P) +source "${DD_RUNTIME_DIR}/scratch.sh" +source "${DD_RUNTIME_DIR}/plan_reader.sh" +source "${DD_RUNTIME_DIR}/step_runner.sh" +source "${DD_RUNTIME_DIR}/cleanup.sh" + +DD_RUNTIME_PREPARED=0 +DD_RUNTIME_FINALIZED=0 +DD_RUNTIME_PHASE=initialization +DD_RUNTIME_RETRY_ARGUMENTS=() +declare -A DD_STEP_STARTED_AT=() + +dd_runtime_log() { + printf '[%s] %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$1" +} + +dd_slurm_run_allocation() { + if [[ $# -ne 2 && $# -ne 5 ]]; then + printf '%s\n' 'allocation runtime requires plan and attempt with optional retry binding' >&2 + return 64 + fi + DD_PLAN_PATH=$1 + DD_ATTEMPT_PATH=$2 + DD_RUNTIME_MANIFEST=${DD_ATTEMPT_PATH}/runtime-manifest.json + DD_RUNTIME_RETRY_ARGUMENTS=() + DD_STEP_STARTED_AT=() + if [[ $# -eq 5 ]]; then + DD_RUNTIME_RETRY_ARGUMENTS=( + --retry-id "$3" + --retry-plan-sha256 "$4" + --effective-resume-mode "$5" + ) + fi + trap dd_runtime_exit EXIT + trap 'exit 130' INT TERM + DD_RUNTIME_PHASE=read_control_plan + dd_read_control_plan "${DD_PLAN_PATH}" + DD_RUNTIME_PHASE=read_plan_secrets + dd_read_plan_secret_names "${DD_PLAN_PATH}" + DD_RUNTIME_CONTAINER_ROOT=${DD_SCRATCH_CONTAINER_ROOT}/runtime + export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST + export DD_RUNTIME_ROOT DD_RUNTIME_CONTAINER_ROOT + + DD_RUNTIME_PHASE=verify_host_context + dd_verify_host_context + DD_RUNTIME_PHASE=start_runtime_timer + dd_start_runtime_timer + + DD_RUNTIME_PREPARED=1 + local host + local -a host_arguments=() + for host in "${DD_ALLOCATION_HOSTS[@]}"; do + host_arguments+=(--node-host "${host}") + done + DD_RUNTIME_PHASE=prepare + dd_run_bound_control_phase prepare \ + --runtime-root "${DD_RUNTIME_ROOT}" \ + --manifest "${DD_RUNTIME_MANIFEST}" \ + "${host_arguments[@]}" + DD_RUNTIME_PHASE=verify_runtime_manifest + dd_verify_runtime_manifest \ + "${DD_RUNTIME_MANIFEST}" \ + "${DD_PLAN_SHA256}" \ + "${DD_SHARD_ID}" \ + "attempt-${DD_ATTEMPT_ORDINAL}" + DD_RUNTIME_PHASE=require_plan_secrets + dd_require_plan_secrets + + DD_RUNTIME_PHASE=client_preflight + dd_runtime_log 'Running client preflight...' + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client_preflight + ((${#DD_STEP_IDS[@]} == 1)) + dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + + DD_RUNTIME_PHASE=server_preflight + dd_runtime_log 'Running serving preflight...' + dd_run_role_steps server_preflight + DD_RUNTIME_PHASE=start_servers + dd_runtime_log 'Starting vLLM servers...' + dd_start_servers + DD_RUNTIME_PHASE=server_readiness + dd_runtime_log 'Waiting for vLLM servers to become ready...' + dd_wait_for_role_readiness server + dd_runtime_log 'vLLM servers are ready.' + DD_RUNTIME_PHASE=start_endpoints + dd_start_endpoints + DD_RUNTIME_PHASE=endpoint_readiness + dd_wait_for_role_readiness endpoint + dd_require_running + DD_RUNTIME_PHASE=ready + dd_run_bound_control_phase ready + + DD_RUNTIME_PHASE=client + dd_runtime_log 'Starting Data Designer client...' + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client + ((${#DD_STEP_IDS[@]} == 1)) + dd_start_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + local client_pid=${DD_LAST_PID} + DD_MANAGED_PIDS+=("${client_pid}") + local client_log_command progress_command + printf -v client_log_command 'tail -F %q %q' "${DD_STEP_STDOUT}" "${DD_STEP_STDERR}" + printf -v progress_command 'cat %q' "${DD_ATTEMPT_PATH}/client-progress.json" + dd_runtime_log "Generation started. Follow client logs: ${client_log_command}" + dd_runtime_log "Inspect generation progress: ${progress_command}" + dd_wait_for_client "${client_pid}" + dd_runtime_log 'Generation finished; finalizing run...' + dd_require_running + + DD_RUNTIME_PHASE=cleanup + dd_cleanup_steps + DD_RUNTIME_PHASE=succeed + dd_run_bound_control_phase succeed + DD_RUNTIME_FINALIZED=1 + DD_RUNTIME_PHASE=completed + dd_runtime_log 'Run completed.' +} + +dd_run_bound_control_phase() { + local operation=$1 + shift + dd_run_control_phase \ + "${operation}" \ + "${DD_RUNTIME_RETRY_ARGUMENTS[@]+"${DD_RUNTIME_RETRY_ARGUMENTS[@]}"}" \ + "$@" +} + +dd_verify_host_context() { + local tool + for tool in bash sha256sum tar jq srun scontrol getent curl find; do + command -v "${tool}" >/dev/null || { + printf 'required allocation tool %q is unavailable\n' "${tool}" >&2 + return 69 + } + done + [[ ${DD_SCRATCH_ROOT:-} == /* && ${DD_RUNTIME_DIR} == "${DD_SCRATCH_ROOT}/runtime" ]] + [[ -d ${DD_SCRATCH_ROOT} && ! -L ${DD_SCRATCH_ROOT} && -O ${DD_SCRATCH_ROOT} ]] + [[ -d ${DD_RUNTIME_DIR} && ! -L ${DD_RUNTIME_DIR} && -O ${DD_RUNTIME_DIR} ]] + [[ -d ${DD_ATTEMPT_PATH} && ! -L ${DD_ATTEMPT_PATH} ]] + [[ -x ${DD_CLIENT_PYTHON} ]] + [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] + [[ ${DD_EXPECTED_NODES} =~ ^[1-9][0-9]*$ ]] + [[ ${DD_CLIENT_NODE_INDEX} =~ ^[0-9]+$ && ${DD_CLIENT_NODE_INDEX} -lt ${DD_EXPECTED_NODES} ]] + [[ ${SLURM_JOB_NUM_NODES:-} == "${DD_EXPECTED_NODES}" ]] + [[ ${SLURM_NODEID:-} == "${DD_CLIENT_NODE_INDEX}" ]] + dd_resolve_allocation_hosts + dd_verify_gpu_count +} + +dd_resolve_allocation_hosts() { + local node_list=${SLURM_JOB_NODELIST:-} + [[ -n ${node_list} && ${node_list} != -* && ${#node_list} -le 4096 ]] || return 65 + [[ ${node_list} != *[$'\t\r\n ']* ]] || return 65 + local expanded + expanded=$(scontrol show hostnames "${node_list}") || return 65 + DD_ALLOCATION_HOSTS=() + local host + while IFS= read -r host; do + DD_ALLOCATION_HOSTS+=("${host}") + done <<<"${expanded}" + ((${#DD_ALLOCATION_HOSTS[@]} == DD_EXPECTED_NODES)) || return 65 + for host in "${DD_ALLOCATION_HOSTS[@]}"; do + [[ ${host} =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ && ${#host} -le 253 ]] || return 65 + done + DD_CLIENT_HOST=${DD_ALLOCATION_HOSTS[DD_CLIENT_NODE_INDEX]} +} + +dd_verify_gpu_count() { + local visible=${CUDA_VISIBLE_DEVICES:-${SLURM_JOB_GPUS:-}} + local count=0 item + local -A devices=() + if [[ -z ${visible//[[:space:]]/} ]]; then + count=0 + elif [[ ${visible} =~ ^gpu(:[^:]+)?:([0-9]+)$ ]]; then + count=${BASH_REMATCH[2]} + else + IFS=, read -r -a device_values <<<"${visible}" + for item in "${device_values[@]+"${device_values[@]}"}"; do + item=${item//[[:space:]]/} + [[ ${item} =~ ^[A-Za-z0-9_.:-]+$ ]] || return 65 + [[ -n ${item} && ! ${devices[${item}]+_} ]] || return 65 + devices[${item}]=1 + ((count += 1)) + done + fi + [[ ${DD_EXPECTED_GPUS} =~ ^[1-9][0-9]*$ && ${count} -ge ${DD_EXPECTED_GPUS} ]] || { + printf '%s\n' 'allocation GPU visibility does not match the resolved plan' >&2 + return 65 + } +} + +dd_start_servers() { + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" server + local step_id deployment delay remaining tail_command + local -A deployment_started=() + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + deployment=${step_id%%-replica-*} + deployment_started[${deployment}]=${deployment_started[${deployment}]:-${SECONDS}} + delay=${DD_STEP_DELAY} + remaining=$((delay - (SECONDS - deployment_started[${deployment}]))) + ((remaining <= 0)) || dd_sleep "${remaining}" + DD_STEP_STARTED_AT[${step_id}]=${SECONDS} + dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + dd_register_required_pid "${DD_LAST_PID}" + dd_require_running + printf -v tail_command 'tail -F %q %q' "${DD_STEP_STDOUT}" "${DD_STEP_STDERR}" + dd_runtime_log "Serving step ${step_id} started. Follow vLLM logs: ${tail_command}" + done +} + +dd_run_role_steps() { + local role=$1 + local step_id + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" "${role}" + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_run_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + done +} + +dd_start_endpoints() { + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" endpoint + local step_id + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + DD_STEP_STARTED_AT[${step_id}]=${SECONDS} + dd_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + dd_register_required_pid "${DD_LAST_PID}" + dd_require_running + done +} + +dd_wait_for_role_readiness() { + local role=$1 + local step_id index started_at probe pending + local -a probe_steps=() probe_hosts=() probe_ports=() probe_paths=() probe_deadlines=() probe_ready=() + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" "${role}" + for step_id in "${DD_STEP_IDS[@]+"${DD_STEP_IDS[@]}"}"; do + dd_read_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + ((${#DD_STEP_PROBE_FIELDS[@]} % 4 == 0 && ${#DD_STEP_PROBE_FIELDS[@]} > 0)) || return 65 + started_at=${DD_STEP_STARTED_AT[${step_id}]:-${SECONDS}} + for ((index = 0; index < ${#DD_STEP_PROBE_FIELDS[@]}; index += 4)); do + probe_steps+=("${step_id}") + probe_hosts+=("${DD_STEP_PROBE_FIELDS[index]}") + probe_ports+=("${DD_STEP_PROBE_FIELDS[index + 1]}") + probe_paths+=("${DD_STEP_PROBE_FIELDS[index + 2]}") + probe_deadlines+=("$((started_at + DD_STEP_PROBE_FIELDS[index + 3]))") + probe_ready+=(0) + done + done + pending=${#probe_steps[@]} + while ((pending > 0)); do + dd_require_running + for ((probe = 0; probe < ${#probe_steps[@]}; probe++)); do + [[ ${probe_ready[probe]} == 0 ]] || continue + ((SECONDS < probe_deadlines[probe])) || { + printf 'runtime step %q readiness timed out\n' "${probe_steps[probe]}" >&2 + return 70 + } + if curl --fail --silent --max-time 1 \ + "http://${probe_hosts[probe]}:${probe_ports[probe]}${probe_paths[probe]}" >/dev/null 2>&1; then + ((SECONDS < probe_deadlines[probe])) || { + printf 'runtime step %q readiness timed out\n' "${probe_steps[probe]}" >&2 + return 70 + } + probe_ready[probe]=1 + pending=$((pending - 1)) + fi + done + ((pending == 0)) || dd_sleep 0.5 + done +} + +dd_runtime_exit() { + local status=$? + trap - EXIT INT TERM + set +e + if ((status != 0)); then + printf 'allocation runtime failed: phase=%s status=%d\n' "${DD_RUNTIME_PHASE}" "${status}" >&2 + fi + dd_cleanup_steps + if ((DD_RUNTIME_PREPARED == 1 && DD_RUNTIME_FINALIZED == 0)); then + dd_run_bound_control_phase fail >/dev/null + fi + dd_cleanup_allocation_scratch >/dev/null 2>&1 || true + dd_stop_runtime_timer + exit "${status}" +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/errors.py new file mode 100644 index 000000000..2a737c717 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/errors.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalized failures for allocation-local Slurm execution.""" + +from __future__ import annotations + +from enum import Enum + + +class SlurmRuntimeErrorCode(str, Enum): + """Stable allocation-runtime failure categories.""" + + INVALID_CONTEXT = "invalid_context" + PREFLIGHT_FAILED = "preflight_failed" + STEP_FAILED = "step_failed" + READINESS_TIMEOUT = "readiness_timeout" + CLIENT_FAILED = "client_failed" + FINALIZATION_FAILED = "finalization_failed" + CLEANUP_FAILED = "cleanup_failed" + + +class SlurmRuntimeError(RuntimeError): + """Bounded allocation-runtime error safe for persisted diagnostics.""" + + def __init__(self, code: SlurmRuntimeErrorCode, message: str) -> None: + if not isinstance(code, SlurmRuntimeErrorCode): + raise TypeError("code must be a SlurmRuntimeErrorCode") + if type(message) is not str or not message or len(message) > 512: + raise ValueError("runtime error message must contain 1 to 512 characters") + if any(ord(character) < 32 or ord(character) == 127 for character in message): + raise ValueError("runtime error message must not contain control characters") + self.code = code + super().__init__(message) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/logs.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/logs.py new file mode 100644 index 000000000..a1742142f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/logs.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execution-scoped runtime log paths and restrictive file creation.""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import replace +from pathlib import Path + +from data_designer.slurm.runtime.models import RuntimeStep +from data_designer.slurm.state.filesystem import ( + ensure_private_child_directory, + open_verified_child_directory, + open_verified_directory, +) + +_LOG_DIRECTORY_NAME = "logs" +_EXECUTION_DIRECTORY_PREFIX = "execution-" + + +def execution_log_directory(attempt_directory: Path, readiness_revision: int) -> Path: + """Return the package-owned log directory for one execution epoch.""" + if type(readiness_revision) is not int or readiness_revision < 1: + raise ValueError("execution readiness revision must be positive") + return attempt_directory / _LOG_DIRECTORY_NAME / f"{_EXECUTION_DIRECTORY_PREFIX}{readiness_revision:08d}" + + +def bind_execution_logs(step: RuntimeStep, log_directory: Path) -> RuntimeStep: + """Bind one immutable step specification to its execution-scoped logs.""" + return replace( + step, + stdout_path=log_directory / f"{step.step_id}.out", + stderr_path=log_directory / f"{step.step_id}.err", + ) + + +@contextmanager +def open_step_logs(step: RuntimeStep) -> Iterator[tuple[int, int]]: + """Create and open one step's restrictive logs below verified directories.""" + log_directory = step.stdout_path.parent + logs_root = log_directory.parent + attempt_directory = logs_root.parent + if logs_root.name != _LOG_DIRECTORY_NAME or not log_directory.name.startswith(_EXECUTION_DIRECTORY_PREFIX): + raise OSError("runtime logs must use a package-owned execution directory") + + with open_verified_directory(attempt_directory, require_private=True) as attempt_descriptor: + ensure_private_child_directory(attempt_descriptor, _LOG_DIRECTORY_NAME, logs_root) + with open_verified_directory(logs_root, require_private=True) as logs_descriptor: + ensure_private_child_directory(logs_descriptor, log_directory.name, log_directory) + with open_verified_child_directory(logs_descriptor, log_directory.name, log_directory) as directory_descriptor: + with ExitStack() as resources: + stdout_descriptor = _create_log_file(directory_descriptor, step.stdout_path.name) + resources.callback(os.close, stdout_descriptor) + stderr_descriptor = _create_log_file(directory_descriptor, step.stderr_path.name) + resources.callback(os.close, stderr_descriptor) + yield stdout_descriptor, stderr_descriptor + + +def _create_log_file(directory_descriptor: int, name: str) -> int: + return os.open( + name, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + dir_fd=directory_descriptor, + ) + + +__all__ = ["bind_execution_logs", "execution_log_directory", "open_step_logs"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py new file mode 100644 index 000000000..60da12f7d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated command manifest consumed by the allocation shell runtime.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path +from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.runtime.network import validate_host_name +from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest + + +class RuntimeProbeSpec(ContractValue): + """One readiness target monitored from the allocation client host.""" + + host: str + port: NetworkPort + path: str + deadline_seconds: PositiveInt + + @field_validator("path") + @classmethod + def validate_path(cls, value: str) -> str: + if not value.startswith("/") or any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("runtime probe path is invalid") + return value + + @field_validator("host") + @classmethod + def validate_host(cls, value: str) -> str: + return validate_host_name(value) + + +class RuntimeStepSpec(ContractValue): + """Native or container command and placement consumed by the Bash step runner.""" + + step_id: Identifier + role: RuntimeStepRole + execution: Literal["native", "container"] + image_path: str | None = None + command: tuple[str, ...] = Field(min_length=1) + cpus: PositiveInt + gpu_indices: tuple[NonNegativeInt, ...] = () + gpu_binding: Literal["mask", "none"] = Field(default="mask", exclude_if=lambda value: value == "mask") + node_hosts: tuple[str, ...] = Field(min_length=1) + kill_on_bad_exit: bool = False + literal_environment: dict[EnvironmentName, str] = Field(default_factory=dict) + secret_environment: dict[EnvironmentName, EnvironmentName] = Field(default_factory=dict) + environment_prefixes: dict[EnvironmentName, str] = Field(default_factory=dict) + container_environment: tuple[EnvironmentName, ...] = () + stdout_path: str + stderr_path: str + launch_delay_seconds: NonNegativeInt = 0 + readiness: tuple[RuntimeProbeSpec, ...] = () + + @field_validator("image_path") + @classmethod + def validate_image_path(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + _stdout_path_is_absolute = field_validator("stdout_path")(validate_absolute_path) + _stderr_path_is_absolute = field_validator("stderr_path")(validate_absolute_path) + + @model_validator(mode="after") + def validate_step(self) -> RuntimeStepSpec: + if any(not argument or "\0" in argument for argument in self.command): + raise ValueError("runtime command is invalid") + if self.gpu_indices != tuple(sorted(set(self.gpu_indices))): + raise ValueError("runtime GPU indices must be sorted and unique") + if self.node_hosts != tuple(dict.fromkeys(self.node_hosts)): + raise ValueError("runtime node hosts must be unique") + for host in self.node_hosts: + validate_host_name(host) + if self.stdout_path == self.stderr_path or Path(self.stdout_path).parent != Path(self.stderr_path).parent: + raise ValueError("runtime log paths must be distinct siblings") + if set(self.environment_prefixes) - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("environment prefixes require a materialized variable") + container_names = set(self.container_environment) + if container_names - (set(self.literal_environment) | set(self.secret_environment)): + raise ValueError("container environment contains an unavailable variable") + if self.execution == "container" and self.image_path is None: + raise ValueError("container runtime steps require an image") + if self.execution == "native" and (self.image_path is not None or self.container_environment): + raise ValueError("native runtime steps cannot contain container settings") + self._validate_placement() + self._validate_readiness() + return self + + def _validate_placement(self) -> None: + server_roles = {RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER} + expected_execution = "container" if self.role in server_roles else "native" + if self.execution != expected_execution: + raise ValueError("runtime step execution does not match its role") + if self.role in server_roles and not self.gpu_indices: + raise ValueError("server runtime steps require GPUs") + if self.role not in server_roles and self.gpu_indices: + raise ValueError("non-server runtime steps cannot request GPUs") + if self.gpu_binding == "none" and not self.gpu_indices: + raise ValueError("unbound GPU steps require GPUs") + if len(self.node_hosts) > 1 and self.role not in server_roles: + raise ValueError("only server runtime steps may span nodes") + if self.kill_on_bad_exit and self.role not in server_roles: + raise ValueError("kill-on-bad-exit is only valid for server runtime steps") + if len(self.node_hosts) > 1 and not self.kill_on_bad_exit: + raise ValueError("multi-node server runtime steps require kill-on-bad-exit") + + def _validate_readiness(self) -> None: + if self.role is RuntimeStepRole.SERVER_PREFLIGHT and self.readiness: + raise ValueError("server preflight steps cannot have readiness probes") + if self.role in {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} and not self.readiness: + raise ValueError("long-running runtime steps require readiness probes") + if self.role in {RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.CLIENT} and self.readiness: + raise ValueError("client runtime steps cannot have readiness probes") + + +class RuntimeBootstrapManifest(ContractRecord): + """Secret-free allocation command manifest.""" + + run_id: Identifier + shard_id: Identifier + attempt_id: Identifier + plan_sha256: Sha256Digest + all_secret_environment_names: tuple[EnvironmentName, ...] + steps: tuple[RuntimeStepSpec, ...] = Field(min_length=4) + + @model_validator(mode="after") + def validate_steps(self) -> RuntimeBootstrapManifest: + step_ids = tuple(step.step_id for step in self.steps) + if len(step_ids) != len(set(step_ids)): + raise ValueError("runtime step identifiers must be unique") + roles = tuple(step.role for step in self.steps) + if roles.count(RuntimeStepRole.CLIENT_PREFLIGHT) != 1 or roles.count(RuntimeStepRole.CLIENT) != 1: + raise ValueError("runtime manifest requires one preflight and generation step") + if RuntimeStepRole.SERVER not in roles or RuntimeStepRole.ENDPOINT not in roles: + raise ValueError("runtime manifest requires server and endpoint steps") + return self + + +__all__ = ["RuntimeBootstrapManifest", "RuntimeProbeSpec", "RuntimeStepSpec"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py new file mode 100644 index 000000000..c91370b38 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Transient typed values for allocation-local process execution.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import MappingProxyType + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.state import AttemptManifest, RetryPlan + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RuntimeStepRole(str, Enum): + """Lifecycle role of one allocation-local Slurm step.""" + + CLIENT_PREFLIGHT = "client_preflight" + SERVER_PREFLIGHT = "server_preflight" + SERVER = "server" + ENDPOINT = "endpoint" + CLIENT = "client" + + +@dataclass(frozen=True, slots=True) +class RuntimeStep: + """One shell-free process specification owned by the common step runner.""" + + step_id: Identifier + role: RuntimeStepRole + command: tuple[str, ...] + environment: Mapping[str, str] + stdout_path: Path + stderr_path: Path + + def __post_init__(self) -> None: + _validate_step_identity(self.step_id) + if not isinstance(self.role, RuntimeStepRole): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime step role is invalid") + _validate_step_command(self.command) + normalized_environment = _validate_step_environment(self.environment) + _validate_step_log_paths(self.stdout_path, self.stderr_path) + object.__setattr__(self, "environment", MappingProxyType(normalized_environment)) + + +def _validate_step_identity(step_id: Identifier) -> None: + try: + normalized_step_id = _IDENTIFIER_ADAPTER.validate_python(step_id, strict=True) + except ValidationError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime step identity is invalid") from error + if normalized_step_id != step_id: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime step identity is not canonical") + + +def _validate_step_command(command: tuple[str, ...]) -> None: + if not command or any(type(argument) is not str or not argument or "\0" in argument for argument in command): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime step command is invalid") + + +def _validate_step_environment(environment: Mapping[str, str]) -> dict[str, str]: + if not isinstance(environment, Mapping): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime environment is invalid") + normalized_environment: dict[str, str] = {} + for name, value in environment.items(): + if type(name) is not str or not name or "=" in name or "\0" in name: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime environment name is invalid") + if type(value) is not str or "\0" in value: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime environment value is invalid") + normalized_environment[name] = value + return normalized_environment + + +def _validate_step_log_paths(stdout_path: Path, stderr_path: Path) -> None: + for path in (stdout_path, stderr_path): + if not path.is_absolute(): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime log paths must be absolute") + if stdout_path.parent != stderr_path.parent or stdout_path == stderr_path: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime log paths must be distinct children of one directory", + ) + + +@dataclass(frozen=True, slots=True) +class RuntimeEndpoint: + """Client-visible logical endpoint published inside one allocation.""" + + model_alias: str + served_model_name: str + host: str + port: int + + def __post_init__(self) -> None: + if type(self.model_alias) is not str or not self.model_alias: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "endpoint model alias is invalid") + if type(self.served_model_name) is not str or not self.served_model_name: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "endpoint model name is invalid") + if self.host != "127.0.0.1": + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "one-node endpoints must use loopback") + if type(self.port) is not int or not 1 <= self.port <= 65535: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "endpoint port is invalid") + + +@dataclass(frozen=True, slots=True) +class AllocationContext: + """Verified plan, shard, attempt, and package-owned attempt directory.""" + + plan: ResolvedSlurmRunPlan + shard: PlannedShard + attempt: AttemptManifest + attempt_directory: Path + retry_plan: RetryPlan | None = None + + def __post_init__(self) -> None: + if not isinstance(self.plan, ResolvedSlurmRunPlan): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation plan is invalid") + if not isinstance(self.shard, PlannedShard) or self.shard not in self.plan.shards: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation shard is not in the plan") + if not isinstance(self.attempt, AttemptManifest): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation attempt is invalid") + expected_directory = ( + Path(self.plan.authored_config.path).parent + / "shards" + / self.shard.shard_id + / "attempts" + / self.attempt.attempt_id + ) + if self.attempt_directory != expected_directory: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt directory does not match persisted identity", + ) + if ( + self.attempt.run_id != self.plan.run_id + or self.attempt.shard_id != self.shard.shard_id + or self.attempt.scheduler is None + or self.attempt.scheduler.array_task_id != self.shard.array_task_index + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt does not match its planned scheduler task", + ) + if self.retry_plan is not None: + selected = tuple( + retry_shard + for retry_shard in self.retry_plan.planned_shards + if retry_shard.shard_id == self.shard.shard_id + and retry_shard.attempt_id == self.attempt.attempt_id + and retry_shard.array_task_index == self.shard.array_task_index + ) + if ( + self.retry_plan.run_id != self.plan.run_id + or self.retry_plan.resolved_plan != self.attempt.resolved_plan + or ( + self.plan.invocation.authored.resume != "if_possible" + and self.retry_plan.effective_resume_mode != self.plan.invocation.authored.resume + ) + or len(selected) != 1 + or selected[0].attempt_ordinal != self.attempt.attempt_ordinal + ): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation attempt does not match its persisted retry plan", + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py new file mode 100644 index 000000000..d17e4311a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared validation for scheduler-derived runtime network identities.""" + +from __future__ import annotations + +import re + +_HOST_NAME_PATTERN = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,251}[A-Za-z0-9])?$") + + +def validate_host_name(host: str) -> str: + """Return a safe scheduler host name or raise ``ValueError``.""" + if type(host) is not str or _HOST_NAME_PATTERN.fullmatch(host) is None: + raise ValueError("allocation host identity is invalid") + return host + + +def validate_network_port(port: int) -> int: + """Return a valid TCP port or raise ``ValueError``.""" + if type(port) is not int or not 1 <= port <= 65535: + raise ValueError("allocation network port is invalid") + return port + + +__all__ = ["validate_host_name", "validate_network_port"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py new file mode 100644 index 000000000..3c9d49576 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict transport records for one coordinated deployment's node workers.""" + +from __future__ import annotations + +import base64 +import binascii +import json +from dataclasses import asdict, dataclass + +from data_designer.slurm.contracts import validate_absolute_path +from data_designer.slurm.runtime.network import validate_host_name, validate_network_port + +_MAXIMUM_SPEC_BYTES = 64 * 1024 + + +@dataclass(frozen=True, slots=True) +class NodeReadinessSpec: + """Bounded loopback readiness for a node-local process.""" + + port: int + path: str + deadline_seconds: int + + def __post_init__(self) -> None: + validate_network_port(self.port) + if ( + type(self.path) is not str + or not self.path.startswith("/") + or any(ord(value) < 32 or ord(value) == 127 for value in self.path) + ): + raise ValueError("node readiness path is invalid") + if type(self.deadline_seconds) is not int or self.deadline_seconds <= 0: + raise ValueError("node readiness deadline is invalid") + + +@dataclass(frozen=True, slots=True) +class NodeProcessSpec: + """One resolved lane process owned by a node worker.""" + + process_id: str + command: tuple[str, ...] + gpu_indices: tuple[int, ...] + launch_delay_seconds: int + readiness: NodeReadinessSpec | None = None + + def __post_init__(self) -> None: + if ( + type(self.process_id) is not str + or not self.process_id + or any(ord(value) < 32 or ord(value) == 127 for value in self.process_id) + ): + raise ValueError("node process identity is invalid") + if ( + type(self.command) is not tuple + or not self.command + or any(type(value) is not str or not value or "\0" in value for value in self.command) + ): + raise ValueError("node process command is invalid") + if ( + type(self.gpu_indices) is not tuple + or not self.gpu_indices + or self.gpu_indices != tuple(sorted(set(self.gpu_indices))) + or any(type(value) is not int or value < 0 for value in self.gpu_indices) + ): + raise ValueError("node process GPU indices are invalid") + if type(self.launch_delay_seconds) is not int or self.launch_delay_seconds < 0: + raise ValueError("node process launch delay is invalid") + if self.readiness is not None and not isinstance(self.readiness, NodeReadinessSpec): + raise ValueError("node process readiness is invalid") + + +@dataclass(frozen=True, slots=True) +class NodeSpec: + """Resolved work and ports for one node in a coordinated step.""" + + node_index: int + host: str + ports: tuple[int, ...] + processes: tuple[NodeProcessSpec, ...] + + def __post_init__(self) -> None: + if type(self.node_index) is not int or self.node_index < 0: + raise ValueError("node index is invalid") + validate_host_name(self.host) + if type(self.ports) is not tuple or type(self.processes) is not tuple: + raise ValueError("node ports and processes must be tuples") + for port in self.ports: + validate_network_port(port) + if self.ports != tuple(sorted(set(self.ports))): + raise ValueError("node ports must be sorted and unique") + if not self.processes or len({process.process_id for process in self.processes}) != len(self.processes): + raise ValueError("node process identities must be present and unique") + if any( + process.readiness is not None and process.readiness.port not in self.ports for process in self.processes + ): + raise ValueError("node readiness port is not assigned to the node") + + +@dataclass(frozen=True, slots=True) +class NodeWorkerSpec: + """Strict, bounded input shared by every task in one deployment step.""" + + schema_version: int + resolved_gpus_per_node: int + required_model_path: str | None + nodes: tuple[NodeSpec, ...] + required_model_paths: tuple[str, ...] = () + allow_gpu_sharing: bool = False + + def __post_init__(self) -> None: + if type(self.allow_gpu_sharing) is not bool: + raise ValueError("GPU sharing policy must be a boolean") + if type(self.schema_version) is not int or self.schema_version != 1: + raise ValueError("node worker schema version is unsupported") + if type(self.resolved_gpus_per_node) is not int or self.resolved_gpus_per_node <= 0: + raise ValueError("resolved GPU count is invalid") + if self.required_model_path is not None and type(self.required_model_path) is not str: + raise ValueError("required model path is invalid") + if self.required_model_path is not None: + validate_absolute_path(self.required_model_path) + if type(self.required_model_paths) is not tuple: + raise ValueError("required model paths must be a tuple") + for path in self.required_model_paths: + if type(path) is not str: + raise ValueError("required model path is invalid") + validate_absolute_path(path) + if ( + type(self.nodes) is not tuple + or not self.nodes + or tuple(node.node_index for node in self.nodes) != tuple(sorted({node.node_index for node in self.nodes})) + ): + raise ValueError("node worker nodes must use sorted unique indices") + process_ids = tuple(process.process_id for node in self.nodes for process in node.processes) + if len({node.host for node in self.nodes}) != len(self.nodes) or len(set(process_ids)) != len(process_ids): + raise ValueError("node worker hosts and process identities must be unique") + for node in self.nodes: + if any(index >= self.resolved_gpus_per_node for process in node.processes for index in process.gpu_indices): + raise ValueError("node process GPU index is outside the allocation") + assigned_gpus = tuple(index for process in node.processes for index in process.gpu_indices) + if len(set(assigned_gpus)) != len(assigned_gpus): + if not self.allow_gpu_sharing: + raise ValueError("node processes must not share GPUs") + if any(process.readiness is None for process in node.processes): + raise ValueError("GPU sharing requires process readiness probes") + + +def encode_node_worker_spec(spec: NodeWorkerSpec) -> str: + """Serialize one validated node-worker specification for an argv boundary.""" + payload = { + "schema_version": spec.schema_version, + "resolved_gpus_per_node": spec.resolved_gpus_per_node, + "required_model_path": spec.required_model_path, + "nodes": [ + { + "node_index": node.node_index, + "host": node.host, + "ports": list(node.ports), + "processes": [ + { + "process_id": process.process_id, + "command": list(process.command), + "gpu_indices": list(process.gpu_indices), + "launch_delay_seconds": process.launch_delay_seconds, + **({"readiness": asdict(process.readiness)} if process.readiness else {}), + } + for process in node.processes + ], + } + for node in spec.nodes + ], + } + if spec.required_model_paths: + payload["required_model_paths"] = list(spec.required_model_paths) + if spec.allow_gpu_sharing: + payload["allow_gpu_sharing"] = True + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + return base64.urlsafe_b64encode(serialized).decode() + + +def decode_node_worker_spec(encoded: str) -> NodeWorkerSpec: + """Decode and validate an untrusted node-worker specification.""" + if type(encoded) is not str or not encoded or len(encoded) > 2 * _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification is invalid") + try: + serialized = base64.b64decode(encoded, altchars=b"-_", validate=True) + if len(serialized) > _MAXIMUM_SPEC_BYTES: + raise ValueError("node worker specification exceeds its size limit") + payload = json.loads(serialized) + except (binascii.Error, json.JSONDecodeError, UnicodeEncodeError) as error: + raise ValueError("node worker specification is invalid") from error + return _parse_worker_spec(payload) + + +def _parse_worker_spec(payload: object) -> NodeWorkerSpec: + keys = {"schema_version", "resolved_gpus_per_node", "required_model_path", "nodes"} + if isinstance(payload, dict): + keys.update(set(payload) & {"required_model_paths", "allow_gpu_sharing"}) + root = _require_mapping(payload, keys) + nodes = tuple(_parse_node(value) for value in _require_list(root["nodes"])) + return NodeWorkerSpec( + schema_version=_require_integer(root["schema_version"]), + resolved_gpus_per_node=_require_integer(root["resolved_gpus_per_node"]), + required_model_path=_require_optional_string(root["required_model_path"]), + nodes=nodes, + required_model_paths=tuple( + _require_string(path) for path in _require_list(root.get("required_model_paths", [])) + ), + allow_gpu_sharing=_require_boolean(root.get("allow_gpu_sharing", False)), + ) + + +def _parse_node(payload: object) -> NodeSpec: + value = _require_mapping(payload, {"node_index", "host", "ports", "processes"}) + return NodeSpec( + node_index=_require_integer(value["node_index"]), + host=_require_string(value["host"]), + ports=tuple(_require_integer(port) for port in _require_list(value["ports"])), + processes=tuple(_parse_process(process) for process in _require_list(value["processes"])), + ) + + +def _parse_process(payload: object) -> NodeProcessSpec: + keys = {"process_id", "command", "gpu_indices", "launch_delay_seconds"} + if isinstance(payload, dict) and "readiness" in payload: + keys.add("readiness") + value = _require_mapping(payload, keys) + return NodeProcessSpec( + process_id=_require_string(value["process_id"]), + command=tuple(_require_string(argument) for argument in _require_list(value["command"])), + gpu_indices=tuple(_require_integer(index) for index in _require_list(value["gpu_indices"])), + launch_delay_seconds=_require_integer(value["launch_delay_seconds"]), + readiness=_parse_readiness(value["readiness"]) if "readiness" in value else None, + ) + + +def _parse_readiness(payload: object) -> NodeReadinessSpec: + value = _require_mapping(payload, {"port", "path", "deadline_seconds"}) + return NodeReadinessSpec( + port=_require_integer(value["port"]), + path=_require_string(value["path"]), + deadline_seconds=_require_integer(value["deadline_seconds"]), + ) + + +def _require_mapping(payload: object, keys: set[str]) -> dict[str, object]: + if not isinstance(payload, dict) or set(payload) != keys or not all(type(key) is str for key in payload): + raise ValueError("node worker object shape is invalid") + return payload + + +def _require_list(payload: object) -> list[object]: + if not isinstance(payload, list): + raise ValueError("node worker list value is invalid") + return payload + + +def _require_string(payload: object) -> str: + if type(payload) is not str: + raise ValueError("node worker string value is invalid") + return payload + + +def _require_optional_string(payload: object) -> str | None: + if payload is None: + return None + return _require_string(payload) + + +def _require_integer(payload: object) -> int: + if type(payload) is not int: + raise ValueError("node worker integer value is invalid") + return payload + + +def _require_boolean(payload: object) -> bool: + if type(payload) is not bool: + raise ValueError("node worker boolean value is invalid") + return payload + + +__all__ = [ + "NodeProcessSpec", + "NodeReadinessSpec", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py new file mode 100644 index 000000000..56186b79d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run the resolved vLLM lanes assigned to one physical allocation node.""" + +from __future__ import annotations + +import argparse +import os +import signal +import socket +import stat +import subprocess +import sys +import time +from collections.abc import Mapping, Sequence +from tempfile import TemporaryDirectory + +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec as NodeProcessSpec, +) +from data_designer.slurm.runtime.node_spec import NodeReadinessSpec +from data_designer.slurm.runtime.node_spec import ( + NodeSpec as NodeSpec, +) +from data_designer.slurm.runtime.node_spec import ( + NodeWorkerSpec as NodeWorkerSpec, +) +from data_designer.slurm.runtime.node_spec import ( + decode_node_worker_spec as decode_node_worker_spec, +) +from data_designer.slurm.runtime.node_spec import ( + encode_node_worker_spec as encode_node_worker_spec, +) +from data_designer.slurm.runtime.probes import HttpReadinessProber + +_POLL_INTERVAL_SECONDS = 0.1 +_TERMINATION_GRACE_SECONDS = 10.0 + + +class NodeProcessSupervisor: + """Own every lane child on one node and stop them as a unit.""" + + def __init__(self, *, environment: Mapping[str, str]) -> None: + self._environment = dict(environment) + self._children: list[subprocess.Popen[bytes]] = [] + self._tokenizer_caches: list[TemporaryDirectory[str]] = [] + self._cleanup_started = False + self._cleanup_complete = False + self._stop_requested = False + + @property + def cleanup_complete(self) -> bool: + """Return whether every registered lane child has exited.""" + return self._cleanup_complete + + def run(self, node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + """Launch local lanes and fail the node task when any lane exits.""" + started_at = time.monotonic() + try: + for index, process in enumerate(node.processes): + early_exit = self._wait_for_launch(process.launch_delay_seconds, started_at) + if early_exit is not None: + return early_exit + if self._stop_requested: + return 1 + self._children.append(self._start_process(process, visible_gpus)) + if process.readiness is not None and any( + set(process.gpu_indices).intersection(later.gpu_indices) for later in node.processes[index + 1 :] + ): + early_exit = self._wait_for_readiness(process.readiness) + if early_exit is not None: + return early_exit + return self._wait_for_first_exit() + finally: + self.cleanup() + + def request_stop(self) -> None: + """Ask the poll loop to enter cleanup at its next boundary.""" + self._stop_requested = True + + def cleanup(self) -> None: + """Idempotently terminate every registered lane child.""" + if self._cleanup_complete: + return + self._cleanup_started = True + live = tuple(child for child in reversed(self._children) if child.poll() is None) + _signal_children(live, signal.SIGTERM) + deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS + for child in live: + _wait_until_deadline(child, deadline) + _signal_children(tuple(child for child in live if child.poll() is None), signal.SIGKILL) + for child in live: + try: + child.wait(timeout=_TERMINATION_GRACE_SECONDS) + except (OSError, subprocess.SubprocessError): + pass + if any(child.poll() is None for child in self._children): + raise RuntimeError("node process cleanup is incomplete") + for cache in self._tokenizer_caches: + cache.cleanup() + self._tokenizer_caches.clear() + self._cleanup_complete = True + + def _start_process(self, process: NodeProcessSpec, visible_gpus: tuple[str, ...]) -> subprocess.Popen[bytes]: + if self._cleanup_started: + raise RuntimeError("cannot launch a node process after cleanup") + environment = dict(self._environment) + environment["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpus[index] for index in process.gpu_indices) + if not environment.get("TIKTOKEN_RS_CACHE_DIR"): + # Harmony writes downloaded vocabularies directly to its cache path. + # Separate concurrent vLLM replicas so one cannot read or remove a + # sibling's partially written vocabulary file. + cache = TemporaryDirectory(prefix="data-designer-tiktoken-", ignore_cleanup_errors=True) + self._tokenizer_caches.append(cache) + environment["TIKTOKEN_RS_CACHE_DIR"] = cache.name + return subprocess.Popen( + process.command, + stdin=subprocess.DEVNULL, + env=environment, + shell=False, + start_new_session=True, + close_fds=True, + ) + + def _wait_for_launch(self, delay_seconds: int, started_at: float) -> int | None: + while not self._stop_requested: + for child in self._children: + returncode = child.poll() + if returncode is not None: + return returncode if returncode != 0 else 1 + remaining = delay_seconds - (time.monotonic() - started_at) + if remaining <= 0: + return + time.sleep(min(remaining, _POLL_INTERVAL_SECONDS)) + + def _wait_for_readiness(self, probe: NodeReadinessSpec) -> int | None: + deadline = time.monotonic() + probe.deadline_seconds + prober = HttpReadinessProber() + while not self._stop_requested: + early_exit = self._wait_for_launch(0, 0) + if early_exit is not None: + return early_exit + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError("shared GPU process readiness timed out") + if prober.is_ready("127.0.0.1", probe.port, probe.path, timeout_seconds=min(1.0, remaining)): + return None + time.sleep(min(remaining, _POLL_INTERVAL_SECONDS)) + return 1 + + def _wait_for_first_exit(self) -> int: + while True: + if self._stop_requested: + return 1 + for child in self._children: + returncode = child.poll() + if returncode is not None: + return returncode if returncode != 0 else 1 + time.sleep(_POLL_INTERVAL_SECONDS) + + +def main(arguments: Sequence[str] | None = None) -> int: + """Preflight or serve the node-local slice of a coordinated deployment.""" + parser = argparse.ArgumentParser(prog="data-designer-slurm-node-worker") + parser.add_argument("operation", choices=("preflight", "serve")) + parser.add_argument("--spec", required=True) + parsed = parser.parse_args(arguments) + try: + spec = decode_node_worker_spec(parsed.spec) + node = _select_node(spec, os.environ) + visible_gpus = _parse_visible_gpus(os.environ.get("CUDA_VISIBLE_DEVICES")) + _verify_node(spec, node, visible_gpus, os.environ) + if parsed.operation == "preflight": + _verify_required_model_path(spec.required_model_path) + for path in spec.required_model_paths: + _verify_required_model_path(path) + _verify_ports(node.ports) + return 0 + return _run_node(node, visible_gpus) + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError): + print("node worker failed at a validated runtime boundary", file=sys.stderr) + return 70 + + +def _run_node(node: NodeSpec, visible_gpus: tuple[str, ...]) -> int: + supervisor = NodeProcessSupervisor(environment=os.environ) + interrupted: list[int] = [] + + def handle_termination(signum: int, frame: object) -> None: + del frame + if not interrupted: + interrupted.append(signum) + supervisor.request_stop() + + previous = {selected: signal.signal(selected, handle_termination) for selected in (signal.SIGINT, signal.SIGTERM)} + try: + return _run_until_exit_or_signal(supervisor, node, visible_gpus, interrupted) + finally: + try: + supervisor.cleanup() + finally: + for selected, handler in previous.items(): + signal.signal(selected, handler) + + +def _run_until_exit_or_signal( + supervisor: NodeProcessSupervisor, + node: NodeSpec, + visible_gpus: tuple[str, ...], + interrupted: list[int], +) -> int: + if interrupted: + return 128 + interrupted[0] + result = supervisor.run(node, visible_gpus) + if interrupted: + return 128 + interrupted[0] + return result + + +def _select_node(spec: NodeWorkerSpec, environment: Mapping[str, str]) -> NodeSpec: + process_id = _parse_non_negative_integer(environment.get("SLURM_PROCID"), "SLURM_PROCID") + if process_id >= len(spec.nodes): + raise ValueError("Slurm process identity is outside the deployment") + return spec.nodes[process_id] + + +def _verify_node( + spec: NodeWorkerSpec, + node: NodeSpec, + visible_gpus: tuple[str, ...], + environment: Mapping[str, str], +) -> None: + if len(visible_gpus) != spec.resolved_gpus_per_node: + raise ValueError("node GPU visibility does not match the resolved plan") + task_count = _parse_non_negative_integer(environment.get("SLURM_NTASKS"), "SLURM_NTASKS") + if task_count != len(spec.nodes): + raise ValueError("Slurm task count does not match the deployment") + scheduler_host = environment.get("SLURMD_NODENAME") + if scheduler_host is not None and scheduler_host != node.host: + raise ValueError("Slurm node identity does not match the deployment") + + +def _verify_ports(ports: tuple[int, ...]) -> None: + reservations: list[socket.socket] = [] + try: + for port in ports: + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + reservations.append(reservation) + reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) + reservation.bind(("0.0.0.0", port)) + finally: + for reservation in reservations: + reservation.close() + + +def _verify_required_model_path(path: str | None) -> None: + if path is None: + return + status = os.stat(path) + if not stat.S_ISREG(status.st_mode) and not stat.S_ISDIR(status.st_mode): + raise OSError("required model path has an unsupported type") + access_mode = os.R_OK | os.X_OK if stat.S_ISDIR(status.st_mode) else os.R_OK + if not os.access(path, access_mode): + raise OSError("required model path is not readable") + + +def _parse_visible_gpus(value: str | None) -> tuple[str, ...]: + if value is None or not value.strip(): + return () + values = tuple(item.strip() for item in value.split(",")) + if not all(values) or len(values) != len(set(values)) or any("\0" in item for item in values): + raise ValueError("node GPU visibility is invalid") + return values + + +def _parse_non_negative_integer(value: str | None, name: str) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise ValueError(f"{name} is unavailable or invalid") + return int(value) + + +def _signal_process_group(process: subprocess.Popen[bytes], selected: signal.Signals) -> None: + try: + os.killpg(process.pid, selected) + except ProcessLookupError: + pass + + +def _signal_children(children: tuple[subprocess.Popen[bytes], ...], selected: signal.Signals) -> None: + for child in children: + try: + _signal_process_group(child, selected) + except OSError: + continue + + +def _wait_until_deadline(process: subprocess.Popen[bytes], deadline: float) -> None: + while process.poll() is None and time.monotonic() < deadline: + time.sleep(_POLL_INTERVAL_SECONDS) + + +if __name__ == "__main__": # pragma: no cover - exercised as a managed step + raise SystemExit(main()) + + +__all__ = [ + "NodeProcessSpec", + "NodeProcessSupervisor", + "NodeSpec", + "NodeWorkerSpec", + "decode_node_worker_spec", + "encode_node_worker_spec", + "main", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py new file mode 100644 index 000000000..edf1593f3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Translate reviewed host paths through resolved container mounts.""" + +from __future__ import annotations + +import posixpath + +from data_designer.slurm.contracts import validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode + +ALLOCATION_SCRATCH_CONTAINER_ROOT = "/run/data-designer-slurm" + + +def get_container_path( + plan: ResolvedSlurmRunPlan, + host_path: str, + *, + require_writable: bool = False, +) -> str: + """Map one absolute host path through the most specific resolved mount.""" + try: + normalized = validate_absolute_path(host_path) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime host path is invalid") from error + candidates = tuple( + mount + for mount in plan.container_mounts + if normalized == mount.source or normalized.startswith(f"{mount.source}/") + ) + if not candidates: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path is not available through a resolved container mount", + ) + mount = max(candidates, key=lambda value: len(value.source)) + if require_writable and mount.read_only: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path requires a writable resolved container mount", + ) + relative_path = posixpath.relpath(normalized, mount.source) + mapped = mount.target if relative_path == "." else posixpath.join(mount.target, relative_path) + return validate_absolute_path(mapped) + + +def get_host_path( + plan: ResolvedSlurmRunPlan, + container_path: str, + *, + require_writable: bool = False, +) -> str: + """Map one absolute container path back through the most specific resolved mount.""" + try: + normalized = validate_absolute_path(container_path) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime container path is invalid") from error + candidates = tuple( + mount + for mount in plan.container_mounts + if normalized == mount.target or normalized.startswith(f"{mount.target}/") + ) + if not candidates: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path is not available through a resolved container mount", + ) + mount = max(candidates, key=lambda value: len(value.target)) + if require_writable and mount.read_only: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "runtime path requires a writable resolved container mount", + ) + relative_path = posixpath.relpath(normalized, mount.target) + mapped = mount.source if relative_path == "." else posixpath.join(mount.source, relative_path) + host_path = validate_absolute_path(mapped) + if get_container_path(plan, host_path, require_writable=require_writable) != normalized: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime path does not map unambiguously through resolved container mounts", + ) + return host_path + + +__all__ = ["ALLOCATION_SCRATCH_CONTAINER_ROOT", "get_container_path", "get_host_path"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh new file mode 100644 index 000000000..f1c78a462 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +dd_read_null_values() { + local array_name=$1 + local index=0 value + unset "${array_name}" + declare -g -a "${array_name}" + printf -v "${array_name}[0]" '%s' "" + unset "${array_name}[0]" + while IFS= read -r -d '' value; do + printf -v "${array_name}[${index}]" '%s' "${value}" + ((index += 1)) + done +} + +dd_read_control_plan() { + local plan=$1 + jq -e ' + .schema_version == 1 + and ([.deployments[].node_indices[]] | length > 0) + and ([.deployments[].node_indices[]] | all(type == "number" and . >= 0)) + and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) + and ([.container_mounts[].target] | all( + . != "/run/data-designer-slurm" + and (startswith("/run/data-designer-slurm/") | not) + and (. as $target | ("/run/data-designer-slurm" | startswith($target + "/") | not)) + )) + ' "${plan}" >/dev/null || return 65 + DD_CLIENT_PYTHON=$(jq -er '.client.runtime.python_executable' "${plan}") + DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") + DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | tostring' "${plan}") + DD_EXPECTED_NODES=$(jq -er '[.deployments[].node_indices[]] | max + 1 | tostring' "${plan}") + DD_CLIENT_NODE_INDEX=$(jq -er '.client.host_node_index | tostring' "${plan}") + DD_GPU_REQUEST_MODE=$(jq -er '.selected_profile.profile.gpu_request_mode' "${plan}") + DD_CONTAINER_MOUNTS=$(jq -jr ' + [.container_mounts[] | .source + ":" + .target + (if .read_only then ":ro" else "" end)] + | join(",") + ' "${plan}") +} + +dd_read_container_path() { + local plan=$1 + local host_path=$2 + local require_writable=$3 + DD_CONTAINER_PATH=$(jq -er \ + --arg path "${host_path}" \ + --argjson require_writable "${require_writable}" ' + [ + .container_mounts[] + | . as $mount + | select(($path == $mount.source) or ($path | startswith($mount.source + "/"))) + | select(($require_writable | not) or ($mount.read_only | not)) + ] + | sort_by(.source | length) + | last + | select(. != null) + | . as $mount + | if $path == $mount.source then + $mount.target + else + $mount.target + ($path | ltrimstr($mount.source)) + end + ' "${plan}") +} + +dd_read_plan_secret_names() { + local plan=$1 + dd_read_null_values DD_REQUIRED_SECRET_NAMES < <( + jq -j '[.. | objects | select(.type? == "secret") | .environment] | unique | .[] | ., "\u0000"' "${plan}" + ) + dd_read_null_values DD_ALL_SECRET_NAMES < <( + jq -j ' + [ + (.. | objects | select(.type? == "secret") | .environment), + ( + .deployments[].authored.server.environment + | to_entries[] + | select(.value.type == "secret") + | .key + ) + ] + | unique + | .[] + | ., "\u0000" + ' "${plan}" + ) +} + +dd_verify_runtime_manifest() { + local manifest=$1 + local plan_sha256=$2 + local shard_id=$3 + local attempt_id=$4 + jq -e \ + --arg plan_sha256 "${plan_sha256}" \ + --arg shard_id "${shard_id}" \ + --arg attempt_id "${attempt_id}" ' + .schema_version == 1 + and .plan_sha256 == $plan_sha256 + and .shard_id == $shard_id + and .attempt_id == $attempt_id + and ([.steps[].step_id] | length == (unique | length)) + ' "${manifest}" >/dev/null +} + +dd_read_step_ids() { + local manifest=$1 + local role=$2 + dd_read_null_values DD_STEP_IDS < <( + jq -j --arg role "${role}" '.steps[] | select(.role == $role) | .step_id, "\u0000"' "${manifest}" + ) +} + +dd_read_step() { + local manifest=$1 + local step_id=$2 + jq -e --arg step_id "${step_id}" '[.steps[] | select(.step_id == $step_id)] | length == 1' \ + "${manifest}" >/dev/null + dd_read_null_values DD_STEP_FIELDS < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .execution, "\u0000", + (.image_path // ""), "\u0000", + (.cpus | tostring), "\u0000", + .stdout_path, "\u0000", + .stderr_path, "\u0000", + (.launch_delay_seconds | tostring), "\u0000", + (.kill_on_bad_exit | tostring), "\u0000", + (.gpu_binding // "mask"), "\u0000" + ' "${manifest}" + ) + DD_STEP_EXECUTION=${DD_STEP_FIELDS[0]} + DD_STEP_IMAGE=${DD_STEP_FIELDS[1]} + DD_STEP_CPUS=${DD_STEP_FIELDS[2]} + DD_STEP_STDOUT=${DD_STEP_FIELDS[3]} + DD_STEP_STDERR=${DD_STEP_FIELDS[4]} + DD_STEP_DELAY=${DD_STEP_FIELDS[5]} + DD_STEP_KILL_ON_BAD_EXIT=${DD_STEP_FIELDS[6]} + DD_STEP_GPU_BINDING=${DD_STEP_FIELDS[7]} + dd_read_null_values DD_STEP_COMMAND < <( + jq -j --arg step_id "${step_id}" '.steps[] | select(.step_id == $step_id) | .command[] | ., "\u0000"' \ + "${manifest}" + ) + dd_read_null_values DD_STEP_GPU_INDICES < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .gpu_indices[] | tostring, "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_NODE_HOSTS < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .node_hosts[] | ., "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_PROBE_FIELDS < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .readiness[] + | .host, "\u0000", + (.port | tostring), "\u0000", + .path, "\u0000", + (.deadline_seconds | tostring), "\u0000" + ' "${manifest}" + ) + dd_read_null_values DD_STEP_CONTAINER_ENV < <( + jq -j --arg step_id "${step_id}" \ + '.steps[] | select(.step_id == $step_id) | .container_environment[] | ., "\u0000"' "${manifest}" + ) + dd_read_null_values DD_STEP_LITERAL_ENV < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .literal_environment + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) + dd_read_null_values DD_STEP_SECRET_ENV < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .secret_environment + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) + dd_read_null_values DD_STEP_ENV_PREFIXES < <( + jq -j --arg step_id "${step_id}" ' + .steps[] + | select(.step_id == $step_id) + | .environment_prefixes + | to_entries[] + | .key, "\u0000", .value, "\u0000" + ' "${manifest}" + ) +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py new file mode 100644 index 000000000..53d6f049a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local port isolation for shared Slurm nodes.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from data_designer.slurm.planning import PortClaim, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import resolve_vllm_server + +_PORT_RANGE_START = 10000 +_PORT_RANGE_END = 65536 +_PORTS_PER_GPU = 256 + + +def resolve_allocation_deployments( + context: AllocationContext, + environment: Mapping[str, str], +) -> tuple[ResolvedVllmServerDeployment, ...]: + """Resolve deployments with ports isolated to one scheduler array element.""" + plan = resolve_allocation_plan(context.plan, environment) + return tuple(resolve_vllm_server(plan, item.deployment_id) for item in plan.deployments) + + +def allocation_ports(context: AllocationContext, environment: Mapping[str, str]) -> tuple[int, ...]: + """Return every allocation-local port in deterministic claim order.""" + plan = resolve_allocation_plan(context.plan, environment) + return tuple(port.port for port in plan.client.ports) + tuple( + port.port for deployment in plan.deployments for port in deployment.ports + ) + + +def resolve_allocation_plan( + plan: ResolvedSlurmRunPlan, + environment: Mapping[str, str], +) -> ResolvedSlurmRunPlan: + """Return a plan with ports bound to the allocation's assigned GPUs.""" + if plan.selected_profile.profile.gpu_request_mode == "visible": + return plan + gpu_ids = _parse_gpu_ids(environment.get("SLURM_JOB_GPUS")) + block_start = _PORT_RANGE_START + min(gpu_ids) * _PORTS_PER_GPU + block_end = block_start + _PORTS_PER_GPU + if block_end > _PORT_RANGE_END: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocated GPU index exceeds the supported allocation port range", + ) + claims = plan.client.ports + tuple(port for deployment in plan.deployments for port in deployment.ports) + claims_by_node: dict[int, set[int]] = {} + for claim in claims: + claims_by_node.setdefault(claim.node_index, set()).add(claim.port) + reserved_by_node: dict[int, set[int]] = {} + otel_port = plan.invocation.effective_run_config.get("otel_metrics_port") + if type(otel_port) is int and block_start <= otel_port < block_end: + reserved_by_node[plan.client.host_node_index] = {otel_port} + mapping: dict[tuple[int, int], int] = {} + for node_index, planned_ports in claims_by_node.items(): + reserved = reserved_by_node.get(node_index, set()) + next_port = block_start + for planned_port in sorted(planned_ports): + while next_port in reserved: + next_port += 1 + if next_port >= block_end: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation port range is exhausted") + mapping[(node_index, planned_port)] = next_port + reserved.add(next_port) + next_port += 1 + + def remap(port: PortClaim) -> PortClaim: + return port.model_copy(update={"port": mapping[(port.node_index, port.port)]}) + + client = plan.client.model_copy(update={"ports": tuple(remap(port) for port in plan.client.ports)}) + deployments = tuple( + deployment.model_copy(update={"ports": tuple(remap(port) for port in deployment.ports)}) + for deployment in plan.deployments + ) + return plan.model_copy(update={"client": client, "deployments": deployments}) + + +def _parse_gpu_ids(value: str | None) -> tuple[int, ...]: + if value is None: + gpu_ids = () + else: + fields = value.split(",") + gpu_ids = tuple(int(field) for field in fields if field.isascii() and field.isdigit()) + if len(gpu_ids) != len(fields): + gpu_ids = () + if not gpu_ids or len(gpu_ids) != len(set(gpu_ids)): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "scheduler environment 'SLURM_JOB_GPUS' is unavailable or invalid", + ) + return gpu_ids + + +__all__ = ["allocation_ports", "resolve_allocation_deployments", "resolve_allocation_plan"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py new file mode 100644 index 000000000..1086a9f6d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fail-fast allocation checks performed before any child process starts.""" + +from __future__ import annotations + +import hashlib +import os +import re +import socket +import stat +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.network import validate_host_name +from data_designer.slurm.runtime.ports import resolve_allocation_plan + +_DIGEST_CHUNK_SIZE = 1024 * 1024 +_GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") + + +@dataclass(frozen=True, slots=True) +class AllocationLayout: + """Verified allocation host identities in planner-index order.""" + + node_hosts: tuple[str, ...] + + def __post_init__(self) -> None: + if ( + type(self.node_hosts) is not tuple + or not self.node_hosts + or len(self.node_hosts) != len(set(self.node_hosts)) + ): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "allocation node identities are invalid") + try: + for host in self.node_hosts: + validate_host_name(host) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, str(error)) from error + + def get_host(self, node_index: int) -> str: + """Return the host assigned to one planner node index.""" + if type(node_index) is not int or node_index < 0: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) + try: + return self.node_hosts[node_index] + except (IndexError, TypeError): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved node index is outside the allocation", + ) from None + + +def validate_allocation_layout(plan: ResolvedSlurmRunPlan, layout: AllocationLayout) -> None: + """Require one scheduler host for every contiguous planner node index.""" + node_indices = { + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + } + if node_indices != set(range(len(layout.node_hosts))): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation host identities do not match the resolved plan", + ) + + +class AllocationPreflight(Protocol): + """Verify one allocation without starting package-managed processes.""" + + def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> None: + """Raise a normalized error when allocation facts disagree with the plan.""" + ... + + +class SystemAllocationPreflight: + """Production verification of scheduler, filesystem, GPU, and port facts.""" + + def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> None: + """Verify every launch-critical fact before model services start.""" + try: + self._verify_scheduler(context, environment) + attempt_directory = context.attempt_directory + self.verify_attempt_directory(attempt_directory) + self.verify_client_runtime(context) + self.verify_artifacts(context) + self.verify_ports(context, environment) + except SlurmRuntimeError: + raise + except (OSError, ValueError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation preflight could not verify launch inputs", + ) from error + + @staticmethod + def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str]) -> None: + node_indices = { + context.plan.client.host_node_index, + *(index for deployment in context.plan.deployments for index in deployment.node_indices), + } + node_count = max(node_indices) + 1 + if node_indices != set(range(node_count)): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "resolved allocation node indices are not complete", + ) + expected = { + "SLURM_ARRAY_JOB_ID": context.attempt.scheduler.array_job_id, + "SLURM_ARRAY_TASK_ID": context.shard.array_task_index, + "SLURM_JOB_NUM_NODES": node_count, + "SLURM_NODEID": context.plan.client.host_node_index, + } + for name, value in expected.items(): + if _parse_non_negative_integer(environment.get(name), name) != value: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + f"scheduler environment {name!r} does not match the resolved plan", + ) + visible_gpus = environment.get("CUDA_VISIBLE_DEVICES") or environment.get("SLURM_JOB_GPUS") + if _parse_gpu_count(visible_gpus) < context.plan.resolved_gpus_per_node: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation GPU visibility does not match the resolved plan", + ) + + @staticmethod + def verify_client_runtime(context: AllocationContext) -> None: + """Require the compute-node interpreter to match the submitted runtime.""" + try: + runtime = ClientRuntimeInspector().inspect() + except ClientRuntimeInspectionError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "client Python runtime is unavailable on the compute node", + ) from error + if runtime != context.plan.client.runtime: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "client Python runtime differs from the resolved plan", + ) + + @staticmethod + def verify_attempt_directory(attempt_directory: Path) -> None: + """Require an attempt workspace accessible only to its owner.""" + status = attempt_directory.lstat() + if not stat.S_ISDIR(status.st_mode) or status.st_mode & 0o077: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "attempt workspace is not a restrictive directory", + ) + + @staticmethod + def verify_artifacts(context: AllocationContext) -> None: + """Verify unique launch artifacts against their planned SHA-256 digests.""" + plan = context.plan + references = [ + ArtifactReference( + path=Path(plan.authored_config.path).with_name("resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ), + plan.runtime_bundle, + plan.client.dependency_lock, + *( + ArtifactReference(path=deployment.image.path, sha256=deployment.image.sha256) + for deployment in plan.deployments + ), + ] + optional_references = [ + plan.builder.source, + context.shard.input_partition, + ] + references.extend(reference for reference in optional_references if reference is not None) + unique_references = {(reference.path, reference.sha256): reference for reference in references} + try: + for reference in unique_references.values(): + _verify_artifact(reference) + except (OSError, ValueError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "allocation preflight could not verify launch inputs", + ) from error + + @staticmethod + def verify_ports(context: AllocationContext, environment: Mapping[str, str]) -> None: + """Verify ports owned by the local client host before nested steps start.""" + plan = resolve_allocation_plan(context.plan, environment) + local_node_index = plan.client.host_node_index + ports = tuple(port.port for port in plan.client.ports if port.node_index == local_node_index) + tuple( + port.port + for deployment in plan.deployments + for port in deployment.ports + if port.node_index == local_node_index + ) + reservations: list[socket.socket] = [] + try: + for port in ports: + reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + reservations.append(reservation) + reservation.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) + reservation.bind(("0.0.0.0", port)) + except OSError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "one or more local allocation ports are unavailable", + ) from error + finally: + for reservation in reservations: + reservation.close() + + +def _verify_artifact(reference: ArtifactReference) -> None: + path = Path(reference.path) + before = path.lstat() + if not stat.S_ISREG(before.st_mode): + raise OSError(f"artifact {path} is not a regular file") + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)) + try: + opened = os.fstat(descriptor) + if _file_identity(before) != _file_identity(opened): + raise OSError(f"artifact {path} changed while it was opened") + digest = hashlib.sha256() + while chunk := os.read(descriptor, _DIGEST_CHUNK_SIZE): + digest.update(chunk) + after = os.fstat(descriptor) + if _file_identity(opened) != _file_identity(after): + raise OSError(f"artifact {path} changed while it was read") + if digest.hexdigest() != reference.sha256: + raise OSError(f"artifact {path} digest does not match the plan") + current = path.lstat() + if not stat.S_ISREG(current.st_mode) or _file_identity(after) != _file_identity(current): + raise OSError(f"artifact {path} was replaced while it was verified") + finally: + os.close(descriptor) + + +def _file_identity(value: os.stat_result) -> tuple[int, int, int, int, int]: + return value.st_dev, value.st_ino, value.st_size, value.st_mtime_ns, value.st_ctime_ns + + +def _parse_non_negative_integer(value: str | None, name: str) -> int: + if value is None or not value.isascii() or not value.isdigit(): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + f"scheduler environment {name!r} is unavailable or invalid", + ) + return int(value) + + +def _parse_gpu_count(value: str | None) -> int: + if value is None or not value.strip(): + return 0 + normalized = value.strip() + match = _GPU_COUNT_PATTERN.fullmatch(normalized) + if match is not None: + return int(match.group(1)) + values = tuple(item.strip() for item in normalized.split(",")) + if not all(values) or len(values) != len(set(values)): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "GPU visibility is invalid") + return len(values) + + +__all__ = [ + "AllocationLayout", + "AllocationPreflight", + "SystemAllocationPreflight", + "validate_allocation_layout", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py new file mode 100644 index 000000000..6d56ea4e3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bounded loopback readiness probes for allocation-local endpoints.""" + +from __future__ import annotations + +import http.client +from typing import Protocol + + +class ReadinessProber(Protocol): + """Probe one loopback HTTP endpoint.""" + + def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: + """Return whether the endpoint produced HTTP 200 within the timeout.""" + ... + + +class HttpReadinessProber: + """Production stdlib HTTP prober with no redirect or proxy behavior.""" + + def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: + """Probe a reviewed loopback address and fully close the connection.""" + if host != "127.0.0.1" or not path.startswith("/") or timeout_seconds <= 0: + return False + connection = http.client.HTTPConnection(host, port, timeout=timeout_seconds) + try: + connection.request("GET", path, headers={"Connection": "close"}) + response = connection.getresponse() + return response.status == 200 + except (OSError, http.client.HTTPException): + return False + finally: + connection.close() + + +__all__ = ["HttpReadinessProber", "ReadinessProber"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py new file mode 100644 index 000000000..69003b0f5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -0,0 +1,464 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local streaming HTTP endpoint for one model alias.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import socket +import time +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from urllib.parse import SplitResult, urlsplit + +from aiohttp import ClientError, ClientResponse, ClientSession, ClientTimeout, TCPConnector, TraceConfig, web + +from data_designer.slurm.runtime.network import validate_host_name + +_LOGGER = logging.getLogger("data_designer.slurm.runtime.proxy") +_MAXIMUM_REQUEST_BYTES = 100 * 1024 * 1024 +_MAXIMUM_UPSTREAM_ATTEMPTS = 3 +_MAXIMUM_CONNECTIONS = 256 +_KEEPALIVE_TIMEOUT_SECONDS = 65.0 +_CONNECT_TIMEOUT_SECONDS = 5.0 +_READ_TIMEOUT_SECONDS = 3600.0 +_STREAM_CHUNK_BYTES = 64 * 1024 +_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) +_HOP_HEADERS = frozenset( + { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "proxy-connection", + "te", + "trailer", + "transfer-encoding", + "upgrade", + } +) + + +@dataclass(frozen=True, slots=True) +class _Backend: + host: str + port: int + + @property + def origin(self) -> str: + return f"http://{self.host}:{self.port}" + + +class _BackendPool: + """Track active requests for deterministic least-connections routing.""" + + def __init__(self, backends: tuple[_Backend, ...], retry_after_seconds: int | None) -> None: + if not backends: + raise ValueError("at least one backend is required") + self.backends = backends + self.retry_after_seconds = retry_after_seconds + self._active = [0] * len(backends) + self._cursor = 0 + + def acquire(self, excluded: frozenset[int]) -> int: + available = tuple(index for index in range(len(self.backends)) if index not in excluded) + if not available: + raise LookupError("no backend remains") + minimum = min(self._active[index] for index in available) + candidates = frozenset(index for index in available if self._active[index] == minimum) + index = next( + candidate + for offset in range(len(self.backends)) + if (candidate := (self._cursor + offset) % len(self.backends)) in candidates + ) + self._cursor = (index + 1) % len(self.backends) + self._active[index] += 1 + return index + + def acquire_index(self, index: int) -> None: + self._active[index] += 1 + + def release(self, index: int) -> None: + if self._active[index] <= 0: + raise RuntimeError("backend activity underflow") + self._active[index] -= 1 + + def active_counts(self) -> tuple[int, ...]: + return tuple(self._active) + + +class _ProxyMetrics: + """Allocation-local counters that never contain request content or secrets.""" + + def __init__(self) -> None: + self.requests = 0 + self.attempts = 0 + self.connection_opened = 0 + self.connection_reused = 0 + self.connection_failed = 0 + self.retries: Counter[str] = Counter() + self.final_statuses: Counter[int] = Counter() + self.backend_requests: Counter[int] = Counter() + + def snapshot(self, pool: _BackendPool, connector: TCPConnector) -> dict[str, object]: + idle_connections = sum(len(connections) for connections in connector._conns.values()) # noqa: SLF001 + return { + "requests": self.requests, + "attempts": self.attempts, + "connections": { + "limit": _MAXIMUM_CONNECTIONS, + "opened": self.connection_opened, + "reused": self.connection_reused, + "failed": self.connection_failed, + "idle": idle_connections, + }, + "retries": dict(sorted(self.retries.items())), + "final_statuses": {str(key): value for key, value in sorted(self.final_statuses.items())}, + "backends": [ + { + "index": index, + "active_requests": active, + "requests": self.backend_requests[index], + } + for index, active in enumerate(pool.active_counts()) + ], + } + + +class _ProxyApplication: + """One logical endpoint with a shared upstream connection pool.""" + + def __init__( + self, + backends: tuple[_Backend, ...], + retry_after_seconds: int | None, + *, + health_path: str = "/health", + ) -> None: + self.pool = _BackendPool(backends, retry_after_seconds) + self.health_path = health_path + self.metrics = _ProxyMetrics() + self.connector: TCPConnector | None = None + self.session: ClientSession | None = None + + def create(self) -> web.Application: + application = web.Application(client_max_size=_MAXIMUM_REQUEST_BYTES) + application.on_startup.append(self._start) + application.on_cleanup.append(self._stop) + application.router.add_route("GET", "/health", self._health) + application.router.add_route("GET", "/metrics", self._metrics) + application.router.add_route("*", "/{path:.*}", self._forward) + return application + + async def _start(self, application: web.Application) -> None: + del application + trace = self._trace_config() + per_backend_limit = max(1, _MAXIMUM_CONNECTIONS // len(self.pool.backends)) + self.connector = TCPConnector( + limit=_MAXIMUM_CONNECTIONS, + limit_per_host=per_backend_limit, + keepalive_timeout=_KEEPALIVE_TIMEOUT_SECONDS, + force_close=False, + socket_factory=_create_socket, + ) + self.session = ClientSession( + connector=self.connector, + timeout=ClientTimeout( + total=_READ_TIMEOUT_SECONDS, + sock_connect=_CONNECT_TIMEOUT_SECONDS, + sock_read=_READ_TIMEOUT_SECONDS, + ), + auto_decompress=False, + raise_for_status=False, + trace_configs=[trace], + ) + + async def _stop(self, application: web.Application) -> None: + del application + if self.session is not None: + await self.session.close() + + def _trace_config(self) -> TraceConfig: + trace = TraceConfig() + + async def connection_opened(*_: object) -> None: + self.metrics.connection_opened += 1 + + async def connection_reused(*_: object) -> None: + self.metrics.connection_reused += 1 + + trace.on_connection_create_end.append(connection_opened) + trace.on_connection_reuseconn.append(connection_reused) + return trace + + async def _health(self, request: web.Request) -> web.Response: + del request + results = await asyncio.gather( + *(self._probe_backend(index, backend) for index, backend in enumerate(self.pool.backends)) + ) + status = 200 if all(results) else 503 + return web.json_response( + { + "status": "ok" if status == 200 else "unavailable", + "backends_ready": sum(results), + "backends": len(results), + }, + status=status, + ) + + async def _probe_backend(self, index: int, backend: _Backend) -> bool: + session = self._require_session() + self.pool.acquire_index(index) + try: + async with session.get( + f"{backend.origin}{self.health_path}", + timeout=ClientTimeout(total=_CONNECT_TIMEOUT_SECONDS), + trace_request_ctx={"backend_index": index}, + ) as response: + await response.read() + return response.status == 200 + except (ClientError, asyncio.TimeoutError): + return False + finally: + self.pool.release(index) + + async def _metrics(self, request: web.Request) -> web.Response: + del request + connector = self._require_connector() + return web.json_response(self.metrics.snapshot(self.pool, connector)) + + async def _forward(self, request: web.Request) -> web.StreamResponse: + started_at = time.monotonic() + self.metrics.requests += 1 + body = await request.read() + excluded: set[int] = set() + outcomes: list[str] = [] + attempt_limit = min(_MAXIMUM_UPSTREAM_ATTEMPTS, len(self.pool.backends)) + final_response: ClientResponse | None = None + final_index: int | None = None + overload_seen = False + overload_retry_after: str | None = None + + for attempt in range(attempt_limit): + index = self.pool.acquire(frozenset(excluded)) + backend = self.pool.backends[index] + self.metrics.attempts += 1 + self.metrics.backend_requests[index] += 1 + try: + response = await self._request_backend(request, backend, body, index) + except (ClientError, asyncio.TimeoutError) as error: + self.metrics.connection_failed += 1 + reason = "timeout" if isinstance(error, asyncio.TimeoutError) else "connection_error" + outcomes.append(reason) + self.pool.release(index) + excluded.add(index) + if attempt + 1 < attempt_limit: + self.metrics.retries[reason] += 1 + continue + if overload_seen: + headers = {} + retry_after = ( + overload_retry_after if overload_retry_after is not None else self.pool.retry_after_seconds + ) + if retry_after is not None: + headers["Retry-After"] = str(retry_after) + result = web.json_response({"error": "backend overloaded"}, status=429, headers=headers) + self._record_final(429, outcomes, started_at) + return result + result = web.json_response({"error": "backend unavailable"}, status=502) + self._record_final(502, outcomes, started_at) + return result + except BaseException: + self.pool.release(index) + raise + + outcomes.append(str(response.status)) + should_retry = response.status in _RETRYABLE_STATUS_CODES and attempt + 1 < attempt_limit + if should_retry: + reason = f"http_{response.status}" + self.metrics.retries[reason] += 1 + if response.status == 429: + overload_seen = True + overload_retry_after = response.headers.get("Retry-After", overload_retry_after) + await _discard_response(response) + self.pool.release(index) + excluded.add(index) + continue + final_response = response + final_index = index + break + + if ( + final_response is None or final_index is None + ): # pragma: no cover - loop always returns or selects a response + raise RuntimeError("proxy attempt loop produced no final response") + try: + result = await self._stream_response(request, final_response) + self._record_final(final_response.status, outcomes, started_at) + return result + finally: + final_response.release() + self.pool.release(final_index) + + async def _request_backend( + self, + request: web.Request, + backend: _Backend, + body: bytes, + index: int, + ) -> ClientResponse: + session = self._require_session() + headers = _filter_request_headers(request.headers) + return await session.request( + request.method, + f"{backend.origin}{request.rel_url}", + data=body, + headers=headers, + allow_redirects=False, + trace_request_ctx={"backend_index": index}, + ) + + async def _stream_response(self, request: web.Request, response: ClientResponse) -> web.StreamResponse: + headers = _filter_response_headers(response.headers) + has_retry_after = any(name.casefold() == "retry-after" for name in headers) + if response.status == 429 and not has_retry_after and self.pool.retry_after_seconds is not None: + headers["Retry-After"] = str(self.pool.retry_after_seconds) + downstream = web.StreamResponse(status=response.status, reason=response.reason, headers=headers) + await downstream.prepare(request) + if request.method != "HEAD": + async for chunk in response.content.iter_chunked(_STREAM_CHUNK_BYTES): + await downstream.write(chunk) + await downstream.write_eof() + return downstream + + def _record_final(self, status: int, outcomes: list[str], started_at: float) -> None: + self.metrics.final_statuses[status] += 1 + _LOGGER.info( + "slurm_proxy_request %s", + json.dumps( + { + "attempts": len(outcomes), + "final_status": status, + "upstream_outcomes": outcomes, + "elapsed_seconds": round(time.monotonic() - started_at, 6), + "failover": len(outcomes) > 1, + }, + separators=(",", ":"), + sort_keys=True, + ), + ) + + def _require_session(self) -> ClientSession: + if self.session is None: + raise RuntimeError("proxy session is not started") + return self.session + + def _require_connector(self) -> TCPConnector: + if self.connector is None: + raise RuntimeError("proxy connector is not started") + return self.connector + + +async def _discard_response(response: ClientResponse) -> None: + response.close() + + +def _filter_request_headers(headers: Mapping[str, str]) -> dict[str, str]: + hop_headers = _get_hop_headers(tuple(headers.items())) + return { + name: value + for name, value in headers.items() + if name.casefold() not in hop_headers and name.casefold() not in {"host", "content-length"} + } + + +def _filter_response_headers(headers: Mapping[str, str]) -> dict[str, str]: + hop_headers = _get_hop_headers(tuple(headers.items())) + return {name: value for name, value in headers.items() if name.casefold() not in hop_headers} + + +def _create_socket( + address_info: tuple[int, int, int, str, tuple[str, int] | tuple[str, int, int, int]], +) -> socket.socket: + family, socket_type, protocol, _, _ = address_info + connection = socket.socket(family=family, type=socket_type, proto=protocol) + connection.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + return connection + + +def main(arguments: Sequence[str] | None = None) -> int: + """Serve one loopback logical endpoint until the step is terminated.""" + parser = argparse.ArgumentParser(prog="data-designer-slurm-proxy") + parser.add_argument("--listen-port", required=True, type=int) + parser.add_argument("--backend", action="append", required=True) + parser.add_argument("--allowed-host", action="append") + parser.add_argument("--retry-after-seconds", type=int) + parser.add_argument("--health-path", default="/health") + parsed = parser.parse_args(arguments) + try: + allowed_hosts = frozenset(validate_host_name(value) for value in (parsed.allowed_host or ("127.0.0.1",))) + except ValueError as error: + parser.error(str(error)) + backends = tuple(_parse_backend(value, allowed_hosts) for value in parsed.backend) + if not 1 <= parsed.listen_port <= 65535: + parser.error("listen port must be between 1 and 65535") + if parsed.retry_after_seconds is not None and parsed.retry_after_seconds <= 0: + parser.error("retry-after seconds must be positive") + if not parsed.health_path.startswith("/") or parsed.health_path.startswith("//"): + parser.error("health path must be absolute") + proxy = _ProxyApplication(backends, parsed.retry_after_seconds, health_path=parsed.health_path) + web.run_app( + proxy.create(), + host="127.0.0.1", + port=parsed.listen_port, + print=None, + access_log=None, + keepalive_timeout=_KEEPALIVE_TIMEOUT_SECONDS, + shutdown_timeout=10.0, + handler_cancellation=True, + ) + return 0 + + +def _parse_backend(value: str, allowed_hosts: frozenset[str] = frozenset({"127.0.0.1"})) -> _Backend: + parsed: SplitResult = urlsplit(value) + host = parsed.hostname + normalized_allowed_hosts = frozenset(allowed_host.casefold() for allowed_host in allowed_hosts) + try: + port = parsed.port + except ValueError as error: + raise argparse.ArgumentTypeError("backend port is invalid") from error + if ( + parsed.scheme != "http" + or host is None + or host.casefold() not in normalized_allowed_hosts + or parsed.username is not None + or parsed.password is not None + or parsed.path not in {"", "/"} + or parsed.query + or parsed.fragment + or port is None + ): + raise argparse.ArgumentTypeError("backends must be allowed HTTP origins") + return _Backend(host, port) + + +def _get_hop_headers(headers: Sequence[tuple[str, str]]) -> frozenset[str]: + nominated = { + token.strip().casefold() + for name, value in headers + if name.casefold() == "connection" + for token in value.split(",") + if token.strip() + } + return _HOP_HEADERS | nominated + + +if __name__ == "__main__": # pragma: no cover - exercised as a managed step + raise SystemExit(main()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py new file mode 100644 index 000000000..53119456e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read and validate #876-owned client and candidate records.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TypeVar + +from pydantic import ValidationError + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ContractRecord +from data_designer.slurm.integration import IntegrationContractError, PlanStateValidator +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.state import AttemptManifest, CandidateOutputManifest +from data_designer.slurm.state.filesystem import open_verified_directory, read_regular_text + +_CLIENT_RESULT_NAME = "client-result.json" +_CANDIDATE_NAME = "output-manifest.json" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_RecordT = TypeVar("_RecordT", bound=ContractRecord) + + +def load_nonempty_client_candidate( + context: AllocationContext, + attempt: AttemptManifest, + *, + attempt_directory: Path | None = None, +) -> tuple[ClientResult, CandidateOutputManifest]: + """Load a nonempty semantic client result and its digest-bound candidate.""" + record_directory = context.attempt_directory if attempt_directory is None else attempt_directory + client_result = _read_record(record_directory, _CLIENT_RESULT_NAME, ClientResult) + if client_result.outcome not in {ClientOutcome.COMPLETE, ClientOutcome.PARTIAL} or not client_result.actual_records: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.CLIENT_FAILED, + f"client generation finished with outcome {client_result.outcome.value!r}", + ) + candidate = _read_record(record_directory, _CANDIDATE_NAME, CandidateOutputManifest) + try: + PlanStateValidator(context.plan).validate_client_candidate( + context.shard, + attempt, + client_result, + candidate, + ) + except IntegrationContractError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "client result and candidate do not match the resolved plan", + ) from error + return client_result, candidate + + +def _read_record(directory: Path, name: str, record_type: type[_RecordT]) -> _RecordT: + display_path = directory / name + try: + with open_verified_directory(directory, require_private=True) as descriptor: + content = read_regular_text( + descriptor, + name, + display_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + return record_type.model_validate_json(content) + except (OSError, UnicodeError, ValidationError, ValueError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + f"runtime result record {name!r} is unavailable or invalid", + ) from error + + +__all__ = ["load_nonempty_client_candidate"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/scratch.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/scratch.sh new file mode 100644 index 000000000..f5af8a2b3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/scratch.sh @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +DD_SCRATCH_ROOT=${DD_SCRATCH_ROOT:-} +DD_SCRATCH_CONTAINER_ROOT=/run/data-designer-slurm +export DD_SCRATCH_CONTAINER_ROOT + +dd_select_scratch_parent() { + local candidate resolved + command -v find >/dev/null || return 69 + for candidate in "${SLURM_TMPDIR:-}" "${TMPDIR:-}" /tmp; do + [[ ${candidate} =~ ^/[A-Za-z0-9._/-]+$ && ${candidate} != / ]] || continue + resolved=$(cd -P -- "${candidate}" 2>/dev/null && pwd -P) || continue + [[ ${resolved} =~ ^/[A-Za-z0-9._/-]+$ && ${resolved} != / ]] || continue + [[ -d ${resolved} && -w ${resolved} && -x ${resolved} ]] || continue + DD_SCRATCH_PARENT=${resolved} + return 0 + done + printf '%s\n' 'no safe writable local scratch directory is available' >&2 + return 73 +} + +dd_initialize_local_scratch() { + local expected=${1:-} + local task_id=${SLURM_ARRAY_TASK_ID:-0} + [[ ${SLURM_JOB_ID:-} =~ ^[1-9][0-9]*$ && ${task_id} =~ ^[0-9]+$ ]] || return 64 + dd_select_scratch_parent || return + local root="${DD_SCRATCH_PARENT}/data-designer-slurm-${SLURM_JOB_ID}-${task_id}" + [[ -z ${expected} || ${expected} == "${root}" ]] || { + printf '%s\n' 'allocation nodes resolved inconsistent local scratch paths' >&2 + return 73 + } + dd_ensure_private_scratch_directory "${root}" || return + DD_SCRATCH_ROOT=${root} + export DD_SCRATCH_ROOT + local directory + for directory in \ + home cache enroot enroot/cache enroot/config enroot/data enroot/tmp; do + dd_ensure_private_scratch_directory "${root}/${directory}" || { + dd_remove_local_scratch "${root}" >/dev/null 2>&1 || true + return 73 + } + done + export HOME="${root}/home" + export ENROOT_CACHE_PATH="${root}/enroot/cache" + export ENROOT_CONFIG_PATH="${root}/enroot/config" + export ENROOT_DATA_PATH="${root}/enroot/data" + export ENROOT_TEMP_PATH="${root}/enroot/tmp" +} + +dd_ensure_private_scratch_directory() { + local path=$1 unsafe + if [[ -e ${path} || -L ${path} ]]; then + [[ -d ${path} && ! -L ${path} && -O ${path} ]] || return 73 + else + (umask 077 && mkdir -- "${path}") || return 73 + fi + chmod 0700 "${path}" || return 73 + [[ -d ${path} && ! -L ${path} && -O ${path} ]] || return 73 + if unsafe=$(find "${path}" -prune -type d -perm /0077 -print 2>/dev/null); then + : + elif unsafe=$(find "${path}" -prune -type d -perm +0077 -print 2>/dev/null); then + : + else + return 69 + fi + [[ -r ${path} && -w ${path} && -x ${path} && -z ${unsafe} ]] || return 73 +} + +dd_extract_runtime_bundle() { + local archive=$1 + local expected_sha256=$2 + local actual_sha256 staging runtime + actual_sha256=$(sha256sum < "${archive}") || return 65 + [[ ${actual_sha256%% *} == "${expected_sha256}" ]] || return 65 + staging="${DD_SCRATCH_ROOT}/.runtime-${expected_sha256}-$$" + runtime="${DD_SCRATCH_ROOT}/runtime" + [[ ! -e ${staging} && ! -L ${staging} && ! -e ${runtime} && ! -L ${runtime} ]] || return 73 + dd_ensure_private_scratch_directory "${staging}" || return + (umask 077 && tar -xzf "${archive}" -C "${staging}") || return 65 + [[ -f ${staging}/entrypoint.sh && ! -L ${staging}/entrypoint.sh ]] || return 65 + [[ -f ${staging}/scratch.sh && ! -L ${staging}/scratch.sh ]] || return 65 + mv -- "${staging}" "${runtime}" || return 73 +} + +dd_stage_allocation_runtime() { + [[ ${SLURM_JOB_NUM_NODES:-} =~ ^[1-9][0-9]*$ ]] || return 64 + local node_command='set -Eeuo pipefail +archive=$1 +expected_sha256=$2 +expected_root=$3 +actual_sha256=$(sha256sum < "${archive}") +[[ ${actual_sha256%% *} == "${expected_sha256}" ]] +{ + printf "set -Eeuo pipefail\\n" + tar -xOf "${archive}" scratch.sh + printf "\ndd_initialize_local_scratch %q\ndd_extract_runtime_bundle %q %q\n" \ + "${expected_root}" "${archive}" "${expected_sha256}" +} | /bin/bash -s -- "${archive}" "${expected_sha256}" "${expected_root}"' + srun \ + "--nodes=${SLURM_JOB_NUM_NODES}" \ + "--ntasks=${SLURM_JOB_NUM_NODES}" \ + --ntasks-per-node=1 \ + --cpus-per-task=1 \ + --gres=none \ + --exact \ + --overlap \ + --unbuffered \ + --kill-on-bad-exit=1 \ + --export=ALL \ + -- \ + /bin/bash -c "${node_command}" dd-scratch \ + "${DD_RUNTIME_ARCHIVE}" "${DD_RUNTIME_SHA256}" "${DD_SCRATCH_ROOT}" +} + +dd_cleanup_allocation_scratch() { + local status=0 + if [[ ${SLURM_JOB_NUM_NODES:-} =~ ^[1-9][0-9]*$ ]] && command -v srun >/dev/null; then + local node_command='set -Eeuo pipefail +archive=$1 +expected_sha256=$2 +expected_root=$3 +actual_sha256=$(sha256sum < "${archive}") +[[ ${actual_sha256%% *} == "${expected_sha256}" ]] +{ + printf "set -Eeuo pipefail\\n" + tar -xOf "${archive}" scratch.sh + printf "\ndd_remove_local_scratch %q\n" "${expected_root}" +} | /bin/bash -s -- "${archive}" "${expected_sha256}" "${expected_root}"' + srun \ + "--nodes=${SLURM_JOB_NUM_NODES}" \ + "--ntasks=${SLURM_JOB_NUM_NODES}" \ + --ntasks-per-node=1 \ + --cpus-per-task=1 \ + --gres=none \ + --exact \ + --overlap \ + --unbuffered \ + --kill-on-bad-exit=1 \ + --time=00:01:00 \ + --wait=0 \ + --export=ALL \ + -- \ + /bin/bash -c "${node_command}" dd-scratch \ + "${DD_RUNTIME_ARCHIVE}" "${DD_RUNTIME_SHA256}" "${DD_SCRATCH_ROOT}" \ + >/dev/null 2>&1 || status=$? + fi + dd_remove_local_scratch "${DD_SCRATCH_ROOT}" || status=$? + return "${status}" +} + +dd_remove_local_scratch() { + local path=${1:-} + local task_id=${SLURM_ARRAY_TASK_ID:-0} + local expected_name parent + [[ ${SLURM_JOB_ID:-} =~ ^[1-9][0-9]*$ && ${task_id} =~ ^[0-9]+$ ]] || return 64 + expected_name="data-designer-slurm-${SLURM_JOB_ID}-${task_id}" + [[ ${path} == /*/${expected_name} && ${path} != *//* && ${path} != */../* && ${path} != */./* ]] || return 73 + parent=$(cd -P -- "${path%/*}" 2>/dev/null && pwd -P) || return 73 + [[ ${path} == "${parent}/${expected_name}" ]] || return 73 + if [[ ! -e ${path} && ! -L ${path} ]]; then + return 0 + fi + [[ -d ${path} && ! -L ${path} && -O ${path} ]] || return 73 + rm -rf -- "${path}" +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py new file mode 100644 index 000000000..f986a49cc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compose local and distributed server entries for the runtime manifest.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Literal + +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.distributed import build_node_worker_spec, build_vllm_process_command +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.manifest import RuntimeProbeSpec, RuntimeStepSpec +from data_designer.slurm.runtime.models import AllocationContext, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import encode_node_worker_spec +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.steps import build_vllm_command +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess + + +def build_server_steps( + deployments: tuple[ResolvedVllmServerDeployment, ...], + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> tuple[RuntimeStepSpec, ...]: + """Build remote preflight entries followed by all serving entries.""" + group_by_id = { + deployment.deployment_id: deployment.authored.resources.shared_node for deployment in context.plan.deployments + } + groups: dict[tuple[str, str], list[ResolvedVllmServerDeployment]] = {} + for deployment in deployments: + group = group_by_id[deployment.deployment_id] + key = ("deployment", deployment.deployment_id) if group is None else ("shared", group) + groups.setdefault(key, []).append(deployment) + preflight_steps: list[RuntimeStepSpec] = [] + serving_steps: list[RuntimeStepSpec] = [] + for members in groups.values(): + first, *others = members + if others or _requires_server_preflight(first, context): + preflight_steps.append( + build_distributed_server_step( + first, + context, + runtime_container_root, + log_directory, + layout, + operation="preflight", + colocated_deployments=tuple(others), + ) + ) + if others: + serving_steps.append( + build_distributed_server_step( + first, + context, + runtime_container_root, + log_directory, + layout, + operation="serve", + colocated_deployments=tuple(others), + ) + ) + else: + serving_steps.extend( + _build_deployment_serving_steps( + first, + context, + runtime_container_root, + log_directory, + layout, + ) + ) + return tuple(preflight_steps + serving_steps) + + +def build_distributed_server_step( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, + *, + operation: Literal["preflight", "serve"], + colocated_deployments: tuple[ResolvedVllmServerDeployment, ...] = (), +) -> RuntimeStepSpec: + """Build one coordinated node-worker step for a deployment.""" + literal_environment, secret_environment, environment_prefixes = _server_environment( + deployment, + runtime_container_root, + ) + worker_spec = build_node_worker_spec(deployment, context.plan, layout, colocated_deployments=colocated_deployments) + readiness = ( + tuple( + RuntimeProbeSpec( + host="127.0.0.1" if len(layout.node_hosts) == 1 else layout.get_host(probe.node_index), + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds + + next( + process.launch_delay_seconds + for process in member.processes + if process.node_index == probe.node_index and process.http_port == probe.port + ), + ) + for member in (deployment, *colocated_deployments) + for probe in member.readiness_probes + ) + if operation == "serve" + else () + ) + step_id = f"{deployment.deployment_id}-{operation}" + if worker_spec.allow_gpu_sharing: + launch_order = { + process.readiness.port: index + for index, process in enumerate(worker_spec.nodes[0].processes) + if process.readiness is not None + } + readiness = tuple(sorted(readiness, key=lambda probe: launch_order[probe.port])) + return RuntimeStepSpec( + step_id=step_id, + role=RuntimeStepRole.SERVER if operation == "serve" else RuntimeStepRole.SERVER_PREFLIGHT, + execution="container", + image_path=deployment.image.path, + command=( + "python3", + f"{runtime_container_root}/data_designer/slurm/runtime/node_worker.py", + operation, + "--spec", + encode_node_worker_spec(worker_spec), + ), + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple( + range(context.plan.resolved_gpus_per_node if colocated_deployments else deployment.gpus_per_node) + ), + gpu_binding="none" if colocated_deployments else "mask", + node_hosts=tuple(layout.get_host(index) for index in deployment.node_indices), + kill_on_bad_exit=True, + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=_server_container_environment(deployment), + stdout_path=(log_directory / f"{step_id}.out").as_posix(), + stderr_path=(log_directory / f"{step_id}.err").as_posix(), + readiness=readiness, + ) + + +def _build_deployment_serving_steps( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> tuple[RuntimeStepSpec, ...]: + if len(deployment.node_indices) > 1 or len(deployment.processes) > 1: + return ( + build_distributed_server_step( + deployment, + context, + runtime_container_root, + log_directory, + layout, + operation="serve", + ), + ) + return tuple( + _build_local_server_step( + deployment, + process, + context, + runtime_container_root, + log_directory, + layout, + ) + for process in deployment.processes + ) + + +def _build_local_server_step( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + context: AllocationContext, + runtime_container_root: str, + log_directory: Path, + layout: AllocationLayout, +) -> RuntimeStepSpec: + if process.pipeline_parallel != 1 or process.http_port is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "single-node runtime received a distributed vLLM process", + ) + literal_environment, secret_environment, environment_prefixes = _server_environment( + deployment, + runtime_container_root, + ) + probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) + probe_host = "127.0.0.1" if len(layout.node_hosts) == 1 else layout.get_host(process.node_index) + command = ( + build_vllm_command(deployment, process, context.plan) + if len(layout.node_hosts) == 1 + else build_vllm_process_command(deployment, process, context.plan, layout) + ) + return RuntimeStepSpec( + step_id=process.process_id, + role=RuntimeStepRole.SERVER, + execution="container", + image_path=deployment.image.path, + command=command, + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple(process.gpu_indices), + node_hosts=(layout.get_host(process.node_index),), + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=_server_container_environment(deployment), + stdout_path=(log_directory / f"{process.process_id}.out").as_posix(), + stderr_path=(log_directory / f"{process.process_id}.err").as_posix(), + launch_delay_seconds=process.launch_delay_seconds, + readiness=( + RuntimeProbeSpec( + host=probe_host, + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ), + ), + ) + + +def _server_environment( + deployment: ResolvedVllmServerDeployment, + runtime_container_root: str, +) -> tuple[dict[str, str], dict[str, str], dict[str, str]]: + literal_environment: dict[str, str] = {"LC_ALL": "C", "PYTHONPATH": runtime_container_root} + secret_environment: dict[str, str] = {} + environment_prefixes: dict[str, str] = {} + for name, binding in deployment.launch_policy.environment.items(): + if isinstance(binding, LiteralEnvironmentBinding): + literal_environment[name] = binding.value + elif isinstance(binding, SecretRef): + secret_environment[name] = binding.environment + else: # pragma: no cover - persisted contracts reject unknown bindings + raise AssertionError(f"unhandled environment binding: {type(binding)!r}") + if "PYTHONPATH" in secret_environment: + literal_environment.pop("PYTHONPATH") + environment_prefixes["PYTHONPATH"] = runtime_container_root + elif "PYTHONPATH" in deployment.launch_policy.environment: + literal_environment["PYTHONPATH"] = os.pathsep.join((runtime_container_root, literal_environment["PYTHONPATH"])) + policy = deployment.launch_policy.queue_backpressure + literal_environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(policy.max_waiting_requests) + literal_environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if policy.retry_after_seconds is None else str(policy.retry_after_seconds) + ) + return literal_environment, secret_environment, environment_prefixes + + +def _server_container_environment(deployment: ResolvedVllmServerDeployment) -> tuple[str, ...]: + return tuple( + sorted( + { + *deployment.launch_policy.environment, + "PYTHONPATH", + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + } + ) + ) + + +def _requires_server_preflight( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, +) -> bool: + return any(node_index != context.plan.client.host_node_index for node_index in deployment.node_indices) + + +__all__ = ["build_distributed_server_step", "build_server_steps"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/signals.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/signals.py new file mode 100644 index 000000000..e9e3a35a0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/signals.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Centralized termination-signal coordination for the allocation runtime.""" + +from __future__ import annotations + +import os +import signal +import threading +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from types import FrameType + +_TERMINATION_SIGNALS = (signal.SIGINT, signal.SIGTERM) + + +class TerminationSignalCoordinator: + """Install, defer, and block allocation termination handling.""" + + def __init__(self) -> None: + self._lock = threading.RLock() + self._cleanup: Callable[[], None] | None = None + self._defer_depth = 0 + self._deferred_signal: signal.Signals | None = None + self._interrupted = False + + @contextmanager + def interrupt_on_termination(self, cleanup: Callable[[], None]) -> Iterator[None]: + """Install one cleanup-first termination handler for an allocation run.""" + if self._cleanup is not None: + raise RuntimeError("termination signal coordination is already active") + previous = {selected: signal.getsignal(selected) for selected in _TERMINATION_SIGNALS} + self._cleanup = cleanup + self._interrupted = False + for selected in _TERMINATION_SIGNALS: + signal.signal(selected, self._handle_termination) + try: + yield + finally: + for selected, handler in previous.items(): + signal.signal(selected, handler) + self._cleanup = None + self._deferred_signal = None + self._defer_depth = 0 + + @contextmanager + def defer_termination(self) -> Iterator[None]: + """Delay termination delivery until a child process is registered.""" + with self._lock: + self._defer_depth += 1 + try: + yield + finally: + deferred = self._finish_defer() + if deferred is not None: + os.kill(os.getpid(), deferred) + + @contextmanager + def block_termination(self) -> Iterator[None]: + """Block termination delivery while owned children are being stopped.""" + previous = signal.pthread_sigmask(signal.SIG_BLOCK, _TERMINATION_SIGNALS) + try: + yield + finally: + signal.pthread_sigmask(signal.SIG_SETMASK, previous) + + def _finish_defer(self) -> signal.Signals | None: + with self._lock: + self._defer_depth -= 1 + if self._defer_depth or self._interrupted: + return None + deferred = self._deferred_signal + self._deferred_signal = None + return deferred + + def _handle_termination(self, signum: int, frame: FrameType | None) -> None: + del frame + selected = signal.Signals(signum) + with self._lock: + if self._defer_depth: + if self._deferred_signal is None: + self._deferred_signal = selected + return + if self._interrupted: + return + self._interrupted = True + cleanup = self._cleanup + try: + if cleanup is not None: + cleanup() + finally: + raise KeyboardInterrupt + + +__all__ = ["TerminationSignalCoordinator"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh new file mode 100644 index 000000000..c4ecbaaec --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +dd_validate_environment_name() { + [[ $1 =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] +} + +dd_require_plan_secrets() { + local name + for name in "${DD_REQUIRED_SECRET_NAMES[@]+"${DD_REQUIRED_SECRET_NAMES[@]}"}"; do + dd_validate_environment_name "${name}" || return 64 + [[ ${!name+x} ]] || { + printf 'required secret environment %q is unavailable\n' "${name}" >&2 + return 78 + } + done +} + +dd_scope_control_environment() { + local name + for name in "${DD_ALL_SECRET_NAMES[@]+"${DD_ALL_SECRET_NAMES[@]}"}"; do + unset "${name}" + done + unset CUDA_VISIBLE_DEVICES + export LC_ALL=C + export PYTHONPATH=${DD_RUNTIME_ROOT} +} + +dd_materialize_step_environment() { + local index name source value prefix current + local -a secret_values=() + for ((index = 0; index < ${#DD_STEP_SECRET_ENV[@]}; index += 2)); do + name=${DD_STEP_SECRET_ENV[index]} + source=${DD_STEP_SECRET_ENV[index + 1]} + dd_validate_environment_name "${name}" && dd_validate_environment_name "${source}" || return 64 + [[ ${!source+x} ]] || return 78 + secret_values+=("${!source}") + done + for name in "${DD_ALL_SECRET_NAMES[@]+"${DD_ALL_SECRET_NAMES[@]}"}"; do + unset "${name}" + done + unset CUDA_VISIBLE_DEVICES + for ((index = 0; index < ${#DD_STEP_LITERAL_ENV[@]}; index += 2)); do + name=${DD_STEP_LITERAL_ENV[index]} + value=${DD_STEP_LITERAL_ENV[index + 1]} + dd_validate_environment_name "${name}" || return 64 + printf -v "${name}" '%s' "${value}" + export "${name}" + done + for ((index = 0; index < ${#DD_STEP_SECRET_ENV[@]}; index += 2)); do + name=${DD_STEP_SECRET_ENV[index]} + printf -v "${name}" '%s' "${secret_values[index / 2]}" + export "${name}" + done + for ((index = 0; index < ${#DD_STEP_ENV_PREFIXES[@]}; index += 2)); do + name=${DD_STEP_ENV_PREFIXES[index]} + prefix=${DD_STEP_ENV_PREFIXES[index + 1]} + dd_validate_environment_name "${name}" || return 64 + current=${!name} + printf -v "${name}" '%s' "${prefix}:${current}" + export "${name}" + done + if [[ -n ${DD_STEP_VISIBLE_GPUS} ]]; then + export CUDA_VISIBLE_DEVICES=${DD_STEP_VISIBLE_GPUS} + fi +} + +dd_build_srun_command() { + local index gpu_mask=0 container_names= node_list= + local container_mounts="${DD_SCRATCH_ROOT}:${DD_SCRATCH_CONTAINER_ROOT}" + ((${#DD_STEP_NODE_HOSTS[@]})) || return 64 + printf -v node_list '%s,' "${DD_STEP_NODE_HOSTS[@]}" + node_list=${node_list%,} + DD_STEP_VISIBLE_GPUS= + DD_SRUN_COMMAND=( + srun + "--nodes=${#DD_STEP_NODE_HOSTS[@]}" + "--ntasks=${#DD_STEP_NODE_HOSTS[@]}" + --ntasks-per-node=1 + "--nodelist=${node_list}" + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_STEP_CPUS}" + ) + if [[ ${DD_STEP_KILL_ON_BAD_EXIT} == true ]]; then + DD_SRUN_COMMAND+=(--kill-on-bad-exit=1) + fi + if ((${#DD_STEP_GPU_INDICES[@]})); then + if [[ ${DD_GPU_REQUEST_MODE} == gres ]]; then + DD_SRUN_COMMAND+=("--gpus-per-task=${#DD_STEP_GPU_INDICES[@]}") + if [[ ${DD_STEP_GPU_BINDING:-mask} == none ]]; then + DD_SRUN_COMMAND+=(--gpu-bind=none) + elif [[ ${DD_STEP_GPU_BINDING:-mask} == mask ]]; then + for index in "${DD_STEP_GPU_INDICES[@]+"${DD_STEP_GPU_INDICES[@]}"}"; do + ((gpu_mask |= 1 << index)) + done + printf -v gpu_mask '0x%x' "${gpu_mask}" + DD_SRUN_COMMAND+=("--gpu-bind=mask_gpu:${gpu_mask}") + else + return 64 + fi + else + local visible_gpus + printf -v visible_gpus '%s,' "${DD_STEP_GPU_INDICES[@]}" + DD_STEP_VISIBLE_GPUS=${visible_gpus%,} + DD_STEP_CONTAINER_ENV+=(CUDA_VISIBLE_DEVICES) + fi + else + DD_SRUN_COMMAND+=(--gres=none) + fi + if [[ ${DD_STEP_EXECUTION} == container ]]; then + [[ -n ${DD_STEP_IMAGE} ]] || return 64 + DD_SRUN_COMMAND+=("--container-image=${DD_STEP_IMAGE}") + [[ -z ${DD_CONTAINER_MOUNTS} ]] || container_mounts="${DD_CONTAINER_MOUNTS},${container_mounts}" + DD_SRUN_COMMAND+=("--container-mounts=${container_mounts}") + if ((${#DD_STEP_CONTAINER_ENV[@]})); then + printf -v container_names '%s,' "${DD_STEP_CONTAINER_ENV[@]}" + DD_SRUN_COMMAND+=("--container-env=${container_names%,}") + fi + elif [[ ${DD_STEP_EXECUTION} != native || -n ${DD_STEP_IMAGE} || ${#DD_STEP_CONTAINER_ENV[@]} -ne 0 ]]; then + return 64 + fi +} + +dd_start_step() { + local manifest=$1 + local step_id=$2 + dd_read_step "${manifest}" "${step_id}" + dd_build_srun_command + ( + dd_materialize_step_environment + exec "${DD_SRUN_COMMAND[@]}" -- "${DD_STEP_COMMAND[@]}" + ) >"${DD_STEP_STDOUT}" 2>"${DD_STEP_STDERR}" & + DD_LAST_PID=$! +} + +dd_run_step() { + local manifest=$1 + local step_id=$2 + dd_start_step "${manifest}" "${step_id}" + local pid=${DD_LAST_PID} + local index=${#DD_MANAGED_PIDS[@]} + local status=0 + DD_MANAGED_PIDS+=("${pid}") + wait "${pid}" || status=$? + unset "DD_MANAGED_PIDS[${index}]" + DD_MANAGED_PIDS=("${DD_MANAGED_PIDS[@]+"${DD_MANAGED_PIDS[@]}"}") + if ((status != 0)); then + printf 'runtime step failed: step_id=%q status=%d stdout=%q stderr=%q\n' \ + "${step_id}" "${status}" "${DD_STEP_STDOUT}" "${DD_STEP_STDERR}" >&2 + fi + return "${status}" +} + +dd_run_control_phase() { + local operation=$1 + shift + local -a command=( + srun + --nodes=1 + --ntasks=1 + "--nodelist=${DD_CLIENT_HOST}" + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_CLIENT_CPUS}" + --gres=none + -- + "${DD_CLIENT_PYTHON}" + -m + data_designer.slurm.runtime.entrypoint + "${operation}" + --plan + "${DD_PLAN_PATH}" + --attempt-dir + "${DD_ATTEMPT_PATH}" + "$@" + ) + ( + dd_scope_control_environment + exec "${command[@]}" + ) +} diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py new file mode 100644 index 000000000..b1b8f8a3c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -0,0 +1,577 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build every allocation-local process as one structured ``srun`` step.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Literal, Protocol + +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) +from data_designer.slurm.contracts import validate_absolute_path +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.paths import ALLOCATION_SCRATCH_CONTAINER_ROOT, get_container_path +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess +from data_designer.slurm.state import AttemptManifest + +_SLURM_ENVIRONMENT_NAMES = ( + "DD_SCRATCH_ROOT", + "SLURM_ARRAY_JOB_ID", + "SLURM_ARRAY_TASK_ID", + "SLURM_CLUSTER_NAME", + "SLURM_CONF", + "SLURM_CPUS_ON_NODE", + "SLURM_JOB_CPUS_PER_NODE", + "SLURM_JOB_GPUS", + "SLURM_JOB_ID", + "SLURM_JOB_NODELIST", + "SLURM_JOB_NUM_NODES", + "SLURM_NODEID", + "SLURM_PROCID", + "SLURM_SUBMIT_DIR", +) + + +class ClientStepBuilder(Protocol): + """Construct #876-owned client commands for the common runtime runner.""" + + def build_preflight_step( + self, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + """Build the zero-GPU client preflight step.""" + ... + + def build_generation_step( + self, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + """Build the zero-GPU client generation step.""" + ... + + +class DefaultClientStepBuilder: + """Invoke the allocation client worker supplied by #876.""" + + def build_preflight_step( + self, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + """Build a deterministic client-worker preflight command.""" + return self._build_step( + "client-preflight", + RuntimeStepRole.CLIENT_PREFLIGHT, + "preflight", + plan, + shard, + attempt, + attempt_directory, + endpoints, + source_environment, + retry_resume_mode, + ) + + def build_generation_step( + self, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + """Build a deterministic client-worker generation command.""" + return self._build_step( + "client-generation", + RuntimeStepRole.CLIENT, + "run", + plan, + shard, + attempt, + attempt_directory, + endpoints, + source_environment, + retry_resume_mode, + ) + + @staticmethod + def _build_step( + step_id: str, + role: RuntimeStepRole, + operation: str, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + retry_resume_mode: Literal["never", "always"] | None, + ) -> RuntimeStep: + command = build_client_command( + operation, + plan, + shard, + attempt, + attempt_directory, + endpoints, + retry_resume_mode, + ) + _, environment = _build_client_environment(plan, source_environment) + return _build_srun_step( + step_id=step_id, + role=role, + image_path=None, + command=command, + environment=environment, + container_environment=(), + plan=plan, + attempt_directory=attempt_directory, + ) + + +def build_client_command( + operation: str, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + retry_resume_mode: Literal["never", "always"] | None, +) -> tuple[str, ...]: + endpoint_arguments = tuple( + argument + for endpoint in endpoints + for argument in ("--endpoint", f"{endpoint.model_alias}=http://{endpoint.host}:{endpoint.port}/v1") + ) + return ( + plan.client.runtime.python_executable, + "-m", + "data_designer.slurm.client.worker", + operation, + "--plan", + plan_path(plan), + "--shard-id", + shard.shard_id, + "--attempt-id", + attempt.attempt_id, + "--attempt-dir", + attempt_directory.as_posix(), + *(() if retry_resume_mode is None else ("--resume-mode", retry_resume_mode)), + *endpoint_arguments, + ) + + +def _build_client_environment( + plan: ResolvedSlurmRunPlan, + source_environment: Mapping[str, str], +) -> tuple[tuple[str, ...], dict[str, str]]: + secret_names = _collect_client_secret_environment_names(plan) + environment = _base_environment(source_environment) + for name in secret_names: + try: + environment[name] = source_environment[name] + except KeyError: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + f"required client secret environment {name!r} is unavailable", + ) from None + try: + scratch_root = _validate_scratch_root(source_environment["DD_SCRATCH_ROOT"]) + except KeyError: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "allocation scratch is unavailable", + ) from None + environment["DD_SCRATCH_ROOT"] = scratch_root + environment["DATA_DESIGNER_SLURM_SCRATCH_ROOT"] = scratch_root + return secret_names, environment + + +def build_vllm_steps( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_root: Path, +) -> tuple[RuntimeStep, ...]: + """Build one structured server step per resolved vLLM process.""" + if any(item.authored.resources.shared_node is not None for item in plan.deployments): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "shared_node requires the coordinated allocation runtime", + ) + return tuple( + _build_vllm_step(deployment, process, plan, attempt_directory, source_environment, runtime_root) + for process in deployment.processes + ) + + +def build_endpoint_steps( + deployments: tuple[ResolvedVllmServerDeployment, ...], + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_proxy_path: Path, + runtime_root: Path | None = None, +) -> tuple[tuple[RuntimeStep, RuntimeEndpoint], ...]: + """Build one package-owned logical-endpoint process per deployment.""" + resolved_runtime_root = runtime_root if runtime_root is not None else runtime_proxy_path.parent + steps: list[tuple[RuntimeStep, RuntimeEndpoint]] = [] + for deployment in deployments: + steps.append( + _build_endpoint_step( + deployment, + plan, + attempt_directory, + source_environment, + runtime_proxy_path, + resolved_runtime_root, + ) + ) + return tuple(steps) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_proxy_path: Path, + runtime_root: Path, +) -> tuple[RuntimeStep, RuntimeEndpoint]: + endpoint = RuntimeEndpoint( + model_alias=deployment.model_alias, + served_model_name=deployment.served_model_name, + host="127.0.0.1", + port=deployment.logical_endpoint.port, + ) + command = build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) + environment = _base_environment(source_environment) + environment["PYTHONPATH"] = runtime_root.as_posix() + step = _build_srun_step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + image_path=None, + command=command, + environment=environment, + container_environment=(), + plan=plan, + attempt_directory=attempt_directory, + ) + return step, endpoint + + +def build_endpoint_command( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + runtime_proxy_path: Path, + port: int, + *, + backend_hosts: tuple[str, ...] | None = None, +) -> tuple[str, ...]: + selected_hosts = backend_hosts or ("127.0.0.1",) * len(deployment.backend_endpoints) + if len(selected_hosts) != len(deployment.backend_endpoints): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "endpoint backend hosts do not match the resolved deployment", + ) + backends = tuple( + f"http://{host}:{backend.port}" + for host, backend in zip(selected_hosts, deployment.backend_endpoints, strict=True) + ) + allowed_host_arguments = ( + tuple(argument for host in dict.fromkeys(selected_hosts) for argument in ("--allowed-host", host)) + if backend_hosts is not None + else () + ) + retry_after_seconds = deployment.launch_policy.queue_backpressure.retry_after_seconds + retry_arguments = ("--retry-after-seconds", str(retry_after_seconds)) if retry_after_seconds is not None else () + try: + host_proxy_path = validate_absolute_path(runtime_proxy_path.as_posix()) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime proxy path is invalid") from error + return ( + plan.client.runtime.python_executable, + host_proxy_path, + "--listen-port", + str(port), + "--health-path", + deployment.launch_policy.readiness_path, + *retry_arguments, + *allowed_host_arguments, + *(argument for backend in backends for argument in ("--backend", backend)), + ) + + +def plan_path(plan: ResolvedSlurmRunPlan) -> str: + """Return the canonical persisted resolved-plan path.""" + return str(Path(plan.authored_config.path).with_name("resolved-plan.json")) + + +def _build_vllm_step( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_root: Path, +) -> RuntimeStep: + _validate_local_vllm_process(process) + command = build_vllm_command(deployment, process, plan) + environment, container_environment = _build_vllm_environment( + deployment, + source_environment, + plan, + runtime_root, + ) + return _build_srun_step( + step_id=process.process_id, + role=RuntimeStepRole.SERVER, + image_path=deployment.image.path, + command=command, + environment=environment, + container_environment=container_environment, + plan=plan, + attempt_directory=attempt_directory, + gpu_indices=tuple(process.gpu_indices), + ) + + +def _validate_local_vllm_process(process: ResolvedVllmProcess) -> None: + if process.pipeline_parallel != 1 or process.node_index != 0 or process.http_port is None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "one-node runtime received a distributed vLLM process", + ) + + +def build_vllm_command( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + plan: ResolvedSlurmRunPlan, +) -> tuple[str, ...]: + if process.http_port is None: # pragma: no cover - validated before command construction + raise AssertionError("vLLM HTTP port is unavailable") + model = get_container_path(plan, deployment.model) if deployment.model.startswith("/") else deployment.model + command: tuple[str, ...] = ( + deployment.executable_path, + "serve", + model, + "--served-model-name", + deployment.served_model_name, + "--host", + "127.0.0.1", + "--port", + str(process.http_port), + "--tensor-parallel-size", + str(process.tensor_parallel), + "--middleware", + "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware", + ) + if deployment.launch_policy.enable_expert_parallel: + command += ("--enable-expert-parallel",) + if deployment.launch_policy.io_processor_plugin is not None: + command += ("--io-processor-plugin", deployment.launch_policy.io_processor_plugin) + return command + deployment.launch_policy.extra_args + + +def _build_vllm_environment( + deployment: ResolvedVllmServerDeployment, + source_environment: Mapping[str, str], + plan: ResolvedSlurmRunPlan, + runtime_root: Path, +) -> tuple[dict[str, str], tuple[str, ...]]: + environment = _base_environment(source_environment) + container_environment: list[str] = [] + for name, binding in deployment.launch_policy.environment.items(): + if isinstance(binding, LiteralEnvironmentBinding): + environment[name] = binding.value + elif isinstance(binding, SecretRef): + try: + environment[name] = source_environment[binding.environment] + except KeyError: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + f"required server secret environment {binding.environment!r} is unavailable", + ) from None + else: # pragma: no cover - persisted contracts reject unknown bindings + raise AssertionError(f"unhandled environment binding: {type(binding)!r}") + container_environment.append(name) + runtime_pythonpath = get_container_path(plan, runtime_root.as_posix()) + configured_pythonpath = environment.get("PYTHONPATH") + environment["PYTHONPATH"] = ( + os.pathsep.join((runtime_pythonpath, configured_pythonpath)) if configured_pythonpath else runtime_pythonpath + ) + queue_policy = deployment.launch_policy.queue_backpressure + environment[MAX_WAITING_REQUESTS_ENVIRONMENT] = str(queue_policy.max_waiting_requests) + environment[RETRY_AFTER_SECONDS_ENVIRONMENT] = ( + "" if queue_policy.retry_after_seconds is None else str(queue_policy.retry_after_seconds) + ) + container_environment.extend(("PYTHONPATH", MAX_WAITING_REQUESTS_ENVIRONMENT, RETRY_AFTER_SECONDS_ENVIRONMENT)) + return environment, tuple(container_environment) + + +def _build_srun_step( + *, + step_id: str, + role: RuntimeStepRole, + image_path: str | None, + command: tuple[str, ...], + environment: Mapping[str, str], + container_environment: tuple[str, ...], + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + gpu_indices: tuple[int, ...] = (), +) -> RuntimeStep: + srun_command = _build_srun_prefix(plan, image_path) + _add_srun_resources(srun_command, gpu_indices) + if image_path is not None: + _add_srun_container_options(srun_command, plan, container_environment, environment) + srun_command.extend(("--", *command)) + log_root = attempt_directory / "logs" + return RuntimeStep( + step_id=step_id, + role=role, + command=tuple(srun_command), + environment=environment, + stdout_path=log_root / f"{step_id}.out", + stderr_path=log_root / f"{step_id}.err", + ) + + +def _build_srun_prefix(plan: ResolvedSlurmRunPlan, image_path: str | None) -> list[str]: + command = [ + "srun", + "--nodes=1", + "--ntasks=1", + "--exact", + "--overlap", + "--unbuffered", + "--export=ALL", + f"--cpus-per-task={plan.client.authored.cpus}", + ] + if image_path is not None: + command.append(f"--container-image={image_path}") + return command + + +def _add_srun_resources(srun_command: list[str], gpu_indices: tuple[int, ...]) -> None: + if gpu_indices: + srun_command.extend( + ( + f"--gpus-per-task={len(gpu_indices)}", + f"--gpu-bind=mask_gpu:{_get_gpu_mask(gpu_indices):#x}", + ) + ) + else: + srun_command.append("--gres=none") + + +def _add_srun_container_options( + srun_command: list[str], + plan: ResolvedSlurmRunPlan, + container_environment: tuple[str, ...], + environment: Mapping[str, str], +) -> None: + mounts = _render_mounts(plan, environment.get("DD_SCRATCH_ROOT")) + if mounts: + srun_command.append(f"--container-mounts={mounts}") + if container_environment: + srun_command.append(f"--container-env={','.join(sorted(set(container_environment)))}") + + +def _get_gpu_mask(gpu_indices: tuple[int, ...]) -> int: + return sum(1 << index for index in gpu_indices) + + +def _render_mounts(plan: ResolvedSlurmRunPlan, scratch_root: str | None) -> str: + rendered: list[str] = [] + for mount in plan.container_mounts: + if any(character in mount.source or character in mount.target for character in (",", ":")): + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "container mount paths cannot contain comma or colon delimiters", + ) + value = f"{mount.source}:{mount.target}" + if mount.read_only: + value += ":ro" + rendered.append(value) + if scratch_root is not None: + rendered.append(f"{_validate_scratch_root(scratch_root)}:{ALLOCATION_SCRATCH_CONTAINER_ROOT}") + return ",".join(rendered) + + +def _validate_scratch_root(value: str) -> str: + try: + validated = validate_absolute_path(value) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation scratch path is invalid") from error + if any(character in validated for character in (",", ":")): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "allocation scratch path is invalid") + return validated + + +def _base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: + environment = { + "PATH": source_environment.get("PATH") or os.defpath, + "LC_ALL": "C", + } + environment.update( + (name, source_environment[name]) for name in _SLURM_ENVIRONMENT_NAMES if name in source_environment + ) + return environment + + +def _collect_client_secret_environment_names(plan: ResolvedSlurmRunPlan) -> tuple[str, ...]: + return collect_secret_environment_names( + (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ) + + +__all__ = [ + "ClientStepBuilder", + "DefaultClientStepBuilder", + "build_client_command", + "build_endpoint_command", + "build_endpoint_steps", + "build_vllm_command", + "build_vllm_steps", + "plan_path", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/supervisor.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/supervisor.py new file mode 100644 index 000000000..3ee61148a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/supervisor.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One owner for allocation-local process launch, observation, and cleanup.""" + +from __future__ import annotations + +import os +import signal +import subprocess +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Protocol + +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.logs import open_step_logs +from data_designer.slurm.runtime.models import RuntimeStep +from data_designer.slurm.runtime.signals import TerminationSignalCoordinator + + +class RuntimeClock(Protocol): + """Monotonic time and sleep boundary used by runtime polling.""" + + def monotonic(self) -> float: + """Return monotonic seconds.""" + ... + + def now(self) -> datetime: + """Return timezone-aware UTC wall-clock time.""" + ... + + def sleep(self, seconds: float) -> None: + """Advance or wait for positive seconds.""" + ... + + +class SystemRuntimeClock: + """Production runtime clock.""" + + def monotonic(self) -> float: + """Return system monotonic seconds.""" + return time.monotonic() + + def now(self) -> datetime: + """Return current UTC wall-clock time.""" + return datetime.now(timezone.utc) + + def sleep(self, seconds: float) -> None: + """Sleep for the requested duration.""" + time.sleep(seconds) + + +class StepProcess(Protocol): + """Minimal process handle required by the allocation supervisor.""" + + @property + def pid(self) -> int: + """Return the child process identity.""" + ... + + def poll(self) -> int | None: + """Return the exit status when complete.""" + ... + + def terminate(self) -> None: + """Request graceful termination.""" + ... + + def kill(self) -> None: + """Force termination.""" + ... + + def wait(self, timeout: float | None = None) -> int: + """Wait for completion and return the exit status.""" + ... + + +class RuntimeStepRunner(Protocol): + """Start structured runtime steps without a shell.""" + + def start(self, step: RuntimeStep) -> StepProcess: + """Start one process and return its handle.""" + ... + + +class SubprocessStepRunner: + """Production process boundary using ``Popen`` with exact argv and environment.""" + + def start(self, step: RuntimeStep) -> StepProcess: + """Start one step with restrictive package-owned log files.""" + with open_step_logs(step) as (stdout_descriptor, stderr_descriptor): + process = subprocess.Popen( + step.command, + stdin=subprocess.DEVNULL, + stdout=stdout_descriptor, + stderr=stderr_descriptor, + env=dict(step.environment), + shell=False, + start_new_session=True, + close_fds=True, + ) + return _SessionStepProcess(process) + + +@dataclass(frozen=True, slots=True) +class _SessionStepProcess: + """Process handle that signals the complete session-owned process group.""" + + process: subprocess.Popen[bytes] + + @property + def pid(self) -> int: + return self.process.pid + + def poll(self) -> int | None: + return self.process.poll() + + def terminate(self) -> None: + self._signal_group(signal.SIGTERM) + + def kill(self) -> None: + self._signal_group(signal.SIGKILL) + + def wait(self, timeout: float | None = None) -> int: + return self.process.wait(timeout=timeout) + + def _signal_group(self, selected: signal.Signals) -> None: + try: + os.killpg(self.pid, selected) + except ProcessLookupError: + if self.poll() is None: + raise + + +@dataclass(frozen=True, slots=True) +class ManagedStep: + """A started structured step and its process identity.""" + + step: RuntimeStep + process: StepProcess + + +class StepSupervisor: + """Own every child process and perform idempotent reverse-order cleanup.""" + + def __init__( + self, + runner: RuntimeStepRunner, + *, + signals: TerminationSignalCoordinator, + clock: RuntimeClock | None = None, + poll_interval_seconds: float = 0.1, + termination_grace_seconds: float = 10.0, + ) -> None: + if poll_interval_seconds <= 0 or termination_grace_seconds <= 0: + raise ValueError("runtime polling and termination intervals must be positive") + self._runner = runner + self._signals = signals + self._clock = clock or SystemRuntimeClock() + self._poll_interval_seconds = poll_interval_seconds + self._termination_grace_seconds = termination_grace_seconds + self._managed: list[ManagedStep] = [] + self._cleanup_started = False + self._cleanup_in_progress = False + self._cleanup_complete = False + + @property + def managed_steps(self) -> tuple[ManagedStep, ...]: + """Return the started process identities in launch order.""" + return tuple(self._managed) + + @property + def cleanup_complete(self) -> bool: + """Return whether every registered process is confirmed exited.""" + return self._cleanup_complete + + def start(self, step: RuntimeStep) -> ManagedStep: + """Start and register one structured step.""" + if self._cleanup_started: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.CLEANUP_FAILED, "cannot start a step after cleanup") + try: + with self._signals.defer_termination(): + managed = ManagedStep(step=step, process=self._runner.start(step)) + self._managed.append(managed) + except (OSError, subprocess.SubprocessError) as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.STEP_FAILED, + f"cannot start runtime step {step.step_id!r}", + ) from error + return managed + + def wait(self, target: ManagedStep, *, required: tuple[ManagedStep, ...] = ()) -> None: + """Wait for one finite step while requiring long-running peers to stay alive.""" + while True: + self.require_running(required) + returncode = target.process.poll() + if returncode is not None: + if returncode != 0: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.STEP_FAILED, + f"runtime step {target.step.step_id!r} exited with status {returncode}", + ) + return + self._clock.sleep(self._poll_interval_seconds) + + @staticmethod + def require_running(required: tuple[ManagedStep, ...]) -> None: + """Fail when any required long-running step has exited.""" + for managed in required: + returncode = managed.process.poll() + if returncode is not None: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.STEP_FAILED, + f"required runtime step {managed.step.step_id!r} exited with status {returncode}", + ) + + def cleanup(self) -> None: + """Terminate every live child exactly once, escalating after a bounded grace period.""" + if self._cleanup_complete or self._cleanup_in_progress: + return + self._cleanup_started = True + self._cleanup_in_progress = True + try: + with self._signals.block_termination(): + failures = self._stop_managed_steps() + self._cleanup_complete = self._are_all_steps_stopped() + finally: + self._cleanup_in_progress = False + if not self._cleanup_complete: + failures.append(RuntimeError("one or more runtime processes are still running")) + if failures: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.CLEANUP_FAILED, + f"cleanup failed for {len(failures)} runtime process operation(s)", + ) from failures[0] + + def _stop_managed_steps(self) -> list[BaseException]: + failures: list[BaseException] = [] + live = tuple(managed for managed in reversed(self._managed) if self._is_running(managed, failures)) + for managed in live: + try: + managed.process.terminate() + except BaseException as error: + failures.append(error) + deadline = self._clock.monotonic() + self._termination_grace_seconds + for managed in live: + self._stop_step(managed, deadline, failures) + return failures + + def _stop_step(self, managed: ManagedStep, deadline: float, failures: list[BaseException]) -> None: + while self._is_running(managed, failures) and self._clock.monotonic() < deadline: + self._clock.sleep(self._poll_interval_seconds) + if self._is_running(managed, failures): + try: + managed.process.kill() + except BaseException as error: + failures.append(error) + try: + managed.process.wait(timeout=self._termination_grace_seconds) + except BaseException as error: + failures.append(error) + + @staticmethod + def _is_running(managed: ManagedStep, failures: list[BaseException]) -> bool: + try: + return managed.process.poll() is None + except BaseException as error: + failures.append(error) + return True + + def _are_all_steps_stopped(self) -> bool: + try: + return all(managed.process.poll() is not None for managed in self._managed) + except BaseException: + return False diff --git a/packages/data-designer-slurm/src/data_designer/slurm/security.py b/packages/data-designer-slurm/src/data_designer/slurm/security.py new file mode 100644 index 000000000..ed9a97ec2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/security.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Security helpers shared by Slurm process boundaries.""" + +from __future__ import annotations + +import re +import unicodedata + +from data_designer.slurm.config.environment import is_secret_bearing_name + +_ASSIGNMENT_START_PATTERN = re.compile( + r"(?P(?P[\"']?)(?P-{0,2}[A-Za-z][A-Za-z0-9_.-]*)(?P=quote)\s*[:=]\s*)" +) +_OPTION_START_PATTERN = re.compile(r"(?P(?P--[A-Za-z][A-Za-z0-9_.-]*)\s+)") +_AUTHORIZATION_PATTERN = re.compile( + r"(?i)(?P(?P[\"']?)\bauthorization(?P=key_quote)\s*[:=]\s*)" + r"[^\r\n]*" +) +_URL_USERINFO_PATTERN = re.compile(r"(?i)(?P\b[A-Za-z][A-Za-z0-9+.-]*://)[^/\s?#]*@") +_TOKEN_PATTERNS = ( + re.compile(r"\bAKIA[0-9A-Z]{16}\b"), + re.compile(r"\bgithub_pat_[A-Za-z0-9_]{20,}\b"), + re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), + re.compile(r"\bnvapi-[A-Za-z0-9_-]{20,}\b"), + re.compile(r"\bsk-(?:proj-)?[A-Za-z0-9_-]{20,}\b"), +) +_REDACTION = "" + + +def redact_sensitive_text(value: str) -> str: + """Redact recognizable credentials without echoing their values.""" + return _redact_sensitive_text(value, replacement=_REDACTION) + + +def redact_sensitive_diagnostic(value: str) -> str: + """Redact secrets both inside and separated by control characters.""" + control_placeholder = _select_placeholder(value, name="control") + redaction_placeholder = _select_placeholder(value, name="redaction") + scan_view = _replace_control_characters(value, replacement=control_placeholder) + redacted = _redact_sensitive_text(scan_view, replacement=redaction_placeholder) + normalized_boundaries = redacted.replace(control_placeholder, " ") + protected = _redact_sensitive_text( + normalized_boundaries, + replacement=redaction_placeholder, + protected_replacement=redaction_placeholder, + ) + return protected.replace(redaction_placeholder, _REDACTION) + + +def _redact_sensitive_text( + value: str, + *, + replacement: str, + protected_replacement: str | None = None, +) -> str: + redacted = _AUTHORIZATION_PATTERN.sub(lambda match: f"{match.group('prefix')}{replacement}", value) + redacted = _URL_USERINFO_PATTERN.sub(lambda match: f"{match.group('scheme')}{replacement}@", redacted) + redacted = _redact_named_values( + redacted, + _ASSIGNMENT_START_PATTERN, + replacement=replacement, + protected_replacement=protected_replacement, + ) + redacted = _redact_named_values( + redacted, + _OPTION_START_PATTERN, + replacement=replacement, + protected_replacement=protected_replacement, + ) + for pattern in _TOKEN_PATTERNS: + redacted = pattern.sub(replacement, redacted) + return redacted + + +def _redact_named_values( + value: str, + start_pattern: re.Pattern[str], + *, + replacement: str, + protected_replacement: str | None, +) -> str: + """Redact secret-bearing named values without letting earlier matches overlap them.""" + parts: list[str] = [] + output_cursor = 0 + search_cursor = 0 + while match := start_pattern.search(value, search_cursor): + if not is_secret_bearing_name(match.group("name").lstrip("-")): + search_cursor = match.end() + continue + if protected_replacement is not None and value.startswith(protected_replacement, match.end()): + search_cursor = match.end() + len(protected_replacement) + continue + value_end = _find_named_value_end(value, match.end()) + parts.append(value[output_cursor : match.start()]) + parts.append(f"{match.group('prefix')}{replacement}") + output_cursor = value_end + search_cursor = value_end + parts.append(value[output_cursor:]) + return "".join(parts) + + +def _find_named_value_end(value: str, start: int) -> int: + """Find the first unambiguous boundary for a secret-bearing value.""" + if start < len(value) and value[start] in {'"', "'"}: + return _find_quoted_value_end(value, start) + index = start + while index < len(value) and not value[index].isspace(): + index += 1 + return index + + +def _find_quoted_value_end(value: str, start: int) -> int: + quote = value[start] + index = start + 1 + while index < len(value): + if value[index] == "\\": + index += 2 + elif value[index] == quote: + return index + 1 + else: + index += 1 + return len(value) + + +def _replace_control_characters(value: str, *, replacement: str) -> str: + """Replace ambiguous non-line boundaries before scanning for secrets.""" + return "".join( + character + if character in "\r\n " or not (character.isspace() or unicodedata.category(character).startswith("C")) + else replacement + for character in value + ) + + +def _select_placeholder(value: str, *, name: str) -> str: + """Return a named printable marker guaranteed absent from the input.""" + placeholder = f"" + while placeholder in value: + placeholder = f"<{placeholder}>" + return placeholder + + +__all__ = ["redact_sensitive_diagnostic", "redact_sensitive_text"] 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 new file mode 100644 index 000000000..e6a3107f1 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public service facades and their dependency contracts.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +from data_designer.slurm.services.benchmark import SlurmBenchmarkBackend, SlurmBenchmarkService +from data_designer.slurm.services.errors import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) +from data_designer.slurm.services.images import SlurmImageManager, SlurmImageResolver, SlurmImageService +from data_designer.slurm.services.profiles import ( + SlurmProfileInitialization, + SlurmProfileMatch, + SlurmProfileService, + SlurmProfileValidation, + create_slurm_profile_service, +) +from data_designer.slurm.services.results import ( + SlurmCollectionExecution, + SlurmPersistedAttemptStatus, + SlurmPersistedRunStatus, + SlurmPersistedShardStatus, + SlurmRetryExecution, + SlurmRunCancellation, + SlurmRunExecution, +) +from data_designer.slurm.services.run import SlurmBatchScriptRenderer, SlurmRunBackend, SlurmRunPlanner, SlurmRunService + +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"), +} + +__all__ = [ + "SlurmBatchScriptRenderer", + "SlurmBenchmarkBackend", + "SlurmBenchmarkService", + "SlurmCollectionExecution", + "SlurmImageManager", + "SlurmImageResolver", + "SlurmImageService", + "SlurmPersistedAttemptStatus", + "SlurmPersistedRunStatus", + "SlurmPersistedShardStatus", + "SlurmProfileInitialization", + "SlurmProfileMatch", + "SlurmProfileService", + "SlurmProfileValidation", + "SlurmRetryExecution", + "SlurmRunArtifactPublisher", + "SlurmRunBackend", + "SlurmRunCancellation", + "SlurmRunExecution", + "SlurmRunPlanner", + "SlurmRunService", + "SlurmServiceError", + "SlurmServiceErrorCode", + "SlurmServiceOperation", + "create_slurm_benchmark_service", + "create_slurm_image_service", + "create_slurm_profile_service", + "create_slurm_run_service", +] + + +def __getattr__(name: str) -> object: + """Lazily import production service wiring.""" + if name in _LAZY_IMPORTS: + module_path, attribute_name = _LAZY_IMPORTS[name] + attribute = getattr(importlib.import_module(module_path), attribute_name) + globals()[name] = attribute + return attribute + raise AttributeError(f"module 'data_designer.slurm.services' has no attribute {name!r}") + + +def __dir__() -> list[str]: + """Return public service exports.""" + return __all__ 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 new file mode 100644 index 000000000..6e01c2d46 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Publish immutable inputs and initial state for one submitted run.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from pydantic import JsonValue + +from data_designer.slurm.client.dependencies import ResolvedClientDependencies +from data_designer.slurm.client.filesystem import ensure_private_directory +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import ArtifactReference, pretty_json +from data_designer.slurm.filesystem import create_restrictive_temporary_file, get_file_facts +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateNotFoundError, +) +from data_designer.slurm.state.filesystem import ( + open_verified_directory, + open_verified_regular_file, + publish_immutable_text, + sync_directory, +) +from data_designer.slurm.state.storage import StateStorage + +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_TEMPORARY_PREFIX = ".artifact." +_TEMPORARY_SUFFIX = ".tmp" + + +class StateRunArtifactPublisher: + """Persist submission inputs through the package-owned state workspace.""" + + def __init__(self, workspace_root: str | Path, clock: Callable[[], datetime]) -> None: + self._workspace_root = Path(workspace_root) + self._clock = clock + + def initialize( + self, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, JsonValue] | None, + *, + force: bool, + ) -> None: + if force: + raise StateConflictError("force cannot replace durable run state") + 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(), + ) + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=created_at, + authored_config=plan.authored_config, + resolved_plan=plan_reference, + shard_count=len(plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + shard_index=shard.shard_index, + record_range=shard.record_range, + input_partition=shard.input_partition, + resume_workspace=shard.resume_workspace, + created_at=created_at, + ) + for shard in plan.shards + ) + writer.initialize_run(authored, plan, run, shards) + try: + self._publish_inputs(writer.run_root, plan, dependencies, builder_payload) + except SlurmStateError: + raise + except Exception as error: + raise SlurmStateError(f"cannot publish immutable inputs for run {plan.run_id!r}") from error + + def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: + writer = SlurmStateWriter(self._workspace_root, plan.run_id) + plan_reference = ArtifactReference( + path=(writer.run_root / "resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ) + for shard in plan.shards: + writer.create_attempt( + AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=job_id, array_task_id=shard.array_task_index), + created_at=submitted_at, + updated_at=submitted_at, + ) + ) + + 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={ + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.CANCELLED, + "updated_at": failed_at, + } + ) + ) + + @staticmethod + def _publish_inputs( + run_root: Path, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, JsonValue] | None, + ) -> None: + _publish_text(run_root, plan.client.dependency_lock, dependencies.lock.serialize_json()) + if (plan.builder.source is None) != (builder_payload is None): + raise SlurmStateError("resolved builder source does not match its staged payload") + if plan.builder.source is not None: + assert builder_payload is not None + _publish_text(run_root, plan.builder.source, pretty_json(builder_payload)) + if (dependencies.lock.source is None) != (dependencies.lock_source is None): + raise SlurmStateError("dependency lock source does not match its staged artifact") + if dependencies.lock.source is not None: + assert dependencies.lock_source is not None + _publish_file(run_root, dependencies.lock.source, dependencies.lock_source) + packages = dependencies.lock.overlay_packages + if len(packages) != len(dependencies.wheel_sources): + raise SlurmStateError("dependency wheel sources do not match the resolved lock") + for package, source in zip(packages, dependencies.wheel_sources, strict=True): + _publish_file(run_root, package.artifact, source) + seed_path = plan.invocation.effective_input_bindings.seed_path + for shard in plan.shards: + if shard.input_partition is None: + continue + _publish_text( + run_root, + shard.input_partition, + pretty_json( + { + "record_range": shard.record_range.model_dump(mode="json"), + "seed_path": seed_path, + } + ), + ) + + +def _publish_text(run_root: Path, reference: ArtifactReference, content: str) -> None: + target = _validate_target(run_root, reference) + if hashlib.sha256(content.encode()).hexdigest() != reference.sha256: + raise SlurmStateError(f"staged artifact {target.name!r} does not match its resolved digest") + ensure_private_directory(target.parent) + with open_verified_directory(target.parent, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + target.name, + content, + target, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + +def _publish_file(run_root: Path, reference: ArtifactReference, source: Path) -> None: + target = _validate_target(run_root, reference) + ensure_private_directory(target.parent) + source_descriptor: int | None = None + try: + source_before = source.lstat() + source_descriptor = os.open( + source, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + ) + source_opened = os.fstat(source_descriptor) + if not stat.S_ISREG(source_opened.st_mode) or get_file_facts(source_before) != get_file_facts(source_opened): + raise OSError(f"source artifact {source} is not a stable regular file") + with open_verified_directory(target.parent, require_private=True) as target_descriptor: + output_descriptor, temporary_name = create_restrictive_temporary_file( + target_descriptor, + prefix=_TEMPORARY_PREFIX, + suffix=_TEMPORARY_SUFFIX, + ) + try: + digest = hashlib.sha256() + try: + while chunk := os.read(source_descriptor, 1024 * 1024): + digest.update(chunk) + remaining = memoryview(chunk) + while remaining: + written = os.write(output_descriptor, remaining) + if written == 0: + raise OSError("artifact copy made no progress") + remaining = remaining[written:] + os.fsync(output_descriptor) + finally: + os.close(output_descriptor) + source_after = os.fstat(source_descriptor) + source_current = source.lstat() + if ( + get_file_facts(source_opened) != get_file_facts(source_after) + or get_file_facts(source_after) != get_file_facts(source_current) + or digest.hexdigest() != reference.sha256 + ): + raise OSError(f"source artifact {source} changed or has an unexpected digest") + try: + os.link( + temporary_name, + target.name, + src_dir_fd=target_descriptor, + dst_dir_fd=target_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + with open_verified_regular_file( + target_descriptor, + target.name, + target, + expected_size=source_after.st_size, + expected_sha256=reference.sha256, + ): + pass + os.unlink(temporary_name, dir_fd=target_descriptor) + temporary_name = None + sync_directory(target_descriptor) + finally: + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=target_descriptor) + except OSError: + pass + finally: + if source_descriptor is not None: + os.close(source_descriptor) + + +def _validate_target(run_root: Path, reference: ArtifactReference) -> Path: + target = Path(reference.path) + try: + target.relative_to(run_root) + except ValueError: + raise SlurmStateError("resolved artifact path is outside the run workspace") from None + return target + + +__all__ = ["StateRunArtifactPublisher"] 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 new file mode 100644 index 000000000..ceb04ec03 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public benchmark facade for Data Designer Slurm.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.services.errors import ( + SlurmServiceOperation, + _invoke_service_backend, + _make_invalid_request_error, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class SlurmBenchmarkBackend(Protocol): + """Process benchmarks through a supported service dependency.""" + + 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, + fail_if_incomplete: bool, + ) -> BenchmarkReport: + """Return one point-in-time benchmark report. + + Any non-``INTERNAL`` service error must contain a caller-safe message. + """ + + +class SlurmBenchmarkService: + """Expose benchmark operations through a package-owned boundary. + + The service borrows its injected dependency and does not manage its lifecycle. + + Args: + backend: Benchmark dependency implementing ``SlurmBenchmarkBackend``. + """ + + def __init__(self, backend: SlurmBenchmarkBackend) -> None: + self._backend = backend + + 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. + + Raises: + SlurmServiceError: If the request is invalid or benchmark launch fails. + """ + 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, + 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(): + raise ValueError("benchmark manifest does not match the requested config") + return manifest + + return _invoke_service_backend(operation, run_benchmark) + + 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. + """ + operation = SlurmServiceOperation.ANALYZE_BENCHMARK + try: + 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 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, + 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: + raise ValueError("benchmark report does not match the requested benchmark") + return report + + return _invoke_service_backend(operation, analyze_benchmark) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py new file mode 100644 index 000000000..ecf336f49 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stable errors raised by public Slurm services.""" + +from __future__ import annotations + +from collections.abc import Callable +from enum import Enum +from typing import TypeVar + +from data_designer.errors import DataDesignerError + +_ResultT = TypeVar("_ResultT") + + +class SlurmServiceErrorCode(str, Enum): + """Machine-readable public service error categories.""" + + INVALID_REQUEST = "invalid_request" + NOT_FOUND = "not_found" + CONFLICT = "conflict" + UNAVAILABLE = "unavailable" + INTERNAL = "internal" + + +class SlurmServiceOperation(str, Enum): + """Public operations with stable error attribution.""" + + PLAN_RUN = "plan_run" + RENDER_ATTEMPT = "render_attempt" + EXECUTE_RUN = "execute_run" + STATUS_RUN = "status_run" + CANCEL_RUN = "cancel_run" + RETRY_RUN = "retry_run" + COLLECT_RUN = "collect_run" + INIT_PROFILE = "init_profile" + VALIDATE_PROFILE = "validate_profile" + RESOLVE_IMAGE = "resolve_image" + ADD_IMAGE = "add_image" + LIST_IMAGES = "list_images" + GET_IMAGE = "get_image" + REMOVE_IMAGE = "remove_image" + RUN_BENCHMARK = "run_benchmark" + ANALYZE_BENCHMARK = "analyze_benchmark" + + +class SlurmServiceError(DataDesignerError): + """Normalized public service failure with a stable code and operation. + + The string form is caller-visible. Non-``INTERNAL`` messages must be + sanitized before construction and must never include raw exception details. + Facades replace dependency-raised ``INTERNAL`` messages with fixed text. + + Args: + code: Machine-readable failure category. + operation: Public operation that failed. + message: Caller-safe text containing 1 to 512 characters. + + Attributes: + code: Machine-readable failure category. + operation: Public operation that failed. + """ + + def __init__( + self, + code: SlurmServiceErrorCode, + operation: SlurmServiceOperation, + message: str, + ) -> None: + if not isinstance(code, SlurmServiceErrorCode): + raise TypeError("code must be a SlurmServiceErrorCode") + if not isinstance(operation, SlurmServiceOperation): + raise TypeError("operation must be a SlurmServiceOperation") + if type(message) is not str or not message or len(message) > 512: + raise ValueError("service error message must contain 1 to 512 characters") + if any(ord(character) < 32 or ord(character) == 127 for character in message): + raise ValueError("service error message must not contain control characters") + self.code = code + self.operation = operation + super().__init__(message) + + def __reduce__( + self, + ) -> tuple[type[SlurmServiceError], tuple[SlurmServiceErrorCode, SlurmServiceOperation, str]]: + return type(self), (self.code, self.operation, str(self)) + + +def _make_invalid_request_error(operation: SlurmServiceOperation, message: str) -> SlurmServiceError: + """Build one normalized invalid-request error.""" + return SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, message) + + +def _invoke_service_backend(operation: SlurmServiceOperation, call: Callable[[], _ResultT]) -> _ResultT: + """Preserve caller-safe failures and redact backend-internal details.""" + try: + return call() + except SlurmServiceError as error: + if error.code is SlurmServiceErrorCode.INTERNAL: + pass + elif error.operation is operation: + raise + else: + raise SlurmServiceError(error.code, operation, str(error)) from None + except Exception: + pass + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + f"{operation.value.replace('_', ' ')} failed", + ) from None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py new file mode 100644 index 000000000..656501ea2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production composition for synchronous Slurm image lifecycle jobs.""" + +from __future__ import annotations + +import codecs +import re +import time +from collections.abc import Callable, Sequence +from datetime import datetime, timedelta, timezone +from uuid import uuid4 + +from data_designer.slurm.config import ImageBuildRequest, SelectedSlurmProfile +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.images.errors import ( + ImageConflictError, + ImageLifecycleError, + ImageRegistryError, + ImageVerificationError, +) +from data_designer.slurm.images.lifecycle import ( + PreparedImageLifecycleJob, + cleanup_prepared_image_lifecycle, + prepare_image_lifecycle_job, + publish_completed_image_lifecycle, + read_image_lifecycle_log, + read_image_lifecycle_phase, + submit_prepared_image_lifecycle, +) +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmAccountingEntry, SlurmQueueEntry +from data_designer.slurm.security import redact_sensitive_diagnostic +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.state import ( + SchedulerJobIdentity, + SchedulerObservation, + SchedulerObservationCollector, + SchedulerState, + SlurmStateError, +) +from data_designer.slurm.state.scheduler import is_scheduler_failure_state, is_scheduler_terminal_state + +LifecycleIdFactory = Callable[[], str] +Clock = Callable[[], datetime] +Sleeper = Callable[[float], None] +ProgressReporter = Callable[[str], None] +LogReporter = Callable[[str], None] + +_POLL_INTERVAL_SECONDS = 30.0 +_HEARTBEAT_INTERVAL = timedelta(minutes=2) +_ACCOUNTING_EXIT_LAG = timedelta(minutes=5) +_LOG_READ_BYTES = 64 * 1024 +_LOG_TOTAL_BYTES = 1024 * 1024 +_LOG_LIMIT_NOTICE = "Image job log display reached its 1 MiB limit; further output is hidden." +_UNSAFE_TERMINAL_CHARACTERS = re.compile(r"[\x00-\x08\x0b-\x1f\x7f-\x9f]") +_PHASE_LABELS = { + "importing_oci": "importing OCI layers and creating SQSH", + "preparing_sqsh": "preparing existing SQSH", + "verifying_sqsh": "verifying SQSH", + "inspecting_image": "inspecting image", + "inspection_complete": "inspection complete", +} + + +class _TerminalLifecycleError(RuntimeError): + pass + + +class _JobLogFollower: + """Best-effort, bounded local output for the exact submitted lifecycle job.""" + + def __init__(self, prepared: PreparedImageLifecycleJob, job_id: int, report: LogReporter) -> None: + self._prepared = prepared + self._job_id = job_id + self._report = report + self._offsets = {"stdout": 0, "stderr": 0} + self._decoders = {stream: codecs.getincrementaldecoder("utf-8")(errors="replace") for stream in self._offsets} + self._pending_lines = {"stdout": "", "stderr": ""} + self._skip_next_lf = {"stdout": False, "stderr": False} + self._source_bytes = 0 + self._displayed_bytes = 0 + self._unavailable_reported = False + self._limit_reported = False + + def drain(self, *, final: bool = False) -> bool: + """Print newly appended bytes; after termination, drain the bounded remainder.""" + before = self._displayed_bytes + if self._limit_reported: + return False + while self._source_bytes < _LOG_TOTAL_BYTES: + received = False + for stream in ("stdout", "stderr"): + allowance = min(_LOG_READ_BYTES, _LOG_TOTAL_BYTES - self._source_bytes) + if allowance == 0: + break + try: + content, offset = read_image_lifecycle_log( + self._prepared, + self._job_id, + stream, + offset=self._offsets[stream], + maximum_bytes=allowance, + ) + except (ImageLifecycleError, OSError, ValueError): + if not self._unavailable_reported: + self._emit("Image job logs are unavailable.") + self._unavailable_reported = True + return self._displayed_bytes > before + self._offsets[stream] = offset + self._source_bytes += len(content) + if content: + received = True + self._append(stream, content) + if self._limit_reported: + return self._displayed_bytes > before + if not final or not received: + break + if final or self._source_bytes >= _LOG_TOTAL_BYTES: + for stream in ("stdout", "stderr"): + self._append(stream, b"", final=True) + if self._pending_lines[stream]: + self._emit_line(stream, self._pending_lines[stream]) + self._pending_lines[stream] = "" + if self._source_bytes >= _LOG_TOTAL_BYTES: + self._emit_limit() + return self._displayed_bytes > before + + def _append(self, stream: str, content: bytes, *, final: bool = False) -> None: + text = self._decoders[stream].decode(content, final=final) + if self._skip_next_lf[stream]: + if text.startswith("\n"): + text = text[1:] + self._skip_next_lf[stream] = False + parts = re.split(r"(\r\n|\r|\n)", text) + pending = self._pending_lines[stream] + for index in range(0, len(parts) - 1, 2): + self._emit_line(stream, pending + parts[index]) + pending = "" + if self._limit_reported: + return + self._pending_lines[stream] = pending + parts[-1] + self._skip_next_lf[stream] = len(parts) > 1 and parts[-2] == "\r" and not parts[-1] + + def _emit_line(self, stream: str, line: str) -> None: + if not line.strip(): + return + sanitized = redact_sensitive_diagnostic(line) + self._emit(f"[image {stream}] {_UNSAFE_TERMINAL_CHARACTERS.sub('?', sanitized)}") + + def _emit(self, message: str) -> None: + if self._limit_reported: + return + available = _LOG_TOTAL_BYTES - len(_LOG_LIMIT_NOTICE.encode()) - 1 - self._displayed_bytes + encoded = message.encode("utf-8") + if len(encoded) + 1 > available: + if available > 1: + self._send(encoded[: available - 1].decode("utf-8", errors="ignore")) + self._emit_limit() + return + self._send(message) + + def _send(self, message: str) -> None: + self._displayed_bytes += len(message.encode("utf-8")) + 1 + try: + self._report(message) + except Exception: + # Log display must never change a submitted job's outcome. + pass + + def _emit_limit(self) -> None: + if not self._limit_reported: + self._limit_reported = True + self._send(_LOG_LIMIT_NOTICE) + + +class _RecordingObservationClient: + def __init__(self, launcher: SlurmCommandClient) -> None: + self._launcher = launcher + self.accounting: tuple[SlurmAccountingEntry, ...] = () + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + return self._launcher.query_queue(selectors) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + self.accounting = self._launcher.query_accounting(selectors) + return self.accounting + + +class SlurmImageLifecycleManager: + """Prepare, submit, reconcile, publish, and register one image.""" + + def __init__( + self, + selected_profile: SelectedSlurmProfile, + launcher: SlurmCommandClient, + *, + lifecycle_id_factory: LifecycleIdFactory | None = None, + clock: Clock | None = None, + sleep: Sleeper | None = None, + progress: ProgressReporter | None = None, + logs: LogReporter | None = None, + ) -> None: + self._profile = selected_profile + self._launcher = launcher + self._lifecycle_id_factory = lifecycle_id_factory or _new_lifecycle_id + self._clock = clock or _utc_now + self._sleep = sleep or time.sleep + self._progress = progress + self._logs = logs + + def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: + """Run one lifecycle job and publish only a successful verified result.""" + operation = SlurmServiceOperation.ADD_IMAGE + try: + prepared = prepare_image_lifecycle_job( + request, + self._profile, + lifecycle_id=self._lifecycle_id_factory(), + ) + except (ImageLifecycleError, OSError, ValueError): + raise _unavailable("image lifecycle job cannot be prepared") from None + + try: + self._report_progress("Submitting image add job...") + receipt = submit_prepared_image_lifecycle(prepared, self._launcher) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + _cleanup_failed_lifecycle(prepared) + raise _unavailable("image lifecycle job cannot be submitted") from None + raise _unavailable( + f"image lifecycle submission outcome is unknown; lifecycle {prepared.plan.lifecycle_id} was retained" + ) from None + except (ImageLifecycleError, SlurmLauncherError, OSError, ValueError): + _cleanup_failed_lifecycle(prepared) + raise _unavailable("image lifecycle job cannot be submitted") from None + + follower = _JobLogFollower(prepared, receipt.job_id, self._logs) if self._logs is not None else None + try: + submitted_at = self._clock() + self._report_progress(f"Image add job {receipt.job_id} submitted; checking every 30 seconds.") + self._wait_for_success(prepared, receipt.job_id, follower, submitted_at=submitted_at) + self._report_progress("Registering image alias...") + registered = publish_completed_image_lifecycle(prepared, replace=replace) + self._report_progress( + f"Image add complete: {registered.name} (total elapsed {_format_elapsed(self._clock() - submitted_at)})." + ) + return registered + except (ImageConflictError, ImageVerificationError): + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "image lifecycle result cannot be verified or registered", + ) from None + except ImageRegistryError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "image registry operation failed", + ) from None + except ImageLifecycleError: + raise _unavailable("image lifecycle result cannot be published") from None + except _TerminalLifecycleError: + _cleanup_failed_lifecycle(prepared) + raise _unavailable(f"image lifecycle job {receipt.job_id} did not complete successfully") from None + except (SlurmLauncherError, SlurmStateError, OSError, ValueError): + self._cancel_and_cleanup(receipt.job_id, prepared, follower) + raise _unavailable(f"image lifecycle job {receipt.job_id} did not complete successfully") from None + except BaseException: + self._cancel_and_cleanup(receipt.job_id, prepared, follower) + raise + + def _wait_for_success( + self, + prepared: PreparedImageLifecycleJob, + job_id: int, + follower: _JobLogFollower | None, + *, + submitted_at: datetime, + ) -> None: + client = _RecordingObservationClient(self._launcher) + observations = SchedulerObservationCollector(client) + previous: SchedulerObservation | None = None + completed_without_accounting_deadline: datetime | None = None + last_activity_at: datetime | None = None + previous_phase: str | None = None + while True: + observed_at = self._clock() + observation = observations.collect( + (job_id,), + observed_at=observed_at, + previous={job_id: previous}, + )[0] + elapsed = _format_elapsed(observed_at - submitted_at) + state = observation.state.value.replace("_", " ") + state_changed = previous is None or observation.state is not previous.state + had_logs = follower.drain() if follower is not None else False + if had_logs: + last_activity_at = observed_at + try: + phase = read_image_lifecycle_phase(prepared) + except (ImageLifecycleError, OSError, ValueError): + phase = None + if phase is not None and phase != previous_phase: + self._report_progress(f"Image add job {job_id}: {_PHASE_LABELS[phase]}.") + previous_phase = phase + last_activity_at = observed_at + heartbeat_due = last_activity_at is None or observed_at - last_activity_at >= _HEARTBEAT_INTERVAL + if state_changed or heartbeat_due: + if not state_changed and observation.state is SchedulerState.RUNNING: + state = "still running" + self._report_progress(f"Image add job {job_id}: {state} (elapsed {elapsed}).") + last_activity_at = observed_at + if observation.state is SchedulerState.COMPLETED: + accounting = next((entry for entry in client.accounting if entry.job_identity == job_id), None) + if accounting is None or accounting.state is not SchedulerState.COMPLETED: + if completed_without_accounting_deadline is None: + self._report_progress("Waiting for successful job exit evidence...") + completed_without_accounting_deadline = observed_at + _ACCOUNTING_EXIT_LAG + elif observed_at > completed_without_accounting_deadline: + raise SlurmStateError("completed image lifecycle job has no exit evidence") + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + continue + if ( + accounting.process_exit_code.exit_status != 0 + or accounting.process_exit_code.termination_signal != 0 + ): + if follower is not None: + follower.drain(final=True) + raise _TerminalLifecycleError("completed image lifecycle job has no successful exit evidence") + if follower is not None: + follower.drain(final=True) + return + if is_scheduler_failure_state(observation.state): + if follower is not None: + follower.drain(final=True) + raise _TerminalLifecycleError("image lifecycle job did not complete successfully") + if observation.state is SchedulerState.UNKNOWN: + raise SlurmStateError("image lifecycle job has unknown scheduler state") + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + + def _report_progress(self, message: str) -> None: + if self._progress is not None: + try: + self._progress(message) + except Exception: + # A closed progress stream must not cancel an already submitted job. + pass + + def _cancel_and_cleanup( + self, + job_id: int, + prepared: PreparedImageLifecycleJob, + follower: _JobLogFollower | None = None, + ) -> None: + terminated = False + try: + self._report_progress("Cancelling image inspection job...") + self._launcher.cancel(job_id) + terminated = self._wait_for_termination(job_id) + except (SlurmLauncherError, SlurmStateError, OSError, ValueError): + pass + if follower is not None: + follower.drain(final=True) + if terminated: + _cleanup_failed_lifecycle(prepared) + + def _wait_for_termination(self, job_id: int) -> bool: + client = _RecordingObservationClient(self._launcher) + observations = SchedulerObservationCollector(client) + previous: SchedulerObservation | None = None + deadline = self._clock() + _ACCOUNTING_EXIT_LAG + while True: + observed_at = self._clock() + observation = observations.collect( + (job_id,), + observed_at=observed_at, + previous={job_id: previous}, + )[0] + self._report_progress("Waiting for image inspection job to terminate...") + accounting = next((entry for entry in client.accounting if entry.job_identity == job_id), None) + if accounting is not None and is_scheduler_terminal_state(accounting.state): + return True + if observation.state is SchedulerState.UNKNOWN or observed_at >= deadline: + return False + previous = observation + self._sleep(_POLL_INTERVAL_SECONDS) + + +def _cleanup_failed_lifecycle(prepared: PreparedImageLifecycleJob) -> None: + try: + cleanup_prepared_image_lifecycle(prepared) + except ImageLifecycleError: + pass + + +def _unavailable(message: str) -> SlurmServiceError: + return SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, SlurmServiceOperation.ADD_IMAGE, message) + + +def _new_lifecycle_id() -> Identifier: + return f"image-{uuid4().hex}" + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def _format_elapsed(duration: timedelta) -> str: + elapsed_seconds = max(0, int(duration.total_seconds())) + hours, remainder = divmod(elapsed_seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return f"{hours:02}:{minutes:02}:{seconds:02}" + + +__all__ = ["Clock", "LifecycleIdFactory", "Sleeper", "SlurmImageLifecycleManager"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py new file mode 100644 index 000000000..84ab6defa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public image-resolution facade for Data Designer Slurm.""" + +from __future__ import annotations + +from typing import Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.config import ImageKind, ImageRef +from data_designer.slurm.config.images import ImageBuildRequest +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.images.records import RegisteredImage, validate_oci_source_for_lifecycle +from data_designer.slurm.planning import ResolvedImage +from data_designer.slurm.services.errors import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, + _invoke_service_backend, + _make_invalid_request_error, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class SlurmImageResolver(Protocol): + """Resolve verified images through a supported service dependency.""" + + def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: + """Return immutable facts for one verified image reference. + + Any non-``INTERNAL`` service error must contain a caller-safe message. + """ + + +class SlurmImageManager(Protocol): + """Manage verified images through a supported service dependency.""" + + def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: + """Build or inspect and register one image.""" + + def list(self) -> tuple[RegisteredImage, ...]: + """Return all registered images in alias order.""" + + def get(self, name: Identifier) -> RegisteredImage: + """Return one registered image.""" + + def remove(self, name: Identifier, *, delete_artifact: bool = False) -> RegisteredImage: + """Unregister one alias and optionally delete its package-owned SQSH.""" + + +class SlurmImageService: + """Expose stable image results through a package-owned boundary. + + The service borrows its injected dependency and does not manage its lifecycle. + + Args: + resolver: Image-resolution dependency implementing ``SlurmImageResolver``. + manager: Optional image-operation dependency. Package-owned factories provide it. + """ + + def __init__(self, resolver: SlurmImageResolver, manager: SlurmImageManager | None = None) -> None: + self._resolver = resolver + self._manager = manager + + def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: + """Resolve one registered image for planning. + + The result must match both the authored reference and expected image kind. + + Raises: + SlurmServiceError: If the request is invalid or resolution fails. + """ + operation = SlurmServiceOperation.RESOLVE_IMAGE + if not isinstance(reference, ImageRef): + raise _make_invalid_request_error(operation, "reference must be an ImageRef") + if not isinstance(expected_kind, ImageKind): + raise _make_invalid_request_error(operation, "expected_kind must be an ImageKind") + + def resolve_image() -> ResolvedImage: + image = self._resolver.resolve(reference, expected_kind=expected_kind) + if not isinstance(image, ResolvedImage): + raise TypeError("image resolver returned an invalid result") + if image.authored_ref != reference: + raise ValueError("resolved image does not match the requested reference") + if image.kind is not expected_kind: + raise ValueError("resolved image does not match the expected kind") + return image + + return _invoke_service_backend(operation, resolve_image) + + def add(self, request: ImageBuildRequest, *, replace: bool = False) -> RegisteredImage: + """Build or inspect and atomically register one image.""" + operation = SlurmServiceOperation.ADD_IMAGE + if not isinstance(request, ImageBuildRequest): + raise _make_invalid_request_error(operation, "request must be an ImageBuildRequest") + if type(replace) is not bool: + raise _make_invalid_request_error(operation, "replace must be a boolean") + try: + validate_oci_source_for_lifecycle(request.source) + except ValueError: + raise _make_invalid_request_error( + operation, + "OCI image source must be a credential-free registry reference without a scheme", + ) from None + manager = self._require_manager(operation) + + def add_image() -> RegisteredImage: + image = manager.add(request, replace=replace) + if not isinstance(image, RegisteredImage) or image.name != request.name or image.kind.value != request.kind: + raise TypeError("image manager returned an invalid registration result") + return image + + return _invoke_service_backend(operation, add_image) + + def list(self) -> tuple[RegisteredImage, ...]: + """Return all registered images in deterministic alias order.""" + operation = SlurmServiceOperation.LIST_IMAGES + manager = self._require_manager(operation) + + def list_images() -> tuple[RegisteredImage, ...]: + images = manager.list() + if not isinstance(images, tuple) or any(not isinstance(image, RegisteredImage) for image in images): + raise TypeError("image manager returned an invalid image list") + names = tuple(image.name for image in images) + if names != tuple(sorted(names)) or len(names) != len(set(names)): + raise ValueError("image manager returned an invalid image order") + return images + + return _invoke_service_backend(operation, list_images) + + def get(self, name: Identifier) -> RegisteredImage: + """Return one registered image by alias.""" + operation = SlurmServiceOperation.GET_IMAGE + normalized_name = _validate_name(name, operation) + manager = self._require_manager(operation) + + def get_image() -> RegisteredImage: + image = manager.get(normalized_name) + if not isinstance(image, RegisteredImage) or image.name != normalized_name: + raise TypeError("image manager returned an invalid image result") + return image + + return _invoke_service_backend(operation, get_image) + + def remove(self, name: Identifier, *, delete_artifact: bool = False) -> RegisteredImage: + """Unregister one alias and optionally delete its package-owned SQSH artifact.""" + operation = SlurmServiceOperation.REMOVE_IMAGE + normalized_name = _validate_name(name, operation) + if type(delete_artifact) is not bool: + raise _make_invalid_request_error(operation, "delete_artifact must be a boolean") + manager = self._require_manager(operation) + + def remove_image() -> RegisteredImage: + image = manager.remove(normalized_name, delete_artifact=delete_artifact) + if not isinstance(image, RegisteredImage) or image.name != normalized_name: + raise TypeError("image manager returned an invalid removal result") + return image + + return _invoke_service_backend(operation, remove_image) + + def _require_manager(self, operation: SlurmServiceOperation) -> SlurmImageManager: + if self._manager is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "image operations require package-owned service construction", + ) + return self._manager + + +def _validate_name(name: object, operation: SlurmServiceOperation) -> Identifier: + try: + return _IDENTIFIER_ADAPTER.validate_python(name, strict=True) + except ValidationError: + raise _make_invalid_request_error(operation, "name must be a valid identifier") from None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py new file mode 100644 index 000000000..7a9d207f7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -0,0 +1,467 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package-owned Slurm profile initialization and validation.""" + +from __future__ import annotations + +import json +import os +import shlex +import socket +import stat +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack, contextmanager +from fnmatch import fnmatchcase +from pathlib import Path + +import yaml +from pydantic import PositiveInt, ValidationError + +from data_designer.slurm.config import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + ProfileSelectionSource, + SlurmConfigLoadError, + SlurmProfile, + SlurmProfileCatalog, + load_profile_catalog, + resolve_profile, +) +from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 +from data_designer.slurm.filesystem import ( + PRIVATE_DIRECTORY_MODE, + create_restrictive_temporary_file, + open_verified_child_directory, + open_verified_directory, +) +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation + +HostnameResolver = Callable[[], tuple[str, ...]] + +_IMAGE_BUILD_CPUS = 2 +_IMAGE_BUILD_MEMORY = "8G" +_IMAGE_BUILD_TIME_LIMIT = "04:00:00" + + +class SlurmProfileInitialization(ContractValue): + """One newly created profile catalog.""" + + profile_file: str + validation_command: str + + +class SlurmProfileMatch(ContractValue): + """One cluster and its matching hostname patterns.""" + + cluster: Identifier + patterns: tuple[str, ...] + + +class SlurmProfileValidation(ContractValue): + """Effective local profile selection and workspace facts.""" + + profile_file: str | None + hostnames: tuple[str, ...] + matched_clusters: tuple[SlurmProfileMatch, ...] + default_cluster: Identifier | None + selected_cluster: Identifier | None + selection_source: ProfileSelectionSource + matched_pattern: str | None + workspace_root: str + image_root: str + registry_file: str + gpus_per_node: PositiveInt + + +class SlurmProfileService: + """Initialize and validate profiles through production loaders and selectors.""" + + def __init__( + self, + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + hostname_resolver: HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, + ) -> None: + self._profile = profile + self._catalog = catalog + self._profile_file = profile_file + self._cluster = cluster + self._launcher = launcher or SlurmCommandClient() + self._hostname_resolver = hostname_resolver or _local_hostnames + self._environ = dict(os.environ if environ is None else environ) + self._home_directory = home_directory + + def initialize( + self, + *, + workspace_root: str | Path, + image_build_partition: str, + cluster: str = "default", + account: str | None = None, + partition: str | None = None, + host_patterns: tuple[str, ...] = (), + ) -> SlurmProfileInitialization: + """Create one deterministic starter catalog without overwriting a file.""" + operation = SlurmServiceOperation.INIT_PROFILE + try: + path = _resolve_profile_path( + self._profile_file, + environ=self._environ, + home_directory=self._home_directory, + ) + payload = _starter_catalog_payload( + workspace_root=Path(workspace_root).expanduser().resolve(), + image_build_partition=image_build_partition, + cluster=cluster, + account=account, + partition=partition, + host_patterns=host_patterns, + ) + SlurmProfileCatalog.model_validate(payload, strict=True) + _create_profile_file(path, _serialize_catalog(payload, suffix=path.suffix)) + except FileExistsError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "profile file already exists", + ) from None + except FileNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile destination parent does not exist", + ) from None + except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile initialization input is invalid", + ) from None + except OSError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "profile file cannot be created", + ) from None + command = shlex.join(("data-designer", "slurm", "profile", "validate", "--profile-file", path.as_posix())) + return SlurmProfileInitialization(profile_file=path.as_posix(), validation_command=command) + + def validate(self) -> SlurmProfileValidation: + """Load, select, and verify one effective production profile.""" + operation = SlurmServiceOperation.VALIDATE_PROFILE + try: + hostnames = _normalize_hostnames(self._hostname_resolver()) + selected = resolve_profile( + profile=self._profile, + catalog=self._catalog, + profile_file=self._profile_file, + cluster=self._cluster, + hostnames=hostnames, + environ=self._environ, + home_directory=self._home_directory, + ) + catalog = self._catalog + if selected.catalog_path is not None: + catalog = load_profile_catalog(selected.catalog_path) + if compute_canonical_json_sha256(catalog.model_dump(mode="json")) != selected.catalog_sha256: + raise SlurmConfigLoadError("profile catalog changed while it was being validated") + _validate_workspace(Path(selected.profile.workspace_root)) + gpus_per_node = self._resolve_gpu_count(selected.profile) + except SlurmServiceError: + raise + except SlurmLauncherError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm is unavailable", + ) from None + except (SlurmConfigLoadError, ValidationError, ValueError, TypeError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "profile configuration cannot be resolved", + ) from None + except OSError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "profile workspace is unavailable", + ) from None + + store = ImageRegistryStore(selected.profile.workspace_root) + return SlurmProfileValidation( + profile_file=selected.catalog_path, + hostnames=hostnames, + matched_clusters=_matching_clusters(catalog, hostnames), + default_cluster=None if catalog is None else catalog.default_cluster, + selected_cluster=selected.cluster_name, + selection_source=selected.selection_source, + matched_pattern=selected.matched_pattern, + workspace_root=selected.profile.workspace_root, + image_root=store.image_root.as_posix(), + registry_file=store.registry_path.as_posix(), + gpus_per_node=gpus_per_node, + ) + + def _resolve_gpu_count(self, profile: SlurmProfile) -> int: + if profile.gpus_per_node != "auto": + return profile.gpus_per_node + counts = tuple(sorted(set(self._launcher.query_gpu_counts(partition=profile.scheduler.partition)))) + if len(counts) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.VALIDATE_PROFILE, + "eligible Slurm nodes do not report one GPU count", + ) + return counts[0] + + +def create_slurm_profile_service( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + hostname_resolver: HostnameResolver | None = None, + environ: Mapping[str, str] | None = None, + home_directory: str | Path | None = None, +) -> SlurmProfileService: + """Create the package-owned profile service.""" + return SlurmProfileService( + profile=profile, + catalog=catalog, + profile_file=profile_file, + cluster=cluster, + launcher=launcher, + hostname_resolver=hostname_resolver, + environ=environ, + home_directory=home_directory, + ) + + +def _starter_catalog_payload( + *, + workspace_root: Path, + image_build_partition: str, + cluster: str, + account: str | None, + partition: str | None, + host_patterns: tuple[str, ...], +) -> dict[str, object]: + scheduler = {key: value for key, value in (("account", account), ("partition", partition)) if value is not None} + profile: dict[str, object] = { + "schema_version": 1, + "host_patterns": list(host_patterns), + "gpus_per_node": "auto", + "workspace_root": workspace_root.as_posix(), + "image_build": { + "partition": image_build_partition, + "cpus_per_task": _IMAGE_BUILD_CPUS, + "memory": _IMAGE_BUILD_MEMORY, + "time_limit": _IMAGE_BUILD_TIME_LIMIT, + }, + "gpu_request_mode": "gres", + } + if scheduler: + profile["scheduler"] = scheduler + return { + "schema_version": 1, + "default_cluster": cluster, + "clusters": {cluster: profile}, + } + + +def _serialize_catalog(payload: dict[str, object], *, suffix: str) -> bytes: + if suffix == ".json": + return (json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n").encode() + return yaml.safe_dump(payload, allow_unicode=True, sort_keys=False).encode() + + +def _create_profile_file(path: Path, content: bytes) -> None: + with _open_or_create_profile_parent(path.parent) as parent_descriptor: + descriptor, temporary_name = create_restrictive_temporary_file( + parent_descriptor, + prefix=f".{path.name}.", + suffix=".tmp", + ) + try: + with os.fdopen(descriptor, "wb") as output: + descriptor = -1 + output.write(content) + output.flush() + os.fsync(output.fileno()) + os.link( + temporary_name, + path.name, + src_dir_fd=parent_descriptor, + dst_dir_fd=parent_descriptor, + follow_symlinks=False, + ) + os.fsync(parent_descriptor) + finally: + if descriptor >= 0: + os.close(descriptor) + try: + os.unlink(temporary_name, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + except FileNotFoundError: + pass + + +@contextmanager +def _open_or_create_profile_parent(parent: Path) -> Iterator[int]: + missing: list[Path] = [] + directory = parent + while not directory.exists(): + missing.append(directory) + directory = directory.parent + with ExitStack() as stack: + descriptor = stack.enter_context(open_verified_directory(directory, resource_name="profile")) + for child in reversed(missing): + created = False + try: + os.mkdir(child.name, PRIVATE_DIRECTORY_MODE, dir_fd=descriptor) + created = True + except FileExistsError: + pass + if created: + _set_created_directory_mode(descriptor, child.name) + parent_descriptor = descriptor + descriptor = stack.enter_context( + open_verified_child_directory( + parent_descriptor, + child.name, + child, + resource_name="profile", + require_private=created, + ) + ) + if created: + os.fsync(descriptor) + os.fsync(parent_descriptor) + yield descriptor + + +def _set_created_directory_mode(parent_descriptor: int, name: str) -> None: + try: + os.chmod(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor, follow_symlinks=False) + except (NotImplementedError, ValueError) as exc: + # Older Linux fchmodat rejects AT_SYMLINK_NOFOLLOW even for real directories. + flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + try: + descriptor = os.open(name, flags, dir_fd=parent_descriptor) + except PermissionError: + pass + else: + try: + os.fchmod(descriptor, PRIVATE_DIRECTORY_MODE) + finally: + os.close(descriptor) + return + + # O_PATH can open a directory whose owner access was removed by umask. + path_flag = getattr(os, "O_PATH", None) + if path_flag is None: + raise OSError("profile directory permissions cannot be set safely") from exc + descriptor = os.open( + name, + path_flag | os.O_DIRECTORY | os.O_NOFOLLOW, + dir_fd=parent_descriptor, + ) + try: + if not stat.S_ISDIR(os.fstat(descriptor).st_mode): + raise OSError("profile directory changed while permissions were being set") + os.chmod(f"/proc/self/fd/{descriptor}", PRIVATE_DIRECTORY_MODE) + finally: + os.close(descriptor) + + +def _validate_workspace(path: Path) -> None: + with open_verified_directory(path, resource_name="profile workspace") as descriptor: + temporary_descriptor, temporary_name = create_restrictive_temporary_file( + descriptor, + prefix=".data-designer-profile-validation-", + suffix=".tmp", + ) + try: + os.close(temporary_descriptor) + temporary_descriptor = -1 + os.unlink(temporary_name, dir_fd=descriptor) + os.fsync(descriptor) + finally: + if temporary_descriptor >= 0: + os.close(temporary_descriptor) + try: + os.unlink(temporary_name, dir_fd=descriptor) + except FileNotFoundError: + pass + + +def _matching_clusters( + catalog: SlurmProfileCatalog | None, + hostnames: tuple[str, ...], +) -> tuple[SlurmProfileMatch, ...]: + if catalog is None: + return () + matches = [] + for cluster, profile in sorted(catalog.clusters.items()): + patterns = tuple( + sorted( + pattern + for pattern in profile.host_patterns + if any(fnmatchcase(hostname, pattern.casefold()) for hostname in hostnames) + ) + ) + if patterns: + matches.append(SlurmProfileMatch(cluster=cluster, patterns=patterns)) + return tuple(matches) + + +def _resolve_profile_path( + explicit_path: str | Path | None, + *, + environ: Mapping[str, str], + home_directory: str | Path | None, +) -> Path: + source = explicit_path + if source is None: + source = environ.get(PROFILE_FILE_ENVIRONMENT) + if source == "": + raise SlurmConfigLoadError(f"{PROFILE_FILE_ENVIRONMENT} must not be empty") + if source is None: + home = Path.home() if home_directory is None else Path(home_directory) + source = home / DEFAULT_PROFILE_FILE_NAME + expanded = Path(source).expanduser() + path = expanded.parent.resolve() / expanded.name + if path.suffix not in {".json", ".yaml", ".yml"}: + raise SlurmConfigLoadError("configuration path must end in .json, .yaml, or .yml") + return path + + +def _normalize_hostnames(hostnames: tuple[str, ...]) -> tuple[str, ...]: + return tuple(dict.fromkeys(hostname.strip().casefold() for hostname in hostnames if hostname.strip())) + + +def _local_hostnames() -> tuple[str, ...]: + return socket.gethostname(), socket.getfqdn() + + +__all__ = [ + "HostnameResolver", + "SlurmProfileInitialization", + "SlurmProfileMatch", + "SlurmProfileService", + "SlurmProfileValidation", + "create_slurm_profile_service", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py new file mode 100644 index 000000000..cecbbbf14 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stable values returned by public Slurm operations.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + AttemptId, + ContractValue, + Identifier, + Sha256Digest, + ShardId, + validate_absolute_path, +) +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CollectionState, + RunManifest, + ShardManifest, + ShardWinner, +) + + +class SlurmRunExecution(ContractValue): + """One rendered dry run or accepted Slurm submission.""" + + run_id: Identifier + state: Literal["dry_run", "submitted"] + plan_sha256: Sha256Digest + shard_count: PositiveInt + job_id: PositiveInt | None = None + batch_script: str | None = None + + @model_validator(mode="after") + def validate_state(self) -> SlurmRunExecution: + if self.state == "dry_run": + if self.job_id is not None or not self.batch_script: + raise ValueError("dry-run execution requires only a rendered batch script") + elif self.job_id is None or self.batch_script is not None: + raise ValueError("submitted execution requires only a Slurm job ID") + return self + + +class SlurmPersistedAttemptStatus(ContractValue): + """Persisted attempt state returned after scheduler reconciliation.""" + + attempt: AttemptManifest + readiness: AttemptReadiness | None = None + + @model_validator(mode="after") + def validate_readiness(self) -> SlurmPersistedAttemptStatus: + if self.readiness is not None and ( + self.readiness.run_id, + self.readiness.shard_id, + self.readiness.attempt_id, + ) != (self.attempt.run_id, self.attempt.shard_id, self.attempt.attempt_id): + raise ValueError("readiness identity does not match its attempt") + return self + + +class SlurmPersistedShardStatus(ContractValue): + """Persisted attempt and winner state for one shard.""" + + shard: ShardManifest + attempts: tuple[SlurmPersistedAttemptStatus, ...] + winner: ShardWinner | None = None + + @model_validator(mode="after") + def validate_shard(self) -> SlurmPersistedShardStatus: + if any(status.attempt.shard_id != self.shard.shard_id for status in self.attempts): + raise ValueError("persisted status contains an attempt for another shard") + if self.winner is not None and ( + self.winner.run_id != self.shard.run_id or self.winner.shard_id != self.shard.shard_id + ): + raise ValueError("persisted winner does not match its shard") + return self + + +class SlurmPersistedRunStatus(ContractValue): + """Status assembled only from durable run records.""" + + run: RunManifest + shards: tuple[SlurmPersistedShardStatus, ...] + + @model_validator(mode="after") + def validate_run(self) -> SlurmPersistedRunStatus: + if len(self.shards) != self.run.shard_count: + raise ValueError("persisted status must contain every run shard") + if tuple(status.shard.shard_index for status in self.shards) != tuple(range(self.run.shard_count)): + raise ValueError("persisted status shards must remain in planned order") + if any(status.shard.run_id != self.run.run_id for status in self.shards): + raise ValueError("persisted status contains a shard for another run") + return self + + +class SlurmRunCancellation(ContractValue): + """Managed Slurm array jobs selected for cancellation.""" + + run_id: Identifier + job_ids: tuple[PositiveInt, ...] + + @model_validator(mode="after") + def validate_jobs(self) -> SlurmRunCancellation: + if self.job_ids != tuple(sorted(set(self.job_ids))): + raise ValueError("cancelled job IDs must be sorted and unique") + return self + + +class SlurmRetryExecution(ContractValue): + """One rendered retry dry run or accepted sparse retry submission.""" + + run_id: Identifier + state: Literal["dry_run", "submitted"] + shard_ids: tuple[ShardId, ...] + attempt_ids: tuple[AttemptId, ...] + effective_resume_mode: Literal["never", "always"] + job_id: PositiveInt | None = None + batch_script: str | None = None + + @model_validator(mode="after") + def validate_execution(self) -> SlurmRetryExecution: + if not self.shard_ids or self.shard_ids != tuple(sorted(set(self.shard_ids))): + raise ValueError("retry shard IDs must be non-empty, sorted, and unique") + if len(self.attempt_ids) != len(self.shard_ids): + raise ValueError("retry attempt IDs must correspond to the selected shards") + if self.state == "dry_run": + if self.job_id is not None or not self.batch_script: + raise ValueError("dry-run retry requires only a rendered batch script") + elif self.job_id is None or self.batch_script is not None: + raise ValueError("submitted retry requires only a Slurm job ID") + return self + + +class SlurmCollectionExecution(ContractValue): + """One accepted or previously active collection submission.""" + + run_id: Identifier + collection_id: Identifier + state: CollectionState + job_id: PositiveInt + output_path: str + num_partitions: PositiveInt + + _output_path_is_absolute = field_validator("output_path")(validate_absolute_path) + + +__all__ = [ + "SlurmCollectionExecution", + "SlurmPersistedAttemptStatus", + "SlurmPersistedRunStatus", + "SlurmPersistedShardStatus", + "SlurmRetryExecution", + "SlurmRunCancellation", + "SlurmRunExecution", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py new file mode 100644 index 000000000..51d01b98d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production retry and collection adapters for the public run service.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from datetime import datetime +from pathlib import Path +from typing import Literal + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.services.results import SlurmCollectionExecution, SlurmRetryExecution +from data_designer.slurm.state import ( + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.outputs import RetryPlan +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RunRetryCollectionBackend: + """Adapt public run operations to persisted retry and collection capabilities.""" + + def __init__( + self, + workspace_root: str, + launcher: SlurmCommandClient, + clock: Callable[[], datetime], + source_environment: Mapping[str, str], + ) -> None: + self._workspace_root = workspace_root + self._launcher = launcher + self._clock = clock + self._export_environment = {"NEMO_DEPLOYMENT_TYPE": source_environment.get("NEMO_DEPLOYMENT_TYPE", "slurm")} + + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry from persisted run state.""" + operation = SlurmServiceOperation.RETRY_RUN + run_id = self._resolve_run_reference(run_or_job_id, operation) + try: + observed_at = self._clock() + coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher, self._export_environment) + resolved_plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) + preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) + if preview is not None and effective_resume_mode not in {None, preview[0].effective_resume_mode}: + preview = None + if preview is not None: + effective_resume_mode = preview[0].effective_resume_mode + shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) + elif effective_resume_mode is None: + preview = coordinator.preview( + shard_ids=shard_ids, + effective_resume_mode="never", + observed_at=observed_at, + ) + effective_resume_mode = self._resolve_if_possible_resume_mode(resolved_plan, preview[0], operation) + shard_ids = tuple(shard.shard_id for shard in preview[0].planned_shards) + if dry_run: + if preview is None or preview[0].effective_resume_mode != effective_resume_mode: + preview = coordinator.preview( + shard_ids=shard_ids, + effective_resume_mode=effective_resume_mode, + observed_at=observed_at, + ) + assert preview is not None + plan, batch_script = preview + return SlurmRetryExecution( + run_id=run_id, + state="dry_run", + shard_ids=tuple(shard.shard_id for shard in plan.planned_shards), + attempt_ids=tuple(shard.attempt_id for shard in plan.planned_shards), + effective_resume_mode=effective_resume_mode, + batch_script=batch_script, + ) + attempts = tuple( + sorted( + coordinator.retry( + shard_ids=shard_ids, + effective_resume_mode=effective_resume_mode, + observed_at=observed_at, + ), + key=lambda attempt: attempt.shard_id, + ) + ) + scheduler_ids = {attempt.scheduler.array_job_id for attempt in attempts if attempt.scheduler is not None} + if not attempts or len(scheduler_ids) != 1 or any(attempt.scheduler is None for attempt in attempts): + raise SlurmStateError("retry did not return one accepted sparse array") + return SlurmRetryExecution( + run_id=run_id, + state="submitted", + shard_ids=tuple(attempt.shard_id for attempt in attempts), + attempt_ids=tuple(attempt.attempt_id for attempt in attempts), + effective_resume_mode=effective_resume_mode, + job_id=scheduler_ids.pop(), + ) + except StateNotFoundError: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") from None + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None + except StateCorruptionError: + raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "retry state is inconsistent") from None + except SlurmStateError as error: + if _has_cause(error, SlurmLauncherError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm retry submission is unavailable", + ) from None + raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "retry state cannot be read") from None + + def collect( + self, + input_path: Identifier | Path, + *, + destination: Path | None, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + operation = SlurmServiceOperation.COLLECT_RUN + run_path = ( + (Path(self._workspace_root) / "runs" / input_path).resolve() if isinstance(input_path, str) else input_path + ) + run_id = self._resolve_run_input_path(run_path, operation) + try: + plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + try: + CollectionDestinationResolver().resolve(plan, destination) + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, str(error)) from None + if num_partitions is not None and num_partitions > plan.invocation.authored.num_records: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "num_partitions must not exceed requested records", + ) + status = SlurmCollectionCoordinator( + self._workspace_root, run_id, self._launcher, self._export_environment + ).submit( + destination=destination, + num_partitions=num_partitions, + submitted_at=self._clock(), + ) + if status.scheduler is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm collection submission was not accepted", + ) + collection_plan = CollectionStorage(StateStorage(Path(self._workspace_root), run_id)).read_plan( + status.collection_id + ) + return SlurmCollectionExecution( + run_id=run_id, + collection_id=status.collection_id, + state=status.state, + job_id=status.scheduler, + output_path=collection_plan.host_destination, + num_partitions=collection_plan.num_partitions, + ) + except SlurmServiceError: + raise + except StateNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + operation, + "run state or shard winner was not found", + ) from None + except StateConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None + except StateCorruptionError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "collection state is inconsistent", + ) from None + except SlurmStateError as error: + if _has_cause(error, SlurmLauncherError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm collection submission is unavailable", + ) from None + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "collection state cannot be read", + ) from None + + @staticmethod + def _resolve_retry_resume_mode( + plan: ResolvedSlurmRunPlan, + requested: Literal["never", "always", "if_possible"], + ) -> Literal["never", "always"] | None: + if requested != "if_possible": + return requested + pinned = plan.invocation.authored.resume + return None if pinned == "if_possible" else pinned + + @staticmethod + def _resolve_if_possible_resume_mode( + plan: ResolvedSlurmRunPlan, + retry_plan: RetryPlan, + operation: SlurmServiceOperation, + ) -> Literal["never", "always"]: + workspaces = {shard.shard_id: Path(shard.resume_workspace.path) for shard in plan.shards} + availability = {_has_resume_data(workspaces[shard.shard_id]) for shard in retry_plan.planned_shards} + if len(availability) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "retry selection mixes resumable and fresh shards; choose --resume never or --resume always", + ) + return "always" if availability.pop() else "never" + + def _resolve_run_reference( + self, + run_or_job_id: Identifier, + operation: SlurmServiceOperation, + ) -> Identifier: + workspace_root = Path(self._workspace_root) + direct = workspace_root / "runs" / run_or_job_id + direct_match = run_or_job_id if direct.is_dir() else None + if direct_match is not None and not run_or_job_id.isdecimal(): + return run_or_job_id + if not run_or_job_id.isdecimal() or int(run_or_job_id) <= 0: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") + job_id = int(run_or_job_id) + matches: list[Identifier] = [] + runs_root = workspace_root / "runs" + if runs_root.is_dir(): + for candidate in sorted(runs_root.iterdir(), key=lambda path: path.name): + if not candidate.is_dir(): + continue + try: + candidate_id = _IDENTIFIER_ADAPTER.validate_python(candidate.name, strict=True) + writer = SlurmStateWriter(workspace_root, candidate_id) + attempts = tuple( + attempt for shard in writer.load_shards() for attempt in writer.load_attempts(shard.shard_id) + ) + except (ValidationError, StateNotFoundError): + continue + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + operation, + "managed run state cannot be searched safely", + ) from None + if any( + attempt.scheduler is not None and attempt.scheduler.array_job_id == job_id for attempt in attempts + ): + matches.append(candidate_id) + targets = set(matches) + if direct_match is not None: + targets.add(direct_match) + if not targets: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "managed Slurm job was not found") + if len(targets) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + operation, + "numeric reference matches multiple managed runs", + ) + return targets.pop() + + def _resolve_run_input_path( + self, + input_path: Path, + operation: SlurmServiceOperation, + ) -> Identifier: + runs_root = (Path(self._workspace_root) / "runs").resolve() + if input_path.parent != runs_root: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "input_path must identify a managed run directory", + ) + try: + run_id = _IDENTIFIER_ADAPTER.validate_python(input_path.name, strict=True) + except ValidationError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "input_path must contain a valid managed run ID", + ) from None + if not input_path.is_dir(): + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") + return run_id + + +def _has_cause(error: BaseException, expected_type: type[BaseException]) -> bool: + cause = error.__cause__ + while cause is not None: + if isinstance(cause, expected_type): + return True + cause = cause.__cause__ + return False + + +def _has_resume_data(path: Path) -> bool: + return not path.is_symlink() and path.is_dir() and next(path.iterdir(), None) is not None + + +__all__ = ["RunRetryCollectionBackend"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py new file mode 100644 index 000000000..4898c2820 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public run-planning facade for Data Designer Slurm.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import partial +from pathlib import Path +from typing import Literal, Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import Identifier, ShardId +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.errors import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, + _invoke_service_backend, + _make_invalid_request_error, +) +from data_designer.slurm.services.results import ( + SlurmCollectionExecution, + SlurmPersistedRunStatus, + SlurmRetryExecution, + SlurmRunCancellation, + SlurmRunExecution, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SHARD_IDS_ADAPTER = TypeAdapter(tuple[ShardId, ...]) + + +class SlurmRunPlanner(Protocol): + """Resolve authored run configs through a supported service dependency.""" + + def plan(self, config: DataDesignerSlurmConfig, *, source_root: Path) -> ResolvedSlurmRunPlan: + """Return the immutable plan for one authored run. + + Any non-``INTERNAL`` service error must contain a caller-safe message. + """ + + +class SlurmBatchScriptRenderer(Protocol): + """Render immutable plans through a supported service dependency.""" + + def __call__(self, plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: + """Return the deterministic batch script for one attempt. + + Any non-``INTERNAL`` service error must contain a caller-safe message. + """ + + +class SlurmRunBackend(Protocol): + """Execute and inspect runs through a supported service dependency.""" + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + ) -> SlurmRunExecution: + """Render or submit one run.""" + + def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: + """Return the durable status for one run.""" + + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: + """Request cancellation of active jobs.""" + + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + """Render or submit one sparse retry.""" + + def collect( + self, + input_path: Identifier | Path, + *, + destination: Path | None, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + + +class SlurmRunService: + """Coordinate public run operations through package-owned boundaries. + + The service borrows its injected dependencies and does not manage their + lifecycle. + + Args: + planner: Run-planning dependency implementing ``SlurmRunPlanner``. + renderer: Batch-rendering dependency implementing ``SlurmBatchScriptRenderer``. + backend: Optional run-operation dependency. Package-owned factories provide it. + """ + + def __init__( + self, + planner: SlurmRunPlanner, + renderer: SlurmBatchScriptRenderer, + backend: SlurmRunBackend | None = None, + ) -> None: + self._planner = planner + self._renderer = renderer + self._backend = backend + + def plan( + self, + config: DataDesignerSlurmConfig, + *, + source_root: str | Path = ".", + ) -> ResolvedSlurmRunPlan: + """Resolve and compile one run without rendering or submission. + + The returned plan must reference the exact serialized authored config. + + Raises: + SlurmServiceError: If the request is invalid or planning fails. + """ + operation = SlurmServiceOperation.PLAN_RUN + if not isinstance(config, DataDesignerSlurmConfig): + raise _make_invalid_request_error(operation, "config must be a DataDesignerSlurmConfig") + if not isinstance(source_root, str | Path): + raise _make_invalid_request_error(operation, "source_root must be a path") + + def build_plan() -> ResolvedSlurmRunPlan: + plan = self._planner.plan(config, source_root=Path(source_root).expanduser().resolve()) + if not isinstance(plan, ResolvedSlurmRunPlan): + raise TypeError("run planner returned an invalid result") + if plan.authored_config.sha256 != config.compute_sha256(): + raise ValueError("run plan does not match the requested authored config") + return plan + + return _invoke_service_backend(operation, build_plan) + + def render_attempt( + self, + resolved_plan: ResolvedSlurmRunPlan, + *, + attempt_ordinal: int = 1, + ) -> str: + """Render one attempt from an already resolved immutable plan. + + Rendering never re-resolves the authored config or submits the script. + + Raises: + SlurmServiceError: If the request is invalid or rendering fails. + """ + operation = SlurmServiceOperation.RENDER_ATTEMPT + if not isinstance(resolved_plan, ResolvedSlurmRunPlan): + raise _make_invalid_request_error(operation, "resolved_plan must be a ResolvedSlurmRunPlan") + if type(attempt_ordinal) is not int or attempt_ordinal <= 0: + raise _make_invalid_request_error(operation, "attempt_ordinal must be a positive integer") + + def render() -> str: + script = self._renderer(resolved_plan, attempt_ordinal=attempt_ordinal) + if type(script) is not str or not script: + raise TypeError("batch renderer returned an invalid script") + return script + + return _invoke_service_backend(operation, render) + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: str | Path = ".", + dry_run: bool = False, + ) -> SlurmRunExecution: + """Render or submit one run through package-owned production wiring.""" + operation = SlurmServiceOperation.EXECUTE_RUN + if not isinstance(config, DataDesignerSlurmConfig): + raise _make_invalid_request_error(operation, "config must be a DataDesignerSlurmConfig") + if not isinstance(source_root, str | Path): + raise _make_invalid_request_error(operation, "source_root must be a path") + if type(dry_run) is not bool: + raise _make_invalid_request_error(operation, "dry_run must be a boolean") + backend = self._require_backend(operation) + + def execute_run() -> SlurmRunExecution: + result = backend.execute( + config, + source_root=Path(source_root).expanduser().resolve(), + dry_run=dry_run, + ) + if not isinstance(result, SlurmRunExecution): + raise TypeError("run backend returned an invalid execution result") + return result + + return _invoke_service_backend(operation, execute_run) + + def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: + """Reconcile scheduler observations and return persisted records.""" + operation = SlurmServiceOperation.STATUS_RUN + normalized_run_id = _validate_run_id(run_id, operation) + backend = self._require_backend(operation) + + def load_status() -> SlurmPersistedRunStatus: + result = backend.status(normalized_run_id) + if not isinstance(result, SlurmPersistedRunStatus) or result.run.run_id != normalized_run_id: + raise TypeError("run backend returned an invalid status result") + return result + + return _invoke_service_backend(operation, load_status) + + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: + """Request cancellation of active jobs; reconciliation updates persisted state.""" + operation = SlurmServiceOperation.CANCEL_RUN + normalized_run_id = _validate_run_id(run_id, operation) + backend = self._require_backend(operation) + + def cancel_run() -> SlurmRunCancellation: + result = backend.cancel(normalized_run_id) + if not isinstance(result, SlurmRunCancellation) or result.run_id != normalized_run_id: + raise TypeError("run backend returned an invalid cancellation result") + return result + + return _invoke_service_backend(operation, cancel_run) + + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: Sequence[ShardId] | None = None, + resume: Literal["never", "always", "if_possible"] = "if_possible", + dry_run: bool = False, + ) -> SlurmRetryExecution: + """Render or submit retry attempts for failed shards.""" + operation = SlurmServiceOperation.RETRY_RUN + normalized_reference = _validate_run_id(run_or_job_id, operation) + normalized_shards = _validate_shard_ids(shard_ids, operation) + if resume not in {"never", "always", "if_possible"}: + raise _make_invalid_request_error(operation, "resume must be 'never', 'always', or 'if_possible'") + if type(dry_run) is not bool: + raise _make_invalid_request_error(operation, "dry_run must be a boolean") + backend = self._require_backend(operation) + + return _invoke_service_backend( + operation, + partial( + _retry_run, + backend, + normalized_reference, + shard_ids=normalized_shards, + resume=resume, + dry_run=dry_run, + ), + ) + + def collect( + self, + input_path: str | Path, + *, + destination: str | Path | None = None, + num_partitions: int | None = None, + ) -> SlurmCollectionExecution: + """Submit or recover collection by run ID or managed run directory.""" + operation = SlurmServiceOperation.COLLECT_RUN + if not isinstance(input_path, str | Path): + raise _make_invalid_request_error(operation, "input_path must be a run ID or path") + if destination is not None and not isinstance(destination, str | Path): + raise _make_invalid_request_error(operation, "destination must be a path") + if num_partitions is not None and (type(num_partitions) is not int or num_partitions <= 0): + raise _make_invalid_request_error(operation, "num_partitions must be a positive integer") + normalized_input = ( + _validate_run_id(input_path, operation) + if isinstance(input_path, str) and "/" not in input_path + else Path(input_path).expanduser().resolve() + ) + normalized_destination = None if destination is None else Path(destination).expanduser().resolve() + backend = self._require_backend(operation) + + return _invoke_service_backend( + operation, + partial( + _collect_run, + backend, + normalized_input, + destination=normalized_destination, + num_partitions=num_partitions, + ), + ) + + def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: + if self._backend is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "run operations require package-owned service construction", + ) + return self._backend + + +def _retry_run( + backend: SlurmRunBackend, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, +) -> SlurmRetryExecution: + result = backend.retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + if not isinstance(result, SlurmRetryExecution): + raise TypeError("run backend returned an invalid retry result") + if shard_ids is not None and result.shard_ids != shard_ids: + raise TypeError("run backend returned retry shards that do not match the request") + if result.state != ("dry_run" if dry_run else "submitted"): + raise TypeError("run backend returned a retry state that does not match the request") + return result + + +def _collect_run( + backend: SlurmRunBackend, + input_path: Identifier | Path, + *, + destination: Path | None, + num_partitions: int | None, +) -> SlurmCollectionExecution: + result = backend.collect( + input_path, + destination=destination, + num_partitions=num_partitions, + ) + if not isinstance(result, SlurmCollectionExecution): + raise TypeError("run backend returned an invalid collection result") + if (destination is not None and result.output_path != destination.as_posix()) or ( + num_partitions is not None and result.num_partitions != num_partitions + ): + raise TypeError("run backend returned collection intent that does not match the request") + return result + + +def _validate_run_id(run_id: object, operation: SlurmServiceOperation) -> Identifier: + try: + return _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except ValidationError: + raise _make_invalid_request_error(operation, "run_id must be a valid identifier") from None + + +def _validate_shard_ids( + shard_ids: Sequence[ShardId] | None, + operation: SlurmServiceOperation, +) -> tuple[ShardId, ...] | None: + if shard_ids is None: + return None + if isinstance(shard_ids, str | bytes): + raise _make_invalid_request_error(operation, "shard_ids must be a sequence of shard identifiers") + try: + normalized = _SHARD_IDS_ADAPTER.validate_python(tuple(shard_ids), strict=True) + except (TypeError, ValidationError): + raise _make_invalid_request_error(operation, "shard_ids must contain valid shard identifiers") from None + if not normalized or len(normalized) != len(set(normalized)): + raise _make_invalid_request_error(operation, "shard_ids must be non-empty and unique when provided") + return tuple(sorted(normalized)) 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 new file mode 100644 index 000000000..a29a7674c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -0,0 +1,1105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package-owned production wiring for public Slurm services.""" + +from __future__ import annotations + +import importlib.metadata +import os +import shlex +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Literal, Protocol, TypeVar +from uuid import uuid4 + +from pydantic import JsonValue, ValidationError + +from data_designer.slurm._errors import format_validation_error +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, + ResolvedClientDependencies, +) +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ensure_private_directory +from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector +from data_designer.slurm.config import ( + DataDesignerSlurmConfig, + ImageBuildRequest, + ImageKind, + ImageRef, + SelectedSlurmProfile, + SlurmConfigLoadError, + SlurmProfile, + SlurmProfileCatalog, + collect_secret_environment_names, + load_builder_payload, + resolve_profile, +) +from data_designer.slurm.contracts import Identifier, ShardId +from data_designer.slurm.images.errors import ( + ImageArtifactRemovalError, + ImageConflictError, + ImageNotFoundError, + SlurmImageError, +) +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.images.service import VerifiedImageRegistry +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +from data_designer.slurm.launcher.scheduler_path import SchedulerPathError, resolve_scheduler_bin_path +from data_designer.slurm.planning import ResolvedImage, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.runtime.bundle import stage_runtime_bundle +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.security import redact_sensitive_diagnostic +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 +from data_designer.slurm.services.results import ( + SlurmCollectionExecution, + SlurmPersistedAttemptStatus, + SlurmPersistedRunStatus, + SlurmPersistedShardStatus, + SlurmRetryExecution, + SlurmRunCancellation, + SlurmRunExecution, +) +from data_designer.slurm.services.retry_collection import RunRetryCollectionBackend +from data_designer.slurm.services.run import SlurmRunService +from data_designer.slurm.serving.resolver import resolve_vllm_server +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptStatus, + AttemptTerminalClassification, + EffectiveAttemptState, + SchedulerState, + SlurmStateError, + SlurmStateReconciler, + SlurmStateWriter, + StateConflictError, + StateNotFoundError, +) + +RunIdFactory = Callable[[], str] +Clock = Callable[[], datetime] +_ResultT = TypeVar("_ResultT") +_MAX_VISIBLE_JOB_IDS = 16 +_ACTIVE_ATTEMPT_STATES = frozenset( + {AttemptLifecycleState.SUBMITTED, AttemptLifecycleState.PENDING, AttemptLifecycleState.RUNNING} +) +_FAILURE_CLASSIFICATIONS = { + SchedulerState.FAILED: AttemptTerminalClassification.FAILED, + SchedulerState.CANCELLED: AttemptTerminalClassification.CANCELLED, + SchedulerState.TIMED_OUT: AttemptTerminalClassification.TIMED_OUT, + SchedulerState.NODE_FAILED: AttemptTerminalClassification.NODE_FAILED, + SchedulerState.PREEMPTED: AttemptTerminalClassification.PREEMPTED, + SchedulerState.REQUEUED: AttemptTerminalClassification.REQUEUED, + SchedulerState.OUT_OF_MEMORY: AttemptTerminalClassification.OUT_OF_MEMORY, +} + + +class SlurmRunArtifactPublisher(Protocol): + """Persist submission-critical run artifacts through the STATE boundary.""" + + def initialize( + self, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, JsonValue] | None, + *, + force: bool, + ) -> None: + """Atomically publish all immutable run inputs before submission.""" + + 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, + job_id: int, + *, + failed_at: datetime, + ) -> None: + """Fail initial attempts owned by the cancelled held job.""" + + +@dataclass(frozen=True, slots=True) +class _PreparedRun: + plan: ResolvedSlurmRunPlan + dependencies: ResolvedClientDependencies + builder_payload: dict[str, JsonValue] | None + batch_script: str + + +class _RunPreparer: + def __init__( + self, + selected_profile: SelectedSlurmProfile, + image_registry: VerifiedImageRegistry, + dependency_resolver: ClientDependencyResolver, + launcher: SlurmCommandClient, + run_id_factory: RunIdFactory, + package_version: str, + runtime_inspector: ClientRuntimeInspector | None = None, + ) -> None: + self._profile = selected_profile + self._images = image_registry + self._dependencies = dependency_resolver + self._launcher = launcher + self._run_id_factory = run_id_factory + self._package_version = package_version + self._runtime_inspector = runtime_inspector or ClientRuntimeInspector() + + @contextmanager + def prepare( + self, + authored: DataDesignerSlurmConfig, + *, + source_root: Path, + operation: SlurmServiceOperation, + ) -> Iterator[_PreparedRun]: + stack = ExitStack() + try: + yield self._prepare(authored, source_root=source_root, operation=operation, stack=stack) + finally: + stack.close() + + def _prepare( + self, + authored: DataDesignerSlurmConfig, + *, + source_root: Path, + operation: SlurmServiceOperation, + stack: ExitStack, + ) -> _PreparedRun: + try: + run_id = self._run_id_factory() + try: + scheduler_bin_path = resolve_scheduler_bin_path(self._profile.profile.scheduler.bin_path) + except SchedulerPathError as error: + raise SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, operation, str(error)) from None + workspace_root = self._profile.profile.workspace_root + run_root = Path(workspace_root) / "runs" / run_id + client_runtime = self._runtime_inspector.inspect() + try: + builder_payload = _resolve_builder_payload(authored, source_root) + except SlurmConfigLoadError as error: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + _safe_builder_load_error(error), + ) from None + _validate_builder_for_submission(authored, builder_payload, operation) + dependencies = stack.enter_context( + self._dependencies.resolve( + authored.client.dependencies, + client_runtime, + run_root=run_root, + source_root=source_root, + ) + ) + runtime_bundle = stage_runtime_bundle(workspace_root) + try: + stack.enter_context(ImageRegistryStore(workspace_root).acquire_run_reference_lock()) + except (SlurmImageError, OSError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "image registry cannot coordinate run publication", + ) from None + deployment_images = tuple( + self._images.resolve_for_planning(deployment.server.image, expected_kind=ImageKind.SERVING) + for deployment in authored.deployments + ) + effective = resolve_slurm_config( + authored, + selected_profile=self._profile, + client_runtime=client_runtime, + deployment_images=deployment_images, + dependency_lock=dependencies.lock, + runtime_bundle=runtime_bundle, + run_id=run_id, + package_version=self._package_version, + scheduler_bin_path=scheduler_bin_path, + resolved_gpus_per_node=self._resolve_gpu_count(authored, operation), + builder_payload=builder_payload, + ) + plan = SlurmRunCompiler.compile(effective) + for deployment in plan.deployments: + resolve_vllm_server(plan, deployment.deployment_id) + return _PreparedRun( + plan=plan, + dependencies=dependencies, + builder_payload=builder_payload, + batch_script=render_generation_attempt_script(plan, attempt_ordinal=1), + ) + except SlurmServiceError: + raise + except ImageNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, operation, "required image is not registered" + ) from None + except SlurmImageError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, operation, "required image cannot be verified" + ) from None + except ClientRuntimeInspectionError as error: + raise SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, str(error)) from None + except (ClientDependencyResolutionError, SlurmConfigLoadError, ValueError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, operation, "run configuration cannot be prepared" + ) from None + except SlurmLauncherError: + raise SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, operation, "Slurm is unavailable") from None + except (SlurmRuntimeError, OSError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, operation, "run artifacts cannot be staged" + ) from None + + def _resolve_gpu_count( + self, + authored: DataDesignerSlurmConfig, + operation: SlurmServiceOperation, + ) -> int | None: + if self._profile.profile.gpus_per_node != "auto": + return None + partition = authored.submission.partition or self._profile.profile.scheduler.partition + counts = tuple(sorted(set(self._launcher.query_gpu_counts(partition=partition)))) + if len(counts) != 1: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "eligible Slurm nodes do not report one GPU count", + ) + return counts[0] + + +class _SystemRunPlanner: + def __init__(self, preparer: _RunPreparer) -> None: + self._preparer = preparer + + def plan(self, config: DataDesignerSlurmConfig, *, source_root: Path) -> ResolvedSlurmRunPlan: + with self._preparer.prepare( + config, source_root=source_root, operation=SlurmServiceOperation.PLAN_RUN + ) as prepared: + return prepared.plan + + +class _SystemRunBackend: + def __init__( + self, + preparer: _RunPreparer, + selected_profile: SelectedSlurmProfile, + launcher: SlurmCommandClient, + publisher: SlurmRunArtifactPublisher, + clock: Clock, + source_environment: Mapping[str, str], + progress: Callable[[str], None] | None = None, + ) -> None: + self._preparer = preparer + self._profile = selected_profile + self._launcher = launcher + self._publisher = publisher + self._clock = clock + self._source_environment = source_environment + self._progress = progress + self._retry_collection = RunRetryCollectionBackend( + selected_profile.profile.workspace_root, + launcher, + clock, + source_environment, + ) + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + ) -> SlurmRunExecution: + with self._preparer.prepare( + config, + source_root=source_root, + operation=SlurmServiceOperation.EXECUTE_RUN, + ) as prepared: + plan = prepared.plan + if dry_run: + return SlurmRunExecution( + run_id=plan.run_id, + state="dry_run", + plan_sha256=plan.compute_sha256(), + shard_count=len(plan.shards), + batch_script=prepared.batch_script, + ) + self._materialize_default_managed_assets(config, plan) + export_environment = self._build_export_environment(config) + self._initialize_run(config, prepared) + self._report_progress("Rendered batch script; invoking sbatch --parsable --hold (script via stdin).") + try: + receipt = self._launcher.submit_script( + prepared.batch_script, + hold=True, + export_environment=export_environment, + ) + except SlurmSubmissionError as error: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + _safe_submission_error(error), + ) from None + except SlurmLauncherError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "Slurm submission is unavailable", + ) from None + try: + self._record_submission(plan, receipt.job_id) + except Exception as error: + try: + self._launcher.cancel(receipt.job_id) + except Exception: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", + ) from error + try: + self._record_submission_failure(plan, receipt.job_id) + except SlurmServiceError as state_error: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + f"Slurm job {receipt.job_id} was cancelled but its partial submission state could not be updated", + ) from state_error + raise + except BaseException: + try: + self._launcher.cancel(receipt.job_id) + except Exception: + pass + else: + try: + self._record_submission_failure(plan, receipt.job_id) + except BaseException: + pass + raise + try: + self._launcher.release(receipt.job_id) + except SlurmLauncherError as error: + try: + self._launcher.cancel(receipt.job_id) + except Exception: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} could not be released or cancelled", + ) from error + try: + self._record_submission_failure(plan, receipt.job_id) + except SlurmServiceError as state_error: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} was cancelled but run {plan.run_id!r} could not be updated", + ) from state_error + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + f"held Slurm job {receipt.job_id} could not be released and was cancelled", + ) from None + run_root = Path(plan.authored_config.path).parent + log_stem = f"slurm-attempt-0001-{receipt.job_id}_0" + status_command = ["data-designer", "slurm", "status", plan.run_id] + if plan.selected_profile.catalog_path is not None: + status_command.extend(("--profile-file", plan.selected_profile.catalog_path)) + if plan.selected_profile.cluster_name is not None: + status_command.extend(("--cluster", plan.selected_profile.cluster_name)) + self._report_progress(f"Submitted. Check status: {shlex.join(status_command)}") + self._report_progress( + "After allocation, follow shard 0 Slurm logs: " + + shlex.join( + ( + "tail", + "-F", + (run_root / f"{log_stem}.out").as_posix(), + (run_root / f"{log_stem}.err").as_posix(), + ) + ) + ) + server_log_root = run_root / "shards" / "shard-00000" / "attempts" / "attempt-0001" / "logs" + server_log_command = ( + f"find {shlex.quote(server_log_root.as_posix())} -type f " + "\\( -name 'deployment-*.out' -o -name 'deployment-*.err' \\) " + "-exec tail -n 10 {} +" + ) + self._report_progress( + "Once serving starts, watch shard 0 vLLM logs (including new executions): " + + shlex.join(("watch", "-n", "2", "-t", server_log_command)) + ) + if len(plan.shards) > 1: + self._report_progress( + "For another shard, replace _0 in the Slurm filenames with its numeric task index, and " + "shard-00000 in the vLLM path with its shard ID (for example, _1 and shard-00001)." + ) + return SlurmRunExecution( + run_id=plan.run_id, + state="submitted", + plan_sha256=plan.compute_sha256(), + shard_count=len(plan.shards), + job_id=receipt.job_id, + ) + + def _report_progress(self, message: str) -> None: + if self._progress is not None: + try: + self._progress(message) + except Exception: + # Display is optional and must not affect submission or reconciliation. + pass + + @staticmethod + def _materialize_default_managed_assets( + config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + ) -> None: + if config.invocation.input_bindings.managed_assets_path is not None: + return + path = plan.invocation.effective_input_bindings.managed_assets_path + assert path is not None + try: + ensure_private_directory(Path(path)) + except (ClientWorkerError, OSError): + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "managed assets workspace cannot be prepared", + ) from None + + def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str, str]: + environment = {"SLURM_EXPORT_ENV": "ALL"} + deployment_type = self._source_environment.get("NEMO_DEPLOYMENT_TYPE", "slurm") + if type(deployment_type) is not str or "\0" in deployment_type: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + "NEMO_DEPLOYMENT_TYPE is invalid", + ) + environment["NEMO_DEPLOYMENT_TYPE"] = deployment_type + if "SLURM_CONF" in self._source_environment: + slurm_conf = self._source_environment["SLURM_CONF"] + if type(slurm_conf) is not str or "\0" in slurm_conf: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + "SLURM_CONF is invalid", + ) + environment["SLURM_CONF"] = slurm_conf + for name in collect_secret_environment_names(config): + try: + value = self._source_environment[name] + except KeyError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + f"required secret environment {name!r} is unavailable", + ) from None + if type(value) is not str or "\0" in value: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.EXECUTE_RUN, + f"required secret environment {name!r} is invalid", + ) + environment[name] = value + return environment + + def _initialize_run( + self, + config: DataDesignerSlurmConfig, + prepared: _PreparedRun, + ) -> None: + try: + self._publisher.initialize( + config, + prepared.plan, + prepared.dependencies, + prepared.builder_payload, + force=False, + ) + except StateNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + SlurmServiceOperation.EXECUTE_RUN, + "required run state was not found", + ) from None + except StateConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.EXECUTE_RUN, + "run state already contains different inputs", + ) from None + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + "run state cannot be initialized", + ) from None + + def _record_submission( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + ) -> None: + try: + self._publisher.record_submission(plan, job_id, submitted_at=self._clock()) + except StateNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + SlurmServiceOperation.EXECUTE_RUN, + "submitted run state was not found", + ) from None + except StateConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.EXECUTE_RUN, + "submitted run state conflicts with persisted state", + ) from None + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + "submission state cannot be recorded", + ) from None + + def _record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + ) -> None: + try: + self._publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) + except (StateNotFoundError, StateConflictError, SlurmStateError): + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.EXECUTE_RUN, + "cancelled submission state cannot be recorded", + ) from None + + def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: + operation = SlurmServiceOperation.STATUS_RUN + try: + writer = SlurmStateWriter(self._profile.profile.workspace_root, run_id) + run = writer.load_run() + persisted_shards = writer.load_shards() + self._reconcile_attempts( + writer, + tuple(attempt for shard in persisted_shards for attempt in writer.load_attempts(shard.shard_id)), + ) + shards = [] + for shard in persisted_shards: + writer.resume_incomplete_finalization(shard.shard_id, published_at=self._clock()) + attempts = tuple( + SlurmPersistedAttemptStatus( + attempt=attempt, + readiness=_load_optional( + lambda shard_id=shard.shard_id, attempt_id=attempt.attempt_id: writer.load_readiness( + shard_id, attempt_id + ) + ), + ) + for attempt in writer.load_attempts(shard.shard_id) + ) + shards.append( + SlurmPersistedShardStatus( + shard=shard, + attempts=attempts, + winner=_load_optional(lambda shard_id=shard.shard_id: writer.load_winner(shard_id)), + ) + ) + return SlurmPersistedRunStatus(run=run, shards=tuple(shards)) + except StateNotFoundError: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "run state was not found") from None + except StateConflictError: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, "run state is inconsistent") from None + except SlurmStateError: + raise SlurmServiceError(SlurmServiceErrorCode.INTERNAL, operation, "run state cannot be read") from None + + def _reconcile_attempts( + self, + writer: SlurmStateWriter, + attempts: tuple[AttemptManifest, ...], + ) -> None: + active = tuple( + attempt for attempt in attempts if attempt.state in _ACTIVE_ATTEMPT_STATES and attempt.scheduler is not None + ) + if not active: + return + try: + reconciled = SlurmStateReconciler( + self._profile.profile.workspace_root, + writer.load_run().run_id, + self._launcher, + ).refresh(observed_at=self._clock()) + except SlurmStateError as error: + if isinstance(error.__cause__, SlurmLauncherError): + return + raise + active_identities = {(attempt.shard_id, attempt.attempt_id) for attempt in active} + for shard in reconciled.shards: + for status in shard.attempts: + if (status.attempt.shard_id, status.attempt.attempt_id) in active_identities: + self._update_reconciled_attempt(writer, status) + + @staticmethod + def _update_reconciled_attempt( + writer: SlurmStateWriter, + status: AttemptStatus, + ) -> None: + attempt = status.attempt + scheduler = status.scheduler + if scheduler is None: # pragma: no cover - active attempts always have scheduler evidence + raise AssertionError("active attempt has no reconciled scheduler evidence") + update: dict[str, object] = {"updated_at": scheduler.observed_at} + if status.effective_state is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: + update["state"] = AttemptLifecycleState.PENDING + elif status.effective_state is EffectiveAttemptState.RUNNING and attempt.state in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + }: + update["state"] = AttemptLifecycleState.RUNNING + elif status.effective_state is EffectiveAttemptState.FAILED: + update.update( + state=AttemptLifecycleState.FAILED, + terminal_classification=_FAILURE_CLASSIFICATIONS.get( + scheduler.state, + AttemptTerminalClassification.UNKNOWN, + ), + ) + else: + return + try: + writer.update_attempt(attempt.model_copy(update=update)) + except StateConflictError: + return + + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: + status = self.status(run_id) + job_ids = tuple( + sorted( + { + attempt.attempt.scheduler.array_job_id + for shard in status.shards + for attempt in shard.attempts + if attempt.attempt.state in _ACTIVE_ATTEMPT_STATES and attempt.attempt.scheduler is not None + } + ) + ) + failures = [] + for job_id in job_ids: + try: + self._launcher.cancel(job_id) + except SlurmLauncherError: + failures.append(job_id) + if failures: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.CANCEL_RUN, + f"failed to cancel managed Slurm jobs: {_format_job_ids(failures)}", + ) + return SlurmRunCancellation(run_id=run_id, job_ids=job_ids) + + def retry( + self, + run_or_job_id: Identifier, + *, + shard_ids: tuple[ShardId, ...] | None, + resume: Literal["never", "always", "if_possible"], + dry_run: bool, + ) -> SlurmRetryExecution: + return self._retry_collection.retry( + run_or_job_id, + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + + def collect( + self, + input_path: Identifier | Path, + *, + destination: Path | None, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + return self._retry_collection.collect( + input_path, + destination=destination, + num_partitions=num_partitions, + ) + + +class _RegistryImageBackend: + def __init__(self, workspace_root: str, lifecycle: SlurmImageLifecycleManager) -> None: + self._verified = VerifiedImageRegistry(workspace_root) + self._store = ImageRegistryStore(workspace_root) + self._lifecycle = lifecycle + + def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: + try: + return self._verified.resolve_for_planning(reference, expected_kind=expected_kind) + except ImageNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + SlurmServiceOperation.RESOLVE_IMAGE, + "image is not registered", + ) from None + except SlurmImageError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RESOLVE_IMAGE, + "registered image cannot be verified", + ) from None + + def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: + return self._lifecycle.add(request, replace=replace) + + def list(self) -> tuple[RegisteredImage, ...]: + return self._invoke_registry(SlurmServiceOperation.LIST_IMAGES, self._store.list_images) + + def get(self, name: Identifier) -> RegisteredImage: + return self._invoke_registry(SlurmServiceOperation.GET_IMAGE, lambda: self._store.get_by_name(name)) + + def remove(self, name: Identifier, *, delete_artifact: bool = False) -> RegisteredImage: + return self._invoke_registry( + SlurmServiceOperation.REMOVE_IMAGE, + lambda: self._verified.unregister(name, delete_artifact=delete_artifact), + ) + + @staticmethod + def _invoke_registry(operation: SlurmServiceOperation, call: Callable[[], _ResultT]) -> _ResultT: + try: + return call() + except ImageNotFoundError: + raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "image is not registered") from None + except ImageConflictError as error: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, str(error)) from None + except ImageArtifactRemovalError as error: + raise SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, operation, str(error)) from None + except SlurmImageError: + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, operation, "image registry operation failed" + ) from None + + +def create_slurm_run_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, + run_id_factory: RunIdFactory | None = None, + clock: Clock | None = None, + package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, + progress: Callable[[str], None] | None = None, +) -> 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, + progress=progress, + ) + + +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, + progress: Callable[[str], None] | None = None, +) -> SlurmRunService: + command_client = launcher or SlurmCommandClient() + selected_clock = clock or _utc_now + preparer = _RunPreparer( + selected, + VerifiedImageRegistry(selected.profile.workspace_root), + dependency_resolver or ClientDependencyResolver(), + command_client, + run_id_factory or _new_run_id, + package_version or importlib.metadata.version("data-designer-slurm"), + ) + backend = _SystemRunBackend( + preparer, + selected, + command_client, + artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock), + selected_clock, + dict(os.environ if source_environment is None else source_environment), + progress, + ) + return SlurmRunService(_SystemRunPlanner(preparer), render_generation_attempt_script, backend) + + +def create_slurm_image_service( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + launcher: SlurmCommandClient | None = None, + lifecycle_id_factory: Callable[[], str] | None = None, + clock: Clock | None = None, + sleep: Callable[[float], None] | None = None, + progress: Callable[[str], None] | None = None, + logs: Callable[[str], None] | None = None, +) -> SlurmImageService: + """Create the production image service for one selected cluster profile.""" + selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + command_client = launcher or SlurmCommandClient() + lifecycle = SlurmImageLifecycleManager( + selected, + command_client, + lifecycle_id_factory=lifecycle_id_factory, + clock=clock, + sleep=sleep, + progress=progress, + logs=logs, + ) + backend = _RegistryImageBackend(selected.profile.workspace_root, lifecycle) + 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, + progress: Callable[[str], None] | 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, + progress=progress, + ) + return SlurmBenchmarkService(backend) + + +def _resolve_builder_payload( + authored: DataDesignerSlurmConfig, + source_root: Path, +) -> dict[str, JsonValue] | None: + if authored.builder.source is None: + return None + return load_builder_payload(source_root / authored.builder.source) + + +def _load_optional(call: Callable[[], _ResultT]) -> _ResultT | None: + try: + return call() + except StateNotFoundError: + return None + + +def _new_run_id() -> str: + return f"run-{uuid4().hex}" + + +def _safe_builder_load_error(error: SlurmConfigLoadError) -> str: + """Expose a bounded, redacted builder diagnostic instead of a generic failure.""" + prefix = "run configuration cannot be prepared: " + detail = " ".join(redact_sensitive_diagnostic(str(error)).split()) or "builder source cannot be loaded" + max_detail_length = 512 - len(prefix) + if len(detail) > max_detail_length: + detail = f"{detail[: max_detail_length - 3]}..." + return f"{prefix}{detail}" + + +def _validate_builder_for_submission( + authored: DataDesignerSlurmConfig, + builder_payload: dict[str, JsonValue] | None, + operation: SlurmServiceOperation, +) -> None: + """Reject invalid built-in builder settings before reserving GPUs. + + A plugin supplied only through the client dependency lock and a local seed + path visible only on compute nodes are checked during allocation preflight. + """ + from data_designer.config.config_builder import BuilderConfig + from data_designer.config.utils.io_helpers import VALID_DATASET_FILE_EXTENSIONS + from data_designer.errors import DataDesignerError + + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + normalized = payload if "data_designer" in payload else {"data_designer": payload} + data_designer = normalized.get("data_designer") + columns = data_designer.get("columns") if isinstance(data_designer, dict) else None + processors = data_designer.get("processors") if isinstance(data_designer, dict) else None + seed_config = data_designer.get("seed_config") if isinstance(data_designer, dict) else None + seed_source = seed_config.get("source") if isinstance(seed_config, dict) else None + builtin_columns = { + "custom", + "embedding", + "expression", + "image", + "llm-code", + "llm-judge", + "llm-structured", + "llm-text", + "sampler", + "seed-dataset", + "validation", + } + builtin_processors = {"drop_columns", "schema_transform"} + + def is_unrecognized_type(value: object, discriminator: str, known: set[str]) -> bool: + return ( + isinstance(value, dict) and isinstance(value.get(discriminator), str) and value[discriminator] not in known + ) + + unknown_column = isinstance(columns, list) and any( + is_unrecognized_type(column, "column_type", builtin_columns) for column in columns + ) + unknown_processor = isinstance(processors, list) and any( + is_unrecognized_type(processor, "processor_type", builtin_processors) for processor in processors + ) + unknown_seed_source = ( + isinstance(seed_source, dict) + and isinstance(seed_source.get("seed_type"), str) + and seed_source["seed_type"] not in {"local", "hf", "df", "directory", "file_contents", "agent_rollout"} + ) + deferred_local_seed_path = ( + isinstance(seed_source, dict) + and seed_source.get("seed_type") == "local" + and isinstance(seed_source.get("path"), str) + and seed_source["path"].lower().endswith(tuple(VALID_DATASET_FILE_EXTENSIONS)) + ) + plugin_dependency = bool(authored.client.dependencies.requirements or authored.client.dependencies.lock_file) + if (plugin_dependency and (unknown_column or unknown_processor or unknown_seed_source)) or deferred_local_seed_path: + # Validate the complete built-in envelope while deferring plugin entries + # and local seed path existence until the client reaches its allocation. + data_designer = dict(data_designer) + if plugin_dependency and unknown_column: + data_designer["columns"] = [ + column for column in columns if not is_unrecognized_type(column, "column_type", builtin_columns) + ] + if plugin_dependency and unknown_processor: + data_designer["processors"] = [ + processor + for processor in processors + if not is_unrecognized_type(processor, "processor_type", builtin_processors) + ] + if (plugin_dependency and unknown_seed_source) or deferred_local_seed_path: + data_designer["seed_config"] = { + **seed_config, + "source": {"seed_type": "hf", "path": "deferred-seed-source"}, + } + normalized = {**normalized, "data_designer": data_designer} + try: + BuilderConfig.model_validate(normalized) + except ValidationError as error: + detail = format_validation_error(error, subject="builder configuration", models=BuilderConfig) + raise SlurmServiceError(SlurmServiceErrorCode.INVALID_REQUEST, operation, _bounded_diagnostic(detail)) from None + except DataDesignerError as error: + detail = " ".join(redact_sensitive_diagnostic(str(error)).split()) + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + _bounded_diagnostic(f"builder configuration failed validation: {detail}"), + ) from None + + +def _bounded_diagnostic(message: str) -> str: + return f"{message[:509]}..." if len(message) > 512 else message + + +def _safe_submission_error(error: SlurmSubmissionError) -> str: + """Preserve bounded, redacted sbatch diagnostics and submission ambiguity.""" + prefix = "Slurm submission outcome is uncertain: " if error.may_have_succeeded else "Slurm submission failed: " + suffix = "; inspect queue and accounting before retrying" if error.may_have_succeeded else "" + detail = " ".join(redact_sensitive_diagnostic(str(error)).split()) or "no diagnostic was returned" + max_detail_length = 512 - len(prefix) - len(suffix) + if len(detail) > max_detail_length: + detail = f"{detail[: max_detail_length - 3]}..." + return f"{prefix}{detail}{suffix}" + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def _format_job_ids(job_ids: list[int]) -> str: + visible = ", ".join(str(job_id) for job_id in job_ids[:_MAX_VISIBLE_JOB_IDS]) + remaining = len(job_ids) - _MAX_VISIBLE_JOB_IDS + return visible if remaining <= 0 else f"{visible}, and {remaining} more" + + +__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/serving/deployment.py b/packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py new file mode 100644 index 000000000..673405d4e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Self-validating resolved server deployment record.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.images import ServingImageInspection +from data_designer.slurm.contracts import ( + ContractValue, + validate_absolute_path, + validate_plain_text, +) +from data_designer.slurm.planning.models import ResolvedImage, ResolvedTopology +from data_designer.slurm.serving.endpoints import ( + ResolvedBackendEndpoint, + ResolvedLogicalEndpoint, + ResolvedReadinessProbe, +) +from data_designer.slurm.serving.vllm import ResolvedVllmLaunchPolicy, ResolvedVllmProcess +from data_designer.slurm.types import Identifier + + +class ResolvedVllmServerDeployment(ContractValue): + """Complete resolved vLLM serving specification for one deployment.""" + + deployment_id: Identifier + server_type: Literal["vllm"] + model_alias: str + model: str + served_model_name: str + image: ResolvedImage + executable_path: str + node_indices: tuple[NonNegativeInt, ...] = Field(min_length=1) + gpus_per_node: PositiveInt + gpu_start_index: NonNegativeInt = Field(default=0, exclude_if=lambda value: value == 0) + topology: ResolvedTopology + launch_policy: ResolvedVllmLaunchPolicy + processes: tuple[ResolvedVllmProcess, ...] = Field(min_length=1) + readiness_probes: tuple[ResolvedReadinessProbe, ...] = Field(min_length=1) + backend_endpoints: tuple[ResolvedBackendEndpoint, ...] = Field(min_length=1) + logical_endpoint: ResolvedLogicalEndpoint + failure_policy: Literal["coordinated"] = "coordinated" + + _executable_path_is_absolute = field_validator("executable_path")(validate_absolute_path) + + @field_validator("model", "served_model_name") + @classmethod + def validate_model_names(cls, value: str) -> str: + return validate_plain_text(value, field_name="model name") + + @model_validator(mode="after") + def validate_resolution(self) -> ResolvedVllmServerDeployment: + inspection = self.image.inspection_facts + if not isinstance(inspection, ServingImageInspection) or inspection.server_type != self.server_type: + raise ValueError("resolved server image inspection must match the server type") + if inspection.executable_path != self.executable_path: + raise ValueError("resolved executable path must match the serving image inspection") + if self.node_indices != tuple(sorted(set(self.node_indices))): + raise ValueError("resolved server nodes must be sorted and unique") + expected_topology = ResolvedTopology.derive( + node_count=len(self.node_indices), + gpus_per_node=self.gpus_per_node, + tensor_parallel=self.topology.tensor_parallel, + nodes_per_replica=self.topology.nodes_per_replica, + ) + if self.topology != expected_topology: + raise ValueError("resolved server topology must match its node and GPU resources") + if self.launch_policy.enable_expert_parallel and self.topology.pipeline_parallel > 1: + raise ValueError("multi-node expert parallel is not supported in v1") + if self.logical_endpoint.model_alias != self.model_alias: + raise ValueError("logical endpoint model alias must match the deployment") + if self.logical_endpoint.served_model_name != self.served_model_name: + raise ValueError("logical endpoint served model name must match the deployment") + if self.logical_endpoint.endpoint_id != f"{self.deployment_id}-logical-endpoint": + raise ValueError("logical endpoint ID must match the deployment") + + deployment_replica_indices = tuple(endpoint.deployment_replica_index for endpoint in self.backend_endpoints) + expected_replicas = tuple(range(self.topology.replica_count)) + if deployment_replica_indices != expected_replicas: + raise ValueError("backend endpoints must use complete ordered replica identities") + backend_ids = tuple(endpoint.backend_id for endpoint in self.backend_endpoints) + expected_backend_ids = tuple( + f"{self.deployment_id}-backend-{deployment_replica_index:05d}" + for deployment_replica_index in expected_replicas + ) + if backend_ids != expected_backend_ids: + raise ValueError("backend endpoint IDs must match the deployment and replica order") + if self.logical_endpoint.backend_ids != backend_ids: + raise ValueError("logical endpoint backends must match the resolved backend order") + if len({process.process_id for process in self.processes}) != len(self.processes): + raise ValueError("resolved process IDs must be unique") + + expected_process_count = self.topology.replica_count * self.topology.pipeline_parallel + if len(self.processes) != expected_process_count: + raise ValueError("resolved process count must match replica and pipeline topology") + process_identities = tuple( + (process.deployment_replica_index, process.pipeline_rank) for process in self.processes + ) + expected_process_identities = tuple( + (deployment_replica_index, pipeline_rank) + for deployment_replica_index in range(self.topology.replica_count) + for pipeline_rank in range(self.topology.pipeline_parallel) + ) + if process_identities != expected_process_identities: + raise ValueError("resolved processes must use complete ordered replica and pipeline identities") + + probes_by_backend = {probe.backend_id: probe for probe in self.readiness_probes} + if len(self.readiness_probes) != len(self.backend_endpoints) or tuple(probes_by_backend) != backend_ids: + raise ValueError("readiness probes must match the resolved backend order") + for endpoint in self.backend_endpoints: + self._validate_replica(endpoint, probes_by_backend[endpoint.backend_id]) + self._validate_network_addresses() + return self + + def _validate_replica(self, endpoint: ResolvedBackendEndpoint, probe: ResolvedReadinessProbe) -> None: + expected_group = endpoint.deployment_replica_index // self.topology.replicas_per_node_group + expected_replica_index_in_node_group = endpoint.deployment_replica_index % self.topology.replicas_per_node_group + expected_head_node = self.node_indices[expected_group * self.topology.nodes_per_replica] + if ( + endpoint.node_group_index != expected_group + or endpoint.replica_index_in_node_group != expected_replica_index_in_node_group + or endpoint.node_index != expected_head_node + or endpoint.served_model_name != self.served_model_name + ): + raise ValueError("backend endpoint must match its resolved replica placement") + if ( + (probe.node_index, probe.port) != (endpoint.node_index, endpoint.port) + or probe.probe_id != f"{self.deployment_id}-readiness-{endpoint.deployment_replica_index:05d}" + or probe.path != self.launch_policy.readiness_path + or probe.deadline_seconds != self.launch_policy.startup_timeout_seconds + ): + raise ValueError("readiness probes must target their resolved backend endpoint") + replica_processes = tuple( + process + for process in self.processes + if process.deployment_replica_index == endpoint.deployment_replica_index + ) + expected_gpu_start = self.gpu_start_index + expected_replica_index_in_node_group * self.topology.tensor_parallel + expected_gpus = tuple(range(expected_gpu_start, expected_gpu_start + self.topology.tensor_parallel)) + expected_delay = ( + 0 + if endpoint.deployment_replica_index == 0 + else self.launch_policy.lead_boot_standoff_seconds + + endpoint.deployment_replica_index * self.launch_policy.rank_launch_stagger_seconds + ) + for process in replica_processes: + expected_node = self.node_indices[expected_group * self.topology.nodes_per_replica + process.pipeline_rank] + expected_process_id = ( + f"{self.deployment_id}-replica-{endpoint.deployment_replica_index:05d}-rank-{process.pipeline_rank:05d}" + ) + if ( + process.process_id != expected_process_id + or process.deployment_id != self.deployment_id + or process.node_group_index != expected_group + or process.replica_index_in_node_group != expected_replica_index_in_node_group + or process.node_index != expected_node + or process.gpu_indices != expected_gpus + or process.tensor_parallel != self.topology.tensor_parallel + or process.pipeline_parallel != self.topology.pipeline_parallel + or process.launch_delay_seconds != expected_delay + ): + raise ValueError("resolved process must match its deployment topology and launch policy") + head = replica_processes[0] + if (head.node_index, head.http_port) != (endpoint.node_index, endpoint.port): + raise ValueError("replica head process must publish its resolved backend endpoint") + if self.topology.pipeline_parallel > 1: + rendezvous = head.rendezvous + if rendezvous is None or any(process.rendezvous != rendezvous for process in replica_processes): + raise ValueError("replica processes must share one rendezvous specification") + if ( + rendezvous.master_node_index != expected_head_node + or rendezvous.timeout_seconds != self.launch_policy.distributed_init_timeout_seconds + ): + raise ValueError("replica rendezvous must match its group head and distributed timeout") + + def _validate_network_addresses(self) -> None: + rendezvous_addresses = tuple( + (process.rendezvous.master_node_index, process.rendezvous.port) + for process in self.processes + if process.pipeline_rank == 0 and process.rendezvous is not None + ) + addresses = ( + tuple((endpoint.node_index, endpoint.port) for endpoint in self.backend_endpoints) + + ((self.logical_endpoint.node_index, self.logical_endpoint.port),) + + rendezvous_addresses + ) + if len(addresses) != len(set(addresses)): + raise ValueError("resolved server network addresses must be unique") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/serving/endpoints.py b/packages/data-designer-slurm/src/data_designer/slurm/serving/endpoints.py new file mode 100644 index 000000000..747804d02 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/endpoints.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Backend, readiness, and logical endpoint records for Slurm serving.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, model_validator + +from data_designer.slurm.contracts import ContractValue +from data_designer.slurm.types import Identifier, NetworkPort + + +class ResolvedBackendEndpoint(ContractValue): + """Internal HTTP endpoint published by one replica's head process.""" + + backend_id: Identifier + deployment_replica_index: NonNegativeInt + node_group_index: NonNegativeInt + replica_index_in_node_group: NonNegativeInt + node_index: NonNegativeInt + port: NetworkPort + served_model_name: str + + +class ResolvedReadinessProbe(ContractValue): + """Bounded readiness probe for one replica-head backend.""" + + probe_id: Identifier + backend_id: Identifier + node_index: NonNegativeInt + port: NetworkPort + path: str + deadline_seconds: PositiveInt + expected_status_code: Literal[200] = 200 + + +class ResolvedLogicalEndpoint(ContractValue): + """Client-visible endpoint aggregating one deployment's healthy backends.""" + + endpoint_id: Identifier + model_alias: str + served_model_name: str + node_index: NonNegativeInt + port: NetworkPort + backend_ids: tuple[Identifier, ...] = Field(min_length=1) + load_balancing: Literal["least_connections"] = "least_connections" + retry_status_codes: tuple[Literal[429], ...] = (429,) + preserve_final_overload_response: Literal[True] = True + require_all_backends_ready: Literal[True] = True + + @model_validator(mode="after") + def validate_backend_ids(self) -> ResolvedLogicalEndpoint: + if len(self.backend_ids) != len(set(self.backend_ids)): + raise ValueError("logical endpoint backend IDs must be unique") + if self.retry_status_codes != (429,): + raise ValueError("logical endpoint must retry HTTP 429 responses") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py new file mode 100644 index 000000000..95f038a37 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure deployment-owned server resolution for Slurm execution.""" + +from __future__ import annotations + +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion, Version +from pydantic import ValidationError + +from data_designer.slurm.config.images import ServingImageInspection +from data_designer.slurm.config.utils import convert_duration_to_seconds +from data_designer.slurm.config.vllm import VllmServerConfig +from data_designer.slurm.planning.models import PortClaim, ResolvedDeployment, ResolvedSlurmRunPlan +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.endpoints import ( + ResolvedBackendEndpoint, + ResolvedLogicalEndpoint, + ResolvedReadinessProbe, +) +from data_designer.slurm.serving.vllm import ( + ResolvedVllmLaunchPolicy, + ResolvedVllmProcess, + ResolvedVllmRendezvous, + VllmProcessRole, +) +from data_designer.slurm.types import Identifier + +_SUPPORTED_VLLM_RUNTIME_RELEASES = ( + SpecifierSet("~=0.21.0"), + SpecifierSet("~=0.22.0"), +) + + +class VllmServerResolutionError(ValueError): + """Raised when planner inputs cannot produce a supported vLLM server specification.""" + + +def resolve_vllm_server( + plan: ResolvedSlurmRunPlan, + deployment_id: Identifier, +) -> ResolvedVllmServerDeployment: + """Resolve one vLLM deployment admitted by the tested runtime matrix.""" + try: + return _resolve_vllm_server(plan, deployment_id) + except ValidationError as error: + raise VllmServerResolutionError("planner inputs produced an inconsistent vLLM server specification") from error + + +def _resolve_vllm_server( + plan: ResolvedSlurmRunPlan, + deployment_id: Identifier, +) -> ResolvedVllmServerDeployment: + deployments = tuple(deployment for deployment in plan.deployments if deployment.deployment_id == deployment_id) + if len(deployments) != 1: + raise VllmServerResolutionError("resolved plan must contain exactly one deployment with the requested ID") + resolved_deployment = deployments[0] + expected_endpoint_id = f"{resolved_deployment.deployment_id}-logical-endpoint" + logical_endpoint_ports = tuple(port for port in plan.client.ports if port.name == expected_endpoint_id) + if len(logical_endpoint_ports) != 1: + raise VllmServerResolutionError("resolved client must contain exactly one logical endpoint for the deployment") + logical_endpoint_port = logical_endpoint_ports[0] + if logical_endpoint_port.name != expected_endpoint_id or logical_endpoint_port.role != "logical_endpoint": + raise VllmServerResolutionError("logical endpoint port must match the resolved deployment") + if logical_endpoint_port.node_index != plan.client.host_node_index: + raise VllmServerResolutionError("logical endpoint port must use the resolved client host") + if (logical_endpoint_port.node_index, logical_endpoint_port.port) in { + (port.node_index, port.port) for port in resolved_deployment.ports + }: + raise VllmServerResolutionError("logical endpoint port must not collide with a deployment port") + authored_deployment = resolved_deployment.authored + server = authored_deployment.server + inspection = resolved_deployment.image.inspection_facts + if not isinstance(inspection, ServingImageInspection) or inspection.server_type != server.type: + raise VllmServerResolutionError("resolved serving image does not match the declared vLLM server") + _validate_vllm_runtime_version(inspection.runtime_version) + if server.enable_expert_parallel and resolved_deployment.topology.pipeline_parallel > 1: + raise VllmServerResolutionError("multi-node expert parallel is not supported in v1") + + http_ports = tuple(port for port in resolved_deployment.ports if port.role == "http") + rendezvous_ports = tuple(port for port in resolved_deployment.ports if port.role == "rendezvous") + processes: list[ResolvedVllmProcess] = [] + backends: list[ResolvedBackendEndpoint] = [] + probes: list[ResolvedReadinessProbe] = [] + nodes_per_replica = resolved_deployment.topology.nodes_per_replica + replicas_per_node_group = resolved_deployment.topology.replicas_per_node_group + tensor_parallel = resolved_deployment.topology.tensor_parallel + + for node_group_index in range(resolved_deployment.topology.node_group_count): + group_start = node_group_index * nodes_per_replica + group_nodes = resolved_deployment.node_indices[group_start : group_start + nodes_per_replica] + for replica_index_in_node_group in range(replicas_per_node_group): + deployment_replica_index = node_group_index * replicas_per_node_group + replica_index_in_node_group + http_port = http_ports[deployment_replica_index] + backend_id = f"{resolved_deployment.deployment_id}-backend-{deployment_replica_index:05d}" + backends.append( + ResolvedBackendEndpoint( + backend_id=backend_id, + deployment_replica_index=deployment_replica_index, + node_group_index=node_group_index, + replica_index_in_node_group=replica_index_in_node_group, + node_index=http_port.node_index, + port=http_port.port, + served_model_name=resolved_deployment.served_model_name, + ) + ) + probes.append( + ResolvedReadinessProbe( + probe_id=f"{resolved_deployment.deployment_id}-readiness-{deployment_replica_index:05d}", + backend_id=backend_id, + node_index=http_port.node_index, + port=http_port.port, + path=server.readiness_path, + deadline_seconds=convert_duration_to_seconds(server.startup_timeout), + ) + ) + rendezvous = _resolve_rendezvous( + resolved_deployment, + rendezvous_ports, + node_group_index=node_group_index, + replica_index_in_node_group=replica_index_in_node_group, + deployment_replica_index=deployment_replica_index, + ) + gpu_start = resolved_deployment.gpu_start_index + replica_index_in_node_group * tensor_parallel + gpu_indices = tuple(range(gpu_start, gpu_start + tensor_parallel)) + launch_delay = _calculate_launch_delay_seconds(server, deployment_replica_index) + for pipeline_rank, node_index in enumerate(group_nodes): + is_head = pipeline_rank == 0 + processes.append( + ResolvedVllmProcess( + process_id=( + f"{resolved_deployment.deployment_id}-replica-{deployment_replica_index:05d}" + f"-rank-{pipeline_rank:05d}" + ), + deployment_id=resolved_deployment.deployment_id, + deployment_replica_index=deployment_replica_index, + node_group_index=node_group_index, + replica_index_in_node_group=replica_index_in_node_group, + pipeline_rank=pipeline_rank, + node_index=node_index, + gpu_indices=gpu_indices, + role=VllmProcessRole.API_SERVER if is_head else VllmProcessRole.FOLLOWER, + tensor_parallel=tensor_parallel, + pipeline_parallel=resolved_deployment.topology.pipeline_parallel, + http_port=http_port.port if is_head else None, + rendezvous=rendezvous, + launch_delay_seconds=launch_delay, + ) + ) + + return ResolvedVllmServerDeployment( + deployment_id=resolved_deployment.deployment_id, + server_type="vllm", + model_alias=authored_deployment.model_alias, + model=authored_deployment.model, + served_model_name=resolved_deployment.served_model_name, + image=resolved_deployment.image, + executable_path=inspection.executable_path, + node_indices=resolved_deployment.node_indices, + gpus_per_node=resolved_deployment.gpus_per_node, + gpu_start_index=resolved_deployment.gpu_start_index, + topology=resolved_deployment.topology, + launch_policy=ResolvedVllmLaunchPolicy( + startup_timeout_seconds=convert_duration_to_seconds(server.startup_timeout), + distributed_init_timeout_seconds=convert_duration_to_seconds(server.distributed_init_timeout), + lead_boot_standoff_seconds=convert_duration_to_seconds(server.lead_boot_standoff), + rank_launch_stagger_seconds=convert_duration_to_seconds(server.rank_launch_stagger), + readiness_path=server.readiness_path, + enable_expert_parallel=server.enable_expert_parallel, + io_processor_plugin=server.io_processor_plugin, + queue_backpressure=server.queue_backpressure, + extra_args=tuple(server.extra_args), + environment=dict(server.environment), + ), + processes=tuple(processes), + readiness_probes=tuple(probes), + backend_endpoints=tuple(backends), + logical_endpoint=ResolvedLogicalEndpoint( + endpoint_id=logical_endpoint_port.name, + model_alias=authored_deployment.model_alias, + served_model_name=resolved_deployment.served_model_name, + node_index=logical_endpoint_port.node_index, + port=logical_endpoint_port.port, + backend_ids=tuple(backend.backend_id for backend in backends), + ), + ) + + +def _resolve_rendezvous( + resolved_deployment: ResolvedDeployment, + rendezvous_ports: tuple[PortClaim, ...], + *, + node_group_index: int, + replica_index_in_node_group: int, + deployment_replica_index: int, +) -> ResolvedVllmRendezvous | None: + if resolved_deployment.topology.pipeline_parallel == 1: + return None + port = rendezvous_ports[deployment_replica_index] + return ResolvedVllmRendezvous( + node_group_index=node_group_index, + replica_index_in_node_group=replica_index_in_node_group, + master_node_index=port.node_index, + port=port.port, + timeout_seconds=convert_duration_to_seconds(resolved_deployment.authored.server.distributed_init_timeout), + ) + + +def _calculate_launch_delay_seconds(server: VllmServerConfig, deployment_replica_index: int) -> int: + if deployment_replica_index == 0: + return 0 + return convert_duration_to_seconds( + server.lead_boot_standoff + ) + deployment_replica_index * convert_duration_to_seconds(server.rank_launch_stagger) + + +def _validate_vllm_runtime_version(runtime_version: str) -> None: + try: + version = Version(runtime_version) + except InvalidVersion as error: + raise VllmServerResolutionError(f"unsupported vLLM runtime version {runtime_version!r}") from error + if not any(version in supported_release for supported_release in _SUPPORTED_VLLM_RUNTIME_RELEASES): + raise VllmServerResolutionError( + f"unsupported vLLM runtime version {runtime_version!r}; supported release lines are 0.21.x and 0.22.x" + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py b/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py new file mode 100644 index 000000000..cb2bc52ea --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolved vLLM launch policy and per-process records for Slurm serving.""" + +from __future__ import annotations + +from enum import Enum + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.environment import EnvironmentBinding +from data_designer.slurm.config.vllm import QueueBackpressureConfig +from data_designer.slurm.config.vllm_validation import ( + validate_vllm_environment_bindings, + validate_vllm_extra_args, + validate_vllm_readiness_path, +) +from data_designer.slurm.contracts import ContractValue +from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort + + +class VllmProcessRole(str, Enum): + """Role of one vLLM process within a resolved replica.""" + + API_SERVER = "api_server" + FOLLOWER = "follower" + + +class ResolvedVllmLaunchPolicy(ContractValue): + """Resolved deployment-wide lifecycle, admission, and launch inputs.""" + + startup_timeout_seconds: PositiveInt + distributed_init_timeout_seconds: PositiveInt + lead_boot_standoff_seconds: NonNegativeInt + rank_launch_stagger_seconds: NonNegativeInt + readiness_path: str + enable_expert_parallel: bool + io_processor_plugin: Identifier | None = Field(default=None, exclude_if=lambda value: value is None) + queue_backpressure: QueueBackpressureConfig + extra_args: tuple[str, ...] = () + environment: dict[EnvironmentName, EnvironmentBinding] = Field(default_factory=dict) + + _readiness_path_is_safe = field_validator("readiness_path")(validate_vllm_readiness_path) + + @field_validator("extra_args") + @classmethod + def validate_extra_args(cls, values: tuple[str, ...]) -> tuple[str, ...]: + validate_vllm_extra_args(values) + return values + + @field_validator("environment") + @classmethod + def validate_environment( + cls, values: dict[EnvironmentName, EnvironmentBinding] + ) -> dict[EnvironmentName, EnvironmentBinding]: + validate_vllm_environment_bindings(values) + return values + + @model_validator(mode="after") + def validate_timeouts(self) -> ResolvedVllmLaunchPolicy: + if self.distributed_init_timeout_seconds > self.startup_timeout_seconds: + raise ValueError("distributed initialization timeout must not exceed startup timeout") + return self + + +class ResolvedVllmRendezvous(ContractValue): + """Planner-owned rendezvous inputs shared by one multi-node replica.""" + + node_group_index: NonNegativeInt + replica_index_in_node_group: NonNegativeInt + master_node_index: NonNegativeInt + port: NetworkPort + timeout_seconds: PositiveInt + + +class ResolvedVllmProcess(ContractValue): + """Typed, shell-free launch specification for one vLLM process.""" + + process_id: Identifier + deployment_id: Identifier + deployment_replica_index: NonNegativeInt + node_group_index: NonNegativeInt + replica_index_in_node_group: NonNegativeInt + pipeline_rank: NonNegativeInt + node_index: NonNegativeInt + gpu_indices: tuple[NonNegativeInt, ...] = Field(min_length=1) + role: VllmProcessRole + tensor_parallel: PositiveInt + pipeline_parallel: PositiveInt + http_port: NetworkPort | None = None + rendezvous: ResolvedVllmRendezvous | None = None + launch_delay_seconds: NonNegativeInt + + @model_validator(mode="after") + def validate_process(self) -> ResolvedVllmProcess: + if self.gpu_indices != tuple(sorted(set(self.gpu_indices))): + raise ValueError("process GPU indices must be sorted and unique") + if len(self.gpu_indices) != self.tensor_parallel: + raise ValueError("process GPU count must equal tensor parallelism") + if self.pipeline_rank >= self.pipeline_parallel: + raise ValueError("process pipeline rank must be below pipeline parallelism") + if self.pipeline_rank == 0: + if self.role is not VllmProcessRole.API_SERVER: + raise ValueError("pipeline rank zero must be the API server") + if self.http_port is None: + raise ValueError("API server processes require an HTTP port") + elif self.role is not VllmProcessRole.FOLLOWER: + raise ValueError("nonzero pipeline ranks must be followers") + elif self.http_port is not None: + raise ValueError("follower processes must not publish an HTTP port") + if self.pipeline_parallel == 1: + if self.pipeline_rank != 0 or self.rendezvous is not None: + raise ValueError("single-node processes must not carry rendezvous inputs") + elif self.rendezvous is None: + raise ValueError("multi-node processes require rendezvous inputs") + elif ( + self.rendezvous.node_group_index != self.node_group_index + or self.rendezvous.replica_index_in_node_group != self.replica_index_in_node_group + ): + raise ValueError( + "process rendezvous identity must match its node group and replica index within the node group" + ) + return self 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 new file mode 100644 index 000000000..15985b2a3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public state contracts for the optional Slurm execution package.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + ContractValue, + Identifier, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, +) +from data_designer.slurm.state.artifacts import compute_candidate_schema_digest +from data_designer.slurm.state.base import ( + SchedulerIdentity, + SchedulerJobIdentity, + StateRecord, + StateValue, +) +from data_designer.slurm.state.collection_records import ( + CollectedOutputFile, + CollectionResult, + CollectionState, + CollectionStatus, +) +from data_designer.slurm.state.errors import ( + SchedulerStateConflictError, + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.observation import ( + SchedulerAccountingRecord, + SchedulerObservationClient, + SchedulerObservationCollector, + SchedulerQueueRecord, +) +from data_designer.slurm.state.outputs import ( + CANDIDATE_OUTPUT_FORMAT, + MAXIMUM_CANDIDATE_OUTPUT_FILES, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + CollectionPlan, + CollectionShard, + RetryPlan, + RetryShard, + ShardWinner, +) +from data_designer.slurm.state.readiness import ( + AttemptReadiness, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + ReasonCode, +) +from data_designer.slurm.state.reconciliation import ( + reconcile_attempt_observation, + validate_readiness_transition, +) +from data_designer.slurm.state.retry_records import RetryState, RetryStatus +from data_designer.slurm.state.scheduler import ( + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, +) +from data_designer.slurm.state.status import ( + AttemptStatus, + EffectiveRunState, + GenerationState, + RunStatus, + ShardStatus, +) +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_manifest, + validate_attempt_set, + validate_attempt_transition, + validate_collection_plan, + validate_scheduler_observation_transition, + validate_shard_attempt_set, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +if TYPE_CHECKING: + from data_designer.slurm.state.collection import SlurmCollectionCoordinator # noqa: F401 + from data_designer.slurm.state.observer import SlurmStateReconciler # noqa: F401 + from data_designer.slurm.state.retry import SlurmRetryCoordinator # noqa: F401 + from data_designer.slurm.state.store import SlurmStateWriter # noqa: F401 + +_LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "SlurmCollectionCoordinator": ("data_designer.slurm.state.collection", "SlurmCollectionCoordinator"), + "SlurmStateReconciler": ("data_designer.slurm.state.observer", "SlurmStateReconciler"), + "SlurmRetryCoordinator": ("data_designer.slurm.state.retry", "SlurmRetryCoordinator"), + "SlurmStateWriter": ("data_designer.slurm.state.store", "SlurmStateWriter"), +} + +__all__ = [ + "ArtifactReference", + "AttemptLifecycleState", + "AttemptManifest", + "AttemptId", + "AttemptReadiness", + "AttemptStatus", + "AttemptTerminalClassification", + "CandidateOutcome", + "CANDIDATE_OUTPUT_FORMAT", + "MAXIMUM_CANDIDATE_OUTPUT_FILES", + "CandidateOutputFile", + "CandidateOutputManifest", + "CollectedOutputFile", + "compute_candidate_schema_digest", + "CollectionPlan", + "CollectionResult", + "CollectionShard", + "CollectionState", + "CollectionStatus", + "ContractRecord", + "ContractValue", + "DeploymentReadiness", + "EffectiveAttemptState", + "EffectiveRunState", + "EndpointPublicationState", + "Identifier", + "GenerationState", + "ProbeEvidence", + "ProbeOutcome", + "ReadinessState", + "ReasonCode", + "RecordRange", + "RetryPlan", + "RetryShard", + "RetryState", + "RetryStatus", + "RunManifest", + "RunStatus", + "ResumeWorkspace", + "SchedulerIdentity", + "SchedulerJobIdentity", + "SchedulerAccountingRecord", + "SchedulerObservationClient", + "SchedulerObservationCollector", + "SchedulerQueueRecord", + "SchedulerObservation", + "SchedulerState", + "SchedulerStateConflictError", + "Sha256Digest", + "ShardManifest", + "ShardStatus", + "ShardId", + "ShardWinner", + "SlurmStateError", + "SlurmCollectionCoordinator", + "SlurmRetryCoordinator", + "SlurmStateReconciler", + "SlurmStateWriter", + "StateConflictError", + "StateContractError", + "StateCorruptionError", + "StateNotFoundError", + "StateRecord", + "StateValue", + "reconcile_attempt_observation", + "validate_attempt_manifest", + "validate_attempt_set", + "validate_attempt_transition", + "validate_collection_plan", + "validate_readiness_transition", + "validate_scheduler_observation_transition", + "validate_shard_attempt_set", + "validate_shard_manifest", + "validate_shard_set", + "validate_shard_winner", +] + + +def __getattr__(name: str) -> object: + """Lazily import state persistence exports when accessed.""" + if name in _LAZY_IMPORTS: + module_path, attribute_name = _LAZY_IMPORTS[name] + attribute = getattr(importlib.import_module(module_path), attribute_name) + globals()[name] = attribute + return attribute + raise AttributeError(f"module 'data_designer.slurm.state' has no attribute {name!r}") + + +def __dir__() -> list[str]: + """Return the public state exports for interactive discovery.""" + return __all__ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py new file mode 100644 index 000000000..c68bb8738 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound validation of candidate dataset artifacts.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import BinaryIO, Callable, Protocol + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.state.filesystem import ( + open_verified_child_directory, + open_verified_directory, + open_verified_regular_file, +) +from data_designer.slurm.state.outputs import CandidateOutputFile, CandidateOutputManifest + + +class SerializedCandidateSchema(Protocol): + """Serialized schema bytes returned by an Arrow-compatible schema.""" + + def to_pybytes(self) -> bytes: + """Return the canonical serialized schema bytes.""" + ... + + +class CandidateSchema(Protocol): + """Structural schema interface used by candidate producers and readers.""" + + def remove_metadata(self) -> CandidateSchema: + """Return the schema without producer-specific metadata.""" + ... + + def serialize(self) -> SerializedCandidateSchema: + """Serialize the normalized schema.""" + ... + + +@dataclass(frozen=True, slots=True) +class _DirectoryBinding: + parent_descriptor: int | None + name: str | None + descriptor: int + display_path: Path + + def rebind(self) -> None: + opened = os.fstat(self.descriptor) + current = ( + self.display_path.lstat() + if self.parent_descriptor is None or self.name is None + else os.stat(self.name, dir_fd=self.parent_descriptor, follow_symlinks=False) + ) + if not stat.S_ISDIR(current.st_mode) or current.st_mode & 0o022 or _identity(current) != _identity(opened): + raise OSError(f"candidate directory {self.display_path} changed during finalization") + + +@dataclass(frozen=True, slots=True) +class _FileBinding: + relative_path: str + descriptor: int + file_facts: tuple[int, int, int, int, int] + + def rebind(self, dataset_descriptor: int, dataset_path: Path) -> None: + parts = PurePosixPath(self.relative_path).parts + with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( + parent_descriptor, + parent_path, + _, + ): + current = os.stat(parts[-1], dir_fd=parent_descriptor, follow_symlinks=False) + if not _is_safe_file(current) or _file_facts(current) != self.file_facts: + raise OSError(f"candidate file {parent_path / parts[-1]} changed during finalization") + + def validate_lease(self) -> None: + current = os.fstat(self.descriptor) + if not _is_safe_file(current) or _file_facts(current) != self.file_facts: + raise OSError(f"candidate file {self.relative_path!r} changed during finalization") + + +@dataclass(frozen=True, slots=True) +class _ArtifactSnapshot: + relative_path: str + directory_identities: tuple[tuple[int, int], ...] + file_facts: tuple[int, int, int, int, int] + + +@dataclass(frozen=True, slots=True) +class VerifiedCandidateArtifacts: + """Bounded live candidate leases plus metadata derived from artifact bytes.""" + + record_counts: tuple[int, ...] + dataset_schema_digest: str + _dataset: _DirectoryBinding + _files: tuple[_FileBinding, ...] + + def rebind(self) -> None: + """Verify every leased file still occupies its original candidate path.""" + self._dataset.rebind() + for output_file in self._files: + output_file.rebind(self._dataset.descriptor, self._dataset.display_path) + for output_file in self._files: + output_file.validate_lease() + + +@dataclass(frozen=True, slots=True) +class CandidateArtifactSnapshot: + """Descriptor-free identity snapshot retained during bounded collection.""" + + record_counts: tuple[int, ...] + dataset_schema_digest: str + _dataset_identity: tuple[int, int] + _artifacts: tuple[_ArtifactSnapshot, ...] + + +class CandidateArtifactVerifier: + """Verify one manifest-bounded candidate and lease its files through publication.""" + + def __init__(self, path_resolver: Callable[[str], Path] = Path) -> None: + self._path_resolver = path_resolver + + @contextmanager + def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandidateArtifacts]: + dataset_path = self._path_resolver(candidate.dataset_path) + with ExitStack() as resources: + dataset_descriptor = resources.enter_context(open_verified_directory(dataset_path, require_private=True)) + dataset = _DirectoryBinding(None, None, dataset_descriptor, dataset_path) + metadata = tuple( + _open_output_file(resources, dataset_descriptor, dataset_path, output_file) + for output_file in candidate.files + ) + schema_digests = tuple(schema_digest for _, schema_digest, _ in metadata) + if not schema_digests or any(digest != schema_digests[0] for digest in schema_digests[1:]): + raise OSError("candidate Parquet files do not share one dataset schema") + yield VerifiedCandidateArtifacts( + record_counts=tuple(record_count for record_count, _, _ in metadata), + dataset_schema_digest=schema_digests[0], + _dataset=dataset, + _files=tuple(binding for _, _, binding in metadata), + ) + + def inspect(self, candidate: CandidateOutputManifest) -> CandidateArtifactSnapshot: + """Inspect one candidate and return identities without retaining descriptors.""" + dataset_path = Path(candidate.dataset_path) + with open_verified_directory(dataset_path, require_private=True) as dataset_descriptor: + record_counts, schema_digest, artifacts = _inspect_candidate_files( + dataset_descriptor, + dataset_path, + candidate.files, + ) + return CandidateArtifactSnapshot( + record_counts=record_counts, + dataset_schema_digest=schema_digest, + _dataset_identity=_identity(os.fstat(dataset_descriptor)), + _artifacts=artifacts, + ) + + def rebind(self, candidate: CandidateOutputManifest, expected: CandidateArtifactSnapshot) -> None: + """Reopen a candidate and require the identities captured before collection.""" + actual = self.inspect(candidate) + if actual != expected: + raise OSError("candidate paths or metadata changed during collection") + + @contextmanager + def open_output( + self, + candidate: CandidateOutputManifest, + output_file: CandidateOutputFile, + ) -> Iterator[BinaryIO]: + """Yield one digest-verified candidate file through a bounded descriptor.""" + if output_file not in candidate.files: + raise ValueError("candidate output file is not declared by the manifest") + dataset_path = Path(candidate.dataset_path) + parts = PurePosixPath(output_file.relative_path).parts + with open_verified_directory(dataset_path, require_private=True) as dataset_descriptor: + with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( + parent_descriptor, + parent_path, + _, + ): + display_path = parent_path / parts[-1] + with open_verified_regular_file( + parent_descriptor, + parts[-1], + display_path, + expected_size=output_file.byte_size, + expected_sha256=output_file.sha256, + require_private=False, + ) as descriptor: + with os.fdopen(os.dup(descriptor), "rb") as source: + yield source + + +def compute_candidate_schema_digest(schema: CandidateSchema) -> str: + """Compute the version-1 digest for an attempt-local candidate schema.""" + payload = schema.remove_metadata().serialize().to_pybytes() + return hashlib.sha256(payload).hexdigest() + + +def _open_output_file( + resources: ExitStack, + dataset_descriptor: int, + dataset_path: Path, + output_file: CandidateOutputFile, +) -> tuple[int, str, _FileBinding]: + parts = PurePosixPath(output_file.relative_path).parts + parent_descriptor, parent_path, _ = resources.enter_context( + _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) + ) + name = parts[-1] + display_path = parent_path / name + descriptor = resources.enter_context( + open_verified_regular_file( + parent_descriptor, + name, + display_path, + expected_size=output_file.byte_size, + expected_sha256=output_file.sha256, + require_private=False, + ) + ) + record_count, schema_digest = _read_parquet_metadata(descriptor, display_path) + binding = _FileBinding( + relative_path=output_file.relative_path, + descriptor=descriptor, + file_facts=_file_facts(os.fstat(descriptor)), + ) + return record_count, schema_digest, binding + + +def _inspect_candidate_files( + dataset_descriptor: int, + dataset_path: Path, + output_files: tuple[CandidateOutputFile, ...], +) -> tuple[tuple[int, ...], str, tuple[_ArtifactSnapshot, ...]]: + metadata = tuple( + _inspect_output_file(dataset_descriptor, dataset_path, output_file) for output_file in output_files + ) + schema_digests = tuple(schema_digest for _, schema_digest, _ in metadata) + if not schema_digests or any(digest != schema_digests[0] for digest in schema_digests[1:]): + raise OSError("candidate Parquet files do not share one dataset schema") + return ( + tuple(record_count for record_count, _, _ in metadata), + schema_digests[0], + tuple(artifact for _, _, artifact in metadata), + ) + + +def _inspect_output_file( + dataset_descriptor: int, + dataset_path: Path, + output_file: CandidateOutputFile, +) -> tuple[int, str, _ArtifactSnapshot]: + parts = PurePosixPath(output_file.relative_path).parts + with _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) as ( + parent_descriptor, + parent_path, + directory_identities, + ): + name = parts[-1] + display_path = parent_path / name + with open_verified_regular_file( + parent_descriptor, + name, + display_path, + expected_size=output_file.byte_size, + expected_sha256=output_file.sha256, + require_private=False, + ) as descriptor: + record_count, schema_digest = _read_parquet_metadata(descriptor, display_path) + artifact = _ArtifactSnapshot( + relative_path=output_file.relative_path, + directory_identities=directory_identities, + file_facts=_file_facts(os.fstat(descriptor)), + ) + return record_count, schema_digest, artifact + + +@contextmanager +def _open_parent_directory( + dataset_descriptor: int, + dataset_path: Path, + parts: tuple[str, ...], +) -> Iterator[tuple[int, Path, tuple[tuple[int, int], ...]]]: + parent_descriptor = os.dup(dataset_descriptor) + parent_path = dataset_path + directory_identities: list[tuple[int, int]] = [] + try: + for part in parts: + child_path = parent_path / part + with open_verified_child_directory( + parent_descriptor, + part, + child_path, + require_private=False, + ) as child_descriptor: + next_descriptor = os.dup(child_descriptor) + os.close(parent_descriptor) + parent_descriptor = next_descriptor + parent_path = child_path + directory_identities.append(_identity(os.fstat(parent_descriptor))) + yield parent_descriptor, parent_path, tuple(directory_identities) + finally: + os.close(parent_descriptor) + + +def _read_parquet_metadata(descriptor: int, display_path: Path) -> tuple[int, str]: + try: + with os.fdopen(os.dup(descriptor), "rb") as source: + parquet_file = lazy.pq.ParquetFile(source) + return parquet_file.metadata.num_rows, compute_candidate_schema_digest(parquet_file.schema_arrow) + except (OSError, lazy.pa.ArrowException) as error: + raise OSError(f"candidate Parquet metadata {display_path} is invalid") from error + + +def _identity(status: os.stat_result) -> tuple[int, int]: + return status.st_dev, status.st_ino + + +def _file_facts(status: os.stat_result) -> tuple[int, int, int, int, int]: + return status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns, status.st_ctime_ns + + +def _is_safe_file(status: os.stat_result) -> bool: + return stat.S_ISREG(status.st_mode) and status.st_nlink == 1 and not status.st_mode & 0o022 + + +__all__ = [ + "CandidateArtifactSnapshot", + "CandidateArtifactVerifier", + "CandidateSchema", + "SerializedCandidateSchema", + "VerifiedCandidateArtifacts", + "compute_candidate_schema_digest", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py new file mode 100644 index 000000000..9e87a5f6d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local binding of a retry task to its persisted attempt.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId, validate_absolute_path +from data_designer.slurm.state.base import SchedulerIdentity +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError +from data_designer.slurm.state.execution import AttemptLifecycleState +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SHARD_ID_ADAPTER = TypeAdapter(ShardId) +_ATTEMPT_ID_ADAPTER = TypeAdapter(AttemptId) + + +def require_attempt_scheduler_identity( + workspace_root: str | Path, + run_id: Identifier, + shard_id: ShardId, + attempt_id: AttemptId, + scheduler: SchedulerIdentity, +) -> None: + """Require the complete persisted attempt chain to name this allocation.""" + root, normalized_run_id, normalized_shard_id, normalized_attempt_id = _validate_identity( + workspace_root, + run_id, + shard_id, + attempt_id, + ) + storage = StateStorage(root, normalized_run_id) + reader = StateReader(storage, normalized_run_id) + run, plan, shard = reader.load_shard_context(normalized_shard_id) + attempts = reader.load_validated_shard_attempts(run, plan, shard) + attempt = reader.get_attempt(attempts, normalized_attempt_id) + if attempt.scheduler != scheduler: + raise StateConflictError("retry allocation does not match the persisted attempt scheduler identity") + if attempt.state is not AttemptLifecycleState.SUBMITTED: + raise StateConflictError("retry allocation requires an unstarted persisted attempt") + + +def _validate_identity( + workspace_root: str | Path, + run_id: Identifier, + shard_id: ShardId, + attempt_id: AttemptId, +) -> tuple[Path, Identifier, ShardId, AttemptId]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + normalized_shard_id = _SHARD_ID_ADAPTER.validate_python(shard_id, strict=True) + normalized_attempt_id = _ATTEMPT_ID_ADAPTER.validate_python(attempt_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid retry allocation identity") from error + return Path(root), normalized_run_id, normalized_shard_id, normalized_attempt_id + + +def main(argv: Sequence[str] | None = None) -> int: + """Validate one allocation identity from explicit scheduler arguments.""" + parser = argparse.ArgumentParser() + parser.add_argument("--workspace-root", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--shard-id", required=True) + parser.add_argument("--attempt-id", required=True) + parser.add_argument("--array-job-id", required=True, type=int) + parser.add_argument("--array-task-id", required=True, type=int) + arguments = parser.parse_args(argv) + require_attempt_scheduler_identity( + arguments.workspace_root, + arguments.run_id, + arguments.shard_id, + arguments.attempt_id, + SchedulerIdentity(array_job_id=arguments.array_job_id, array_task_id=arguments.array_task_id), + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = ["main", "require_attempt_scheduler_identity"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py new file mode 100644 index 000000000..dc3ffcf19 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import TypeAlias + +from pydantic import NonNegativeInt, PositiveInt + +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + Sha256Digest, + validate_absolute_path, + validate_relative_path, +) + +StateValue = ContractValue +StateRecord = ContractRecord + + +def validate_utc_timestamp(value: datetime) -> datetime: + """Validate that a timestamp is timezone-aware and expressed in UTC.""" + if value.tzinfo is None or value.utcoffset() is None: + raise ValueError("timestamp must include timezone information") + if value.utcoffset() != timedelta(0): + raise ValueError("timestamp must be in UTC") + return value + + +def validate_optional_utc_timestamp(value: datetime | None) -> datetime | None: + """Validate an optional timestamp when it is present.""" + if value is None: + return None + return validate_utc_timestamp(value) + + +class SchedulerIdentity(StateValue): + """Slurm array job and task identity assigned to one attempt.""" + + array_job_id: PositiveInt + array_task_id: NonNegativeInt + + +SchedulerJobIdentity: TypeAlias = PositiveInt | SchedulerIdentity + + +__all__ = [ + "ArtifactReference", + "Identifier", + "SchedulerIdentity", + "SchedulerJobIdentity", + "Sha256Digest", + "StateRecord", + "StateValue", + "validate_absolute_path", + "validate_optional_utc_timestamp", + "validate_relative_path", + "validate_utc_timestamp", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py new file mode 100644 index 000000000..3028107c5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process preparation and reconciliation of CPU collection jobs.""" + +from __future__ import annotations + +import sys +from collections.abc import Mapping +from datetime import datetime, timezone +from importlib.metadata import version +from pathlib import Path +from typing import Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, validate_absolute_path +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.collection import render_collection_script +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt, SlurmSubmissionMatch +from data_designer.slurm.runtime.bundle import stage_runtime_bundle +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.state.collection_filesystem import ( + derive_collection_staging_directory, + prepare_collection_destination, + remove_collection_stage, +) +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import ( + derive_collection_state, + validate_collection_result, + validate_collection_status_transition, +) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.observation import SchedulerObservationClient, SchedulerObservationCollector +from data_designer.slurm.state.outputs import CollectionPlan +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.scheduler import SchedulerState +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class CollectionScheduler(SchedulerObservationClient, Protocol): + """Scheduler operations required by collection submission and refresh.""" + + def submit_script( + self, script: str, *, export_environment: Mapping[str, str] | None = None + ) -> SlurmJobSubmissionReceipt: + """Submit one rendered CPU collection job.""" + ... + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact collection submission name.""" + ... + + +class SlurmCollectionCoordinator: + """Persist, submit, and refresh winner-driven collection jobs.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: CollectionScheduler | None = None, + export_environment: Mapping[str, str] | None = None, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + self._export_environment = None if export_environment is None else dict(export_environment) + self._state = StateStorage(root, normalized_run_id) + self._reader = StateReader(self._state, normalized_run_id) + self._collections = CollectionStorage(self._state) + self._inputs = CollectionInputResolver(self._state, self._reader) + self._destinations = CollectionDestinationResolver() + self._collector = SchedulerObservationCollector(self._scheduler) + self._run_id = normalized_run_id + self._python_executable = _resolve_python_executable() + + def submit( + self, + *, + destination: str | Path | None = None, + num_partitions: int | None = None, + submitted_at: datetime | None = None, + ) -> CollectionStatus: + """Validate winners and submit a CPU-only collection job.""" + timestamp = datetime.now(timezone.utc) if submitted_at is None else submitted_at + try: + with self._collections.acquire_lock(): + self._collections.discard_incomplete_tail() + current = self._get_current_status() + if current is not None: + current_plan = self._load_bound_plan(current) + if current.state is CollectionState.PREPARED: + current = self._reconcile_prepared_collection(current_plan, current, timestamp) + if current.state is not CollectionState.FAILED: + self._validate_existing_destination(current, destination) + if num_partitions is not None and current_plan.num_partitions != num_partitions: + raise StateConflictError("persisted collection partition count does not match the request") + return current + resolved_destination = self._destinations.validate_persisted( + self._reader.load_resolved_plan(), + current_plan, + ) + remove_collection_stage( + Path(current_plan.host_destination), + current.staging_directory, + Path(resolved_destination.authorized_root), + ) + run, resolved_plan, _ = self._reader.load_context() + requested_destination = ( + current_plan.host_destination if destination is None and current is not None else destination + ) + resolved_destination = self._destinations.resolve(resolved_plan, requested_destination) + effective_partitions = num_partitions + if effective_partitions is None: + effective_partitions = current_plan.num_partitions if current is not None else 1 + planned_shards = self._inputs.get_winner_shards() + runtime_bundle = stage_runtime_bundle(self._state.workspace_root) + collection_plan = CollectionPlan( + schema_version=1, + collection_id=self._collections.get_next_collection_id(), + run_id=run.run_id, + created_at=timestamp, + resolved_plan=run.resolved_plan, + planned_shards=planned_shards, + python_executable=self._python_executable, + package_version=version("data-designer-slurm"), + runtime_bundle=runtime_bundle, + host_destination=resolved_destination.host_path, + num_partitions=effective_partitions, + ) + self._inputs.resolve(collection_plan) + prepare_collection_destination( + Path(collection_plan.host_destination), + Path(resolved_destination.authorized_root), + ) + self._collections.ensure_collection(collection_plan.collection_id) + self._collections.publish_plan(collection_plan) + prepared = CollectionStatus( + schema_version=1, + collection_id=collection_plan.collection_id, + run_id=run.run_id, + collection_plan=self._collections.get_plan_reference(collection_plan), + staging_directory=derive_collection_staging_directory(collection_plan), + revision=1, + updated_at=timestamp, + state=CollectionState.PREPARED, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, + ) + self._collections.publish_status(prepared) + script = render_collection_script(resolved_plan, collection_plan, resolved_destination) + return self._submit_prepared(collection_plan, prepared, script, timestamp) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, SlurmRuntimeError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot submit collection for run {self._run_id!r}") from error + + def refresh( + self, + *, + collection_id: Identifier | None = None, + observed_at: datetime | None = None, + ) -> CollectionStatus: + """Reconcile one persisted collection from scheduler and publication evidence.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + with self._collections.acquire_lock(): + selected_id = self._get_selected_collection_id(collection_id) + previous = self._collections.read_status(selected_id) + plan = self._load_bound_plan(previous) + resolved_plan, _ = self._inputs.resolve(plan) + destination = self._destinations.validate_persisted(resolved_plan, plan) + if previous.state is CollectionState.SUCCEEDED: + self._load_valid_result(plan, previous) + return previous + recovered = self._load_optional_result(plan) + if recovered is not None: + return self._publish_succeeded(plan, previous, recovered, timestamp) + if previous.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + previous.staging_directory, + Path(destination.authorized_root), + ) + return previous + if previous.state is CollectionState.PREPARED: + previous = self._reconcile_prepared_collection(plan, previous, timestamp) + if previous.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + previous.staging_directory, + Path(destination.authorized_root), + ) + return previous + assert previous.scheduler is not None + observations = self._collector.collect( + (previous.scheduler,), + observed_at=timestamp, + previous={previous.scheduler: previous.scheduler_observation}, + ) + observation = observations[0] + state = derive_collection_state(observation.state) + if observation.state is SchedulerState.COMPLETED: + state = CollectionState.FAILED + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=timestamp, + state=state, + scheduler_observation=observation, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + if current.state is CollectionState.FAILED: + remove_collection_stage( + Path(plan.host_destination), + current.staging_directory, + Path(destination.authorized_root), + ) + return current + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot refresh collection for run {self._run_id!r}") from error + + def _submit_prepared( + self, + plan: CollectionPlan, + prepared: CollectionStatus, + script: str, + submitted_at: datetime, + ) -> CollectionStatus: + try: + receipt = ( + self._scheduler.submit_script(script) + if self._export_environment is None + else self._scheduler.submit_script(script, export_environment=self._export_environment) + ) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + failed = _updated_status( + prepared, + revision=2, + updated_at=submitted_at, + state=CollectionState.FAILED, + reconciliation_deadline=None, + ) + validate_collection_status_transition(prepared, failed) + self._collections.replace_status(failed) + raise SlurmStateError(f"cannot submit collection {plan.collection_id!r}") from error + except SlurmLauncherError as error: + raise SlurmStateError(f"cannot submit collection {plan.collection_id!r}") from error + submitted = _updated_status( + prepared, + revision=2, + updated_at=submitted_at, + state=CollectionState.SUBMITTED, + scheduler=receipt.job_id, + reconciliation_deadline=None, + ) + validate_collection_status_transition(prepared, submitted) + self._collections.replace_status(submitted) + return submitted + + def _reconcile_prepared_collection( + self, + plan: CollectionPlan, + prepared: CollectionStatus, + observed_at: datetime, + ) -> CollectionStatus: + assert prepared.reconciliation_deadline is not None + job_id = resolve_prepared_submission( + self._scheduler, + PreparedSubmission( + job_name=plan.submission_job_name, + submitted_after=plan.created_at, + reconciliation_deadline=prepared.reconciliation_deadline, + expected_array_task_ids=None, + ), + observed_at=observed_at, + ) + state = CollectionState.SUBMITTED if job_id is not None else CollectionState.FAILED + current = _updated_status( + prepared, + revision=prepared.revision + 1, + updated_at=observed_at, + state=state, + scheduler=job_id, + reconciliation_deadline=None, + ) + validate_collection_status_transition(prepared, current) + self._collections.replace_status(current) + return current + + def _get_current_status(self) -> CollectionStatus | None: + collection_ids = self._collections.list_collection_ids() + return None if not collection_ids else self._collections.read_status(collection_ids[-1]) + + def _validate_existing_destination( + self, + status: CollectionStatus, + requested_destination: str | Path | None, + ) -> None: + plan = self._load_bound_plan(status) + resolved_plan = self._reader.load_resolved_plan() + resolved = self._destinations.validate_persisted(resolved_plan, plan) + if requested_destination is None: + if status.state is CollectionState.SUCCEEDED: + self._load_valid_result(plan, status) + return + if Path(requested_destination).as_posix() != resolved.host_path: + raise StateConflictError("persisted collection destination does not match the requested destination") + requested = self._destinations.resolve(resolved_plan, requested_destination) + if requested != resolved: + raise StateConflictError("persisted collection destination does not match the requested destination") + if plan.host_destination != resolved.host_path: + raise StateCorruptionError("persisted collection destination does not match the resolved plan") + if status.state is CollectionState.SUCCEEDED: + self._load_valid_result(plan, status) + + def _get_selected_collection_id(self, collection_id: Identifier | None) -> Identifier: + collection_ids = self._collections.list_collection_ids() + if not collection_ids: + raise StateConflictError("run has no persisted collection") + selected = collection_ids[-1] if collection_id is None else collection_id + if selected not in collection_ids: + raise StateConflictError("requested collection is not persisted for this run") + return selected + + def _load_optional_result(self, plan: CollectionPlan) -> CollectionResult | None: + try: + return self._load_valid_result(plan) + except FileNotFoundError: + return None + + def _load_valid_result( + self, + plan: CollectionPlan, + status: CollectionStatus | None = None, + ) -> CollectionResult: + _, candidates = self._inputs.resolve(plan) + result = self._collections.read_result(plan) + validated = validate_collection_result( + plan, + result, + expected_records=sum(candidate.actual_records for candidate in candidates), + ) + if status is not None and status.result != self._collections.get_result_reference(plan, validated): + raise StateCorruptionError("collection status does not bind its published result") + self._collections.verify_result_files(plan, validated, verify_digests=False) + return validated + + def _publish_succeeded( + self, + plan: CollectionPlan, + previous: CollectionStatus, + result: CollectionResult, + timestamp: datetime, + ) -> CollectionStatus: + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=timestamp, + state=CollectionState.SUCCEEDED, + result=self._collections.get_result_reference(plan, result), + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return current + + def _load_bound_plan(self, status: CollectionStatus) -> CollectionPlan: + plan = self._collections.read_plan(status.collection_id) + if status.collection_plan != self._collections.get_plan_reference(plan): + raise StateCorruptionError("collection status does not bind its persisted collection plan") + if status.staging_directory != derive_collection_staging_directory(plan): + raise StateCorruptionError("collection status does not bind its exact staging directory") + return plan + + +def _resolve_python_executable() -> str: + raw = Path(sys.executable).absolute().as_posix() + try: + return validate_absolute_path(raw) + except ValueError as error: + raise SlurmStateError("collection Python executable must be a normalized absolute path") from error + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted collection location") from error + return Path(root), normalized_run_id + + +def _updated_status(previous: CollectionStatus, **updates: object) -> CollectionStatus: + payload = previous.model_dump(mode="python") + payload.update(updates) + return CollectionStatus.model_validate(payload) + + +__all__ = ["CollectionScheduler", "SlurmCollectionCoordinator"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py new file mode 100644 index 000000000..fdc239f79 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py @@ -0,0 +1,527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Same-parent staging and atomic no-overwrite collection publication.""" + +from __future__ import annotations + +import ctypes +import errno +import hashlib +import os +import stat +import sys +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + +from data_designer.slurm.contracts import is_path_below +from data_designer.slurm.filesystem import PRIVATE_DIRECTORY_MODE, open_verified_directory +from data_designer.slurm.state.errors import StateConflictError +from data_designer.slurm.state.filesystem import open_verified_child_directory, open_verified_regular_file +from data_designer.slurm.state.outputs import CollectionPlan + +_RENAME_NOREPLACE = 1 +_RENAME_EXCL = 0x00000004 +_UNSUPPORTED_NO_REPLACE_ERRORS = frozenset( + { + errno.EINVAL, + errno.ENOSYS, + errno.ENOTSUP, + errno.EOPNOTSUPP, + } +) + + +@dataclass(frozen=True, slots=True) +class StagedFile: + """Expected immutable bytes in a collection staging directory.""" + + name: str + sha256: str + byte_size: int + + +@dataclass(slots=True) +class StagedCollection: + """Private sibling directory that becomes visible only after publication.""" + + path: Path + destination: Path + _parent_descriptor: int + _parent_identity: tuple[int, int] + _stage_identity: tuple[int, int] + _published: bool = False + + def publish(self, expected_files: tuple[StagedFile, ...]) -> None: + """Atomically rename the complete stage while refusing any collision.""" + if self._published: + return + self._rebind() + _verify_stage_files(self, expected_files) + _rename_without_overwrite( + self._parent_descriptor, + self.path.name, + self._parent_descriptor, + self.destination.name, + ) + self._published = True + try: + self._rebind_published_destination() + os.fsync(self._parent_descriptor) + self._rebind_published_destination() + except OSError: + self._restore_stage() + raise + + def _rebind(self) -> None: + self._rebind_parent() + stage = os.stat(self.path.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + if not stat.S_ISDIR(stage.st_mode) or _identity(stage) != self._stage_identity: + raise OSError(f"collection stage {self.path} changed") + + def _rebind_parent(self) -> None: + opened = os.fstat(self._parent_descriptor) + current = self.destination.parent.lstat() + if ( + not stat.S_ISDIR(current.st_mode) + or _identity(opened) != self._parent_identity + or _identity(current) != self._parent_identity + ): + raise OSError(f"collection destination parent {self.destination.parent} changed") + + def _rebind_published_destination(self) -> None: + self._rebind_parent() + opened_view = os.stat(self.destination.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + path_view = self.destination.lstat() + if ( + not stat.S_ISDIR(opened_view.st_mode) + or not stat.S_ISDIR(path_view.st_mode) + or _identity(opened_view) != self._stage_identity + or _identity(path_view) != self._stage_identity + ): + raise OSError(f"published collection destination {self.destination} changed") + + def _restore_stage(self) -> None: + published = os.stat(self.destination.name, dir_fd=self._parent_descriptor, follow_symlinks=False) + if not stat.S_ISDIR(published.st_mode) or _identity(published) != self._stage_identity: + raise OSError(f"published collection destination {self.destination} changed before rollback") + _rename_without_overwrite( + self._parent_descriptor, + self.destination.name, + self._parent_descriptor, + self.path.name, + ) + self._published = False + os.fsync(self._parent_descriptor) + + +def derive_collection_staging_directory(plan: CollectionPlan) -> str: + """Derive one collision-resistant, persisted stage identity from the plan.""" + return f".dd-collection-{plan.compute_sha256()[:32]}.tmp" + + +def remove_collection_stage(destination: Path, staging_directory: str, authorized_root: Path) -> None: + """Remove only the exact persisted stage for one terminal collection.""" + try: + with _open_authorized_parent(destination, authorized_root) as parent_descriptor: + _require_restrictive_parent(parent_descriptor) + _recover_private_reservation(parent_descriptor, staging_directory, destination.name) + _remove_existing_stage(parent_descriptor, destination.parent, staging_directory) + except _MissingDestinationParent: + return + + +def prepare_collection_destination(destination: Path, authorized_root: Path) -> None: + """Create and validate the destination parent without creating the dataset.""" + with _open_authorized_parent(destination, authorized_root, create_missing=True) as parent_descriptor: + _require_restrictive_parent(parent_descriptor) + _require_absent(parent_descriptor, destination.name, destination) + + +@contextmanager +def stage_collection( + destination: Path, + staging_directory: str, + authorized_root: Path, +) -> Iterator[StagedCollection]: + """Yield a private sibling stage and remove it unless atomically published.""" + if destination.parent == destination: + raise StateConflictError("collection destination cannot be the filesystem root") + with _open_authorized_parent(destination, authorized_root) as parent_descriptor: + parent_status = _require_restrictive_parent(parent_descriptor) + _recover_private_reservation(parent_descriptor, staging_directory, destination.name) + _require_absent(parent_descriptor, destination.name, destination) + _remove_existing_stage(parent_descriptor, destination.parent, staging_directory) + stage_name = _create_stage_directory(parent_descriptor, staging_directory) + stage_path = destination.parent / stage_name + staged = StagedCollection( + path=stage_path, + destination=destination, + _parent_descriptor=parent_descriptor, + _parent_identity=_identity(parent_status), + _stage_identity=_identity(os.stat(stage_name, dir_fd=parent_descriptor, follow_symlinks=False)), + ) + try: + yield staged + finally: + if not staged._published: + _remove_stage(staged) + + +class _MissingDestinationParent(FileNotFoundError): + """A child beneath an existing authorized root has not been created yet.""" + + +@contextmanager +def _open_authorized_parent( + destination: Path, + authorized_root: Path, + *, + create_missing: bool = False, +) -> Iterator[int]: + destination_text = destination.as_posix() + root_text = authorized_root.as_posix() + if destination_text == root_text or not is_path_below(destination_text, root_text): + raise StateConflictError("collection destination must be below its authorized mount root") + relative_parent = PurePosixPath(destination.parent.relative_to(authorized_root)).parts + with ExitStack() as resources: + descriptor = resources.enter_context(open_verified_directory(authorized_root, resource_name="collection")) + current_path = authorized_root + for part in relative_parent: + current_path /= part + if create_missing: + _ensure_private_child_directory(descriptor, part) + try: + descriptor = resources.enter_context( + open_verified_child_directory( + descriptor, + part, + current_path, + require_private=False, + ) + ) + except FileNotFoundError as error: + raise _MissingDestinationParent(current_path) from error + yield descriptor + + +def _ensure_private_child_directory(parent_descriptor: int, name: str) -> None: + try: + os.mkdir(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor) + except FileExistsError: + return + os.fsync(parent_descriptor) + + +def _require_restrictive_parent(parent_descriptor: int) -> os.stat_result: + status = os.fstat(parent_descriptor) + if status.st_mode & 0o022: + raise StateConflictError("collection destination parent must not be group- or world-writable") + return status + + +def _require_absent(parent_descriptor: int, name: str, display_path: Path) -> None: + try: + os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + raise StateConflictError(f"collection destination {display_path} already exists") + + +def _create_stage_directory(parent_descriptor: int, name: str) -> str: + os.mkdir(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + return name + + +def _remove_existing_stage(parent_descriptor: int, parent_path: Path, name: str) -> None: + try: + status = os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISDIR(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"persisted collection stage {parent_path / name} is not a private directory") + _delete_stage_directory(parent_descriptor, name, parent_path / name, _identity(status)) + + +def _remove_stage(staged: StagedCollection) -> None: + try: + current = os.stat(staged.path.name, dir_fd=staged._parent_descriptor, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISDIR(current.st_mode) or _identity(current) != staged._stage_identity: + raise OSError(f"collection stage {staged.path} changed before cleanup") + _delete_stage_directory( + staged._parent_descriptor, + staged.path.name, + staged.path, + staged._stage_identity, + ) + + +def _delete_stage_directory( + parent_descriptor: int, + name: str, + display_path: Path, + expected_identity: tuple[int, int], +) -> None: + with open_verified_child_directory(parent_descriptor, name, display_path) as descriptor: + if _identity(os.fstat(descriptor)) != expected_identity: + raise OSError(f"collection stage {display_path} changed before cleanup") + _clear_directory(descriptor, display_path) + current = os.stat(name, dir_fd=parent_descriptor, follow_symlinks=False) + if _identity(current) != expected_identity: + raise OSError(f"collection stage {display_path} changed during cleanup") + os.rmdir(name, dir_fd=parent_descriptor) + os.fsync(parent_descriptor) + + +def _clear_directory(directory_descriptor: int, display_path: Path) -> None: + for name in os.listdir(directory_descriptor): + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + child_path = display_path / name + if stat.S_ISDIR(status.st_mode): + with open_verified_child_directory(directory_descriptor, name, child_path) as child_descriptor: + if _identity(os.fstat(child_descriptor)) != _identity(status): + raise OSError(f"collection staging directory {child_path} changed before cleanup") + _clear_directory(child_descriptor, child_path) + current = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if _identity(current) != _identity(status): + raise OSError(f"collection staging directory {child_path} changed during cleanup") + os.rmdir(name, dir_fd=directory_descriptor) + else: + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +def _verify_stage_files(staged: StagedCollection, expected_files: tuple[StagedFile, ...]) -> None: + expected_names = tuple(file.name for file in expected_files) + if len(expected_names) != len(set(expected_names)): + raise OSError("collection stage expectation contains duplicate paths") + with open_verified_child_directory( + staged._parent_descriptor, + staged.path.name, + staged.path, + ) as stage_descriptor: + if set(os.listdir(stage_descriptor)) != set(expected_names): + raise OSError("collection stage inventory changed before publication") + for expected in expected_files: + with open_verified_regular_file( + stage_descriptor, + expected.name, + staged.path / expected.name, + expected_size=expected.byte_size, + expected_sha256=expected.sha256, + require_private=False, + ): + pass + + +def _rename_without_overwrite( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, +) -> None: + try: + _atomic_rename_without_overwrite( + source_directory, + source_name, + destination_directory, + destination_name, + ) + except OSError as error: + if error.errno not in _UNSUPPORTED_NO_REPLACE_ERRORS: + raise + _rename_over_private_reservation( + source_directory, + source_name, + destination_directory, + destination_name, + ) + + +def _atomic_rename_without_overwrite( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, +) -> None: + library = ctypes.CDLL(None, use_errno=True) + source = os.fsencode(source_name) + destination = os.fsencode(destination_name) + if sys.platform.startswith("linux") and hasattr(library, "renameat2"): + result = library.renameat2( + source_directory, + ctypes.c_char_p(source), + destination_directory, + ctypes.c_char_p(destination), + _RENAME_NOREPLACE, + ) + elif sys.platform == "darwin" and hasattr(library, "renameatx_np"): + result = library.renameatx_np( + source_directory, + ctypes.c_char_p(source), + destination_directory, + ctypes.c_char_p(destination), + _RENAME_EXCL, + ) + else: + raise OSError(errno.ENOTSUP, "atomic no-overwrite directory rename is unavailable") + if result == 0: + return + error_number = ctypes.get_errno() + if error_number in {errno.EEXIST, errno.ENOTEMPTY}: + raise StateConflictError(f"collection destination {destination_name!r} already exists") + raise OSError(error_number, os.strerror(error_number), destination_name) + + +def _rename_over_private_reservation( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, +) -> None: + """Publish through an exclusive empty directory on filesystems without no-replace rename.""" + marker_name = _reservation_marker_name(source_name, destination_name) + reservation: os.stat_result | None = None + remove_marker = False + try: + try: + os.mkdir(destination_name, PRIVATE_DIRECTORY_MODE, dir_fd=destination_directory) + except FileExistsError: + raise StateConflictError(f"collection destination {destination_name!r} already exists") from None + reservation = os.stat(destination_name, dir_fd=destination_directory, follow_symlinks=False) + if not stat.S_ISDIR(reservation.st_mode) or reservation.st_mode & 0o077: + raise OSError(f"collection destination reservation {destination_name!r} is not private") + os.fsync(destination_directory) + # An unmarked directory is ambiguous after a crash and must never be removed by recovery. + marker_bytes = _reservation_marker_bytes(source_name, destination_name, _reservation_identity(reservation)) + _create_reservation_marker(destination_directory, marker_name, marker_bytes) + remove_marker = True + current = os.stat(destination_name, dir_fd=destination_directory, follow_symlinks=False) + if _identity(current) != _identity(reservation): + raise OSError(f"collection destination reservation {destination_name!r} changed") + os.rename( + source_name, + destination_name, + src_dir_fd=source_directory, + dst_dir_fd=destination_directory, + ) + except BaseException: + if reservation is not None: + _remove_unchanged_reservation(destination_directory, destination_name, reservation) + raise + finally: + if remove_marker: + os.unlink(marker_name, dir_fd=destination_directory) + os.fsync(destination_directory) + + +def _reservation_marker_name(source_name: str, destination_name: str) -> str: + identity = hashlib.sha256(f"{source_name}\n{destination_name}\n".encode("utf-8")).hexdigest()[:32] + return f".dd-collection-reservation-{identity}" + + +def _reservation_marker_bytes(source_name: str, destination_name: str, identity: tuple[int, int, int]) -> bytes: + return f"{source_name}\n{destination_name}\n{identity[0]}\n{identity[1]}\n{identity[2]}\n".encode("utf-8") + + +def _create_reservation_marker(directory: int, name: str, content: bytes) -> None: + descriptor = os.open( + name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), 0o600, dir_fd=directory + ) + try: + if os.write(descriptor, content) != len(content): + raise OSError("collection reservation marker write was incomplete") + os.fsync(descriptor) + except BaseException: + os.unlink(name, dir_fd=directory) + raise + finally: + os.close(descriptor) + os.fsync(directory) + + +def _recover_private_reservation(directory: int, source_name: str, destination_name: str) -> None: + marker_name = _reservation_marker_name(source_name, destination_name) + try: + marker = os.stat(marker_name, dir_fd=directory, follow_symlinks=False) + except FileNotFoundError: + return + if ( + not stat.S_ISREG(marker.st_mode) + or marker.st_mode & 0o077 + or marker.st_nlink != 1 + or marker.st_uid != os.getuid() + ): + raise OSError(f"collection reservation marker {marker_name!r} is unsafe") + descriptor = os.open(marker_name, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), dir_fd=directory) + try: + if _identity(os.fstat(descriptor)) != _identity(marker): + raise OSError(f"collection reservation marker {marker_name!r} changed") + marker_content = os.read(descriptor, 1024) + finally: + os.close(descriptor) + marker_prefix = f"{source_name}\n{destination_name}\n".encode("utf-8") + if not marker_content.startswith(marker_prefix): + raise OSError(f"collection reservation marker {marker_name!r} is invalid") + identity_parts = marker_content[len(marker_prefix) :].split(b"\n") + if len(identity_parts) != 4 or identity_parts[3] or not all(part.isdigit() for part in identity_parts[:3]): + raise OSError(f"collection reservation marker {marker_name!r} is invalid") + marker_identity = int(identity_parts[0]), int(identity_parts[1]), int(identity_parts[2]) + try: + reservation = os.stat(destination_name, dir_fd=directory, follow_symlinks=False) + except FileNotFoundError: + reservation = None + if ( + reservation is not None + and stat.S_ISDIR(reservation.st_mode) + and not reservation.st_mode & 0o077 + and _reservation_identity(reservation) == marker_identity + ): + opened = os.open( + destination_name, os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0), dir_fd=directory + ) + try: + if _identity(os.fstat(opened)) == _identity(reservation) and not os.listdir(opened): + _remove_unchanged_reservation(directory, destination_name, reservation) + finally: + os.close(opened) + os.unlink(marker_name, dir_fd=directory) + os.fsync(directory) + + +def _remove_unchanged_reservation( + destination_directory: int, + destination_name: str, + reservation: os.stat_result, +) -> None: + try: + current = os.stat(destination_name, dir_fd=destination_directory, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISDIR(current.st_mode) or _reservation_identity(current) != _reservation_identity(reservation): + return + os.rmdir(destination_name, dir_fd=destination_directory) + os.fsync(destination_directory) + + +def _identity(status: os.stat_result) -> tuple[int, int]: + return status.st_dev, status.st_ino + + +def _reservation_identity(status: os.stat_result) -> tuple[int, int, int]: + return status.st_dev, status.st_ino, status.st_ctime_ns + + +__all__ = [ + "StagedCollection", + "StagedFile", + "derive_collection_staging_directory", + "prepare_collection_destination", + "remove_collection_stage", + "stage_collection", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py new file mode 100644 index 000000000..75319d0e6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Winner-only resolution of deterministic collection inputs.""" + +from __future__ import annotations + +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier +from data_designer.slurm.state.collection_schema import require_compatible_candidate_schemas +from data_designer.slurm.state.collection_validation import validate_collection_inputs +from data_designer.slurm.state.errors import StateCorruptionError, StateNotFoundError +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan, CollectionShard, ShardWinner +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import StateContractError, validate_collection_plan + + +class CollectionInputResolver: + """Resolve planned winner chains without scanning attempt output trees.""" + + def __init__(self, storage: StateStorage, reader: StateReader) -> None: + self._storage = storage + self._reader = reader + self._finalizer = WinnerFinalizer(storage, reader) + self._verifier = CandidateArtifactVerifier() + + def get_winner_shards(self) -> tuple[CollectionShard, ...]: + """Return one canonical winner reference for every ordered run shard.""" + run, plan, shards = self._reader.load_context() + planned: list[CollectionShard] = [] + for shard in shards: + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + raise StateNotFoundError(f"shard {shard.shard_id!r} has no winner") + planned.append( + CollectionShard( + shard_id=shard.shard_id, + winner_manifest=ArtifactReference( + path=self._storage.get_winner_path(shard.shard_id).as_posix(), + sha256=winner.compute_sha256(), + ), + ) + ) + return tuple(planned) + + def resolve( + self, collection_plan: CollectionPlan + ) -> tuple[ResolvedSlurmRunPlan, tuple[CandidateOutputManifest, ...]]: + """Validate the complete winner chain and return ordered candidate manifests.""" + run, plan, shards = self._reader.load_context() + winners: list[ShardWinner] = [] + candidates: list[CandidateOutputManifest] = [] + for planned_collection_shard, shard in zip(collection_plan.planned_shards, shards, strict=True): + if planned_collection_shard.shard_id != shard.shard_id: + raise StateCorruptionError("collection plan does not preserve planned shard order") + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + raise StateNotFoundError(f"shard {shard.shard_id!r} has no winner") + expected_reference = ArtifactReference( + path=self._storage.get_winner_path(shard.shard_id).as_posix(), + sha256=winner.compute_sha256(), + ) + if planned_collection_shard.winner_manifest != expected_reference: + raise StateCorruptionError(f"collection winner changed for shard {shard.shard_id!r}") + attempt = self._reader.get_attempt(attempts, winner.attempt_id) + result = self._reader.load_optional_attempt_result(plan, shard, attempt) + if result is None: + raise StateCorruptionError(f"winning attempt {winner.attempt_id!r} has no result records") + winners.append(winner) + candidates.append(result[1]) + try: + validate_collection_plan(run, collection_plan, shards, tuple(winners)) + validate_collection_inputs(plan, collection_plan, tuple(candidates)) + except StateContractError as error: + raise StateCorruptionError("collection inputs violate persisted run intent") from error + if len({candidate.dataset_schema_digest for candidate in candidates}) != 1: + require_compatible_candidate_schemas(tuple(candidates), self._verifier) + return plan, tuple(candidates) + + +__all__ = ["CollectionInputResolver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py new file mode 100644 index 000000000..aa2d2cda2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bounded-memory merging of validated Parquet shard winners.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import BinaryIO, Protocol + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.filesystem import get_file_facts +from data_designer.slurm.state.artifacts import CandidateArtifactSnapshot, CandidateArtifactVerifier +from data_designer.slurm.state.collection_filesystem import StagedCollection, StagedFile +from data_designer.slurm.state.collection_records import CollectedOutputFile, CollectionResult +from data_designer.slurm.state.collection_schema import require_compatible_candidate_schemas +from data_designer.slurm.state.filesystem import open_verified_directory, publish_immutable_text, sync_directory +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan + +_BATCH_SIZE = 65_536 +_RESULT_FILENAME = "collection-result.json" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class CollectionPartitionWriter(Protocol): + """Bounded writer for one deterministic output partition.""" + + def write(self, batch: object) -> None: + """Append one Arrow record batch.""" + ... + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + """Seal and describe the exact file descriptor that received the records.""" + ... + + +class _ArrowBatchWriter(Protocol): + def write_batch(self, batch: object) -> None: ... + + def close(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _PartitionLayout: + stage_path: Path + schema: object + record_counts: tuple[int, ...] + + +class CollectionMerger: + """Merge ordered winner files while retaining only bounded data and descriptors.""" + + def __init__( + self, + *, + verifier: CandidateArtifactVerifier | None = None, + completed_at: datetime | None = None, + ) -> None: + self._verifier = verifier if verifier is not None else CandidateArtifactVerifier() + self._completed_at = completed_at + + def merge( + self, + collection_plan: CollectionPlan, + candidates: tuple[CandidateOutputManifest, ...], + staged: StagedCollection, + ) -> CollectionResult: + """Write deterministic partitions, rebind inputs, and publish the stage.""" + snapshots = tuple(self._inspect_candidate(candidate) for candidate in candidates) + schema = ( + require_compatible_candidate_schemas(candidates, self._verifier) + if len({candidate.dataset_schema_digest for candidate in candidates}) != 1 + else None + ) + batches = self._iter_batches(candidates) + try: + first_batch = next(batches, None) + if first_batch is None: + raise OSError("collection inputs contain no records") + layout = _PartitionLayout( + stage_path=staged.path, + schema=schema if schema is not None else first_batch.schema.remove_metadata(), + record_counts=_partition_record_counts( + sum(candidate.actual_records for candidate in candidates), + collection_plan.num_partitions, + ), + ) + files = self._write_partitions(layout, first_batch, batches) + finally: + batches.close() + completion_time = self._completed_at if self._completed_at is not None else _utc_now() + result = CollectionResult( + schema_version=1, + collection_id=collection_plan.collection_id, + run_id=collection_plan.run_id, + completed_at=completion_time, + collection_plan_sha256=collection_plan.compute_sha256(), + actual_records=sum(output.record_count for output in files), + files=files, + ) + result_bytes = self._write_result(staged.path, result) + for candidate, snapshot in zip(candidates, snapshots, strict=True): + self._verifier.rebind(candidate, snapshot) + staged.publish(_stage_expectations(files, result, result_bytes)) + return result + + def _inspect_candidate(self, candidate: CandidateOutputManifest) -> CandidateArtifactSnapshot: + snapshot = self._verifier.inspect(candidate) + if snapshot.record_counts != tuple(output.record_count for output in candidate.files): + raise OSError("candidate Parquet row counts changed before collection") + if snapshot.dataset_schema_digest != candidate.dataset_schema_digest: + raise OSError("candidate Parquet schema changed before collection") + return snapshot + + def _iter_batches(self, candidates: tuple[CandidateOutputManifest, ...]) -> Generator[object, None, None]: + for candidate in candidates: + for output_file in candidate.files: + with self._verifier.open_output(candidate, output_file) as source: + parquet_file = lazy.pq.ParquetFile(source) + yield from parquet_file.iter_batches(batch_size=_BATCH_SIZE) + + def _write_partitions( + self, + layout: _PartitionLayout, + first_batch: object, + remaining_batches: Generator[object, None, None], + ) -> tuple[CollectedOutputFile, ...]: + cursor = _BatchCursor(first_batch, remaining_batches) + outputs: list[CollectedOutputFile] = [] + for partition_index, record_count in enumerate(layout.record_counts): + relative_path = f"part-{partition_index:05d}.parquet" + output_path = layout.stage_path / relative_path + with _open_parquet_writer(output_path, layout.schema) as writer: + cursor.write_records(writer, record_count) + outputs.append(writer.finish(relative_path, record_count)) + if cursor.has_remaining_records(): + raise OSError("collection inputs contain more rows than declared") + return tuple(outputs) + + @staticmethod + def _write_result(stage_path: Path, result: CollectionResult) -> bytes: + serialized = result.serialize_json() + with open_verified_directory(stage_path, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + _RESULT_FILENAME, + serialized, + stage_path / _RESULT_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + sync_directory(descriptor) + return serialized.encode("utf-8") + + +class _BatchCursor: + def __init__(self, first_batch: object, remaining: Generator[object, None, None]) -> None: + self._batch = first_batch + self._remaining = remaining + self._offset = 0 + + def write_records(self, writer: CollectionPartitionWriter, record_count: int) -> None: + remaining = record_count + while remaining: + available = self._batch.num_rows - self._offset + if available == 0: + self._advance() + continue + size = min(remaining, available) + writer.write(self._batch.slice(self._offset, size)) + self._offset += size + remaining -= size + + def has_remaining_records(self) -> bool: + if self._offset < self._batch.num_rows: + return True + return next(self._remaining, None) is not None + + def _advance(self) -> None: + next_batch = next(self._remaining, None) + if next_batch is None: + raise OSError("collection inputs contain fewer rows than declared") + self._batch = next_batch + self._offset = 0 + + +class _BoundOutput: + def __init__(self, path: Path, output: BinaryIO) -> None: + self.path = path + self._output = output + + def describe(self, relative_path: str, record_count: int) -> CollectedOutputFile: + _sync_output(self._output) + return _describe_open_output(self._output.fileno(), self.path, relative_path, record_count) + + +class _ParquetWriter: + def __init__(self, writer: _ArrowBatchWriter, output: _BoundOutput, schema: object) -> None: + self._writer = writer + self._output = output + self._schema = schema + self._closed = False + + def write(self, batch: object) -> None: + normalized = batch.replace_schema_metadata(None) + if not normalized.schema.equals(self._schema): + normalized = normalized.cast(self._schema) + self._writer.write_batch(normalized) + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + self.close() + return self._output.describe(relative_path, record_count) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._writer.close() + + +@contextmanager +def _open_parquet_writer(output_path: Path, schema: object) -> Generator[CollectionPartitionWriter, None, None]: + with output_path.open("x+b") as output: + os.fchmod(output.fileno(), 0o600) + partition_writer = _ParquetWriter( + lazy.pq.ParquetWriter(output, schema), _BoundOutput(output_path, output), schema + ) + try: + yield partition_writer + finally: + partition_writer.close() + _sync_output(output) + + +def _sync_output(output: BinaryIO) -> None: + output.flush() + os.fsync(output.fileno()) + + +def _partition_record_counts(record_count: int, partition_count: int) -> tuple[int, ...]: + floor_count = record_count // partition_count + return tuple( + record_count - floor_count * (partition_count - 1) if index == partition_count - 1 else floor_count + for index in range(partition_count) + ) + + +def _describe_open_output( + descriptor: int, + output_path: Path, + relative_path: str, + record_count: int, +) -> CollectedOutputFile: + before = os.fstat(descriptor) + _require_safe_output(before, output_path) + digest = hashlib.sha256() + offset = 0 + while block := os.pread(descriptor, 1024 * 1024, offset): + digest.update(block) + offset += len(block) + after = os.fstat(descriptor) + path_status = output_path.lstat() + if get_file_facts(before) != get_file_facts(after) or get_file_facts(after) != get_file_facts(path_status): + raise OSError(f"collection output {output_path} changed while it was being described") + _require_safe_output(after, output_path) + return CollectedOutputFile( + relative_path=relative_path, + sha256=digest.hexdigest(), + byte_size=after.st_size, + record_count=record_count, + modified_at_ns=after.st_mtime_ns, + changed_at_ns=after.st_ctime_ns, + ) + + +def _require_safe_output(status: os.stat_result, output_path: Path) -> None: + if not stat.S_ISREG(status.st_mode) or status.st_nlink != 1 or status.st_mode & 0o077: + raise OSError(f"collection output {output_path} is not a private single-link regular file") + + +def _stage_expectations( + files: tuple[CollectedOutputFile, ...], + result: CollectionResult, + result_bytes: bytes, +) -> tuple[StagedFile, ...]: + return ( + *(StagedFile(file.relative_path, file.sha256, file.byte_size) for file in files), + StagedFile(_RESULT_FILENAME, result.compute_sha256(), len(result_bytes)), + ) + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +__all__ = ["CollectionMerger", "CollectionPartitionWriter"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py new file mode 100644 index 000000000..b8774661f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Persisted collection lifecycle and output records.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Annotated + +from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ArtifactReference, Identifier, Sha256Digest, validate_relative_path +from data_designer.slurm.state.base import ( + SchedulerJobIdentity, + StateRecord, + StateValue, + validate_optional_utc_timestamp, + validate_utc_timestamp, +) +from data_designer.slurm.state.scheduler import SchedulerObservation + + +class CollectionState(str, Enum): + """Persisted lifecycle for one CPU collection job.""" + + PREPARED = "prepared" + SUBMITTED = "submitted" + PENDING = "pending" + RUNNING = "running" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class CollectedOutputFile(StateValue): + """One deterministic file in a published collected dataset.""" + + relative_path: str + sha256: Sha256Digest + byte_size: NonNegativeInt + record_count: NonNegativeInt + modified_at_ns: NonNegativeInt + changed_at_ns: NonNegativeInt + + _relative_path_is_safe = field_validator("relative_path")(validate_relative_path) + + +class CollectionResult(StateRecord): + """Immutable proof of a completely staged collection output.""" + + collection_id: Identifier + run_id: Identifier + completed_at: datetime + collection_plan_sha256: Sha256Digest + actual_records: NonNegativeInt + files: tuple[CollectedOutputFile, ...] = Field(min_length=1) + + _completed_at_is_utc = field_validator("completed_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_files(self) -> CollectionResult: + paths = tuple(output.relative_path for output in self.files) + if len(paths) != len(set(paths)): + raise ValueError("collected output paths must be unique") + if sum(output.record_count for output in self.files) != self.actual_records: + raise ValueError("collected output row counts must equal actual_records") + return self + + +class CollectionStatus(StateRecord): + """Atomically replaced scheduler and publication state for one collection.""" + + collection_id: Identifier + run_id: Identifier + collection_plan: ArtifactReference + staging_directory: Annotated[ + str, + StringConstraints(pattern=r"^\.dd-collection-[0-9a-f]{32}\.tmp$"), + ] + revision: PositiveInt + updated_at: datetime + state: CollectionState + scheduler: SchedulerJobIdentity | None = None + scheduler_observation: SchedulerObservation | None = None + result: ArtifactReference | None = None + reconciliation_deadline: datetime | None = None + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) + + @model_validator(mode="after") + def validate_evidence(self) -> CollectionStatus: + if self.state is CollectionState.PREPARED: + if self.scheduler is not None or self.scheduler_observation is not None or self.result is not None: + raise ValueError("prepared collection cannot contain scheduler or result evidence") + if self.reconciliation_deadline is None or self.reconciliation_deadline <= self.updated_at: + raise ValueError("prepared collection requires a future reconciliation deadline") + return self + if self.reconciliation_deadline is not None: + raise ValueError("settled collection cannot contain a reconciliation deadline") + if self.state is not CollectionState.FAILED and self.scheduler is None: + raise ValueError("submitted collection states require a scheduler identity") + if self.scheduler is None and self.scheduler_observation is not None: + raise ValueError("collection observation requires a scheduler identity") + if self.scheduler_observation is not None and self.scheduler_observation.scheduler != self.scheduler: + raise ValueError("collection scheduler observation identity does not match") + if (self.state is CollectionState.SUCCEEDED) != (self.result is not None): + raise ValueError("collection result is required exactly for succeeded state") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_schema.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_schema.py new file mode 100644 index 000000000..fcf572a11 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_schema.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Verify collection schemas while allowing reordered struct children.""" + +from __future__ import annotations + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier, compute_candidate_schema_digest +from data_designer.slurm.state.errors import StateCorruptionError +from data_designer.slurm.state.outputs import CandidateOutputManifest + + +def require_compatible_candidate_schemas( + candidates: tuple[CandidateOutputManifest, ...], + verifier: CandidateArtifactVerifier, +) -> object: + """Return the first verified schema if only struct child order differs.""" + reference: object | None = None + for candidate in candidates: + if not candidate.files: + raise StateCorruptionError("collection candidate has no Parquet files") + try: + with verifier.open_output(candidate, candidate.files[0]) as source: + schema = lazy.pq.ParquetFile(source).schema_arrow.remove_metadata() + except (OSError, lazy.pa.ArrowException) as error: + raise StateCorruptionError("collection candidate Parquet schema cannot be verified") from error + if compute_candidate_schema_digest(schema) != candidate.dataset_schema_digest: + raise StateCorruptionError("collection candidate Parquet schema does not match its manifest") + if reference is None: + reference = schema + elif not schemas_match_except_struct_order(reference, schema): + raise StateCorruptionError("collection candidates have incompatible schemas") + if reference is None: + raise StateCorruptionError("collection has no candidates") + return reference + + +def schemas_match_except_struct_order(reference: object, actual: object) -> bool: + """Require identical top-level fields and types, ignoring struct child order.""" + if len(reference) != len(actual): + return False + return all(_fields_match(left, right) for left, right in zip(reference, actual, strict=True)) + + +def _fields_match(left: object, right: object) -> bool: + if left.name != right.name or left.nullable != right.nullable or left.metadata != right.metadata: + return False + return _types_match(left.type, right.type) + + +def _types_match(left: object, right: object) -> bool: + if left.id != right.id: + return False + if lazy.pa.types.is_struct(left): + left_children = {field.name: field for field in left} + right_children = {field.name: field for field in right} + return ( + len(left_children) == len(left) + and len(right_children) == len(right) + and left_children.keys() == right_children.keys() + and all(_fields_match(left_children[name], right_children[name]) for name in left_children) + ) + if lazy.pa.types.is_list(left) or lazy.pa.types.is_large_list(left): + return _fields_match(left.value_field, right.value_field) + if lazy.pa.types.is_fixed_size_list(left): + return left.list_size == right.list_size and _fields_match(left.value_field, right.value_field) + if lazy.pa.types.is_map(left): + return ( + left.keys_sorted == right.keys_sorted + and _fields_match(left.key_field, right.key_field) + and _fields_match(left.item_field, right.item_field) + ) + return left.equals(right) + + +__all__ = ["require_compatible_candidate_schemas", "schemas_match_except_struct_order"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py new file mode 100644 index 000000000..3d53de2e0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py @@ -0,0 +1,333 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound persistence for collection plans and lifecycle.""" + +from __future__ import annotations + +import os +import re +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.filesystem import get_file_facts +from data_designer.slurm.state.collection_records import CollectedOutputFile, CollectionResult, CollectionStatus +from data_designer.slurm.state.filesystem import ( + acquire_file_lock, + ensure_private_child_directory, + is_state_temporary_name, + open_verified_child_directory, + open_verified_directory, + open_verified_regular_file, + publish_immutable_text, + replace_text, +) +from data_designer.slurm.state.outputs import CollectionPlan +from data_designer.slurm.state.storage import StateStorage + +_COLLECTIONS_DIRECTORY = "collections" +_COLLECTION_LOCK = "collection.lock" +_WORKER_LOCK = "worker.lock" +_PLAN_FILENAME = "plan.json" +_STATUS_FILENAME = "status.json" +_RESULT_FILENAME = "collection-result.json" +_COLLECTION_PATTERN = re.compile(r"^collection-[0-9]{4,}$") +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class CollectionStorage: + """Persist collection state without adding collection concerns to run storage.""" + + def __init__(self, state_storage: StateStorage) -> None: + self._state = state_storage + self.collections_root = state_storage.run_root / _COLLECTIONS_DIRECTORY + + @contextmanager + def acquire_lock(self) -> Iterator[None]: + """Serialize collection preparation, refresh, and worker publication.""" + with self._state.open_run_directory() as run_descriptor: + with acquire_file_lock( + run_descriptor, + _COLLECTION_LOCK, + self._state.run_root / _COLLECTION_LOCK, + ): + yield + + @contextmanager + def acquire_worker_lock(self, collection_id: Identifier) -> Iterator[None]: + """Serialize bulk I/O owners without blocking collection metadata refresh.""" + with self._open_collection_directory(collection_id) as descriptor: + with acquire_file_lock( + descriptor, + _WORKER_LOCK, + self.get_collection_root(collection_id) / _WORKER_LOCK, + ): + yield + + def get_next_collection_id(self) -> Identifier: + """Return the next monotonic identity after validating existing directories.""" + names = self.list_collection_ids() + return f"collection-{len(names) + 1:04d}" + + def discard_incomplete_tail(self) -> None: + """Discard one trailing collection journal that predates submission.""" + try: + with self._open_collections_directory() as collections_descriptor: + collection_ids = _validated_collection_ids(tuple(os.listdir(collections_descriptor))) + if not collection_ids: + return + collection_id = collection_ids[-1] + with open_verified_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) as collection_descriptor: + if _record_exists(collection_descriptor, _STATUS_FILENAME): + return + _discard_prepared_files( + collection_descriptor, + self.get_collection_root(collection_id), + ) + os.rmdir(collection_id, dir_fd=collections_descriptor) + os.fsync(collections_descriptor) + except FileNotFoundError: + return + + def list_collection_ids(self) -> tuple[Identifier, ...]: + """List only a complete monotonic set of managed collection directories.""" + try: + with self._open_collections_directory() as descriptor: + names = tuple(os.listdir(descriptor)) + except FileNotFoundError: + return () + return _validated_collection_ids(names) + + def ensure_collection(self, collection_id: Identifier) -> None: + """Create one private collection state directory.""" + with self._state.open_run_directory() as run_descriptor: + ensure_private_child_directory(run_descriptor, _COLLECTIONS_DIRECTORY, self.collections_root) + with open_verified_child_directory( + run_descriptor, + _COLLECTIONS_DIRECTORY, + self.collections_root, + ) as collections_descriptor: + ensure_private_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) + + def get_collection_root(self, collection_id: Identifier) -> Path: + return self.collections_root / collection_id + + def get_plan_path(self, collection_id: Identifier) -> Path: + return self.get_collection_root(collection_id) / _PLAN_FILENAME + + def get_status_path(self, collection_id: Identifier) -> Path: + return self.get_collection_root(collection_id) / _STATUS_FILENAME + + def get_result_path(self, plan: CollectionPlan) -> Path: + return Path(plan.host_destination) / _RESULT_FILENAME + + def get_result_reference(self, plan: CollectionPlan, result: CollectionResult) -> ArtifactReference: + """Return the canonical host-view reference for a published result.""" + return ArtifactReference(path=self.get_result_path(plan).as_posix(), sha256=result.compute_sha256()) + + def publish_plan(self, plan: CollectionPlan) -> None: + self._require_run_id(plan.run_id) + with self._open_collection_directory(plan.collection_id) as descriptor: + publish_immutable_text( + descriptor, + _PLAN_FILENAME, + plan.serialize_json(), + self.get_plan_path(plan.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_plan(self, collection_id: Identifier) -> CollectionPlan: + with self._open_collection_directory(collection_id) as descriptor: + plan = self._state.read_record( + descriptor, + _PLAN_FILENAME, + self.get_plan_path(collection_id), + CollectionPlan, + ) + if plan.collection_id != collection_id or plan.run_id != self._state.run_id: + raise OSError("collection plan identity does not match its persisted location") + return plan + + def get_plan_reference(self, plan: CollectionPlan) -> ArtifactReference: + """Return the canonical persisted reference for an immutable collection plan.""" + return ArtifactReference(path=self.get_plan_path(plan.collection_id).as_posix(), sha256=plan.compute_sha256()) + + def publish_status(self, status: CollectionStatus) -> None: + self._require_run_id(status.run_id) + with self._open_collection_directory(status.collection_id) as descriptor: + publish_immutable_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_status_path(status.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_status(self, status: CollectionStatus) -> None: + self._require_run_id(status.run_id) + with self._open_collection_directory(status.collection_id) as descriptor: + replace_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_status_path(status.collection_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_status(self, collection_id: Identifier) -> CollectionStatus: + with self._open_collection_directory(collection_id) as descriptor: + status = self._state.read_record( + descriptor, + _STATUS_FILENAME, + self.get_status_path(collection_id), + CollectionStatus, + ) + if status.collection_id != collection_id or status.run_id != self._state.run_id: + raise OSError("collection status identity does not match its persisted location") + return status + + def read_result(self, plan: CollectionPlan) -> CollectionResult: + """Read a result through the immutable host destination.""" + destination = Path(plan.host_destination) + with open_verified_directory(destination, require_private=True) as descriptor: + return self._state.read_record( + descriptor, + _RESULT_FILENAME, + destination / _RESULT_FILENAME, + CollectionResult, + ) + + def verify_result_files( + self, + plan: CollectionPlan, + result: CollectionResult, + *, + verify_digests: bool = True, + ) -> None: + """Verify exact inventory with bounded metadata or full output digests.""" + destination = Path(plan.host_destination) + expected_names = tuple(output.relative_path for output in result.files) + (_RESULT_FILENAME,) + if any("/" in name for name in expected_names): + raise OSError("collection result inventory must contain only direct child files") + with open_verified_directory(destination, require_private=True) as descriptor: + if set(os.listdir(descriptor)) != set(expected_names): + raise OSError("published collection inventory does not match its result manifest") + for output in result.files: + if verify_digests: + with open_verified_regular_file( + descriptor, + output.relative_path, + destination / output.relative_path, + expected_size=output.byte_size, + expected_sha256=output.sha256, + require_private=False, + ): + pass + else: + _verify_output_metadata(descriptor, output, destination) + result_bytes = result.serialize_json().encode("utf-8") + with open_verified_regular_file( + descriptor, + _RESULT_FILENAME, + destination / _RESULT_FILENAME, + expected_size=len(result_bytes), + expected_sha256=result.compute_sha256(), + ): + pass + + @contextmanager + def _open_collections_directory(self) -> Iterator[int]: + with self._state.open_run_directory() as run_descriptor: + with open_verified_child_directory( + run_descriptor, + _COLLECTIONS_DIRECTORY, + self.collections_root, + ) as collections_descriptor: + yield collections_descriptor + + @contextmanager + def _open_collection_directory(self, collection_id: Identifier) -> Iterator[int]: + with self._open_collections_directory() as collections_descriptor: + with open_verified_child_directory( + collections_descriptor, + collection_id, + self.get_collection_root(collection_id), + ) as descriptor: + yield descriptor + + def _require_run_id(self, run_id: Identifier) -> None: + if run_id != self._state.run_id: + raise OSError("collection record run identity does not match storage") + + +def _verify_output_metadata( + directory_descriptor: int, + output: CollectedOutputFile, + destination: Path, +) -> None: + before = os.stat(output.relative_path, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + not stat.S_ISREG(before.st_mode) + or before.st_nlink != 1 + or before.st_mode & 0o022 + or before.st_size != output.byte_size + or before.st_mtime_ns != output.modified_at_ns + or before.st_ctime_ns != output.changed_at_ns + ): + raise OSError(f"collected output {destination / output.relative_path} is not a safe regular file") + descriptor = os.open( + output.relative_path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + try: + after = os.fstat(descriptor) + rebound = os.stat(output.relative_path, dir_fd=directory_descriptor, follow_symlinks=False) + finally: + os.close(descriptor) + if get_file_facts(before) != get_file_facts(after) or get_file_facts(after) != get_file_facts(rebound): + raise OSError(f"collected output {destination / output.relative_path} changed during validation") + + +def _validated_collection_ids(names: tuple[str, ...]) -> tuple[Identifier, ...]: + if any(_COLLECTION_PATTERN.fullmatch(name) is None for name in names): + raise OSError("collection state contains an unowned directory") + ordered = tuple(sorted(names, key=lambda name: int(name.rsplit("-", maxsplit=1)[1]))) + expected = tuple(f"collection-{index:04d}" for index in range(1, len(ordered) + 1)) + if ordered != expected: + raise OSError("collection state identities are not a complete monotonic sequence") + return ordered + + +def _record_exists(directory_descriptor: int, name: str) -> bool: + try: + os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return False + return True + + +def _discard_prepared_files(directory_descriptor: int, display_path: Path) -> None: + names = tuple(os.listdir(directory_descriptor)) + if any(name != _PLAN_FILENAME and not is_state_temporary_name(name) for name in names): + raise OSError(f"incomplete collection journal {display_path} contains an unowned entry") + for name in names: + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"incomplete collection journal entry {display_path / name} is unsafe") + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +__all__ = ["CollectionStorage"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py new file mode 100644 index 000000000..e98146db2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-record validation for collection lifecycle and inputs.""" + +from __future__ import annotations + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan +from data_designer.slurm.state.scheduler import SchedulerState +from data_designer.slurm.state.validation import StateContractError + +_TERMINAL_COLLECTION_STATES = frozenset({CollectionState.SUCCEEDED, CollectionState.FAILED}) + + +def validate_collection_inputs( + resolved_plan: ResolvedSlurmRunPlan, + collection_plan: CollectionPlan, + candidates: tuple[CandidateOutputManifest, ...], +) -> tuple[CandidateOutputManifest, ...]: + """Require one compatible, nonempty candidate for every planned shard.""" + _require(len(candidates) == len(collection_plan.planned_shards), "collection candidate set is incomplete") + expected_shards = tuple(shard.shard_id for shard in collection_plan.planned_shards) + actual_shards = tuple(candidate.shard_id for candidate in candidates) + _require(actual_shards == expected_shards, "collection candidates do not match ordered planned shards") + _require(all(candidate.winner_eligible for candidate in candidates), "collection candidate is empty") + + _require( + all(candidate.provenance_digest == resolved_plan.compute_sha256() for candidate in candidates), + "collection candidate provenance does not match the resolved plan", + ) + + actual_records = sum(candidate.actual_records for candidate in candidates) + _require(collection_plan.num_partitions <= actual_records, "collection partitions exceed available records") + return candidates + + +def validate_collection_status_transition( + previous: CollectionStatus, + current: CollectionStatus, +) -> CollectionStatus: + """Validate immutable identity, monotonic revisions, and terminal evidence.""" + _require(current.collection_id == previous.collection_id, "collection identity cannot change") + _require(current.run_id == previous.run_id, "collection run identity cannot change") + _require(current.collection_plan == previous.collection_plan, "collection plan identity cannot change") + _require(current.staging_directory == previous.staging_directory, "collection staging identity cannot change") + _require(previous.state not in _TERMINAL_COLLECTION_STATES, "terminal collection status is immutable") + _require(current.revision == previous.revision + 1, "collection status revision must increase by one") + _require(current.updated_at >= previous.updated_at, "collection status timestamp cannot move backward") + if previous.scheduler is not None: + _require(current.scheduler == previous.scheduler, "collection scheduler identity cannot change") + if previous.scheduler_observation is not None and current.scheduler_observation is not None: + _require( + current.scheduler_observation.observed_at >= previous.scheduler_observation.observed_at, + "collection scheduler observation cannot move backward", + ) + return current + + +def validate_collection_result( + collection_plan: CollectionPlan, + result: CollectionResult, + *, + expected_records: int, +) -> CollectionResult: + """Bind a collected result to its immutable plan and exact row count.""" + _require(result.collection_id == collection_plan.collection_id, "collection result identity does not match") + _require(result.run_id == collection_plan.run_id, "collection result run identity does not match") + _require( + result.collection_plan_sha256 == collection_plan.compute_sha256(), + "collection result does not bind the collection plan", + ) + _require(result.completed_at >= collection_plan.created_at, "collection completion precedes plan creation") + _require(result.actual_records == expected_records, "collected result has the wrong row count") + _require(len(result.files) == collection_plan.num_partitions, "collected result has the wrong partition count") + expected_paths = tuple(f"part-{index:05d}.parquet" for index in range(collection_plan.num_partitions)) + _require( + tuple(output.relative_path for output in result.files) == expected_paths, + "collected result files do not match deterministic partition intent", + ) + return result + + +def derive_collection_state(observation_state: SchedulerState) -> CollectionState: + """Map normalized scheduler evidence to collection lifecycle state.""" + if observation_state is SchedulerState.PENDING: + return CollectionState.PENDING + if observation_state is SchedulerState.RUNNING: + return CollectionState.RUNNING + if observation_state is SchedulerState.ACCOUNTING_LAG: + return CollectionState.ACCOUNTING_LAG + if observation_state is SchedulerState.UNKNOWN: + return CollectionState.UNKNOWN + return CollectionState.FAILED + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) + + +__all__ = [ + "derive_collection_state", + "validate_collection_inputs", + "validate_collection_result", + "validate_collection_status_transition", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py new file mode 100644 index 000000000..0fc4cd39e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-only worker for one persisted collection plan.""" + +from __future__ import annotations + +import argparse +import os +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from pathlib import Path +from time import sleep + +from pydantic import TypeAdapter, ValidationError + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.contracts import Identifier, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.collection_filesystem import derive_collection_staging_directory, stage_collection +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_merge import CollectionMerger +from data_designer.slurm.state.collection_records import CollectionResult, CollectionState, CollectionStatus +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import ( + validate_collection_result, + validate_collection_status_transition, +) +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.outputs import CandidateOutputManifest, CollectionPlan +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SCHEDULER_BINDING_WAIT_ATTEMPTS = 300 + + +class SlurmCollectionWorker: + """Execute bulk collection only inside its recorded CPU Slurm job.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + collection_id: Identifier, + *, + environment: Mapping[str, str] | None = None, + ) -> None: + root, normalized_run_id, normalized_collection_id = _validate_location(workspace_root, run_id, collection_id) + self._state = StateStorage(root, normalized_run_id) + self._collections = CollectionStorage(self._state) + self._reader = StateReader(self._state, normalized_run_id) + self._inputs = CollectionInputResolver(self._state, self._reader) + self._destinations = CollectionDestinationResolver() + self._environment = dict(os.environ if environment is None else environment) + self._run_id = normalized_run_id + self._collection_id = normalized_collection_id + + def run(self, *, completed_at: datetime | None = None) -> CollectionResult: + """Validate, merge, and atomically publish one collection output.""" + try: + self._wait_for_scheduler_binding() + started_at = _utc_now() if completed_at is None else completed_at + with self._collections.acquire_worker_lock(self._collection_id): + plan, status = self._begin_collection(started_at) + if status.state is CollectionState.SUCCEEDED: + return self._load_valid_result(plan, status) + resolved_plan, candidates = self._inputs.resolve(plan) + self._destinations.validate_persisted(resolved_plan, plan) + recovered = self._load_optional_result(plan) + if recovered is not None: + return self._publish_success_status(plan, recovered, recovered.completed_at) + return self._execute_collection(plan, resolved_plan, candidates, completed_at) + except (StateConflictError, SlurmStateError): + raise + except (OSError, ValueError, lazy.pa.ArrowException) as error: + raise SlurmStateError(f"collection {self._collection_id!r} failed") from error + + def _begin_collection(self, started_at: datetime) -> tuple[CollectionPlan, CollectionStatus]: + with self._collections.acquire_lock(): + status = self._collections.read_status(self._collection_id) + plan = self._load_bound_plan(status) + self._validate_identity(plan.run_id, status.run_id) + self._require_scheduler_job(status) + if status.state is CollectionState.SUCCEEDED: + return plan, status + return plan, self._advance_status(status, CollectionState.RUNNING, started_at) + + def _execute_collection( + self, + plan: CollectionPlan, + resolved_plan: ResolvedSlurmRunPlan, + candidates: tuple[CandidateOutputManifest, ...], + completed_at: datetime | None, + ) -> CollectionResult: + try: + destination = self._destinations.validate_persisted(resolved_plan, plan) + with stage_collection( + Path(plan.host_destination), + derive_collection_staging_directory(plan), + Path(destination.authorized_root), + ) as staged: + merger = CollectionMerger(completed_at=completed_at) + result = merger.merge( + plan, + candidates, + staged, + ) + except (SlurmStateError, OSError, ValueError, lazy.pa.ArrowException) as error: + return self._recover_or_fail(plan, completed_at, error) + return self._publish_success_status(plan, result, result.completed_at) + + def _recover_or_fail( + self, + plan: CollectionPlan, + completed_at: datetime | None, + error: Exception, + ) -> CollectionResult: + try: + recovered = self._load_optional_result(plan) + except (SlurmStateError, OSError, ValueError) as recovery_error: + self._publish_failed_status(self._completion_time(completed_at)) + raise recovery_error from error + if recovered is not None: + return self._publish_success_status(plan, recovered, recovered.completed_at) + self._publish_failed_status(self._completion_time(completed_at)) + raise error + + def _load_optional_result(self, plan: CollectionPlan) -> CollectionResult | None: + try: + return self._load_valid_result(plan) + except FileNotFoundError: + return None + + def _load_valid_result( + self, + plan: CollectionPlan, + status: CollectionStatus | None = None, + ) -> CollectionResult: + result = self._collections.read_result(plan) + _, candidates = self._inputs.resolve(plan) + validated = validate_collection_result( + plan, + result, + expected_records=sum(candidate.actual_records for candidate in candidates), + ) + if status is not None and status.result != self._collections.get_result_reference(plan, validated): + raise StateCorruptionError("collection status does not bind its published result") + self._collections.verify_result_files(plan, validated) + return validated + + def _publish_success_status( + self, + plan: CollectionPlan, + result: CollectionResult, + updated_at: datetime, + ) -> CollectionResult: + with self._collections.acquire_lock(): + previous = self._collections.read_status(self._collection_id) + bound_plan = self._load_bound_plan(previous) + if bound_plan != plan: + raise StateCorruptionError("collection plan changed during worker execution") + self._require_scheduler_job(previous) + result_reference = self._collections.get_result_reference(plan, result) + if previous.state is CollectionState.SUCCEEDED: + if previous.result != result_reference: + raise StateCorruptionError("collection status does not bind its published result") + return result + current = CollectionStatus( + schema_version=1, + collection_id=previous.collection_id, + run_id=previous.run_id, + collection_plan=previous.collection_plan, + staging_directory=previous.staging_directory, + revision=previous.revision + 1, + updated_at=max(updated_at, previous.updated_at), + state=CollectionState.SUCCEEDED, + scheduler=previous.scheduler, + scheduler_observation=previous.scheduler_observation, + result=result_reference, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return result + + def _publish_failed_status(self, updated_at: datetime) -> None: + with self._collections.acquire_lock(): + previous = self._collections.read_status(self._collection_id) + self._load_bound_plan(previous) + self._require_scheduler_job(previous) + self._advance_status(previous, CollectionState.FAILED, updated_at) + + def _advance_status( + self, + previous: CollectionStatus, + state: CollectionState, + updated_at: datetime, + ) -> CollectionStatus: + current = _updated_status( + previous, + revision=previous.revision + 1, + updated_at=max(updated_at, previous.updated_at), + state=state, + ) + validate_collection_status_transition(previous, current) + self._collections.replace_status(current) + return current + + def _require_scheduler_job(self, status: CollectionStatus) -> None: + scheduler = status.scheduler + if type(scheduler) is not int: + raise StateConflictError("collection worker requires an ordinary Slurm job identity") + observed_job_id = self._environment.get("SLURM_JOB_ID") + if observed_job_id != str(scheduler): + raise StateConflictError("collection worker must run inside its recorded Slurm job") + + def _wait_for_scheduler_binding(self) -> None: + if self._environment.get("SLURM_JOB_ID") is None: + raise StateConflictError("collection worker must run inside its recorded Slurm job") + for attempt in range(_SCHEDULER_BINDING_WAIT_ATTEMPTS): + status = self._collections.read_status(self._collection_id) + if status.scheduler is not None: + self._require_scheduler_job(status) + return + if status.state is not CollectionState.PREPARED: + raise StateConflictError("collection worker requires an ordinary Slurm job identity") + if attempt + 1 < _SCHEDULER_BINDING_WAIT_ATTEMPTS: + sleep(1) + raise StateConflictError("collection scheduler identity was not published before allocation startup") + + def _validate_identity(self, plan_run_id: Identifier, status_run_id: Identifier) -> None: + if plan_run_id != self._run_id or status_run_id != self._run_id: + raise StateConflictError("collection records do not match the requested run") + + def _load_bound_plan(self, status: CollectionStatus) -> CollectionPlan: + plan = self._collections.read_plan(status.collection_id) + if status.collection_plan != self._collections.get_plan_reference(plan): + raise StateCorruptionError("collection status does not bind its persisted collection plan") + if status.staging_directory != derive_collection_staging_directory(plan): + raise StateCorruptionError("collection status does not bind its exact staging directory") + return plan + + def _completion_time(self, completed_at: datetime | None) -> datetime: + return _utc_now() if completed_at is None else completed_at + + +def _updated_status(previous: CollectionStatus, **updates: object) -> CollectionStatus: + payload = previous.model_dump(mode="python") + payload.update(updates) + return CollectionStatus.model_validate(payload) + + +def _validate_location( + workspace_root: str | Path, + run_id: Identifier, + collection_id: Identifier, +) -> tuple[Path, Identifier, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + normalized_collection_id = _IDENTIFIER_ADAPTER.validate_python(collection_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted collection worker location") from error + return Path(root), normalized_run_id, normalized_collection_id + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def main(argv: Sequence[str] | None = None) -> int: + """Run one allocation-local collection from explicit persisted identities.""" + parser = argparse.ArgumentParser() + parser.add_argument("--workspace-root", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--collection-id", required=True) + arguments = parser.parse_args(argv) + SlurmCollectionWorker( + arguments.workspace_root, + arguments.run_id, + arguments.collection_id, + ).run() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = ["SlurmCollectionWorker", "main"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py new file mode 100644 index 000000000..d06bfefcd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Profile-authorized host collection destinations.""" + +from __future__ import annotations + +import posixpath +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.contracts import is_path_below, paths_overlap, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.errors import StateConflictError +from data_designer.slurm.state.outputs import CollectionPlan + + +@dataclass(frozen=True, slots=True) +class CollectionDestination: + """One normalized host output directory and its authorized writable root.""" + + host_path: str + authorized_root: str + + +class CollectionDestinationResolver: + """Resolve the pinned output root through the selected writable mount map.""" + + def resolve( + self, + plan: ResolvedSlurmRunPlan, + requested_destination: str | Path | None = None, + ) -> CollectionDestination: + """Return one authorized host output path.""" + workspace_root = plan.selected_profile.profile.workspace_root + raw_destination = ( + posixpath.join(workspace_root, "runs", plan.run_id, "merge-output") + if requested_destination is None + else Path(requested_destination).as_posix() + ) + try: + host_path = validate_absolute_path(raw_destination) + except ValueError as error: + raise StateConflictError("collection destination must be a normalized absolute path") from error + _validate_destination_overlap(plan, host_path) + authorized = {workspace_root, *(mount.source for mount in plan.container_mounts if not mount.read_only)} + writable = tuple(root for root in authorized if host_path == root or is_path_below(host_path, root)) + if not writable: + raise StateConflictError("collection destination is not covered by a profile-authorized writable mount") + authorized_root = max(writable, key=len) + if host_path == authorized_root: + raise StateConflictError("collection destination must be below its writable mount source") + return CollectionDestination(host_path, authorized_root) + + def validate_persisted( + self, + resolved_plan: ResolvedSlurmRunPlan, + collection_plan: CollectionPlan, + ) -> CollectionDestination: + """Reauthorize persisted collection intent against its pinned run plan.""" + destination = self.resolve(resolved_plan, collection_plan.host_destination) + if collection_plan.run_id != resolved_plan.run_id: + raise StateConflictError("collection run identity does not match the resolved plan") + if collection_plan.num_partitions > resolved_plan.invocation.authored.num_records: + raise StateConflictError("collection partitions exceed requested records") + if collection_plan.host_destination != destination.host_path: + raise StateConflictError("collection host destination no longer matches resolved intent") + return destination + + +def _validate_destination_overlap(plan: ResolvedSlurmRunPlan, host_path: str) -> None: + workspace_root = plan.selected_profile.profile.workspace_root + reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) + if any(paths_overlap(host_path, path) for path in reserved): + raise StateConflictError("collection destination must not overlap package-managed workspace state") + managed_assets_path = plan.invocation.effective_input_bindings.managed_assets_path + assert managed_assets_path is not None + if paths_overlap(host_path, managed_assets_path): + raise StateConflictError("collection destination must not overlap managed assets") + runs_root = posixpath.join(workspace_root, "runs") + run_output_root = posixpath.join(runs_root, plan.run_id, "merge-output") + if paths_overlap(host_path, runs_root) and not ( + host_path == run_output_root or is_path_below(host_path, run_output_root) + ): + raise StateConflictError("collection destination must not overlap package-managed run state") + + +__all__ = ["CollectionDestination", "CollectionDestinationResolver"] 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 new file mode 100644 index 000000000..c3d5bb528 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical errors for persisted Slurm run state.""" + +from __future__ import annotations + + +class SlurmStateError(RuntimeError): + """Base error for state storage and publication.""" + + +class StateNotFoundError(SlurmStateError): + """Raised when requested persisted state does not exist.""" + + +class StateConflictError(SlurmStateError): + """Raised when a write conflicts with immutable or newer state.""" + + +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/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py new file mode 100644 index 000000000..927401001 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/execution.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + Identifier, + RecordRange, + ResumeWorkspace, + ShardId, +) +from data_designer.slurm.state.base import ( + SchedulerIdentity, + StateRecord, + validate_utc_timestamp, +) + + +class RunManifest(StateRecord): + """Identity and immutable authored/resolved inputs for a Slurm run.""" + + run_id: Identifier + created_at: datetime + authored_config: ArtifactReference + resolved_plan: ArtifactReference + shard_count: PositiveInt + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + +class ShardManifest(StateRecord): + """Stable shard identity and planner-owned input partition reference.""" + + run_id: Identifier + shard_id: ShardId + shard_index: NonNegativeInt + record_range: RecordRange + input_partition: ArtifactReference | None = None + resume_workspace: ResumeWorkspace + created_at: datetime + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + +class AttemptLifecycleState(str, Enum): + CREATED = "created" + SUBMITTED = "submitted" + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class AttemptTerminalClassification(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + TIMED_OUT = "timed_out" + NODE_FAILED = "node_failed" + PREEMPTED = "preempted" + REQUEUED = "requeued" + OUT_OF_MEMORY = "out_of_memory" + UNKNOWN = "unknown" + + +class AttemptManifest(StateRecord): + """Attempt identity, lifecycle, scheduler identity, and output reference.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + resolved_plan: ArtifactReference + state: AttemptLifecycleState + terminal_classification: AttemptTerminalClassification | None = None + scheduler: SchedulerIdentity | None = None + candidate_output: ArtifactReference | None = None + created_at: datetime + updated_at: datetime + + _timestamps_are_utc = field_validator("created_at", "updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_lifecycle(self) -> AttemptManifest: + if self.updated_at < self.created_at: + raise ValueError("updated_at must not precede created_at") + + if self.state is not AttemptLifecycleState.CREATED and self.scheduler is None: + raise ValueError("submitted and later attempts require scheduler identity") + + terminal = self.state in { + AttemptLifecycleState.SUCCEEDED, + AttemptLifecycleState.FAILED, + } + if terminal != (self.terminal_classification is not None): + raise ValueError("terminal classification must be present exactly for terminal attempts") + + if self.state is AttemptLifecycleState.SUCCEEDED: + if self.terminal_classification is not AttemptTerminalClassification.SUCCEEDED: + raise ValueError("successful attempts require a succeeded terminal classification") + if self.candidate_output is None: + raise ValueError("successful attempts require a candidate output reference") + elif self.terminal_classification is AttemptTerminalClassification.SUCCEEDED: + raise ValueError("failed attempts cannot have a succeeded terminal classification") + + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py new file mode 100644 index 000000000..9394889ae --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py @@ -0,0 +1,430 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound filesystem operations for persisted Slurm run state.""" + +from __future__ import annotations + +import hashlib +import os +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.filesystem import ( + PRIVATE_DIRECTORY_MODE, + acquire_restrictive_file_lock, + create_restrictive_temporary_file, + get_file_facts, + is_managed_temporary_name, +) +from data_designer.slurm.filesystem import ( + open_verified_child_directory as open_shared_child_directory, +) +from data_designer.slurm.filesystem import ( + open_verified_directory as open_shared_directory, +) + +_READ_RETRY_ATTEMPTS = 10 +_TEMPORARY_PREFIX = ".state." +_TEMPORARY_SUFFIX = ".tmp" + + +@contextmanager +def acquire_file_lock(directory_descriptor: int, name: str, display_path: Path) -> Iterator[None]: + """Acquire an exclusive package-owned lock relative to a verified directory.""" + with acquire_restrictive_file_lock( + directory_descriptor, + name, + display_path, + resource_name="state", + ): + yield + + +def ensure_private_child_directory(parent_descriptor: int, name: str, display_path: Path) -> None: + """Create or restrict one child of an already verified directory.""" + created = False + try: + os.mkdir(name, PRIVATE_DIRECTORY_MODE, dir_fd=parent_descriptor) + created = True + except FileExistsError: + pass + with open_verified_child_directory( + parent_descriptor, + name, + display_path, + require_private=False, + ) as descriptor: + os.fchmod(descriptor, PRIVATE_DIRECTORY_MODE) + os.fsync(descriptor) + if created: + os.fsync(parent_descriptor) + + +@contextmanager +def open_verified_directory(path: Path, *, require_private: bool = False) -> Iterator[int]: + """Open one real directory and reject path replacement during the open.""" + with open_shared_directory(path, resource_name="state", require_private=require_private) as descriptor: + yield descriptor + + +@contextmanager +def open_verified_child_directory( + parent_descriptor: int, + name: str, + display_path: Path, + *, + require_private: bool = True, +) -> Iterator[int]: + """Open one real child directory relative to a verified parent.""" + with open_shared_child_directory( + parent_descriptor, + name, + display_path, + resource_name="state", + require_private=require_private, + ) as descriptor: + yield descriptor + + +def read_regular_text( + directory_descriptor: int, + name: str, + display_path: Path, + *, + maximum_size: int, +) -> str: + """Read one restrictive regular UTF-8 file without following a symlink.""" + for _ in range(_READ_RETRY_ATTEMPTS): + descriptor: int | None = None + try: + before_open = _repair_interrupted_publication( + directory_descriptor, + name, + display_path, + os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False), + ) + if not _is_restrictive_regular(before_open): + raise OSError(f"state record {display_path} is not a restrictive regular file") + if before_open.st_size > maximum_size: + raise OSError(f"state record {display_path} exceeds its size limit") + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + after_open = os.fstat(descriptor) + if not _is_restrictive_regular(after_open): + raise OSError(f"state record {display_path} is not a restrictive regular file") + if get_file_facts(before_open) != get_file_facts(after_open): + continue + record_file = os.fdopen(descriptor, "r", encoding="utf-8", newline="") + descriptor = None + with record_file: + content = record_file.read(maximum_size + 1) + if len(content.encode("utf-8")) > maximum_size: + raise OSError(f"state record {display_path} exceeds its size limit") + after_read = os.fstat(record_file.fileno()) + if not _is_restrictive_regular(after_read): + raise OSError(f"state record {display_path} is not a restrictive regular file") + if get_file_facts(after_open) != get_file_facts(after_read): + continue + return content + finally: + if descriptor is not None: + os.close(descriptor) + raise OSError(f"state record {display_path} changed too often to read consistently") + + +@contextmanager +def open_verified_regular_file( + directory_descriptor: int, + name: str, + display_path: Path, + *, + expected_size: int, + expected_sha256: str, + require_private: bool = True, +) -> Iterator[int]: + """Hold one digest-verified regular file without following a symlink.""" + descriptor, before_open = _open_digest_verified_file( + directory_descriptor, + name, + display_path, + expected_size, + expected_sha256, + require_private, + ) + body_failed = False + try: + yield descriptor + except BaseException: + body_failed = True + raise + finally: + _close_verified_file( + directory_descriptor, + name, + display_path, + descriptor, + before_open, + require_private, + body_failed, + ) + + +def _open_digest_verified_file( + directory_descriptor: int, + name: str, + display_path: Path, + expected_size: int, + expected_sha256: str, + require_private: bool, +) -> tuple[int, os.stat_result]: + before_open = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not _is_safe_regular(before_open, require_private) or before_open.st_size != expected_size: + raise OSError(f"candidate file {display_path} is not the expected safe regular file") + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), + dir_fd=directory_descriptor, + ) + try: + _verify_opened_file(descriptor, before_open, display_path, expected_sha256, require_private) + return descriptor, before_open + except BaseException: + os.close(descriptor) + raise + + +def _verify_opened_file( + descriptor: int, + before_open: os.stat_result, + display_path: Path, + expected_sha256: str, + require_private: bool, +) -> None: + after_open = os.fstat(descriptor) + if not _is_safe_regular(after_open, require_private) or get_file_facts(before_open) != get_file_facts(after_open): + raise OSError(f"candidate file {display_path} changed while it was being opened") + digest = hashlib.sha256() + while chunk := os.read(descriptor, 1024 * 1024): + digest.update(chunk) + if digest.hexdigest() != expected_sha256: + raise OSError(f"candidate file {display_path} digest does not match its manifest") + os.lseek(descriptor, 0, os.SEEK_SET) + + +def _close_verified_file( + directory_descriptor: int, + name: str, + display_path: Path, + descriptor: int, + before_open: os.stat_result, + require_private: bool, + body_failed: bool, +) -> None: + try: + _verify_open_file_binding( + directory_descriptor, + name, + display_path, + descriptor, + before_open, + require_private, + ) + except OSError: + if not body_failed: + raise + finally: + os.close(descriptor) + + +def _verify_open_file_binding( + directory_descriptor: int, + name: str, + display_path: Path, + descriptor: int, + before_open: os.stat_result, + require_private: bool, +) -> None: + after_read = os.fstat(descriptor) + current = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if ( + not _is_safe_regular(after_read, require_private) + or not _is_safe_regular(current, require_private) + or get_file_facts(before_open) != get_file_facts(after_read) + or get_file_facts(after_read) != get_file_facts(current) + ): + raise OSError(f"candidate file {display_path} changed while it was being verified") + + +def publish_immutable_text( + directory_descriptor: int, + name: str, + content: str, + display_path: Path, + *, + maximum_size: int, +) -> bool: + """Publish one immutable record without replacing an existing name.""" + temporary_name = _write_temporary_text(directory_descriptor, content, maximum_size=maximum_size) + try: + try: + os.link( + temporary_name, + name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + follow_symlinks=False, + ) + except FileExistsError: + existing = read_regular_text( + directory_descriptor, + name, + display_path, + maximum_size=maximum_size, + ) + if existing == content: + sync_directory(directory_descriptor) + return False + raise FileExistsError(f"state record {display_path} already contains different bytes") from None + try: + os.unlink(temporary_name, dir_fd=directory_descriptor) + except FileNotFoundError: + pass + temporary_name = None + sync_directory(directory_descriptor) + return True + finally: + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=directory_descriptor) + sync_directory(directory_descriptor) + except OSError: + pass + + +def replace_text( + directory_descriptor: int, + name: str, + content: str, + display_path: Path, + *, + maximum_size: int, +) -> None: + """Atomically replace one existing restrictive regular record.""" + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"state record {display_path} is not a restrictive regular file") + temporary_name = _write_temporary_text(directory_descriptor, content, maximum_size=maximum_size) + try: + os.replace( + temporary_name, + name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + ) + temporary_name = None + sync_directory(directory_descriptor) + finally: + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=directory_descriptor) + except OSError: + pass + + +def sync_directory(directory_descriptor: int) -> None: + """Flush prior changes to an already verified directory.""" + os.fsync(directory_descriptor) + + +def create_state_temporary_file(directory_descriptor: int) -> tuple[int, str]: + """Create one restrictive file using the persisted-state temporary convention.""" + return create_restrictive_temporary_file( + directory_descriptor, + prefix=_TEMPORARY_PREFIX, + suffix=_TEMPORARY_SUFFIX, + ) + + +def is_state_temporary_name(name: str) -> bool: + """Return whether a name belongs to persisted-state temporary publication.""" + return is_managed_temporary_name( + name, + prefix=_TEMPORARY_PREFIX, + suffix=_TEMPORARY_SUFFIX, + ) + + +def _write_temporary_text(directory_descriptor: int, content: str, *, maximum_size: int) -> str: + encoded = content.encode("utf-8") + if len(encoded) > maximum_size: + raise OSError("state record exceeds its size limit") + descriptor: int | None = None + temporary_name: str | None = None + try: + descriptor, temporary_name = create_state_temporary_file(directory_descriptor) + output = os.fdopen(descriptor, "wb") + descriptor = None + with output: + output.write(encoded) + output.flush() + os.fsync(output.fileno()) + return temporary_name + except BaseException: + if descriptor is not None: + os.close(descriptor) + if temporary_name is not None: + try: + os.unlink(temporary_name, dir_fd=directory_descriptor) + except OSError: + pass + raise + + +def _repair_interrupted_publication( + directory_descriptor: int, + name: str, + display_path: Path, + status: os.stat_result, +) -> os.stat_result: + if _is_restrictive_regular(status): + return status + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077 or status.st_nlink != 2: + return status + + temporary_names: list[str] = [] + for candidate_name in os.listdir(directory_descriptor): + if not is_state_temporary_name(candidate_name): + continue + try: + candidate_status = os.stat(candidate_name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + continue + if (candidate_status.st_dev, candidate_status.st_ino) == (status.st_dev, status.st_ino): + temporary_names.append(candidate_name) + if len(temporary_names) != 1: + return os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + + try: + os.unlink(temporary_names[0], dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + except FileNotFoundError: + pass + repaired = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not _is_restrictive_regular(repaired): + raise OSError(f"state record {display_path} changed while recovering interrupted publication") + return repaired + + +def _is_restrictive_regular(status: os.stat_result) -> bool: + return stat.S_ISREG(status.st_mode) and status.st_nlink == 1 and not status.st_mode & 0o077 + + +def _is_safe_regular(status: os.stat_result, require_private: bool) -> bool: + forbidden_mode = 0o077 if require_private else 0o022 + return stat.S_ISREG(status.st_mode) and status.st_nlink == 1 and not status.st_mode & forbidden_mode diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py new file mode 100644 index 000000000..3d18e035d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -0,0 +1,425 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataset-workspace ownership and immutable shard-winner policy.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Literal + +from pydantic import ValidationError + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.contracts import AttemptId, ShardId +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier, VerifiedCandidateArtifacts +from data_designer.slurm.state.errors import ( + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_winner, +) + + +@dataclass(frozen=True, slots=True) +class _WinnerResolution: + winner: ShardWinner + candidate: CandidateOutputManifest + plan: ResolvedSlurmRunPlan + persisted_attempt: AttemptManifest + successful_attempt: AttemptManifest + already_published: bool + + +class WinnerFinalizer: + """Compose persisted state, artifact verification, and winner mutation.""" + + def __init__(self, storage: StateStorage, reader: StateReader) -> None: + self._storage = storage + self._reader = reader + self._artifacts = CandidateArtifactVerifier(storage.get_local_path) + + @contextmanager + def acquire_dataset_workspace( + self, + shard_id: ShardId, + attempt_id: AttemptId, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Iterator[Path]: + with self._storage.acquire_resume_lock(shard_id): + dataset_path = self._prepare_workspace_or_normalize(shard_id, attempt_id, resume_mode) + yield dataset_path + + def finalize_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None = None, + ) -> ShardWinner: + try: + with self._storage.acquire_resume_lock(shard_id): + return self._finalize_with_dataset_lease(shard_id, attempt_id, published_at, completed_at) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (PlanStateContractError, StateContractError, ValidationError) as error: + raise StateConflictError("attempt result is not eligible for winner publication") from error + except FileNotFoundError as error: + raise StateNotFoundError( + f"attempt {attempt_id!r} does not contain complete finalization records" + ) from error + except OSError as error: + raise SlurmStateError(f"cannot finalize winner for shard {shard_id!r}") from error + + def load_winner(self, shard_id: ShardId) -> ShardWinner: + try: + run, plan, shard = self._reader.load_shard_context(shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self.load_optional_winner(run, plan, shard, attempts) + if winner is None: + raise StateNotFoundError(f"shard {shard_id!r} has no winner") + return winner + except (StateCorruptionError, StateNotFoundError): + raise + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"shard {shard_id!r} has an invalid winner chain") from error + except (FileNotFoundError, OSError) as error: + raise StateCorruptionError(f"shard {shard_id!r} has unreadable winner state") from error + + def require_no_winner( + self, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempts: tuple[AttemptManifest, ...], + ) -> None: + if self.load_optional_winner(run, plan, shard, attempts) is not None: + raise StateConflictError(f"shard {shard.shard_id!r} already has an immutable winner") + + def load_optional_winner( + self, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempts: tuple[AttemptManifest, ...], + ) -> ShardWinner | None: + try: + winner = self._storage.read_winner(shard.shard_id) + except FileNotFoundError: + return None + except OSError as error: + raise StateCorruptionError("persisted winner record is unsafe or unreadable") from error + attempt = self._get_winner_attempt(attempts, winner) + try: + self._validate_persisted_winner(run, plan, shard, attempt, winner) + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError("persisted winner chain is invalid") from error + return winner + + def _finalize_with_dataset_lease( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None, + ) -> ShardWinner: + resolution = self._resolve_under_state_locks(shard_id, attempt_id, published_at, completed_at) + if resolution.already_published: + return resolution.winner + with ExitStack() as resources: + artifacts = self._open_candidate_artifacts(resources, resolution) + self._validate_artifact_metadata(resolution.candidate, artifacts) + return self._publish_verified_resolution( + shard_id, + attempt_id, + published_at, + completed_at, + resolution, + artifacts, + ) + + def _open_candidate_artifacts( + self, + resources: ExitStack, + resolution: _WinnerResolution, + ) -> VerifiedCandidateArtifacts: + try: + return resources.enter_context(self._artifacts.verify(resolution.candidate)) + except OSError as error: + raise StateContractError("candidate output files are unavailable or invalid") from error + + def _publish_verified_resolution( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None, + expected: _WinnerResolution, + artifacts: VerifiedCandidateArtifacts, + ) -> ShardWinner: + with self._storage.acquire_shard_lock(shard_id): + current = self._resolve_winner(shard_id, attempt_id, published_at, completed_at) + if current.already_published: + return current.winner + if ( + current.winner != expected.winner + or current.candidate != expected.candidate + or current.successful_attempt != expected.successful_attempt + ): + raise StateContractError("attempt finalization records changed during verification") + try: + artifacts.rebind() + except OSError as error: + raise StateContractError("candidate output paths changed during finalization") from error + if current.persisted_attempt != current.successful_attempt: + self._storage.replace_attempt(current.successful_attempt) + try: + self._storage.publish_winner(expected.winner) + except BaseException: + winner_was_published = self._winner_was_published(expected.winner) + if not winner_was_published: + self._storage.replace_attempt(current.persisted_attempt) + elif completed_at is not None: + return expected.winner + raise + return expected.winner + + def _resolve_under_state_locks( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None, + ) -> _WinnerResolution: + with self._storage.acquire_shard_lock(shard_id): + return self._resolve_winner(shard_id, attempt_id, published_at, completed_at) + + def _resolve_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + completed_at: datetime | None, + ) -> _WinnerResolution: + run, plan, shard = self._reader.load_shard_context(shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + persisted_attempt = self._reader.get_attempt(attempts, attempt_id) + existing = self.load_optional_winner(run, plan, shard, attempts) + client_result, candidate = self._load_finalization_records(persisted_attempt) + if existing is not None: + if existing.attempt_id != attempt_id: + raise StateConflictError(f"shard {shard_id!r} already has an immutable winner") + return _WinnerResolution( + existing, + candidate, + plan, + persisted_attempt, + persisted_attempt, + True, + ) + successful_attempt = self._successful_attempt(persisted_attempt, completed_at) + winner = self._build_winner(run, shard, successful_attempt, client_result, published_at) + self._validate_finalization_chain( + run, + plan, + shard, + successful_attempt, + client_result, + candidate, + winner, + ) + return _WinnerResolution( + winner, + candidate, + plan, + persisted_attempt, + successful_attempt, + False, + ) + + def _winner_was_published(self, expected: ShardWinner) -> bool: + try: + persisted = self._storage.read_winner(expected.shard_id) + except FileNotFoundError: + return False + if persisted != expected: + raise StateCorruptionError("persisted winner changed during finalization") + return True + + @staticmethod + def _successful_attempt(attempt: AttemptManifest, completed_at: datetime | None) -> AttemptManifest: + if completed_at is None or attempt.state is AttemptLifecycleState.SUCCEEDED: + return attempt + successful = attempt.model_copy( + update={ + "state": AttemptLifecycleState.SUCCEEDED, + "terminal_classification": AttemptTerminalClassification.SUCCEEDED, + "updated_at": max(completed_at, attempt.updated_at), + } + ) + validate_attempt_transition(attempt, successful) + return successful + + def _prepare_dataset_workspace( + self, + shard_id: ShardId, + attempt_id: AttemptId, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Path: + with self._storage.acquire_shard_lock(shard_id): + run, plan, shard = self._reader.load_shard_context(shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + attempt = self._reader.get_attempt(attempts, attempt_id) + self.require_no_winner(run, plan, shard, attempts) + self._validate_workspace_mode(plan, attempt, resume_mode) + if resume_mode == "if_possible": + return self._storage.get_local_path(shard.resume_workspace.path) + dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, resume_mode) + expected_path = ( + self._storage.get_local_path(shard.resume_workspace.path) + if resume_mode == "always" + else self._storage.get_local_path( + self._storage.logical_run_root / "shards" / shard_id / "attempts" / attempt_id / "dataset" + ) + ) + if dataset_path != expected_path: + raise StateContractError("dataset workspace path does not match the resolved plan") + return dataset_path + + def _prepare_workspace_or_normalize( + self, + shard_id: ShardId, + attempt_id: AttemptId, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Path: + try: + return self._prepare_dataset_workspace(shard_id, attempt_id, resume_mode) + except (StateConflictError, StateCorruptionError, StateNotFoundError, SlurmStateError): + raise + except (PlanStateContractError, StateContractError) as error: + raise StateConflictError("attempt cannot acquire the requested dataset workspace") from error + except OSError as error: + raise SlurmStateError(f"cannot prepare dataset workspace for attempt {attempt_id!r}") from error + + def _load_finalization_records( + self, + attempt: AttemptManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + try: + return self._storage.read_finalization_records(attempt.shard_id, attempt.attempt_id) + except FileNotFoundError: + raise + except OSError as error: + raise StateCorruptionError("attempt finalization records are unsafe or unreadable") from error + + def _validate_persisted_winner( + self, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempt: AttemptManifest, + winner: ShardWinner, + ) -> None: + try: + client_result, candidate = self._load_finalization_records(attempt) + except FileNotFoundError as error: + raise StateCorruptionError("persisted winner is missing its finalization records") from error + self._validate_finalization_chain(run, plan, shard, attempt, client_result, candidate, winner) + + @staticmethod + def _build_winner( + run: RunManifest, + shard: ShardManifest, + attempt: AttemptManifest, + client_result: ClientResult, + published_at: datetime, + ) -> ShardWinner: + reference = client_result.candidate_output_manifest + if reference is None: + raise StateContractError("nonempty client result has no candidate reference") + return ShardWinner( + schema_version=1, + run_id=run.run_id, + shard_id=shard.shard_id, + attempt_id=attempt.attempt_id, + attempt_ordinal=attempt.attempt_ordinal, + candidate_manifest=reference, + published_at=published_at, + ) + + @staticmethod + def _validate_finalization_chain( + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + winner: ShardWinner, + ) -> None: + PersistedPlanStateValidator(plan).validate_finalization_chain( + plan.shards[shard.shard_index], + attempt, + client_result, + candidate, + winner, + ) + validate_shard_winner(run, shard, attempt, candidate, winner) + + @staticmethod + def _validate_artifact_metadata( + candidate: CandidateOutputManifest, + artifacts: VerifiedCandidateArtifacts, + ) -> None: + declared_counts = tuple(output_file.record_count for output_file in candidate.files) + if artifacts.record_counts != declared_counts or sum(artifacts.record_counts) != candidate.actual_records: + raise StateContractError("candidate Parquet row counts do not match the output manifest") + if artifacts.dataset_schema_digest != candidate.dataset_schema_digest: + raise StateContractError("candidate Parquet schema does not match the output manifest") + + @staticmethod + def _validate_workspace_mode( + plan: ResolvedSlurmRunPlan, + attempt: AttemptManifest, + resume_mode: Literal["never", "always", "if_possible"], + ) -> None: + if attempt.state not in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + }: + raise StateContractError("dataset workspace requires an active submitted attempt") + requested = plan.invocation.authored.resume + if requested != "if_possible" and resume_mode != requested: + raise StateContractError("resume mode does not match the resolved plan") + + def _get_winner_attempt( + self, + attempts: tuple[AttemptManifest, ...], + winner: ShardWinner, + ) -> AttemptManifest: + try: + return self._reader.get_attempt(attempts, winner.attempt_id) + except StateNotFoundError as error: + raise StateCorruptionError("persisted winner references an unknown attempt") from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization_validation.py new file mode 100644 index 000000000..e47cc1f57 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization_validation.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Semantic validation for complete Slurm candidate chains.""" + +from __future__ import annotations + +import posixpath + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, +) +from data_designer.slurm.state.outputs import CANDIDATE_OUTPUT_FORMAT, CandidateOutputManifest, ShardWinner + + +class FinalizationContractError(ValueError): + """Raised when a candidate chain disagrees with its resolved plan.""" + + +class FinalizationChainValidator: + """Validate one successful attempt through immutable winner publication.""" + + def __init__(self, plan: ResolvedSlurmRunPlan) -> None: + self._plan = plan + + def validate( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + winner: ShardWinner, + ) -> ShardWinner: + """Validate the identities, outputs, references, and chronology of one chain.""" + self._validate_winner_outcome(attempt) + self.validate_attempt_result(planned_shard, attempt, client_result, candidate) + self._validate_winner_identity(attempt, winner) + self._validate_winner_reference(attempt, client_result, winner) + self._validate_winner_timestamp(attempt, client_result, winner) + return winner + + def validate_attempt_result( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + """Validate producer-owned records before immutable publication.""" + self._validate_result_outcomes(attempt, client_result, candidate) + self._validate_result_identities(planned_shard, attempt, client_result, candidate) + self._validate_output_contract(planned_shard, client_result, candidate) + self._validate_dataset_paths(planned_shard, attempt, client_result, candidate) + self._validate_result_reference(planned_shard, attempt, client_result, candidate) + self._validate_attempt_reference(attempt, client_result) + self._validate_result_timestamps(attempt, client_result, candidate) + + @staticmethod + def _validate_result_outcomes( + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + _require( + attempt.state in {AttemptLifecycleState.RUNNING, AttemptLifecycleState.SUCCEEDED}, + "attempt results require a running or successful attempt", + ) + _require( + client_result.outcome in {ClientOutcome.COMPLETE, ClientOutcome.PARTIAL} + and bool(client_result.actual_records), + "only nonempty client results may be finalized", + ) + + @staticmethod + def _validate_winner_outcome(attempt: AttemptManifest) -> None: + _require(attempt.state is AttemptLifecycleState.SUCCEEDED, "only successful attempts may be finalized") + _require( + attempt.terminal_classification is AttemptTerminalClassification.SUCCEEDED, + "only successfully classified attempts may be finalized", + ) + + def _validate_result_identities( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + for record_name, run_id in ( + ("client result", client_result.run_id), + ("candidate", candidate.run_id), + ): + _require(run_id == self._plan.run_id, f"{record_name} run_id does not match the resolved plan") + for record_name, shard_id in ( + ("client result", client_result.shard_id), + ("candidate", candidate.shard_id), + ): + _require(shard_id == planned_shard.shard_id, f"{record_name} shard_id does not match the planned shard") + for record_name, attempt_id in ( + ("client result", client_result.attempt_id), + ("candidate", candidate.attempt_id), + ): + _require(attempt_id == attempt.attempt_id, f"{record_name} attempt_id does not match the attempt") + _require( + candidate.attempt_ordinal == attempt.attempt_ordinal, + "candidate and attempt ordinals must match", + ) + + def _validate_winner_identity(self, attempt: AttemptManifest, winner: ShardWinner) -> None: + _require(winner.run_id == self._plan.run_id, "winner run_id does not match the resolved plan") + _require(winner.shard_id == attempt.shard_id, "winner shard_id does not match the attempt") + _require(winner.attempt_id == attempt.attempt_id, "winner attempt_id does not match the attempt") + _require(winner.attempt_ordinal == attempt.attempt_ordinal, "winner and attempt ordinals must match") + + def _validate_output_contract( + self, + planned_shard: PlannedShard, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + _require( + client_result.requested_records == candidate.requested_records == planned_shard.requested_records, + "client and candidate requested records must match the planned shard", + ) + _require( + client_result.actual_records == candidate.actual_records, + "client and candidate actual records must match", + ) + _require(candidate.winner_eligible, "only nonempty candidate outputs may be finalized") + expected_suffix = f".{CANDIDATE_OUTPUT_FORMAT}" + _require( + all(output_file.relative_path.endswith(expected_suffix) for output_file in candidate.files), + "candidate output file extensions must use the attempt-local Parquet format", + ) + _require( + all(output_file.byte_size > 0 for output_file in candidate.files if output_file.record_count > 0), + "non-empty candidate output files must contain bytes", + ) + _require( + candidate.provenance_digest == self._plan.compute_sha256(), + "candidate provenance digest does not match the resolved plan", + ) + _require( + client_result.requested_resume_mode == self._plan.invocation.authored.resume, + "client requested resume mode does not match the resolved plan", + ) + + def _validate_dataset_paths( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + expected_path = self._get_expected_dataset_path(planned_shard, attempt, client_result) + _require(client_result.dataset_path == expected_path, "client dataset path does not match planned intent") + _require(candidate.dataset_path == expected_path, "candidate dataset path does not match planned intent") + + @staticmethod + def _validate_result_reference( + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + expected_path = posixpath.join( + posixpath.dirname(planned_shard.resume_workspace.path), + "attempts", + attempt.attempt_id, + "output-manifest.json", + ) + reference = client_result.candidate_output_manifest + _require(reference is not None, "nonempty client result has no candidate manifest reference") + _require(reference.path == expected_path, "candidate manifest path does not match planned intent") + _require(reference.sha256 == candidate.compute_sha256(), "client candidate digest does not match the manifest") + + @staticmethod + def _validate_winner_reference( + attempt: AttemptManifest, + client_result: ClientResult, + winner: ShardWinner, + ) -> None: + reference = client_result.candidate_output_manifest + _require(reference is not None, "nonempty client result has no candidate manifest reference") + _require(attempt.candidate_output == reference, "attempt candidate reference does not match client result") + _require(winner.candidate_manifest == reference, "winner candidate reference does not match client result") + + @staticmethod + def _validate_attempt_reference(attempt: AttemptManifest, client_result: ClientResult) -> None: + reference = client_result.candidate_output_manifest + _require(reference is not None, "nonempty client result has no candidate manifest reference") + _require( + attempt.candidate_output is None or attempt.candidate_output == reference, + "attempt candidate reference does not match client result", + ) + + @staticmethod + def _validate_result_timestamps( + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + _require(candidate.created_at >= attempt.created_at, "candidate creation cannot precede attempt creation") + _require( + client_result.completed_at >= candidate.created_at, "client completion cannot precede candidate creation" + ) + + @staticmethod + def _validate_winner_timestamp( + attempt: AttemptManifest, + client_result: ClientResult, + winner: ShardWinner, + ) -> None: + _require( + attempt.updated_at >= client_result.completed_at, + "attempt completion cannot precede client completion", + ) + _require( + winner.published_at >= attempt.updated_at, + "winner publication cannot precede attempt completion", + ) + + @staticmethod + def _get_expected_dataset_path( + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + ) -> str: + if client_result.effective_resume_mode == "always": + return planned_shard.resume_workspace.path + return posixpath.join( + posixpath.dirname(planned_shard.resume_workspace.path), + "attempts", + attempt.attempt_id, + "dataset", + ) + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise FinalizationContractError(message) + + +__all__ = ["FinalizationChainValidator", "FinalizationContractError"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py new file mode 100644 index 000000000..a3102658e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observation.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize active and accounting evidence into scheduler observations.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta +from typing import Callable, Protocol, TypeVar + +from data_designer.slurm.state.base import SchedulerJobIdentity, validate_utc_timestamp +from data_designer.slurm.state.errors import SlurmStateError +from data_designer.slurm.state.scheduler import ( + SchedulerObservation, + SchedulerState, + is_scheduler_terminal_state, +) +from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition + +_ACCOUNTING_LAG_WINDOW = timedelta(minutes=5) +_PREEMPTION_REQUEUE_WINDOW = timedelta(minutes=5) +_SchedulerRecordT = TypeVar("_SchedulerRecordT", bound=object) +_SchedulerQueryError = OSError | RuntimeError | ValueError + + +class SchedulerQueueRecord(Protocol): + """Normalized active-queue record consumed by reconciliation.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + + +class SchedulerAccountingRecord(Protocol): + """Normalized accounting record consumed by reconciliation.""" + + job_identity: SchedulerJobIdentity + state: SchedulerState + + +class SchedulerObservationClient(Protocol): + """Query normalized active and accounting scheduler records.""" + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SchedulerQueueRecord, ...]: + """Return active queue records for the requested identities.""" + ... + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SchedulerAccountingRecord, ...]: + """Return accounting records for the requested identities.""" + ... + + +class SchedulerObservationCollector: + """Apply terminal-accounting precedence and bounded lag semantics.""" + + def __init__(self, client: SchedulerObservationClient) -> None: + self._client = client + + def collect( + self, + selectors: Sequence[SchedulerJobIdentity], + *, + observed_at: datetime, + previous: Mapping[SchedulerJobIdentity, SchedulerObservation | None] | None = None, + observation_floors: Mapping[SchedulerJobIdentity, datetime] | None = None, + ) -> tuple[SchedulerObservation, ...]: + """Return one deterministic observation for every requested identity.""" + validate_utc_timestamp(observed_at) + requested = tuple(dict.fromkeys(selectors)) + if not requested: + return () + prior = {} if previous is None else previous + observation_times = _resolve_observation_times(requested, observed_at, observation_floors) + queue, accounting = self._query_scheduler(requested) + queue_by_identity = self._index_records(queue, requested, source="active queue") + accounting_by_identity = self._index_records(accounting, requested, source="accounting") + return tuple( + self._resolve_observation( + identity, + observation_times[identity], + queue_by_identity.get(identity), + accounting_by_identity.get(identity), + prior.get(identity), + ) + for identity in requested + ) + + def _query_scheduler( + self, + selectors: tuple[SchedulerJobIdentity, ...], + ) -> tuple[tuple[SchedulerQueueRecord, ...], tuple[SchedulerAccountingRecord, ...]]: + queue, queue_error = _capture_scheduler_query(self._client.query_queue, selectors) + accounting, accounting_error = _capture_scheduler_query(self._client.query_accounting, selectors) + error = queue_error if queue_error is not None else accounting_error + if error is not None: + raise SlurmStateError("cannot query normalized scheduler observations") from error + return queue, accounting + + @staticmethod + def _index_records( + records: Sequence[SchedulerQueueRecord | SchedulerAccountingRecord], + requested: tuple[SchedulerJobIdentity, ...], + *, + source: str, + ) -> dict[SchedulerJobIdentity, SchedulerState]: + expected = set(requested) + indexed: dict[SchedulerJobIdentity, SchedulerState] = {} + for record in records: + identity = record.job_identity + if identity not in expected: + raise SlurmStateError(f"{source} returned an unrequested scheduler identity") + if identity in indexed: + raise SlurmStateError(f"{source} returned a duplicate scheduler identity") + if not isinstance(record.state, SchedulerState): + raise SlurmStateError(f"{source} returned an invalid normalized scheduler state") + indexed[identity] = record.state + return indexed + + @staticmethod + def _resolve_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + queue_state: SchedulerState | None, + accounting_state: SchedulerState | None, + previous: SchedulerObservation | None, + ) -> SchedulerObservation: + state = _select_observed_state(queue_state, accounting_state) + if ( + previous is not None + and is_scheduler_terminal_state(previous.state) + and (accounting_state is None or not is_scheduler_terminal_state(accounting_state)) + ): + state = previous.state + if state is None: + observation = _resolve_missing_observation(identity, observed_at, previous) + elif state is SchedulerState.PREEMPTED and queue_state is None: + observation = _resolve_preemption_observation(identity, observed_at, previous) + else: + observation = SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=state, + ) + if previous is not None: + try: + validate_scheduler_observation_transition(previous, observation) + except StateContractError as error: + raise SlurmStateError("scheduler observation violates persisted chronology") from error + return observation + + +def _resolve_observation_times( + requested: tuple[SchedulerJobIdentity, ...], + observed_at: datetime, + observation_floors: Mapping[SchedulerJobIdentity, datetime] | None, +) -> dict[SchedulerJobIdentity, datetime]: + floors = {} if observation_floors is None else observation_floors + if not set(floors).issubset(requested): + raise SlurmStateError("scheduler observation floors contain an unrequested identity") + resolved: dict[SchedulerJobIdentity, datetime] = {} + for identity in requested: + floor = floors.get(identity, observed_at) + validate_utc_timestamp(floor) + resolved[identity] = max(observed_at, floor) + return resolved + + +def _select_observed_state( + queue_state: SchedulerState | None, + accounting_state: SchedulerState | None, +) -> SchedulerState | None: + if accounting_state is not None and is_scheduler_terminal_state(accounting_state): + return accounting_state + if queue_state is not None: + return queue_state + return accounting_state + + +def _resolve_preemption_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + previous: SchedulerObservation | None, +) -> SchedulerObservation: + deadline = ( + previous.reconciliation_deadline + if previous is not None + and previous.state is SchedulerState.PREEMPTED + and previous.reconciliation_deadline is not None + else observed_at + _PREEMPTION_REQUEUE_WINDOW + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.FAILED if observed_at > deadline else SchedulerState.PREEMPTED, + reconciliation_deadline=None if observed_at > deadline else deadline, + ) + + +def _resolve_missing_observation( + identity: SchedulerJobIdentity, + observed_at: datetime, + previous: SchedulerObservation | None, +) -> SchedulerObservation: + if ( + previous is not None + and previous.state is SchedulerState.PREEMPTED + and previous.reconciliation_deadline is not None + ): + return _resolve_preemption_observation(identity, observed_at, previous) + if previous is not None and previous.state is SchedulerState.ACCOUNTING_LAG: + deadline = previous.reconciliation_deadline + if deadline is None: + raise SlurmStateError("persisted accounting lag has no reconciliation deadline") + if observed_at > deadline: + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.UNKNOWN, + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.ACCOUNTING_LAG, + reconciliation_deadline=deadline, + ) + if previous is not None and ( + previous.state is SchedulerState.UNKNOWN or is_scheduler_terminal_state(previous.state) + ): + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=previous.state, + ) + return SchedulerObservation( + schema_version=1, + scheduler=identity, + observed_at=observed_at, + state=SchedulerState.ACCOUNTING_LAG, + reconciliation_deadline=observed_at + _ACCOUNTING_LAG_WINDOW, + ) + + +def _capture_scheduler_query( + query: Callable[[Sequence[SchedulerJobIdentity]], tuple[_SchedulerRecordT, ...]], + selectors: tuple[SchedulerJobIdentity, ...], +) -> tuple[tuple[_SchedulerRecordT, ...], _SchedulerQueryError | None]: + try: + return query(selectors), None + except (OSError, RuntimeError, ValueError) as error: + return (), error + + +__all__ = [ + "SchedulerAccountingRecord", + "SchedulerObservationClient", + "SchedulerObservationCollector", + "SchedulerQueueRecord", +] 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 new file mode 100644 index 000000000..c0f8a4f58 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process reconciliation of one persisted Slurm run.""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId, validate_absolute_path +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, + StateNotFoundError, +) +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.observation import SchedulerObservationClient, SchedulerObservationCollector +from data_designer.slurm.state.outputs import ShardWinner +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.reconciliation import reconcile_attempt_observation +from data_designer.slurm.state.scheduler import EffectiveAttemptState, SchedulerObservation +from data_designer.slurm.state.status import ( + AttemptStatus, + RunStatus, + ShardStatus, + derive_generation_state, + derive_run_state, + derive_shard_state, +) +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import StateContractError, validate_scheduler_observation_transition + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +@dataclass(frozen=True, slots=True) +class _ShardSnapshot: + run: RunManifest + plan: ResolvedSlurmRunPlan + shard: ShardManifest + attempts: tuple[AttemptManifest, ...] + + +@dataclass(frozen=True, slots=True) +class _ObservationBatch: + previous: dict[SchedulerIdentity, SchedulerObservation | None] + current: dict[SchedulerJobIdentity, SchedulerObservation] + + +class SlurmStateReconciler: + """Refresh persisted run status from normalized scheduler observations. + + Each refresh reconstructs state from the compute-visible workspace. No + controller memory participates in status, wait, or benchmark refreshes. + + Args: + workspace_root: Selected compute-visible workspace root. + run_id: Stable application-owned run identity. + scheduler: Client returning normalized active and accounting records. + """ + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: SchedulerObservationClient, + ) -> None: + normalized_root, normalized_run_id = _validate_location(workspace_root, run_id) + self._storage = StateStorage(normalized_root, normalized_run_id) + self._reader = StateReader(self._storage, normalized_run_id) + self._finalizer = WinnerFinalizer(self._storage, self._reader) + self._collector = SchedulerObservationCollector(scheduler) + self._run_id = normalized_run_id + + @property + def run_root(self) -> Path: + """Return the workspace-derived root for this run.""" + return self._storage.run_root + + def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: + """Persist current scheduler evidence and return complete run status. + + Raises: + SlurmStateError: If scheduler evidence cannot be queried or state + cannot be reconstructed safely. + """ + return self._observe(observed_at=observed_at, persist=True) + + def observe(self, *, observed_at: datetime | None = None) -> RunStatus: + """Return current scheduler status without persisting observations.""" + return self._observe(observed_at=observed_at, persist=False) + + def _observe(self, *, observed_at: datetime | None, persist: bool) -> RunStatus: + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + _validate_observed_at(timestamp) + run, plan, shards = self._reader.load_context() + if timestamp < run.created_at: + raise SlurmStateError("observation timestamp cannot precede run creation") + attempts_by_shard = self._reader.load_validated_attempts(run, plan, shards) + previous = self._load_previous_observations(attempts_by_shard) + selectors = tuple(previous.keys()) + current = self._collector.collect( + selectors, + observed_at=timestamp, + previous=previous, + observation_floors=self._load_observation_floors(plan, attempts_by_shard, previous), + ) + batch = _ObservationBatch( + previous=previous, + current={observation.scheduler: observation for observation in current}, + ) + shard_statuses = tuple( + self._refresh_shard( + _ShardSnapshot(run, plan, shard, attempts_by_shard[shard.shard_id]), + batch, + persist=persist, + ) + for shard in shards + ) + latest_evidence = max( + ( + attempt.readiness.updated_at + for shard in shard_statuses + for attempt in shard.attempts + if attempt.readiness + ), + default=timestamp, + ) + return self._compose_run_status(run, max(timestamp, latest_evidence), shard_statuses) + + def _compose_run_status( + self, + run: RunManifest, + observed_at: datetime, + shards: tuple[ShardStatus, ...], + ) -> RunStatus: + try: + return RunStatus( + run=run, + observed_at=observed_at, + shards=shards, + effective_state=derive_run_state(shards), + ) + except ValidationError as error: + raise StateCorruptionError(f"cannot reconcile run {self._run_id!r}") from error + + def _load_previous_observations( + self, + attempts_by_shard: dict[ShardId, tuple[AttemptManifest, ...]], + ) -> dict[SchedulerIdentity, SchedulerObservation | None]: + previous: dict[SchedulerIdentity, SchedulerObservation | None] = {} + for attempts in attempts_by_shard.values(): + for attempt in attempts: + if attempt.scheduler is not None: + previous[attempt.scheduler] = self._reader.load_optional_scheduler_observation(attempt) + return previous + + def _load_observation_floors( + self, + plan: ResolvedSlurmRunPlan, + attempts_by_shard: dict[ShardId, tuple[AttemptManifest, ...]], + previous: dict[SchedulerIdentity, SchedulerObservation | None], + ) -> dict[SchedulerIdentity, datetime]: + floors: dict[SchedulerIdentity, datetime] = {} + for attempts in attempts_by_shard.values(): + for attempt in attempts: + scheduler = attempt.scheduler + if scheduler is None: + continue + floor = attempt.updated_at + prior = previous[scheduler] + if prior is not None: + floor = max(floor, prior.observed_at) + readiness = self._reader.load_optional_readiness(plan, attempt) + if readiness is not None: + floor = max(floor, readiness.updated_at) + floors[scheduler] = floor + return floors + + def _refresh_shard( + self, + expected: _ShardSnapshot, + batch: _ObservationBatch, + *, + persist: bool, + ) -> ShardStatus: + try: + shard_lock = self._storage.acquire_shard_lock(expected.shard.shard_id) if persist else nullcontext() + with shard_lock: + current_run, current_plan, current_shard = self._reader.load_shard_context(expected.shard.shard_id) + attempts = self._reader.load_validated_shard_attempts(current_run, current_plan, current_shard) + self._require_unchanged_context( + expected, + _ShardSnapshot(current_run, current_plan, current_shard, attempts), + ) + winner = self._finalizer.load_optional_winner( + expected.run, + expected.plan, + expected.shard, + attempts, + ) + statuses = tuple( + self._build_attempt_status( + expected, + batch, + attempt, + winner, + persist=persist, + ) + for attempt in attempts + ) + self._validate_winner_scheduler_consistency(winner, statuses) + return ShardStatus( + shard=expected.shard, + attempts=statuses, + winner=winner, + effective_state=derive_shard_state(statuses, winner), + ) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (OSError, StateContractError, ValidationError) as error: + raise StateCorruptionError(f"cannot reconcile shard {expected.shard.shard_id!r}") from error + + def _build_attempt_status( + self, + snapshot: _ShardSnapshot, + batch: _ObservationBatch, + attempt: AttemptManifest, + winner: ShardWinner | None, + *, + persist: bool, + ) -> AttemptStatus: + readiness = self._reader.load_optional_readiness(snapshot.plan, attempt) + result = self._reader.load_optional_attempt_result(snapshot.plan, snapshot.shard, attempt) + scheduler = None + if attempt.scheduler is not None: + scheduler = batch.current[attempt.scheduler] + if persist: + self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) + effective_state = reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=max(scheduler.observed_at, readiness.updated_at) + if readiness is not None + else scheduler.observed_at, + ) + else: + if attempt.state is not AttemptLifecycleState.CREATED: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has no scheduler identity") + effective_state = EffectiveAttemptState.PENDING + if attempt.candidate_output is not None and result is None: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} is missing its generation results") + client_result, candidate = (None, None) if result is None else result + is_winner = winner is not None and winner.attempt_id == attempt.attempt_id + generation_state = derive_generation_state( + effective_state, + has_candidate=candidate is not None, + is_winner=is_winner, + ) + return AttemptStatus( + attempt=attempt, + readiness=readiness, + scheduler=scheduler, + client_result=client_result, + candidate_output=candidate, + effective_state=effective_state, + generation_state=generation_state, + is_winner=is_winner, + ) + + def _persist_observation( + self, + attempt: AttemptManifest, + expected_previous: SchedulerObservation | None, + current: SchedulerObservation, + ) -> None: + persisted = self._reader.load_optional_scheduler_observation(attempt) + if persisted != expected_previous: + raise StateConflictError("scheduler evidence changed during reconciliation; refresh again") + if persisted == current: + self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) + return + if persisted is None: + self._storage.publish_scheduler_observation(attempt.shard_id, attempt.attempt_id, current) + return + validate_scheduler_observation_transition(persisted, current) + self._storage.replace_scheduler_observation(attempt.shard_id, attempt.attempt_id, current) + + @staticmethod + def _require_unchanged_context( + expected: _ShardSnapshot, + current: _ShardSnapshot, + ) -> None: + if current != expected: + raise StateConflictError("persisted state changed during reconciliation; refresh again") + + @staticmethod + def _validate_winner_scheduler_consistency( + winner: ShardWinner | None, + statuses: tuple[AttemptStatus, ...], + ) -> None: + if winner is None: + 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 SchedulerStateConflictError("persisted winner conflicts with terminal scheduler evidence") + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + normalized_root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted run location") from error + return Path(normalized_root), normalized_run_id + + +def _validate_observed_at(observed_at: datetime) -> None: + if not isinstance(observed_at, datetime) or observed_at.tzinfo is None or observed_at.utcoffset() is None: + raise SlurmStateError("observation timestamp must be timezone-aware UTC") + if observed_at.utcoffset().total_seconds() != 0: + raise SlurmStateError("observation timestamp must be timezone-aware UTC") + + +__all__ = ["SlurmStateReconciler"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py new file mode 100644 index 000000000..f1232753e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + Identifier, + Sha256Digest, + ShardId, + validate_absolute_path, + validate_relative_path, +) +from data_designer.slurm.state.base import ( + StateRecord, + StateValue, + validate_utc_timestamp, +) + +CANDIDATE_OUTPUT_FORMAT: Literal["parquet"] = "parquet" +# Finalization retains one file and one parent-directory descriptor per entry. +MAXIMUM_CANDIDATE_OUTPUT_FILES: int = 64 + + +class CandidateOutcome(str, Enum): + COMPLETE = "complete" + PARTIAL = "partial" + EMPTY = "empty" + + +class CandidateOutputFile(StateValue): + """One immutable file contained in a candidate output.""" + + relative_path: str + sha256: Sha256Digest + byte_size: NonNegativeInt + record_count: NonNegativeInt + + _relative_path_is_safe = field_validator("relative_path")(validate_relative_path) + + +class CandidateOutputManifest(StateRecord): + """Attempt-local output that may become the immutable shard winner.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + created_at: datetime + dataset_path: str + requested_records: PositiveInt + actual_records: NonNegativeInt + outcome: CandidateOutcome + files: tuple[CandidateOutputFile, ...] = Field(max_length=MAXIMUM_CANDIDATE_OUTPUT_FILES) + dataset_schema_digest: Sha256Digest + provenance_digest: Sha256Digest + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + _dataset_path_is_safe = field_validator("dataset_path")(validate_absolute_path) + + @property + def winner_eligible(self) -> bool: + """Whether this nonempty candidate may be published as the shard winner.""" + return (self.outcome is CandidateOutcome.COMPLETE and self.actual_records == self.requested_records) or ( + self.outcome is CandidateOutcome.PARTIAL and 0 < self.actual_records < self.requested_records + ) + + @model_validator(mode="after") + def validate_output(self) -> CandidateOutputManifest: + if self.actual_records > self.requested_records: + raise ValueError("actual_records must not exceed requested_records") + + expected_outcome = ( + CandidateOutcome.EMPTY + if self.actual_records == 0 + else CandidateOutcome.COMPLETE + if self.actual_records == self.requested_records + else CandidateOutcome.PARTIAL + ) + if self.outcome is not expected_outcome: + raise ValueError(f"outcome must be {expected_outcome.value!r} for this record count") + + relative_paths = [output_file.relative_path for output_file in self.files] + if len(relative_paths) != len(set(relative_paths)): + raise ValueError("candidate output file paths must be unique") + if self.actual_records > 0 and not self.files: + raise ValueError("non-empty candidate outputs require at least one file") + if sum(output_file.record_count for output_file in self.files) != self.actual_records: + raise ValueError("candidate output file record counts must equal actual_records") + return self + + +class ShardWinner(StateRecord): + """Immutable pointer selecting exactly one candidate for a shard.""" + + run_id: Identifier + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + candidate_manifest: ArtifactReference + published_at: datetime + + _published_at_is_utc = field_validator("published_at")(validate_utc_timestamp) + + +class CollectionShard(StateValue): + """Winner manifest selected for one shard in a collection plan.""" + + shard_id: ShardId + winner_manifest: ArtifactReference + + +class RetryShard(StateValue): + """One failed planned shard selected for a new attempt.""" + + shard_id: ShardId + attempt_id: AttemptId + attempt_ordinal: PositiveInt + array_task_index: NonNegativeInt + + +class RetryPlan(StateRecord): + """Immutable failed-shard selection used for one retry submission.""" + + retry_id: Identifier + run_id: Identifier + created_at: datetime + resolved_plan: ArtifactReference + planned_shards: tuple[RetryShard, ...] = Field(min_length=1) + effective_resume_mode: Literal["never", "always"] + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + + @property + def submission_job_name(self) -> Identifier: + """Return the immutable scheduler lookup key for this retry.""" + return f"dd-retry-{self.compute_sha256()[:32]}" + + @model_validator(mode="after") + def validate_shards(self) -> RetryPlan: + shard_ids = tuple(shard.shard_id for shard in self.planned_shards) + task_indices = tuple(shard.array_task_index for shard in self.planned_shards) + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("retry shard IDs must be unique") + if len(task_indices) != len(set(task_indices)): + raise ValueError("retry array-task indices must be unique") + if task_indices != tuple(sorted(task_indices)): + raise ValueError("retry shards must be ordered by array-task index") + return self + + +class CollectionPlan(StateRecord): + """Immutable inputs and destinations for deterministic collection.""" + + collection_id: Identifier + run_id: Identifier + created_at: datetime + resolved_plan: ArtifactReference + planned_shards: tuple[CollectionShard, ...] = Field(min_length=1) + python_executable: str + package_version: str | None = Field(default=None, min_length=1, exclude_if=lambda value: value is None) + runtime_bundle: ArtifactReference | None = Field(default=None, exclude_if=lambda value: value is None) + host_destination: str + num_partitions: PositiveInt + overwrite: Literal[False] = False + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + _paths_are_safe = field_validator("python_executable", "host_destination")(validate_absolute_path) + + @property + def submission_job_name(self) -> Identifier: + """Return the immutable scheduler lookup key for this collection.""" + return f"dd-collect-{self.compute_sha256()[:32]}" + + @model_validator(mode="after") + def validate_shards(self) -> CollectionPlan: + if (self.package_version is None) != (self.runtime_bundle is None): + raise ValueError("collection package version and runtime bundle must be set together") + shard_ids = [shard.shard_id for shard in self.planned_shards] + winner_paths = [shard.winner_manifest.path for shard in self.planned_shards] + if len(shard_ids) != len(set(shard_ids)): + raise ValueError("collection shard IDs must be unique") + if len(winner_paths) != len(set(winner_paths)): + raise ValueError("collection winner manifest paths must be unique") + return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py new file mode 100644 index 000000000..1c1480641 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validation that binds persisted run state to one resolved Slurm plan.""" + +from __future__ import annotations + +import posixpath +from collections.abc import Mapping +from types import MappingProxyType + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.contracts import ArtifactReference, ShardId +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.finalization_validation import FinalizationChainValidator, FinalizationContractError +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.readiness import AttemptReadiness, ReadinessState +from data_designer.slurm.state.validation import StateContractError, validate_shard_manifest, validate_shard_set + + +class PlanStateContractError(ValueError): + """Raised when persisted state disagrees with its resolved plan.""" + + +class PersistedPlanStateValidator: + """Validate persisted records against one resolved plan.""" + + def __init__(self, plan: ResolvedSlurmRunPlan) -> None: + self._plan: ResolvedSlurmRunPlan = plan + run_root = posixpath.dirname(plan.authored_config.path) + self._plan_reference: ArtifactReference = ArtifactReference( + path=posixpath.join(run_root, "resolved-plan.json"), + sha256=plan.compute_sha256(), + ) + self._shards_by_id: Mapping[ShardId, PlannedShard] = MappingProxyType( + {planned_shard.shard_id: planned_shard for planned_shard in plan.shards} + ) + self._finalization_validator = FinalizationChainValidator(plan) + + @property + def plan(self) -> ResolvedSlurmRunPlan: + """Return the resolved plan used for validation.""" + return self._plan + + def validate_plan_shards( + self, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> tuple[ShardManifest, ...]: + """Validate the complete ordered state shard set against planned intent.""" + self._validate_run(run) + try: + validate_shard_set(run, shards) + except StateContractError as error: + raise PlanStateContractError(str(error)) from error + _require(len(shards) == len(self.plan.shards), "state shards must exactly match the planned shard count") + for planned, persisted in zip(self.plan.shards, shards, strict=True): + _require(planned.shard_id == persisted.shard_id, "state shard identity does not match planned order") + self.validate_plan_shard(run, persisted) + return shards + + def validate_plan_shard( + self, + run: RunManifest, + persisted: ShardManifest, + ) -> ShardManifest: + """Validate one persisted shard against its canonical planned shard.""" + self._validate_run(run) + try: + validate_shard_manifest(run, persisted) + except StateContractError as error: + raise PlanStateContractError(str(error)) from error + planned = self._get_planned_shard(persisted.shard_id) + _require(persisted.shard_index == planned.shard_index, "state shard index does not match planned order") + _require(persisted.record_range == planned.record_range, "state shard record range does not match the plan") + _require( + persisted.input_partition == planned.input_partition, "state shard input partition does not match the plan" + ) + _require( + persisted.resume_workspace == planned.resume_workspace, + "state shard resume workspace does not match the plan", + ) + return persisted + + def validate_initial_readiness( + self, + attempt: AttemptManifest, + readiness: AttemptReadiness, + ) -> AttemptReadiness: + """Anchor the first readiness snapshot to a fully validated planned attempt.""" + self.validate_readiness_snapshot(attempt, readiness) + _require(readiness.revision == 1, "initial readiness must use revision 1") + _require(readiness.state is ReadinessState.PENDING, "initial readiness must be pending") + return readiness + + def validate_readiness_snapshot( + self, + attempt: AttemptManifest, + readiness: AttemptReadiness, + ) -> AttemptReadiness: + """Bind any persisted readiness revision to its planned attempt and deployments.""" + planned_shard = self._get_planned_shard(attempt.shard_id) + self.validate_planned_attempt(planned_shard, attempt) + _require(readiness.run_id == attempt.run_id, "readiness run_id does not match the attempt") + _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match the attempt") + _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match the attempt") + _require(readiness.updated_at >= attempt.created_at, "readiness cannot precede attempt creation") + expected = tuple( + (deployment.deployment_id, deployment.authored.model_alias, deployment.topology.replica_count) + for deployment in self.plan.deployments + ) + actual = tuple( + (deployment.deployment_id, deployment.model_alias, deployment.expected_backends) + for deployment in readiness.deployments + ) + _require(actual == expected, "readiness deployments do not match the resolved plan") + return readiness + + def validate_planned_attempt( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + ) -> AttemptManifest: + """Validate attempt identity and scheduler task ownership against one shard.""" + _require(attempt.run_id == self.plan.run_id, "attempt run_id does not match the resolved plan") + self._validate_plan_reference(attempt.resolved_plan) + canonical_shard = self._get_planned_shard(attempt.shard_id) + _require(canonical_shard == planned_shard, "planned shard is not the canonical shard for the attempt") + _require( + attempt.attempt_id == f"attempt-{attempt.attempt_ordinal:04d}", "attempt ID does not match its ordinal" + ) + _require( + attempt.state is not AttemptLifecycleState.CREATED, + "planned attempts must be submitted before scheduler task validation", + ) + _require(attempt.scheduler is not None, "planned attempts require scheduler array-task identity") + _require( + attempt.scheduler.array_task_id == planned_shard.array_task_index, + "attempt scheduler array task does not match the planned shard", + ) + return attempt + + def validate_finalization_chain( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + winner: ShardWinner, + ) -> ShardWinner: + """Validate a complete semantic result through immutable winner publication.""" + self.validate_planned_attempt(planned_shard, attempt) + try: + return self._finalization_validator.validate(planned_shard, attempt, client_result, candidate, winner) + except FinalizationContractError as error: + raise PlanStateContractError(str(error)) from error + + def validate_attempt_result( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + """Validate producer-owned result records before immutable publication.""" + self._validate_attempt_result_contract(planned_shard, attempt, client_result, candidate) + + def validate_client_candidate( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> CandidateOutputManifest: + """Validate producer-owned records before terminal winner publication.""" + self._validate_attempt_result_contract(planned_shard, attempt, client_result, candidate) + return candidate + + def _validate_attempt_result_contract( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + self.validate_planned_attempt(planned_shard, attempt) + try: + self._finalization_validator.validate_attempt_result( + planned_shard, + attempt, + client_result, + candidate, + ) + except FinalizationContractError as error: + raise PlanStateContractError(str(error)) from error + + def _get_planned_shard(self, shard_id: ShardId) -> PlannedShard: + planned_shard = self._shards_by_id.get(shard_id) + if planned_shard is None: + raise PlanStateContractError("attempt shard_id does not identify a planned shard") + return planned_shard + + def _validate_run(self, run: RunManifest) -> None: + _require(run.run_id == self.plan.run_id, "run manifest identity does not match the resolved plan") + _require( + run.authored_config == self.plan.authored_config, "run authored config does not match the resolved plan" + ) + self._validate_plan_reference(run.resolved_plan) + _require(run.shard_count == len(self.plan.shards), "run shard count does not match the resolved plan") + + def _validate_plan_reference(self, reference: ArtifactReference) -> None: + _require( + reference.path == self._plan_reference.path, "resolved plan reference path does not match planned intent" + ) + _require( + reference.sha256 == self._plan_reference.sha256, + "resolved plan reference digest does not match plan bytes", + ) + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise PlanStateContractError(message) + + +__all__ = ["PersistedPlanStateValidator", "PlanStateContractError"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py new file mode 100644 index 000000000..eeb23e66c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated reads of persisted Slurm state.""" + +from __future__ import annotations + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.errors import StateCorruptionError, StateNotFoundError +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.outputs import CandidateOutputManifest +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError +from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import SchedulerObservation +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_manifest, + validate_attempt_set, + validate_shard_attempt_set, +) + + +class StateReader: + """Compose physical records into plan-validated state snapshots.""" + + def __init__(self, storage: StateStorage, run_id: Identifier) -> None: + self._storage = storage + self._run_id = run_id + + def load_run(self) -> RunManifest: + try: + run = self._storage.read_run() + if ( + run.run_id != self._run_id + or run.authored_config.path != self._storage.logical_authored_config_path.as_posix() + or run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() + ): + raise StateCorruptionError(f"run {self._run_id!r} manifest does not match its persisted location") + return run + except FileNotFoundError as error: + raise StateNotFoundError(f"run {self._run_id!r} is not initialized") from error + except StateCorruptionError: + raise + except OSError as error: + raise StateCorruptionError(f"cannot load persisted run {self._run_id!r}") from error + + def load_authored_config(self, run: RunManifest | None = None) -> DataDesignerSlurmConfig: + bound_run = self.load_run() if run is None else run + try: + authored_config = self._storage.read_authored_config() + except (FileNotFoundError, OSError) as error: + raise StateCorruptionError(f"run {self._run_id!r} has no valid authored config") from error + if ( + bound_run.authored_config.path != self._storage.logical_authored_config_path.as_posix() + or bound_run.authored_config.sha256 != authored_config.compute_sha256() + ): + raise StateCorruptionError(f"run {self._run_id!r} authored config does not match its manifest") + return authored_config + + def load_resolved_plan(self, run: RunManifest | None = None) -> ResolvedSlurmRunPlan: + bound_run = self.load_run() if run is None else run + try: + plan = self._storage.read_resolved_plan() + except (FileNotFoundError, OSError) as error: + raise StateCorruptionError(f"run {self._run_id!r} has no valid resolved plan") from error + if ( + bound_run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() + or bound_run.resolved_plan.sha256 != plan.compute_sha256() + ): + raise StateCorruptionError(f"run {self._run_id!r} resolved plan does not match its manifest") + if bound_run.authored_config != plan.authored_config: + raise StateCorruptionError(f"run {self._run_id!r} resolved plan does not bind its authored config") + self.load_authored_config(bound_run) + return plan + + def load_shards( + self, + run: RunManifest | None = None, + plan: ResolvedSlurmRunPlan | None = None, + ) -> tuple[ShardManifest, ...]: + bound_run = self.load_run() if run is None else run + bound_plan = self.load_resolved_plan(bound_run) if plan is None else plan + try: + shards = self._storage.read_shards(bound_run.shard_count) + PersistedPlanStateValidator(bound_plan).validate_plan_shards(bound_run, shards) + return shards + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"run {self._run_id!r} has invalid persisted shards") from error + except StateCorruptionError: + raise + except (FileNotFoundError, OSError) as error: + raise StateCorruptionError(f"run {self._run_id!r} has unreadable persisted shards") from error + + def load_context(self) -> tuple[RunManifest, ResolvedSlurmRunPlan, tuple[ShardManifest, ...]]: + run = self.load_run() + plan = self.load_resolved_plan(run) + shards = self.load_shards(run, plan) + return run, plan, shards + + def load_shard_context(self, shard_id: ShardId) -> tuple[RunManifest, ResolvedSlurmRunPlan, ShardManifest]: + """Load and validate only the immutable context needed by one shard.""" + run = self.load_run() + plan = self.load_resolved_plan(run) + try: + shard = self._storage.read_shard(shard_id) + PersistedPlanStateValidator(plan).validate_plan_shard(run, shard) + return run, plan, shard + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"shard {shard_id!r} is invalid") from error + except StateCorruptionError: + raise + except FileNotFoundError as error: + raise StateNotFoundError(f"shard {shard_id!r} is not persisted") from error + except OSError as error: + raise StateCorruptionError(f"shard {shard_id!r} is unreadable") from error + + def load_attempts( + self, + shard_id: ShardId, + context: tuple[RunManifest, ResolvedSlurmRunPlan, ShardManifest] | None = None, + ) -> tuple[AttemptManifest, ...]: + try: + run, plan, shard = self.load_shard_context(shard_id) if context is None else context + if shard.shard_id != shard_id: + raise StateContractError("shard context does not match the requested shard") + return self.load_validated_shard_attempts(run, plan, shard) + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"shard {shard_id!r} has invalid attempts") from error + except (StateCorruptionError, StateNotFoundError): + raise + except (FileNotFoundError, OSError) as error: + raise StateCorruptionError(f"cannot load attempts for shard {shard_id!r}") from error + + def load_attempt(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptManifest: + return self.get_attempt(self.load_attempts(shard_id), attempt_id) + + def load_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptReadiness: + context = self.load_shard_context(shard_id) + _, plan, _ = context + attempt = self.get_attempt(self.load_attempts(shard_id, context), attempt_id) + readiness = self.load_optional_readiness(plan, attempt) + if readiness is None: + raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") + return readiness + + def load_optional_readiness( + self, + plan: ResolvedSlurmRunPlan, + attempt: AttemptManifest, + ) -> AttemptReadiness | None: + """Load validated readiness when the runtime has published it.""" + try: + readiness = self._storage.read_readiness(attempt.shard_id, attempt.attempt_id) + PersistedPlanStateValidator(plan).validate_readiness_snapshot(attempt, readiness) + return readiness + except FileNotFoundError: + return None + except PlanStateContractError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has invalid persisted readiness") from error + except OSError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable readiness") from error + + def load_optional_scheduler_observation(self, attempt: AttemptManifest) -> SchedulerObservation | None: + """Load and identity-check the most recent scheduler observation.""" + try: + observation = self._storage.read_scheduler_observation(attempt.shard_id, attempt.attempt_id) + except FileNotFoundError: + return None + except OSError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable scheduler evidence") from error + if attempt.scheduler is None or observation.scheduler != attempt.scheduler: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has mismatched scheduler evidence") + if observation.observed_at < attempt.created_at: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has scheduler evidence before its creation") + return observation + + def load_optional_attempt_result( + self, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempt: AttemptManifest, + ) -> tuple[ClientResult, CandidateOutputManifest] | None: + """Load a committed producer result pair, if the attempt has one.""" + if attempt.candidate_output is None: + # The client also writes partial results to this directory. Only + # the attempt's candidate reference commits a result for state use. + return None + try: + client_result, candidate = self._storage.read_finalization_records(shard.shard_id, attempt.attempt_id) + except FileNotFoundError: + return None + except OSError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has unreadable generation results") from error + try: + PersistedPlanStateValidator(plan).validate_attempt_result( + plan.shards[shard.shard_index], + attempt, + client_result, + candidate, + ) + except PlanStateContractError as error: + raise StateCorruptionError(f"attempt {attempt.attempt_id!r} has invalid generation results") from error + return client_result, candidate + + def load_validated_attempts( + self, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shards: tuple[ShardManifest, ...], + ) -> dict[ShardId, tuple[AttemptManifest, ...]]: + try: + attempts_by_shard = {shard.shard_id: self._storage.read_attempts(shard.shard_id) for shard in shards} + all_attempts = tuple(attempt for shard in shards for attempt in attempts_by_shard[shard.shard_id]) + for shard in shards: + for attempt in attempts_by_shard[shard.shard_id]: + self.validate_attempt_against_plan(run, plan, shard, attempt) + validate_attempt_set(run, shards, all_attempts) + return attempts_by_shard + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"run {self._run_id!r} has invalid persisted attempts") from error + + def load_validated_shard_attempts( + self, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + ) -> tuple[AttemptManifest, ...]: + """Read and plan-validate attempts for one shard only.""" + try: + attempts = self._storage.read_attempts(shard.shard_id) + for attempt in attempts: + self.validate_attempt_against_plan(run, plan, shard, attempt) + validate_shard_attempt_set(run, shard, attempts) + return attempts + except (PlanStateContractError, StateContractError) as error: + raise StateCorruptionError(f"shard {shard.shard_id!r} has invalid persisted attempts") from error + + @staticmethod + def validate_attempt_against_plan( + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shard: ShardManifest, + attempt: AttemptManifest, + ) -> None: + validate_attempt_manifest(run, shard, attempt) + planned_shard = plan.shards[shard.shard_index] + if planned_shard.shard_id != attempt.shard_id: + raise StateContractError("attempt shard does not match resolved plan order") + if attempt.attempt_id != f"attempt-{attempt.attempt_ordinal:04d}": + raise StateContractError("attempt ID does not match its ordinal") + if attempt.scheduler is not None and attempt.scheduler.array_task_id != planned_shard.array_task_index: + raise StateContractError("attempt scheduler task does not match the resolved plan shard") + if attempt.state is not AttemptLifecycleState.CREATED: + PersistedPlanStateValidator(plan).validate_planned_attempt(planned_shard, attempt) + + @staticmethod + def get_shard(shards: tuple[ShardManifest, ...], shard_id: ShardId) -> ShardManifest: + try: + return next(shard for shard in shards if shard.shard_id == shard_id) + except StopIteration: + raise StateNotFoundError(f"shard {shard_id!r} is not persisted") from None + + @staticmethod + def get_attempt(attempts: tuple[AttemptManifest, ...], attempt_id: AttemptId) -> AttemptManifest: + try: + return next(attempt for attempt in attempts if attempt.attempt_id == attempt_id) + except StopIteration: + raise StateNotFoundError(f"attempt {attempt_id!r} is not persisted") from None 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 new file mode 100644 index 000000000..27fc325a2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum +from typing import Annotated + +from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId +from data_designer.slurm.state.base import StateRecord, StateValue, validate_utc_timestamp + +ReasonCode = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=64, + pattern=r"^[a-z][a-z0-9_]*$", + ), +] + + +class ReadinessState(str, Enum): + PENDING = "pending" + RESTARTING = "restarting" + STARTING = "starting" + READY = "ready" + FAILED = "failed" + STOPPED = "stopped" + + +class EndpointPublicationState(str, Enum): + PENDING = "pending" + PUBLISHED = "published" + FAILED = "failed" + + +class ProbeOutcome(str, Enum): + SUCCESS = "success" + FAILURE = "failure" + + +class ProbeEvidence(StateValue): + """Bounded, redacted evidence from one readiness probe.""" + + observed_at: datetime + outcome: ProbeOutcome + reason_code: ReasonCode + redacted_message: Annotated[str, StringConstraints(max_length=512)] + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + + @field_validator("redacted_message") + @classmethod + def validate_redacted_message(cls, value: str) -> str: + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError("redacted_message must not contain control characters") + return value + + +class DeploymentReadiness(StateValue): + """Readiness state for one authored-order model deployment.""" + + deployment_id: Identifier + model_alias: str + state: ReadinessState + expected_backends: PositiveInt + ready_backends: NonNegativeInt + endpoint_publication: EndpointPublicationState + last_probe: ProbeEvidence | None = None + + @model_validator(mode="after") + def validate_counts_and_state(self) -> DeploymentReadiness: + _validate_deployment_common(self) + _validate_deployment_state(self) + return self + + +class AttemptReadiness(StateRecord): + """Revisioned readiness snapshot for all deployments in one attempt.""" + + run_id: Identifier + shard_id: ShardId + 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) + return self + + +def _validate_deployment_common(deployment: DeploymentReadiness) -> None: + if deployment.ready_backends > deployment.expected_backends: + raise ValueError("ready_backends must not exceed expected_backends") + if deployment.endpoint_publication is EndpointPublicationState.FAILED and deployment.state not in { + ReadinessState.FAILED, + ReadinessState.STOPPED, + }: + raise ValueError("failed endpoint publication requires a failed or stopped deployment") + + +def _validate_deployment_state(deployment: DeploymentReadiness) -> None: + if deployment.state in {ReadinessState.PENDING, ReadinessState.RESTARTING}: + _validate_inactive_deployment(deployment) + elif deployment.state is ReadinessState.STARTING: + _validate_starting_deployment(deployment) + elif deployment.state is ReadinessState.READY: + _validate_ready_deployment(deployment) + elif deployment.state is ReadinessState.STOPPED and deployment.ready_backends != 0: + raise ValueError("stopped deployments cannot have ready backends") + + +def _validate_inactive_deployment(deployment: DeploymentReadiness) -> None: + if deployment.ready_backends != 0: + raise ValueError("pending or restarting deployments cannot have ready backends") + if deployment.endpoint_publication is not EndpointPublicationState.PENDING: + raise ValueError("pending or restarting deployments require pending endpoint publication") + + +def _validate_starting_deployment(deployment: DeploymentReadiness) -> None: + if ( + deployment.ready_backends == deployment.expected_backends + and deployment.endpoint_publication is EndpointPublicationState.PUBLISHED + ): + raise ValueError("fully ready published deployments must use the ready state") + + +def _validate_ready_deployment(deployment: DeploymentReadiness) -> None: + if deployment.ready_backends != deployment.expected_backends: + raise ValueError("ready deployments require every expected backend") + if deployment.endpoint_publication is not EndpointPublicationState.PUBLISHED: + raise ValueError("ready deployments require a published endpoint") + + +def _validate_deployment_uniqueness(deployments: tuple[DeploymentReadiness, ...]) -> None: + deployment_ids = [deployment.deployment_id for deployment in deployments] + model_aliases = [deployment.model_alias for deployment in deployments] + if len(deployment_ids) != len(set(deployment_ids)): + raise ValueError("deployment IDs must be unique") + if len(model_aliases) != len(set(model_aliases)): + raise ValueError("model aliases must be unique") + + +def _validate_probe_chronology(deployments: tuple[DeploymentReadiness, ...], updated_at: datetime) -> None: + if any( + deployment.last_probe is not None and deployment.last_probe.observed_at > updated_at + for deployment in deployments + ): + raise ValueError("probe observations must not be later than the readiness snapshot") + + +def _validate_attempt_state(state: ReadinessState, deployments: tuple[DeploymentReadiness, ...]) -> None: + deployment_states = tuple(deployment.state for deployment in deployments) + if state in {ReadinessState.PENDING, ReadinessState.RESTARTING}: + if any(deployment_state is not state for deployment_state in deployment_states): + raise ValueError(f"{state.value} attempts require every deployment to be {state.value}") + elif state is ReadinessState.STARTING: + if any(item in {ReadinessState.FAILED, ReadinessState.STOPPED} for item in deployment_states): + raise ValueError("starting attempts cannot contain failed or stopped deployments") + if all(item is ReadinessState.READY for item in deployment_states): + raise ValueError("attempts with every deployment ready must use the ready state") + elif state is ReadinessState.READY and any(item is not ReadinessState.READY for item in deployment_states): + raise ValueError("ready attempts require every deployment to be ready") + elif state is ReadinessState.FAILED and not any(item is ReadinessState.FAILED for item in deployment_states): + raise ValueError("failed attempts require at least one failed deployment") + elif state is ReadinessState.STOPPED and any(item is not ReadinessState.STOPPED for item in deployment_states): + raise ValueError("stopped attempts require every deployment to be stopped") 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 new file mode 100644 index 000000000..3dd9ab644 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timedelta + +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest +from data_designer.slurm.state.readiness import ( + AttemptReadiness, + EndpointPublicationState, + ReadinessState, +) +from data_designer.slurm.state.scheduler import ( + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, + is_scheduler_failure_state, +) +from data_designer.slurm.state.validation import StateContractError + +_ALLOWED_READINESS_TRANSITIONS: dict[ReadinessState, frozenset[ReadinessState]] = { + ReadinessState.PENDING: frozenset( + { + ReadinessState.PENDING, + ReadinessState.RESTARTING, + ReadinessState.STARTING, + ReadinessState.FAILED, + ReadinessState.STOPPED, + } + ), + ReadinessState.STARTING: frozenset( + { + ReadinessState.RESTARTING, + ReadinessState.STARTING, + ReadinessState.READY, + ReadinessState.FAILED, + ReadinessState.STOPPED, + } + ), + ReadinessState.RESTARTING: frozenset( + {ReadinessState.RESTARTING, ReadinessState.STARTING, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.READY: frozenset( + {ReadinessState.RESTARTING, ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.FAILED: frozenset({ReadinessState.FAILED, ReadinessState.STOPPED}), + ReadinessState.STOPPED: frozenset({ReadinessState.STOPPED}), +} + +_ALLOWED_ENDPOINT_TRANSITIONS: dict[EndpointPublicationState, frozenset[EndpointPublicationState]] = { + EndpointPublicationState.PENDING: frozenset( + { + EndpointPublicationState.PENDING, + EndpointPublicationState.PUBLISHED, + EndpointPublicationState.FAILED, + } + ), + EndpointPublicationState.PUBLISHED: frozenset({EndpointPublicationState.PUBLISHED}), + EndpointPublicationState.FAILED: frozenset({EndpointPublicationState.FAILED}), +} + + +def validate_readiness_transition( + previous: AttemptReadiness, + current: AttemptReadiness, +) -> AttemptReadiness: + """Validate monotonic revision, identity, order, and readiness transitions.""" + _require(previous.run_id == current.run_id, "readiness run_id cannot change") + _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], + f"attempt readiness cannot move from {previous.state.value} to {current.state.value}", + ) + _require( + len(previous.deployments) == len(current.deployments), + "readiness deployment count cannot change", + ) + + restarting = current.state is ReadinessState.RESTARTING + for old_deployment, new_deployment in zip(previous.deployments, current.deployments, strict=True): + _require( + old_deployment.deployment_id == new_deployment.deployment_id, + "readiness deployment order or ID cannot change", + ) + _require( + old_deployment.model_alias == new_deployment.model_alias, + "readiness deployment model alias cannot change", + ) + _require( + old_deployment.expected_backends == new_deployment.expected_backends, + "readiness expected backend count cannot change", + ) + if restarting: + _require( + old_deployment.state not in {ReadinessState.FAILED, ReadinessState.STOPPED}, + f"deployment {old_deployment.deployment_id!r} cannot restart from terminal readiness", + ) + _require( + new_deployment.state is ReadinessState.RESTARTING + and new_deployment.ready_backends == 0 + and new_deployment.endpoint_publication is EndpointPublicationState.PENDING, + f"deployment {old_deployment.deployment_id!r} restart state is not reset", + ) + else: + _require( + new_deployment.state in _ALLOWED_READINESS_TRANSITIONS[old_deployment.state], + ( + f"deployment {old_deployment.deployment_id!r} cannot move from " + f"{old_deployment.state.value} to {new_deployment.state.value}" + ), + ) + _require( + new_deployment.endpoint_publication + in _ALLOWED_ENDPOINT_TRANSITIONS[old_deployment.endpoint_publication], + f"deployment {old_deployment.deployment_id!r} endpoint publication cannot move backward", + ) + old_probe = old_deployment.last_probe + new_probe = new_deployment.last_probe + if old_probe is not None and not restarting: + if new_probe is None: + raise StateContractError( + f"deployment {old_deployment.deployment_id!r} probe evidence cannot be removed" + ) + _require( + new_probe.observed_at >= old_probe.observed_at, + f"deployment {old_deployment.deployment_id!r} probe observation cannot move backward", + ) + return current + + +def reconcile_attempt_observation( + attempt: AttemptManifest, + readiness: AttemptReadiness | None, + scheduler: SchedulerObservation, + *, + current_time: datetime, +) -> EffectiveAttemptState: + """Apply scheduler terminal precedence without treating readiness as success.""" + _require_utc(current_time, "current_time") + _require(current_time >= scheduler.observed_at, "current_time cannot precede scheduler observation") + _require(current_time >= attempt.updated_at, "current_time cannot precede attempt update") + if readiness is not None: + _require(current_time >= readiness.updated_at, "current_time cannot precede readiness update") + _require(readiness.run_id == attempt.run_id, "readiness run_id does not match attempt") + _require(readiness.shard_id == attempt.shard_id, "readiness shard_id does not match attempt") + _require(readiness.attempt_id == attempt.attempt_id, "readiness attempt_id does not match attempt") + _require(attempt.scheduler is not None, "attempt has no scheduler identity") + _require(scheduler.scheduler == attempt.scheduler, "scheduler identity does not match attempt") + _require(scheduler.observed_at >= attempt.created_at, "scheduler observation cannot precede attempt creation") + if readiness is not None: + _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") + + if is_scheduler_failure_state(scheduler.state): + return EffectiveAttemptState.FAILED + if attempt.state is AttemptLifecycleState.FAILED: + return EffectiveAttemptState.FAILED + if attempt.state is AttemptLifecycleState.SUCCEEDED: + return EffectiveAttemptState.SUCCEEDED + + if scheduler.state is SchedulerState.COMPLETED: + return EffectiveAttemptState.FAILED + if scheduler.state is SchedulerState.ACCOUNTING_LAG: + deadline = scheduler.reconciliation_deadline + _require(deadline is not None, "accounting lag has no reconciliation deadline") + if current_time <= deadline: + return EffectiveAttemptState.ACCOUNTING_LAG + return EffectiveAttemptState.UNKNOWN + if scheduler.state is SchedulerState.UNKNOWN: + return EffectiveAttemptState.UNKNOWN + if readiness is not None and readiness.state is ReadinessState.FAILED: + return EffectiveAttemptState.FAILED + if scheduler.state is SchedulerState.PREEMPTED: + deadline = scheduler.reconciliation_deadline + if deadline is not None and current_time > deadline: + return EffectiveAttemptState.FAILED + return EffectiveAttemptState.PENDING + if scheduler.state in {SchedulerState.PENDING, SchedulerState.REQUEUED}: + return EffectiveAttemptState.PENDING + if scheduler.state is SchedulerState.RUNNING: + return EffectiveAttemptState.RUNNING + raise AssertionError(f"unhandled scheduler state: {scheduler.state}") # pragma: no cover + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) + + +def _require_utc(value: datetime, field_name: str) -> None: + _require( + value.tzinfo is not None and value.utcoffset() == timedelta(0), + f"{field_name} must be in UTC", + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/results.py b/packages/data-designer-slurm/src/data_designer/slurm/state/results.py new file mode 100644 index 000000000..5c15c568a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/results.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated publication and attempt-binding boundary for producer-owned results.""" + +from __future__ import annotations + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.state.errors import ( + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.execution import AttemptManifest +from data_designer.slurm.state.outputs import CandidateOutputManifest +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import StateContractError, validate_attempt_transition + + +class AttemptResultPublisher: + """Publish one result pair and bind its candidate reference as the commit marker.""" + + def __init__(self, storage: StateStorage, reader: StateReader) -> None: + self._storage = storage + self._reader = reader + + def publish( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + """Publish a result pair and bind the attempt under its shard lock.""" + self._validate_record_location(client_result, candidate) + try: + with self._storage.acquire_shard_lock(candidate.shard_id): + run, plan, shard = self._reader.load_shard_context(candidate.shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + attempt = self._reader.get_attempt(attempts, candidate.attempt_id) + PersistedPlanStateValidator(plan).validate_attempt_result( + plan.shards[shard.shard_index], + attempt, + client_result, + candidate, + ) + bound_attempt = self._bind_candidate_reference(attempt, client_result) + self._storage.publish_finalization_records(client_result, candidate) + if bound_attempt != attempt: + self._storage.replace_attempt(bound_attempt) + else: + self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) + return client_result, candidate + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (FileExistsError, PlanStateContractError, StateContractError) as error: + raise StateConflictError("attempt result does not match persisted run intent") from error + except OSError as error: + raise SlurmStateError(f"cannot publish result for attempt {candidate.attempt_id!r}") from error + + def _validate_record_location( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + if not isinstance(client_result, ClientResult) or not isinstance(candidate, CandidateOutputManifest): + raise StateConflictError("attempt result records have invalid types") + if client_result.run_id != self._storage.run_id or candidate.run_id != self._storage.run_id: + raise StateConflictError("attempt result run identity does not match the state writer") + if client_result.shard_id != candidate.shard_id or client_result.attempt_id != candidate.attempt_id: + raise StateConflictError("attempt result records do not identify the same attempt") + + @staticmethod + def _bind_candidate_reference(attempt: AttemptManifest, client_result: ClientResult) -> AttemptManifest: + reference = client_result.candidate_output_manifest + if reference is None: + raise StateContractError("nonempty client result has no candidate manifest reference") + if attempt.candidate_output == reference: + return attempt + bound_attempt = attempt.model_copy(update={"candidate_output": reference}) + validate_attempt_transition(attempt, bound_attempt) + return bound_attempt + + +__all__ = ["AttemptResultPublisher"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py new file mode 100644 index 000000000..8faef34c3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process failed-shard retry orchestration.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager +from datetime import datetime, timezone +from pathlib import Path +from typing import Literal, Protocol + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.contracts import Identifier, ShardId, validate_absolute_path +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import SlurmJobSubmissionReceipt, SlurmSubmissionMatch +from data_designer.slurm.launcher.renderer import render_generation_retry_script +from data_designer.slurm.state.base import SchedulerIdentity +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError, StateCorruptionError +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.observation import SchedulerObservationClient +from data_designer.slurm.state.observer import SlurmStateReconciler +from data_designer.slurm.state.outputs import RetryPlan, RetryShard +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.retry_records import RetryState, RetryStatus, validate_retry_status_transition +from data_designer.slurm.state.retry_storage import RetryStorage +from data_designer.slurm.state.scheduler import EffectiveAttemptState +from data_designer.slurm.state.status import RunStatus, ShardStatus +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.submission_recovery import ( + SUBMISSION_VISIBILITY_WINDOW, + PreparedSubmission, + resolve_prepared_submission, +) +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_attempt_set, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RetryScheduler(SchedulerObservationClient, Protocol): + """Scheduler operations required by fresh-process retry.""" + + def submit_script( + self, script: str, *, export_environment: Mapping[str, str] | None = None + ) -> SlurmJobSubmissionReceipt: + """Submit one rendered retry array.""" + ... + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact retry submission name.""" + ... + + +class SlurmRetryCoordinator: + """Select retryable shards and durably submit their next attempts.""" + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + scheduler: RetryScheduler | None = None, + export_environment: Mapping[str, str] | None = None, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + self._export_environment = None if export_environment is None else dict(export_environment) + self._state = StateStorage(root, normalized_run_id) + self._reader = StateReader(self._state, normalized_run_id) + self._retries = RetryStorage(self._state) + self._finalizer = WinnerFinalizer(self._state, self._reader) + self._reconciler = SlurmStateReconciler(root, normalized_run_id, self._scheduler) + self._run_id = normalized_run_id + + def retry( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + effective_resume_mode: Literal["never", "always"], + observed_at: datetime | None = None, + ) -> tuple[AttemptManifest, ...]: + """Refresh state, submit exactly the retryable selection, and publish attempts.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + if effective_resume_mode not in {"never", "always"}: + raise StateConflictError("effective resume mode must be 'never' or 'always'") + with self._retries.acquire_lock(): + self._retries.discard_incomplete_tail() + self._settle_pending_retry(timestamp) + status = self._reconciler.refresh(observed_at=timestamp) + active = self._load_active_retry(status, shard_ids, effective_resume_mode) + if active is not None: + return active + selected = _select_retryable_shards(status, shard_ids) + plan = self._build_retry_plan(status, selected, effective_resume_mode, timestamp) + script = render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + with self._acquire_selection_locks(selected): + self._require_fresh_selection(status, selected) + self._persist_prepared_retry(plan, timestamp) + receipt = self._submit(plan, script, timestamp) + attempts, superseded = self._publish_attempts(plan, receipt.job_id, timestamp) + self._settle_retry(plan, receipt.job_id, timestamp, superseded=superseded) + return attempts + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot retry persisted run {self._run_id!r}") from error + + def preview( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + effective_resume_mode: Literal["never", "always"], + observed_at: datetime | None = None, + ) -> tuple[RetryPlan, str]: + """Render a retry plan without persisting or submitting it.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + if effective_resume_mode not in {"never", "always"}: + raise StateConflictError("effective resume mode must be 'never' or 'always'") + status = self._reconciler.observe(observed_at=timestamp) + selected = _select_retryable_shards(status, shard_ids) + plan = self._build_retry_plan(status, selected, effective_resume_mode, timestamp) + return plan, render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot preview retry for persisted run {self._run_id!r}") from error + + def preview_active( + self, + *, + shard_ids: Sequence[ShardId] | None = None, + observed_at: datetime | None = None, + ) -> tuple[RetryPlan, str] | None: + """Render the latest active retry without persisting observations.""" + timestamp = datetime.now(timezone.utc) if observed_at is None else observed_at + try: + status = self._reconciler.observe(observed_at=timestamp) + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return None + retry_status = self._retries.read_optional_status(retry_ids[-1]) + if retry_status is None: + return None + plan = self._load_bound_plan(retry_status) + if self._load_active_retry(status, shard_ids, plan.effective_resume_mode) is None: + return None + return plan, render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) + except (StateConflictError, StateCorruptionError, SlurmStateError): + raise + except (OSError, ValidationError, ValueError) as error: + raise SlurmStateError(f"cannot preview active retry for persisted run {self._run_id!r}") from error + + def _settle_pending_retry(self, updated_at: datetime) -> None: + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return + latest_id = retry_ids[-1] + status = self._retries.read_status(latest_id) + plan = self._load_bound_plan(status) + if status.state is RetryState.PREPARED: + status = self._reconcile_prepared_retry(plan, status, updated_at) + if status.state is not RetryState.SUBMITTED: + return + assert status.array_job_id is not None + shard_ids = tuple(shard.shard_id for shard in plan.planned_shards) + run_status = self._reconciler.refresh(observed_at=updated_at) + selected = tuple(_get_shard_status(run_status, shard_id) for shard_id in shard_ids) + with self._acquire_selection_locks(selected): + _, superseded = self._publish_attempts(plan, status.array_job_id, plan.created_at) + self._settle_retry(plan, status.array_job_id, updated_at, superseded=superseded) + + def _reconcile_prepared_retry( + self, + plan: RetryPlan, + status: RetryStatus, + updated_at: datetime, + ) -> RetryStatus: + assert status.reconciliation_deadline is not None + job_id = resolve_prepared_submission( + self._scheduler, + PreparedSubmission( + job_name=plan.submission_job_name, + submitted_after=plan.created_at, + reconciliation_deadline=status.reconciliation_deadline, + expected_array_task_ids=tuple(shard.array_task_index for shard in plan.planned_shards), + ), + observed_at=updated_at, + ) + if job_id is not None: + return self._publish_submitted_status(plan, job_id, updated_at) + self._fail_retry(plan, updated_at) + return self._retries.read_status(plan.retry_id) + + def _build_retry_plan( + self, + status: RunStatus, + selected: tuple[ShardStatus, ...], + effective_resume_mode: Literal["never", "always"], + created_at: datetime, + ) -> RetryPlan: + plan = self._reader.load_resolved_plan(status.run) + requested_resume = plan.invocation.authored.resume + if requested_resume != "if_possible" and effective_resume_mode != requested_resume: + raise StateConflictError("effective resume mode does not match the pinned resolved plan") + return RetryPlan( + schema_version=1, + retry_id=self._retries.get_next_retry_id(), + run_id=status.run.run_id, + created_at=created_at, + resolved_plan=status.run.resolved_plan, + planned_shards=tuple( + RetryShard( + shard_id=shard_status.shard.shard_id, + attempt_id=f"attempt-{len(shard_status.attempts) + 1:04d}", + attempt_ordinal=len(shard_status.attempts) + 1, + array_task_index=plan.shards[shard_status.shard.shard_index].array_task_index, + ) + for shard_status in selected + ), + effective_resume_mode=effective_resume_mode, + ) + + def _load_active_retry( + self, + status: RunStatus, + requested_shard_ids: Sequence[ShardId] | None, + effective_resume_mode: Literal["never", "always"], + ) -> tuple[AttemptManifest, ...] | None: + retry_ids = self._retries.list_retry_ids() + if not retry_ids: + return None + retry_status = self._retries.read_status(retry_ids[-1]) + retry_plan = self._load_bound_plan(retry_status) + if retry_status.state is not RetryState.COMPLETED: + return None + planned_ids = tuple(shard.shard_id for shard in retry_plan.planned_shards) + if requested_shard_ids is None and any( + _is_retryable(shard) and shard.shard.shard_id not in planned_ids for shard in status.shards + ): + return None + if retry_plan.effective_resume_mode != effective_resume_mode or not self._matches_requested_shards( + requested_shard_ids, planned_ids + ): + return None + return self._get_active_attempts(status, retry_plan) + + @staticmethod + def _matches_requested_shards( + requested_shard_ids: Sequence[ShardId] | None, + planned_ids: tuple[ShardId, ...], + ) -> bool: + if requested_shard_ids is None: + return True + requested = tuple(requested_shard_ids) + return len(requested) == len(set(requested)) and set(requested) == set(planned_ids) + + @staticmethod + def _get_active_attempts(status: RunStatus, retry_plan: RetryPlan) -> tuple[AttemptManifest, ...] | None: + statuses_by_shard = {shard.shard.shard_id: shard for shard in status.shards} + attempts: list[AttemptManifest] = [] + for planned in retry_plan.planned_shards: + shard_status = statuses_by_shard.get(planned.shard_id) + if shard_status is None or not shard_status.attempts: + return None + matching = next( + (item for item in shard_status.attempts if item.attempt.attempt_id == planned.attempt_id), + None, + ) + if shard_status.winner is not None: + if ( + shard_status.winner.attempt_id != planned.attempt_id + or matching is None + or matching.effective_state is not EffectiveAttemptState.SUCCEEDED + ): + return None + attempts.append(matching.attempt) + continue + latest = shard_status.attempts[-1] + if latest.attempt.attempt_id != planned.attempt_id or latest.effective_state not in { + EffectiveAttemptState.PENDING, + EffectiveAttemptState.RUNNING, + EffectiveAttemptState.ACCOUNTING_LAG, + }: + return None + attempts.append(latest.attempt) + return tuple(attempts) + + @contextmanager + def _acquire_selection_locks(self, selected: tuple[ShardStatus, ...]) -> Iterator[None]: + with ExitStack() as resources: + for shard_status in selected: + resources.enter_context(self._state.acquire_resume_and_shard_locks(shard_status.shard.shard_id)) + yield + + def _require_fresh_selection(self, status: RunStatus, selected: tuple[ShardStatus, ...]) -> None: + for shard_status in selected: + run, plan, shard = self._reader.load_shard_context(shard_status.shard.shard_id) + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + expected_attempts = tuple(attempt_status.attempt for attempt_status in shard_status.attempts) + expected_observations = tuple(attempt_status.scheduler for attempt_status in shard_status.attempts) + current_observations = tuple( + self._reader.load_optional_scheduler_observation(attempt) for attempt in attempts + ) + if run != status.run or shard != shard_status.shard or attempts != expected_attempts: + raise StateConflictError("persisted shard changed after retry reconciliation; retry again") + if current_observations != expected_observations: + raise StateConflictError("scheduler evidence changed after retry reconciliation; retry again") + if winner is not None: + raise StateConflictError(f"shard {shard.shard_id!r} already has an immutable winner") + + def _persist_prepared_retry(self, plan: RetryPlan, timestamp: datetime) -> None: + self._retries.ensure_retry(plan.retry_id) + self._retries.publish_plan(plan) + self._retries.publish_status( + RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=1, + updated_at=timestamp, + state=RetryState.PREPARED, + reconciliation_deadline=timestamp + SUBMISSION_VISIBILITY_WINDOW, + ) + ) + + def _submit(self, plan: RetryPlan, script: str, timestamp: datetime) -> SlurmJobSubmissionReceipt: + try: + receipt = ( + self._scheduler.submit_script(script) + if self._export_environment is None + else self._scheduler.submit_script(script, export_environment=self._export_environment) + ) + except SlurmSubmissionError as error: + if not error.may_have_succeeded: + self._fail_retry(plan, timestamp) + raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error + except SlurmLauncherError as error: + raise SlurmStateError(f"cannot submit retry {plan.retry_id!r}") from error + self._publish_submitted_status(plan, receipt.job_id, timestamp) + return receipt + + def _publish_submitted_status(self, plan: RetryPlan, array_job_id: int, timestamp: datetime) -> RetryStatus: + plan_reference = self._retries.get_plan_reference(plan) + submitted = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=plan_reference, + revision=2, + updated_at=timestamp, + state=RetryState.SUBMITTED, + array_job_id=array_job_id, + ) + validate_retry_status_transition(self._retries.read_status(plan.retry_id), submitted) + self._retries.replace_status(submitted) + return submitted + + def _publish_attempts( + self, + retry_plan: RetryPlan, + array_job_id: int, + timestamp: datetime, + ) -> tuple[tuple[AttemptManifest, ...], bool]: + run = self._reader.load_run() + plan = self._reader.load_resolved_plan(run) + if retry_plan.run_id != run.run_id or retry_plan.resolved_plan != run.resolved_plan: + raise StateConflictError("retry plan does not bind the current persisted run") + prepared: list[tuple[AttemptManifest, bool]] = [] + superseded = False + for selected in retry_plan.planned_shards: + shard = self._reader.load_shard_context(selected.shard_id)[2] + attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + winner = self._finalizer.load_optional_winner(run, plan, shard, attempts) + attempt = AttemptManifest( + schema_version=1, + run_id=run.run_id, + shard_id=selected.shard_id, + attempt_id=selected.attempt_id, + attempt_ordinal=selected.attempt_ordinal, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity( + array_job_id=array_job_id, + array_task_id=selected.array_task_index, + ), + created_at=timestamp, + updated_at=timestamp, + ) + existing = next((item for item in attempts if item.attempt_id == selected.attempt_id), None) + if existing is not None: + try: + validate_attempt_transition(attempt, existing) + except StateContractError as error: + raise StateConflictError( + f"retry attempt {selected.attempt_id!r} contains incompatible state" + ) from error + if winner is not None and winner.attempt_id != selected.attempt_id: + superseded = True + continue + prepared.append((existing, False)) + continue + if winner is not None: + superseded = True + continue + self._finalizer.require_no_winner(run, plan, shard, attempts) + if selected.attempt_ordinal != len(attempts) + 1: + raise StateConflictError("retry attempt ordinal is no longer next for its shard") + self._reader.validate_attempt_against_plan(run, plan, shard, attempt) + validate_shard_attempt_set(run, shard, attempts + (attempt,)) + prepared.append((attempt, True)) + published: list[AttemptManifest] = [] + for attempt, requires_publication in prepared: + if requires_publication: + self._state.publish_attempt(attempt) + published.append(attempt) + return tuple(published), superseded + + def _settle_retry( + self, + plan: RetryPlan, + array_job_id: int, + timestamp: datetime, + *, + superseded: bool, + ) -> None: + if superseded: + self._fail_retry(plan, timestamp, array_job_id=array_job_id) + else: + self._complete_retry(plan, array_job_id, timestamp) + + def _fail_retry(self, plan: RetryPlan, timestamp: datetime, *, array_job_id: int | None = None) -> None: + previous = self._retries.read_status(plan.retry_id) + failed = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=previous.revision + 1, + updated_at=timestamp, + state=RetryState.FAILED, + array_job_id=array_job_id, + ) + validate_retry_status_transition(previous, failed) + self._retries.replace_status(failed) + + def _complete_retry(self, plan: RetryPlan, array_job_id: int, timestamp: datetime) -> None: + previous = self._retries.read_status(plan.retry_id) + if previous.state is RetryState.COMPLETED: + return + completed = RetryStatus( + schema_version=1, + retry_id=plan.retry_id, + run_id=plan.run_id, + retry_plan=self._retries.get_plan_reference(plan), + revision=previous.revision + 1, + updated_at=timestamp, + state=RetryState.COMPLETED, + array_job_id=array_job_id, + ) + validate_retry_status_transition(previous, completed) + self._retries.replace_status(completed) + + def _load_bound_plan(self, status: RetryStatus) -> RetryPlan: + plan = self._retries.read_plan(status.retry_id) + if status.retry_plan != self._retries.get_plan_reference(plan): + raise StateCorruptionError("retry status does not bind its persisted retry plan") + return plan + + +def _select_retryable_shards( + status: RunStatus, + shard_ids: Sequence[ShardId] | None, +) -> tuple[ShardStatus, ...]: + requested = None if shard_ids is None else tuple(shard_ids) + if requested is not None and len(requested) != len(set(requested)): + raise StateConflictError("explicit retry shard IDs must be unique") + known = {shard.shard.shard_id for shard in status.shards} + if requested is not None and not set(requested).issubset(known): + raise StateConflictError("explicit retry selection contains an unknown shard") + selected = tuple( + shard + for shard in status.shards + if (requested is None or shard.shard.shard_id in requested) and _is_retryable(shard) + ) + expected_count = len( + tuple(shard for shard in status.shards if requested is None or shard.shard.shard_id in requested) + ) + if requested is not None and len(selected) != expected_count: + raise StateConflictError("explicit retry selection includes a sealed or nonterminal shard") + if not selected: + raise StateConflictError("run has no retryable shards") + return selected + + +def _is_retryable(shard: ShardStatus) -> bool: + return ( + shard.winner is None + and bool(shard.attempts) + and shard.attempts[-1].effective_state in {EffectiveAttemptState.FAILED, EffectiveAttemptState.UNKNOWN} + ) + + +def _get_shard_status(status: RunStatus, shard_id: ShardId) -> ShardStatus: + shard_status = next((shard for shard in status.shards if shard.shard.shard_id == shard_id), None) + if shard_status is None: + raise StateCorruptionError(f"retry plan references unknown shard {shard_id!r}") + return shard_status + + +def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: + try: + root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted retry location") from error + return Path(root), normalized_run_id + + +__all__ = ["RetryScheduler", "SlurmRetryCoordinator"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py new file mode 100644 index 000000000..d5045c20c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Persisted retry submission lifecycle.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import PositiveInt, field_validator, model_validator + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.state.base import StateRecord, validate_optional_utc_timestamp, validate_utc_timestamp + + +class RetryState(str, Enum): + """Durable state of one failed-shard retry request.""" + + PREPARED = "prepared" + SUBMITTED = "submitted" + COMPLETED = "completed" + FAILED = "failed" + + +class RetryStatus(StateRecord): + """Atomically replaced submission progress for one retry plan.""" + + retry_id: Identifier + run_id: Identifier + retry_plan: ArtifactReference + revision: PositiveInt + updated_at: datetime + state: RetryState + array_job_id: PositiveInt | None = None + reconciliation_deadline: datetime | None = None + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) + + @model_validator(mode="after") + def validate_scheduler_identity(self) -> RetryStatus: + if self.state in {RetryState.SUBMITTED, RetryState.COMPLETED} and self.array_job_id is None: + raise ValueError("submitted retry state requires an array job identity") + if self.state is RetryState.PREPARED: + if self.array_job_id is not None: + raise ValueError("prepared retry state cannot contain an array job identity") + if self.reconciliation_deadline is None or self.reconciliation_deadline <= self.updated_at: + raise ValueError("prepared retry state requires a future reconciliation deadline") + elif self.reconciliation_deadline is not None: + raise ValueError("settled retry state cannot contain a reconciliation deadline") + return self + + +def validate_retry_status_transition(previous: RetryStatus, current: RetryStatus) -> RetryStatus: + """Require immutable retry identity and one-way submission progress.""" + if previous.retry_id != current.retry_id or previous.run_id != current.run_id: + raise ValueError("retry status identity cannot change") + if previous.retry_plan != current.retry_plan: + raise ValueError("retry status plan identity cannot change") + if current.revision != previous.revision + 1: + raise ValueError("retry status revision must increase by one") + if current.updated_at < previous.updated_at: + raise ValueError("retry status timestamp cannot move backward") + allowed = { + RetryState.PREPARED: {RetryState.SUBMITTED, RetryState.FAILED}, + RetryState.SUBMITTED: {RetryState.COMPLETED, RetryState.FAILED}, + RetryState.COMPLETED: set(), + RetryState.FAILED: set(), + } + if current.state not in allowed[previous.state]: + raise ValueError(f"retry status cannot move from {previous.state.value!r} to {current.state.value!r}") + return current + + +__all__ = ["RetryState", "RetryStatus", "validate_retry_status_transition"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py new file mode 100644 index 000000000..8847516e2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Descriptor-bound persistence for retry plans and submission progress.""" + +from __future__ import annotations + +import os +import re +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.state.filesystem import ( + acquire_file_lock, + ensure_private_child_directory, + is_state_temporary_name, + open_verified_child_directory, + publish_immutable_text, + replace_text, +) +from data_designer.slurm.state.outputs import RetryPlan +from data_designer.slurm.state.retry_records import RetryStatus +from data_designer.slurm.state.storage import StateStorage + +_RETRIES_DIRECTORY = "retries" +_RETRY_LOCK = "retry.lock" +_PLAN_FILENAME = "plan.json" +_STATUS_FILENAME = "status.json" +_RETRY_PATTERN = re.compile(r"^retry-[0-9]{4,}$") +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 + + +class RetryStorage: + """Persist retry journals separately from run and attempt storage.""" + + def __init__(self, state_storage: StateStorage) -> None: + self._state = state_storage + self.retries_root = state_storage.run_root / _RETRIES_DIRECTORY + + @contextmanager + def acquire_lock(self) -> Iterator[None]: + """Serialize retry selection and submission for one run.""" + with self._state.open_run_directory() as run_descriptor: + with acquire_file_lock(run_descriptor, _RETRY_LOCK, self._state.run_root / _RETRY_LOCK): + yield + + def get_next_retry_id(self) -> Identifier: + """Return the next monotonic retry identity.""" + retry_ids = self.list_retry_ids() + if retry_ids and self.read_optional_status(retry_ids[-1]) is None: + return retry_ids[-1] + return f"retry-{len(retry_ids) + 1:04d}" + + def discard_incomplete_tail(self) -> None: + """Discard one trailing journal that cannot have reached submission.""" + try: + with self._open_retries_directory() as retries_descriptor: + retry_ids = _validated_retry_ids(tuple(os.listdir(retries_descriptor))) + if not retry_ids: + return + retry_id = retry_ids[-1] + with open_verified_child_directory( + retries_descriptor, + retry_id, + self.get_retry_root(retry_id), + ) as retry_descriptor: + if _record_exists(retry_descriptor, _STATUS_FILENAME): + return + _discard_prepared_files(retry_descriptor, self.get_retry_root(retry_id)) + os.rmdir(retry_id, dir_fd=retries_descriptor) + os.fsync(retries_descriptor) + except FileNotFoundError: + return + + def list_retry_ids(self) -> tuple[Identifier, ...]: + """List a complete monotonic set of managed retry directories.""" + try: + with self._open_retries_directory() as descriptor: + names = tuple(os.listdir(descriptor)) + except FileNotFoundError: + return () + return _validated_retry_ids(names) + + def ensure_retry(self, retry_id: Identifier) -> None: + """Create one private retry journal directory.""" + with self._state.open_run_directory() as run_descriptor: + ensure_private_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) + with open_verified_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) as descriptor: + ensure_private_child_directory(descriptor, retry_id, self.get_retry_root(retry_id)) + + def get_retry_root(self, retry_id: Identifier) -> Path: + return self.retries_root / retry_id + + def get_plan_path(self, retry_id: Identifier) -> Path: + """Return the canonical immutable retry-plan path.""" + return self.get_retry_root(retry_id) / _PLAN_FILENAME + + def get_plan_reference(self, plan: RetryPlan) -> ArtifactReference: + """Return the exact path and digest bound by retry status.""" + logical_path = self._state.logical_run_root / _RETRIES_DIRECTORY / plan.retry_id / _PLAN_FILENAME + return ArtifactReference(path=logical_path.as_posix(), sha256=plan.compute_sha256()) + + def publish_plan(self, plan: RetryPlan) -> None: + self._require_run_id(plan.run_id) + with self._open_retry_directory(plan.retry_id) as descriptor: + publish_immutable_text( + descriptor, + _PLAN_FILENAME, + plan.serialize_json(), + self.get_plan_path(plan.retry_id), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_plan(self, retry_id: Identifier) -> RetryPlan: + with self._open_retry_directory(retry_id) as descriptor: + plan = self._state.read_record( + descriptor, + _PLAN_FILENAME, + self.get_plan_path(retry_id), + RetryPlan, + ) + if plan.retry_id != retry_id or plan.run_id != self._state.run_id: + raise OSError("retry plan identity does not match its persisted location") + return plan + + def publish_status(self, status: RetryStatus) -> None: + self._require_run_id(status.run_id) + with self._open_retry_directory(status.retry_id) as descriptor: + publish_immutable_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_retry_root(status.retry_id) / _STATUS_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_status(self, status: RetryStatus) -> None: + self._require_run_id(status.run_id) + with self._open_retry_directory(status.retry_id) as descriptor: + replace_text( + descriptor, + _STATUS_FILENAME, + status.serialize_json(), + self.get_retry_root(status.retry_id) / _STATUS_FILENAME, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def read_status(self, retry_id: Identifier) -> RetryStatus: + with self._open_retry_directory(retry_id) as descriptor: + status = self._state.read_record( + descriptor, + _STATUS_FILENAME, + self.get_retry_root(retry_id) / _STATUS_FILENAME, + RetryStatus, + ) + if status.retry_id != retry_id or status.run_id != self._state.run_id: + raise OSError("retry status identity does not match its persisted location") + return status + + def read_optional_status(self, retry_id: Identifier) -> RetryStatus | None: + """Return retry status when its journal reached status publication.""" + try: + return self.read_status(retry_id) + except FileNotFoundError: + return None + + @contextmanager + def _open_retries_directory(self) -> Iterator[int]: + with self._state.open_run_directory() as run_descriptor: + with open_verified_child_directory(run_descriptor, _RETRIES_DIRECTORY, self.retries_root) as descriptor: + yield descriptor + + @contextmanager + def _open_retry_directory(self, retry_id: Identifier) -> Iterator[int]: + with self._open_retries_directory() as retries_descriptor: + with open_verified_child_directory( + retries_descriptor, + retry_id, + self.get_retry_root(retry_id), + ) as descriptor: + yield descriptor + + def _require_run_id(self, run_id: Identifier) -> None: + if run_id != self._state.run_id: + raise OSError("retry record run identity does not match storage") + + +def _validated_retry_ids(names: tuple[str, ...]) -> tuple[Identifier, ...]: + if any(_RETRY_PATTERN.fullmatch(name) is None for name in names): + raise OSError("retry state contains an unowned directory") + ordered = tuple(sorted(names, key=lambda name: int(name.rsplit("-", maxsplit=1)[1]))) + expected = tuple(f"retry-{index:04d}" for index in range(1, len(ordered) + 1)) + if ordered != expected: + raise OSError("retry identities are not a complete monotonic sequence") + return ordered + + +def _record_exists(directory_descriptor: int, name: str) -> bool: + try: + os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + except FileNotFoundError: + return False + return True + + +def _discard_prepared_files(directory_descriptor: int, display_path: Path) -> None: + names = tuple(os.listdir(directory_descriptor)) + if any(name != _PLAN_FILENAME and not is_state_temporary_name(name) for name in names): + raise OSError(f"incomplete retry journal {display_path} contains an unowned entry") + for name in names: + status = os.stat(name, dir_fd=directory_descriptor, follow_symlinks=False) + if not stat.S_ISREG(status.st_mode) or status.st_mode & 0o077: + raise OSError(f"incomplete retry journal entry {display_path / name} is unsafe") + os.unlink(name, dir_fd=directory_descriptor) + os.fsync(directory_descriptor) + + +__all__ = ["RetryStorage"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py new file mode 100644 index 000000000..6961a1295 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import field_validator, model_validator + +from data_designer.slurm.state.base import ( + SchedulerJobIdentity, + StateRecord, + validate_optional_utc_timestamp, + validate_utc_timestamp, +) + + +class SchedulerState(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + TIMED_OUT = "timed_out" + NODE_FAILED = "node_failed" + PREEMPTED = "preempted" + REQUEUED = "requeued" + OUT_OF_MEMORY = "out_of_memory" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + + +class SchedulerObservation(StateRecord): + """Normalized scheduler observation used for deterministic reconciliation.""" + + scheduler: SchedulerJobIdentity + observed_at: datetime + state: SchedulerState + reconciliation_deadline: datetime | None = None + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + _reconciliation_deadline_is_utc = field_validator("reconciliation_deadline")(validate_optional_utc_timestamp) + + @model_validator(mode="after") + def validate_reconciliation_deadline(self) -> SchedulerObservation: + if self.state is SchedulerState.ACCOUNTING_LAG: + if self.reconciliation_deadline is None: + raise ValueError("accounting lag requires a reconciliation deadline") + elif self.state is not SchedulerState.PREEMPTED and self.reconciliation_deadline is not None: + raise ValueError("only accounting lag or preemption may have a reconciliation deadline") + if self.reconciliation_deadline is not None and self.reconciliation_deadline < self.observed_at: + raise ValueError("reconciliation deadline must not precede the observation") + return self + + +class EffectiveAttemptState(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + + +def is_scheduler_failure_state(state: SchedulerState) -> bool: + """Return whether a scheduler state is terminal failure evidence.""" + return state in { + SchedulerState.FAILED, + SchedulerState.CANCELLED, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.OUT_OF_MEMORY, + } + + +def is_scheduler_terminal_state(state: SchedulerState) -> bool: + """Return whether a scheduler state is terminal accounting evidence.""" + return state is SchedulerState.COMPLETED or is_scheduler_failure_state(state) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/status.py b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py new file mode 100644 index 000000000..2d8aeaa5a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/status.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process status values derived from persisted and scheduler evidence.""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from pydantic import Field, field_validator, model_validator + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.state.base import StateValue, validate_utc_timestamp +from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import EffectiveAttemptState, SchedulerObservation + + +class GenerationState(str, Enum): + """Effective progress of one attempt's dataset generation.""" + + NOT_STARTED = "not_started" + ACTIVE = "active" + CANDIDATE_READY = "candidate_ready" + WON = "won" + FAILED = "failed" + UNKNOWN = "unknown" + + +class EffectiveRunState(str, Enum): + """Aggregated state of all planned shards in one run.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" + + +class AttemptStatus(StateValue): + """Validated persisted and observed evidence for one attempt.""" + + attempt: AttemptManifest + readiness: AttemptReadiness | None + scheduler: SchedulerObservation | None + client_result: ClientResult | None + candidate_output: CandidateOutputManifest | None + effective_state: EffectiveAttemptState + generation_state: GenerationState + is_winner: bool = False + + @model_validator(mode="after") + def validate_evidence(self) -> AttemptStatus: + if self.attempt.scheduler is None: + if self.scheduler is not None: + raise ValueError("created attempts cannot have scheduler evidence") + elif self.scheduler is None or self.scheduler.scheduler != self.attempt.scheduler: + raise ValueError("attempt status requires matching scheduler evidence") + if self.readiness is not None and ( + self.readiness.run_id, + self.readiness.shard_id, + self.readiness.attempt_id, + ) != (self.attempt.run_id, self.attempt.shard_id, self.attempt.attempt_id): + raise ValueError("attempt status readiness identity does not match") + if (self.client_result is None) != (self.candidate_output is None): + raise ValueError("attempt status requires a complete generation-result pair") + if self.client_result is not None and self.candidate_output is not None: + expected = (self.attempt.run_id, self.attempt.shard_id, self.attempt.attempt_id) + if ( + self.client_result.run_id, + self.client_result.shard_id, + self.client_result.attempt_id, + ) != expected: + raise ValueError("client result identity does not match the attempt") + if ( + self.candidate_output.run_id, + self.candidate_output.shard_id, + self.candidate_output.attempt_id, + ) != expected: + raise ValueError("candidate output identity does not match the attempt") + expected_generation = derive_generation_state( + self.effective_state, + has_candidate=self.candidate_output is not None, + is_winner=self.is_winner, + ) + if self.generation_state is not expected_generation: + raise ValueError("generation state does not match the attempt evidence") + return self + + +class ShardStatus(StateValue): + """Effective state and attempt history for one planned shard.""" + + shard: ShardManifest + attempts: tuple[AttemptStatus, ...] + winner: ShardWinner | None + effective_state: EffectiveAttemptState + + @model_validator(mode="after") + def validate_status(self) -> ShardStatus: + if any(status.attempt.shard_id != self.shard.shard_id for status in self.attempts): + raise ValueError("shard status contains an attempt for another shard") + ordinals = tuple(status.attempt.attempt_ordinal for status in self.attempts) + if ordinals != tuple(range(1, len(self.attempts) + 1)): + raise ValueError("shard status attempts must be in complete ordinal order") + winning_attempts = tuple(status for status in self.attempts if status.is_winner) + if self.winner is None: + if winning_attempts: + raise ValueError("shard status marks a winner without a winner record") + elif ( + self.winner.shard_id != self.shard.shard_id + or len(winning_attempts) != 1 + or winning_attempts[0].attempt.attempt_id != self.winner.attempt_id + ): + raise ValueError("shard winner does not match its observed attempt") + if self.effective_state is not derive_shard_state(self.attempts, self.winner): + raise ValueError("effective shard state does not match its evidence") + return self + + +class RunStatus(StateValue): + """Fresh-process status for every planned shard in one run.""" + + run: RunManifest + observed_at: datetime + shards: tuple[ShardStatus, ...] = Field(min_length=1) + effective_state: EffectiveRunState + + _observed_at_is_utc = field_validator("observed_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_status(self) -> RunStatus: + if self.observed_at < self.run.created_at: + raise ValueError("run observation cannot precede run creation") + if len(self.shards) != self.run.shard_count: + raise ValueError("run status must include every planned shard") + if tuple(status.shard.shard_index for status in self.shards) != tuple(range(self.run.shard_count)): + raise ValueError("run status shards must be in planned order") + if any(status.shard.run_id != self.run.run_id for status in self.shards): + raise ValueError("run status contains a shard for another run") + if self.effective_state is not derive_run_state(self.shards): + raise ValueError("effective run state does not match its shard evidence") + return self + + +def derive_generation_state( + effective_state: EffectiveAttemptState, + *, + has_candidate: bool, + is_winner: bool, +) -> GenerationState: + """Derive generation progress without treating readiness as success.""" + if effective_state is EffectiveAttemptState.FAILED: + return GenerationState.FAILED + if effective_state is EffectiveAttemptState.UNKNOWN: + return GenerationState.UNKNOWN + if is_winner: + return GenerationState.WON + if has_candidate: + return GenerationState.CANDIDATE_READY + if effective_state is EffectiveAttemptState.RUNNING: + return GenerationState.ACTIVE + return GenerationState.NOT_STARTED + + +def derive_shard_state( + attempts: tuple[AttemptStatus, ...], + winner: ShardWinner | None, +) -> EffectiveAttemptState: + """Derive one shard's effective state from its newest attempt and winner.""" + if winner is not None: + winning = next((status for status in attempts if status.attempt.attempt_id == winner.attempt_id), None) + if winning is not None and winning.effective_state is EffectiveAttemptState.SUCCEEDED: + return EffectiveAttemptState.SUCCEEDED + if not attempts: + return EffectiveAttemptState.PENDING + latest = attempts[-1].effective_state + return EffectiveAttemptState.UNKNOWN if latest is EffectiveAttemptState.SUCCEEDED else latest + + +def derive_run_state(shards: tuple[ShardStatus, ...]) -> EffectiveRunState: + """Aggregate shard states without declaring a partially active run terminal.""" + states = tuple(shard.effective_state for shard in shards) + if all(state is EffectiveAttemptState.SUCCEEDED for state in states): + return EffectiveRunState.SUCCEEDED + if any(state is EffectiveAttemptState.RUNNING for state in states): + return EffectiveRunState.RUNNING + if any(state is EffectiveAttemptState.ACCOUNTING_LAG for state in states): + return EffectiveRunState.ACCOUNTING_LAG + if any(state is EffectiveAttemptState.PENDING for state in states): + return EffectiveRunState.PENDING + if any(state is EffectiveAttemptState.UNKNOWN for state in states): + return EffectiveRunState.UNKNOWN + return EffectiveRunState.FAILED + + +__all__ = [ + "AttemptStatus", + "EffectiveRunState", + "GenerationState", + "RunStatus", + "ShardStatus", + "derive_generation_state", + "derive_run_state", + "derive_shard_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 new file mode 100644 index 000000000..606aab299 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -0,0 +1,617 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Physical storage for persisted Slurm state records.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Callable, Literal, TypeVar + +from pydantic import ValidationError + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import AttemptId, ContractRecord, Identifier, ShardId, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.errors import SlurmStateError, StateCorruptionError, StateNotFoundError +from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.filesystem import ( + acquire_file_lock, + ensure_private_child_directory, + is_state_temporary_name, + open_verified_child_directory, + open_verified_directory, + publish_immutable_text, + read_regular_text, + replace_text, + sync_directory, +) +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.scheduler import SchedulerObservation + +_RUN_FILENAME = "run.json" +_AUTHORED_CONFIG_FILENAME = "authored-config.json" +_RESOLVED_PLAN_FILENAME = "resolved-plan.json" +_SHARDS_DIRECTORY_NAME = "shards" +_ATTEMPTS_DIRECTORY_NAME = "attempts" +_SHARD_FILENAME = "shard.json" +_SHARD_LOCK_FILENAME = "shard.lock" +_ATTEMPT_FILENAME = "attempt.json" +_READINESS_FILENAME = "readiness.json" +_SCHEDULER_OBSERVATION_FILENAME = "scheduler.json" +_CLIENT_RESULT_FILENAME = "client-result.json" +_CANDIDATE_OUTPUT_FILENAME = "output-manifest.json" +_WINNER_FILENAME = "winner.json" +_DATASET_DIRECTORY_NAME = "dataset" +_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) +_LocalPathResolver = Callable[[str], str | Path] + + +class StateStorage: + """Own descriptor-bound paths, locking, and record serialization.""" + + def __init__( + self, + workspace_root: Path, + run_id: Identifier, + *, + logical_workspace_root: Path | None = None, + local_path_resolver: _LocalPathResolver | None = None, + ) -> None: + self.workspace_root = workspace_root + self.logical_workspace_root = logical_workspace_root or workspace_root + self.run_id = run_id + self.runs_root = workspace_root / "runs" + self.locks_root = self.runs_root / _LOCK_DIRECTORY_NAME + self.run_root = self.runs_root / run_id + self.logical_run_root = self.logical_workspace_root / "runs" / run_id + self._local_path_resolver = local_path_resolver + + @property + def authored_config_path(self) -> Path: + return self.run_root / _AUTHORED_CONFIG_FILENAME + + @property + def resolved_plan_path(self) -> Path: + return self.run_root / _RESOLVED_PLAN_FILENAME + + @property + def logical_authored_config_path(self) -> Path: + return self.logical_run_root / _AUTHORED_CONFIG_FILENAME + + @property + def logical_resolved_plan_path(self) -> Path: + return self.logical_run_root / _RESOLVED_PLAN_FILENAME + + def get_local_path(self, logical_path: str | Path) -> Path: + logical_path = Path(logical_path) + try: + relative_path = logical_path.relative_to(self.logical_workspace_root) + except ValueError as error: + raise StateCorruptionError("persisted path is outside the selected workspace") from error + if self._local_path_resolver is not None: + try: + return Path(validate_absolute_path(Path(self._local_path_resolver(logical_path.as_posix())).as_posix())) + except Exception as error: + raise StateCorruptionError("persisted path has no valid local mapping") from error + return self.workspace_root / relative_path + + def get_shard_path(self, shard_id: str) -> Path: + return self.run_root / _SHARDS_DIRECTORY_NAME / shard_id + + def get_attempt_path(self, shard_id: str, attempt_id: str) -> Path: + return self.get_shard_path(shard_id) / _ATTEMPTS_DIRECTORY_NAME / attempt_id + + def get_readiness_path(self, shard_id: str, attempt_id: str) -> Path: + return self.get_attempt_path(shard_id, attempt_id) / _READINESS_FILENAME + + def get_scheduler_observation_path(self, shard_id: str, attempt_id: str) -> Path: + return self.get_attempt_path(shard_id, attempt_id) / _SCHEDULER_OBSERVATION_FILENAME + + def get_winner_path(self, shard_id: str) -> Path: + return self.get_shard_path(shard_id) / _WINNER_FILENAME + + def ensure_storage(self) -> None: + with open_verified_directory(self.workspace_root) as workspace_descriptor: + ensure_private_child_directory(workspace_descriptor, "runs", self.runs_root) + with open_verified_child_directory(workspace_descriptor, "runs", self.runs_root) as runs_descriptor: + ensure_private_child_directory(runs_descriptor, _LOCK_DIRECTORY_NAME, self.locks_root) + + @contextmanager + def acquire_run_lock(self) -> Iterator[None]: + try: + 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"run-{self.run_id}.lock" + with acquire_file_lock(locks_descriptor, lock_name, self.locks_root / lock_name): + yield + 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: + with self.open_shard_directory(shard_id) as shard_descriptor: + with acquire_file_lock( + shard_descriptor, + _SHARD_LOCK_FILENAME, + self.get_shard_path(shard_id) / _SHARD_LOCK_FILENAME, + ): + yield + except FileNotFoundError as error: + raise StateNotFoundError(f"shard {shard_id!r} is not persisted") from error + + @contextmanager + def acquire_resume_lock(self, shard_id: ShardId) -> Iterator[None]: + with ExitStack() as contexts: + try: + shard_descriptor = contexts.enter_context(self.open_shard_directory(shard_id)) + contexts.enter_context( + acquire_file_lock( + shard_descriptor, + _RESUME_LOCK_FILENAME, + self.get_shard_path(shard_id) / _RESUME_LOCK_FILENAME, + ) + ) + except FileNotFoundError as error: + raise StateNotFoundError(f"shard {shard_id!r} is not persisted") from error + except OSError as error: + raise SlurmStateError(f"cannot lock dataset workspace for shard {shard_id!r}") from error + yield + + @contextmanager + def acquire_resume_and_shard_locks(self, shard_id: ShardId) -> Iterator[None]: + """Acquire dataset-resume then shard state locks in the canonical order.""" + with self.acquire_resume_lock(shard_id): + with self.acquire_shard_lock(shard_id): + yield + + def publish_initial_state( + self, + authored_config: DataDesignerSlurmConfig, + resolved_plan: ResolvedSlurmRunPlan, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> None: + with open_verified_directory(self.runs_root, require_private=True) as runs_descriptor: + ensure_private_child_directory(runs_descriptor, self.run_id, self.run_root) + with open_verified_child_directory(runs_descriptor, self.run_id, self.run_root) as run_descriptor: + ensure_private_child_directory( + run_descriptor, + _SHARDS_DIRECTORY_NAME, + self.run_root / _SHARDS_DIRECTORY_NAME, + ) + self._publish_immutable_record(run_descriptor, _AUTHORED_CONFIG_FILENAME, authored_config) + self._publish_immutable_record(run_descriptor, _RESOLVED_PLAN_FILENAME, resolved_plan) + with open_verified_child_directory( + run_descriptor, + _SHARDS_DIRECTORY_NAME, + self.run_root / _SHARDS_DIRECTORY_NAME, + ) as shards_descriptor: + for shard in shards: + self._publish_initial_shard(shards_descriptor, shard) + self._publish_immutable_record(run_descriptor, _RUN_FILENAME, run) + + def read_run(self) -> RunManifest: + with self.open_run_directory() as run_descriptor: + return self.read_record(run_descriptor, _RUN_FILENAME, self.run_root / _RUN_FILENAME, RunManifest) + + def read_authored_config(self) -> DataDesignerSlurmConfig: + with self.open_run_directory() as run_descriptor: + return self.read_record( + run_descriptor, + _AUTHORED_CONFIG_FILENAME, + self.authored_config_path, + DataDesignerSlurmConfig, + ) + + def read_resolved_plan(self) -> ResolvedSlurmRunPlan: + with self.open_run_directory() as run_descriptor: + return self.read_record( + run_descriptor, + _RESOLVED_PLAN_FILENAME, + self.resolved_plan_path, + ResolvedSlurmRunPlan, + ) + + def read_shards(self, shard_count: int) -> tuple[ShardManifest, ...]: + with self.open_run_directory() as run_descriptor: + with open_verified_child_directory( + run_descriptor, + _SHARDS_DIRECTORY_NAME, + self.run_root / _SHARDS_DIRECTORY_NAME, + ) as shards_descriptor: + expected_names = tuple(f"shard-{index:05d}" for index in range(shard_count)) + if set(os.listdir(shards_descriptor)) != set(expected_names): + raise StateCorruptionError(f"run {self.run_id!r} has an incomplete shard directory set") + return tuple(self._read_shard(shards_descriptor, shard_name) for shard_name in expected_names) + + def read_shard(self, shard_id: ShardId) -> ShardManifest: + """Read one shard without scanning unrelated shard state.""" + with self.open_run_directory() as run_descriptor: + with open_verified_child_directory( + run_descriptor, + _SHARDS_DIRECTORY_NAME, + self.run_root / _SHARDS_DIRECTORY_NAME, + ) as shards_descriptor: + return self._read_shard(shards_descriptor, shard_id) + + def read_attempts(self, shard_id: ShardId) -> tuple[AttemptManifest, ...]: + with self.open_shard_directory(shard_id) as shard_descriptor: + with open_verified_child_directory( + shard_descriptor, + _ATTEMPTS_DIRECTORY_NAME, + self.get_shard_path(shard_id) / _ATTEMPTS_DIRECTORY_NAME, + ) as attempts_descriptor: + names = tuple(os.listdir(attempts_descriptor)) + if any(_ATTEMPT_NAME_PATTERN.fullmatch(name) is None for name in names): + raise StateCorruptionError(f"shard {shard_id!r} contains an invalid attempt directory") + ordered_names = tuple(sorted(names, key=lambda name: int(name.rsplit("-", maxsplit=1)[1]))) + attempts, incomplete_names = self._read_attempt_directories( + attempts_descriptor, + shard_id, + ordered_names, + ) + self._validate_attempt_sequence(shard_id, attempts, incomplete_names) + return attempts + + def publish_attempt(self, attempt: AttemptManifest) -> None: + with self.open_shard_directory(attempt.shard_id) as shard_descriptor: + with open_verified_child_directory( + shard_descriptor, + _ATTEMPTS_DIRECTORY_NAME, + self.get_shard_path(attempt.shard_id) / _ATTEMPTS_DIRECTORY_NAME, + ) as attempts_descriptor: + attempt_root = self.get_attempt_path(attempt.shard_id, attempt.attempt_id) + ensure_private_child_directory(attempts_descriptor, attempt.attempt_id, attempt_root) + with open_verified_child_directory( + attempts_descriptor, + attempt.attempt_id, + attempt_root, + ) as attempt_descriptor: + self._publish_immutable_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) + + def replace_attempt(self, attempt: AttemptManifest) -> None: + with self.open_attempt_directory(attempt.shard_id, attempt.attempt_id) as attempt_descriptor: + self._replace_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) + + def read_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptReadiness: + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + return self.read_record( + attempt_descriptor, + _READINESS_FILENAME, + self.get_readiness_path(shard_id, attempt_id), + AttemptReadiness, + ) + + def publish_readiness(self, readiness: AttemptReadiness) -> None: + with self.open_attempt_directory(readiness.shard_id, readiness.attempt_id) as attempt_descriptor: + self._publish_immutable_record(attempt_descriptor, _READINESS_FILENAME, readiness) + + def replace_readiness(self, readiness: AttemptReadiness) -> None: + with self.open_attempt_directory(readiness.shard_id, readiness.attempt_id) as attempt_descriptor: + self._replace_record(attempt_descriptor, _READINESS_FILENAME, readiness) + + def read_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + ) -> SchedulerObservation: + """Read one attempt's latest reconciled scheduler evidence.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + return self.read_record( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + path, + SchedulerObservation, + ) + + def publish_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + observation: SchedulerObservation, + ) -> None: + """Publish the first scheduler observation for one attempt.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + publish_immutable_text( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + observation.serialize_json(), + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def replace_scheduler_observation( + self, + shard_id: ShardId, + attempt_id: AttemptId, + observation: SchedulerObservation, + ) -> None: + """Atomically replace one attempt's scheduler observation.""" + path = self.get_scheduler_observation_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + replace_text( + attempt_descriptor, + _SCHEDULER_OBSERVATION_FILENAME, + observation.serialize_json(), + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def sync_attempt_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> None: + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + sync_directory(attempt_descriptor) + + def ensure_dataset_directory( + self, + shard_id: ShardId, + attempt_id: AttemptId, + effective_resume_mode: Literal["never", "always"], + ) -> Path: + if effective_resume_mode == "always": + dataset_path = self.get_shard_path(shard_id) / _DATASET_DIRECTORY_NAME + with self.open_shard_directory(shard_id) as descriptor: + ensure_private_child_directory(descriptor, _DATASET_DIRECTORY_NAME, dataset_path) + logical_path = self.logical_run_root / _SHARDS_DIRECTORY_NAME / shard_id / _DATASET_DIRECTORY_NAME + return self.get_local_path(logical_path) + dataset_path = self.get_attempt_path(shard_id, attempt_id) / _DATASET_DIRECTORY_NAME + with self.open_attempt_directory(shard_id, attempt_id) as descriptor: + ensure_private_child_directory(descriptor, _DATASET_DIRECTORY_NAME, dataset_path) + logical_path = ( + self.logical_run_root + / _SHARDS_DIRECTORY_NAME + / shard_id + / _ATTEMPTS_DIRECTORY_NAME + / attempt_id + / _DATASET_DIRECTORY_NAME + ) + return self.get_local_path(logical_path) + + def read_finalization_records( + self, + shard_id: ShardId, + attempt_id: AttemptId, + ) -> tuple[ClientResult, CandidateOutputManifest]: + attempt_path = self.get_attempt_path(shard_id, attempt_id) + with self.open_attempt_directory(shard_id, attempt_id) as descriptor: + client_result = self.read_record( + descriptor, + _CLIENT_RESULT_FILENAME, + attempt_path / _CLIENT_RESULT_FILENAME, + ClientResult, + ) + candidate = self.read_record( + descriptor, + _CANDIDATE_OUTPUT_FILENAME, + attempt_path / _CANDIDATE_OUTPUT_FILENAME, + CandidateOutputManifest, + ) + return client_result, candidate + + def publish_finalization_records( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: + """Convergently publish one producer-owned attempt result pair.""" + with self.open_attempt_directory(candidate.shard_id, candidate.attempt_id) as descriptor: + self._publish_immutable_record(descriptor, _CANDIDATE_OUTPUT_FILENAME, candidate) + self._publish_immutable_record(descriptor, _CLIENT_RESULT_FILENAME, client_result) + + def read_winner(self, shard_id: ShardId) -> ShardWinner: + with self.open_shard_directory(shard_id) as descriptor: + return self.read_record(descriptor, _WINNER_FILENAME, self.get_winner_path(shard_id), ShardWinner) + + def publish_winner(self, winner: ShardWinner) -> None: + with self.open_shard_directory(winner.shard_id) as descriptor: + self._publish_immutable_record(descriptor, _WINNER_FILENAME, winner) + + def read_record( + self, + directory_descriptor: int, + name: str, + display_path: Path, + record_type: type[_RecordT], + ) -> _RecordT: + try: + content = read_regular_text( + directory_descriptor, + name, + display_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + _validate_supported_record_version(content, display_path) + record = record_type.model_validate_json(content) + if record.serialize_json() != content: + raise StateCorruptionError( + f"persisted state record {display_path.name!r} is not canonical for schema_version 1; " + "persisted schema migrations are not supported" + ) + return record + except StateCorruptionError: + raise + except (RecursionError, UnicodeError, ValueError, ValidationError) as error: + raise StateCorruptionError(f"persisted state record {display_path.name!r} is invalid") from error + + @contextmanager + def open_run_directory(self) -> Iterator[int]: + with open_verified_directory(self.runs_root, require_private=True) as runs_descriptor: + with open_verified_child_directory(runs_descriptor, self.run_id, self.run_root) as run_descriptor: + yield run_descriptor + + @contextmanager + def open_shard_directory(self, shard_id: ShardId) -> Iterator[int]: + with self.open_run_directory() as run_descriptor: + with open_verified_child_directory( + run_descriptor, + _SHARDS_DIRECTORY_NAME, + self.run_root / _SHARDS_DIRECTORY_NAME, + ) as shards_descriptor: + with open_verified_child_directory( + shards_descriptor, + shard_id, + self.get_shard_path(shard_id), + ) as shard_descriptor: + yield shard_descriptor + + @contextmanager + def open_attempt_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> Iterator[int]: + with self.open_shard_directory(shard_id) as shard_descriptor: + with open_verified_child_directory( + shard_descriptor, + _ATTEMPTS_DIRECTORY_NAME, + self.get_shard_path(shard_id) / _ATTEMPTS_DIRECTORY_NAME, + ) as attempts_descriptor: + with open_verified_child_directory( + attempts_descriptor, + attempt_id, + self.get_attempt_path(shard_id, attempt_id), + ) as attempt_descriptor: + yield attempt_descriptor + + def _publish_initial_shard(self, shards_descriptor: int, shard: ShardManifest) -> None: + shard_root = self.get_shard_path(shard.shard_id) + ensure_private_child_directory(shards_descriptor, shard.shard_id, shard_root) + with open_verified_child_directory(shards_descriptor, shard.shard_id, shard_root) as shard_descriptor: + ensure_private_child_directory( + shard_descriptor, + _ATTEMPTS_DIRECTORY_NAME, + shard_root / _ATTEMPTS_DIRECTORY_NAME, + ) + self._publish_immutable_record(shard_descriptor, _SHARD_FILENAME, shard) + + def _read_shard(self, shards_descriptor: int, shard_id: str) -> ShardManifest: + shard_root = self.get_shard_path(shard_id) + with open_verified_child_directory(shards_descriptor, shard_id, shard_root) as shard_descriptor: + return self.read_record(shard_descriptor, _SHARD_FILENAME, shard_root / _SHARD_FILENAME, ShardManifest) + + def _read_attempt_directories( + self, + attempts_descriptor: int, + shard_id: ShardId, + names: tuple[str, ...], + ) -> tuple[tuple[AttemptManifest, ...], tuple[str, ...]]: + published: list[AttemptManifest] = [] + incomplete: list[str] = [] + for name in names: + attempt = self._read_attempt_if_published(attempts_descriptor, shard_id, name) + if attempt is None: + incomplete.append(name) + else: + published.append(attempt) + return tuple(published), tuple(incomplete) + + def _read_attempt_if_published( + self, + attempts_descriptor: int, + shard_id: ShardId, + attempt_id: str, + ) -> AttemptManifest | None: + try: + return self._read_attempt(attempts_descriptor, shard_id, attempt_id) + except FileNotFoundError: + pass + + attempt_root = self.get_attempt_path(shard_id, attempt_id) + with open_verified_child_directory(attempts_descriptor, attempt_id, attempt_root) as attempt_descriptor: + unexpected = tuple(name for name in os.listdir(attempt_descriptor) if not is_state_temporary_name(name)) + if unexpected == (): + return None + if unexpected == (_ATTEMPT_FILENAME,): + return self._read_attempt(attempts_descriptor, shard_id, attempt_id) + raise StateCorruptionError(f"attempt {attempt_id!r} contains unpublished state records") + + def _read_attempt(self, attempts_descriptor: int, shard_id: ShardId, attempt_id: str) -> AttemptManifest: + attempt_root = self.get_attempt_path(shard_id, attempt_id) + with open_verified_child_directory(attempts_descriptor, attempt_id, attempt_root) as attempt_descriptor: + return self.read_record( + attempt_descriptor, + _ATTEMPT_FILENAME, + attempt_root / _ATTEMPT_FILENAME, + AttemptManifest, + ) + + @staticmethod + def _validate_attempt_sequence( + shard_id: ShardId, + attempts: tuple[AttemptManifest, ...], + incomplete_names: tuple[str, ...], + ) -> None: + ordinals = tuple(attempt.attempt_ordinal for attempt in attempts) + if ordinals != tuple(range(1, len(attempts) + 1)): + raise StateCorruptionError(f"shard {shard_id!r} attempts are not a complete monotonic sequence") + if tuple(attempt.attempt_id for attempt in attempts) != tuple(f"attempt-{ordinal:04d}" for ordinal in ordinals): + raise StateCorruptionError(f"shard {shard_id!r} attempt IDs do not match their ordinals") + expected_incomplete = f"attempt-{len(attempts) + 1:04d}" + if incomplete_names not in ((), (expected_incomplete,)): + raise StateCorruptionError(f"shard {shard_id!r} contains an invalid incomplete attempt") + + def _publish_immutable_record(self, directory_descriptor: int, name: str, record: ContractRecord) -> None: + publish_immutable_text( + directory_descriptor, + name, + record.serialize_json(), + self._get_record_path(name, record), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def _replace_record(self, directory_descriptor: int, name: str, record: ContractRecord) -> None: + replace_text( + directory_descriptor, + name, + record.serialize_json(), + self._get_record_path(name, record), + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + + def _get_record_path(self, name: str, record: ContractRecord) -> Path: + if name == _RUN_FILENAME: + return self.run_root / name + if name in (_AUTHORED_CONFIG_FILENAME, _RESOLVED_PLAN_FILENAME): + return self.run_root / name + if isinstance(record, ShardManifest): + return self.get_shard_path(record.shard_id) / name + if isinstance(record, ShardWinner): + return self.get_winner_path(record.shard_id) + if isinstance(record, (AttemptManifest, AttemptReadiness, ClientResult, CandidateOutputManifest)): + return self.get_attempt_path(record.shard_id, record.attempt_id) / name + raise TypeError(f"unsupported persisted record type: {type(record).__name__}") + + +def _validate_supported_record_version(content: str, display_path: Path) -> None: + payload = json.loads(content) + if not isinstance(payload, dict) or payload.get("schema_version") != 1: + raise StateCorruptionError(f"persisted state record {display_path.name!r} uses an unsupported schema_version") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py new file mode 100644 index 000000000..270264e99 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -0,0 +1,437 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validated state transitions for one persisted Slurm run.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from datetime import datetime, timedelta +from pathlib import Path +from typing import Literal + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId, validate_absolute_path +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state.errors import ( + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.outputs import CandidateOutputManifest, RetryPlan, ShardWinner +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.readiness import AttemptReadiness +from data_designer.slurm.state.reconciliation import validate_readiness_transition +from data_designer.slurm.state.results import AttemptResultPublisher +from data_designer.slurm.state.retry_storage import RetryStorage +from data_designer.slurm.state.storage import StateStorage +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_attempt_set, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_SHARD_ID_ADAPTER = TypeAdapter(ShardId) +_ATTEMPT_ID_ADAPTER = TypeAdapter(AttemptId) + + +class SlurmStateWriter: + """Persist and reload one run's immutable and revisioned state. + + The writer owns state-machine and plan validation while its composed + storage collaborator owns descriptor-bound paths, locks, serialization, + and publication. Attempt and readiness mutations lock and validate only + the target shard; full-run validation remains available for status and + audit snapshots. + + Args: + workspace_root: Workspace root visible to this process. + run_id: Stable application-owned run identity. + logical_workspace_root: Optional host-side root persisted in plan records. + local_path_resolver: Optional mapping for persisted paths with nested mounts. + """ + + def __init__( + self, + workspace_root: str | Path, + run_id: Identifier, + *, + logical_workspace_root: str | Path | None = None, + local_path_resolver: Callable[[str], str | Path] | None = None, + ) -> None: + try: + normalized_root = validate_absolute_path(Path(workspace_root).as_posix()) + normalized_logical_root = validate_absolute_path( + Path(logical_workspace_root if logical_workspace_root is not None else workspace_root).as_posix() + ) + normalized_run_id = _IDENTIFIER_ADAPTER.validate_python(run_id, strict=True) + except (ValidationError, ValueError) as error: + raise SlurmStateError("invalid persisted run location") from error + self._storage = StateStorage( + Path(normalized_root), + normalized_run_id, + logical_workspace_root=Path(normalized_logical_root), + local_path_resolver=local_path_resolver, + ) + self._reader = StateReader(self._storage, normalized_run_id) + self._results = AttemptResultPublisher(self._storage, self._reader) + self._finalizer = WinnerFinalizer(self._storage, self._reader) + self._retries = RetryStorage(self._storage) + self._run_id = normalized_run_id + + @property + def run_root(self) -> Path: + """Return the workspace-derived root for this run.""" + return self._storage.run_root + + def initialize_run( + self, + authored_config: DataDesignerSlurmConfig, + resolved_plan: ResolvedSlurmRunPlan, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> RunManifest: + """Convergently publish immutable run intent, plan, and shards.""" + self._validate_initial_state(authored_config, resolved_plan, run, shards) + try: + self._storage.ensure_storage() + with self._storage.acquire_run_lock(): + self._storage.publish_initial_state(authored_config, resolved_plan, run, shards) + return run + except StateConflictError: + raise + except FileExistsError as error: + raise StateConflictError(f"run {self._run_id!r} already contains different immutable state") from error + except OSError as error: + raise SlurmStateError(f"cannot initialize persisted run {self._run_id!r}") from error + + def load_run(self) -> RunManifest: + """Load the committed immutable run manifest.""" + return self._reader.load_run() + + def load_authored_config(self) -> DataDesignerSlurmConfig: + """Load and digest-verify the run's immutable authored config.""" + return self._reader.load_authored_config() + + def load_resolved_plan(self) -> ResolvedSlurmRunPlan: + """Load and digest-verify the run's immutable resolved plan.""" + return self._reader.load_resolved_plan() + + def load_retry_plan(self, retry_id: Identifier) -> RetryPlan: + """Load one retry plan and verify its status binds the same immutable record.""" + try: + normalized_retry_id = _IDENTIFIER_ADAPTER.validate_python(retry_id, strict=True) + plan = self._retries.read_plan(normalized_retry_id) + status = self._retries.read_status(normalized_retry_id) + if status.retry_plan != self._retries.get_plan_reference(plan): + raise StateCorruptionError("retry status does not bind its persisted retry plan") + return plan + except (StateCorruptionError, StateNotFoundError): + raise + except (FileNotFoundError, ValidationError) as error: + raise StateNotFoundError(f"retry {retry_id!r} is unavailable") from error + except OSError as error: + raise StateCorruptionError(f"retry {retry_id!r} is unsafe or unreadable") from error + + def load_shards(self) -> tuple[ShardManifest, ...]: + """Load and validate the complete ordered shard set.""" + return self._reader.load_shards() + + def create_attempt(self, attempt: AttemptManifest) -> AttemptManifest: + """Publish the next monotonically numbered attempt for one shard.""" + self._validate_attempt_location(attempt) + try: + return self._create_attempt_with_locks(attempt) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (PlanStateContractError, StateContractError) as error: + raise StateConflictError("attempt does not match persisted run intent") from error + except (FileNotFoundError, OSError) as error: + raise SlurmStateError(f"cannot create attempt {attempt.attempt_id!r}") from error + + def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: + """Atomically replace one attempt after validating its monotonic transition.""" + self._validate_attempt_location(attempt) + try: + return self._update_attempt_with_locks(attempt) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (PlanStateContractError, StateContractError) as error: + raise StateConflictError("attempt update is not a valid monotonic transition") from error + except (FileNotFoundError, OSError) as error: + raise SlurmStateError(f"cannot update attempt {attempt.attempt_id!r}") from error + + def load_attempts(self, shard_id: ShardId) -> tuple[AttemptManifest, ...]: + """Load one shard's attempts in ordinal order.""" + normalized_shard_id = self._validate_shard_id(shard_id) + return self._reader.load_attempts(normalized_shard_id) + + def load_attempt(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptManifest: + """Load one persisted attempt.""" + normalized_shard_id = self._validate_shard_id(shard_id) + normalized_attempt_id = self._validate_attempt_id(attempt_id) + return self._reader.load_attempt(normalized_shard_id, normalized_attempt_id) + + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: + """Create or atomically replace one validated readiness snapshot.""" + self._validate_readiness_location(readiness) + try: + return self._write_readiness_with_lock(readiness) + except (StateConflictError, StateCorruptionError, StateNotFoundError): + raise + except (PlanStateContractError, StateContractError) as error: + raise StateConflictError("readiness update is not a valid monotonic transition") from error + except (FileNotFoundError, OSError) as error: + raise SlurmStateError(f"cannot persist readiness for attempt {readiness.attempt_id!r}") from error + + def load_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptReadiness: + """Load one attempt's latest readiness snapshot.""" + normalized_shard_id = self._validate_shard_id(shard_id) + normalized_attempt_id = self._validate_attempt_id(attempt_id) + return self._reader.load_readiness(normalized_shard_id, normalized_attempt_id) + + def publish_attempt_result( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + """Publish one result pair and atomically bind its candidate to the attempt.""" + return self._results.publish(client_result, candidate) + + @contextmanager + def acquire_dataset_workspace( + self, + shard_id: ShardId, + attempt_id: AttemptId, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Iterator[Path]: + """Yield one validated dataset path while holding its shard lease.""" + normalized_shard_id = self._validate_shard_id(shard_id) + normalized_attempt_id = self._validate_attempt_id(attempt_id) + if type(resume_mode) is not str or resume_mode not in {"never", "always", "if_possible"}: + raise StateConflictError("resume mode must be 'never', 'always', or 'if_possible'") + with self._finalizer.acquire_dataset_workspace( + normalized_shard_id, + normalized_attempt_id, + resume_mode, + ) as dataset_path: + yield dataset_path + + def finalize_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + *, + published_at: datetime, + completed_at: datetime | None = None, + ) -> ShardWinner: + """Validate artifacts and optionally commit attempt success with its winner.""" + normalized_shard_id = self._validate_shard_id(shard_id) + normalized_attempt_id = self._validate_attempt_id(attempt_id) + if ( + not isinstance(published_at, datetime) + or published_at.tzinfo is None + or published_at.utcoffset() != timedelta(0) + ): + raise StateConflictError("winner publication timestamp must be timezone-aware UTC") + if completed_at is not None and ( + not isinstance(completed_at, datetime) + or completed_at.tzinfo is None + or completed_at.utcoffset() != timedelta(0) + ): + raise StateConflictError("attempt completion timestamp must be timezone-aware UTC") + return self._finalizer.finalize_winner( + normalized_shard_id, + normalized_attempt_id, + published_at, + completed_at, + ) + + def load_winner(self, shard_id: ShardId) -> ShardWinner: + """Load and validate one shard's immutable winner chain.""" + return self._finalizer.load_winner(self._validate_shard_id(shard_id)) + + def resume_incomplete_finalization( + self, + shard_id: ShardId, + *, + published_at: datetime, + ) -> ShardWinner | None: + """Finish winner publication for a successful attempt after process interruption.""" + normalized_shard_id = self._validate_shard_id(shard_id) + try: + return self._finalizer.load_winner(normalized_shard_id) + except StateNotFoundError: + pass + successful = tuple( + attempt + for attempt in self._reader.load_attempts(normalized_shard_id) + if attempt.state is AttemptLifecycleState.SUCCEEDED + ) + if not successful: + return None + if len(successful) != 1: + raise StateCorruptionError(f"shard {normalized_shard_id!r} has multiple successful attempts") + attempt = successful[0] + return self.finalize_winner( + normalized_shard_id, + attempt.attempt_id, + published_at=max(published_at, attempt.updated_at), + ) + + def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: + with self._storage.acquire_resume_and_shard_locks(attempt.shard_id): + run, plan, shard = self._reader.load_shard_context(attempt.shard_id) + shard_attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + existing = next((item for item in shard_attempts if item.attempt_id == attempt.attempt_id), None) + if existing is not None: + if existing != attempt: + raise StateConflictError(f"attempt {attempt.attempt_id!r} already contains different state") + self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) + return existing + self._finalizer.require_no_winner(run, plan, shard, shard_attempts) + expected_ordinal = len(shard_attempts) + 1 + if attempt.attempt_ordinal != expected_ordinal or attempt.attempt_id != f"attempt-{expected_ordinal:04d}": + raise StateConflictError("attempt identity does not match the next shard ordinal") + self._reader.validate_attempt_against_plan(run, plan, shard, attempt) + validate_shard_attempt_set(run, shard, shard_attempts + (attempt,)) + self._storage.publish_attempt(attempt) + return attempt + + def _update_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: + with self._storage.acquire_shard_lock(attempt.shard_id): + run, plan, shard = self._reader.load_shard_context(attempt.shard_id) + shard_attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + previous = self._reader.get_attempt(shard_attempts, attempt.attempt_id) + if previous == attempt: + self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) + return previous + validate_attempt_transition(previous, attempt) + if previous.candidate_output is None and attempt.candidate_output is not None: + raise StateContractError("candidate output must be bound by result publication") + self._reader.validate_attempt_against_plan(run, plan, shard, attempt) + updated_attempts = tuple( + attempt if item.shard_id == attempt.shard_id and item.attempt_id == attempt.attempt_id else item + for item in shard_attempts + ) + validate_shard_attempt_set(run, shard, updated_attempts) + self._storage.replace_attempt(attempt) + return attempt + + def _write_readiness_with_lock(self, readiness: AttemptReadiness) -> AttemptReadiness: + with self._storage.acquire_shard_lock(readiness.shard_id): + run, plan, shard = self._reader.load_shard_context(readiness.shard_id) + shard_attempts = self._reader.load_validated_shard_attempts(run, plan, shard) + attempt = self._reader.get_attempt(shard_attempts, readiness.attempt_id) + previous = self._load_optional_readiness(readiness) + if previous is None: + PersistedPlanStateValidator(plan).validate_initial_readiness(attempt, readiness) + self._storage.publish_readiness(readiness) + return readiness + self._validate_persisted_readiness(plan, attempt, previous) + if previous == readiness: + self._storage.sync_attempt_directory(readiness.shard_id, readiness.attempt_id) + return previous + validate_readiness_transition(previous, readiness) + self._storage.replace_readiness(readiness) + return readiness + + def _load_optional_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness | None: + try: + return self._storage.read_readiness(readiness.shard_id, readiness.attempt_id) + except FileNotFoundError: + return None + + @staticmethod + def _validate_persisted_readiness( + plan: ResolvedSlurmRunPlan, + attempt: AttemptManifest, + readiness: AttemptReadiness, + ) -> None: + try: + PersistedPlanStateValidator(plan).validate_readiness_snapshot(attempt, readiness) + except PlanStateContractError as error: + raise StateCorruptionError(f"attempt {readiness.attempt_id!r} has invalid persisted readiness") from error + + def _validate_initial_state( + self, + authored_config: DataDesignerSlurmConfig, + resolved_plan: ResolvedSlurmRunPlan, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> None: + try: + self._validate_initial_record_types(authored_config, resolved_plan, run, shards) + self._validate_initial_bindings(authored_config, resolved_plan, run) + PersistedPlanStateValidator(resolved_plan).validate_plan_shards(run, shards) + except (PlanStateContractError, StateContractError) as error: + raise StateConflictError("run initialization does not match resolved plan intent") from error + + @staticmethod + def _validate_initial_record_types( + authored_config: DataDesignerSlurmConfig, + resolved_plan: ResolvedSlurmRunPlan, + run: RunManifest, + shards: tuple[ShardManifest, ...], + ) -> None: + if not isinstance(authored_config, DataDesignerSlurmConfig): + raise StateContractError("authored config has an invalid type") + if not isinstance(resolved_plan, ResolvedSlurmRunPlan): + raise StateContractError("resolved plan has an invalid type") + if not isinstance(run, RunManifest): + raise StateContractError("run manifest has an invalid type") + if not isinstance(shards, tuple) or any(not isinstance(shard, ShardManifest) for shard in shards): + raise StateContractError("shard manifests have an invalid type") + + def _validate_initial_bindings( + self, + authored_config: DataDesignerSlurmConfig, + resolved_plan: ResolvedSlurmRunPlan, + run: RunManifest, + ) -> None: + if run.run_id != self._run_id or resolved_plan.run_id != self._run_id: + raise StateContractError("run identity does not match the state writer") + if resolved_plan.selected_profile.profile.workspace_root != self._storage.logical_workspace_root.as_posix(): + raise StateContractError("resolved plan workspace does not match the state writer") + if resolved_plan.authored_config.sha256 != authored_config.compute_sha256(): + raise StateContractError("authored config digest does not match the resolved plan") + if run.authored_config != resolved_plan.authored_config: + raise StateContractError("run authored config does not match the resolved plan") + if run.authored_config.path != self._storage.logical_authored_config_path.as_posix(): + raise StateContractError("run authored config reference does not match its persisted location") + if ( + run.resolved_plan.path != self._storage.logical_resolved_plan_path.as_posix() + or run.resolved_plan.sha256 != resolved_plan.compute_sha256() + ): + raise StateContractError("run resolved plan reference does not match persisted plan bytes") + + def _validate_attempt_location(self, attempt: AttemptManifest) -> None: + if not isinstance(attempt, AttemptManifest) or attempt.run_id != self._run_id: + raise StateConflictError("attempt identity does not match the state writer") + + def _validate_readiness_location(self, readiness: AttemptReadiness) -> None: + if not isinstance(readiness, AttemptReadiness) or readiness.run_id != self._run_id: + raise StateConflictError("readiness identity does not match the state writer") + + @staticmethod + def _validate_shard_id(shard_id: ShardId) -> ShardId: + try: + return _SHARD_ID_ADAPTER.validate_python(shard_id, strict=True) + except ValidationError as error: + raise SlurmStateError("invalid shard identity") from error + + @staticmethod + def _validate_attempt_id(attempt_id: AttemptId) -> AttemptId: + try: + return _ATTEMPT_ID_ADAPTER.validate_python(attempt_id, strict=True) + except ValidationError as error: + raise SlurmStateError("invalid attempt identity") from error diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py new file mode 100644 index 000000000..99a5a2d04 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared recovery policy for ambiguous Slurm submission receipts.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Protocol + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.launcher.models import SlurmSubmissionMatch +from data_designer.slurm.state.errors import SlurmStateError, StateConflictError + +SUBMISSION_VISIBILITY_WINDOW = timedelta(minutes=5) + + +class SubmissionLookup(Protocol): + """Scheduler lookup needed to recover one immutable submission plan.""" + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + """Return allocations matching one exact submission name.""" + ... + + +@dataclass(frozen=True) +class PreparedSubmission: + """Immutable scheduler correlation facts for one prepared operation.""" + + job_name: Identifier + submitted_after: datetime + reconciliation_deadline: datetime + expected_array_task_ids: tuple[int, ...] | None + + +def resolve_prepared_submission( + scheduler: SubmissionLookup, + prepared: PreparedSubmission, + *, + observed_at: datetime, +) -> int | None: + """Return one recovered job ID, or ``None`` after definitive bounded absence.""" + try: + matches = scheduler.query_submissions_by_name( + prepared.job_name, + submitted_after=prepared.submitted_after, + ) + except SlurmLauncherError as error: + raise SlurmStateError("cannot reconcile ambiguous Slurm submission") from error + if len(matches) > 1: + raise StateConflictError("multiple scheduler jobs match the prepared submission") + if matches: + match = matches[0] + if match.array_task_ids == prepared.expected_array_task_ids: + return match.job_id + if _is_partial_array_view(match.array_task_ids, prepared.expected_array_task_ids): + if observed_at <= prepared.reconciliation_deadline: + raise StateConflictError("prepared submission is still being reconciled") + return None + raise StateConflictError("scheduler job shape does not match the prepared submission") + if observed_at <= prepared.reconciliation_deadline: + raise StateConflictError("prepared submission is still being reconciled") + return None + + +def _is_partial_array_view( + observed: tuple[int, ...] | None, + expected: tuple[int, ...] | None, +) -> bool: + return observed is not None and expected is not None and set(observed) < set(expected) + + +__all__ = [ + "SUBMISSION_VISIBILITY_WINDOW", + "PreparedSubmission", + "SubmissionLookup", + "resolve_prepared_submission", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py new file mode 100644 index 000000000..1e20fc7c4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.outputs import ( + CandidateOutputManifest, + CollectionPlan, + ShardWinner, +) +from data_designer.slurm.state.scheduler import ( + SchedulerObservation, + SchedulerState, + is_scheduler_terminal_state, +) + +_ATTEMPT_STATE_ORDER = { + AttemptLifecycleState.CREATED: 0, + AttemptLifecycleState.SUBMITTED: 1, + AttemptLifecycleState.PENDING: 2, + AttemptLifecycleState.RUNNING: 3, +} +_TERMINAL_ATTEMPT_STATES = frozenset( + { + AttemptLifecycleState.SUCCEEDED, + AttemptLifecycleState.FAILED, + } +) + + +class StateContractError(ValueError): + """Raised when related state records violate a cross-record contract.""" + + +def validate_shard_manifest(run: RunManifest, shard: ShardManifest) -> ShardManifest: + """Validate a shard against the immutable run identity and bounds.""" + _require(shard.run_id == run.run_id, "shard run_id does not match run manifest") + _require(shard.shard_index < run.shard_count, "shard_index is outside the run shard count") + _require(shard.created_at >= run.created_at, "shard creation cannot precede run creation") + return shard + + +def validate_shard_set( + run: RunManifest, + shards: tuple[ShardManifest, ...], +) -> tuple[ShardManifest, ...]: + """Validate ordered run state shards without duplicating resolved-plan coverage.""" + _require(len(shards) == run.shard_count, "shard set must include exactly the run shard count") + for shard in shards: + validate_shard_manifest(run, shard) + + shard_ids = tuple(shard.shard_id for shard in shards) + resume_workspace_paths = tuple(shard.resume_workspace.path for shard in shards) + shard_indices = tuple(shard.shard_index for shard in shards) + _require(len(set(shard_ids)) == len(shards), "shard IDs must be unique") + _require(len(set(resume_workspace_paths)) == len(shards), "resume workspace paths must be unique") + _require( + shard_indices == tuple(range(run.shard_count)), + "shards must be ordered by a complete zero-based shard index", + ) + for previous, current in zip(shards, shards[1:]): + _require( + current.record_range.start_index >= previous.record_range.end_index_exclusive, + "shard record ranges must not overlap", + ) + return shards + + +def validate_attempt_manifest( + run: RunManifest, + shard: ShardManifest, + attempt: AttemptManifest, +) -> AttemptManifest: + """Validate an attempt against its run, shard, and resolved plan.""" + validate_shard_manifest(run, shard) + _require(attempt.run_id == run.run_id, "attempt run_id does not match run manifest") + _require(attempt.shard_id == shard.shard_id, "attempt shard_id does not match shard manifest") + _require( + attempt.resolved_plan == run.resolved_plan, + "attempt resolved plan does not match run manifest", + ) + _require(attempt.created_at >= shard.created_at, "attempt creation cannot precede shard creation") + return attempt + + +def validate_attempt_set( + run: RunManifest, + shards: tuple[ShardManifest, ...], + attempts: tuple[AttemptManifest, ...], +) -> tuple[AttemptManifest, ...]: + """Validate attempt identities and scheduler ownership across a run.""" + validate_shard_set(run, shards) + shard_by_id = {shard.shard_id: shard for shard in shards} + + for shard in shards: + validate_shard_attempt_set( + run, + shard, + tuple(attempt for attempt in attempts if attempt.shard_id == shard.shard_id), + ) + for attempt in attempts: + shard = shard_by_id.get(attempt.shard_id) + _require(shard is not None, f"attempt references unknown shard {attempt.shard_id!r}") + + attempt_ids = tuple((attempt.shard_id, attempt.attempt_id) for attempt in attempts) + shard_ordinals = tuple((attempt.shard_id, attempt.attempt_ordinal) for attempt in attempts) + scheduler_identities = tuple(attempt.scheduler for attempt in attempts if attempt.scheduler is not None) + _require(len(set(attempt_ids)) == len(attempts), "attempt IDs must be unique within each shard") + _require( + len(set(shard_ordinals)) == len(attempts), + "attempt ordinals must be unique within each shard", + ) + _require( + len(set(scheduler_identities)) == len(scheduler_identities), + "scheduler identities must be unique across attempts", + ) + return attempts + + +def validate_shard_attempt_set( + run: RunManifest, + shard: ShardManifest, + attempts: tuple[AttemptManifest, ...], +) -> tuple[AttemptManifest, ...]: + """Validate attempt identities and scheduler ownership within one shard.""" + validate_shard_manifest(run, shard) + for attempt in attempts: + validate_attempt_manifest(run, shard, attempt) + attempt_ids = tuple(attempt.attempt_id for attempt in attempts) + ordinals = tuple(attempt.attempt_ordinal for attempt in attempts) + scheduler_identities = tuple(attempt.scheduler for attempt in attempts if attempt.scheduler is not None) + _require(len(set(attempt_ids)) == len(attempts), "attempt IDs must be unique within each shard") + _require(len(set(ordinals)) == len(attempts), "attempt ordinals must be unique within each shard") + _require( + len(set(scheduler_identities)) == len(scheduler_identities), + "scheduler identities must be unique within each shard", + ) + return attempts + + +def validate_attempt_transition( + previous: AttemptManifest, + current: AttemptManifest, +) -> AttemptManifest: + """Validate immutable identity and monotonic lifecycle updates for an attempt.""" + for field_name in ( + "run_id", + "shard_id", + "attempt_id", + "attempt_ordinal", + "resolved_plan", + "created_at", + ): + _require( + getattr(previous, field_name) == getattr(current, field_name), + f"attempt {field_name} cannot change", + ) + _require(current.updated_at >= previous.updated_at, "attempt updated_at cannot move backward") + + if previous.state in _TERMINAL_ATTEMPT_STATES: + _require(current == previous, "terminal attempt manifests are immutable") + return current + + if current.state not in _TERMINAL_ATTEMPT_STATES: + _require( + _ATTEMPT_STATE_ORDER[current.state] >= _ATTEMPT_STATE_ORDER[previous.state], + f"attempt state cannot move from {previous.state.value} to {current.state.value}", + ) + if previous.scheduler is not None: + _require(current.scheduler == previous.scheduler, "attempt scheduler identity cannot change") + if previous.candidate_output is not None: + _require(current.candidate_output == previous.candidate_output, "attempt candidate output cannot change") + return current + + +def validate_scheduler_observation_transition( + previous: SchedulerObservation, + current: SchedulerObservation, +) -> SchedulerObservation: + """Validate scheduler identity, chronology, and a fixed accounting-lag deadline.""" + _require(current.scheduler == previous.scheduler, "scheduler identity cannot change between observations") + _require(current.observed_at >= previous.observed_at, "scheduler observed_at cannot move backward") + if is_scheduler_terminal_state(previous.state): + _require(current.state is previous.state, "terminal scheduler evidence cannot change") + + if previous.state is SchedulerState.ACCOUNTING_LAG: + deadline = previous.reconciliation_deadline + _require(deadline is not None, "accounting lag has no reconciliation deadline") + if current.state is SchedulerState.ACCOUNTING_LAG: + _require( + current.reconciliation_deadline == deadline, + "accounting-lag reconciliation deadline cannot change", + ) + elif current.state is SchedulerState.UNKNOWN: + _require( + current.observed_at > deadline, + "accounting lag cannot become unknown before its reconciliation deadline expires", + ) + if previous.state is SchedulerState.PREEMPTED and previous.reconciliation_deadline is not None: + if current.state is SchedulerState.PREEMPTED and current.reconciliation_deadline is not None: + _require( + current.reconciliation_deadline == previous.reconciliation_deadline, + "preemption reconciliation deadline cannot change", + ) + return current + + +def validate_shard_winner( + run: RunManifest, + shard: ShardManifest, + attempt: AttemptManifest, + candidate: CandidateOutputManifest, + winner: ShardWinner, + *, + existing_winner: ShardWinner | None = None, +) -> ShardWinner: + """Validate first-writer winner publication without performing persistence.""" + validate_attempt_manifest(run, shard, attempt) + _require(existing_winner is None, "a shard winner is immutable once published") + + for record_name, record_run_id in ( + ("candidate", candidate.run_id), + ("winner", winner.run_id), + ): + _require(record_run_id == run.run_id, f"{record_name} run_id does not match run manifest") + for record_name, record_shard_id in ( + ("candidate", candidate.shard_id), + ("winner", winner.shard_id), + ): + _require(record_shard_id == shard.shard_id, f"{record_name} shard_id does not match shard manifest") + + _require(candidate.attempt_id == attempt.attempt_id, "candidate attempt_id does not match attempt") + _require(winner.attempt_id == attempt.attempt_id, "winner attempt_id does not match attempt") + _require( + candidate.attempt_ordinal == attempt.attempt_ordinal == winner.attempt_ordinal, + "candidate and winner attempt ordinals must match the attempt", + ) + _require(attempt.state is AttemptLifecycleState.SUCCEEDED, "only successful attempts may win") + _require( + attempt.terminal_classification is AttemptTerminalClassification.SUCCEEDED, + "only successfully classified attempts may win", + ) + _require(candidate.winner_eligible, "candidate output does not satisfy winner policy") + _require( + candidate.requested_records == shard.record_range.record_count, + "candidate requested record count does not match shard record range", + ) + _require(candidate.created_at >= attempt.created_at, "candidate creation cannot precede attempt creation") + _require(attempt.updated_at >= candidate.created_at, "attempt completion cannot precede candidate creation") + _require(winner.published_at >= attempt.updated_at, "winner publication cannot precede attempt completion") + _require( + winner.candidate_manifest.sha256 == candidate.compute_sha256(), + "winner candidate digest does not match candidate manifest", + ) + _require( + attempt.candidate_output == winner.candidate_manifest, + "attempt candidate output reference does not match winner", + ) + return winner + + +def validate_collection_plan( + run: RunManifest, + plan: CollectionPlan, + shards: tuple[ShardManifest, ...], + winners: tuple[ShardWinner, ...], +) -> CollectionPlan: + """Validate collection against the exact run shard set and winner digests.""" + validate_shard_set(run, shards) + _require(plan.run_id == run.run_id, "collection run_id does not match run manifest") + _require(plan.resolved_plan == run.resolved_plan, "collection resolved plan does not match run manifest") + _require(plan.created_at >= run.created_at, "collection creation cannot precede run creation") + + expected_shard_ids = tuple(shard.shard_id for shard in shards) + planned_shard_ids = tuple(shard.shard_id for shard in plan.planned_shards) + _require( + planned_shard_ids == expected_shard_ids, + "collection planned shards must exactly match the ordered run shard set", + ) + _require(len(winners) == len(shards), "winner set must include exactly the run shard count") + + winner_by_shard = {winner.shard_id: winner for winner in winners} + _require(len(winner_by_shard) == len(winners), "winner shard IDs must be unique") + _require(set(winner_by_shard) == set(expected_shard_ids), "winner shard IDs must exactly match the run shard set") + for planned_shard in plan.planned_shards: + winner = winner_by_shard[planned_shard.shard_id] + _require(winner.run_id == run.run_id, "winner run_id does not match run manifest") + _require(plan.created_at >= winner.published_at, "collection creation cannot precede winner publication") + _require( + planned_shard.winner_manifest.sha256 == winner.compute_sha256(), + f"winner digest mismatch for shard {planned_shard.shard_id!r}", + ) + return plan + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise StateContractError(message) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/types.py b/packages/data-designer-slurm/src/data_designer/slurm/types.py new file mode 100644 index 000000000..2eafa605f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/types.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared constrained scalar types for Slurm configuration and records.""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import BeforeValidator, Field, StringConstraints + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] + + +def _normalize_partition_selection(value: object) -> object: + if isinstance(value, list): + if not value or not all(isinstance(name, str) for name in value): + raise ValueError("partitions must be a nonempty list of names") + value = ",".join(value) + if isinstance(value, str) and len(set(value.split(","))) != len(value.split(",")): + raise ValueError("partition names must be unique") + return value + + +PartitionSelection = Annotated[ + str, + BeforeValidator(_normalize_partition_selection), + StringConstraints( + min_length=1, + max_length=1024, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}(?:,[A-Za-z0-9][A-Za-z0-9._-]{0,127})*$", + ), +] +ShardId = Annotated[str, StringConstraints(pattern=r"^shard-[0-9]{5,}$")] +AttemptId = Annotated[str, StringConstraints(pattern=r"^attempt-[0-9]{4,}$")] +SchemaVersion = Literal[1] +EnvironmentName = Annotated[str, StringConstraints(pattern=r"^[A-Za-z_][A-Za-z0-9_]*$")] +Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] +Duration = Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:s|m|h|d)$")] +NonNegativeDuration = Annotated[str, StringConstraints(pattern=r"^(?:0|[1-9][0-9]*)(?:s|m|h|d)$")] +NetworkPort = Annotated[int, Field(ge=1024, le=65535)] + +__all__ = [ + "AttemptId", + "Duration", + "EnvironmentName", + "Identifier", + "NetworkPort", + "NonNegativeDuration", + "PartitionSelection", + "SchemaVersion", + "Sha256Digest", + "ShardId", +] 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..3c1c2a3b1 --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py @@ -0,0 +1,101 @@ +# 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 all(case.child_run_config.array_tasks.count == 1 for case in compiled.cases) + assert all(case.child_run_config.array_tasks.max_concurrent is None 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 + + +def test_default_record_policy_scales_high_concurrency( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + payload = benchmark_config.model_dump(mode="json") + payload.pop("record_policy") + payload["concurrency_values"] = [512, 1024, 8192] + payload["deployment_cases"] = payload["deployment_cases"][:1] + config = DataDesignerSlurmBenchmarkConfig.model_validate(payload) + + compiled = BenchmarkCompiler.compile(config, authored_run) + + assert [case.child_run_config.invocation.num_records for case in compiled.cases] == [1000, 1024, 5000] + assert all(case.child_run_config.array_tasks.count == 1 for case in compiled.cases) + + payload["record_policy"] = {"type": "adaptive", "records_per_concurrency": 4.0} + override = DataDesignerSlurmBenchmarkConfig.model_validate(payload) + overridden = BenchmarkCompiler.compile(override, authored_run) + + assert [case.child_run_config.invocation.num_records for case in overridden.cases] == [2048, 4096, 5000] 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..6fb48843d --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -0,0 +1,781 @@ +# 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, + FixedRecordPolicy, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.services import ( + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) + +NOW = datetime(2026, 8, 19, 12, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + ("gpu_start_index", "topology_digest"), + [ + (None, "19380747e9269b80f4d64da44d67a1e010eb61923c34365792df18aab01bc9b7"), + (0, "59b41934fbfddc75c4e8d268533737d38c11580e6a041ca27e00ecd457e83dad"), + (1, "205f985652cf46653c02b5958d95e3b2fb2775dcec56a014ecdde9bd86093080"), + ], +) +def test_shared_benchmark_counts_physical_nodes_and_allocated_gpu_budget( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + shared_node_config: DataDesignerSlurmConfig, + shared_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, + gpu_start_index: int | None, + topology_digest: str, +) -> None: + payload = shared_node_config.model_dump(mode="python") + payload["deployments"][0]["resources"]["gpus_per_node"] = 2 + payload["deployments"][1]["resources"]["gpu_start_index"] = gpu_start_index + base_run = DataDesignerSlurmConfig.model_validate(payload) + benchmark_payload = benchmark_config.model_dump(mode="python") + benchmark_payload.update( + base_run={"inline": base_run.model_dump(mode="python")}, + model_aliases=["writer", "judge"], + concurrency_values=[1], + record_policy={"type": "fixed", "records": 4}, + deployment_cases=[{"name": "shared", "deployments": {"writer": {"nodes": 1, "nodes_per_replica": 1}}}], + ) + config = DataDesignerSlurmBenchmarkConfig.model_validate(benchmark_payload) + compiled = BenchmarkCompiler.compile(config, base_run) + store, manifest = _manifest(tmp_path, config, compiled) + case = compiled.cases[0] + expected = shared_node_plan + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + case.child_run_config, + selected_profile=expected.selected_profile, + client_runtime=expected.client.runtime, + deployment_images=tuple(deployment.image for deployment in expected.deployments), + dependency_lock=dependency_lock_single, + runtime_bundle=expected.runtime_bundle, + run_id=case.child_run_id, + package_version=expected.package_version, + ) + ) + observer = _Observer( + { + case.child_run_id: BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=plan, + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements(4, 1, 2, 3, 1), + ) + } + ) + report = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, base_run, manifest, store.manifest_reference(manifest), refresh_state=False + ) + assert sum(deployment.gpus_per_node for deployment in plan.deployments) == 6 + assert report.cases[0].nodes_per_job == 1 + assert report.cases[0].gpus_per_job == 8 + assert report.cases[0].gpu_hours_per_job == 32 + assert report.cases[0].topology_digest == topology_digest + + +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): + assert (self.benchmark_root / "benchmark.json").is_file() + self.calls.append((self.run_id, config, source_root, dry_run)) + 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): + self.second_check.wait(timeout=1) + return super().execute(config, source_root=source_root, dry_run=dry_run) + + +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() + progress: list[str] = [] + 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__, + progress=progress.append, + ) + + 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 progress == [ + "Benchmark manifest saved; checking 2 child runs.", + "Benchmark case 1/2: submitted.", + "Benchmark case 2/2: submitted.", + "Benchmark launch: 2 submitted, 0 already submitted, 0 failed.", + "Benchmark manifest saved; checking 2 child runs.", + "Benchmark case 1/2: already submitted.", + "Benchmark case 2/2: already submitted.", + "Benchmark launch: 0 submitted, 2 already submitted, 0 failed.", + ] + + +def test_progress_output_failure_cannot_change_benchmark_submissions( + 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 + + def broken_progress(_: str) -> None: + raise RuntimeError("progress stream failed") + + 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__, + progress=broken_progress, + ) + + backend.run(config, source_root=tmp_path, force=False) + + assert len(calls) == len(compiled.cases) + assert submissions == {case.child_run_id for case in compiled.cases} + + +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} + progress: list[str] = [] + 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__, + progress=progress.append, + ) + + with pytest.raises(SlurmServiceError, match="1 of 2") as failure: + backend.run(config, source_root=tmp_path, force=True) + assert len(calls) == 2 + assert compiled.benchmark_id in str(failure.value) + assert missing_id in str(failure.value) + assert "submission unavailable" in str(failure.value) + assert ( + f"Benchmark case 1/2 submission failed: {compiled.cases[0].case_id} " + f"(run {missing_id}): submission unavailable; continuing." + ) in progress + assert "Benchmark launch: 1 submitted, 0 already submitted, 1 failed." in progress + + 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_uses_actual_rows_from_nonempty_partial_case( + 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] + requested_records = case.child_run_config.invocation.num_records + actual_records = requested_records - 9 + 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( + actual_records, + 60, + 120, + 180, + 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, + ) + + result = report.cases[0] + assert result.outcome is BenchmarkOutcome.SUCCEEDED + assert result.requested_records == requested_records + assert result.actual_records == actual_records + assert result.rows_per_second == actual_records / 120 + assert result.feasible is True + assert result.target_jobs is not 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_accepts_saved_benchmark_from_previous_compiler( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + legacy = BenchmarkCompiler._compile(config, authored_run, legacy=True) + assert legacy.benchmark_id == "benchmark-996c47990258ceb4a3ca7b245ddee971" + assert legacy.cases[0].child_run_id == "run-996c47990258ceb4-0000-89d3e0eccd48" + assert all(case.child_run_config.array_tasks.count == 2 for case in legacy.cases) + store, manifest = _manifest(tmp_path, config, legacy) + store.publish(config, authored_run, manifest) + saved_config, saved_base_run, saved_manifest = store.load() + observer = _Observer( + {case.child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING) for case in legacy.cases} + ) + + report = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + saved_config, + saved_base_run, + saved_manifest, + store.manifest_reference(saved_manifest), + refresh_state=False, + ) + + assert report.benchmark_id == legacy.benchmark_id + assert tuple(case.outcome for case in report.cases) == (BenchmarkOutcome.MISSING,) * len(legacy.cases) + assert tuple(case.requested_records for case in report.cases) == (1000,) * len(legacy.cases) + + +def test_analysis_accepts_single_task_case_below_base_array_count( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run).model_copy( + update={"record_policy": FixedRecordPolicy(type="fixed", records=1)} + ) + compiled = BenchmarkCompiler.compile(config, authored_run) + assert authored_run.array_tasks.count == 2 + assert all(case.child_run_config.array_tasks.count == 1 for case in compiled.cases) + store, manifest = _manifest(tmp_path, config, compiled) + observer = _Observer( + {case.child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING) for case in compiled.cases} + ) + + report = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert tuple(case.requested_records for case in report.cases) == (1,) * len(compiled.cases) + + +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/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py new file mode 100644 index 000000000..6ce539561 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import platform +import sys +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from packaging.tags import interpreter_name, interpreter_version + +from data_designer.config import ResumeMode +from data_designer.slurm.client.environment import PreparedClientEnvironment +from data_designer.slurm.client.records import ClientInstallerOutcome +from data_designer.slurm.contracts import InstalledDistribution, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.ports import resolve_allocation_plan + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" + + +@dataclass(frozen=True) +class ClientWorkerCase: + plan: ResolvedSlurmRunPlan + plan_path: Path + lock: ResolvedDependencyLock + attempt_dir: Path + scratch_root: Path + prepared: PreparedClientEnvironment + endpoints: dict[str, str] + + +class FakeCreationResults: + def __init__( + self, + dataset_path: Path, + *, + requested: int, + actual: int, + resume: ResumeMode, + ) -> None: + self.dataset_path = dataset_path + self.requested_num_records = requested + self.actual_num_records = actual + self.early_shutdown = actual < requested + self.requested_resume_mode = resume + self.effective_resume_mode = ResumeMode.NEVER + self.artifact_storage: SimpleNamespace = SimpleNamespace(final_dataset_path=dataset_path / "parquet-files") + + def export(self, path: Path, *, format: str) -> Path: + assert format == "parquet" + pq.write_table(pa.table({"value": tuple(range(self.actual_num_records))}), path) + return path + + +class FakeDataDesigner: + def __init__( + self, + *, + actual_records: int | None = None, + create_error: BaseException | None = None, + create_error_after_batch: BaseException | None = None, + effective_resume: ResumeMode = ResumeMode.NEVER, + batch_count: int = 1, + schema_metadata: dict[bytes, bytes] | None = None, + mixed_batch_schemas: bool = False, + **kwargs: object, + ) -> None: + self.actual_records = actual_records + self.create_error = create_error + self.create_error_after_batch = create_error_after_batch + self.effective_resume = effective_resume + self.batch_count = batch_count + self.schema_metadata = schema_metadata + self.mixed_batch_schemas = mixed_batch_schemas + self.initialization = kwargs + self.run_config = None + self.validated_builder = None + + def set_run_config(self, run_config: object) -> None: + self.run_config = run_config + + def validate(self, config_builder: object) -> None: + self.validated_builder = config_builder + + def create( + self, + config_builder: object, + *, + num_records: int, + dataset_name: str, + resume: ResumeMode, + artifact_path: Path, + on_batch_complete: object, + ) -> FakeCreationResults: + del config_builder + if self.create_error is not None: + raise self.create_error + actual = num_records if self.actual_records is None else self.actual_records + dataset_path = artifact_path / dataset_name + dataset_path.mkdir(parents=True, exist_ok=True) + dataset_path.chmod(0o700) + batch_dir = dataset_path / "parquet-files" + batch_dir.mkdir(exist_ok=True) + for index in range(self.batch_count): + start = actual * index // self.batch_count + end = actual * (index + 1) // self.batch_count + values = tuple(range(start, end)) + if self.mixed_batch_schemas and index == 1: + table = pa.table({"value": tuple(float(value) for value in values)}) + else: + table = pa.table({"value": values}) + if self.schema_metadata is not None: + table = table.replace_schema_metadata(self.schema_metadata) + batch_path = batch_dir / f"batch_{index:05d}.parquet" + pq.write_table(table, batch_path) + on_batch_complete(batch_path) + if self.create_error_after_batch is not None: + raise self.create_error_after_batch + results = FakeCreationResults(dataset_path, requested=num_records, actual=actual, resume=resume) + results.effective_resume_mode = self.effective_resume + return results + + +@pytest.fixture(autouse=True) +def allocation_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SLURM_JOB_GPUS", "0") + + +@pytest.fixture +def client_worker_case(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ClientWorkerCase: + workspace = tmp_path / "workspace" + payload = json.loads( + (GOLDEN_DIRECTORY / "single_node_plan.json").read_text().replace("/workspace/primary", workspace.as_posix()) + ) + python_abi = f"{interpreter_name()}{interpreter_version()}" + runtime = payload["client"]["runtime"] + runtime.update( + { + "python_abi": python_abi, + "python_executable": sys.executable, + "python_implementation": platform.python_implementation().lower(), + "python_version": platform.python_version(), + } + ) + runtime["runtime_sha256"] = compute_canonical_json_sha256( + {key: value for key, value in runtime.items() if key != "runtime_sha256"} + ) + lock_payload = json.loads((GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text()) + lock_payload["python_abi"] = python_abi + lock_payload["client_runtime_sha256"] = runtime["runtime_sha256"] + lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) + payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + mount = {"source": workspace.as_posix(), "target": workspace.as_posix(), "read_only": False} + payload["selected_profile"]["profile"]["container_mounts"] = [mount] + payload["container_mounts"] = [mount] + payload["selected_profile"]["profile_sha256"] = compute_canonical_json_sha256( + payload["selected_profile"]["profile"] + ) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + plan_path = workspace / "runs" / plan.run_id / "resolved-plan.json" + plan_path.parent.mkdir(parents=True) + plan_path.write_text(plan.serialize_json()) + lock_path = Path(plan.client.dependency_lock.path) + lock_path.write_text(lock.serialize_json()) + Path(plan.invocation.effective_input_bindings.managed_assets_path).mkdir(parents=True) + shard = plan.shards[0] + attempt_dir = Path(shard.resume_workspace.path).parent / "attempts" / "attempt-0001" + attempt_dir.mkdir(parents=True) + scratch_root = tmp_path / "scratch" + monkeypatch.setenv("DATA_DESIGNER_SLURM_SCRATCH_ROOT", scratch_root.as_posix()) + overlay_path = scratch_root / "client-env" / "site-packages" + overlay_path.mkdir(parents=True) + prepared = PreparedClientEnvironment( + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_dir=attempt_dir, + scratch_root=scratch_root, + overlay_path=overlay_path, + dependency_lock=plan.client.dependency_lock, + client_runtime_sha256=plan.client.runtime.runtime_sha256, + python_abi=lock.python_abi, + installer_outcome=ClientInstallerOutcome.NOT_REQUIRED, + installed_distributions=tuple( + InstalledDistribution(name=distribution.name, version=distribution.version) + for distribution in lock.base_distributions + ), + ) + allocation_plan = resolve_allocation_plan(plan, {"SLURM_JOB_GPUS": "0"}) + endpoints = { + allocation_plan.deployments[ + 0 + ].authored.model_alias: f"http://127.0.0.1:{allocation_plan.client.ports[0].port}/v1" + } + return ClientWorkerCase(plan, plan_path, lock, attempt_dir, scratch_root, prepared, endpoints) diff --git a/packages/data-designer-slurm/tests/client/test_dependency_resolution.py b/packages/data-designer-slurm/tests/client/test_dependency_resolution.py new file mode 100644 index 000000000..bc5dec95f --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_dependency_resolution.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path + +import pytest + +from data_designer.slurm.client.dependencies import ( + ClientDependencyResolutionError, + ClientDependencyResolver, +) +from data_designer.slurm.config import ClientDependencies +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import LockedPackage, ResolvedDependencyLock, ResolvedSlurmRunPlan + + +def test_empty_requirements_resolve_without_invoking_pip( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + def unexpected_runner(*_args: object) -> None: + raise AssertionError("pip must not run") + + resolver = ClientDependencyResolver(command_runner=unexpected_runner) + dependencies = ClientDependencies(requirements=[]) + + with resolver.resolve( + dependencies, + single_node_plan.client.runtime, + run_root=tmp_path / "workspace/runs/run-001", + source_root=tmp_path, + ) as resolved: + assert resolved.lock.authored_requirements == () + assert resolved.lock.overlay_packages == () + assert resolved.lock.client_runtime_sha256 == single_node_plan.client.runtime.runtime_sha256 + assert resolved.wheel_sources == () + assert resolved.lock_source is None + + +def test_inline_requirements_resolve_pure_wheels_and_omit_image_packages( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + token = "DUMMY_INDEX_TOKEN" + calls: list[tuple[tuple[str, ...], dict[str, str]]] = [] + + def download(command: tuple[str, ...], _cwd: Path, environment: dict[str, str]) -> None: + calls.append((command, environment)) + destination = Path(next(item.removeprefix("--dest=") for item in command if item.startswith("--dest="))) + (destination / "example_plugin-1.2.0-py3-none-any.whl").write_bytes(b"plugin") + (destination / "pip-26.1-py3-none-any.whl").write_bytes(b"image package") + + dependencies = ClientDependencies( + requirements=["example-plugin==1.2.0"], + index_credentials={ + "private-index": {"type": "secret", "environment": "PACKAGE_INDEX_TOKEN"}, + }, + ) + resolver = ClientDependencyResolver( + command_runner=download, + environ={ + "PACKAGE_INDEX_TOKEN": token, + "PIP_CONFIG_FILE": "/ambient/pip.conf", + "PIP_EXTRA_INDEX_URL": "https://ambient.example/simple", + "PIP_INDEX_URL": "https://ambient.example/simple", + "SAFE_VALUE": "preserved", + }, + python_executable="/python", + ) + run_root = tmp_path / "workspace/runs/run-001" + + with resolver.resolve( + dependencies, + single_node_plan.client.runtime, + run_root=run_root, + source_root=tmp_path, + ) as resolved: + assert calls[0][0][0:4] == ("/python", "-m", "pip", "download") + assert token not in calls[0][0] + assert "--index-url=https://pypi.org/simple" in calls[0][0] + assert calls[0][1] == {"PIP_CONFIG_FILE": os.devnull, "SAFE_VALUE": "preserved"} + assert tuple(package.name for package in resolved.lock.overlay_packages) == ("example-plugin",) + package = resolved.lock.overlay_packages[0] + assert package.artifact.path == (run_root / "dependencies/example_plugin-1.2.0-py3-none-any.whl").as_posix() + assert package.artifact.sha256 == hashlib.sha256(b"plugin").hexdigest() + assert len(resolved.wheel_sources) == 1 + assert resolved.wheel_sources[0].is_file() + + +def test_inline_resolution_rejects_platform_specific_wheels( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + def download(command: tuple[str, ...], _cwd: Path, _environment: dict[str, str]) -> None: + destination = Path(next(item.removeprefix("--dest=") for item in command if item.startswith("--dest="))) + (destination / "example_plugin-1.2.0-cp312-cp312-manylinux_2_28_x86_64.whl").write_bytes(b"plugin") + + resolver = ClientDependencyResolver(command_runner=download) + + with ( + pytest.raises(ClientDependencyResolutionError, match="platform-independent"), + resolver.resolve( + ClientDependencies(requirements=["example-plugin==1.2.0"]), + single_node_plan.client.runtime, + run_root=tmp_path / "workspace/runs/run-001", + source_root=tmp_path, + ), + ): + pass + + +def test_supplied_lock_is_verified_and_rebound_to_run_inputs( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + wheel = tmp_path / "example_plugin-1.2.0-py3-none-any.whl" + wheel.write_bytes(b"plugin") + runtime = single_node_plan.client.runtime + supplied = ResolvedDependencyLock( + schema_version=1, + resolver_version="external-1", + python_abi=runtime.python_abi, + client_runtime_sha256=runtime.runtime_sha256, + authored_requirements=("example-plugin==1.2.0",), + base_distributions=tuple(sorted(runtime.distributions, key=lambda item: item.name)), + overlay_packages=( + LockedPackage( + name="example-plugin", + version="1.2.0", + artifact=ArtifactReference( + path=wheel.as_posix(), + sha256=hashlib.sha256(b"plugin").hexdigest(), + ), + ), + ), + ) + lock_file = tmp_path / "lock.json" + lock_file.write_text(supplied.serialize_json(), encoding="utf-8") + run_root = tmp_path / "workspace/runs/run-001" + + with ClientDependencyResolver().resolve( + ClientDependencies(requirements=None, lock_file="lock.json"), + single_node_plan.client.runtime, + run_root=run_root, + source_root=tmp_path, + ) as resolved: + assert resolved.lock.authored_source == "lock.json" + assert resolved.lock.source == ArtifactReference( + path=(run_root / "inputs/lock.json").as_posix(), + sha256=hashlib.sha256(supplied.serialize_json().encode()).hexdigest(), + ) + assert ( + resolved.lock.overlay_packages[0].artifact.path + == (run_root / "dependencies/example_plugin-1.2.0-py3-none-any.whl").as_posix() + ) + assert resolved.wheel_sources == (wheel,) + assert resolved.lock_source == lock_file + + +def test_supplied_lock_rejects_platform_specific_wheels( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + wheel = tmp_path / "example_plugin-1.2.0-cp312-cp312-manylinux_2_28_x86_64.whl" + wheel.write_bytes(b"plugin") + runtime = single_node_plan.client.runtime + supplied = ResolvedDependencyLock( + schema_version=1, + resolver_version="external-1", + python_abi=runtime.python_abi, + client_runtime_sha256=runtime.runtime_sha256, + authored_requirements=("example-plugin==1.2.0",), + base_distributions=tuple(sorted(runtime.distributions, key=lambda item: item.name)), + overlay_packages=( + LockedPackage( + name="example-plugin", + version="1.2.0", + artifact=ArtifactReference( + path=wheel.as_posix(), + sha256=hashlib.sha256(b"plugin").hexdigest(), + ), + ), + ), + ) + lock_file = tmp_path / "lock.json" + lock_file.write_text(supplied.serialize_json(), encoding="utf-8") + + with ( + pytest.raises(ClientDependencyResolutionError, match="artifact differs"), + ClientDependencyResolver().resolve( + ClientDependencies(requirements=None, lock_file="lock.json"), + single_node_plan.client.runtime, + run_root=tmp_path / "workspace/runs/run-001", + source_root=tmp_path, + ), + ): + pass + + +def test_resolution_redacts_runner_failure_and_missing_credentials( + tmp_path: Path, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + def fail(*_args: object) -> None: + raise RuntimeError("DUMMY_SECRET_VALUE") + + resolver = ClientDependencyResolver(command_runner=fail) + dependencies = ClientDependencies(requirements=["example-plugin==1.2.0"]) + with pytest.raises(ClientDependencyResolutionError, match="resolution failed") as captured: + with resolver.resolve( + dependencies, + single_node_plan.client.runtime, + run_root=tmp_path / "workspace/runs/run-001", + source_root=tmp_path, + ): + pass + assert "DUMMY_SECRET_VALUE" not in str(captured.value) + + credentialed = ClientDependencies( + requirements=["example-plugin==1.2.0"], + index_credentials={"private": {"type": "secret", "environment": "MISSING_TOKEN"}}, + ) + with pytest.raises(ClientDependencyResolutionError, match="credential is unavailable"): + with resolver.resolve( + credentialed, + single_node_plan.client.runtime, + run_root=tmp_path / "workspace/runs/run-002", + source_root=tmp_path, + ): + pass diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py new file mode 100644 index 000000000..7146a5c10 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -0,0 +1,556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import importlib +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import cast +from unittest.mock import Mock + +import pytest +from conftest import ClientWorkerCase, FakeDataDesigner + +from data_designer.slurm.client import environment as environment_module +from data_designer.slurm.client.environment import ClientEnvironmentBuilder, activate_environment, inspect_distributions +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.execution import ClientWorker +from data_designer.slurm.client.plugins import discover_plugins +from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome +from data_designer.slurm.config import SlurmProfile +from data_designer.slurm.contracts import InstalledDistribution, canonical_json, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan + + +def test_runtime_inventory_parser_preserves_vcs_provenance() -> None: + provenance_sha256 = "a" * 64 + + assert environment_module._parse_distributions( # noqa: SLF001 + [{"name": "data-designer", "version": "1.0.0", "provenance_sha256": provenance_sha256}] + ) == (InstalledDistribution(name="data-designer", version="1.0.0", provenance_sha256=provenance_sha256),) + + +def test_environment_prepares_empty_verified_overlay(client_worker_case: ClientWorkerCase) -> None: + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.base_distributions if path is None else () + + prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( + client_worker_case.plan_path, + shard_id=client_worker_case.plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + + assert prepared.installer_outcome is ClientInstallerOutcome.NOT_REQUIRED + assert prepared.installed_distributions == client_worker_case.lock.base_distributions + assert prepared.overlay_path == client_worker_case.scratch_root / "client-env/site-packages" + assert not (client_worker_case.attempt_dir / "client-env").exists() + + +def test_environment_activation_keeps_home_and_cache_in_scratch( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sys, "path", sys.path.copy()) + + activate_environment(client_worker_case.prepared) + + assert os.environ["DATA_DESIGNER_HOME"] == (client_worker_case.scratch_root / "data-designer-home").as_posix() + assert os.environ["XDG_CACHE_HOME"] == (client_worker_case.scratch_root / "cache").as_posix() + assert not (client_worker_case.attempt_dir / "data-designer-home").exists() + + +def test_environment_rejects_symlink_scratch_root( + client_worker_case: ClientWorkerCase, + tmp_path: Path, +) -> None: + outside = tmp_path / "outside" + outside.mkdir() + scratch = tmp_path / "scratch-link" + scratch.symlink_to(outside, target_is_directory=True) + + with pytest.raises(ClientWorkerError) as error: + ClientEnvironmentBuilder().prepare( + client_worker_case.plan_path, + shard_id=client_worker_case.plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=scratch, + ) + + assert error.value.code is ClientErrorCode.INVALID_INPUT + + +def test_native_client_ignores_container_mount_mapping(client_worker_case: ClientWorkerCase) -> None: + physical_workspace = client_worker_case.plan_path.parents[2] + payload = cast(dict[str, object], client_worker_case.plan.model_dump(mode="json")) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mount = {"source": physical_workspace.as_posix(), "target": "/container/workspace", "read_only": False} + profile_payload["container_mounts"] = [mount] + payload["container_mounts"] = [mount] + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.base_distributions if path is None else () + + prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( + client_worker_case.plan_path, + shard_id=plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.dataset_path.startswith(physical_workspace.as_posix()) + assert result.candidate_output_manifest is not None + assert result.candidate_output_manifest.path.startswith(physical_workspace.as_posix()) + + +def test_native_environment_rejects_container_only_plan_path( + client_worker_case: ClientWorkerCase, + tmp_path: Path, +) -> None: + physical_workspace = client_worker_case.plan_path.parents[2] + physical_plan = tmp_path / "fast-plan" / "resolved-plan.json" + logical_workspace = "/host/workspace" + payload = cast( + dict[str, object], + json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), + ) + logical_plan = f"{logical_workspace}/runs/{payload['run_id']}/resolved-plan.json" + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mounts = [ + {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, + { + "source": logical_plan, + "target": physical_plan.as_posix(), + "read_only": False, + }, + ] + profile_payload["container_mounts"] = mounts + payload["container_mounts"] = mounts + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + physical_plan.parent.mkdir(parents=True) + physical_plan.write_text(plan.serialize_json()) + shard_id = plan.shards[0].shard_id + attempt_dir = physical_workspace / "runs" / plan.run_id / "shards" / shard_id / "attempts" / "attempt-0001" + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.base_distributions if path is None else () + + with pytest.raises(ClientWorkerError) as error: + ClientEnvironmentBuilder(inventory=inventory).prepare( + physical_plan, + shard_id=shard_id, + attempt_id="attempt-0001", + attempt_dir=attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + + assert error.value.code is ClientErrorCode.INVALID_INPUT + + +@pytest.mark.parametrize("user_site_enabled", (False, True)) +def test_inspect_distributions_uses_active_environment_paths( + monkeypatch: pytest.MonkeyPatch, + user_site_enabled: bool, +) -> None: + calls: list[dict[str, object]] = [] + + def distributions(**kwargs: object) -> tuple[()]: + calls.append(kwargs) + return () + + submit_path = "/submit/venv/site-packages" + interpreter_path = "/control/venv/site-packages" + user_site_path = "/workspace/alice/.local/site-packages" + pth_path = "/installed/editable-source" + monkeypatch.setattr(sys, "path", [submit_path, interpreter_path, user_site_path, pth_path]) + monkeypatch.setattr(environment_module.site, "getsitepackages", lambda: [interpreter_path]) + monkeypatch.setattr(environment_module.site, "getusersitepackages", lambda: user_site_path) + monkeypatch.setattr(environment_module.site, "ENABLE_USER_SITE", user_site_enabled) + monkeypatch.setattr( + environment_module.subprocess, + "check_output", + lambda *args, **kwargs: json.dumps([submit_path, interpreter_path, pth_path]), + ) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join((submit_path, user_site_path))) + monkeypatch.setattr("data_designer.slurm.client.environment.importlib.metadata.distributions", distributions) + + assert inspect_distributions(None) == () + expected_paths = [interpreter_path, pth_path] + assert calls == [{"path": expected_paths}] + + +@pytest.mark.parametrize("user_site_enabled", (False, True)) +@pytest.mark.parametrize("path_kind", ("repeated", "symlink", "separate")) +def test_inspect_distributions_deduplicates_install_paths( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, user_site_enabled: bool, path_kind: str +) -> None: + site_path = tmp_path / "site-packages" + metadata = site_path / "example-1.0.0.dist-info" + metadata.mkdir(parents=True) + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: example\nVersion: 1.0.0\n") + alias_path = tmp_path / "alias" + if path_kind == "symlink": + alias_path.symlink_to(site_path, target_is_directory=True) + elif path_kind == "repeated": + alias_path = site_path + else: + shutil.copytree(site_path, alias_path) + paths = [str(site_path), str(alias_path)] + monkeypatch.setattr(sys, "path", paths) + monkeypatch.setattr(environment_module.site, "ENABLE_USER_SITE", user_site_enabled) + monkeypatch.setattr(environment_module.site, "getsitepackages", lambda: paths) + monkeypatch.setattr(environment_module.site, "getusersitepackages", lambda: str(tmp_path / "user-site")) + monkeypatch.setattr(environment_module.subprocess, "check_output", lambda *args, **kwargs: json.dumps(paths)) + monkeypatch.delenv("PYTHONPATH", raising=False) + + if path_kind == "separate": + with pytest.raises(ClientWorkerError, match="installed distribution inventory is invalid"): + inspect_distributions(None) + else: + assert inspect_distributions(None) == (InstalledDistribution(name="example", version="1.0.0"),) + + +def test_inspect_distributions_refreshes_overlay_with_unchanged_directory_mtime(tmp_path: Path) -> None: + original_stat = tmp_path.stat() + assert inspect_distributions(tmp_path) == () + distribution = tmp_path / "cache_probe-1.0.0.dist-info" + distribution.mkdir() + (distribution / "METADATA").write_text( + "Metadata-Version: 2.1\nName: cache-probe\nVersion: 1.0.0\n", + encoding="utf-8", + ) + os.utime(tmp_path, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns)) + + assert inspect_distributions(tmp_path) == (InstalledDistribution(name="cache-probe", version="1.0.0"),) + + +def test_inspect_distributions_rejects_unhashed_direct_url(monkeypatch: pytest.MonkeyPatch) -> None: + distribution = Mock( + metadata={"Name": "example"}, + version="1.0.0", + ) + distribution.read_text.return_value = '{"url":"https://example.test/example.whl"}' + monkeypatch.setattr( + "data_designer.slurm.client.environment.importlib.metadata.distributions", + lambda **kwargs: (distribution,), + ) + + with pytest.raises(ClientWorkerError) as error: + inspect_distributions(None) + + assert error.value.code is ClientErrorCode.DEPENDENCY_CONFLICT + + +def test_inspect_distributions_accepts_immutable_git_provenance(monkeypatch: pytest.MonkeyPatch) -> None: + direct_url = { + "url": "https://github.com/example/project.git", + "vcs_info": { + "vcs": "git", + "commit_id": "a" * 40, + "requested_revision": "feature-branch", + }, + } + distribution = Mock(metadata={"Name": "example"}, version="1.0.0") + distribution.read_text.return_value = json.dumps(direct_url) + monkeypatch.setattr( + "data_designer.slurm.client.environment.importlib.metadata.distributions", + lambda **kwargs: (distribution,), + ) + + installed = inspect_distributions(None) + + assert installed == ( + InstalledDistribution( + name="example", + version="1.0.0", + provenance_sha256=hashlib.sha256(canonical_json(direct_url)).hexdigest(), + ), + ) + + +@pytest.mark.parametrize( + "direct_url", + ( + {"url": "file:///workspace/project", "dir_info": {"editable": True}}, + { + "url": "https://github.com/example/project.git", + "vcs_info": {"vcs": "git", "commit_id": "abcdef", "requested_revision": "main"}, + }, + ), +) +def test_inspect_distributions_rejects_mutable_install_provenance( + monkeypatch: pytest.MonkeyPatch, + direct_url: dict[str, object], +) -> None: + distribution = Mock(metadata={"Name": "example"}, version="1.0.0") + distribution.read_text.return_value = json.dumps(direct_url) + monkeypatch.setattr( + "data_designer.slurm.client.environment.importlib.metadata.distributions", + lambda **kwargs: (distribution,), + ) + + with pytest.raises(ClientWorkerError) as error: + inspect_distributions(None) + + assert error.value.code is ClientErrorCode.DEPENDENCY_CONFLICT + + +def test_environment_rejects_client_runtime_inventory_conflict(client_worker_case: ClientWorkerCase) -> None: + conflicting = (InstalledDistribution(name="unexpected", version="1.0"),) + + with pytest.raises(ClientWorkerError) as error: + ClientEnvironmentBuilder(inventory=lambda _: conflicting).prepare( + client_worker_case.plan_path, + shard_id=client_worker_case.plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + + assert error.value.code is ClientErrorCode.DEPENDENCY_CONFLICT + + +@pytest.mark.parametrize( + ("shard_id", "attempt_id", "attempt_dir"), + ( + ("../escaped", "attempt-0001", "../escaped/attempts/attempt-0001"), + ("shard-00000", "attempt-0000", "shard-00000/attempts/attempt-0000"), + ), +) +def test_environment_rejects_invalid_identity_before_creating_attempt( + client_worker_case: ClientWorkerCase, + shard_id: str, + attempt_id: str, + attempt_dir: str, +) -> None: + path = client_worker_case.plan_path.parent / "shards" / attempt_dir + + with pytest.raises(ClientWorkerError) as error: + ClientEnvironmentBuilder().prepare( + client_worker_case.plan_path, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_dir=path, + scratch_root=client_worker_case.scratch_root, + ) + + assert error.value.code is ClientErrorCode.INVALID_INPUT + assert not path.exists() + + +def test_environment_rejects_missing_locked_wheel(client_worker_case: ClientWorkerCase) -> None: + plan_payload = client_worker_case.plan.model_dump(mode="json") + lock_payload = client_worker_case.lock.model_dump(mode="json") + wheel_path = client_worker_case.plan_path.parent / "dependencies" / "missing_plugin-1.0.0-py3-none-any.whl" + lock_payload["authored_requirements"] = ["missing-plugin==1.0.0"] + lock_payload["overlay_packages"] = [ + { + "name": "missing-plugin", + "version": "1.0.0", + "artifact": {"path": wheel_path.as_posix(), "sha256": "a" * 64}, + } + ] + lock = client_worker_case.lock.model_validate_json(json.dumps(lock_payload)) + Path(client_worker_case.plan.client.dependency_lock.path).write_text(lock.serialize_json()) + plan_payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(plan_payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return lock.base_distributions if path is None else () + + with pytest.raises(ClientWorkerError) as error: + ClientEnvironmentBuilder(inventory=inventory).prepare( + client_worker_case.plan_path, + shard_id=plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + + assert error.value.code is ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING + + +def test_environment_installs_verified_wheel_overlay(client_worker_case: ClientWorkerCase) -> None: + plan_payload = client_worker_case.plan.model_dump(mode="json") + lock_payload = client_worker_case.lock.model_dump(mode="json") + wheel_path = client_worker_case.plan_path.parent / "dependencies" / "example_plugin-1.0.0-py3-none-any.whl" + wheel_path.parent.mkdir() + wheel_path.write_bytes(b"immutable wheel artifact") + lock_payload["authored_requirements"] = ["example-plugin==1.0.0"] + lock_payload["overlay_packages"] = [ + { + "name": "example-plugin", + "version": "1.0.0", + "artifact": { + "path": wheel_path.as_posix(), + "sha256": hashlib.sha256(wheel_path.read_bytes()).hexdigest(), + }, + } + ] + lock = client_worker_case.lock.model_validate_json(json.dumps(lock_payload)) + Path(client_worker_case.plan.client.dependency_lock.path).write_text(lock.serialize_json()) + plan_payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(plan_payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + commands: list[tuple[str, ...]] = [] + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + if path is None: + return lock.base_distributions + return (InstalledDistribution(name="example-plugin", version="1.0.0"),) if commands else () + + prepared = ClientEnvironmentBuilder(inventory=inventory, command_runner=commands.append).prepare( + client_worker_case.plan_path, + shard_id=plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_dir=client_worker_case.attempt_dir, + scratch_root=client_worker_case.scratch_root, + ) + + assert prepared.installer_outcome is ClientInstallerOutcome.INSTALLED + assert "--no-index" in commands[0] + assert "--no-deps" in commands[0] + + +def test_overlay_reuse_checks_installed_files_and_rejects_strays(tmp_path: Path) -> None: + target = tmp_path / "scratch" / "client-env" / "site-packages" + expected = (InstalledDistribution(name="example-plugin", version="1.0.0"),) + calls: list[tuple[str, ...]] = [] + + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return expected if calls else () + + def install(command: tuple[str, ...]) -> None: + calls.append(command) + (target / "example_plugin.py").write_text("original") + + builder = ClientEnvironmentBuilder(inventory=inventory, command_runner=install) + assert builder._install_overlay(Path(sys.executable), (tmp_path / "plugin.whl",), expected, target) is ( # noqa: SLF001 + ClientInstallerOutcome.INSTALLED + ) + assert builder._install_overlay(Path(sys.executable), (tmp_path / "plugin.whl",), expected, target) is ( # noqa: SLF001 + ClientInstallerOutcome.REUSED + ) + + generated_bytecode = target / "__pycache__" / "example_plugin.cpython-311.pyc" + generated_bytecode.parent.mkdir() + generated_bytecode.write_bytes(b"unverified cached code") + assert builder._install_overlay(Path(sys.executable), (tmp_path / "plugin.whl",), expected, target) is ( # noqa: SLF001 + ClientInstallerOutcome.REUSED + ) + assert not generated_bytecode.exists() + + (target / "example_plugin.py").write_text("changed") + with pytest.raises(ClientWorkerError, match="overlay files differ"): + builder._install_overlay(Path(sys.executable), (tmp_path / "plugin.whl",), expected, target) # noqa: SLF001 + + (target / "example_plugin.py").write_text("original") + (target / "unexpected.py").write_text("stray") + with pytest.raises(ClientWorkerError, match="overlay files differ"): + builder._install_overlay(Path(sys.executable), (tmp_path / "plugin.whl",), expected, target) # noqa: SLF001 + + +def test_discover_plugins_loads_verified_entry_point( + fake_plugin_overlay: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(fake_plugin_overlay.as_posix()) + importlib.invalidate_caches() + + plugins = discover_plugins((InstalledDistribution(name="fake-data-designer-plugin", version="1.0.0"),)) + + assert len(plugins) == 1 + assert plugins[0].plugin_name == "fake-slurm-column" + + +def test_worker_bootstrap_loads_overlay_plugin_before_config_import( + fake_plugin_overlay: Path, + tmp_path: Path, +) -> None: + script = """ +import sys +from pathlib import Path + +import data_designer.slurm.client.worker as worker + +assert "data_designer.config.column_types" not in sys.modules +from data_designer.slurm.client.environment import PreparedClientEnvironment +from data_designer.slurm.client.records import ClientInstallerOutcome +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution + +prepared = PreparedClientEnvironment( + run_id="run-test", + shard_id="shard-00000", + attempt_id="attempt-0001", + attempt_dir=Path(sys.argv[2]), + scratch_root=Path(sys.argv[3]), + overlay_path=Path(sys.argv[1]), + dependency_lock=ArtifactReference(path=(Path(sys.argv[2]) / "dependency-lock.json").as_posix(), sha256="a" * 64), + client_runtime_sha256="b" * 64, + python_abi="test", + installer_outcome=ClientInstallerOutcome.REUSED, + installed_distributions=(InstalledDistribution(name="fake-data-designer-plugin", version="1.0.0"),), +) +worker.activate_environment(prepared) + +from data_designer.slurm.client.plugins import discover_plugins + +plugins = discover_plugins(prepared.installed_distributions) +assert plugins[0].plugin_name == "fake-slurm-column" + +from data_designer.config import DataDesignerConfigBuilder + +builder = DataDesignerConfigBuilder.from_config( + {"data_designer": {"columns": [{"name": "custom", "column_type": "fake-slurm-column"}], "model_configs": []}} +) +assert builder.get_column_configs()[0].column_type == "fake-slurm-column" +""" + scratch_root = tmp_path / "scratch" + scratch_root.mkdir() + result = subprocess.run( + [ + sys.executable, + "-c", + script, + fake_plugin_overlay.as_posix(), + (tmp_path / "attempt").as_posix(), + scratch_root.as_posix(), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr diff --git a/packages/data-designer-slurm/tests/client/test_filesystem.py b/packages/data-designer-slurm/tests/client/test_filesystem.py new file mode 100644 index 000000000..a45ad4f92 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_filesystem.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import data_designer.slurm.client.filesystem as filesystem_module +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ensure_private_directory, publish_private_text, replace_private_text +from data_designer.slurm.client.records import ClientErrorCode + + +def test_ensure_private_directory_rejects_symlink(tmp_path: Path) -> None: + target = tmp_path / "target" + target.mkdir() + link = tmp_path / "link" + link.symlink_to(target, target_is_directory=True) + + with pytest.raises(ClientWorkerError) as error: + ensure_private_directory(link) + + assert error.value.code is ClientErrorCode.INVALID_INPUT + + +def test_ensure_private_directory_does_not_follow_intermediate_symlink(tmp_path: Path) -> None: + target = tmp_path / "target" + target.mkdir() + link = tmp_path / "link" + link.symlink_to(target, target_is_directory=True) + + with pytest.raises(ClientWorkerError) as error: + ensure_private_directory(link / "nested") + + assert error.value.code is ClientErrorCode.INVALID_INPUT + assert not (target / "nested").exists() + + +@pytest.mark.parametrize("operation_name", ("read_regular_bytes", "compute_file_sha256")) +def test_regular_file_operation_rejects_swap_to_symlink( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + operation_name: str, +) -> None: + path = tmp_path / "input.json" + path.write_text("original") + replacement = tmp_path / "replacement.json" + replacement.write_text("replacement") + original_open = filesystem_module.os.open + swapped = False + + def swapping_open(name: object, flags: int, mode: int = 0o777, *, dir_fd: int | None = None) -> int: + nonlocal swapped + if name == path.name and dir_fd is not None and not swapped: + swapped = True + path.unlink() + path.symlink_to(replacement) + return original_open(name, flags, mode, dir_fd=dir_fd) + + monkeypatch.setattr(filesystem_module.os, "open", swapping_open) + + with pytest.raises(ClientWorkerError) as error: + getattr(filesystem_module, operation_name)(path, missing_code=ClientErrorCode.INVALID_INPUT) + + assert error.value.code is ClientErrorCode.INVALID_INPUT + + +def test_publish_private_text_is_convergent_and_immutable(tmp_path: Path) -> None: + path = tmp_path / "record.json" + publish_private_text(path, "same") + publish_private_text(path, "same") + + with pytest.raises(ClientWorkerError) as error: + publish_private_text(path, "different") + + assert error.value.code is ClientErrorCode.OUTPUT_INVALID + assert path.read_text() == "same" + + +@pytest.mark.parametrize("interruption", (KeyboardInterrupt, SystemExit)) +def test_replace_private_text_cleans_temporary_file_on_interruption( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + interruption: type[BaseException], +) -> None: + def interrupt_write(_descriptor: int, _payload: object) -> int: + raise interruption + + monkeypatch.setattr(filesystem_module.os, "write", interrupt_write) + + with pytest.raises(interruption): + replace_private_text(tmp_path / "record.json", "value") + + assert not tuple(tmp_path.glob(".client.*.tmp")) diff --git a/packages/data-designer-slurm/tests/client/test_process.py b/packages/data-designer-slurm/tests/client/test_process.py new file mode 100644 index 000000000..98450e5fa --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_process.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.client.process import ClientWorkerProcess + + +def test_client_worker_process_uses_a_fresh_python_interpreter() -> None: + commands: list[tuple[str, ...]] = [] + process = ClientWorkerProcess( + executable="/client/python", + executor=lambda command: commands.append(command) or 17, + ) + + assert process.run(("run", "--plan", "/workspace/resolved-plan.json")) == 17 + assert commands == [ + ( + "/client/python", + "-m", + "data_designer.slurm.client.worker", + "run", + "--plan", + "/workspace/resolved-plan.json", + ) + ] diff --git a/packages/data-designer-slurm/tests/client/test_runtime.py b/packages/data-designer-slurm/tests/client/test_runtime.py new file mode 100644 index 000000000..787fc5adf --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_runtime.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import venv +from pathlib import Path + +import pytest + +from data_designer.slurm.client import environment as environment_module +from data_designer.slurm.client import runtime as runtime_module +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.records import ClientErrorCode +from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector +from data_designer.slurm.contracts import InstalledDistribution, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.preflight import SystemAllocationPreflight +from data_designer.slurm.state import AttemptManifest + + +@pytest.fixture +def distribution_sites(tmp_path: Path) -> tuple[Path, Path]: + base, user = tmp_path / "site-packages", tmp_path / "user-site" + base.mkdir() + user.mkdir() + for name in ( + "aiohttp", + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + ): + version = "3.14.3" if name == "aiohttp" else "1.0.0" + metadata = base / f"{name.replace('-', '_')}-{version}.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text(f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n") + metadata = user / "irrelevant-1.0.0.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: irrelevant\nVersion: 1.0.0\n") + return base, user + + +@pytest.fixture +def runtime_sites(distribution_sites: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + base, user = distribution_sites + monkeypatch.setattr(sys, "path", [str(base), str(user)]) + monkeypatch.setattr(environment_module.site, "getsitepackages", lambda: [str(base)]) + monkeypatch.setattr(environment_module.site, "getusersitepackages", lambda: str(user)) + monkeypatch.setattr(environment_module.site, "ENABLE_USER_SITE", True) + monkeypatch.setattr(environment_module.subprocess, "check_output", lambda *args, **kwargs: json.dumps([str(base)])) + return base, user + + +@pytest.mark.parametrize("required_user_only", (False, True)) +def test_runtime_inspector_excludes_user_site_pth_paths( + distribution_sites: tuple[Path, Path], + tmp_path: Path, + required_user_only: bool, +) -> None: + base, extra = distribution_sites + environment = tmp_path / "venv" + venv.EnvBuilder(with_pip=False, system_site_packages=True, symlinks=True).create(environment) + python = environment / "bin/python" + user_base = tmp_path / "user" + variables = dict(os.environ, PYTHONUSERBASE=str(user_base)) + variables.pop("PYTHONNOUSERSITE", None) + variables.pop("PYTHONPATH", None) + paths = json.loads( + subprocess.check_output( + ( + str(python), + "-c", + "import importlib.metadata,json,site; print(json.dumps([site.getsitepackages()[0],site.getusersitepackages(), [item.metadata['Name'].lower().replace('_','-') for item in importlib.metadata.distributions()]]))", + ), + text=True, + env=variables, + ) + ) + installed, user_site = map(Path, paths[:2]) + if required_user_only: + metadata = base / "data_designer-1.0.0.dist-info" + metadata.rename(extra / metadata.name) + for metadata in base.iterdir(): + if metadata.name.split("-")[0].replace("_", "-") not in paths[2]: + shutil.copytree(metadata, installed / metadata.name) + user_site.mkdir(parents=True) + (user_site / "extra.pth").write_text(f"{extra}\n") + working_directory = tmp_path / "submit" + working_directory.mkdir() + metadata = working_directory / "submit_only-1.0.0.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: submit-only\nVersion: 1.0.0\n") + packages = Path(__file__).parents[3] + dependencies = ( + Path(sys.executable).parents[1] / f"lib/python{sys.version_info.major}.{sys.version_info.minor}/site-packages" + ) + variables["PYTHONPATH"] = os.pathsep.join( + ( + *( + str(packages / name / "src") + for name in ("data-designer-config", "data-designer-engine", "data-designer", "data-designer-slurm") + ), + str(dependencies), + ) + ) + script = tmp_path / "inspect_runtime.py" + script.write_text( + "from data_designer.slurm.client.runtime import ClientRuntimeInspector\nprint(ClientRuntimeInspector().inspect().model_dump_json())" + ) + submitted = subprocess.run( + (str(python), str(script)), cwd=working_directory, env=variables, capture_output=True, text=True + ) + if required_user_only: + assert submitted.returncode != 0 + assert "user-site disabled: data-designer" in submitted.stderr + else: + assert submitted.returncode == 0, submitted.stderr + assert {"irrelevant", "submit-only"}.isdisjoint( + item["name"] for item in json.loads(submitted.stdout)["distributions"] + ) + allocated = subprocess.run( + (str(python), "-s", str(script)), cwd=environment, env=variables, capture_output=True, text=True + ) + assert allocated.returncode == 0, allocated.stderr + assert json.loads(allocated.stdout) == json.loads(submitted.stdout) + + +@pytest.mark.parametrize( + "failure", + (OSError("private"), subprocess.CalledProcessError(1, "private"), subprocess.TimeoutExpired("private", 10)), +) +def test_runtime_inspector_reports_isolated_path_failure( + runtime_sites: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + failure: Exception, +) -> None: + def fail(*args: object, **kwargs: object) -> str: + raise failure + + monkeypatch.setattr(environment_module.subprocess, "check_output", fail) + with pytest.raises(ClientRuntimeInspectionError, match="user-site disabled") as error: + ClientRuntimeInspector().inspect() + assert "private" not in str(error.value) + + +@pytest.mark.parametrize("pythonpath", (False, True)) +def test_runtime_inspection_matches_batch_without_user_site( + runtime_sites: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + single_node_plan: ResolvedSlurmRunPlan, + attempt_manifest: AttemptManifest, + pythonpath: bool, +) -> None: + base, user = runtime_sites + if pythonpath: + monkeypatch.setenv("PYTHONPATH", str(user)) + else: + monkeypatch.delenv("PYTHONPATH", raising=False) + submitted = ClientRuntimeInspector().inspect() + assert "irrelevant" not in {item.name for item in submitted.distributions} + monkeypatch.setattr(sys, "path", [str(base)]) + monkeypatch.setattr(environment_module.site, "ENABLE_USER_SITE", False) + plan = single_node_plan.model_copy( + update={"client": single_node_plan.client.model_copy(update={"runtime": submitted})} + ) + attempt_directory = ( + Path(plan.authored_config.path).parent + / "shards" + / plan.shards[0].shard_id + / "attempts" + / attempt_manifest.attempt_id + ) + context = AllocationContext( + plan=plan, shard=plan.shards[0], attempt=attempt_manifest, attempt_directory=attempt_directory + ) + SystemAllocationPreflight.verify_client_runtime(context) + + +@pytest.mark.parametrize("name", ("aiohttp", "data-designer", "pip")) +def test_runtime_inspector_rejects_required_user_only_distribution( + runtime_sites: tuple[Path, Path], + name: str, +) -> None: + base, user = runtime_sites + metadata = next(base.glob(f"{name.replace('-', '_')}-*.dist-info")) + metadata.rename(user / metadata.name) + with pytest.raises(ClientRuntimeInspectionError, match="user-site disabled") as failure: + ClientRuntimeInspector().inspect() + assert name in str(failure.value) + assert "install" in str(failure.value) + + +def test_runtime_inspector_binds_active_interpreter_and_required_distributions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distributions = tuple( + InstalledDistribution(name=name, version="1.0.0") + for name in ( + "aiohttp", + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + ) + ) + distributions = ( + InstalledDistribution(name="aiohttp", version="3.14.3"), + *(distribution for distribution in distributions if distribution.name != "aiohttp"), + ) + monkeypatch.setattr(runtime_module, "inspect_distributions", lambda path: distributions) + + runtime = ClientRuntimeInspector().inspect() + + assert runtime.python_executable == sys.executable + assert runtime.distributions == distributions + assert runtime.runtime_sha256 == compute_canonical_json_sha256( + runtime.model_dump(mode="json", exclude={"runtime_sha256"}) + ) + + +def test_runtime_inspector_rejects_missing_required_distribution(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + runtime_module, + "inspect_distributions", + lambda path: (InstalledDistribution(name="pip", version="1.0.0"),), + ) + + with pytest.raises(ClientRuntimeInspectionError, match="required client distributions"): + ClientRuntimeInspector().inspect() + + +def test_runtime_inspector_preserves_sanitized_inventory_error(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_inventory(_path: Path | None) -> tuple[InstalledDistribution, ...]: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, + "mutable installed distribution 'example-plugin' is forbidden", + ) + + monkeypatch.setattr(runtime_module, "inspect_distributions", fail_inventory) + + with pytest.raises(ClientRuntimeInspectionError, match="example-plugin"): + ClientRuntimeInspector().inspect() + + +def test_runtime_inspector_rejects_incompatible_proxy_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + distributions = tuple( + InstalledDistribution(name=name, version="3.13.0" if name == "aiohttp" else "1.0.0") + for name in ( + "aiohttp", + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + ) + ) + monkeypatch.setattr(runtime_module, "inspect_distributions", lambda path: distributions) + + with pytest.raises(ClientRuntimeInspectionError, match="aiohttp>=3.14.3,<4"): + ClientRuntimeInspector().inspect() + + +def test_runtime_inspector_rejects_mixed_data_designer_versions(monkeypatch: pytest.MonkeyPatch) -> None: + distributions = tuple( + InstalledDistribution( + name=name, + version=("3.14.3" if name == "aiohttp" else "0.9.3" if name == "data-designer-slurm" else "0.9.2"), + ) + for name in ( + "aiohttp", + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + "pip", + ) + ) + monkeypatch.setattr(runtime_module, "inspect_distributions", lambda path: distributions) + + with pytest.raises(ClientRuntimeInspectionError, match="must share one version"): + ClientRuntimeInspector().inspect() diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py new file mode 100644 index 000000000..a3bb85ffe --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -0,0 +1,784 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import subprocess +import sys +from dataclasses import replace +from functools import partial +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pyarrow.parquet as pq +import pytest +from conftest import ClientWorkerCase, FakeDataDesigner + +import data_designer.slurm.client.execution as execution_module +import data_designer.slurm.client.worker as worker_module +from data_designer.config import ResumeMode +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.execution import ClientWorker +from data_designer.slurm.client.records import ( + ClientEnvironmentManifest, + ClientEnvironmentOutcome, + ClientErrorCode, + ClientOutcome, + ClientProgress, + ClientProgressPhase, + ClientResult, +) +from data_designer.slurm.contracts import compute_serialized_json_sha256 +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.state import CandidateOutputManifest, compute_candidate_schema_digest +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier + + +def test_preflight_materializes_endpoint_and_ready_environment(client_worker_case: ClientWorkerCase) -> None: + designers: list[FakeDataDesigner] = [] + + def factory(**kwargs: object) -> FakeDataDesigner: + designer = FakeDataDesigner(**kwargs) + designers.append(designer) + return designer + + manifest = ClientWorker(data_designer_factory=factory).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert manifest.outcome is ClientEnvironmentOutcome.READY + assert designers[0].validated_builder is not None + model_config = designers[0].validated_builder.model_configs[0] + assert model_config.model == client_worker_case.plan.deployments[0].served_model_name + assert model_config.provider == "slurm-generator" + assert designers[0].initialization["model_providers"][0].endpoint == next( + iter(client_worker_case.endpoints.values()) + ) + + +def test_preflight_rejects_endpoint_from_another_gpu_allocation(client_worker_case: ClientWorkerCase) -> None: + alias = client_worker_case.plan.deployments[0].authored.model_alias + + with pytest.raises(ClientWorkerError, match="runtime model endpoint is invalid"): + ClientWorker(data_designer_factory=FakeDataDesigner, environment={"SLURM_JOB_GPUS": "0"}).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints={alias: "http://127.0.0.1:10256/v1"}, + plugins=(), + ) + + +def test_preflight_rejects_missing_managed_assets(client_worker_case: ClientWorkerCase) -> None: + Path(client_worker_case.plan.invocation.effective_input_bindings.managed_assets_path).rmdir() + + with pytest.raises(ClientWorkerError) as error: + ClientWorker(data_designer_factory=FakeDataDesigner).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.CONFIG_INVALID + manifest = ClientEnvironmentManifest.model_validate_json( + (client_worker_case.attempt_dir / "client-environment.json").read_text() + ) + assert manifest.outcome is ClientEnvironmentOutcome.FAILED + + +def test_preflight_normalizes_context_construction_failure(client_worker_case: ClientWorkerCase) -> None: + def factory(**_: object) -> FakeDataDesigner: + raise RuntimeError("credential=do-not-persist") + + with pytest.raises(ClientWorkerError) as error: + ClientWorker(data_designer_factory=factory).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.CONFIG_INVALID + persisted = (client_worker_case.attempt_dir / "client-environment.json").read_text() + assert "do-not-persist" not in persisted + assert ClientEnvironmentManifest.model_validate_json(persisted).error_code == ClientErrorCode.CONFIG_INVALID.value + + +def test_main_preserves_equals_in_model_alias( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + environment_builder = Mock() + environment_builder.prepare.return_value = client_worker_case.prepared + worker = Mock() + modules = { + "data_designer.slurm.client.plugins": SimpleNamespace(discover_plugins=lambda _: ()), + "data_designer.slurm.client.execution": SimpleNamespace(ClientWorker=lambda: worker), + } + monkeypatch.setattr(worker_module, "ClientEnvironmentBuilder", lambda: environment_builder) + monkeypatch.setattr(worker_module, "activate_environment", lambda _: None) + monkeypatch.setattr(worker_module, "importlib", SimpleNamespace(import_module=modules.__getitem__)) + endpoint = "http://127.0.0.1:17000/v1" + + result = worker_module.main( + [ + "preflight", + "--plan", + client_worker_case.plan_path.as_posix(), + "--shard-id", + client_worker_case.prepared.shard_id, + "--attempt-id", + client_worker_case.prepared.attempt_id, + "--attempt-dir", + client_worker_case.attempt_dir.as_posix(), + "--resume-mode", + "always", + "--endpoint", + f"judge=v2={endpoint}", + ] + ) + + assert result == 0 + assert worker.preflight.call_args.kwargs["endpoints"] == {"judge=v2": endpoint} + assert worker.preflight.call_args.kwargs["retry_resume"] is ResumeMode.ALWAYS + + +@pytest.mark.parametrize( + ("record_delta", "expected_outcome"), + ((0, ClientOutcome.COMPLETE), (-1, ClientOutcome.PARTIAL)), +) +def test_run_persists_semantic_result_and_candidate( + client_worker_case: ClientWorkerCase, + record_delta: int, + expected_outcome: ClientOutcome, +) -> None: + actual_records = client_worker_case.plan.shards[0].requested_records + record_delta + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, actual_records=actual_records)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.outcome is expected_outcome + assert result.actual_records == actual_records + persisted = ClientResult.model_validate_json((client_worker_case.attempt_dir / "client-result.json").read_text()) + candidate = CandidateOutputManifest.model_validate_json( + (client_worker_case.attempt_dir / "output-manifest.json").read_text() + ) + progress = ClientProgress.model_validate_json((client_worker_case.attempt_dir / "client-progress.json").read_text()) + assert persisted == result + assert candidate.actual_records == actual_records + assert tuple(file.relative_path for file in candidate.files) == ("parquet-files/batch_00000.parquet",) + assert not (client_worker_case.attempt_dir / "dataset/part-00000.parquet").exists() + assert result.candidate_output_manifest.sha256 == candidate.compute_sha256() + assert progress.phase is ClientProgressPhase.COMPLETE + + +def test_progress_metadata_error_does_not_abort_generation( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, batch_count=2)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + original = pq.read_metadata + reads = 0 + + def fail_once(path: Path) -> object: + nonlocal reads + reads += 1 + if reads == 1: + raise OSError("transient metadata read failure") + return original(path) + + monkeypatch.setattr(execution_module.lazy.pq, "read_metadata", fail_once) + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.outcome is ClientOutcome.COMPLETE + assert result.actual_records == client_worker_case.plan.shards[0].requested_records + + +def test_progress_retries_transient_batch_metadata_errors( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, batch_count=3)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + original_read = pq.read_metadata + original_best_effort = execution_module._ProgressWriter.best_effort # noqa: SLF001 + failed = False + generating_counts: list[int | None] = [] + + def fail_second_batch_once(path: Path) -> object: + nonlocal failed + if path.name == "batch_00001.parquet" and not failed: + failed = True + raise OSError("transient metadata read failure") + return original_read(path) + + def observe_progress( + writer: object, + phase: ClientProgressPhase, + *, + completed_records: int | None = None, + error_code: ClientErrorCode | None = None, + ) -> None: + if phase is ClientProgressPhase.GENERATING: + generating_counts.append(completed_records) + original_best_effort(writer, phase, completed_records=completed_records, error_code=error_code) + + monkeypatch.setattr(execution_module.lazy.pq, "read_metadata", fail_second_batch_once) + monkeypatch.setattr(execution_module._ProgressWriter, "best_effort", observe_progress) # noqa: SLF001 + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert failed + assert result.outcome is ClientOutcome.COMPLETE + assert generating_counts[-1] == result.actual_records + + +def test_run_normalizes_parquet_schema_metadata(client_worker_case: ClientWorkerCase) -> None: + worker = ClientWorker( + data_designer_factory=partial(FakeDataDesigner, schema_metadata={b"producer": b"data-designer"}) + ) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + candidate = CandidateOutputManifest.model_validate_json( + (client_worker_case.attempt_dir / "output-manifest.json").read_text() + ) + schema = pq.read_schema(client_worker_case.attempt_dir / "dataset/parquet-files/batch_00000.parquet") + assert candidate.dataset_schema_digest == compute_candidate_schema_digest(schema) + + +@pytest.mark.parametrize("batch_count", (2, 65)) +def test_run_uses_batch_files_when_candidate_is_bounded( + client_worker_case: ClientWorkerCase, + batch_count: int, +) -> None: + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, batch_count=batch_count)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + candidate = CandidateOutputManifest.model_validate_json( + (client_worker_case.attempt_dir / "output-manifest.json").read_text() + ) + if batch_count == 2: + assert tuple(file.relative_path for file in candidate.files) == ( + "parquet-files/batch_00000.parquet", + "parquet-files/batch_00001.parquet", + ) + assert not (client_worker_case.attempt_dir / "dataset/part-00000.parquet").exists() + with CandidateArtifactVerifier().verify(candidate) as verified: + assert sum(verified.record_counts) == candidate.actual_records + assert verified.dataset_schema_digest == candidate.dataset_schema_digest + else: + assert tuple(file.relative_path for file in candidate.files) == ("part-00000.parquet",) + assert (client_worker_case.attempt_dir / "dataset/part-00000.parquet").is_file() + + +def test_run_exports_batches_with_different_schemas(client_worker_case: ClientWorkerCase) -> None: + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, batch_count=2, mixed_batch_schemas=True)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + candidate = CandidateOutputManifest.model_validate_json( + (client_worker_case.attempt_dir / "output-manifest.json").read_text() + ) + assert tuple(file.relative_path for file in candidate.files) == ("part-00000.parquet",) + + +@pytest.mark.parametrize( + ("failure", "expected_code"), + ( + (RuntimeError("credential=do-not-persist"), ClientErrorCode.GENERATION_FAILED), + (KeyboardInterrupt(), ClientErrorCode.INTERRUPTED), + ), +) +def test_run_normalizes_generation_failures( + client_worker_case: ClientWorkerCase, + failure: BaseException, + expected_code: ClientErrorCode, +) -> None: + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, create_error=failure)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + with pytest.raises(ClientWorkerError) as error: + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is expected_code + persisted = (client_worker_case.attempt_dir / "client-result.json").read_text() + assert "do-not-persist" not in persisted + assert ClientResult.model_validate_json(persisted).error_code == expected_code.value + + +def test_run_rejects_overcount_as_invalid_output(client_worker_case: ClientWorkerCase) -> None: + actual_records = client_worker_case.plan.shards[0].requested_records + 1 + worker = ClientWorker(data_designer_factory=partial(FakeDataDesigner, actual_records=actual_records)) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + with pytest.raises(ClientWorkerError) as error: + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.OUTPUT_INVALID + persisted = ClientResult.model_validate_json((client_worker_case.attempt_dir / "client-result.json").read_text()) + assert persisted.error_code == ClientErrorCode.OUTPUT_INVALID.value + + +def test_preflight_rejects_required_resume_without_workspace(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "always" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + with pytest.raises(ClientWorkerError) as error: + ClientWorker(data_designer_factory=FakeDataDesigner).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.CONFIG_INVALID + + +def test_run_relocates_completed_resume_downgrade(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "always" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + resume_path = Path(plan.shards[0].resume_workspace.path) + resume_path.mkdir() + (resume_path / "partial").touch() + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.effective_resume_mode == "never" + assert not resume_path.exists() + assert (client_worker_case.attempt_dir / "dataset").is_dir() + + +def test_if_possible_without_workspace_uses_attempt_dataset(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.requested_resume_mode == "if_possible" + assert result.effective_resume_mode == "never" + assert result.dataset_path == (client_worker_case.attempt_dir / "dataset").as_posix() + + +@pytest.mark.parametrize("retry_resume", [ResumeMode.NEVER, ResumeMode.ALWAYS]) +def test_retry_resume_mode_controls_if_possible_generation( + client_worker_case: ClientWorkerCase, + retry_resume: ResumeMode, +) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + if retry_resume is ResumeMode.ALWAYS: + resume_path = Path(plan.shards[0].resume_workspace.path) + resume_path.mkdir() + (resume_path / "partial").touch() + worker = ClientWorker( + data_designer_factory=partial(FakeDataDesigner, effective_resume=retry_resume), + ) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=retry_resume, + ) + + result = worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=retry_resume, + ) + + expected_path = ( + Path(plan.shards[0].resume_workspace.path) + if retry_resume is ResumeMode.ALWAYS + else client_worker_case.attempt_dir / "dataset" + ) + assert result.requested_resume_mode == "if_possible" + assert result.effective_resume_mode == retry_resume.value + assert result.dataset_path == expected_path.as_posix() + + +def test_retry_rejects_effective_resume_mode_drift(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + resume_path = Path(plan.shards[0].resume_workspace.path) + resume_path.mkdir() + (resume_path / "partial").touch() + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=ResumeMode.ALWAYS, + ) + + with pytest.raises(ClientWorkerError) as error: + worker.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + retry_resume=ResumeMode.ALWAYS, + ) + + assert error.value.code is ClientErrorCode.OUTPUT_INVALID + + +def test_if_possible_interruption_preserves_workspace_for_retry(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["resume"] = "if_possible" + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + interrupted = ClientWorker( + data_designer_factory=partial(FakeDataDesigner, create_error_after_batch=KeyboardInterrupt()) + ) + interrupted.preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + with pytest.raises(ClientWorkerError): + interrupted.run( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + resume_path = Path(plan.shards[0].resume_workspace.path) + assert (resume_path / "parquet-files/batch_00000.parquet").is_file() + attempt_dir = resume_path.parent / "attempts" / "attempt-0002" + scratch_root = client_worker_case.scratch_root.parent / "retry-scratch" + overlay_path = scratch_root / "client-env" / "site-packages" + overlay_path.mkdir(parents=True) + prepared = replace( + client_worker_case.prepared, + attempt_id="attempt-0002", + attempt_dir=attempt_dir, + scratch_root=scratch_root, + overlay_path=overlay_path, + ) + retry = ClientWorker(data_designer_factory=partial(FakeDataDesigner, effective_resume=ResumeMode.ALWAYS)) + retry.preflight( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + result = retry.run( + client_worker_case.plan_path, + prepared=prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert result.effective_resume_mode == "always" + assert result.dataset_path == resume_path.as_posix() + + +def test_preflight_rejects_builder_with_missing_plugin(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + builder = payload["builder"]["inline"] + builder["data_designer"]["columns"] = [{"name": "custom", "column_type": "missing-plugin"}] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(builder) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + with pytest.raises(ClientWorkerError) as error: + ClientWorker(data_designer_factory=FakeDataDesigner).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.CONFIG_INVALID + + +def test_preflight_rejects_plugin_secondary_model_alias( + client_worker_case: ClientWorkerCase, + fake_plugin_overlay: Path, +) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + builder = payload["builder"]["inline"] + builder["data_designer"]["columns"] = [ + { + "name": "custom", + "column_type": "fake-slurm-column", + "model_alias": "generator", + "judge_model_alias": "missing", + } + ] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(builder) + lock_payload = client_worker_case.lock.model_dump(mode="json") + wheel_path = ( + client_worker_case.plan_path.parent / "dependencies" / "fake_data_designer_plugin-1.0.0-py3-none-any.whl" + ) + lock_payload["authored_requirements"] = ["fake-data-designer-plugin==1.0.0"] + lock_payload["overlay_packages"] = [ + { + "name": "fake-data-designer-plugin", + "version": "1.0.0", + "artifact": {"path": wheel_path.as_posix(), "sha256": "a" * 64}, + } + ] + lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) + Path(client_worker_case.plan.client.dependency_lock.path).write_text(lock.serialize_json()) + payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + script = """ +import json +import sys +from pathlib import Path + +import data_designer.slurm.client.worker as worker +from data_designer.slurm.client.environment import PreparedClientEnvironment +from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution + +plan_path = Path(sys.argv[1]) +attempt_dir = Path(sys.argv[2]) +plan = json.loads(plan_path.read_text()) +lock = json.loads(Path(plan["client"]["dependency_lock"]["path"]).read_text()) +installed = tuple( + sorted( + ( + *(InstalledDistribution(**item) for item in lock["base_distributions"]), + *(InstalledDistribution(name=item["name"], version=item["version"]) for item in lock["overlay_packages"]), + ), + key=lambda item: item.name, + ) +) +prepared = PreparedClientEnvironment( + run_id=plan["run_id"], + shard_id=plan["shards"][0]["shard_id"], + attempt_id="attempt-0001", + attempt_dir=attempt_dir, + scratch_root=Path(sys.argv[4]), + overlay_path=Path(sys.argv[3]), + dependency_lock=ArtifactReference(**plan["client"]["dependency_lock"]), + client_runtime_sha256=plan["client"]["runtime"]["runtime_sha256"], + python_abi=lock["python_abi"], + installer_outcome=ClientInstallerOutcome.REUSED, + installed_distributions=installed, +) +worker.activate_environment(prepared) + +from data_designer.slurm.client.plugins import discover_plugins + +plugins = discover_plugins(installed) +from data_designer.slurm.client.execution import ClientWorker +from data_designer.slurm.client.errors import ClientWorkerError + +try: + ClientWorker().preflight( + plan_path, + prepared=prepared, + endpoints={"generator": "http://127.0.0.1:10000/v1"}, + plugins=plugins, + ) +except ClientWorkerError as error: + assert error.code is ClientErrorCode.CONFIG_INVALID + assert error.redacted_message == "builder references unavailable model aliases" +else: + raise AssertionError("preflight accepted a missing secondary model alias") +""" + scratch_root = client_worker_case.attempt_dir.parent / "scratch" + scratch_root.mkdir() + result = subprocess.run( + [ + sys.executable, + "-c", + script, + client_worker_case.plan_path.as_posix(), + client_worker_case.attempt_dir.as_posix(), + fake_plugin_overlay.as_posix(), + scratch_root.as_posix(), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +def test_preflight_rejects_profiler_model_alias(client_worker_case: ClientWorkerCase) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + builder = payload["builder"]["inline"] + builder["data_designer"]["profilers"] = [{"model_alias": "missing"}] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(builder) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + + with pytest.raises(ClientWorkerError) as error: + ClientWorker(data_designer_factory=FakeDataDesigner).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert error.value.code is ClientErrorCode.CONFIG_INVALID + + +def test_preflight_keeps_mcp_secret_out_of_records( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload = client_worker_case.plan.model_dump(mode="json") + payload["invocation"]["authored"]["mcp_providers"] = [ + { + "provider_type": "sse", + "name": "private-mcp", + "endpoint": "https://mcp.example.test/events", + "api_key": {"type": "secret", "environment": "WORKER_SECRET"}, + } + ] + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + client_worker_case.plan_path.write_text(plan.serialize_json()) + monkeypatch.setenv("WORKER_SECRET", "private-value") + designers: list[FakeDataDesigner] = [] + + def factory(**kwargs: object) -> FakeDataDesigner: + designer = FakeDataDesigner(**kwargs) + designers.append(designer) + return designer + + ClientWorker(data_designer_factory=factory).preflight( + client_worker_case.plan_path, + prepared=client_worker_case.prepared, + endpoints=client_worker_case.endpoints, + plugins=(), + ) + + assert designers[0].initialization["mcp_providers"][0].api_key == "private-value" + assert "private-value" not in (client_worker_case.attempt_dir / "client-environment.json").read_text() diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py new file mode 100644 index 000000000..618a091b1 --- /dev/null +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -0,0 +1,396 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from data_designer.config import DataDesignerConfigBuilder, LLMTextColumnConfig, ModelConfig +from data_designer.slurm.config import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + ClientConfig, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ProfileSelectionSource, + SlurmConfigBuilderError, + SlurmConfigLoadError, + SlurmProfileCatalog, + load_benchmark_config, + load_builder_payload, + load_profile_catalog, + load_run_config, + resolve_profile, +) + + +def _config_builder(*, prompt: str | None = None) -> DataDesignerSlurmConfigBuilder: + data_designer = DataDesignerConfigBuilder( + model_configs=[ModelConfig(alias="generator", model="example/generator", provider="openai")] + ) + if prompt is not None: + data_designer.add_column(LLMTextColumnConfig(name="generated", prompt=prompt, model_alias="generator")) + return ( + DataDesignerSlurmConfigBuilder.from_config_builder(data_designer, name="generated-run") + .with_invocation( + num_records=8, + dataset_name="generated", + model_concurrency={"generator": 8}, + ) + .with_client() + .with_deployment( + { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "topology": {"tensor_parallel": 8}, + } + ) + ) + + +def test_builder_builds_without_file_discovery_or_serialization(tmp_path: Path) -> None: + builder = _config_builder() + + config = builder.build() + + assert config.name == "generated-run" + assert config.builder.inline is not None + assert config.deployments[0].model_alias == "generator" + assert config.client.cpus == 16 + assert not tuple(tmp_path.iterdir()) + + +def test_client_cpu_default_is_16() -> None: + assert ClientConfig().cpus == 16 + + +def test_builder_requires_complete_authored_intent() -> None: + builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") + + with pytest.raises(SlurmConfigBuilderError, match="invocation, deployment"): + builder.build() + + +@pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) +def test_builder_write_config_round_trips_supported_formats(tmp_path: Path, suffix: str) -> None: + builder = _config_builder() + path = tmp_path / f"run{suffix}" + + builder.write_config(path) + + 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_config_file_extension(tmp_path: Path) -> None: + with pytest.raises(SlurmConfigBuilderError, match="must end"): + _config_builder().write_config(tmp_path / "run.txt") + + +@pytest.mark.parametrize( + ("method", "values"), + [ + ("with_invocation", {"num_records": 0, "dataset_name": "generated"}), + ("with_array_tasks", {"count": 1, "max_concurrent": 2}), + ("with_submission", {"time_limit": "invalid"}), + ], +) +def test_builder_normalizes_invalid_authored_values( + method: str, + values: dict[str, object], +) -> None: + with pytest.raises(SlurmConfigBuilderError): + getattr(_config_builder(), method)(**values) + + +def test_builder_validation_errors_identify_the_invalid_field() -> None: + with pytest.raises(SlurmConfigBuilderError, match="num_records: must be greater"): + _config_builder().with_invocation(num_records=0, dataset_name="generated") + + +def test_builder_validation_errors_hide_secret_inputs() -> None: + secret = "super-secret-token" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_client(cpus=secret) # type: ignore[arg-type] + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_builder_validation_errors_hide_dynamic_mapping_keys() -> None: + secret = "sk_live_ABC123XYZ" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_invocation( + num_records=1, + dataset_name="generated", + model_concurrency={secret: 0}, + ) + + assert "model_concurrency" in str(error.value) + assert secret not in str(error.value) + + +def test_builder_validation_errors_preserve_discriminated_union_locations() -> None: + secret = "super-secret-token" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_invocation( + num_records=1, + dataset_name="generated", + mcp_providers=[ + { + "provider_type": "streamable_http", + "name": "example", + "endpoint": f"https://example.com/?token={secret}", + } + ], + ) + + assert "mcp_providers[0].endpoint" in str(error.value) + assert secret not in str(error.value) + + +def test_builder_custom_validation_errors_hide_secret_values() -> None: + secret = "super-secret-token" + + with pytest.raises(SlurmConfigBuilderError) as error: + _config_builder().with_client( + dependencies={"requirements": [f"pkg=={secret}"]}, + ) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_builder_normalizes_write_failures(tmp_path: Path) -> None: + path = tmp_path / "run.json" + path.mkdir() + + with pytest.raises(SlurmConfigBuilderError, match="cannot write"): + _config_builder().write_config(path) + + +@pytest.mark.parametrize( + ("suffix", "contents", "message"), + [ + (".json", '{"schema_version": 1, "schema_version": 1}', "duplicate"), + (".yaml", "schema_version: 1\nschema_version: 1\n", "duplicate"), + (".yaml", "defaults: &defaults\n schema_version: 1\nrun: *defaults\n", "anchors"), + (".yaml", "schema_version: 1\nname: ${RUN_NAME}\n", "interpolation"), + (".yaml", "schema_version: 1\nbuilder:\n source: ${HOME}/builder.json\n", "interpolation"), + ], +) +def test_strict_loader_rejects_ambiguous_yaml_and_json( + tmp_path: Path, + suffix: str, + contents: str, + message: str, +) -> None: + path = tmp_path / f"run{suffix}" + path.write_text(contents) + + with pytest.raises(SlurmConfigLoadError, match=message): + load_run_config(path) + + +def test_strict_loader_rejects_non_object_and_unknown_extension(tmp_path: Path) -> None: + json_path = tmp_path / "run.json" + json_path.write_text("[]") + + with pytest.raises(SlurmConfigLoadError, match="root must be an object"): + load_run_config(json_path) + with pytest.raises(SlurmConfigLoadError, match="must end"): + load_run_config(tmp_path / "run.toml") + + +def test_loader_validation_errors_hide_secret_inputs(tmp_path: Path) -> None: + secret = "super-secret-token" + payload = _config_builder().build().model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["api_key"] = secret + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(SlurmConfigLoadError) as error: + load_run_config(path) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_loader_validation_errors_identify_the_invalid_field(tmp_path: Path) -> None: + payload = _config_builder().build().model_dump(mode="json") + payload["invocation"]["num_records"] = 0 + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(SlurmConfigLoadError, match=r"invocation\.num_records: must be greater"): + load_run_config(path) + + +def test_loader_validation_errors_preserve_nested_schema_locations(tmp_path: Path) -> None: + secret = "super-secret-token" + payload = _config_builder().build().model_dump(mode="json") + payload["deployments"][0]["server"]["image"] = {"path": f"/images/{secret}.txt"} + path = tmp_path / "run.json" + path.write_text(json.dumps(payload)) + + with pytest.raises(SlurmConfigLoadError) as error: + load_run_config(path) + + assert "deployments[0].server.image.path" in str(error.value) + assert secret not in str(error.value) + + +def test_loader_parse_errors_hide_source_values(tmp_path: Path) -> None: + secret = "super-secret-token" + path = tmp_path / "run.yaml" + path.write_text(f"api_key: {secret}: x\n") + + with pytest.raises(SlurmConfigLoadError, match="invalid YAML at line 1") as error: + load_run_config(path) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +@pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) +def test_loader_preserves_literal_interpolation_inside_builder_payload(tmp_path: Path, suffix: str) -> None: + path = tmp_path / f"run{suffix}" + builder = _config_builder(prompt="Use the literal ${HOME} value") + + builder.write_config(path) + + assert load_run_config(path) == builder.build() + + +def test_builder_payload_loader_validates_complete_config(tmp_path: Path) -> None: + payload = _config_builder(prompt="Use the literal ${HOME} value").build().builder.inline + assert payload is not None + path = tmp_path / "builder.json" + path.write_text(json.dumps(payload)) + + assert load_builder_payload(path) == payload + + +def test_builder_payload_loader_rejects_partial_config(tmp_path: Path) -> None: + path = tmp_path / "builder.json" + path.write_text('{"library_version":"1.0"}') + + with pytest.raises(SlurmConfigLoadError, match="failed validation"): + load_builder_payload(path) + + +def test_profile_source_and_selection_precedence( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, +) -> None: + explicit_path = tmp_path / "explicit.json" + environment_path = tmp_path / "environment.yaml" + default_path = tmp_path / DEFAULT_PROFILE_FILE_NAME + explicit_path.write_text(profile_catalog.serialize_json()) + environment_path.write_text(yaml.safe_dump(profile_catalog.model_dump(mode="json"), sort_keys=True)) + default_path.write_text(profile_catalog.serialize_json()) + + explicit = resolve_profile( + profile_file=explicit_path, + cluster="lab", + hostname_resolver=lambda: pytest.fail("explicit cluster selection must not resolve hostnames"), + environ={PROFILE_FILE_ENVIRONMENT: str(environment_path)}, + ) + environment = resolve_profile( + hostname_resolver=lambda: ("LAB-LOGIN-1",), + environ={PROFILE_FILE_ENVIRONMENT: str(environment_path)}, + ) + default = resolve_profile( + hostname_resolver=lambda: ("unmatched",), + environ={}, + home_directory=tmp_path, + ) + + assert (explicit.cluster_name, explicit.selection_source) == ("lab", ProfileSelectionSource.EXPLICIT) + assert explicit.catalog_path == explicit_path.as_posix() + assert (environment.cluster_name, environment.selection_source) == ( + "lab", + ProfileSelectionSource.HOSTNAME, + ) + assert (default.cluster_name, default.selection_source) == ("primary", ProfileSelectionSource.DEFAULT) + assert load_profile_catalog(default_path) == profile_catalog + + +def test_injected_profile_bypasses_catalog_lookup(profile_catalog: SlurmProfileCatalog) -> None: + selected = resolve_profile( + profile=profile_catalog.clusters["primary"], + hostname_resolver=lambda: pytest.fail("hostname lookup must not run"), + environ={PROFILE_FILE_ENVIRONMENT: "/missing/profile.json"}, + ) + + assert selected.selection_source is ProfileSelectionSource.INJECTED + assert selected.catalog_path is None + + +def test_profile_resolution_rejects_conflicting_or_empty_sources( + profile_catalog: SlurmProfileCatalog, +) -> None: + with pytest.raises(SlurmConfigLoadError, match="mutually exclusive"): + resolve_profile(catalog=profile_catalog, profile_file="profile.json") + with pytest.raises(SlurmConfigLoadError, match="must not be empty"): + resolve_profile(environ={PROFILE_FILE_ENVIRONMENT: ""}) + with pytest.raises(SlurmConfigLoadError, match="cluster selection"): + resolve_profile(profile=profile_catalog.clusters["primary"], cluster="primary") + with pytest.raises(SlurmConfigLoadError, match="unknown cluster"): + resolve_profile(catalog=profile_catalog, cluster="missing") + + +def test_profile_resolution_normalizes_ambiguous_hostname_errors( + profile_catalog: SlurmProfileCatalog, +) -> None: + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["host_patterns"] = ["*-login-*"] + catalog = SlurmProfileCatalog.model_validate(payload) + + with pytest.raises(SlurmConfigLoadError, match="multiple clusters"): + resolve_profile(catalog=catalog, hostnames=("primary-login-1",)) + + +def test_profile_resolution_hides_validation_inputs(profile_catalog: SlurmProfileCatalog) -> None: + secret = "super-secret-token: x" + catalog = profile_catalog.model_copy( + update={ + "clusters": {secret: profile_catalog.clusters["primary"]}, + "default_cluster": secret, + } + ) + + with pytest.raises(SlurmConfigLoadError, match="profile selection failed validation") as error: + resolve_profile(catalog=catalog, cluster=secret) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_json_builder_output_is_stable(tmp_path: Path) -> None: + path = tmp_path / "run.json" + builder = _config_builder() + + builder.write_config(path) + + assert path.read_text() == builder.build().serialize_json() + assert json.loads(path.read_text())["schema_version"] == 1 + assert DataDesignerSlurmConfig.model_validate_json(path.read_text()) == builder.build() diff --git a/packages/data-designer-slurm/tests/config/test_partition_selection.py b/packages/data-designer-slurm/tests/config/test_partition_selection.py new file mode 100644 index 000000000..dda5786dd --- /dev/null +++ b/packages/data-designer-slurm/tests/config/test_partition_selection.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config.profiles import SchedulerProfile +from data_designer.slurm.config.run import SubmissionConfig +from data_designer.slurm.planning.models import ResolvedSubmission + + +@pytest.mark.parametrize("selection", ["gpu-a,gpu-b", ["gpu-a", "gpu-b"]]) +def test_gpu_partition_selection_normalizes_to_sbatch_value(selection: object) -> None: + assert SchedulerProfile.model_validate({"partition": selection}).partition == "gpu-a,gpu-b" + assert SubmissionConfig.model_validate({"partition": selection}).partition == "gpu-a,gpu-b" + assert ( + ResolvedSubmission.model_validate( + {"job_name": "demo", "time_limit": "01:00:00", "partition": selection} + ).partition + == "gpu-a,gpu-b" + ) + + +@pytest.mark.parametrize( + "selection", + [[], ["gpu-a", "gpu-a"], "gpu-a,gpu-a", "gpu-a,,gpu-b", "gpu-a, gpu-b", "gpu-a;echo", ["gpu-a", "bad name"]], +) +def test_gpu_partition_selection_rejects_ambiguous_or_unsafe_values(selection: object) -> None: + for model in (SchedulerProfile, SubmissionConfig, ResolvedSubmission): + payload = {"partition": selection} + if model is ResolvedSubmission: + payload.update(job_name="demo", time_limit="01:00:00") + with pytest.raises(ValidationError): + model.model_validate(payload) diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py new file mode 100644 index 000000000..ec7950ebe --- /dev/null +++ b/packages/data-designer-slurm/tests/conftest.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast + +import pytest +from slurm_test_fakes import ( + FakeClock, + FakeCommandResponse, + FakeLogicalEndpoint, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, + FakeVllmBackend, +) + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ImageInspectionRecord, + ServerDeploymentConfig, + SlurmProfileCatalog, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerIdentity, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + +TEST_DIRECTORY = Path(__file__).parent +CONTRACT_GOLDEN_DIRECTORY = TEST_DIRECTORY / "contracts" / "golden" +INTEGRATION_GOLDEN_PATH = TEST_DIRECTORY / "integration" / "golden" / "finalization_chain.json" +SLURM_GOLDEN_DIRECTORY = TEST_DIRECTORY / "slurm_test_fakes" / "golden" / "slurm" +STATE_GOLDEN_DIRECTORY = TEST_DIRECTORY / "state" / "golden" + + +@pytest.fixture +def authored_run() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "authored_run.json").read_text()) + + +@pytest.fixture +def authored_run_single() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json").read_text() + ) + + +@pytest.fixture +def profile_catalog() -> SlurmProfileCatalog: + return SlurmProfileCatalog.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "profile_catalog.json").read_text()) + + +@pytest.fixture +def dependency_lock() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "dependency_lock.json").read_text()) + + +@pytest.fixture +def dependency_lock_single() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text() + ) + + +@pytest.fixture +def single_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text()) + + +@pytest.fixture +def shared_node_config(authored_run_single: DataDesignerSlurmConfig) -> DataDesignerSlurmConfig: + payload = authored_run_single.model_dump(mode="python") + models = (("writer", "Qwen/Qwen2.5-0.5B-Instruct"), ("judge", "microsoft/Phi-3-mini-4k-instruct")) + payload["builder"]["inline"]["data_designer"]["model_configs"] = [ + {"alias": alias, "model": model, "provider": "openai"} for alias, model in models + ] + builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.yaml").with_invocation( + num_records=4, dataset_name="judged", model_concurrency={"writer": 1, "judge": 1} + ) + for alias, model in models: + deployment = authored_run_single.deployments[0].model_dump(mode="python") + deployment.update( + model_alias=alias, model=model, resources={"nodes": 1, "shared_node": "models", "gpus_per_node": 4} + ) + deployment["topology"]["tensor_parallel"] = 2 + builder.with_deployment(ServerDeploymentConfig.model_validate(deployment)) + config = builder.build().model_dump(mode="python") + config["builder"] = payload["builder"] + return DataDesignerSlurmConfig.model_validate(config) + + +@pytest.fixture +def shared_node_plan( + shared_node_config: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, +) -> ResolvedSlurmRunPlan: + return SlurmRunCompiler.compile( + resolve_slurm_config( + shared_node_config, + selected_profile=single_node_plan.selected_profile, + client_runtime=single_node_plan.client.runtime, + deployment_images=(single_node_plan.deployments[0].image,) * 2, + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=single_node_plan.run_id, + package_version=single_node_plan.package_version, + ) + ) + + +@pytest.fixture +def multi_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "multi_node_plan.json").read_text()) + + +@pytest.fixture +def client_result() -> ClientResult: + return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.json").read_text()) + + +@pytest.fixture +def benchmark_config() -> DataDesignerSlurmBenchmarkConfig: + return DataDesignerSlurmBenchmarkConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "benchmark_config.json").read_text() + ) + + +@pytest.fixture +def benchmark_manifest() -> BenchmarkManifest: + return BenchmarkManifest.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_manifest.json").read_text()) + + +@pytest.fixture +def benchmark_report() -> BenchmarkReport: + return BenchmarkReport.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_report.json").read_text()) + + +@pytest.fixture +def client_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "client_image_inspection.json").read_text() + ) + + +@pytest.fixture +def serving_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "serving_image_inspection.json").read_text() + ) + + +@pytest.fixture +def collection_plan() -> CollectionPlan: + return CollectionPlan.model_validate_json((STATE_GOLDEN_DIRECTORY / "collection_plan.json").read_text()) + + +@pytest.fixture +def accounting_lag_observation() -> SchedulerObservation: + return SchedulerObservation.model_validate_json((STATE_GOLDEN_DIRECTORY / "accounting_lag.json").read_text()) + + +@pytest.fixture +def finalization_chain_payload() -> dict[str, object]: + return cast(dict[str, object], json.loads(INTEGRATION_GOLDEN_PATH.read_text())) + + +@pytest.fixture +def run_manifest(finalization_chain_payload: dict[str, object]) -> RunManifest: + return RunManifest.model_validate_json(json.dumps(finalization_chain_payload["run"])) + + +@pytest.fixture +def shard_manifests(finalization_chain_payload: dict[str, object]) -> tuple[ShardManifest, ...]: + return tuple( + ShardManifest.model_validate_json(json.dumps(payload)) + for payload in cast(list[object], finalization_chain_payload["shards"]) + ) + + +@pytest.fixture +def attempt_manifest(finalization_chain_payload: dict[str, object]) -> AttemptManifest: + return AttemptManifest.model_validate_json(json.dumps(finalization_chain_payload["attempt"])) + + +@pytest.fixture +def attempt_readiness(finalization_chain_payload: dict[str, object]) -> AttemptReadiness: + return AttemptReadiness.model_validate_json(json.dumps(finalization_chain_payload["readiness"])) + + +@pytest.fixture +def finalization_client_result(finalization_chain_payload: dict[str, object]) -> ClientResult: + return ClientResult.model_validate_json(json.dumps(finalization_chain_payload["client_result"])) + + +@pytest.fixture +def candidate_output_manifest(finalization_chain_payload: dict[str, object]) -> CandidateOutputManifest: + return CandidateOutputManifest.model_validate_json(json.dumps(finalization_chain_payload["candidate"])) + + +@pytest.fixture +def shard_winner(finalization_chain_payload: dict[str, object]) -> ShardWinner: + return ShardWinner.model_validate_json(json.dumps(finalization_chain_payload["winner"])) + + +@pytest.fixture +def fake_clock() -> FakeClock: + """Return an isolated explicitly controlled clock.""" + return FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=100.0) + + +@pytest.fixture +def fake_slurm_runner() -> FakeSlurmRunner: + """Return an isolated Slurm runner with one array and one bounded sinfo query.""" + return FakeSlurmRunner( + arrays=( + FakeSlurmArray( + tasks=( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)), + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + ), + ) + ), + ), + sinfo_responses={ + ("sinfo", "--noheader", "--format=%P|%G"): FakeCommandResponse( + stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + ) + }, + ) + + +@pytest.fixture +def fake_plugin_overlay() -> Path: + """Return the installed-layout fake Data Designer plugin overlay.""" + return TEST_DIRECTORY / "fixtures" / "fake_plugin_overlay" + + +@pytest.fixture +def fake_logical_endpoint() -> FakeLogicalEndpoint: + """Return an isolated two-backend logical endpoint.""" + return FakeLogicalEndpoint( + "http://127.0.0.1:31000", + ( + FakeVllmBackend("http://127.0.0.1:31001", rank=0), + FakeVllmBackend("http://127.0.0.1:31002", rank=1), + ), + ) diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json new file mode 100644 index 000000000..6bcdc0505 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -0,0 +1,139 @@ +{ + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] + } + }, + "source": null + }, + "client": { + "cpus": 32, + "dependencies": { + "index_credentials": { + "private-index": { + "environment": "PACKAGE_INDEX_TOKEN", + "type": "secret" + } + }, + "lock_file": null, + "requirements": [ + "data-designer-speech==0.2.0" + ] + } + }, + "deployments": [ + { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 2 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": { + "HF_TOKEN": { + "environment": "HF_TOKEN", + "type": "secret" + } + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "image": { + "name": "vllm-0-21", + "path": null + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 2, + "tensor_parallel": 8 + } + }, + { + "model": "/models/judge", + "model_alias": "judge", + "resources": { + "nodes": 1 + }, + "served_model_name": "judge-api", + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "dataset_name": "training-data", + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 64, + "judge": 32 + }, + "num_records": 100, + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + } + }, + "name": "two-model-generation", + "schema_version": 1, + "submission": { + "account": null, + "comment": null, + "job_name": "dd-two-model", + "partition": null, + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json new file mode 100644 index 000000000..2c0485901 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -0,0 +1,85 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": null + }, + "builder": { + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] + } + }, + "source": null + }, + "client": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + } + }, + "deployments": [ + { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 1 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": "vllm-0-21", + "path": null + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 8 + } + } + ], + "invocation": { + "dataset_name": "single-node", + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 8 + }, + "num_records": 8, + "resume": "never", + "run_config": {} + }, + "name": "single-node-generation", + "schema_version": 1, + "submission": { + "account": null, + "comment": null, + "job_name": "data-designer", + "partition": null, + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json new file mode 100644 index 000000000..dce19282d --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json @@ -0,0 +1,46 @@ +{ + "analysis": { + "target_runtime": "4h", + "target_total_records": 1000000 + }, + "base_run": { + "inline": null, + "source": "run.yaml" + }, + "concurrency_values": [ + 32, + 64, + 128 + ], + "deployment_cases": [ + { + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 1 + } + }, + "name": "two-independent-replicas" + }, + { + "deployments": { + "generator": { + "nodes": 2, + "nodes_per_replica": 2 + } + }, + "name": "one-two-node-replica" + } + ], + "model_aliases": [ + "generator" + ], + "name": "generator-scaling", + "record_policy": { + "base_records": 1000, + "max_records": 5000, + "records_per_concurrency": 1.0, + "type": "adaptive" + }, + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json new file mode 100644 index 000000000..d72acb06f --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json @@ -0,0 +1,26 @@ +{ + "benchmark_config": { + "path": "/workspace/primary/benchmarks/benchmark-001/config.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "benchmark_id": "benchmark-001", + "children": [ + { + "case_id": "two-independent-replicas-c32", + "child_authored_config": { + "path": "/workspace/primary/runs/run-benchmark-001-c32/authored-config.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "child_run_id": "run-benchmark-001-c32" + }, + { + "case_id": "one-two-node-replica-c32", + "child_authored_config": { + "path": "/workspace/primary/runs/run-benchmark-002-c32/authored-config.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "child_run_id": "run-benchmark-002-c32" + } + ], + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json new file mode 100644 index 000000000..99fbf17ac --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json @@ -0,0 +1,58 @@ +{ + "analysis_id": "analysis-001", + "benchmark_id": "benchmark-001", + "benchmark_manifest": { + "path": "/workspace/primary/benchmarks/benchmark-001/benchmark.json", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "cases": [ + { + "actual_records": 1000, + "boot_seconds": 60.0, + "case_id": "two-independent-replicas-c32", + "child_run_id": "run-benchmark-001-c32", + "feasible": true, + "generation_seconds": 120.0, + "gpu_hours_per_job": 0.8, + "gpus_per_job": 16, + "nodes_per_job": 2, + "outcome": "succeeded", + "request_count": 1000, + "requested_records": 1000, + "rows_per_second": 8.333333333333334, + "target_jobs": 1000, + "token_count": 200000, + "topology_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "total_gpu_hours": 800.0, + "wall_seconds": 180.0 + }, + { + "actual_records": null, + "boot_seconds": null, + "case_id": "one-two-node-replica-c32", + "child_run_id": "run-benchmark-002-c32", + "feasible": null, + "generation_seconds": null, + "gpu_hours_per_job": null, + "gpus_per_job": 16, + "nodes_per_job": 2, + "outcome": "pending", + "request_count": null, + "requested_records": 1000, + "rows_per_second": null, + "target_jobs": null, + "token_count": null, + "topology_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "total_gpu_hours": null, + "wall_seconds": null + } + ], + "created_at": "2026-08-19T12:00:00Z", + "recommendations": [ + { + "case_id": "two-independent-replicas-c32", + "kind": "minimum_gpu_hours" + } + ], + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json new file mode 100644 index 000000000..17f97038c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -0,0 +1,39 @@ +{ + "inspection": { + "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "data-designer-config", + "version": "0.9.2" + }, + { + "name": "data-designer-engine", + "version": "0.9.2" + }, + { + "name": "data-designer-slurm", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "installer_path": "/usr/bin/pip", + "installer_version": "26.1", + "kind": "client", + "python_abi": "cp312", + "python_implementation": "cpython", + "python_version": "3.12.12" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/client_result.json b/packages/data-designer-slurm/tests/contracts/golden/client_result.json new file mode 100644 index 000000000..f3b5edac6 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_result.json @@ -0,0 +1,20 @@ +{ + "actual_records": 50, + "attempt_id": "attempt-0001", + "candidate_output_manifest": { + "path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "completed_at": "2026-08-19T12:00:00Z", + "dataset_path": "/workspace/primary/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": false, + "effective_resume_mode": "never", + "error_code": null, + "outcome": "complete", + "redacted_message": null, + "requested_records": 50, + "requested_resume_mode": "if_possible", + "run_id": "run-001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json new file mode 100644 index 000000000..1f44f2042 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -0,0 +1,47 @@ +{ + "authored_requirements": [ + "data-designer-speech==0.2.0" + ], + "authored_source": null, + "base_distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "data-designer-config", + "version": "0.9.2" + }, + { + "name": "data-designer-engine", + "version": "0.9.2" + }, + { + "name": "data-designer-slurm", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "client_runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1", + "overlay_packages": [ + { + "artifact": { + "path": "/workspace/primary/runs/run-001/dependencies/data_designer_speech-0.2.0.whl", + "sha256": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff" + }, + "name": "data-designer-speech", + "version": "0.2.0" + } + ], + "python_abi": "cp312", + "resolver_version": "resolver-1", + "schema_version": 1, + "source": null +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json new file mode 100644 index 000000000..b75c6b608 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -0,0 +1,36 @@ +{ + "authored_requirements": [], + "authored_source": null, + "base_distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "data-designer-config", + "version": "0.9.2" + }, + { + "name": "data-designer-engine", + "version": "0.9.2" + }, + { + "name": "data-designer-slurm", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "client_runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1", + "overlay_packages": [], + "python_abi": "cp312", + "resolver_version": "resolver-1", + "schema_version": 1, + "source": null +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json new file mode 100644 index 000000000..2b912541e --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -0,0 +1,445 @@ +{ + "array_tasks": { + "count": 2, + "max_concurrent": 2 + }, + "authored_config": { + "path": "/workspace/primary/runs/run-001/authored-config.json", + "sha256": "8c453e1064bf37c675b8ed62f1fe7b89234c7854e2c2f0955ae0c83938c73506" + }, + "builder": { + "authored_source": null, + "content_sha256": "2d4536fc8ff5a32e06269769bb563cb354c53cee95f45da31c9c6781130b4222", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator", + "judge" + ], + "source": null + }, + "client": { + "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": { + "private-index": { + "environment": "PACKAGE_INDEX_TOKEN", + "type": "secret" + } + }, + "lock_file": null, + "requirements": [ + "data-designer-speech==0.2.0" + ] + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-001/dependency-lock.json", + "sha256": "d318c2a934694f9c5186236c98a4e66c58b5572bcb4d1aec1309344fea0766be" + }, + "gpu_count": 0, + "host_node_index": 0, + "ports": [ + { + "name": "deployment-00000-logical-endpoint", + "node_index": 0, + "port": 17000, + "role": "logical_endpoint" + }, + { + "name": "deployment-00001-logical-endpoint", + "node_index": 0, + "port": 17001, + "role": "logical_endpoint" + } + ], + "runtime": { + "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "data-designer-config", + "version": "0.9.2" + }, + { + "name": "data-designer-engine", + "version": "0.9.2" + }, + { + "name": "data-designer-slurm", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "python_abi": "cp312", + "python_executable": "/opt/data-designer/bin/python", + "python_implementation": "cpython", + "python_version": "3.12.12", + "runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1" + } + }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "deployments": [ + { + "authored": { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 2 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": { + "HF_TOKEN": { + "environment": "HF_TOKEN", + "type": "secret" + } + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "image": { + "name": "vllm-0-21", + "path": null + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 2, + "tensor_parallel": 8 + } + }, + "deployment_id": "deployment-00000", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "node_indices": [ + 0, + 1 + ], + "ports": [ + { + "name": "deployment-00000-http-00000", + "node_index": 0, + "port": 18000, + "role": "http" + }, + { + "name": "deployment-00000-rendezvous-00000", + "node_index": 0, + "port": 19000, + "role": "rendezvous" + } + ], + "served_model_name": "example/generator", + "topology": { + "gpus_per_replica": 16, + "node_group_count": 1, + "nodes_per_replica": 2, + "pipeline_parallel": 2, + "replica_count": 1, + "replicas_per_node_group": 1, + "tensor_parallel": 8 + } + }, + { + "authored": { + "model": "/models/judge", + "model_alias": "judge", + "resources": { + "nodes": 1 + }, + "served_model_name": "judge-api", + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 1 + } + }, + "deployment_id": "deployment-00001", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": null, + "path": "/images/vllm-0-22.sqsh" + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.22.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "9999999999999999999999999999999999999999999999999999999999999999" + }, + "path": "/images/vllm-0-22.sqsh", + "sha256": "9999999999999999999999999999999999999999999999999999999999999999" + }, + "node_indices": [ + 2 + ], + "ports": [ + { + "name": "deployment-00001-http-00000", + "node_index": 2, + "port": 18000, + "role": "http" + }, + { + "name": "deployment-00001-http-00001", + "node_index": 2, + "port": 18001, + "role": "http" + }, + { + "name": "deployment-00001-http-00002", + "node_index": 2, + "port": 18002, + "role": "http" + }, + { + "name": "deployment-00001-http-00003", + "node_index": 2, + "port": 18003, + "role": "http" + }, + { + "name": "deployment-00001-http-00004", + "node_index": 2, + "port": 18004, + "role": "http" + }, + { + "name": "deployment-00001-http-00005", + "node_index": 2, + "port": 18005, + "role": "http" + }, + { + "name": "deployment-00001-http-00006", + "node_index": 2, + "port": 18006, + "role": "http" + }, + { + "name": "deployment-00001-http-00007", + "node_index": 2, + "port": 18007, + "role": "http" + } + ], + "served_model_name": "judge-api", + "topology": { + "gpus_per_replica": 1, + "node_group_count": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "replica_count": 8, + "replicas_per_node_group": 8, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "authored": { + "dataset_name": "training-data", + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 64, + "judge": 32 + }, + "num_records": 100, + "resume": "if_possible", + "run_config": { + "buffer_size": 8192 + } + }, + "effective_input_bindings": { + "managed_assets_path": "/workspace/primary/managed-assets", + "seed_path": null + }, + "effective_run_config": { + "async_trace": false, + "buffer_size": 8192, + "disable_early_shutdown": true, + "display_tui": false, + "jinja_rendering_engine": "secure", + "max_concurrent_row_groups": 3, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "max_in_flight_tasks": 1024, + "non_inference_max_parallel_workers": 4, + "otel_metrics_port": null, + "preserve_dropped_columns": true, + "progress_interval": 5.0, + "request_admission": null, + "shutdown_error_rate": 1.0, + "shutdown_error_window": 10, + "write_scheduler_events": false + } + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-001", + "runtime_bundle": { + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", + "cluster_name": "primary", + "matched_pattern": null, + "profile": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" + }, + "scheduler": { + "account": "research", + "bin_path": null, + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", + "schema_version": 1, + "selection_source": "explicit" + }, + "shards": [ + { + "array_task_index": 0, + "input_partition": null, + "record_range": { + "end_index_exclusive": 50, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-001/shards/shard-00000/dataset" + }, + "shard_id": "shard-00000", + "shard_index": 0 + }, + { + "array_task_index": 1, + "input_partition": null, + "record_range": { + "end_index_exclusive": 100, + "start_index": 50 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-001/shards/shard-00001/dataset" + }, + "shard_id": "shard-00001", + "shard_index": 1 + } + ], + "submission": { + "account": "research", + "comment": null, + "job_name": "dd-two-model", + "partition": "batch", + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json new file mode 100644 index 000000000..baf71e3b7 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -0,0 +1,57 @@ +{ + "clusters": { + "lab": { + "container_mounts": [], + "gpu_request_mode": "visible", + "gpus_per_node": "auto", + "host_patterns": [ + "lab-login-*" + ], + "image_build": { + "cpus_per_task": 2, + "memory": "8G", + "partition": "cpu", + "time_limit": "02:00:00" + }, + "scheduler": { + "account": "lab", + "bin_path": null, + "mem_per_gpu": null, + "partition": "gpu" + }, + "schema_version": 1, + "workspace_root": "/workspace/lab" + }, + "primary": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" + }, + "scheduler": { + "account": "research", + "bin_path": null, + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + } + }, + "default_cluster": "primary", + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json new file mode 100644 index 000000000..8cb5348e2 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json @@ -0,0 +1,11 @@ +{ + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" +} diff --git a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json new file mode 100644 index 000000000..5938747c5 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -0,0 +1,278 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": null + }, + "authored_config": { + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "bf9d45d571f5f288c74a3332cdd23576df60ce7bb3d934774f50bdce4b984fba" + }, + "builder": { + "authored_source": null, + "content_sha256": "0bac3a88498774a30b64f0a62617511c8e75601ca54e0b775d86484bb13788d8", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator" + ], + "source": null + }, + "client": { + "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-single/dependency-lock.json", + "sha256": "df1cf487016d89f1ab4d0d2df5551d0d899fd695804cab0cd3d99a7a936e2f3f" + }, + "gpu_count": 0, + "host_node_index": 0, + "ports": [ + { + "name": "deployment-00000-logical-endpoint", + "node_index": 0, + "port": 17000, + "role": "logical_endpoint" + } + ], + "runtime": { + "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "data-designer-config", + "version": "0.9.2" + }, + { + "name": "data-designer-engine", + "version": "0.9.2" + }, + { + "name": "data-designer-slurm", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "python_abi": "cp312", + "python_executable": "/opt/data-designer/bin/python", + "python_implementation": "cpython", + "python_version": "3.12.12", + "runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1" + } + }, + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "deployments": [ + { + "authored": { + "model": "example/generator", + "model_alias": "generator", + "resources": { + "nodes": 1 + }, + "served_model_name": null, + "server": { + "distributed_init_timeout": "10m", + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "image": { + "name": "vllm-0-21", + "path": null + }, + "lead_boot_standoff": "60s", + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "rank_launch_stagger": "5s", + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 8 + } + }, + "deployment_id": "deployment-00000", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "node_indices": [ + 0 + ], + "ports": [ + { + "name": "deployment-00000-http-00000", + "node_index": 0, + "port": 18000, + "role": "http" + } + ], + "served_model_name": "example/generator", + "topology": { + "gpus_per_replica": 8, + "node_group_count": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "replica_count": 1, + "replicas_per_node_group": 1, + "tensor_parallel": 8 + } + } + ], + "invocation": { + "authored": { + "dataset_name": "single-node", + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 8 + }, + "num_records": 8, + "resume": "never", + "run_config": {} + }, + "effective_input_bindings": { + "managed_assets_path": "/workspace/primary/managed-assets", + "seed_path": null + }, + "effective_run_config": { + "async_trace": false, + "buffer_size": 16384, + "disable_early_shutdown": true, + "display_tui": false, + "jinja_rendering_engine": "secure", + "max_concurrent_row_groups": 3, + "max_conversation_correction_steps": 0, + "max_conversation_restarts": 0, + "max_in_flight_tasks": 1024, + "non_inference_max_parallel_workers": 4, + "otel_metrics_port": null, + "preserve_dropped_columns": true, + "progress_interval": 5.0, + "request_admission": null, + "shutdown_error_rate": 1.0, + "shutdown_error_window": 10, + "write_scheduler_events": false + } + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-single", + "runtime_bundle": { + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", + "cluster_name": "primary", + "matched_pattern": null, + "profile": { + "container_mounts": [ + { + "read_only": false, + "source": "/workspace", + "target": "/workspace" + } + ], + "gpu_request_mode": "gres", + "gpus_per_node": 8, + "host_patterns": [ + "login*.primary.example", + "primary-login-*" + ], + "image_build": { + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" + }, + "scheduler": { + "account": "research", + "bin_path": null, + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", + "schema_version": 1, + "selection_source": "explicit" + }, + "shards": [ + { + "array_task_index": 0, + "input_partition": null, + "record_range": { + "end_index_exclusive": 8, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + }, + "shard_id": "shard-00000", + "shard_index": 0 + } + ], + "submission": { + "account": "research", + "comment": null, + "job_name": "data-designer", + "partition": "batch", + "time_limit": "03:55:00" + } +} diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py new file mode 100644 index 000000000..33120368b --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -0,0 +1,813 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.config import DataDesignerConfigBuilder, HuggingFaceSeedSource, ModelConfig +from data_designer.slurm.config import ( + ArrayTasksConfig, + BenchmarkBaseRun, + BuilderInput, + ClientDependencies, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageBuildRequest, + ImageInspectionRecord, + ImageRef, + LiteralEnvironmentBinding, + LocalStdioMCPProviderConfig, + QueueBackpressureConfig, + RemoteMCPProviderConfig, + SecretRef, + ServerDeploymentConfig, + SubmissionConfig, + VllmServerConfig, +) + + +@pytest.mark.parametrize("version", [None, 0, 2, "1"]) +def test_run_config_requires_supported_version(authored_run: DataDesignerSlurmConfig, version: object) -> None: + payload = authored_run.model_dump(mode="json") + if version is None: + payload.pop("schema_version") + else: + payload["schema_version"] = version + + with pytest.raises(ValidationError): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_run_config_rejects_unknown_fields(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["placement"] = {"gpu_ids": [0]} + + with pytest.raises(ValidationError, match="placement"): + DataDesignerSlurmConfig.model_validate(payload) + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"name": "image", "path": "/images/image.sqsh"}, + {"path": "relative.sqsh"}, + {"path": "/images/image.tar"}, + ], +) +def test_image_ref_requires_one_registered_alias_or_absolute_sqsh(payload: dict[str, str]) -> None: + with pytest.raises(ValidationError): + ImageRef.model_validate(payload) + + +@pytest.mark.parametrize( + "payload", + [ + {"requirements": None}, + {"requirements": [], "lock_file": "lock.json"}, + {"requirements": ["-e ./plugin"]}, + {"requirements": ["plugin @ git+https://example.test/plugin.git"]}, + {"requirements": ["plugin @ https://example.test/plugin.whl"]}, + {"requirements": ["plugin @ https://user:secret@example.test/plugin.whl#sha256=" + "a" * 64]}, + {"requirements": ["plugin @ https://example.test/plugin.whl?token=secret#sha256=" + "a" * 64]}, + {"requirements": ["plugin @ https://example.test:invalid/plugin.whl#sha256=" + "a" * 64]}, + {"requirements": ["my_pkg==1", "my-pkg==2"]}, + {"requirements": ["not valid !!!"]}, + {"requirements": ["plugin=="]}, + {"requirements": ["plugin==1; python_version >= '3.12'"]}, + {"requirements": None, "lock_file": "../lock.json"}, + ], +) +def test_client_dependencies_reject_mutable_or_ambiguous_sources(payload: dict[str, object]) -> None: + with pytest.raises(ValidationError): + ClientDependencies.model_validate(payload) + + +def test_client_dependencies_accept_digest_bound_wheel() -> None: + dependencies = ClientDependencies(requirements=["plugin @ https://example.test/plugin.whl#sha256=" + "a" * 64]) + + assert dependencies.requirements is not None + + +def test_rejected_dependency_does_not_appear_in_validation_error() -> None: + sentinel = "VERY_SECRET_VALUE" + + with pytest.raises(ValidationError) as error: + ClientDependencies(requirements=[f"not valid {sentinel} !!!"]) + + assert sentinel not in str(error.value) + + +def test_secret_reference_serializes_only_external_binding() -> None: + secret = SecretRef(type="secret", environment="HF_TOKEN") + sentinel = "VERY_SECRET_VALUE" + + assert secret.model_dump(mode="json") == {"type": "secret", "environment": "HF_TOKEN"} + with pytest.raises(ValidationError) as error: + SecretRef.model_validate({"type": "secret", "environment": "HF_TOKEN", "value": sentinel}) + assert sentinel not in str(error.value) + + +def test_rejected_secret_argument_does_not_appear_in_validation_error() -> None: + sentinel = "VERY_SECRET_VALUE" + + with pytest.raises(ValidationError) as error: + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=[f"--api-key={sentinel}"], + ) + + assert sentinel not in str(error.value) + + +def test_literal_environment_rejects_control_characters() -> None: + with pytest.raises(ValidationError, match="control"): + LiteralEnvironmentBinding(type="literal", value="line\nbreak") + + +def test_secret_shaped_environment_requires_external_reference() -> None: + literal = LiteralEnvironmentBinding(type="literal", value="plaintext-secret") + + with pytest.raises(ValidationError, match="external secret references"): + LocalStdioMCPProviderConfig( + provider_type="stdio", + name="provider", + command="provider", + environment={"API_TOKEN": literal}, + ) + with pytest.raises(ValidationError, match="external secret references"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + environment={"MODEL_PASSWORD": literal}, + ) + + +@pytest.mark.parametrize( + "endpoint", + [ + "https://user:password@example.test/mcp", + "https://example.test/mcp?token=plaintext-secret", + "https://example.test/mcp#secret", + ], +) +def test_remote_mcp_endpoint_rejects_embedded_credentials(endpoint: str) -> None: + with pytest.raises(ValidationError, match="must not embed"): + RemoteMCPProviderConfig(provider_type="sse", name="provider", endpoint=endpoint) + + +@pytest.mark.parametrize("endpoint", ["https:///missing-host", "https://example.test:invalid/mcp"]) +def test_remote_mcp_endpoint_requires_valid_host_and_port(endpoint: str) -> None: + with pytest.raises(ValidationError, match=r"HTTP\(S\)"): + RemoteMCPProviderConfig(provider_type="sse", name="provider", endpoint=endpoint) + + +@pytest.mark.parametrize( + "argument", + [ + "--api-key", + "--api-key plaintext-secret", + " --api-key plaintext-secret", + "--access-token=plaintext-secret", + "password", + ], +) +def test_stdio_mcp_rejects_secret_shaped_arguments(argument: str) -> None: + with pytest.raises(ValidationError, match="secret-shaped"): + LocalStdioMCPProviderConfig( + provider_type="stdio", + name="provider", + command="provider", + args=[argument], + ) + + +def test_rejected_stdio_mcp_secret_argument_does_not_appear_in_validation_error() -> None: + sentinel = "VALUE_THAT_MUST_NOT_APPEAR" + + with pytest.raises(ValidationError, match="secret-shaped") as error: + LocalStdioMCPProviderConfig( + provider_type="stdio", + name="provider", + command="provider", + args=[f"--api-key {sentinel}"], + ) + + assert sentinel not in str(error.value) + + +def test_vllm_defaults_and_backpressure_override() -> None: + server = VllmServerConfig(type="vllm", image=ImageRef(name="vllm")) + overridden = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + lead_boot_standoff="0s", + rank_launch_stagger="12s", + queue_backpressure=QueueBackpressureConfig(max_waiting_requests=0, retry_after_seconds=None), + ) + + assert server.lead_boot_standoff == "60s" + assert server.rank_launch_stagger == "5s" + assert server.queue_backpressure.model_dump() == {"max_waiting_requests": 128, "retry_after_seconds": 1} + assert overridden.lead_boot_standoff == "0s" + assert overridden.rank_launch_stagger == "12s" + assert overridden.queue_backpressure.model_dump() == {"max_waiting_requests": 0, "retry_after_seconds": None} + + +def test_vllm_rejects_nonpositive_retry_after_seconds() -> None: + with pytest.raises(ValidationError, match="greater than 0"): + QueueBackpressureConfig(retry_after_seconds=0) + + +@pytest.mark.parametrize("plugin", [None, "deberta_gliner_io"]) +def test_vllm_io_processor_plugin_round_trip(plugin: str | None) -> None: + server = VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), io_processor_plugin=plugin) + + assert VllmServerConfig.model_validate_json(server.model_dump_json()) == server + assert ("io_processor_plugin" in server.model_dump()) == (plugin is not None) + + +@pytest.mark.parametrize("plugin", ["", " ", "../plugin", "plugin module", "--plugin", "plugin;exit", "x" * 129]) +def test_vllm_rejects_invalid_io_processor_plugin(plugin: str) -> None: + with pytest.raises(ValidationError, match="io_processor_plugin"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), io_processor_plugin=plugin) + + +@pytest.mark.parametrize( + "argument", + [ + "--api-key=plaintext-secret", + "--api-key plaintext-secret", + "-asc=2", + "-dp", + "-dpa=127.0.0.1", + "-dpb", + "-dcp=2", + "-dpe", + "-dph", + "-dpl", + "-dpm", + "-dpn=1", + "-dpp", + "-dpr", + "-ep", + "-n", + "-n+2", + "-n2", + "-pcp=2", + "-pp=2", + "-r", + "-r0", + "-r1", + "-tp", + "--api-server-count=2", + "--config=/tmp/vllm.yaml", + "--cpu-distributed-timeout-seconds=120", + "--cp-kv-cache-interleave-size=2", + "--cpunodebind=0", + "--data-parallel-address", + "--data-parallel-backend=mp", + "--data-parallel-external-lb", + "--data-parallel-hybrid-lb", + "--data_parallel_master_ip=127.0.0.1", + "--data-parallel-rank=1", + "--data-parallel-rpc-port 29502", + "--data-parallel-size", + "--data-parallel-size-local", + "--data-parallel-start-rank=1", + "--data-parallel-supervisor-port=9256", + "--dcp-comm-backend=a2a", + "--dcp-kv-cache-interleave-size=2", + "--decode-context-parallel-size=2", + "--default-mm-loras={}", + "--distributed-executor-backend=mp", + "--distributed-init-address", + "--distributed-timeout-seconds", + "--dp-supervisor-probe-interval-s=1", + "--ec-transfer-config={}", + "--enable-elastic-ep", + "--enable-expert-parallel", + "--enable-lora", + "--enable-ssl-refresh", + "--grpc", + "--headless", + "--host=0.0.0.0", + "--io-processor-plugin=custom", + "--kv-events-config={}", + "--kv-transfer-config={}", + "--logits-processors=custom.LogitsProcessor", + "--lora-modules=adapter=/models/adapter", + "--max-cpu-loras=2", + "--max-loras=2", + "--master-addr", + "--master-port=29501", + "--middleware", + "--model", + "--nnodes", + "--node-rank=1", + "--no-data-parallel-external-lb", + "--no-data-parallel-hybrid-lb", + "--no-enable-elastic-ep", + "--no-enable-expert-parallel", + "--no-enable-lora", + "--no-enable-ssl-refresh", + "--no-numa-bind", + "--numa-bind", + "--numa_bind_nodes=[0,1]", + "--physcpubind=0-7", + "--pipeline-parallel-size=2", + "--port", + "--port 9000", + "--prefill-context-parallel-size=2", + "--reasoning-parser-plugin=/tmp/reasoning.py", + "--root-path=/v1", + "--served-model-name=example/model", + "--spec-model=example/draft-model", + "--speculative-config={}", + "--ssl-ca-certs=/tmp/ca.pem", + "--ssl-cert-reqs=2", + "--ssl-certfile=/tmp/cert.pem", + "--ssl-ciphers=ECDHE-RSA-AES256-GCM-SHA384", + "--ssl-keyfile=/tmp/key.pem", + "--tensor-parallel-size", + "--tensor_parallel_size=2", + "--tensor=2", + "--tool-parser-plugin=/tmp/tools.py", + "--uds=/tmp/vllm.sock", + "--weight-transfer-config={}", + "--worker-cls=custom.Worker", + "--worker-extension-cls=custom.Extension", + ], +) +def test_vllm_rejects_runtime_owned_arguments(argument: str) -> None: + with pytest.raises(ValidationError, match="owned"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[argument]) + + +@pytest.mark.parametrize("abbreviated_flag", ["-a", "-as", "-dc", "-e", "-pc", "-t"]) +def test_vllm_rejects_runtime_owned_short_flag_abbreviations(abbreviated_flag: str) -> None: + with pytest.raises(ValidationError, match="owned"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[abbreviated_flag, "4"]) + + +def test_vllm_allows_builtin_parser_selection() -> None: + config = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=["--reasoning_parser", "deepseek_r1", "--tool-call-parser", "hermes"], + ) + + assert config.extra_args == ["--reasoning_parser", "deepseek_r1", "--tool-call-parser", "hermes"] + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--structured-outputs-config.reasoning-parser-plugin=/tmp/reasoning.py"], + ["--structured-outputs-config", '{"reasoning_parser_plugin":"/tmp/reasoning.py"}'], + ['--structured-outputs-config={"reasoning_parser_plugin":"/tmp/reasoning.py"}'], + ['--additional-config={"nested":{"worker_cls":"custom.Worker"}}'], + ['--additional-config={"partition_key":"group"}'], + ], +) +def test_vllm_does_not_interpret_plugin_hooks_nested_in_json_arguments(extra_args: list[str]) -> None: + config = VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + assert config.extra_args == extra_args + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--hf-overrides", '{"api_key":"plaintext-secret"}'], + ['--hf-overrides={"nested":{"access_token":"plaintext-secret"}}'], + ], +) +def test_vllm_rejects_plaintext_secrets_in_plugin_owned_json_arguments(extra_args: list[str]) -> None: + with pytest.raises(ValidationError, match="plaintext values under secret-bearing keys"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + +def test_vllm_allows_nonsecret_json_arguments() -> None: + config = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=['--hf-overrides={"model_type":"custom"}'], + ) + + assert config.extra_args == ['--hf-overrides={"model_type":"custom"}'] + + +def test_vllm_preserves_argv_values_with_whitespace_and_templates() -> None: + extra_args = ["--hf-overrides", '{"model_type": "custom"}', "--chat-template={{messages}}"] + + config = VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + assert config.extra_args == extra_args + + +@pytest.mark.parametrize("argument", [" --max-model-len", "32768 ", " "]) +def test_vllm_rejects_argv_with_outer_whitespace(argument: str) -> None: + with pytest.raises(ValidationError, match="leading or trailing whitespace"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[argument]) + + +@pytest.mark.parametrize("argument", ["=", "=x"]) +def test_vllm_rejects_empty_option_names(argument: str) -> None: + with pytest.raises(ValidationError, match="option name must not be empty"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=[argument]) + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--hf-overrides", "{'model_type':'custom'}"], + ['--hf-overrides={"model_type":"custom"'], + ], +) +def test_vllm_preserves_structured_values_for_runtime_validation(extra_args: list[str]) -> None: + config = VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + assert config.extra_args == extra_args + + +@pytest.mark.parametrize( + "environment_name", + [ + "CUDA_VISIBLE_DEVICES", + "ENABLE_EP", + "HF_HOME", + "GROUP_RANK", + "LOCAL_RANK", + "MASTER_ADDR", + "MASTER_PORT", + "NVIDIA_VISIBLE_DEVICES", + "PYTHON_EXEC", + "RANK", + "ROLE_WORLD_SIZE", + "SLURM_PROCID", + "TORCHELASTIC_RUN_ID", + "VLLM_ALLOW_RUNTIME_LORA_UPDATING", + "VLLM_API_KEY", + "VLLM_CACHE_ROOT", + "VLLM_LEAD_BOOT_STANDOFF_SECONDS", + "VLLM_DP_RANK", + "VLLM_HOST_IP", + "VLLM_LORA_RESOLVER_HF_REPO_LIST", + "VLLM_MODEL_REDIRECT_PATH", + "VLLM_MOONCAKE_BOOTSTRAP_PORT", + "VLLM_NIXL_SIDE_CHANNEL_PORT", + "VLLM_PLUGINS", + "VLLM_PORT", + "VLLM_RAY_PER_WORKER_GPUS", + "VLLM_RPC_BASE_PATH", + "WORLD_SIZE", + ], +) +def test_vllm_rejects_runtime_owned_environment_names(environment_name: str) -> None: + with pytest.raises(ValidationError, match="owned by the compiler or runtime"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + environment={environment_name: SecretRef(type="secret", environment="EXTERNAL_VALUE")}, + ) + + +def test_vllm_allows_explicit_tuning_environment() -> None: + config = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + environment={ + "NCCL_DEBUG": LiteralEnvironmentBinding(type="literal", value="INFO"), + "VLLM_WORKER_MULTIPROC_METHOD": LiteralEnvironmentBinding(type="literal", value="spawn"), + }, + ) + + assert set(config.environment) == {"NCCL_DEBUG", "VLLM_WORKER_MULTIPROC_METHOD"} + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--hf-token=plaintext-secret"], + ["--hf-token", "plaintext-secret"], + ["--hf-token plaintext-secret"], + ], +) +def test_vllm_rejects_secret_shaped_arguments(extra_args: list[str]) -> None: + with pytest.raises(ValidationError, match="secret-shaped"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=extra_args) + + +def test_vllm_rejects_non_tokenized_arguments() -> None: + with pytest.raises(ValidationError, match="one token"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=["--max-model-len 32768"], + ) + + +def test_vllm_rejects_duplicate_arguments() -> None: + with pytest.raises(ValidationError, match="duplicate"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=["--max-model-len", "32768", "--max-model-len=16384"], + ) + + with pytest.raises(ValidationError, match="duplicate or conflicting"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=["--max_model_len=4096", "--max-model-len=8192"], + ) + + +def test_vllm_rejects_conflicting_boolean_arguments() -> None: + with pytest.raises(ValidationError, match="duplicate or conflicting"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + extra_args=["--enable-prefix-caching", "--no-enable-prefix-caching"], + ) + + +def test_vllm_rejects_option_terminator() -> None: + with pytest.raises(ValidationError, match="option terminators"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), extra_args=["--"]) + + +def test_vllm_rejects_distributed_timeout_beyond_startup_timeout() -> None: + with pytest.raises(ValidationError, match="must not exceed"): + VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + startup_timeout="10m", + distributed_init_timeout="11m", + ) + + +@pytest.mark.parametrize("field", ["lead_boot_standoff", "rank_launch_stagger"]) +def test_vllm_rejects_negative_launch_timing(field: str) -> None: + with pytest.raises(ValidationError): + VllmServerConfig.model_validate({"type": "vllm", "image": {"name": "vllm"}, field: "-1s"}) + + +def test_deployment_rejects_invalid_topology() -> None: + payload = { + "model_alias": "generator", + "model": "example/generator", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + "resources": {"nodes": 3}, + "topology": {"tensor_parallel": 8, "nodes_per_replica": 2}, + } + + with pytest.raises(ValidationError, match="divide"): + ServerDeploymentConfig.model_validate(payload) + + payload["resources"]["nodes"] = 2 + payload["topology"]["nodes_per_replica"] = 1 + payload["server"]["enable_expert_parallel"] = True + + deployment = ServerDeploymentConfig.model_validate(payload) + + assert deployment.topology.nodes_per_replica == 1 + + payload["topology"]["nodes_per_replica"] = 2 + with pytest.raises(ValidationError, match="expert"): + ServerDeploymentConfig.model_validate(payload) + + +def test_model_alias_preserves_public_data_designer_values() -> None: + alias = "judge/v2" + ModelConfig(alias=alias, model="example/judge", provider="openai") + + deployment = ServerDeploymentConfig.model_validate( + { + "model_alias": alias, + "model": "example/judge", + "server": {"type": "vllm", "image": {"name": "vllm"}}, + } + ) + + assert deployment.model_alias == alias + + +def test_run_rejects_duplicate_alias_and_unknown_concurrency(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["deployments"][1]["model_alias"] = "generator" + with pytest.raises(ValidationError, match="aliases"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["invocation"]["model_concurrency"]["missing"] = 1 + with pytest.raises(ValidationError, match="undeclared"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["server_configs"] = [] + + with pytest.raises(ValidationError, match="retired") as caught: + DataDesignerSlurmConfig.model_validate(payload) + + assert "builder input contains retired fields" in str(caught.value) + + +@pytest.mark.parametrize( + "secret_key", + [ + "api_key", + "accessToken", + "client-secret", + "consumer_key", + "license_key", + "password", + "private_key", + "signing_key", + "ssh_key", + "subscription_key", + ], +) +def test_builder_input_rejects_plaintext_secrets_in_plugin_owned_fields(secret_key: str) -> None: + inline = {"columns": [], "plugin_config": {secret_key: "opaque-plugin-value"}} + + with pytest.raises(ValidationError, match="plaintext values under secret-bearing keys"): + BuilderInput(inline=inline) + + +def test_rejected_embedded_secret_does_not_appear_in_validation_error() -> None: + sentinel = "VERY_SECRET_VALUE" + + with pytest.raises(ValidationError) as error: + BuilderInput(inline={"columns": [], "plugin_config": {"api_key": sentinel}}) + + assert sentinel not in str(error.value) + + +@pytest.mark.parametrize("plugin_key", ["sort_key", "partition_key", "primary_key", "idempotency_key"]) +def test_builder_input_allows_non_secret_plugin_keys(plugin_key: str) -> None: + inline = {"columns": [{"column_type": "plugin", plugin_key: "value"}]} + + assert BuilderInput(inline=inline).inline == inline + + +def test_builder_input_accepts_exported_and_shorthand_configs() -> None: + exported = DataDesignerConfigBuilder(model_configs=[]).get_builder_config().to_dict() + + assert BuilderInput(inline=exported).inline == exported + assert BuilderInput(inline={"columns": []}).inline == {"columns": []} + + +def test_builder_input_accepts_null_secret_fields_from_canonical_export() -> None: + builder = DataDesignerConfigBuilder(model_configs=[]).with_seed_dataset( + HuggingFaceSeedSource(path="datasets/example/seed/*.parquet") + ) + exported = builder.get_builder_config().to_dict() + + assert exported["data_designer"]["seed_config"]["source"]["token"] is None + assert BuilderInput(inline=exported).inline == exported + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_builder_input_rejects_non_finite_json(value: float) -> None: + with pytest.raises(ValidationError): + BuilderInput(inline={"columns": [], "value": value}) + + +@pytest.mark.parametrize( + "inline", + [ + {"data_designer": {}, "library_version": 1}, + {"data_designer": {}, "unknown": True}, + {"model_configs": []}, + ], +) +def test_builder_input_rejects_invalid_serialized_shapes(inline: dict[str, object]) -> None: + with pytest.raises(ValidationError, match="complete serialized"): + BuilderInput.model_validate({"inline": inline}) + + +def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["invocation"]["run_config"]["buffer_size"] = 0 + with pytest.raises(ValidationError, match="buffer_size"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["invocation"]["run_config"]["display_tui"] = True + with pytest.raises(ValidationError, match="display_tui"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["output"] = {"format": "jsonl", "partitions": 1} + with pytest.raises(ValidationError, match="output"): + DataDesignerSlurmConfig.model_validate(payload) + + payload = authored_run.model_dump(mode="json") + payload["array_tasks"]["count"] = 101 + payload["array_tasks"]["max_concurrent"] = 1 + with pytest.raises(ValidationError, match="requested records"): + DataDesignerSlurmConfig.model_validate(payload) + + +@pytest.mark.parametrize("readiness_path", ["health", "//other-host/health", "/health check"]) +def test_small_config_values_validate_at_boundary(readiness_path: str) -> None: + assert ArrayTasksConfig(count=2).max_concurrent is None + with pytest.raises(ValidationError, match="concurrency"): + ArrayTasksConfig(count=2, max_concurrent=3) + with pytest.raises(ValidationError, match="minutes"): + SubmissionConfig(time_limit="00:60:00") + with pytest.raises(ValidationError, match="absolute URL path"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), readiness_path=readiness_path) + + +@pytest.mark.parametrize( + "source", + ["nvcr.io/example/vllm:latest", "relative.sqsh"], +) +def test_image_build_request_rejects_mutable_or_relative_source(source: str) -> None: + with pytest.raises(ValidationError): + ImageBuildRequest(name="vllm", kind="serving", source=source) + + +def test_image_inspection_rejects_duplicate_distribution_names() -> None: + payload = { + "schema_version": 1, + "inspector_version": "v1", + "sqsh_sha256": "a" * 64, + "inspection": { + "kind": "client", + "python_implementation": "cpython", + "python_version": "3.12.1", + "python_abi": "cp312", + "distributions": [ + {"name": "plugin", "version": "1"}, + {"name": "plugin", "version": "2"}, + ], + "installer_path": "/usr/bin/pip", + "installer_version": "1", + }, + } + + with pytest.raises(ValidationError, match="unique"): + ImageInspectionRecord.model_validate_json(json.dumps(payload)) + + +def test_benchmark_config_rejects_duplicate_axes() -> None: + payload = { + "schema_version": 1, + "name": "bench", + "base_run": "./run.yaml", + "model_aliases": ["generator", "generator"], + "concurrency_values": [32, 32], + "deployment_cases": [{"name": "case", "deployments": {"generator": {"nodes": 1, "nodes_per_replica": 1}}}], + "record_policy": {"type": "fixed", "records": 100}, + "analysis": {"target_total_records": 1000, "target_runtime": "1h"}, + } + + with pytest.raises(ValidationError, match="aliases"): + DataDesignerSlurmBenchmarkConfig.model_validate(payload) + + payload["model_aliases"] = ["generator"] + with pytest.raises(ValidationError, match="concurrency"): + DataDesignerSlurmBenchmarkConfig.model_validate(payload) + + +def test_benchmark_base_run_normalizes_local_source() -> None: + base_run = BenchmarkBaseRun.model_validate("./run.yaml") + + assert base_run.source == "run.yaml" + + +def test_config_models_do_not_mutate_input(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + original = deepcopy(payload) + + DataDesignerSlurmConfig.model_validate(payload) + + assert payload == original + + +def test_config_models_are_deeply_immutable(authored_run: DataDesignerSlurmConfig) -> None: + inline = authored_run.builder.inline + assert inline is not None + data_designer = inline["data_designer"] + assert isinstance(data_designer, dict) + columns = data_designer["columns"] + assert isinstance(columns, list) + + with pytest.raises(TypeError, match="frozen list"): + authored_run.deployments.clear() + with pytest.raises(TypeError, match="frozen dictionary"): + authored_run.invocation.run_config["buffer_size"] = 1 + with pytest.raises(TypeError, match="frozen list"): + columns.append({}) diff --git a/packages/data-designer-slurm/tests/contracts/test_golden_records.py b/packages/data-designer-slurm/tests/contracts/test_golden_records.py new file mode 100644 index 000000000..712b158c9 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_golden_records.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageInspectionRecord, + SlurmProfileCatalog, +) +from data_designer.slurm.contracts import AuthoredConfig, ContractRecord +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +@pytest.mark.parametrize( + ("record_type", "filename"), + [ + (DataDesignerSlurmConfig, "authored_run.json"), + (DataDesignerSlurmConfig, "authored_run_single.json"), + (SlurmProfileCatalog, "profile_catalog.json"), + (ImageInspectionRecord, "client_image_inspection.json"), + (ImageInspectionRecord, "serving_image_inspection.json"), + (ResolvedDependencyLock, "dependency_lock.json"), + (ResolvedDependencyLock, "dependency_lock_single.json"), + (ResolvedSlurmRunPlan, "single_node_plan.json"), + (ResolvedSlurmRunPlan, "multi_node_plan.json"), + (ClientResult, "client_result.json"), + (DataDesignerSlurmBenchmarkConfig, "benchmark_config.json"), + (BenchmarkManifest, "benchmark_manifest.json"), + (BenchmarkReport, "benchmark_report.json"), + ], +) +def test_golden_record_round_trip(record_type: type[BaseModel], filename: str) -> None: + fixture = (GOLDEN_DIR / filename).read_text() + record = record_type.model_validate_json(fixture) + + assert record_type.model_validate_json(record.model_dump_json()) == record + if isinstance(record, (AuthoredConfig, ContractRecord)): + assert record_type.model_validate_json(record.serialize_json()) == record + assert record.compute_sha256() == hashlib.sha256(record.serialize_json().encode()).hexdigest() + assert record.serialize_json() == fixture + + +def test_golden_records_are_sanitized() -> None: + contents = "\n".join(path.read_text().casefold() for path in GOLDEN_DIR.glob("*.json")) + + assert "nvidia" not in contents + assert "secret-value" not in contents + assert "lustre" not in contents + + +def test_canonical_serialization_ignores_mapping_order(authored_run: DataDesignerSlurmConfig) -> None: + payload = authored_run.model_dump(mode="json") + payload["invocation"]["model_concurrency"] = {"judge": 32, "generator": 64} + reordered = DataDesignerSlurmConfig.model_validate(payload) + + assert authored_run.serialize_json() == reordered.serialize_json() + assert authored_run.compute_sha256() == reordered.compute_sha256() diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py new file mode 100644 index 000000000..6601b8db9 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -0,0 +1,551 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig +from data_designer.slurm.contracts import compute_serialized_json_sha256 +from data_designer.slurm.planning import ( + ArtifactReference, + ResolvedBuilderInput, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, +) +from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.validation import validate_resolved_plan + + +def test_resolved_topology_derives_all_resource_fields() -> None: + topology = ResolvedTopology.derive( + node_count=4, + gpus_per_node=8, + tensor_parallel=4, + nodes_per_replica=2, + ) + + assert topology == ResolvedTopology( + tensor_parallel=4, + nodes_per_replica=2, + pipeline_parallel=2, + node_group_count=2, + replicas_per_node_group=2, + replica_count=4, + gpus_per_replica=8, + ) + + +@pytest.mark.parametrize( + "inputs", + [ + {"node_count": 0, "gpus_per_node": 8, "tensor_parallel": 4, "nodes_per_replica": 2}, + {"node_count": 3, "gpus_per_node": 8, "tensor_parallel": 4, "nodes_per_replica": 2}, + {"node_count": 4, "gpus_per_node": 8, "tensor_parallel": 3, "nodes_per_replica": 2}, + ], +) +def test_resolved_topology_rejects_invalid_derivation_inputs(inputs: dict[str, int]) -> None: + with pytest.raises(ValueError): + ResolvedTopology.derive(**inputs) + + +def test_multi_node_plan_matches_authored_inputs( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + assert validate_resolved_plan(authored_run, dependency_lock, multi_node_plan) is multi_node_plan + assert multi_node_plan.authored_config.sha256 == authored_run.compute_sha256() + assert [deployment.topology.replica_count for deployment in multi_node_plan.deployments] == [1, 8] + assert multi_node_plan.client.gpu_count == 0 + + +def test_single_node_plan_matches_authored_inputs( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + assert validate_resolved_plan(authored_run_single, dependency_lock_single, single_node_plan) is single_node_plan + assert single_node_plan.authored_config.sha256 == authored_run_single.compute_sha256() + assert [deployment.topology.replica_count for deployment in single_node_plan.deployments] == [1] + assert single_node_plan.client.gpu_count == 0 + + +def test_plan_canonical_json_is_byte_stable(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = json.loads(multi_node_plan.serialize_json()) + payload["invocation"]["authored"]["model_concurrency"] = {"judge": 32, "generator": 64} + reordered = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + assert reordered.serialize_canonical_json() == multi_node_plan.serialize_canonical_json() + + +def test_resolved_plan_is_deeply_immutable(multi_node_plan: ResolvedSlurmRunPlan) -> None: + inline = multi_node_plan.builder.inline + assert inline is not None + data_designer = inline["data_designer"] + assert isinstance(data_designer, dict) + columns = data_designer["columns"] + assert isinstance(columns, list) + + with pytest.raises(TypeError, match="frozen dictionary"): + multi_node_plan.invocation.effective_run_config["buffer_size"] = 1 + with pytest.raises(TypeError, match="frozen list"): + columns.append({}) + + +def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.builder.model_dump(mode="json") + inline = payload["inline"] + assert isinstance(inline, dict) + inline["plugin_config"] = {"api_key": "VERY_SECRET_VALUE"} + payload["content_sha256"] = compute_serialized_json_sha256(inline) + + with pytest.raises(ValidationError, match="plaintext values under secret-bearing keys") as error: + ResolvedBuilderInput.model_validate_json(json.dumps(payload)) + + assert "VERY_SECRET_VALUE" not in str(error.value) + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload.pop("schema_version"), + lambda payload: payload.update(schema_version=2), + lambda payload: payload.update(unknown=True), + lambda payload: payload.update(resolved_gpus_per_node=4), + lambda payload: payload["client"].update(host_node_index=1), + lambda payload: payload["deployments"][1].update(node_indices=[1]), + lambda payload: payload["deployments"][0].update(deployment_id="unrelated-runtime-name"), + lambda payload: payload["deployments"][0]["ports"][1].update(port=18000), + lambda payload: payload["deployments"][1].update(ports=payload["deployments"][1]["ports"][:1]), + lambda payload: payload["client"].update(ports=payload["client"]["ports"][:1]), + lambda payload: payload.update(shards=payload["shards"][:1]), + lambda payload: payload["shards"][1]["record_range"].update(start_index=49), + lambda payload: payload["shards"][1].update(shard_id="shard-00002"), + lambda payload: payload["shards"][1].update(resume_workspace=payload["shards"][0]["resume_workspace"]), + lambda payload: payload.update(run_id="run-other"), + lambda payload: payload.update(container_mounts=[]), + lambda payload: payload["deployments"][0]["topology"].update(replica_count=2), + lambda payload: payload["invocation"]["effective_input_bindings"].update( + managed_assets_path="/workspace/other-assets" + ), + lambda payload: payload["runtime_bundle"].update(path="/tmp/" + "e" * 64 + ".tar.gz"), + lambda payload: payload["runtime_bundle"].update(path="/workspace/primary/runtime/runtime.tar.gz"), + ], +) +def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, mutator: object) -> None: + payload = deepcopy(multi_node_plan.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"] = {} + + with pytest.raises(ValidationError, match="fully materialized"): + ResolvedSlurmRunPlan.model_validate(payload) + + +def test_plan_requires_default_non_inference_worker_count( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + + with pytest.raises(ValidationError, match="Data Designer default"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_preserves_explicit_non_inference_worker_override( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["authored"]["run_config"]["non_inference_max_parallel_workers"] = 16 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 16 + + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + assert plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 16 + + +def test_plan_materializes_default_managed_assets_path(single_node_plan: ResolvedSlurmRunPlan) -> None: + assert single_node_plan.invocation.authored.input_bindings.managed_assets_path is None + assert single_node_plan.invocation.effective_input_bindings.managed_assets_path == ( + "/workspace/primary/managed-assets" + ) + + +def test_plan_rejects_managed_assets_overlapping_workspace_state( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = single_node_plan.model_dump(mode="json") + managed_assets_path = "/workspace/primary/runs" + payload["invocation"]["authored"]["input_bindings"]["managed_assets_path"] = managed_assets_path + payload["invocation"]["effective_input_bindings"]["managed_assets_path"] = managed_assets_path + + with pytest.raises(ValidationError, match="managed_assets_path"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_otel_port_collision(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["otel_metrics_port"] = 18000 + + with pytest.raises(ValidationError, match="OTEL metrics port collides"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"] = payload["builder"]["inline"]["data_designer"][ + "model_configs" + ][:1] + payload["builder"]["model_aliases"] = ["generator"] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) + + with pytest.raises(ValidationError, match="deployment alias"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_multi_node_tp4_requires_rendezvous_per_replica_lane(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.deployments[0].model_dump(mode="json") + payload["authored"]["topology"]["tensor_parallel"] = 4 + payload["topology"].update( + tensor_parallel=4, + replicas_per_node_group=2, + replica_count=2, + gpus_per_replica=8, + ) + payload["ports"] = [ + { + "name": "deployment-00000-http-00000", + "role": "http", + "node_index": 0, + "port": 18000, + }, + { + "name": "deployment-00000-http-00001", + "role": "http", + "node_index": 0, + "port": 18001, + }, + { + "name": "deployment-00000-rendezvous-00000", + "role": "rendezvous", + "node_index": 0, + "port": 19000, + }, + ] + + with pytest.raises(ValidationError, match="rendezvous"): + ResolvedDeployment.model_validate_json(json.dumps(payload)) + + payload["ports"].append( + { + "name": "deployment-00000-rendezvous-00001", + "role": "rendezvous", + "node_index": 0, + "port": 19001, + } + ) + assert ResolvedDeployment.model_validate_json(json.dumps(payload)).topology.replica_count == 2 + + +def test_sourced_builder_validation_requires_resolved_payload( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + sourced_authored = authored_run.model_copy(update={"builder": BuilderInput(source="builder.json")}) + builder_digest = compute_serialized_json_sha256(authored_run.builder.inline) + payload = multi_node_plan.model_dump(mode="json") + payload["authored_config"]["sha256"] = sourced_authored.compute_sha256() + payload["builder"] = { + "authored_source": "builder.json", + "source": {"path": "/workspace/primary/runs/run-001/builder-config.json", "sha256": builder_digest}, + "inline": None, + "content_sha256": builder_digest, + "model_aliases": ["generator", "judge"], + } + sourced_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + with pytest.raises(SlurmPlanContractError, match="resolved payload"): + validate_resolved_plan(sourced_authored, dependency_lock, sourced_plan) + + assert ( + validate_resolved_plan( + sourced_authored, + dependency_lock, + sourced_plan, + builder_payload=authored_run.builder.inline, + ) + is sourced_plan + ) + + +@pytest.mark.parametrize( + ("builder_update", "message"), + [ + ({"content_sha256": "a" * 64}, "builder digest"), + ({"model_aliases": ("other",)}, "model aliases"), + ], +) +def test_inline_builder_validation_rederives_identity( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + builder_update: dict[str, object], + message: str, +) -> None: + invalid = multi_node_plan.model_copy(update={"builder": multi_node_plan.builder.model_copy(update=builder_update)}) + + with pytest.raises(SlurmPlanContractError, match=message): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +@pytest.mark.parametrize( + "update", + [ + {"time_limit": "invalid"}, + {"comment": "bad\ncomment"}, + ], +) +def test_resolved_submission_preserves_authored_validation(update: dict[str, object]) -> None: + payload = {"job_name": "data-designer", "time_limit": "03:55:00", **update} + + with pytest.raises(ValidationError): + ResolvedSubmission.model_validate(payload) + + +def test_resolved_image_rejects_digest_mismatch(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["deployments"][0]["image"]["inspection"]["sqsh_sha256"] = "a" * 64 + + with pytest.raises(ValidationError, match="digest"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload["overlay_packages"].append( + { + "name": "data-designer", + "version": "0.9.2", + "artifact": {"path": "/wheels/data_designer.whl", "sha256": "a" * 64}, + } + ), + lambda payload: payload.update(base_distributions=list(reversed(payload["base_distributions"]))), + lambda payload: payload["overlay_packages"][0]["artifact"].update(path="/wheels/plugin.tar.gz"), + lambda payload: payload.update(authored_source="lock.json"), + lambda payload: payload.update( + authored_source="lock.yaml", + source={"path": "/workspace/lock.yaml", "sha256": "a" * 64}, + ), + ], +) +def test_dependency_lock_rejects_invalid_boundaries( + dependency_lock: ResolvedDependencyLock, + mutator: object, +) -> None: + payload = deepcopy(dependency_lock.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_requires_every_authored_package(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["overlay_packages"] = [] + + with pytest.raises(ValidationError, match="missing from the dependency lock"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_requires_compatible_versions(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["overlay_packages"][0]["version"] = "0.3.0" + + with pytest.raises(ValidationError, match="does not satisfy"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + +def test_dependency_lock_binds_direct_wheel_digest(dependency_lock: ResolvedDependencyLock) -> None: + payload = dependency_lock.model_dump(mode="json") + payload["authored_requirements"] = [ + "data-designer-speech @ https://example.test/data_designer_speech.whl#sha256=" + "a" * 64 + ] + + with pytest.raises(ValidationError, match="locked overlay artifact"): + ResolvedDependencyLock.model_validate_json(json.dumps(payload)) + + payload["overlay_packages"][0]["artifact"]["sha256"] = "a" * 64 + assert ResolvedDependencyLock.model_validate_json(json.dumps(payload)).overlay_packages + + +def test_cross_record_validation_rejects_authored_digest( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid = multi_node_plan.model_copy( + update={ + "authored_config": ArtifactReference( + path=multi_node_plan.authored_config.path, + sha256="0" * 64, + ) + } + ) + + with pytest.raises(SlurmPlanContractError, match="authored config digest"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_rejects_dependency_lock_digest( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": ArtifactReference( + path=multi_node_plan.client.dependency_lock.path, + sha256="0" * 64, + ) + } + ) + invalid = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(SlurmPlanContractError, match="dependency lock digest"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_binds_authored_lock_source( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + dependencies = ClientDependencies(requirements=None, lock_file="locks/user-lock.json") + authored = authored_run.model_copy( + update={"client": authored_run.client.model_copy(update={"dependencies": dependencies})} + ) + plan_payload = multi_node_plan.model_dump(mode="json") + plan_payload["authored_config"]["sha256"] = authored.compute_sha256() + plan_payload["client"]["authored"] = authored.client.model_dump(mode="json") + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(plan_payload)) + + with pytest.raises(SlurmPlanContractError, match="authored lock file"): + validate_resolved_plan(authored, dependency_lock, plan) + + lock_payload = dependency_lock.model_dump(mode="json") + lock_payload.update( + authored_source="locks/user-lock.json", + source={ + "path": "/workspace/primary/runs/run-001/inputs/user-lock.json", + "sha256": "a" * 64, + }, + ) + matching_lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) + unexpected_source_plan = multi_node_plan.model_copy( + update={ + "client": multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": matching_lock.compute_sha256()} + ) + } + ) + } + ) + with pytest.raises(SlurmPlanContractError, match="present for authored requirements"): + validate_resolved_plan(authored_run, matching_lock, unexpected_source_plan) + + matching_plan = plan.model_copy( + update={ + "client": plan.client.model_copy( + update={ + "dependency_lock": plan.client.dependency_lock.model_copy( + update={"sha256": matching_lock.compute_sha256()} + ) + } + ) + } + ) + + assert validate_resolved_plan(authored, matching_lock, matching_plan) is matching_plan + + +def test_cross_record_validation_rejects_python_abi( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid_lock = dependency_lock.model_copy(update={"python_abi": "cp311"}) + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": invalid_lock.compute_sha256()} + ) + } + ) + invalid_plan = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(SlurmPlanContractError, match="Python ABI"): + validate_resolved_plan(authored_run, invalid_lock, invalid_plan) + + +def test_cross_record_validation_rejects_runtime_inventory( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid_lock = dependency_lock.model_copy(update={"base_distributions": ()}) + client = multi_node_plan.client.model_copy( + update={ + "dependency_lock": multi_node_plan.client.dependency_lock.model_copy( + update={"sha256": invalid_lock.compute_sha256()} + ) + } + ) + invalid_plan = multi_node_plan.model_copy(update={"client": client}) + + with pytest.raises(SlurmPlanContractError, match="runtime inspection"): + validate_resolved_plan(authored_run, invalid_lock, invalid_plan) + + +def test_cross_record_validation_rejects_changed_invocation( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invocation = multi_node_plan.invocation.model_copy( + update={"authored": authored_run.invocation.model_copy(update={"dataset_name": "other"})} + ) + invalid = multi_node_plan.model_copy(update={"invocation": invocation}) + + with pytest.raises(SlurmPlanContractError, match="invocation"): + validate_resolved_plan(authored_run, dependency_lock, invalid) + + +def test_cross_record_validation_preserves_explicit_run_config( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["invocation"]["effective_run_config"]["buffer_size"] = 16384 + invalid = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + with pytest.raises(SlurmPlanContractError, match="run_config.buffer_size"): + validate_resolved_plan(authored_run, dependency_lock, invalid) diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py new file mode 100644 index 000000000..93e7aecd2 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import ( + ProfileSelectionSource, + SlurmProfile, + SlurmProfileCatalog, + injected_profile, + select_profile, + validate_selected_profile, +) + + +def test_profile_selection_precedence(profile_catalog: SlurmProfileCatalog) -> None: + explicit = select_profile(profile_catalog, cluster="lab", hostnames=("primary-login-1",)) + hostname = select_profile(profile_catalog, hostnames=("PRIMARY-LOGIN-1", "host.example")) + default = select_profile(profile_catalog, hostnames=("unmatched",)) + + assert (explicit.cluster_name, explicit.selection_source) == ("lab", ProfileSelectionSource.EXPLICIT) + assert (hostname.cluster_name, hostname.selection_source, hostname.matched_pattern) == ( + "primary", + ProfileSelectionSource.HOSTNAME, + "primary-login-*", + ) + assert (default.cluster_name, default.selection_source) == ("primary", ProfileSelectionSource.DEFAULT) + + +def test_injected_profile_has_no_catalog_provenance(profile_catalog: SlurmProfileCatalog) -> None: + selected = injected_profile(profile_catalog.clusters["lab"]) + + assert selected.selection_source is ProfileSelectionSource.INJECTED + assert selected.cluster_name is None + assert selected.catalog_sha256 is None + with pytest.raises(ValueError, match="no source catalog"): + validate_selected_profile(profile_catalog, selected) + + +def test_catalog_selection_digest_validation(profile_catalog: SlurmProfileCatalog) -> None: + selected = select_profile(profile_catalog, cluster="primary") + + assert validate_selected_profile(profile_catalog, selected) is selected + changed = profile_catalog.model_copy(update={"default_cluster": "lab"}) + with pytest.raises(ValueError, match="catalog digest"): + validate_selected_profile(changed, selected) + + +def test_catalog_selection_revalidates_provenance(profile_catalog: SlurmProfileCatalog) -> None: + forged_default = select_profile(profile_catalog, cluster="lab").model_copy( + update={"selection_source": ProfileSelectionSource.DEFAULT} + ) + with pytest.raises(ValueError, match="catalog default"): + validate_selected_profile(profile_catalog, forged_default) + + forged_pattern = select_profile(profile_catalog, hostnames=("primary-login-1",)).model_copy( + update={"matched_pattern": "lab-*"} + ) + with pytest.raises(ValueError, match="pattern"): + validate_selected_profile(profile_catalog, forged_pattern) + + +def test_unselected_profile_edit_keeps_selected_profile_digest(profile_catalog: SlurmProfileCatalog) -> None: + first = select_profile(profile_catalog, cluster="primary") + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["workspace_root"] = "/workspace/other-lab" + changed = SlurmProfileCatalog.model_validate(payload) + second = select_profile(changed, cluster="primary") + + assert first.profile_sha256 == second.profile_sha256 + assert first.catalog_sha256 != second.catalog_sha256 + + +def test_hostname_selection_rejects_ambiguous_clusters(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.model_dump(mode="json") + payload["clusters"]["lab"]["host_patterns"] = ["primary-*"] + catalog = SlurmProfileCatalog.model_validate(payload) + + with pytest.raises(ValueError, match="multiple"): + select_profile(catalog, hostnames=("primary-login-1",)) + + +def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfileCatalog) -> None: + with pytest.raises(ValueError, match="unknown cluster"): + select_profile(profile_catalog, cluster="missing") + + +@pytest.mark.parametrize( + "mutator", + [ + lambda payload: payload.pop("schema_version"), + lambda payload: payload.update(default_cluster="missing"), + lambda payload: payload["clusters"]["lab"].update(host_patterns=["primary-login-*"]), + lambda payload: payload["clusters"]["primary"].update(extra="unknown"), + lambda payload: payload["clusters"]["primary"].update(workspace_root="relative"), + lambda payload: payload["clusters"]["primary"]["scheduler"].update(bin_path="/opt/slurm:/usr/bin"), + lambda payload: payload["clusters"]["primary"].update(host_patterns=["login[broken"]), + lambda payload: payload["clusters"]["primary"]["image_build"].update(cpus_per_task=0), + lambda payload: payload["clusters"]["primary"]["image_build"].update(memory="0G"), + lambda payload: payload["clusters"]["primary"]["image_build"].update(time_limit="00:00:00"), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[ + {"source": "/one", "target": "/same"}, + {"source": "/two", "target": "/same"}, + ] + ), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[{"source": "/source", "target": "/run"}] + ), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[{"source": "/source", "target": "/run/data-designer-slurm"}] + ), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[{"source": "/source", "target": "/run/data-designer-slurm/nested"}] + ), + ], +) +def test_profile_catalog_rejects_invalid_boundaries( + profile_catalog: SlurmProfileCatalog, + mutator: object, +) -> None: + payload = deepcopy(profile_catalog.model_dump(mode="json")) + mutator(payload) + + with pytest.raises(ValidationError): + SlurmProfileCatalog.model_validate(payload) + + +def test_profile_requires_explicit_version(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.clusters["primary"].model_dump(mode="json") + payload.pop("schema_version") + + with pytest.raises(ValidationError, match="schema_version"): + SlurmProfile.model_validate(payload) + + +def test_profile_rejects_mem_per_gpu_without_gres(profile_catalog: SlurmProfileCatalog) -> None: + payload = profile_catalog.clusters["primary"].model_dump(mode="json") + payload["gpu_request_mode"] = "visible" + payload["scheduler"]["mem_per_gpu"] = "80G" + + with pytest.raises(ValidationError, match="requires GRES"): + SlurmProfile.model_validate(payload) diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py new file mode 100644 index 000000000..83d725e1b --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from copy import deepcopy +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.contracts import ( + ArtifactReference as CommonArtifactReference, +) +from data_designer.slurm.contracts import ( + RecordRange as CommonRecordRange, +) +from data_designer.slurm.contracts import ( + ResumeWorkspace as CommonResumeWorkspace, +) +from data_designer.slurm.planning import ( + ArtifactReference as PlanningArtifactReference, +) +from data_designer.slurm.planning import ( + RecordRange as PlanningRecordRange, +) +from data_designer.slurm.planning import ( + ResumeWorkspace as PlanningResumeWorkspace, +) + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +def test_planning_reexports_common_contract_types() -> None: + assert PlanningArtifactReference is CommonArtifactReference + assert PlanningRecordRange is CommonRecordRange + assert PlanningResumeWorkspace is CommonResumeWorkspace + + +@pytest.mark.parametrize("actual_records", [0, 25]) +def test_client_result_allows_partial_and_failure_facts_to_differ(actual_records: int) -> None: + partial = { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": actual_records, + "outcome": "partial", + "dataset_path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": True, + "requested_resume_mode": "if_possible", + "effective_resume_mode": "never", + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, + } + failed = { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0002", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": None, + "outcome": "failed", + "early_shutdown": None, + "requested_resume_mode": "if_possible", + "effective_resume_mode": None, + "error_code": "generation_error", + "redacted_message": "generation failed", + } + + assert ClientResult.model_validate_json(json.dumps(partial)).actual_records == actual_records + assert ClientResult.model_validate_json(json.dumps(failed)).dataset_path is None + + +@pytest.mark.parametrize( + "mutation", + [ + {"actual_records": 51}, + {"outcome": "complete", "actual_records": 49}, + {"outcome": "failed", "candidate_output_manifest": {"path": "/x", "sha256": "a" * 64}}, + {"outcome": "failed", "candidate_output_manifest": None, "error_code": None}, + {"effective_resume_mode": "always"}, + {"early_shutdown": None}, + {"early_shutdown": True}, + {"dataset_path": "/workspace/other/dataset"}, + { + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0002/output-manifest.json", + "sha256": "a" * 64, + } + }, + {"completed_at": "2026-08-19T12:00:00+01:00"}, + {"redacted_message": "bad\nmessage"}, + ], +) +def test_client_result_rejects_inconsistent_semantics( + mutation: dict[str, object], + client_result_payload: dict[str, object], +) -> None: + payload = deepcopy(client_result_payload) + payload.update(mutation) + + with pytest.raises(ValidationError): + ClientResult.model_validate_json(json.dumps(payload)) + + +@pytest.fixture +def client_result_payload() -> dict[str, object]: + return { + "schema_version": 1, + "run_id": "run-001", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "completed_at": "2026-08-19T12:00:00Z", + "requested_records": 50, + "actual_records": 50, + "outcome": "complete", + "dataset_path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": False, + "requested_resume_mode": "never", + "effective_resume_mode": "never", + "candidate_output_manifest": { + "path": "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "a" * 64, + }, + } + + +@pytest.mark.parametrize( + ("effective_resume_mode", "dataset_path"), + [ + ("never", "/workspace/runs/run-001/shards/shard-00000/attempts/attempt-0001/dataset"), + ("always", "/workspace/runs/run-001/shards/shard-00000/dataset"), + ], +) +def test_if_possible_uses_effective_resume_dataset_location( + effective_resume_mode: str, + dataset_path: str, + client_result_payload: dict[str, object], +) -> None: + payload = deepcopy(client_result_payload) + payload.update( + requested_resume_mode="if_possible", + effective_resume_mode=effective_resume_mode, + dataset_path=dataset_path, + ) + + assert ClientResult.model_validate_json(json.dumps(payload)).dataset_path == dataset_path + + +def test_always_allows_engine_resume_downgrade(client_result_payload: dict[str, object]) -> None: + payload = deepcopy(client_result_payload) + payload["requested_resume_mode"] = "always" + + assert ClientResult.model_validate_json(json.dumps(payload)).effective_resume_mode == "never" + + +def test_benchmark_manifest_rejects_duplicate_child_identity() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "benchmark_config": {"path": "/workspace/config.json", "sha256": "a" * 64}, + "children": [ + { + "case_id": "case", + "child_run_id": "run", + "child_authored_config": { + "path": "/workspace/runs/run/authored-config.json", + "sha256": "b" * 64, + }, + }, + { + "case_id": "case", + "child_run_id": "run-2", + "child_authored_config": { + "path": "/workspace/runs/run-2/authored-config.json", + "sha256": "c" * 64, + }, + }, + ], + } + + with pytest.raises(ValidationError, match="case IDs"): + BenchmarkManifest.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_rejects_unknown_or_incomplete_recommendations() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "analysis_id": "analysis", + "benchmark_manifest": {"path": "/workspace/benchmark.json", "sha256": "a" * 64}, + "created_at": "2026-08-19T12:00:00Z", + "cases": [ + { + "case_id": "case", + "child_run_id": "run", + "outcome": "pending", + "topology_digest": "b" * 64, + "requested_records": 100, + "gpus_per_job": 8, + "nodes_per_job": 1, + } + ], + "recommendations": [{"kind": "pareto", "case_id": "missing"}], + } + + with pytest.raises(ValidationError, match="unknown cases"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + payload["recommendations"] = [{"kind": "pareto", "case_id": "case"}] + with pytest.raises(ValidationError, match="successful feasible"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +def test_successful_benchmark_case_requires_metrics() -> None: + payload = { + "schema_version": 1, + "benchmark_id": "bench", + "analysis_id": "analysis", + "benchmark_manifest": {"path": "/workspace/benchmark.json", "sha256": "a" * 64}, + "created_at": "2026-08-19T12:00:00Z", + "cases": [ + { + "case_id": "case", + "child_run_id": "run", + "outcome": "succeeded", + "topology_digest": "b" * 64, + "requested_records": 100, + "gpus_per_job": 8, + "nodes_per_job": 1, + } + ], + } + + with pytest.raises(ValidationError, match="complete"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("actual_records", 0), + ("actual_records", 1001), + ("generation_seconds", 0), + ("wall_seconds", 0), + ("rows_per_second", 0), + ], +) +def test_successful_benchmark_case_requires_complete_positive_output(field: str, value: int) -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + payload["cases"][0][field] = value + + with pytest.raises(ValidationError): + 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]) + second_case.update(case_id="second-case", child_run_id="second-run", topology_digest="f" * 64) + payload["cases"].append(second_case) + payload["recommendations"] = [ + {"kind": "pareto", "case_id": payload["cases"][0]["case_id"]}, + {"kind": "pareto", "case_id": second_case["case_id"]}, + ] + + assert len(BenchmarkReport.model_validate_json(json.dumps(payload)).recommendations) == 2 + + payload["recommendations"] = [ + {"kind": "minimum_jobs", "case_id": payload["cases"][0]["case_id"]}, + {"kind": "minimum_jobs", "case_id": second_case["case_id"]}, + ] + with pytest.raises(ValidationError, match="minimum"): + BenchmarkReport.model_validate_json(json.dumps(payload)) + + +def test_benchmark_report_rejects_duplicate_child_runs() -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + payload["cases"][1]["child_run_id"] = payload["cases"][0]["child_run_id"] + + with pytest.raises(ValidationError, match="child run IDs"): + BenchmarkReport.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA new file mode 100644 index 000000000..1e52ff779 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA @@ -0,0 +1,4 @@ +Metadata-Version: 2.3 +Name: fake-data-designer-plugin +Version: 1.0.0 +Requires-Dist: data-designer diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt new file mode 100644 index 000000000..19e752123 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[data_designer.plugins] +fake-slurm-column = fake_data_designer_plugin.plugin:plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt new file mode 100644 index 000000000..7ca87bde3 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt @@ -0,0 +1 @@ +fake_data_designer_plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py new file mode 100644 index 000000000..aa059a7a1 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal Data Designer plugin used by Slurm client-environment tests.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py new file mode 100644 index 000000000..ea90846c8 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from data_designer.engine.column_generators.generators.base import ColumnGeneratorCellByCell +from fake_data_designer_plugin.plugin import FakePluginConfig + + +class FakePluginImplementation(ColumnGeneratorCellByCell[FakePluginConfig]): + """Add a deterministic marker through the real generator contract.""" + + def generate(self, data: dict[str, Any]) -> dict[str, Any]: + """Return the complete record with the plugin-owned column.""" + return {**data, self.config.name: "plugin-marker"} diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py new file mode 100644 index 000000000..dea69fc26 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Literal + +from data_designer.config.base import SingleColumnConfig +from data_designer.plugins import Plugin, PluginType + + +class FakePluginConfig(SingleColumnConfig): + """Minimal custom column configuration.""" + + column_type: Literal["fake-slurm-column"] = "fake-slurm-column" + model_alias: str | None = None + judge_model_alias: str | None = None + + @property + def required_columns(self) -> list[str]: + return [] + + @property + def side_effect_columns(self) -> list[str]: + return [] + + def get_model_aliases(self) -> list[str]: + return [alias for alias in (self.model_alias, self.judge_model_alias) if alias is not None] + + +plugin = Plugin( + config_qualified_name="fake_data_designer_plugin.plugin.FakePluginConfig", + impl_qualified_name="fake_data_designer_plugin.implementation.FakePluginImplementation", + plugin_type=PluginType.COLUMN_GENERATOR, +) diff --git a/packages/data-designer-slurm/tests/images/test_image_registry_records.py b/packages/data-designer-slurm/tests/images/test_image_registry_records.py new file mode 100644 index 000000000..cd883f943 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_image_registry_records.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.config import ClientImageInspection, ImageInspectionRecord, ImageKind, InstalledDistribution +from data_designer.slurm.images.records import ImageRegistryDocument, RegisteredImage + +_SQSH_SHA256 = "a" * 64 + + +def test_registered_image_exposes_inspected_kind() -> None: + image = _get_registered_image("alpha") + + assert image.kind is ImageKind.CLIENT + assert image.immutable_facts == (_SQSH_SHA256, None, image.inspection) + + +def test_registry_document_accepts_aliases_with_identical_facts() -> None: + alpha = _get_registered_image("alpha") + beta = alpha.model_copy(update={"name": "beta"}) + + document = ImageRegistryDocument(schema_version=1, images=(alpha, beta)) + + assert document.images == (alpha, beta) + + +def _get_registered_image(name: str) -> RegisteredImage: + inspection = ImageInspectionRecord( + schema_version=1, + inspector_version="inspector-1", + sqsh_sha256=_SQSH_SHA256, + inspection=ClientImageInspection( + kind="client", + python_implementation="cpython", + python_version="3.13.3", + python_abi="cp313", + distributions=(InstalledDistribution(name="data-designer", version="0.9.2"),), + installer_path="/usr/bin/pip", + installer_version="26.1", + ), + ) + return RegisteredImage( + schema_version=1, + name=name, + path="/workspace/images/shared.sqsh", + sqsh_sha256=_SQSH_SHA256, + inspection=inspection, + ) diff --git a/packages/data-designer-slurm/tests/images/test_inspection.py b/packages/data-designer-slurm/tests/images/test_inspection.py new file mode 100644 index 000000000..0b682c6e7 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_inspection.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.metadata +import re +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from slurm_test_fakes import FakeInspectionEnvironment + +from data_designer.slurm.config import ImageInspectionRecord, InstalledDistribution +from data_designer.slurm.images.errors import ImageInspectionError +from data_designer.slurm.images.inspection import ( + ClientImageInspector, + ServingImageInspector, + SystemInspectionEnvironment, +) +from data_designer.slurm.images.resources import inspect_image as resource_inspector + + +def _get_client_distributions() -> tuple[InstalledDistribution, ...]: + return ( + InstalledDistribution(name="aiohttp", version="3.14.3"), + InstalledDistribution(name="data-designer", version="0.9.2"), + InstalledDistribution(name="data-designer-config", version="0.9.2"), + InstalledDistribution(name="data-designer-engine", version="0.9.2"), + InstalledDistribution(name="data-designer-slurm", version="0.9.2"), + InstalledDistribution(name="pip", version="26.1"), + ) + + +def test_client_inspector_produces_digest_bound_golden_facts( + client_image_inspection: ImageInspectionRecord, +) -> None: + environment = FakeInspectionEnvironment( + distributions=_get_client_distributions(), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + + inspected = ClientImageInspector(environment).inspect(client_image_inspection.sqsh_sha256) + + assert inspected == client_image_inspection + + +def test_serving_inspector_produces_digest_bound_golden_facts( + serving_image_inspection: ImageInspectionRecord, +) -> None: + environment = FakeInspectionEnvironment( + distribution_versions={"vllm": "0.21.0"}, + executables={"vllm": "/usr/local/bin/vllm"}, + ) + + inspected = ServingImageInspector(environment).inspect(serving_image_inspection.sqsh_sha256) + + assert inspected == serving_image_inspection + + +@pytest.mark.parametrize( + "environment", + ( + FakeInspectionEnvironment( + distributions=_get_client_distributions(), + ), + FakeInspectionEnvironment( + distributions=tuple( + distribution for distribution in _get_client_distributions() if distribution.name != "pip" + ), + executables={"pip": "/usr/bin/pip"}, + ), + ), + ids=("missing-installer", "missing-installer-distribution"), +) +def test_client_inspector_rejects_missing_installer_capabilities( + environment: FakeInspectionEnvironment, +) -> None: + with pytest.raises(ImageInspectionError, match="required"): + ClientImageInspector(environment).inspect("a" * 64) + + +def test_client_inspector_rejects_generic_python_image_without_data_designer() -> None: + environment = FakeInspectionEnvironment( + distributions=(InstalledDistribution(name="pip", version="26.1"),), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + + with pytest.raises(ImageInspectionError, match="data-designer"): + ClientImageInspector(environment).inspect("a" * 64) + + +@pytest.mark.parametrize( + "missing_name", + ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm"), +) +def test_client_inspector_rejects_incomplete_data_designer_package_set(missing_name: str) -> None: + environment = FakeInspectionEnvironment( + distributions=tuple( + distribution for distribution in _get_client_distributions() if distribution.name != missing_name + ), + executables={"pip": "/usr/bin/pip"}, + ) + + with pytest.raises(ImageInspectionError, match=missing_name): + ClientImageInspector(environment).inspect("a" * 64) + + +@pytest.mark.parametrize( + "environment", + ( + FakeInspectionEnvironment(distribution_versions={"vllm": "0.21.0"}), + FakeInspectionEnvironment(executables={"vllm": "/usr/local/bin/vllm"}), + ), + ids=("missing-executable", "missing-runtime-distribution"), +) +def test_serving_inspector_rejects_missing_runtime_capabilities( + environment: FakeInspectionEnvironment, +) -> None: + with pytest.raises(ImageInspectionError, match="required"): + ServingImageInspector(environment).inspect("a" * 64) + + +def test_client_inspector_normalizes_invalid_facts_to_canonical_error() -> None: + environment = FakeInspectionEnvironment( + distributions=_get_client_distributions() + + ( + InstalledDistribution(name="plugin", version="1"), + InstalledDistribution(name="plugin", version="2"), + ), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + + with pytest.raises(ImageInspectionError, match="invalid facts"): + ClientImageInspector(environment).inspect("a" * 64) + + +def test_client_inspector_sorts_distribution_inventory() -> None: + environment = FakeInspectionEnvironment( + distributions=tuple(reversed(_get_client_distributions())), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + + inspection = ClientImageInspector(environment).inspect("a" * 64).inspection + + assert tuple(distribution.name for distribution in inspection.distributions) == tuple( + distribution.name for distribution in _get_client_distributions() + ) + + +def test_client_inspector_rejects_installer_version_outside_distribution_inventory() -> None: + environment = FakeInspectionEnvironment( + distributions=_get_client_distributions(), + distribution_versions={"pip": "unexpected"}, + executables={"pip": "/usr/bin/pip"}, + ) + + with pytest.raises(ImageInspectionError, match="does not match"): + ClientImageInspector(environment).inspect("a" * 64) + + +def test_system_inspection_environment_normalizes_distribution_inventory( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distributions = ( + SimpleNamespace(metadata={"Name": "My_Package"}, version="1.2.3"), + SimpleNamespace(metadata={"Name": "another.package"}, version="2"), + SimpleNamespace(metadata={"Name": "my-package"}, version="1.2.3"), + ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=distributions)) + + inspected = SystemInspectionEnvironment().list_distributions() + + assert tuple((distribution.name, distribution.version) for distribution in inspected) == ( + ("another-package", "2"), + ("my-package", "1.2.3"), + ) + + +@pytest.mark.parametrize( + "distributions", + ( + ( + SimpleNamespace(metadata={"Name": "plugin"}, version="1"), + SimpleNamespace(metadata={"Name": "plugin"}, version="2"), + ), + (SimpleNamespace(metadata={}, version="1"),), + ), + ids=("conflicting-versions", "missing-name"), +) +def test_system_inspection_environment_rejects_invalid_distribution_inventory( + monkeypatch: pytest.MonkeyPatch, + distributions: tuple[SimpleNamespace, ...], +) -> None: + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=distributions)) + + with pytest.raises(ImageInspectionError): + SystemInspectionEnvironment().list_distributions() + + +def test_system_inspection_environment_normalizes_missing_distribution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=())) + + with pytest.raises(ImageInspectionError, match="not installed"): + SystemInspectionEnvironment().get_distribution_console_script("missing") + + +def test_package_and_standalone_serving_inspectors_share_distribution_owned_console_script_semantics( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + executable_path = tmp_path / "active-environment" / "bin" / "vllm" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + shadow_path = tmp_path / "shadow" / "bin" / "vllm" + shadow_path.parent.mkdir(parents=True) + shadow_path.write_text("#!/bin/sh\n") + shadow_path.chmod(0o700) + distribution = SimpleNamespace( + metadata={"Name": "vllm"}, + version="0.21.0", + entry_points=(SimpleNamespace(group="console_scripts", name="vllm"),), + files=(Path("../../../bin/vllm"),), + locate_file=lambda _installed_file: executable_path, + ) + monkeypatch.setenv("PATH", shadow_path.parent.as_posix()) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=(distribution,))) + + package_record = ServingImageInspector().inspect("d" * 64) + standalone_record = ImageInspectionRecord.model_validate(resource_inspector.inspect_image("serving", "d" * 64)) + + assert package_record == standalone_record + assert package_record.inspection.executable_path == executable_path.as_posix() # type: ignore[union-attr] + + +def test_system_inspection_environment_rejects_unowned_console_script( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distribution = SimpleNamespace( + metadata={"Name": "pip"}, + version="26.1", + entry_points=(SimpleNamespace(group="console_scripts", name="pip"),), + files=(Path("pip"),), + locate_file=lambda installed_file: installed_file, + ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=(distribution,))) + + with pytest.raises(ImageInspectionError, match="does not own"): + SystemInspectionEnvironment().get_distribution_console_script("pip") + + +def test_package_and_standalone_client_inspectors_reject_duplicate_same_version_pip( + monkeypatch: pytest.MonkeyPatch, +) -> None: + distributions = ( + *( + SimpleNamespace(metadata={"Name": name}, version="1.0.0") + for name in ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm") + ), + SimpleNamespace(metadata={"Name": "pip"}, version="26.1"), + SimpleNamespace(metadata={"Name": "Pip"}, version="26.1"), + ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=distributions)) + + with pytest.raises(ImageInspectionError, match="exactly once"): + ClientImageInspector().inspect("a" * 64) + with pytest.raises(RuntimeError, match="exactly once"): + resource_inspector.inspect_image("client", "a" * 64) + + +def test_system_inspection_environment_reports_current_python_facts() -> None: + environment = SystemInspectionEnvironment() + + assert environment.get_python_implementation() + assert re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", environment.get_python_version()) + assert re.fullmatch(r"[A-Za-z0-9._-]+", environment.get_python_abi()) diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py new file mode 100644 index 000000000..8216b4830 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -0,0 +1,1704 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import stat +import subprocess +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from pathlib import Path +from threading import Barrier +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from pydantic import ValidationError +from slurm_test_fakes import FakeSlurmJob, FakeSlurmRunner + +import data_designer.slurm.images.lifecycle as image_lifecycle +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildProfile, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + SchedulerProfile, + SelectedSlurmProfile, + ServingImageInspection, + SlurmProfile, + injected_profile, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.images.errors import ( + ImageConflictError, + ImageLifecycleError, + ImageRegistryError, + ImageVerificationError, +) +from data_designer.slurm.images.lifecycle import ( + PreparedImageLifecycleJob, + cleanup_prepared_image_lifecycle, + prepare_image_lifecycle_job, + publish_completed_image_lifecycle, + render_image_lifecycle_script, + submit_prepared_image_lifecycle, +) +from data_designer.slurm.images.records import ImageLifecycleOperation, ImageLifecyclePlan +from data_designer.slurm.images.resources import inspect_image as resource_inspector +from data_designer.slurm.images.service import VerifiedImageRegistry +from data_designer.slurm.launcher.client import SlurmCommandClient + + +def test_prepare_oci_import_stages_checksum_bound_job_beneath_selected_workspace(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source_digest = "a" * 64 + request = ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/team/vllm@sha256:{source_digest}", + ) + + prepared = prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-0001") + + job_directory = workspace / "images" / ".tmp" / "jobs" / "image-job-0001" + assert prepared.plan.operation is ImageLifecycleOperation.IMPORT_OCI + assert prepared.plan.source_oci_digest == source_digest + assert prepared.plan.sqsh_path == (job_directory / "candidate.sqsh").as_posix() + assert prepared.plan.inspection_output_path == (job_directory / "output" / "inspection.json").as_posix() + assert prepared.plan_file.path == (job_directory / "image-lifecycle-plan.json").as_posix() + assert prepared.script_file.path == (job_directory / "image-lifecycle.sbatch").as_posix() + assert ImageLifecyclePlan.model_validate_json(Path(prepared.plan_file.path).read_text()) == prepared.plan + for artifact in ( + prepared.plan.inspector_script, + prepared.plan.enroot_rc, + prepared.plan_file, + prepared.script_file, + ): + content = Path(artifact.path).read_bytes() + assert hashlib.sha256(content).hexdigest() == artifact.sha256 + assert stat.S_IMODE(Path(prepared.plan.inspector_script.path).stat().st_mode) == 0o500 + assert stat.S_IMODE(Path(prepared.plan.enroot_rc.path).stat().st_mode) == 0o500 + assert stat.S_IMODE(job_directory.stat().st_mode) == 0o700 + assert stat.S_IMODE((job_directory / "output").stat().st_mode) == 0o700 + assert stat.S_IMODE(Path(prepared.plan_file.path).stat().st_mode) == 0o600 + assert stat.S_IMODE(Path(prepared.script_file.path).stat().st_mode) == 0o500 + + +def test_prepare_existing_sqsh_inspects_in_place_without_oci_identity(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = tmp_path / "external" / "client.sqsh" + request = ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()) + + prepared = prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-0002") + + assert prepared.plan.operation is ImageLifecycleOperation.INSPECT_SQSH + assert prepared.plan.sqsh_path == image_path.as_posix() + assert prepared.plan.source_oci_digest is None + script = Path(prepared.script_file.path).read_text() + assert "enroot import" not in script + assert f'readonly DD_IMAGE_SQSH="{image_path.as_posix()}"' in script + assert 'verify_enroot_compatibility 3 5 "existing SQSH inspection"' in script + assert "enroot start --root" in script + assert "inspect_image.py:x-create=file,bind,ro" in script + assert "/output:x-create=dir,bind" in script + assert f'DD_INSPECTION_SQSH="{prepared.plan.job_directory}/verified.sqsh"' in script + + +def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_path: Path) -> None: + workspace = tmp_path / "workspace-safe" + source = f"registry.example.test/team/image:release@sha256:{'b' * 64}" + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="serving", kind="serving", source=source), + _get_selected_profile(workspace), + lifecycle_id="image-job-0003", + ) + + script = render_image_lifecycle_script(prepared.plan) + + assert "#SBATCH --account=research" in script + assert "#SBATCH --partition=image-build" in script + assert "#SBATCH --nodes=1" in script + assert "#SBATCH --cpus-per-task=2" in script + assert "#SBATCH --mem=8G" in script + assert "#SBATCH --time=03:55:00" in script + assert "#SBATCH --gres=" not in script + expected_uri = f"docker://registry.example.test#team/image:sha256:{'b' * 64}" + assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in script + assert 'enroot import -o "${DD_IMAGE_SQSH}" "${DD_OCI_SOURCE}"' in script + assert "inspect_image.py:x-create=file,bind,ro" in script + assert "/opt/data-designer-slurm/output:x-create=dir,bind" in script + assert 'export HOME="${DD_SCRATCH_ROOT}/home"' in script + assert script.index("trap 'exit 143' TERM") < script.index("\ndd_initialize_local_scratch\n") + assert 'export ENROOT_CONFIG_PATH="${root}/enroot/config"' in script + assert 'export ENROOT_MAX_PROCESSORS="${SLURM_CPUS_PER_TASK}"' in script + assert script.index('export HOME="${DD_SCRATCH_ROOT}/home"') < script.rindex("\nverify_enroot_compatibility 3 5") + assert 'version="$(enroot version)"' in script + assert 'verify_enroot_compatibility 3 5 "digest-pinned OCI imports"' in script + assert f'--mount "{prepared.plan.job_directory}:' not in script + completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) + assert completed.returncode == 0, completed.stderr + + +@pytest.mark.parametrize( + ("source_kind", "enroot_version"), + (("oci", "3.5.0"), ("oci", "4.0.0"), ("existing", "3.5.0")), + ids=("oci-import-enroot-3.5", "oci-import-enroot-4", "existing-sqsh-enroot-3.5"), +) +def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( + tmp_path: Path, + source_kind: str, + enroot_version: str, +) -> None: + if source_kind == "oci": + source = f"registry.example.test/client@sha256:{'a' * 64}" + else: + image_path = tmp_path / "client.sqsh" + image_path.write_bytes(b"client image") + source = image_path.as_posix() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id=f"image-job-{source_kind}-execution", + ) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_enroot = fake_bin / "enroot" + fake_enroot.write_text( + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + ': "${HOME:?HOME must be set}"\n' + 'printf \'%s\\n\' "${HOME}" >> "${DD_TEST_HOME_LOG}"\n' + 'printf \'%s\\n\' "$1" >> "${DD_TEST_ENROOT_LOG}"\n' + 'if [[ "$1" == "version" ]]; then\n' + f" printf '%s\\n' '{enroot_version}'\n" + 'elif [[ "$1" == "import" ]]; then\n' + ' [[ "$4" == docker://*":sha256:"* ]]\n' + ' [[ "$4" != *"@sha256:"* ]]\n' + " printf '%s\\n' 'imported image' > \"$3\"\n" + 'elif [[ "$1" == "start" ]]; then\n' + ' [[ " $* " == *" --root "* ]]\n' + ' [[ "$*" == *":x-create=file,bind,ro"* ]]\n' + ' [[ "$*" == *":x-create=dir,bind"* ]]\n' + " printf '%s\\n' '{}' > \"${DD_TEST_INSPECTION_OUTPUT}\"\n" + "fi\n" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + home_log = tmp_path / "home.log" + enroot_log = tmp_path / "enroot.log" + + completed = subprocess.run( + ("bash",), + input=script, + capture_output=True, + text=True, + check=False, + env={ + "DD_TEST_ENROOT_LOG": enroot_log.as_posix(), + "DD_TEST_HOME_LOG": home_log.as_posix(), + "DD_TEST_INSPECTION_OUTPUT": prepared.plan.inspection_output_path, + "SLURM_CPUS_PER_TASK": "2", + "SLURM_JOB_ID": "5101", + }, + ) + + assert completed.returncode == 0, completed.stderr + expected_home = Path(home_log.read_text().strip()) + assert expected_home.name == "home" + assert expected_home.parent.name == "data-designer-slurm-5101-0" + assert not expected_home.parent.exists() + assert Path(prepared.plan.inspection_output_path).read_text() == "{}\n" + assert (Path(prepared.plan.job_directory) / "phase").read_text() == "inspection_complete\n" + if source_kind == "oci": + assert enroot_log.read_text().splitlines() == ["version", "import", "create", "start", "remove"] + assert Path(prepared.plan.sqsh_path).read_text() == "imported image\n" + else: + assert enroot_log.read_text().splitlines() == ["version", "create", "start", "remove"] + + +@pytest.mark.parametrize( + ("source_kind", "enroot_version", "required_version"), + (("oci", "3.4.1", "3.5"), ("existing", "3.4.1", "3.5")), + ids=("digest-import-before-3.5", "existing-sqsh-before-3.5"), +) +def test_rendered_image_lifecycle_rejects_unsupported_enroot_before_image_operations( + tmp_path: Path, + source_kind: str, + enroot_version: str, + required_version: str, +) -> None: + if source_kind == "oci": + source = f"registry.example.test/client@sha256:{'a' * 64}" + else: + image_path = tmp_path / "client.sqsh" + image_path.write_bytes(b"client image") + source = image_path.as_posix() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id=f"image-job-{source_kind}-old-enroot", + ) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_enroot = fake_bin / "enroot" + fake_enroot.write_text( + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + 'printf \'%s\\n\' "$1" >> "${DD_TEST_ENROOT_LOG}"\n' + 'if [[ "$1" == "version" ]]; then\n' + f" printf '%s\\n' '{enroot_version}'\n" + "fi\n" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + enroot_log = tmp_path / "enroot.log" + + completed = subprocess.run( + ("bash",), + input=script, + capture_output=True, + text=True, + check=False, + env={ + "DD_TEST_ENROOT_LOG": enroot_log.as_posix(), + "SLURM_CPUS_PER_TASK": "2", + "SLURM_JOB_ID": "5101", + }, + ) + + assert completed.returncode == 78 + assert f"Enroot {required_version} or newer" in completed.stderr + assert enroot_log.read_text() == "version\n" + + +@pytest.mark.parametrize( + "source", + ( + f"user:token@registry.example.test/image@sha256:{'b' * 64}", + f"docker://registry.example.test/image@sha256:{'b' * 64}", + f"registry.example.test/image?token=value@sha256:{'b' * 64}", + ), + ids=("credentials", "scheme", "query"), +) +def test_prepare_rejects_oci_sources_that_could_persist_credentials(tmp_path: Path, source: str) -> None: + request = ImageBuildRequest(name="serving", kind="serving", source=source) + + with pytest.raises(ImageLifecycleError, match="credential-free"): + prepare_image_lifecycle_job( + request, + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-secret", + ) + + assert not (tmp_path / "workspace" / "images").exists() + + +@pytest.mark.parametrize( + ("source", "expected_uri"), + ( + ( + f"vllm/vllm-openai@sha256:{'a' * 64}", + f"docker://registry-1.docker.io#vllm/vllm-openai:sha256:{'a' * 64}", + ), + ( + f"example@sha256:{'a' * 64}", + f"docker://registry-1.docker.io#library/example:sha256:{'a' * 64}", + ), + ( + f"docker.io/vllm/vllm-openai:v0.22.0@sha256:{'a' * 64}", + f"docker://registry-1.docker.io#vllm/vllm-openai:sha256:{'a' * 64}", + ), + ( + f"vllm/vllm-openai:v0.22.0@sha256:{'a' * 64}", + f"docker://registry-1.docker.io#vllm/vllm-openai:sha256:{'a' * 64}", + ), + ( + f"registry.example.test:5000/team/image@sha256:{'a' * 64}", + f"docker://registry.example.test:5000#team/image:sha256:{'a' * 64}", + ), + ( + f"registry.example.test:5000/team/image:release@sha256:{'a' * 64}", + f"docker://registry.example.test:5000#team/image:sha256:{'a' * 64}", + ), + ), + ids=( + "docker-hub-namespace", + "docker-hub-library", + "explicit-docker-hub", + "versioned-vllm", + "explicit-registry", + "tagged-source", + ), +) +def test_renderer_normalizes_digest_qualified_sources_for_enroot( + tmp_path: Path, + source: str, + expected_uri: str, +) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="serving", kind="serving", source=source), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-uri", + ) + + assert f'readonly DD_OCI_SOURCE="{expected_uri}"' in render_image_lifecycle_script(prepared.plan) + + +def test_prepare_refuses_duplicate_and_invalid_lifecycle_ids(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + profile = _get_selected_profile(workspace) + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + original = prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + original_job_directory = Path(original.plan.job_directory) + sentinel = original_job_directory / "keep-existing-job" + sentinel.write_text("active") + + with pytest.raises(ImageLifecycleError, match="cannot prepare"): + prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + with pytest.raises(ImageLifecycleError, match="ID is invalid"): + prepare_image_lifecycle_job(request, profile, lifecycle_id="../escape") # type: ignore[arg-type] + assert sentinel.read_text() == "active" + assert Path(original.plan_file.path).is_file() + assert Path(original.script_file.path).is_file() + assert not (workspace / "images" / ".tmp" / "escape").exists() + + +@pytest.mark.parametrize( + "workspace_name", ("workspace unsafe", "workspace\\unsafe", "workspace:unsafe", "workspace,unsafe") +) +def test_prepare_rejects_workspace_paths_that_enroot_cannot_mount( + tmp_path: Path, + workspace_name: str, +) -> None: + workspace = tmp_path / workspace_name + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + + with pytest.raises(ImageLifecycleError, match="Enroot mount"): + prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-unsafe") + + assert not (workspace / "images").exists() + + +def test_rendered_oci_import_rejects_dangling_candidate_symlink(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="client", + kind="client", + source=f"registry.example.test/client@sha256:{'a' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-dangling-candidate", + ) + candidate = Path(prepared.plan.sqsh_path) + outside = tmp_path / "outside.sqsh" + candidate.symlink_to(outside) + + completed = subprocess.run( + ("bash",), + input=render_image_lifecycle_script(prepared.plan), + capture_output=True, + text=True, + check=False, + env={"SLURM_CPUS_PER_TASK": "2", "SLURM_JOB_ID": "5101"}, + ) + + assert completed.returncode == 73 + assert candidate.is_symlink() + assert not outside.exists() + + +def test_prepare_rejects_symlinked_package_owned_image_root(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + (workspace / "images").symlink_to(outside, target_is_directory=True) + request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) + + with pytest.raises(ImageLifecycleError, match="cannot prepare"): + prepare_image_lifecycle_job(request, _get_selected_profile(workspace), lifecycle_id="image-job-symlink") + + assert tuple(outside.iterdir()) == () + + +@pytest.mark.parametrize( + ("field", "value", "match"), + ( + ("operation", "inspect_sqsh", "require an import"), + ("sqsh_path", "/workspace/other.sqsh", "attempt-local"), + ("source_oci_digest", "f" * 64, "does not match"), + ("inspection_output_path", "/workspace/inspection.json", "output directory"), + ), +) +def test_image_lifecycle_plan_rejects_mismatched_oci_facts( + tmp_path: Path, + field: str, + value: str, + match: str, +) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/vllm@sha256:{'e' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-contract", + ) + payload = prepared.plan.model_dump(mode="json") + payload[field] = value + + with pytest.raises(ValidationError, match=match): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_image_lifecycle_plan_rejects_non_package_inspector_path(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-runtime-path", + ) + payload = prepared.plan.model_dump(mode="json") + payload["inspector_script"]["path"] = f"{prepared.plan.job_directory}/other.py" + + with pytest.raises(ValidationError, match="package-owned job paths"): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_image_lifecycle_plan_rejects_credential_bearing_oci_source(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/vllm@sha256:{'e' * 64}", + ), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-source-contract", + ) + payload = prepared.plan.model_dump(mode="json") + payload["request"]["source"] = f"user:token@registry.example.test/vllm@sha256:{'e' * 64}" + + with pytest.raises(ValidationError, match="credential-free"): + ImageLifecyclePlan.model_validate_json(json.dumps(payload)) + + +def test_submit_prepared_image_lifecycle_uses_isolated_sbatch_client(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0005", + ) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + + receipt = submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) + + assert receipt.job_id == 5101 + assert runner.calls == [ + ( + "sbatch", + "--parsable", + "--export=ALL", + ) + ] + assert runner.inputs == [render_image_lifecycle_script(prepared.plan)] + + +@pytest.mark.parametrize("artifact_name", ("plan_file", "script_file"), ids=("plan", "script")) +def test_submit_rejects_modified_prepared_files(tmp_path: Path, artifact_name: str) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0006", + ) + artifact = getattr(prepared, artifact_name) + artifact_path = Path(artifact.path) + artifact_path.chmod(0o700) + artifact_path.write_text("modified\n") + + with pytest.raises(ImageLifecycleError, match=f"{artifact_name.removesuffix('_file')} no longer matches"): + submit_prepared_image_lifecycle(prepared, SlurmCommandClient(FakeSlurmRunner())) + + +def test_submit_rejects_recreated_prepared_job_directory(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-recreated-before-submit", + ) + job_directory = Path(prepared.plan.job_directory) + detached_job_directory = job_directory.with_name("detached-before-submit") + job_directory.rename(detached_job_directory) + shutil.copytree(detached_job_directory, job_directory) + runner = FakeSlurmRunner() + + with pytest.raises(ImageLifecycleError, match="directory no longer matches its identity"): + submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) + + assert runner.calls == [] + + +def test_submit_rejects_modified_script_even_when_artifact_digest_is_rebound(tmp_path: Path) -> None: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-0007", + ) + script_path = Path(prepared.script_file.path) + script_path.chmod(0o700) + modified_content = b"#!/usr/bin/env bash\nexit 0\n" + script_path.write_bytes(modified_content) + rebound = replace( + prepared, + script_file=ArtifactReference( + path=script_path.as_posix(), + sha256=hashlib.sha256(modified_content).hexdigest(), + ), + ) + + with pytest.raises(ImageLifecycleError, match="script no longer matches"): + submit_prepared_image_lifecycle(rebound, SlurmCommandClient(FakeSlurmRunner())) + + +def test_submit_uses_verified_script_bytes_when_path_changes_during_submission( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("NEMO_DEPLOYMENT_TYPE", raising=False) + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="image-job-submission-race", + ) + expected_script = render_image_lifecycle_script(prepared.plan) + runner = _MutatingSubmissionRunner(Path(prepared.script_file.path)) + + receipt = submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) + + assert receipt.job_id == 5101 + assert runner.command == ("sbatch", "--parsable", "--export=ALL") + assert runner.environment == {"NEMO_DEPLOYMENT_TYPE": "slurm"} + assert runner.input_text == expected_script + assert Path(prepared.script_file.path).read_text() == "replaced after verification\n" + + +@pytest.mark.parametrize("source_kind", ("oci", "existing"), ids=("oci-import", "existing-sqsh")) +def test_publish_completed_lifecycle_registers_digest_bound_image_in_fresh_process( + tmp_path: Path, + source_kind: str, +) -> None: + workspace = tmp_path / "workspace" + content = f"{source_kind} image".encode() + if source_kind == "oci": + source_digest = "a" * 64 + source = f"registry.example.test/client@sha256:{source_digest}" + else: + source_digest = None + source = (tmp_path / "external" / "client.sqsh").as_posix() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=source), + _get_selected_profile(workspace), + lifecycle_id=f"publish-{source_kind}", + ) + source_path = Path(prepared.plan.sqsh_path) + source_path.parent.mkdir(parents=True, exist_ok=True) + source_path.write_bytes(content) + _write_completed_inspection(prepared, _get_client_inspection(content)) + + registered = publish_completed_image_lifecycle(prepared) + + expected_path = source_path + if source_kind == "oci": + expected_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(content).hexdigest()}.sqsh" + assert registered.path == expected_path.as_posix() + assert registered.source_oci_digest == source_digest + assert expected_path.read_bytes() == content + assert not Path(prepared.plan.job_directory).exists() + resolved = VerifiedImageRegistry(workspace).resolve_for_planning( + ImageRef(name="client"), + expected_kind=ImageKind.CLIENT, + ) + assert resolved.path == expected_path.as_posix() + assert resolved.sha256 == hashlib.sha256(content).hexdigest() + + +def test_publish_versioned_vllm_rejects_inspected_version_mismatch(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + content = b"unexpected vllm version" + source = f"vllm/vllm-openai:v0.22.0@sha256:{'a' * 64}" + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="vllm-v0-22-0", kind="serving", source=source), + _get_selected_profile(workspace), + lifecycle_id="version-mismatch", + ) + candidate = Path(prepared.plan.sqsh_path) + candidate.parent.mkdir(parents=True, exist_ok=True) + candidate.write_bytes(content) + _write_completed_inspection(prepared, _get_serving_inspection(content, "0.21.0")) + + with pytest.raises(ImageVerificationError, match="does not match the requested image version"): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not (workspace / "images" / "artifacts").exists() + + +def test_publish_versioned_vllm_records_digest_and_verified_version(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + content = b"expected vllm version" + source_digest = "a" * 64 + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="vllm-v0-22-0", kind="serving", source=f"vllm/vllm-openai:v0.22.0@sha256:{source_digest}" + ), + _get_selected_profile(workspace), + lifecycle_id="version-match", + ) + candidate = Path(prepared.plan.sqsh_path) + candidate.parent.mkdir(parents=True, exist_ok=True) + candidate.write_bytes(content) + _write_completed_inspection(prepared, _get_serving_inspection(content, "0.22.0")) + + image = publish_completed_image_lifecycle(prepared) + + assert image.source_oci_digest == source_digest + assert isinstance(image.inspection.inspection, ServingImageInspection) + assert image.inspection.inspection.runtime_version == "0.22.0" + + +def test_publish_rejects_recreated_job_directory_with_substituted_result(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="recreated-before-publication", + content=b"original", + ) + _write_completed_inspection(prepared, _get_client_inspection(b"original")) + job_directory = Path(prepared.plan.job_directory) + detached_job_directory = job_directory.with_name("detached-before-publication") + job_directory.rename(detached_job_directory) + shutil.copytree(detached_job_directory, job_directory) + Path(prepared.plan.sqsh_path).write_bytes(b"substituted") + _write_completed_inspection(prepared, _get_client_inspection(b"substituted")) + + with pytest.raises(ImageLifecycleError, match="directory no longer matches its identity"): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert Path(prepared.plan.sqsh_path).read_bytes() == b"substituted" + assert detached_job_directory.is_dir() + + +def test_publish_reads_original_job_descriptor_after_directory_replacement(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="replaced-during-publication", + content=b"original", + ) + _write_completed_inspection(prepared, _get_client_inspection(b"original")) + job_directory = Path(prepared.plan.job_directory) + detached_job_directory = job_directory.with_name("detached-during-publication") + original_read_regular_file = image_lifecycle._read_regular_file + replaced = False + + def replace_job_directory_before_read( + directory_descriptor: int, + name: str, + display_path: Path, + *, + maximum_size: int | None = None, + ) -> bytes: + nonlocal replaced + if not replaced: + replaced = True + job_directory.rename(detached_job_directory) + shutil.copytree(detached_job_directory, job_directory) + Path(prepared.plan.sqsh_path).write_bytes(b"substituted") + _write_completed_inspection(prepared, _get_client_inspection(b"substituted")) + return original_read_regular_file( + directory_descriptor, + name, + display_path, + maximum_size=maximum_size, + ) + + with ( + patch( + "data_designer.slurm.images.lifecycle._read_regular_file", + side_effect=replace_job_directory_before_read, + ), + pytest.raises(ImageLifecycleError, match="cannot clean image lifecycle job"), + ): + publish_completed_image_lifecycle(prepared) + + registered = VerifiedImageRegistry(workspace).list_images() + assert len(registered) == 1 + assert Path(registered[0].path).read_bytes() == b"original" + assert Path(prepared.plan.sqsh_path).read_bytes() == b"substituted" + assert detached_job_directory.is_dir() + + +@pytest.mark.parametrize( + ("inspection_payload", "match"), + ( + (b"not-json\n", "invalid"), + (b"{}\n", "invalid"), + (b"x" * (1024 * 1024 + 1), "too large"), + ), + ids=("invalid-json", "incomplete-record", "oversized-record"), +) +def test_publish_rejects_invalid_inspection_and_cleans_temporary_state( + tmp_path: Path, + inspection_payload: bytes, + match: str, +) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="invalid-output", content=b"candidate") + Path(prepared.plan.inspection_output_path).write_bytes(inspection_payload) + + with pytest.raises(ImageLifecycleError, match=match): + publish_completed_image_lifecycle(prepared) + + assert not Path(prepared.plan.job_directory).exists() + assert VerifiedImageRegistry(workspace).list_images() == () + assert not (workspace / "images" / "artifacts").exists() + + +def test_publish_rejects_fifo_inspection_without_blocking(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="fifo-inspection-output", + content=b"candidate", + ) + os.mkfifo(prepared.plan.inspection_output_path) + + with pytest.raises(ImageLifecycleError, match="not regular"): + publish_completed_image_lifecycle(prepared) + + assert not Path(prepared.plan.job_directory).exists() + assert VerifiedImageRegistry(workspace).list_images() == () + + +def test_publish_rejects_inspection_mutated_during_read(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="inspection-mutated-during-read", + content=b"candidate", + ) + inspection_path = Path(prepared.plan.inspection_output_path) + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + inspection_inode = inspection_path.stat().st_ino + original_fstat = os.fstat + inspection_fstat_count = 0 + + def mutate_before_final_inspection_fstat(descriptor: int) -> os.stat_result: + nonlocal inspection_fstat_count + status = original_fstat(descriptor) + if status.st_ino == inspection_inode: + inspection_fstat_count += 1 + if inspection_fstat_count == 2: + inspection_path.write_text("mutated") + status = original_fstat(descriptor) + return status + + with ( + patch("data_designer.slurm.images.lifecycle.os.fstat", side_effect=mutate_before_final_inspection_fstat), + pytest.raises(ImageLifecycleError, match="changed while it was being read"), + ): + publish_completed_image_lifecycle(prepared) + + assert not Path(prepared.plan.job_directory).exists() + assert VerifiedImageRegistry(workspace).list_images() == () + + +def test_cleanup_failure_preserves_primary_publication_error_and_keeps_output_unusable(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="cleanup-failure", content=b"candidate") + Path(prepared.plan.inspection_output_path).write_text("not-json") + + with ( + patch("data_designer.slurm.images.lifecycle.os.listdir", side_effect=OSError("injected cleanup failure")), + pytest.raises(ImageLifecycleError, match="inspection output is invalid"), + ): + publish_completed_image_lifecycle(prepared) + + assert Path(prepared.plan.job_directory).exists() + assert VerifiedImageRegistry(workspace).list_images() == () + assert not (workspace / "images" / "artifacts").exists() + cleanup_prepared_image_lifecycle(prepared) + + +def test_explicit_cleanup_normalizes_failure(tmp_path: Path) -> None: + prepared = _prepare_completed_oci_lifecycle( + tmp_path / "workspace", + lifecycle_id="explicit-cleanup-failure", + content=b"candidate", + ) + + with ( + patch("data_designer.slurm.images.lifecycle.os.listdir", side_effect=OSError("injected cleanup failure")), + pytest.raises(ImageLifecycleError, match="cannot clean image lifecycle job"), + ): + cleanup_prepared_image_lifecycle(prepared) + + cleanup_prepared_image_lifecycle(prepared) + + +def test_cleanup_refuses_replaced_job_directory(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="replaced-cleanup-directory", + content=b"candidate", + ) + job_directory = Path(prepared.plan.job_directory) + detached_job_directory = job_directory.with_name("detached-cleanup-directory") + job_directory.rename(detached_job_directory) + job_directory.mkdir() + replacement_sentinel = job_directory / "must-not-delete" + replacement_sentinel.write_text("replacement") + + with pytest.raises(ImageLifecycleError, match="cannot clean"): + cleanup_prepared_image_lifecycle(prepared) + + assert replacement_sentinel.read_text() == "replacement" + assert detached_job_directory.is_dir() + replacement_sentinel.unlink() + job_directory.rmdir() + detached_job_directory.rename(job_directory) + cleanup_prepared_image_lifecycle(prepared) + + +def test_cleanup_does_not_descend_into_directory_replaced_after_identity_check(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="replacement-during-cleanup", + content=b"candidate", + ) + job_directory = Path(prepared.plan.job_directory) + detached_job_directory = job_directory.with_name("detached-during-cleanup") + replacement_sentinel = job_directory / "must-not-delete" + original_listdir = os.listdir + replaced = False + + def replace_before_recursive_cleanup(directory_descriptor: int) -> list[str]: + nonlocal replaced + if not replaced: + replaced = True + job_directory.rename(detached_job_directory) + job_directory.mkdir() + replacement_sentinel.write_text("replacement") + return original_listdir(directory_descriptor) + + with ( + patch("data_designer.slurm.images.lifecycle.os.listdir", side_effect=replace_before_recursive_cleanup), + pytest.raises(ImageLifecycleError, match="cannot clean"), + ): + cleanup_prepared_image_lifecycle(prepared) + + assert replacement_sentinel.read_text() == "replacement" + assert detached_job_directory.is_dir() + replacement_sentinel.unlink() + job_directory.rmdir() + detached_job_directory.rename(job_directory) + cleanup_prepared_image_lifecycle(prepared) + + +def test_successful_publication_reports_cleanup_failure_without_invalidating_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="successful-cleanup-failure", + content=b"candidate", + ) + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + + with ( + patch("data_designer.slurm.images.lifecycle.os.listdir", side_effect=OSError("injected cleanup failure")), + pytest.raises(ImageLifecycleError, match="cannot clean image lifecycle job"), + ): + publish_completed_image_lifecycle(prepared) + + assert Path(prepared.plan.job_directory).exists() + resolved = VerifiedImageRegistry(workspace).resolve_for_planning( + ImageRef(name="client"), + expected_kind=ImageKind.CLIENT, + ) + assert Path(resolved.path).read_bytes() == b"candidate" + cleanup_prepared_image_lifecycle(prepared) + + +@pytest.mark.parametrize("mismatch", ("digest", "kind")) +def test_publish_rejects_mismatched_inspection_without_exposing_artifact( + tmp_path: Path, + mismatch: str, +) -> None: + workspace = tmp_path / "workspace" + content = b"candidate" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id=f"mismatch-{mismatch}", content=content) + inspection = ( + _get_client_inspection(b"different") if mismatch == "digest" else _get_serving_inspection(content, "0.21.0") + ) + _write_completed_inspection(prepared, inspection) + + with pytest.raises(ImageVerificationError, match="digest|kind"): + publish_completed_image_lifecycle(prepared) + + assert not Path(prepared.plan.job_directory).exists() + assert VerifiedImageRegistry(workspace).list_images() == () + assert not (workspace / "images" / "artifacts").exists() + + +def test_publish_requires_explicit_alias_replacement_and_preserves_prior_state(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="first", content=b"first") + _write_completed_inspection(first, _get_client_inspection(b"first")) + original = publish_completed_image_lifecycle(first) + second = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="second", content=b"second") + _write_completed_inspection(second, _get_client_inspection(b"second")) + + with pytest.raises(ImageConflictError, match="already registered"): + publish_completed_image_lifecycle(second) + + resolved = VerifiedImageRegistry(workspace).resolve_for_planning( + ImageRef(name="client"), + expected_kind=ImageKind.CLIENT, + ) + assert resolved.path == original.path + assert Path(original.path).read_bytes() == b"first" + assert not Path(second.plan.job_directory).exists() + assert tuple((workspace / "images" / "artifacts").glob("*.sqsh")) == (Path(original.path),) + + +def test_explicit_alias_replacement_publishes_new_artifact_without_deleting_prior_artifact(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="first-replace", content=b"first") + _write_completed_inspection(first, _get_client_inspection(b"first")) + original = publish_completed_image_lifecycle(first) + second = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="second-replace", content=b"second") + _write_completed_inspection(second, _get_client_inspection(b"second")) + + replacement = publish_completed_image_lifecycle(second, replace=True) + + assert replacement.path != original.path + assert Path(replacement.path).read_bytes() == b"second" + assert Path(original.path).read_bytes() == b"first" + assert VerifiedImageRegistry(workspace).list_images() == (replacement,) + + +@pytest.mark.parametrize("existing_content", (b"candidate", b"conflict"), ids=("identical", "conflicting")) +def test_artifact_path_collision_is_deterministic(tmp_path: Path, existing_content: bytes) -> None: + workspace = tmp_path / "workspace" + content = b"candidate" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="artifact-collision", content=content) + inspection = _get_client_inspection(content) + _write_completed_inspection(prepared, inspection) + artifact_path = workspace / "images" / "artifacts" / f"client-{inspection.sqsh_sha256}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(existing_content) + + if existing_content == content: + registered = publish_completed_image_lifecycle(prepared) + assert registered.path == artifact_path.as_posix() + assert artifact_path.read_bytes() == content + else: + with pytest.raises(ImageConflictError, match="conflicting bytes"): + publish_completed_image_lifecycle(prepared) + assert VerifiedImageRegistry(workspace).list_images() == () + assert artifact_path.read_bytes() == existing_content + assert not Path(prepared.plan.job_directory).exists() + + +@pytest.mark.parametrize("collision_kind", ("directory", "fifo", "symlink")) +def test_nonregular_artifact_path_collision_is_deterministic(tmp_path: Path, collision_kind: str) -> None: + workspace = tmp_path / "workspace" + content = b"candidate" + prepared = _prepare_completed_oci_lifecycle( + workspace, lifecycle_id="nonregular-artifact-collision", content=content + ) + inspection = _get_client_inspection(content) + _write_completed_inspection(prepared, inspection) + artifact_path = workspace / "images" / "artifacts" / f"client-{inspection.sqsh_sha256}.sqsh" + artifact_path.parent.mkdir(parents=True) + if collision_kind == "directory": + artifact_path.mkdir() + elif collision_kind == "fifo": + os.mkfifo(artifact_path) + else: + symlink_target = tmp_path / "symlink-target.sqsh" + symlink_target.write_bytes(b"outside") + artifact_path.symlink_to(symlink_target) + original_status = artifact_path.lstat() + + with pytest.raises(ImageConflictError, match="cannot be reused"): + publish_completed_image_lifecycle(prepared) + + current_status = artifact_path.lstat() + assert (current_status.st_dev, current_status.st_ino, stat.S_IFMT(current_status.st_mode)) == ( + original_status.st_dev, + original_status.st_ino, + stat.S_IFMT(original_status.st_mode), + ) + assert VerifiedImageRegistry(workspace).list_images() == () + assert not Path(prepared.plan.job_directory).exists() + + +def test_registry_failure_rolls_back_new_artifact_and_keeps_alias_unpublished(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="registry-failure", content=b"candidate") + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + original_replace = os.replace + + def fail_registry_publish(source: str, destination: str, **kwargs: int) -> None: + if destination == "registry.yaml": + assert not Path(prepared.plan.sqsh_path).exists() + raise OSError("injected registry publication failure") + original_replace(source, destination, **kwargs) + + with ( + patch("data_designer.slurm.images.registry.os.replace", side_effect=fail_registry_publish), + pytest.raises(ImageRegistryError, match="cannot persist"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not tuple((workspace / "images" / "artifacts").glob("*.sqsh")) + assert not Path(prepared.plan.job_directory).exists() + + +def test_registry_failure_does_not_remove_replaced_artifact_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="registry-failure-replaced-artifact", + content=b"candidate", + ) + inspection = _get_client_inspection(b"candidate") + _write_completed_inspection(prepared, inspection) + artifact_path = workspace / "images" / "artifacts" / f"client-{inspection.sqsh_sha256}.sqsh" + original_replace = os.replace + + def replace_artifact_before_registry_failure(source: str, destination: str, **kwargs: int) -> None: + if destination == "registry.yaml": + artifact_path.unlink() + artifact_path.write_bytes(b"replacement") + raise OSError("injected registry publication failure") + original_replace(source, destination, **kwargs) + + with ( + patch("data_designer.slurm.images.registry.os.replace", side_effect=replace_artifact_before_registry_failure), + pytest.raises(ImageRegistryError, match="cannot persist"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert artifact_path.read_bytes() == b"replacement" + assert not Path(prepared.plan.job_directory).exists() + + +def test_artifact_directory_sync_failure_removes_new_artifact_and_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="artifact-sync-failure", content=b"candidate") + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + artifact_directory = workspace / "images" / "artifacts" + original_fsync = os.fsync + + def fail_artifact_directory_sync(descriptor: int) -> None: + descriptor_status = os.fstat(descriptor) + if artifact_directory.exists(): + artifact_status = artifact_directory.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == ( + artifact_status.st_dev, + artifact_status.st_ino, + ): + raise OSError("injected sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.service.os.fsync", side_effect=fail_artifact_directory_sync), + pytest.raises(ImageVerificationError, match="cannot publish image artifact"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not tuple((workspace / "images" / "artifacts").glob("*.sqsh")) + assert not Path(prepared.plan.job_directory).exists() + + +def test_artifact_identity_lookup_failure_removes_new_artifact_and_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="artifact-identity-failure", + content=b"candidate", + ) + inspection = _get_client_inspection(b"candidate") + _write_completed_inspection(prepared, inspection) + artifact_path = workspace / "images" / "artifacts" / f"client-{inspection.sqsh_sha256}.sqsh" + original_stat = os.stat + failed = False + + def fail_artifact_identity_lookup(path: str, **kwargs: object) -> os.stat_result: + nonlocal failed + if path == artifact_path.name and not failed: + failed = True + raise OSError("injected artifact identity failure") + return original_stat(path, **kwargs) + + with ( + patch("data_designer.slurm.images.service.os.stat", side_effect=fail_artifact_identity_lookup), + pytest.raises(ImageVerificationError, match="cannot publish image artifact"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not artifact_path.exists() + assert not Path(prepared.plan.job_directory).exists() + + +def test_candidate_directory_sync_failure_removes_new_artifact_and_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="candidate-sync-failure", + content=b"candidate", + ) + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + candidate_directory = Path(prepared.plan.sqsh_path).parent + original_fsync = os.fsync + failed = False + + def fail_candidate_directory_sync(descriptor: int) -> None: + nonlocal failed + descriptor_status = os.fstat(descriptor) + candidate_directory_status = candidate_directory.stat() + if not failed and (descriptor_status.st_dev, descriptor_status.st_ino) == ( + candidate_directory_status.st_dev, + candidate_directory_status.st_ino, + ): + failed = True + raise OSError("injected candidate sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.service.os.fsync", side_effect=fail_candidate_directory_sync), + pytest.raises(ImageVerificationError, match="cannot publish image artifact"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not tuple((workspace / "images" / "artifacts").glob("*.sqsh")) + assert not Path(prepared.plan.job_directory).exists() + + +def test_artifact_file_sync_failure_removes_new_artifact_and_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, lifecycle_id="artifact-file-sync-failure", content=b"candidate" + ) + _write_completed_inspection(prepared, _get_client_inspection(b"candidate")) + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'candidate').hexdigest()}.sqsh" + original_fsync = os.fsync + + def fail_regular_file_sync(descriptor: int) -> None: + descriptor_status = os.fstat(descriptor) + if artifact_path.exists(): + artifact_status = artifact_path.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == ( + artifact_status.st_dev, + artifact_status.st_ino, + ): + raise OSError("injected file sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.service.os.fsync", side_effect=fail_regular_file_sync), + pytest.raises(ImageVerificationError, match="cannot synchronize image artifact"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert not tuple((workspace / "images" / "artifacts").glob("*.sqsh")) + assert not Path(prepared.plan.job_directory).exists() + + +def test_registry_directory_sync_failure_restores_prior_alias_and_artifact(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="before-registry-sync", content=b"first") + _write_completed_inspection(first, _get_client_inspection(b"first")) + original = publish_completed_image_lifecycle(first) + second = _prepare_completed_oci_lifecycle(workspace, lifecycle_id="registry-sync-failure", content=b"second") + _write_completed_inspection(second, _get_client_inspection(b"second")) + original_fsync = os.fsync + image_root = workspace / "images" + image_root_sync_count = 0 + + def fail_published_registry_sync(descriptor: int) -> None: + nonlocal image_root_sync_count + descriptor_status = os.fstat(descriptor) + image_root_status = image_root.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == ( + image_root_status.st_dev, + image_root_status.st_ino, + ): + image_root_sync_count += 1 + if image_root_sync_count == 2: + raise OSError("injected registry sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.registry.os.fsync", side_effect=fail_published_registry_sync), + pytest.raises(ImageRegistryError, match="cannot persist"), + ): + publish_completed_image_lifecycle(second, replace=True) + + registry = VerifiedImageRegistry(workspace) + assert registry.list_images() == (original,) + assert Path(original.path).read_bytes() == b"first" + assert tuple((workspace / "images" / "artifacts").glob("*.sqsh")) == (Path(original.path),) + assert not Path(second.plan.job_directory).exists() + + +def test_committed_marker_sync_failure_preserves_artifact_for_recovery(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id="committed-marker-sync-failure", + content=b"candidate", + ) + inspection = _get_client_inspection(b"candidate") + _write_completed_inspection(prepared, inspection) + artifact_path = workspace / "images" / "artifacts" / f"client-{inspection.sqsh_sha256}.sqsh" + image_root = workspace / "images" + image_root_sync_count = 0 + original_fsync = os.fsync + + def fail_committed_marker_sync(descriptor: int) -> None: + nonlocal image_root_sync_count + descriptor_status = os.fstat(descriptor) + image_root_status = image_root.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == ( + image_root_status.st_dev, + image_root_status.st_ino, + ): + image_root_sync_count += 1 + if image_root_sync_count == 3: + raise OSError("injected committed marker sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.registry.os.fsync", side_effect=fail_committed_marker_sync), + pytest.raises(ImageRegistryError, match="commit state requires recovery"), + ): + publish_completed_image_lifecycle(prepared) + + committed_marker = image_root / ".registry.committed.yaml" + assert committed_marker.is_file() + assert artifact_path.read_bytes() == b"candidate" + registered = VerifiedImageRegistry(workspace).list_images() + assert len(registered) == 1 + assert registered[0].path == artifact_path.as_posix() + assert not committed_marker.exists() + assert not Path(prepared.plan.job_directory).exists() + + +def test_existing_sqsh_replacement_after_registry_rename_rolls_back_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = tmp_path / "external" / "client.sqsh" + image_path.parent.mkdir() + image_path.write_bytes(b"client") + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _get_selected_profile(workspace), + lifecycle_id="existing-replacement", + ) + _write_completed_inspection(prepared, _get_client_inspection(b"client")) + replacement = tmp_path / "replacement.sqsh" + replacement.write_bytes(b"replacement") + original_replace = os.replace + + def replace_after_registry_publish(source: str, destination: str, **kwargs: int) -> None: + original_replace(source, destination, **kwargs) + if destination == "registry.yaml": + replacement.replace(image_path) + + with ( + patch("data_designer.slurm.images.registry.os.replace", side_effect=replace_after_registry_publish), + pytest.raises(ImageVerificationError, match="no longer matches"), + ): + publish_completed_image_lifecycle(prepared) + + assert VerifiedImageRegistry(workspace).list_images() == () + assert image_path.read_bytes() == b"replacement" + assert not Path(prepared.plan.job_directory).exists() + + +def test_concurrent_same_alias_publications_have_one_consistent_winner(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + barrier = Barrier(2) + + def publish(content: bytes) -> str: + prepared = _prepare_completed_oci_lifecycle( + workspace, + lifecycle_id=f"concurrent-{content.decode()}", + content=content, + ) + _write_completed_inspection(prepared, _get_client_inspection(content)) + barrier.wait() + return publish_completed_image_lifecycle(prepared).path + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(publish, content) for content in (b"first", b"second")) + + successes = tuple(future.result() for future in futures if future.exception() is None) + failures = tuple(future.exception() for future in futures if future.exception() is not None) + assert len(successes) == 1 + assert len(failures) == 1 + assert isinstance(failures[0], ImageConflictError) + registered = VerifiedImageRegistry(workspace).list_images() + assert len(registered) == 1 + assert registered[0].path == successes[0] + assert hashlib.sha256(Path(successes[0]).read_bytes()).hexdigest() == registered[0].sqsh_sha256 + assert tuple((workspace / "images" / "artifacts").glob("*.sqsh")) == (Path(successes[0]),) + + +def test_concurrent_aliases_for_same_existing_sqsh_preserve_both_bindings(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = tmp_path / "external" / "shared.sqsh" + image_path.parent.mkdir() + image_path.write_bytes(b"shared") + inspection = _get_client_inspection(b"shared") + barrier = Barrier(2) + + def publish(name: str) -> str: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name=name, kind="client", source=image_path.as_posix()), + _get_selected_profile(workspace), + lifecycle_id=f"shared-{name}", + ) + _write_completed_inspection(prepared, inspection) + barrier.wait() + return publish_completed_image_lifecycle(prepared).name + + with ThreadPoolExecutor(max_workers=2) as executor: + names = tuple(executor.map(publish, ("alpha", "beta"))) + + assert names == ("alpha", "beta") + registered = VerifiedImageRegistry(workspace).list_images() + assert tuple(image.name for image in registered) == ("alpha", "beta") + assert {image.path for image in registered} == {image_path.as_posix()} + + +def test_two_supported_serving_images_publish_and_resolve_independently(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + expected_versions = {"vllm-021": "0.21.3", "vllm-022": "0.22.1"} + + for name, runtime_version in expected_versions.items(): + content = f"{name}-{runtime_version}".encode() + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name=name, + kind="serving", + source=f"registry.example.test/{name}@sha256:{hashlib.sha256(name.encode()).hexdigest()}", + ), + _get_selected_profile(workspace), + lifecycle_id=f"publish-{name}", + ) + Path(prepared.plan.sqsh_path).write_bytes(content) + _write_completed_inspection(prepared, _get_serving_inspection(content, runtime_version)) + publish_completed_image_lifecycle(prepared) + + registry = VerifiedImageRegistry(workspace) + resolved = { + name: registry.resolve_for_planning(ImageRef(name=name), expected_kind=ImageKind.SERVING) + for name in expected_versions + } + assert {name: image.inspection_facts.runtime_version for name, image in resolved.items()} == expected_versions + assert resolved["vllm-021"].path != resolved["vllm-022"].path + assert resolved["vllm-021"].sha256 != resolved["vllm-022"].sha256 + + +def test_rendered_existing_inspection_rejects_path_replacement_and_cleans_output(tmp_path: Path) -> None: + image_path = tmp_path / "client.sqsh" + image_path.write_bytes(b"client image") + replacement = tmp_path / "replacement.sqsh" + replacement.write_bytes(b"replacement image") + prepared = prepare_image_lifecycle_job( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _get_selected_profile(tmp_path / "workspace"), + lifecycle_id="existing-path-replacement", + ) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_enroot = fake_bin / "enroot" + fake_enroot.write_text( + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + 'if [[ "$1" == "version" ]]; then\n' + " printf '%s\\n' '4.0.0'\n" + 'elif [[ "$1" == "create" ]]; then\n' + ' cp "${@: -1}" "${DD_TEST_INSPECTED_SQSH}"\n' + ' mv "${DD_TEST_REPLACEMENT}" "${DD_TEST_SOURCE}"\n' + 'elif [[ "$1" == "start" ]]; then\n' + " printf '%s\\n' '{}' > \"${DD_TEST_INSPECTION_OUTPUT}\"\n" + "fi\n" + ) + fake_enroot.chmod(0o700) + script = render_image_lifecycle_script(prepared.plan).replace( + 'export PATH="', + f'export PATH="{fake_bin.as_posix()}:', + 1, + ) + + completed = subprocess.run( + ("bash",), + input=script, + capture_output=True, + text=True, + check=False, + env={ + "DD_TEST_INSPECTION_OUTPUT": prepared.plan.inspection_output_path, + "DD_TEST_INSPECTED_SQSH": (tmp_path / "inspected.sqsh").as_posix(), + "DD_TEST_REPLACEMENT": replacement.as_posix(), + "DD_TEST_SOURCE": image_path.as_posix(), + "SLURM_CPUS_PER_TASK": "2", + "SLURM_JOB_ID": "5101", + }, + ) + + assert completed.returncode == 74 + assert "changed during inspection" in completed.stderr + assert (tmp_path / "inspected.sqsh").read_bytes() == b"client image" + assert image_path.read_bytes() == b"replacement image" + assert not Path(prepared.plan.inspection_output_path).exists() + assert not (Path(prepared.plan.job_directory) / "verified.sqsh").exists() + + +def test_standalone_client_inspector_binds_pip_to_its_distribution( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + executable_path = tmp_path / "active-environment" / "bin" / "pip" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + shadow_path = tmp_path / "shadow" / "bin" / "pip" + shadow_path.parent.mkdir(parents=True) + shadow_path.write_text("#!/bin/sh\n") + shadow_path.chmod(0o700) + monkeypatch.setenv("PATH", shadow_path.parent.as_posix()) + installer_distribution = SimpleNamespace( + metadata={"Name": "pip"}, + version="26.1", + entry_points=(SimpleNamespace(group="console_scripts", name="pip"),), + files=(Path("../../../bin/pip"),), + locate_file=lambda _installed_file: executable_path, + ) + distributions = ( + *( + SimpleNamespace(metadata={"Name": name}, version="1.0.0") + for name in ("data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm") + ), + installer_distribution, + ) + with patch.object(resource_inspector.importlib.metadata, "distributions", return_value=distributions): + record = resource_inspector.inspect_image("client", "c" * 64) + output_path = tmp_path / "inspection.json" + + resource_inspector.write_inspection(output_path, record) + + inspection = ImageInspectionRecord.model_validate_json(output_path.read_text()) + assert inspection.inspection.kind == "client" + assert inspection.sqsh_sha256 == "c" * 64 + assert inspection.inspection.installer_path == executable_path.as_posix() # type: ignore[union-attr] + assert inspection.inspection.installer_version == "26.1" # type: ignore[union-attr] + assert stat.S_IMODE(output_path.stat().st_mode) == 0o600 + + +def test_standalone_serving_inspector_binds_version_and_executable_to_one_distribution(tmp_path: Path) -> None: + executable_path = tmp_path / "active-environment" / "bin" / "vllm" + executable_path.parent.mkdir(parents=True) + executable_path.write_text("#!/bin/sh\n") + executable_path.chmod(0o700) + installed_file = Path("../../../bin/vllm") + entry_point = SimpleNamespace(group="console_scripts", name="vllm") + distribution = SimpleNamespace( + metadata={"Name": "vllm"}, + version="0.21.0", + entry_points=(entry_point,), + files=(installed_file,), + locate_file=lambda _installed_file: executable_path, + ) + + with ( + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=(distribution,)), + ): + payload = resource_inspector.inspect_image("serving", "d" * 64) + + inspection = ImageInspectionRecord.model_validate_json(json.dumps(payload)) + assert inspection.inspection.kind == "serving" + assert inspection.inspection.runtime_version == "0.21.0" # type: ignore[union-attr] + assert inspection.inspection.executable_path == executable_path.as_posix() # type: ignore[union-attr] + + +@pytest.mark.parametrize( + ("entry_points", "files", "match"), + ( + ((), (), "does not expose one console script"), + ((SimpleNamespace(group="console_scripts", name="vllm"),), None, "installed-file inventory"), + ((SimpleNamespace(group="console_scripts", name="vllm"),), (Path("vllm"),), "does not own one executable"), + ), +) +def test_standalone_serving_inspector_rejects_unverifiable_distribution_console_script( + entry_points: tuple[SimpleNamespace, ...], + files: tuple[Path, ...] | None, + match: str, +) -> None: + distribution = SimpleNamespace( + metadata={"Name": "vllm"}, + version="0.21.0", + entry_points=entry_points, + files=files, + locate_file=lambda installed_file: installed_file, + ) + + with ( + patch.object(resource_inspector.importlib.metadata, "distributions", return_value=(distribution,)), + pytest.raises(RuntimeError, match=match), + ): + resource_inspector.inspect_image("serving", "d" * 64) + + +@pytest.mark.parametrize( + ("kind", "sqsh_sha256", "match"), + (("unknown", "f" * 64, "image kind"), ("client", "not-a-digest", "SHA-256")), +) +def test_standalone_inspector_rejects_invalid_invocation(kind: str, sqsh_sha256: str, match: str) -> None: + with pytest.raises(ValueError, match=match): + resource_inspector.inspect_image(kind, sqsh_sha256) + + +def _get_selected_profile(workspace: Path) -> SelectedSlurmProfile: + profile = SlurmProfile( + schema_version=1, + scheduler=SchedulerProfile(account="research", partition="gpu"), + gpus_per_node=8, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile( + partition="image-build", + cpus_per_task=2, + memory="8G", + time_limit="03:55:00", + ), + ) + return injected_profile(profile) + + +def _prepare_completed_oci_lifecycle( + workspace: Path, + *, + lifecycle_id: str, + content: bytes, +) -> PreparedImageLifecycleJob: + prepared = prepare_image_lifecycle_job( + ImageBuildRequest( + name="client", + kind="client", + source=f"registry.example.test/client@sha256:{'a' * 64}", + ), + _get_selected_profile(workspace), + lifecycle_id=lifecycle_id, + ) + Path(prepared.plan.sqsh_path).write_bytes(content) + return prepared + + +def _write_completed_inspection( + prepared: PreparedImageLifecycleJob, + inspection: ImageInspectionRecord, +) -> None: + Path(prepared.plan.inspection_output_path).write_text(inspection.model_dump_json()) + + +def _get_client_inspection(content: bytes) -> ImageInspectionRecord: + distributions = ( + (InstalledDistribution(name="aiohttp", version="3.14.3"),) + + tuple( + InstalledDistribution(name=name, version="0.9.2") + for name in ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + ) + ) + + (InstalledDistribution(name="pip", version="26.1"),) + ) + return ImageInspectionRecord( + schema_version=1, + inspector_version="inspector-1", + sqsh_sha256=hashlib.sha256(content).hexdigest(), + inspection=ClientImageInspection( + kind="client", + python_implementation="cpython", + python_version="3.13.3", + python_abi="cp313", + distributions=distributions, + installer_path="/usr/bin/pip", + installer_version="26.1", + ), + ) + + +def _get_serving_inspection(content: bytes, runtime_version: str) -> ImageInspectionRecord: + return ImageInspectionRecord( + schema_version=1, + inspector_version="inspector-1", + sqsh_sha256=hashlib.sha256(content).hexdigest(), + inspection=ServingImageInspection( + kind="serving", + server_type="vllm", + runtime_version=runtime_version, + executable_path="/usr/local/bin/vllm", + ), + ) + + +class _MutatingSubmissionRunner: + def __init__(self, script_path: Path) -> None: + self._script_path = script_path + self.command: tuple[str, ...] | None = None + self.input_text: str | None = None + self.environment: dict[str, str] = {} + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + self.command = tuple(command) + self.input_text = input_text + self.environment = dict(environment or {}) + self._script_path.chmod(0o700) + self._script_path.write_text("replaced after verification\n") + return subprocess.CompletedProcess(command, 0, stdout="5101\n", stderr="") diff --git a/packages/data-designer-slurm/tests/images/test_registry.py b/packages/data-designer-slurm/tests/images/test_registry.py new file mode 100644 index 000000000..93c794985 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_registry.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import secrets +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml +from slurm_test_fakes import FakeInspectionEnvironment + +from data_designer.slurm.config import ( + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, +) +from data_designer.slurm.images.errors import ( + ImageConflictError, + ImageNotFoundError, + ImageVerificationError, +) +from data_designer.slurm.images.inspection import ClientImageInspector, ServingImageInspector +from data_designer.slurm.images.records import ImageRegistryDocument +from data_designer.slurm.images.service import ( + VerifiedImageRegistry, + compute_sqsh_file_sha256, +) + + +def _create_registry(workspace: Path) -> VerifiedImageRegistry: + return VerifiedImageRegistry(workspace) + + +def _get_registry_path(workspace: Path) -> Path: + return workspace / "images" / "registry.yaml" + + +def test_register_and_resolve_existing_sqsh_by_alias_and_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = _write_sqsh(tmp_path / "artifacts" / "client.sqsh", b"client-image") + inspection = _inspect_client(image_path) + service = _create_registry(workspace) + + registered = service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + inspection, + ) + + by_alias = service.resolve_for_planning(ImageRef(name="client"), expected_kind=ImageKind.CLIENT) + by_path = service.resolve_for_planning(ImageRef(path=image_path.as_posix()), expected_kind=ImageKind.CLIENT) + persisted = ImageRegistryDocument.model_validate_json( + json.dumps(yaml.safe_load(_get_registry_path(workspace).read_text())) + ) + assert by_alias.path == by_path.path + assert by_alias.sha256 == by_path.sha256 + assert by_alias.inspection == by_path.inspection + assert by_alias.path == image_path.as_posix() + assert registered == persisted.images[0] + assert _get_registry_path(workspace) == workspace / "images" / "registry.yaml" + + +@pytest.mark.parametrize("aiohttp_version", (None, "3.14.2", "4.0.0", "invalid")) +def test_client_image_without_compatible_proxy_dependency_cannot_be_planned( + tmp_path: Path, aiohttp_version: str | None +) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client-image") + service = _create_registry(tmp_path / "workspace") + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_client(image_path, aiohttp_version=aiohttp_version), + ) + + with pytest.raises(ImageVerificationError, match="client image must install aiohttp"): + service.resolve_for_planning(ImageRef(name="client"), expected_kind=ImageKind.CLIENT) + + +def test_registry_lists_aliases_deterministically(tmp_path: Path) -> None: + service = _create_registry(tmp_path / "workspace") + for name in ("zeta", "alpha"): + image_path = _write_sqsh(tmp_path / f"{name}.sqsh", name.encode()) + service.register_existing( + ImageBuildRequest(name=name, kind="serving", source=image_path.as_posix()), + _inspect_serving(image_path), + ) + + assert tuple(image.name for image in service.list_images()) == ("alpha", "zeta") + + +def test_unregister_removes_only_alias(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client-image") + service = _create_registry(tmp_path / "workspace") + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_client(image_path), + ) + + removed = service.unregister("client") + + assert removed.path == image_path.as_posix() + assert image_path.read_bytes() == b"client-image" + with pytest.raises(ImageNotFoundError, match="not registered"): + service.resolve_for_planning(ImageRef(name="client"), expected_kind=ImageKind.CLIENT) + + +def test_registration_requires_explicit_replace_for_alias_updates(tmp_path: Path) -> None: + first_path = _write_sqsh(tmp_path / "first.sqsh", b"first") + second_path = _write_sqsh(tmp_path / "second.sqsh", b"second") + service = _create_registry(tmp_path / "workspace") + request = ImageBuildRequest(name="serving", kind="serving", source=first_path.as_posix()) + service.register_existing(request, _inspect_serving(first_path)) + + replacement = ImageBuildRequest(name="serving", kind="serving", source=second_path.as_posix()) + with pytest.raises(ImageConflictError, match="already registered"): + service.register_existing(replacement, _inspect_serving(second_path)) + + service.register_existing(replacement, _inspect_serving(second_path), replace=True) + + assert ( + service.resolve_for_planning(ImageRef(name="serving"), expected_kind=ImageKind.SERVING).path + == second_path.as_posix() + ) + assert first_path.exists() + + +def test_failed_replacement_keeps_previous_alias(tmp_path: Path) -> None: + first_path = _write_sqsh(tmp_path / "first.sqsh", b"first") + second_path = _write_sqsh(tmp_path / "second.sqsh", b"second") + service = _create_registry(tmp_path / "workspace") + service.register_existing( + ImageBuildRequest(name="serving", kind="serving", source=first_path.as_posix()), + _inspect_serving(first_path), + ) + second_inspection = _inspect_serving(second_path) + original_token_hex = secrets.token_hex + + def mutate_replacement_before_registry_write(nbytes: int | None = None) -> str: + second_path.write_bytes(b"modified") + return original_token_hex(nbytes) + + with ( + patch( + "data_designer.slurm.filesystem.secrets.token_hex", + side_effect=mutate_replacement_before_registry_write, + ), + pytest.raises(ImageVerificationError, match="no longer matches"), + ): + service.register_existing( + ImageBuildRequest(name="serving", kind="serving", source=second_path.as_posix()), + second_inspection, + replace=True, + ) + + reloaded = _create_registry(tmp_path / "workspace") + assert ( + reloaded.resolve_for_planning(ImageRef(name="serving"), expected_kind=ImageKind.SERVING).path + == first_path.as_posix() + ) + + +def test_registration_does_not_publish_alias_when_sqsh_changes_before_atomic_publish(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + inspection = _inspect_client(image_path) + service = _create_registry(workspace) + original_token_hex = secrets.token_hex + + def mutate_image_before_registry_write(nbytes: int | None = None) -> str: + image_path.write_bytes(b"modified") + return original_token_hex(nbytes) + + with ( + patch( + "data_designer.slurm.filesystem.secrets.token_hex", + side_effect=mutate_image_before_registry_write, + ), + pytest.raises(ImageVerificationError, match="no longer matches"), + ): + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + inspection, + ) + + assert _create_registry(workspace).list_images() == () + + +def test_hash_rejects_path_replaced_while_reading(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + replacement = _write_sqsh(tmp_path / "replacement.sqsh", b"replacement") + original_lstat = Path.lstat + lstat_calls = 0 + + def replace_before_final_path_check(path: Path) -> os.stat_result: + nonlocal lstat_calls + lstat_calls += 1 + if lstat_calls == 2: + replacement.replace(image_path) + return original_lstat(path) + + with ( + patch.object(Path, "lstat", replace_before_final_path_check), + pytest.raises(ImageVerificationError, match="changed while it was being verified"), + ): + compute_sqsh_file_sha256(image_path) + + +def test_hash_uses_exact_sqsh_bytes(tmp_path: Path) -> None: + content = b"exact SQSH bytes\x00\xff" + image_path = _write_sqsh(tmp_path / "client.sqsh", content) + + assert compute_sqsh_file_sha256(image_path) == hashlib.sha256(content).hexdigest() + + +@pytest.mark.parametrize("mutation", (b"modified", b""), ids=("modified", "truncated")) +def test_resolution_rejects_modified_registered_sqsh(tmp_path: Path, mutation: bytes) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"original") + service = _create_registry(tmp_path / "workspace") + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_client(image_path), + ) + image_path.write_bytes(mutation) + + with pytest.raises(ImageVerificationError, match="no longer matches"): + service.resolve_for_planning(ImageRef(name="client"), expected_kind=ImageKind.CLIENT) + + +def test_resolution_rejects_kind_mismatch(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + service = _create_registry(tmp_path / "workspace") + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_client(image_path), + ) + + with pytest.raises(ImageVerificationError, match="does not match"): + service.resolve_for_planning(ImageRef(name="client"), expected_kind=ImageKind.SERVING) + + +def test_registration_rejects_missing_symlink_and_wrong_inspection(tmp_path: Path) -> None: + target = _write_sqsh(tmp_path / "target.sqsh", b"target") + symlink = tmp_path / "link.sqsh" + symlink.symlink_to(target) + service = _create_registry(tmp_path / "workspace") + + with pytest.raises(ImageVerificationError, match="does not exist"): + service.register_existing( + ImageBuildRequest(name="missing", kind="client", source=(tmp_path / "missing.sqsh").as_posix()), + _inspect_client(target), + ) + with pytest.raises(ImageVerificationError, match="regular non-symlink"): + service.register_existing( + ImageBuildRequest(name="symlink", kind="client", source=symlink.as_posix()), + _inspect_client(target), + ) + with pytest.raises(ImageVerificationError, match="digest"): + service.register_existing( + ImageBuildRequest(name="mismatch", kind="client", source=target.as_posix()), + ClientImageInspector(_get_client_environment()).inspect("f" * 64), + ) + assert service.list_images() == () + + +def test_registration_rejects_kind_mismatched_inspection(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "image.sqsh", b"image") + service = _create_registry(tmp_path / "workspace") + + with pytest.raises(ImageVerificationError, match="kind"): + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_serving(image_path), + ) + assert service.list_images() == () + + +def test_registration_rejects_unsupported_inspector_version(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + inspection = _inspect_client(image_path).model_copy(update={"inspector_version": "inspector-2"}) + service = _create_registry(tmp_path / "workspace") + + with pytest.raises(ImageVerificationError, match="unsupported image inspector version"): + service.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + inspection, + ) + + assert service.list_images() == () + + +def test_existing_registration_rejects_oci_source(tmp_path: Path) -> None: + service = _create_registry(tmp_path / "workspace") + request = ImageBuildRequest( + name="serving", + kind="serving", + source=f"registry.example.test/serving@sha256:{'a' * 64}", + ) + + with pytest.raises(ImageVerificationError, match="CPU Slurm image lifecycle"): + service.register_existing(request, ServingImageInspector(_get_serving_environment()).inspect("a" * 64)) + + +def test_direct_path_must_be_registered(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "unregistered.sqsh", b"unregistered") + service = _create_registry(tmp_path / "workspace") + + with pytest.raises(ImageNotFoundError, match="not registered"): + service.resolve_for_planning(ImageRef(path=image_path.as_posix()), expected_kind=ImageKind.CLIENT) + + +def _write_sqsh(path: Path, content: bytes) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + return path + + +def _get_client_environment(aiohttp_version: str | None = "3.14.3") -> FakeInspectionEnvironment: + return FakeInspectionEnvironment( + distributions=( + *((InstalledDistribution(name="aiohttp", version=aiohttp_version),) if aiohttp_version is not None else ()), + InstalledDistribution(name="data-designer", version="0.9.2"), + InstalledDistribution(name="data-designer-config", version="0.9.2"), + InstalledDistribution(name="data-designer-engine", version="0.9.2"), + InstalledDistribution(name="data-designer-slurm", version="0.9.2"), + InstalledDistribution(name="pip", version="26.1"), + ), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + + +def _inspect_client(path: Path, *, aiohttp_version: str | None = "3.14.3") -> ImageInspectionRecord: + return ClientImageInspector(_get_client_environment(aiohttp_version)).inspect(compute_sqsh_file_sha256(path)) + + +def _inspect_serving(path: Path) -> ImageInspectionRecord: + return ServingImageInspector(_get_serving_environment()).inspect(compute_sqsh_file_sha256(path)) + + +def _get_serving_environment() -> FakeInspectionEnvironment: + return FakeInspectionEnvironment( + distribution_versions={"vllm": "0.21.0"}, + executables={"vllm": "/usr/local/bin/vllm"}, + ) diff --git a/packages/data-designer-slurm/tests/images/test_registry_store.py b/packages/data-designer-slurm/tests/images/test_registry_store.py new file mode 100644 index 000000000..2568dd656 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -0,0 +1,912 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import stat +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier, Event +from unittest.mock import patch + +import pytest +import yaml +from slurm_test_fakes import FakeInspectionEnvironment + +import data_designer.slurm.images.service as image_service +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, InstalledDistribution +from data_designer.slurm.images import filesystem as image_filesystem +from data_designer.slurm.images.errors import ( + ImageArtifactRemovalError, + ImageConflictError, + ImageRegistryError, + ImageVerificationError, +) +from data_designer.slurm.images.inspection import ClientImageInspector, ServingImageInspector +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.images.service import VerifiedImageRegistry, compute_sqsh_file_sha256 + + +@pytest.mark.parametrize( + "content", + ( + b"not-yaml: [\n", + b"\xff", + b"schema_version: 1\nimages: &recursive\n - *recursive\n", + b"images: []\n", + b"schema_version: 2\nimages: []\n", + b"schema_version: 1\nschema_version: 1\nimages: []\n", + ), + ids=( + "invalid-yaml", + "invalid-utf8", + "recursive-alias", + "missing-version", + "unsupported-version", + "duplicate-key", + ), +) +def test_registry_normalizes_corrupt_persisted_state(tmp_path: Path, content: bytes) -> None: + workspace = tmp_path / "workspace" + _write_registry_content(workspace, content) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_persisted_non_sqsh_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + payload = image.model_dump(mode="json") + payload["path"] = (tmp_path / "client.bin").as_posix() + _write_registry_images(workspace, (payload,)) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_persisted_inspection_digest_mismatch(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + payload = image.model_dump(mode="json") + payload["sqsh_sha256"] = "f" * 64 + _write_registry_images(workspace, (payload,)) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_persisted_unsorted_aliases(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + zeta = _get_registered_client_image(_write_sqsh(tmp_path / "zeta.sqsh", b"zeta"), name="zeta") + alpha = _get_registered_client_image(_write_sqsh(tmp_path / "alpha.sqsh", b"alpha"), name="alpha") + _write_registry_images(workspace, (zeta.model_dump(mode="json"), alpha.model_dump(mode="json"))) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_persisted_duplicate_aliases(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + payload = image.model_dump(mode="json") + _write_registry_images(workspace, (payload, payload)) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_persisted_conflicting_facts_for_one_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = _write_sqsh(tmp_path / "shared.sqsh", b"shared") + client = _get_registered_client_image(image_path, name="client") + serving_inspection = _inspect_serving(image_path) + serving = RegisteredImage( + schema_version=1, + name="serving", + path=image_path.as_posix(), + sqsh_sha256=serving_inspection.sqsh_sha256, + inspection=serving_inspection, + ) + _write_registry_images( + workspace, + (client.model_dump(mode="json"), serving.model_dump(mode="json")), + ) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_allows_aliases_with_identical_facts_for_one_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "shared.sqsh", b"shared"), name="alpha") + alias = image.model_copy(update={"name": "beta"}) + store = ImageRegistryStore(workspace) + + store.register(image, verify_before_publish=lambda _image: None) + store.register(alias, verify_before_publish=lambda _image: None) + + assert store.list_images() == (image, alias) + + +def test_unregister_can_delete_unshared_imported_artifact(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + + removed = store.unregister("client", delete_artifact=True) + + assert removed == image + assert store.list_images() == () + assert not artifact_path.exists() + + +@pytest.mark.parametrize("run_name", ("run-example", "experiment-1")) +def test_unregister_keeps_artifact_referenced_by_persisted_run(tmp_path: Path, run_name: str) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + run_root = workspace / "runs" / run_name + run_root.mkdir(parents=True) + (run_root / "resolved-plan.json").write_text( + json.dumps({"deployments": [{"image": {"path": artifact_path.as_posix()}}]}) + ) + + with pytest.raises(ImageConflictError, match="referenced by run"): + store.unregister("client", delete_artifact=True) + + assert store.get_by_name("client") == image + assert artifact_path.exists() + + +@pytest.mark.parametrize("image_path", (None, "", "relative/image.sqsh", 3)) +def test_unregister_rejects_corrupt_run_image_path(tmp_path: Path, image_path: object) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + run_root = workspace / "runs" / "run-corrupt" + run_root.mkdir(parents=True) + (run_root / "resolved-plan.json").write_text(json.dumps({"deployments": [{"image": {"path": image_path}}]})) + + with pytest.raises(ImageConflictError, match="run plan cannot be checked"): + store.unregister("client", delete_artifact=True) + + assert store.get_by_name("client") == image + assert artifact_path.exists() + + +@pytest.mark.parametrize("invalid_plan", ("symlink", "oversized")) +def test_unregister_keeps_artifact_when_run_plan_cannot_be_safely_read( + tmp_path: Path, + invalid_plan: str, +) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + run_root = workspace / "runs" / "run-example" + run_root.mkdir(parents=True) + plan_path = run_root / "resolved-plan.json" + if invalid_plan == "symlink": + source = tmp_path / "outside.json" + source.write_text(json.dumps({"deployments": []})) + plan_path.symlink_to(source) + else: + with plan_path.open("wb") as plan: + plan.truncate(16 * 1024 * 1024 + 1) + + with pytest.raises(ImageConflictError, match="run plan cannot be checked"): + store.unregister("client", delete_artifact=True) + + assert store.get_by_name("client") == image + assert artifact_path.exists() + + +def test_artifact_removal_waits_for_run_plan_publication(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + started = Event() + + def remove() -> None: + started.set() + store.unregister("client", delete_artifact=True) + + with ThreadPoolExecutor(max_workers=1) as executor: + with store.acquire_run_reference_lock(): + pending = executor.submit(remove) + assert started.wait(timeout=2) + run_root = workspace / "runs" / "run-concurrent" + run_root.mkdir(parents=True) + (run_root / "resolved-plan.json").write_text( + json.dumps({"deployments": [{"image": {"path": artifact_path.as_posix()}}]}) + ) + assert not pending.done() + with pytest.raises(ImageConflictError, match="referenced by run"): + pending.result(timeout=2) + + assert store.get_by_name("client") == image + assert artifact_path.exists() + + +def test_run_reference_lock_allows_concurrent_publishers(tmp_path: Path) -> None: + store = ImageRegistryStore(tmp_path / "workspace") + entered = Event() + + def publish() -> None: + with store.acquire_run_reference_lock(): + entered.set() + + with ThreadPoolExecutor(max_workers=1) as executor: + with store.acquire_run_reference_lock(): + pending = executor.submit(publish) + assert entered.wait(timeout=2) + pending.result(timeout=2) + + +def test_artifact_reference_scan_does_not_block_registry_reads(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + scanning = Event() + resume = Event() + scan = store._require_no_persisted_run_reference + + def delayed_scan(path: str) -> None: + scanning.set() + assert resume.wait(timeout=3) + scan(path) + + monkeypatch.setattr(store, "_require_no_persisted_run_reference", delayed_scan) + publishing = Event() + + def publish() -> None: + with store.acquire_run_reference_lock(): + publishing.set() + + with ThreadPoolExecutor(max_workers=3) as executor: + removal = executor.submit(store.unregister, "client", delete_artifact=True) + assert scanning.wait(timeout=2) + listing = executor.submit(store.list_images) + publisher = executor.submit(publish) + try: + assert listing.result(timeout=2) == (image,) + assert not publishing.wait(timeout=0.05) + finally: + resume.set() + assert removal.result(timeout=2) == image + publisher.result(timeout=2) + assert publishing.is_set() + + +def test_unregister_refuses_to_delete_shared_or_external_artifacts(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + imported = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(imported, verify_before_publish=lambda _image: None) + store.register(imported.model_copy(update={"name": "other"}), verify_before_publish=lambda _image: None) + + with pytest.raises(ImageConflictError, match="another image alias"): + store.unregister("client", delete_artifact=True) + + assert len(store.list_images()) == 2 + assert artifact_path.exists() + + external_path = _write_sqsh(tmp_path / "external.sqsh", b"external") + external = _get_registered_client_image(external_path, name="external") + store.register(external, verify_before_publish=lambda _image: None) + with pytest.raises(ImageConflictError, match="package-owned"): + store.unregister("external", delete_artifact=True) + assert store.get_by_name("external") == external + assert external_path.exists() + + +def test_unregister_reports_partial_artifact_removal_failure(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(b'client').hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(b"client") + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + real_unlink = os.unlink + + def fail_artifact_unlink(name: str, *, dir_fd: int | None = None) -> None: + if name == artifact_path.name: + raise OSError("injected artifact removal failure") + real_unlink(name, dir_fd=dir_fd) + + with ( + patch("data_designer.slurm.images.registry.os.unlink", side_effect=fail_artifact_unlink), + pytest.raises(ImageArtifactRemovalError, match="alias was removed"), + ): + store.unregister("client", delete_artifact=True) + + assert store.list_images() == () + assert artifact_path.exists() + + +def test_unregister_refuses_symlinked_artifact_without_removing_alias(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + content = b"client" + artifact_path = workspace / "images" / "artifacts" / f"client-{hashlib.sha256(content).hexdigest()}.sqsh" + artifact_path.parent.mkdir(parents=True) + artifact_path.write_bytes(content) + image = _get_registered_client_image(artifact_path).model_copy(update={"source_oci_digest": "a" * 64}) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + outside = _write_sqsh(tmp_path / "outside.sqsh", b"outside") + artifact_path.unlink() + artifact_path.symlink_to(outside) + + with pytest.raises(ImageConflictError, match="unavailable for deletion"): + store.unregister("client", delete_artifact=True) + + assert store.get_by_name("client") == image + assert outside.read_bytes() == b"outside" + + +def test_registry_rejects_new_alias_with_conflicting_facts_for_one_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_path = _write_sqsh(tmp_path / "shared.sqsh", b"shared") + client = _get_registered_client_image(image_path, name="client") + serving_inspection = _inspect_serving(image_path) + serving = RegisteredImage( + schema_version=1, + name="serving", + path=image_path.as_posix(), + sqsh_sha256=serving_inspection.sqsh_sha256, + inspection=serving_inspection, + ) + store = ImageRegistryStore(workspace) + store.register(client, verify_before_publish=lambda _image: None) + + with pytest.raises(ImageConflictError, match="different immutable facts"): + store.register(serving, verify_before_publish=lambda _image: None) + + +def test_registry_rejects_symlinked_state_file(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + registry_path = _get_registry_path(workspace) + registry_path.parent.mkdir(parents=True) + target = tmp_path / "outside.yaml" + target.write_text("schema_version: 1\nimages: []\n") + registry_path.symlink_to(target) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_reads_through_symlinked_image_root(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + outside_workspace = tmp_path / "outside-workspace" + outside_registry = VerifiedImageRegistry(outside_workspace) + outside_registry.register_existing( + ImageBuildRequest(name="client", kind="client", source=image_path.as_posix()), + _inspect_client(image_path), + ) + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "images").symlink_to(outside_workspace / "images", target_is_directory=True) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_nonregular_state_file_without_blocking(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + registry_path = _get_registry_path(workspace) + registry_path.parent.mkdir(parents=True) + os.mkfifo(registry_path) + + with pytest.raises(ImageRegistryError, match="cannot load"): + ImageRegistryStore(workspace).list_images() + + +def test_registry_rejects_fifo_substituted_between_check_and_open(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + registry_path = store.registry_path + original_stat = os.stat + replaced = False + + def replace_registry_after_stat(path: str, **kwargs: object) -> os.stat_result: + nonlocal replaced + status = original_stat(path, **kwargs) + if path == "registry.yaml" and not replaced: + replaced = True + registry_path.unlink() + os.mkfifo(registry_path) + return status + + with ( + patch("data_designer.slurm.images.filesystem.os.stat", side_effect=replace_registry_after_stat), + pytest.raises(ImageRegistryError, match="cannot load"), + ): + store.list_images() + + +def test_registry_rejects_state_mutated_during_read(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + store = ImageRegistryStore(workspace) + store.register(image, verify_before_publish=lambda _image: None) + registry_path = store.registry_path + registry_inode = registry_path.stat().st_ino + original_fstat = os.fstat + registry_fstat_count = 0 + + def mutate_before_final_registry_fstat(descriptor: int) -> os.stat_result: + nonlocal registry_fstat_count + status = original_fstat(descriptor) + if status.st_ino == registry_inode: + registry_fstat_count += 1 + if registry_fstat_count == 2: + registry_path.write_text("mutated") + status = original_fstat(descriptor) + return status + + with ( + patch("data_designer.slurm.images.filesystem.os.fstat", side_effect=mutate_before_final_registry_fstat), + pytest.raises(ImageRegistryError, match="cannot load"), + ): + store.list_images() + + +def test_sqsh_reader_rejects_fifo_substituted_between_check_and_open(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + original_lstat = Path.lstat + replaced = False + + def replace_sqsh_after_lstat(path: Path) -> os.stat_result: + nonlocal replaced + status = original_lstat(path) + if path == image_path and not replaced: + replaced = True + image_path.unlink() + os.mkfifo(image_path) + return status + + with ( + patch.object(Path, "lstat", replace_sqsh_after_lstat), + pytest.raises(ImageVerificationError, match="changed while it was being opened"), + ): + compute_sqsh_file_sha256(image_path) + + +def test_sqsh_reader_rejects_path_replaced_during_hash(tmp_path: Path) -> None: + image_path = _write_sqsh(tmp_path / "client.sqsh", b"client") + replacement = _write_sqsh(tmp_path / "replacement.sqsh", b"replacement") + original_hash_descriptor = image_service._hash_descriptor + + def replace_path_after_hash(descriptor: int) -> str: + digest = original_hash_descriptor(descriptor) + replacement.replace(image_path) + return digest + + with ( + patch("data_designer.slurm.images.service._hash_descriptor", side_effect=replace_path_after_hash), + pytest.raises(ImageVerificationError, match="changed while it was being verified"), + ): + compute_sqsh_file_sha256(image_path) + + +def test_registry_rejects_symlinked_lock_file_without_changing_target(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + lock_directory = workspace / "images" / ".locks" + lock_directory.mkdir(parents=True) + target = tmp_path / "outside.lock" + target.write_text("outside") + target.chmod(0o640) + (lock_directory / "alias-client.lock").symlink_to(target) + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + + with pytest.raises(ImageRegistryError, match="cannot lock"): + ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None) + + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + + +def test_registry_rejects_nonregular_lock_file(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + lock_directory = workspace / "images" / ".locks" + lock_directory.mkdir(parents=True) + os.mkfifo(lock_directory / "alias-client.lock") + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + + with pytest.raises(ImageRegistryError, match="cannot lock"): + ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None) + + +def test_registry_lock_preserves_operation_oserror(tmp_path: Path) -> None: + lock_directory = tmp_path / "locks" + lock_directory.mkdir(mode=0o700) + lock_path = lock_directory / "registry.lock" + + with image_filesystem.open_verified_directory(lock_directory) as directory_descriptor: + with pytest.raises(OSError, match="registry operation failed") as caught: + with image_filesystem.acquire_file_lock(directory_descriptor, lock_path.name, lock_path): + raise OSError("registry operation failed") + + assert type(caught.value) is OSError + + +@pytest.mark.parametrize("directory", ("images", "locks")) +def test_registry_rejects_symlinked_storage_directories(tmp_path: Path, directory: str) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + image_root = workspace / "images" + outside = tmp_path / "outside" + outside.mkdir() + if directory == "images": + image_root.symlink_to(outside, target_is_directory=True) + else: + image_root.mkdir() + (image_root / ".locks").symlink_to(outside, target_is_directory=True) + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + + with pytest.raises(ImageRegistryError, match="cannot initialize"): + ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None) + + assert tuple(outside.iterdir()) == () + + +def test_registry_uses_restrictive_atomic_workspace_storage(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_root = workspace / "images" + lock_directory = image_root / ".locks" + lock_directory.mkdir(parents=True) + image_root.chmod(0o775) + lock_directory.chmod(0o775) + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + + ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None) + + registry_path = _get_registry_path(workspace) + assert stat.S_IMODE(registry_path.stat().st_mode) == 0o600 + assert stat.S_IMODE(image_root.stat().st_mode) == 0o700 + assert stat.S_IMODE(lock_directory.stat().st_mode) == 0o700 + assert not tuple(registry_path.parent.glob(".registry.*.tmp")) + assert not (registry_path.parent / ".registry.rollback.yaml").exists() + assert not (registry_path.parent / ".registry.committed.yaml").exists() + + +def test_registration_blocks_readers_until_final_verification(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + store = ImageRegistryStore(workspace) + verification_started = Event() + finish_verification = Event() + + def verify_before_publish(_image: RegisteredImage) -> None: + verification_started.set() + assert finish_verification.wait(timeout=5) + + with ThreadPoolExecutor(max_workers=2) as executor: + writer = executor.submit( + store.register, + image, + verify_before_publish=verify_before_publish, + ) + assert verification_started.wait(timeout=5) + reader = executor.submit(ImageRegistryStore(workspace).list_images) + assert not reader.done() + finish_verification.set() + assert writer.result().name == "client" + assert tuple(registered.name for registered in reader.result()) == ("client",) + + assert tuple(registered.name for registered in ImageRegistryStore(workspace).list_images()) == ("client",) + + +def test_existing_sqsh_hashing_does_not_block_registry_readers(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first_path = _write_sqsh(tmp_path / "first.sqsh", b"first") + second_path = _write_sqsh(tmp_path / "second.sqsh", b"second") + registry = VerifiedImageRegistry(workspace) + first = registry.register_existing( + ImageBuildRequest(name="first", kind="client", source=first_path.as_posix()), + _inspect_client(first_path), + ) + second_inspection = _inspect_client(second_path) + hash_started = Event() + finish_hash = Event() + original_hash_descriptor = image_service._hash_descriptor + + def block_second_hash(descriptor: int) -> str: + descriptor_status = os.fstat(descriptor) + second_status = second_path.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == (second_status.st_dev, second_status.st_ino): + hash_started.set() + assert finish_hash.wait(timeout=5) + return original_hash_descriptor(descriptor) + + with ( + patch("data_designer.slurm.images.service._hash_descriptor", side_effect=block_second_hash), + ThreadPoolExecutor(max_workers=2) as executor, + ): + writer = executor.submit( + registry.register_existing, + ImageBuildRequest(name="second", kind="client", source=second_path.as_posix()), + second_inspection, + ) + assert hash_started.wait(timeout=5) + reader = executor.submit(VerifiedImageRegistry(workspace).list_images) + assert reader.result(timeout=5) == (first,) + finish_hash.set() + assert writer.result(timeout=5).name == "second" + + +def test_failed_registration_runs_rollback_before_releasing_target_lock(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = _get_registered_client_image(_write_sqsh(tmp_path / "shared.sqsh", b"shared"), name="first") + second = first.model_copy(update={"name": "second"}) + store = ImageRegistryStore(workspace) + rollback_started = Event() + finish_rollback = Event() + + def fail_after_publish(_image: RegisteredImage) -> None: + raise ImageVerificationError("injected post-publication failure") + + def rollback_after_failure(_image: RegisteredImage) -> None: + rollback_started.set() + assert finish_rollback.wait(timeout=5) + + with ThreadPoolExecutor(max_workers=2) as executor: + failed_writer = executor.submit( + store.register, + first, + verify_before_publish=lambda _image: None, + verify_after_publish=fail_after_publish, + rollback_after_failure=rollback_after_failure, + ) + assert rollback_started.wait(timeout=5) + succeeding_writer = executor.submit( + store.register, + second, + verify_before_publish=lambda _image: None, + ) + assert not succeeding_writer.done() + finish_rollback.set() + with pytest.raises(ImageVerificationError, match="injected post-publication failure"): + failed_writer.result() + assert succeeding_writer.result() == second + + assert ImageRegistryStore(workspace).list_images() == (second,) + + +def test_fresh_reader_recovers_previous_registry_after_interrupted_publication(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + original = _get_registered_client_image(_write_sqsh(tmp_path / "original.sqsh", b"original")) + replacement = _get_registered_client_image(_write_sqsh(tmp_path / "replacement.sqsh", b"replacement")) + store = ImageRegistryStore(workspace) + store.register(original, verify_before_publish=lambda _image: None) + registry_path = store.registry_path + previous_content = registry_path.read_bytes() + (registry_path.parent / ".registry.rollback.yaml").write_bytes(previous_content) + registry_path.write_text( + yaml.safe_dump( + {"schema_version": 1, "images": [replacement.model_dump(mode="json")]}, + sort_keys=True, + ) + ) + + assert _list_image_names_in_fresh_process(workspace) == (original.name,) + assert not (registry_path.parent / ".registry.rollback.yaml").exists() + assert ImageRegistryStore(workspace).list_images() == (original,) + + +def test_fresh_reader_keeps_committed_registry_after_interrupted_marker_cleanup(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + original = _get_registered_client_image(_write_sqsh(tmp_path / "original.sqsh", b"original")) + replacement = _get_registered_client_image(_write_sqsh(tmp_path / "replacement.sqsh", b"replacement")) + store = ImageRegistryStore(workspace) + store.register(original, verify_before_publish=lambda _image: None) + registry_path = store.registry_path + (registry_path.parent / ".registry.committed.yaml").write_bytes(registry_path.read_bytes()) + registry_path.write_text( + yaml.safe_dump( + {"schema_version": 1, "images": [replacement.model_dump(mode="json")]}, + sort_keys=True, + ) + ) + + assert _list_image_names_in_fresh_process(workspace) == (replacement.name,) + assert not (registry_path.parent / ".registry.committed.yaml").exists() + assert ImageRegistryStore(workspace).list_images() == (replacement,) + + +def test_committed_marker_sync_failure_reports_failure_and_keeps_recovery_state(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + store = ImageRegistryStore(workspace) + image_root_sync_count = 0 + original_fsync = os.fsync + + def fail_committed_marker_sync(descriptor: int) -> None: + nonlocal image_root_sync_count + descriptor_status = os.fstat(descriptor) + image_root_status = store.image_root.stat() + if (descriptor_status.st_dev, descriptor_status.st_ino) == ( + image_root_status.st_dev, + image_root_status.st_ino, + ): + image_root_sync_count += 1 + if image_root_sync_count == 3: + raise OSError("injected committed marker sync failure") + original_fsync(descriptor) + + with ( + patch("data_designer.slurm.images.registry.os.fsync", side_effect=fail_committed_marker_sync), + pytest.raises(ImageRegistryError, match="commit state requires recovery"), + ): + store.register(image, verify_before_publish=lambda _image: None) + + committed_marker = store.image_root / ".registry.committed.yaml" + assert committed_marker.is_file() + assert _list_image_names_in_fresh_process(workspace) == (image.name,) + assert not committed_marker.exists() + assert store.list_images() == (image,) + + +def test_registry_transaction_stays_bound_when_image_root_path_is_replaced(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + image_root = workspace / "images" + detached_image_root = workspace / "detached-images" + replacement_image_root = workspace / "replacement-images" + replacement_image_root.mkdir(parents=True) + image = _get_registered_client_image(_write_sqsh(tmp_path / "client.sqsh", b"client")) + + def replace_image_root(_image: RegisteredImage) -> None: + image_root.rename(detached_image_root) + replacement_image_root.rename(image_root) + + ImageRegistryStore(workspace).register(image, verify_before_publish=replace_image_root) + + assert not (image_root / "registry.yaml").exists() + relocated_registry = yaml.safe_load((detached_image_root / "registry.yaml").read_text()) + assert relocated_registry["images"][0]["name"] == "client" + + +def test_concurrent_alias_writers_preserve_both_registrations(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + barrier = Barrier(2) + + def register(name: str) -> str: + image = _get_registered_client_image(_write_sqsh(tmp_path / f"{name}.sqsh", name.encode()), name=name) + barrier.wait() + return ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None).name + + with ThreadPoolExecutor(max_workers=2) as executor: + names = tuple(executor.map(register, ("first", "second"))) + + assert names == ("first", "second") + assert tuple(image.name for image in ImageRegistryStore(workspace).list_images()) == ("first", "second") + + +def test_concurrent_writers_cannot_publish_conflicting_aliases(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + barrier = Barrier(2) + + def register(content: bytes) -> str: + image = _get_registered_client_image(_write_sqsh(tmp_path / f"{content.decode()}.sqsh", content), name="shared") + barrier.wait() + return ImageRegistryStore(workspace).register(image, verify_before_publish=lambda _image: None).path + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(register, content) for content in (b"first", b"second")) + + failures = tuple(future.exception() for future in futures if future.exception() is not None) + outcomes = tuple(future.result() for future in futures if future.exception() is None) + assert len(outcomes) == 1 + assert len(failures) == 1 + assert isinstance(failures[0], ImageConflictError) + assert ImageRegistryStore(workspace).list_images()[0].path == outcomes[0] + + +def _write_registry_content(workspace: Path, content: bytes) -> None: + registry_path = _get_registry_path(workspace) + registry_path.parent.mkdir(parents=True) + registry_path.write_bytes(content) + + +def _list_image_names_in_fresh_process(workspace: Path) -> tuple[str, ...]: + completed = subprocess.run( + ( + sys.executable, + "-c", + ( + "import json,sys; " + "from data_designer.slurm.images.registry import ImageRegistryStore; " + "print(json.dumps([image.name for image in ImageRegistryStore(sys.argv[1]).list_images()]))" + ), + workspace.as_posix(), + ), + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return tuple(json.loads(completed.stdout)) + + +def _write_registry_images(workspace: Path, images: tuple[dict[str, object], ...]) -> None: + content = yaml.safe_dump({"schema_version": 1, "images": list(images)}, sort_keys=True).encode() + _write_registry_content(workspace, content) + + +def _get_registry_path(workspace: Path) -> Path: + return workspace / "images" / "registry.yaml" + + +def _write_sqsh(path: Path, content: bytes) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + return path + + +def _get_registered_client_image(path: Path, *, name: str = "client") -> RegisteredImage: + inspection = _inspect_client(path) + return RegisteredImage( + schema_version=1, + name=name, + path=path.as_posix(), + sqsh_sha256=inspection.sqsh_sha256, + inspection=inspection, + ) + + +def _inspect_client(path: Path) -> ImageInspectionRecord: + environment = FakeInspectionEnvironment( + distributions=( + InstalledDistribution(name="data-designer", version="0.9.2"), + InstalledDistribution(name="data-designer-config", version="0.9.2"), + InstalledDistribution(name="data-designer-engine", version="0.9.2"), + InstalledDistribution(name="data-designer-slurm", version="0.9.2"), + InstalledDistribution(name="pip", version="26.1"), + ), + distribution_versions={"pip": "26.1"}, + executables={"pip": "/usr/bin/pip"}, + ) + return ClientImageInspector(environment).inspect(compute_sqsh_file_sha256(path)) + + +def _inspect_serving(path: Path) -> ImageInspectionRecord: + environment = FakeInspectionEnvironment( + distribution_versions={"vllm": "0.21.0"}, + executables={"vllm": "/usr/local/bin/vllm"}, + ) + return ServingImageInspector(environment).inspect(compute_sqsh_file_sha256(path)) diff --git a/packages/data-designer-slurm/tests/images/test_vllm_source.py b/packages/data-designer-slurm/tests/images/test_vllm_source.py new file mode 100644 index 000000000..4078add49 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_vllm_source.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from http.client import IncompleteRead +from urllib.error import HTTPError, URLError +from urllib.request import Request + +import pytest + +import data_designer.slurm.images.vllm_source as vllm_source + + +def test_resolves_official_version_to_immutable_digest(monkeypatch: pytest.MonkeyPatch) -> None: + requests: list[Request] = [] + + def request(resource: Request) -> tuple[bytes, dict[str, str]]: + requests.append(resource) + if len(requests) == 1: + return b'{"token":"private-token"}', {} + return b"", {"Docker-Content-Digest": f"sha256:{'a' * 64}"} + + monkeypatch.setattr(vllm_source, "_request", request) + + pinned = vllm_source.resolve_versioned_vllm_source("vllm/vllm-openai:v0.22.0") + + assert pinned == f"vllm/vllm-openai:v0.22.0@sha256:{'a' * 64}" + assert requests[0].full_url.startswith("https://auth.docker.io/token?") + assert requests[1].full_url == "https://registry-1.docker.io/v2/vllm/vllm-openai/manifests/v0.22.0" + assert requests[1].get_method() == "HEAD" + assert requests[1].get_header("Authorization") == "Bearer private-token" + assert vllm_source.pinned_vllm_version(pinned) == "0.22.0" + + +@pytest.mark.parametrize( + "source", + ( + "vllm/vllm-openai:latest", + "other/vllm-openai:v0.22.0", + "vllm/vllm-openai:v0.22.0-rc1", + "https://vllm/vllm-openai:v0.22.0", + ), +) +def test_rejects_non_release_sources(source: str) -> None: + assert not vllm_source.is_versioned_vllm_tag(source) + with pytest.raises(ValueError, match="official versioned vLLM"): + vllm_source.resolve_versioned_vllm_source(source) + + +@pytest.mark.parametrize( + "failure", + ( + URLError("private-token"), + HTTPError("url", 500, "private-token", {}, None), + IncompleteRead(b"private-token", 100), + ), +) +def test_resolution_failure_is_sanitized(monkeypatch: pytest.MonkeyPatch, failure: Exception) -> None: + def request(_resource: Request) -> tuple[bytes, dict[str, str]]: + raise failure + + monkeypatch.setattr(vllm_source, "_request", request) + with pytest.raises(vllm_source.VllmSourceResolutionError) as caught: + vllm_source.resolve_versioned_vllm_source("vllm/vllm-openai:v0.22.0") + assert "private-token" not in str(caught.value) + + +def test_resolution_rejects_missing_digest(monkeypatch: pytest.MonkeyPatch) -> None: + responses = iter(((b'{"token":"token"}', {}), (b"", {"Docker-Content-Digest": "not-a-digest"}))) + monkeypatch.setattr(vllm_source, "_request", lambda _resource: next(responses)) + + with pytest.raises(vllm_source.VllmSourceResolutionError, match="manifest digest"): + vllm_source.resolve_versioned_vllm_source("vllm/vllm-openai:v0.22.0") + + +def test_missing_version_has_distinct_sanitized_error(monkeypatch: pytest.MonkeyPatch) -> None: + responses = iter(((b'{"token":"token"}', {}), HTTPError("url", 404, "private-token", {}, None))) + + def request(_resource: Request) -> tuple[bytes, dict[str, str]]: + result = next(responses) + if isinstance(result, HTTPError): + raise result + return result + + monkeypatch.setattr(vllm_source, "_request", request) + with pytest.raises(vllm_source.VllmTagNotFoundError, match="not found") as caught: + vllm_source.resolve_versioned_vllm_source("vllm/vllm-openai:v0.22.0") + assert "private-token" not in str(caught.value) + + +def test_token_endpoint_not_found_is_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + def request(_resource: Request) -> tuple[bytes, dict[str, str]]: + raise HTTPError("url", 404, "private-token", {}, None) + + monkeypatch.setattr(vllm_source, "_request", request) + with pytest.raises(vllm_source.VllmSourceResolutionError, match="could not be resolved") as caught: + vllm_source.resolve_versioned_vllm_source("vllm/vllm-openai:v0.22.0") + assert not isinstance(caught.value, vllm_source.VllmTagNotFoundError) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json new file mode 100644 index 000000000..eeb5661b1 --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -0,0 +1,128 @@ +{ + "attempt": { + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_output": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" + }, + "created_at": "2026-08-19T12:00:02Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" + }, + "run_id": "run-single", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "succeeded", + "terminal_classification": "succeeded", + "updated_at": "2026-08-19T12:05:02Z" + }, + "candidate": { + "actual_records": 8, + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "created_at": "2026-08-19T12:05:00Z", + "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", + "dataset_schema_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "files": [ + { + "byte_size": 4096, + "record_count": 8, + "relative_path": "part-00000.parquet", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + } + ], + "outcome": "complete", + "provenance_digest": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75", + "requested_records": 8, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + }, + "client_result": { + "actual_records": 8, + "attempt_id": "attempt-0001", + "candidate_output_manifest": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" + }, + "completed_at": "2026-08-19T12:05:01Z", + "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", + "early_shutdown": false, + "effective_resume_mode": "never", + "outcome": "complete", + "requested_records": 8, + "requested_resume_mode": "never", + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + }, + "readiness": { + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": null, + "model_alias": "generator", + "ready_backends": 0, + "state": "pending" + } + ], + "revision": 1, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "pending", + "updated_at": "2026-08-19T12:00:03Z" + }, + "run": { + "authored_config": { + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "bf9d45d571f5f288c74a3332cdd23576df60ce7bb3d934774f50bdce4b984fba" + }, + "created_at": "2026-08-19T12:00:00Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" + }, + "run_id": "run-single", + "schema_version": 1, + "shard_count": 1 + }, + "shards": [ + { + "created_at": "2026-08-19T12:00:01Z", + "input_partition": null, + "record_range": { + "end_index_exclusive": 8, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/dataset" + }, + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000", + "shard_index": 0 + } + ], + "winner": { + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_manifest": { + "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" + }, + "published_at": "2026-08-19T12:05:03Z", + "run_id": "run-single", + "schema_version": 1, + "shard_id": "shard-00000" + } +} diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py new file mode 100644 index 000000000..dcc4be0ce --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -0,0 +1,741 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import posixpath +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import TypeVar + +import pytest + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference, ContractValue, RecordRange, ResumeWorkspace +from data_designer.slurm.integration import IntegrationContractError, PlanStateValidator +from data_designer.slurm.planning import ( + ArtifactReference as PlanningArtifactReference, +) +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + ArtifactReference as StateArtifactReference, +) +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputManifest, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerIdentity, + ShardManifest, + ShardWinner, +) + +TEST_ROOT = Path(__file__).parents[1] +CONTRACT_GOLDEN_DIR = TEST_ROOT / "contracts" / "golden" +INTEGRATION_GOLDEN_DIR = Path(__file__).parent / "golden" +CREATED_AT = datetime(2026, 8, 19, 12, 0, tzinfo=timezone.utc) +_RecordT = TypeVar("_RecordT", bound=ContractValue) + + +@dataclass(frozen=True) +class IntegrationRecords: + plan: ResolvedSlurmRunPlan + validator: PlanStateValidator + run: RunManifest + shards: tuple[ShardManifest, ...] + attempt: AttemptManifest + readiness: AttemptReadiness + client_result: ClientResult + candidate: CandidateOutputManifest + winner: ShardWinner + + +@pytest.fixture +def records() -> IntegrationRecords: + plan = _load_plan("single_node_plan.json") + payload = json.loads((INTEGRATION_GOLDEN_DIR / "finalization_chain.json").read_text()) + return IntegrationRecords( + plan=plan, + validator=PlanStateValidator(plan), + run=_load_record(RunManifest, payload["run"]), + shards=tuple(_load_record(ShardManifest, shard) for shard in payload["shards"]), + attempt=_load_record(AttemptManifest, payload["attempt"]), + readiness=_load_record(AttemptReadiness, payload["readiness"]), + client_result=_load_record(ClientResult, payload["client_result"]), + candidate=_load_record(CandidateOutputManifest, payload["candidate"]), + winner=_load_record(ShardWinner, payload["winner"]), + ) + + +def test_shared_contract_types_retain_exact_identity() -> None: + assert PlanningArtifactReference is ArtifactReference + assert StateArtifactReference is ArtifactReference + + +def test_golden_records_validate_every_plan_state_join(records: IntegrationRecords) -> None: + planned_shard = records.plan.shards[0] + + assert records.validator.validate_plan_shards(records.run, records.shards) is records.shards + assert records.validator.validate_initial_readiness(records.attempt, records.readiness) is records.readiness + assert records.validator.validate_planned_attempt(planned_shard, records.attempt) is records.attempt + assert ( + records.validator.validate_finalization_chain( + planned_shard, + records.attempt, + records.client_result, + records.candidate, + records.winner, + ) + is records.winner + ) + + +def test_plan_shards_reject_missing_and_extra_state(records: IntegrationRecords) -> None: + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + records.validator.validate_plan_shards(records.run, ()) + with pytest.raises(IntegrationContractError, match="exactly the run shard count"): + records.validator.validate_plan_shards(records.run, records.shards + records.shards) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("run_id", "run manifest identity"), + ("authored_config", "authored config"), + ("plan_path", "reference path"), + ("plan_digest", "reference digest"), + ("shard_count", "run shard count"), + ("shard_id", "shard identity"), + ("shard_index", "outside the run shard count"), + ("record_range", "record range"), + ("input_partition", "input partition"), + ("resume_workspace", "resume workspace"), + ("created_at", "cannot precede"), + ], +) +def test_plan_shards_reject_one_field_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + run = records.run + shard = records.shards[0] + if mutation == "run_id": + run = run.model_copy(update={"run_id": "run-other"}) + elif mutation == "authored_config": + run = run.model_copy( + update={ + "authored_config": ArtifactReference( + path=run.authored_config.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "plan_path": + run = run.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=posixpath.join(posixpath.dirname(run.resolved_plan.path), "other-plan.json"), + sha256=run.resolved_plan.sha256, + ) + } + ) + elif mutation == "plan_digest": + run = run.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=run.resolved_plan.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "shard_count": + run = run.model_copy(update={"shard_count": 2}) + elif mutation == "shard_id": + shard = shard.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "shard_index": + shard = shard.model_copy(update={"shard_index": 1}) + elif mutation == "record_range": + shard = shard.model_copy(update={"record_range": RecordRange(start_index=1, end_index_exclusive=8)}) + elif mutation == "input_partition": + shard = shard.model_copy( + update={ + "input_partition": ArtifactReference( + path=posixpath.join(posixpath.dirname(shard.resume_workspace.path), "input-partition.json"), + sha256="a" * 64, + ) + } + ) + elif mutation == "resume_workspace": + shard = shard.model_copy( + update={"resume_workspace": ResumeWorkspace(path=f"{shard.resume_workspace.path}-other")} + ) + else: + shard = shard.model_copy(update={"created_at": run.created_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + records.validator.validate_plan_shards(run, (shard,)) + + +def test_plan_shards_reject_reordered_state() -> None: + plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(plan) + run, shards = _state_shards_for_plan(plan) + + with pytest.raises(IntegrationContractError, match="ordered"): + validator.validate_plan_shards(run, tuple(reversed(shards))) + + +def test_initial_readiness_rejects_plan_order_alias_and_backend_count(records: IntegrationRecords) -> None: + deployment = records.readiness.deployments[0] + wrong_alias = deployment.model_copy(update={"model_alias": "other"}) + readiness = records.readiness.model_copy(update={"deployments": (wrong_alias,)}) + with pytest.raises(IntegrationContractError, match="deployments"): + records.validator.validate_initial_readiness(records.attempt, readiness) + + wrong_count = deployment.model_copy(update={"expected_backends": 2}) + readiness = records.readiness.model_copy(update={"deployments": (wrong_count,)}) + with pytest.raises(IntegrationContractError, match="deployments"): + records.validator.validate_initial_readiness(records.attempt, readiness) + + multi_plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(multi_plan) + attempt = _attempt_for_plan(multi_plan) + readiness = _pending_readiness(multi_plan, attempt) + reordered = readiness.model_copy(update={"deployments": tuple(reversed(readiness.deployments))}) + with pytest.raises(IntegrationContractError, match="deployments"): + validator.validate_initial_readiness(attempt, reordered) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("run_id", "readiness run_id"), + ("shard_id", "readiness shard_id"), + ("attempt_id", "readiness attempt_id"), + ("updated_at", "cannot precede"), + ("deployment_id", "deployments"), + ], +) +def test_initial_readiness_rejects_one_field_identity_and_time_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + readiness = records.readiness + if mutation == "run_id": + readiness = readiness.model_copy(update={"run_id": "run-other"}) + elif mutation == "shard_id": + readiness = readiness.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "attempt_id": + readiness = readiness.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "updated_at": + readiness = readiness.model_copy(update={"updated_at": records.attempt.created_at - timedelta(seconds=1)}) + else: + deployment = readiness.deployments[0].model_copy(update={"deployment_id": "deployment-99999"}) + readiness = readiness.model_copy(update={"deployments": (deployment,)}) + + with pytest.raises(IntegrationContractError, match=message): + records.validator.validate_initial_readiness(records.attempt, readiness) + + +def test_initial_readiness_requires_first_pending_revision(records: IntegrationRecords) -> None: + with pytest.raises(IntegrationContractError, match="revision 1"): + records.validator.validate_initial_readiness( + records.attempt, + records.readiness.model_copy(update={"revision": 2}), + ) + with pytest.raises(IntegrationContractError, match="must be pending"): + records.validator.validate_initial_readiness( + records.attempt, + records.readiness.model_copy(update={"state": ReadinessState.READY}), + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("unplanned_shard", "planned shard"), + ("wrong_array_task", "array task"), + ("nondeterministic_attempt_id", "ordinal"), + ("unsubmitted_attempt", "scheduler"), + ("unsubmitted_attempt_with_scheduler", "submitted"), + ], +) +def test_initial_readiness_rejects_invalid_planned_attempt( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + readiness = records.readiness + if mutation == "unplanned_shard": + attempt = attempt.model_copy(update={"shard_id": "shard-99999"}) + readiness = readiness.model_copy(update={"shard_id": attempt.shard_id}) + elif mutation == "wrong_array_task": + attempt = attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) + elif mutation == "nondeterministic_attempt_id": + attempt = attempt.model_copy(update={"attempt_id": "attempt-0002"}) + readiness = readiness.model_copy(update={"attempt_id": attempt.attempt_id}) + elif mutation == "unsubmitted_attempt": + attempt = attempt.model_copy( + update={ + "state": AttemptLifecycleState.CREATED, + "scheduler": None, + "terminal_classification": None, + "candidate_output": None, + } + ) + else: + attempt = AttemptManifest.model_validate( + attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.CREATED, + "terminal_classification": None, + "candidate_output": None, + } + ) + + with pytest.raises(IntegrationContractError, match=message): + records.validator.validate_initial_readiness(attempt, readiness) + + +def test_planned_attempt_rejects_identity_task_and_unsubmitted_state(records: IntegrationRecords) -> None: + planned_shard = records.plan.shards[0] + wrong_run = records.attempt.model_copy(update={"run_id": "run-other"}) + with pytest.raises(IntegrationContractError, match="attempt run_id"): + records.validator.validate_planned_attempt(planned_shard, wrong_run) + + wrong_task = records.attempt.model_copy(update={"scheduler": SchedulerIdentity(array_job_id=4101, array_task_id=1)}) + with pytest.raises(IntegrationContractError, match="array task"): + records.validator.validate_planned_attempt(planned_shard, wrong_task) + + wrong_id = records.attempt.model_copy(update={"attempt_id": "attempt-0002"}) + with pytest.raises(IntegrationContractError, match="ordinal"): + records.validator.validate_planned_attempt(planned_shard, wrong_id) + + no_scheduler = records.attempt.model_copy( + update={ + "state": AttemptLifecycleState.CREATED, + "scheduler": None, + "terminal_classification": None, + "candidate_output": None, + } + ) + with pytest.raises(IntegrationContractError, match="scheduler"): + records.validator.validate_planned_attempt(planned_shard, no_scheduler) + + created_with_scheduler = AttemptManifest.model_validate( + no_scheduler.model_dump(mode="python") | {"scheduler": records.attempt.scheduler} + ) + with pytest.raises(IntegrationContractError, match="submitted"): + records.validator.validate_planned_attempt(planned_shard, created_with_scheduler) + + multi_plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(multi_plan) + with pytest.raises(IntegrationContractError, match="canonical shard"): + validator.validate_planned_attempt(multi_plan.shards[1], _attempt_for_plan(multi_plan)) + + +def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: + plan = _load_plan("multi_node_plan.json") + validator = PlanStateValidator(plan) + attempts = tuple(_attempt_for_planned_shard(plan, planned_shard) for planned_shard in plan.shards) + + assert ( + tuple( + validator.validate_planned_attempt(planned_shard, attempt) + for planned_shard, attempt in zip(plan.shards, attempts, strict=True) + ) + == attempts + ) + + +def test_client_candidate_validation_supports_preterminal_runtime_state(records: IntegrationRecords) -> None: + attempt = AttemptManifest.model_validate( + records.attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.RUNNING, + "terminal_classification": None, + "candidate_output": None, + } + ) + + assert ( + records.validator.validate_client_candidate( + records.plan.shards[0], + attempt, + records.client_result, + records.candidate, + ) + is records.candidate + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("partial_result", "actual records must match"), + ("failed_result", "nonempty client results"), + ("failed_attempt", "successful attempts"), + ("wrong_terminal_classification", "successfully classified"), + ("ineligible_candidate", "nonempty candidate outputs"), + ("stale_winner", "winner attempt_id"), + ("digest_mismatch", "digest"), + ("dataset_path_mismatch", "dataset path"), + ("wrong_file_extension", "output file extensions"), + ("empty_file_with_records", "must contain bytes"), + ("client_before_candidate", "client completion"), + ("winner_before_attempt", "winner publication"), + ], +) +def test_finalization_rejects_invalid_chains( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + client_result = records.client_result + candidate = records.candidate + winner = records.winner + if mutation == "partial_result": + client_result = client_result.model_copy( + update={"outcome": ClientOutcome.PARTIAL, "actual_records": 4, "early_shutdown": True} + ) + elif mutation == "failed_result": + client_result = ClientResult.model_validate( + client_result.model_dump(mode="python") + | { + "actual_records": None, + "outcome": ClientOutcome.FAILED, + "dataset_path": None, + "early_shutdown": None, + "effective_resume_mode": None, + "candidate_output_manifest": None, + "error_code": "generation_failed", + "redacted_message": "generation failed", + } + ) + elif mutation == "failed_attempt": + attempt = AttemptManifest.model_validate( + attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.FAILED, + "terminal_classification": AttemptTerminalClassification.FAILED, + "candidate_output": None, + } + ) + elif mutation == "wrong_terminal_classification": + attempt = attempt.model_copy(update={"terminal_classification": AttemptTerminalClassification.FAILED}) + elif mutation == "ineligible_candidate": + candidate = candidate.model_copy(update={"outcome": CandidateOutcome.PARTIAL}) + elif mutation == "stale_winner": + winner = winner.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "digest_mismatch": + candidate = candidate.model_copy(update={"provenance_digest": "c" * 64}) + elif mutation == "dataset_path_mismatch": + candidate = candidate.model_copy(update={"dataset_path": "/workspace/other/dataset"}) + elif mutation == "wrong_file_extension": + candidate = candidate.model_copy( + update={"files": (candidate.files[0].model_copy(update={"relative_path": "part-00000.txt"}),)} + ) + elif mutation == "empty_file_with_records": + candidate = candidate.model_copy(update={"files": (candidate.files[0].model_copy(update={"byte_size": 0}),)}) + elif mutation == "client_before_candidate": + client_result = client_result.model_copy(update={"completed_at": candidate.created_at - timedelta(seconds=1)}) + else: + winner = winner.model_copy(update={"published_at": records.attempt.updated_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + records.validator.validate_finalization_chain( + records.plan.shards[0], + attempt, + client_result, + candidate, + winner, + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("client_run_id", "client result run_id"), + ("candidate_run_id", "candidate run_id"), + ("winner_run_id", "winner run_id"), + ("client_shard_id", "client result shard_id"), + ("candidate_shard_id", "candidate shard_id"), + ("winner_shard_id", "winner shard_id"), + ("client_attempt_id", "client result attempt_id"), + ("candidate_attempt_id", "candidate attempt_id"), + ("candidate_ordinal", "attempt ordinal"), + ("winner_ordinal", "attempt ordinal"), + ("client_requested_records", "requested records"), + ("candidate_requested_records", "requested records"), + ("client_actual_records", "actual records"), + ("candidate_actual_records", "actual records"), + ("requested_resume_mode", "resume mode"), + ("effective_resume_mode", "client dataset path"), + ("client_dataset_path", "client dataset path"), + ("missing_candidate_reference", "no candidate manifest"), + ("candidate_reference_path", "candidate manifest path"), + ("attempt_candidate_reference", "attempt candidate reference"), + ("winner_candidate_reference", "winner candidate reference"), + ("candidate_before_attempt", "candidate creation"), + ("attempt_before_client", "attempt completion"), + ], +) +def test_finalization_rejects_one_field_join_drift( + records: IntegrationRecords, + mutation: str, + message: str, +) -> None: + attempt = records.attempt + client_result = records.client_result + candidate = records.candidate + winner = records.winner + candidate_reference = client_result.candidate_output_manifest + assert candidate_reference is not None + if mutation == "client_run_id": + client_result = client_result.model_copy(update={"run_id": "run-other"}) + elif mutation == "candidate_run_id": + candidate = candidate.model_copy(update={"run_id": "run-other"}) + elif mutation == "winner_run_id": + winner = winner.model_copy(update={"run_id": "run-other"}) + elif mutation == "client_shard_id": + client_result = client_result.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "candidate_shard_id": + candidate = candidate.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "winner_shard_id": + winner = winner.model_copy(update={"shard_id": "shard-99999"}) + elif mutation == "client_attempt_id": + client_result = client_result.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "candidate_attempt_id": + candidate = candidate.model_copy(update={"attempt_id": "attempt-0002"}) + elif mutation == "candidate_ordinal": + candidate = candidate.model_copy(update={"attempt_ordinal": 2}) + elif mutation == "winner_ordinal": + winner = winner.model_copy(update={"attempt_ordinal": 2}) + elif mutation == "client_requested_records": + client_result = client_result.model_copy(update={"requested_records": 7}) + elif mutation == "candidate_requested_records": + candidate = candidate.model_copy(update={"requested_records": 7}) + elif mutation == "client_actual_records": + client_result = client_result.model_copy(update={"actual_records": 7}) + elif mutation == "candidate_actual_records": + candidate = candidate.model_copy(update={"actual_records": 7}) + elif mutation == "requested_resume_mode": + client_result = client_result.model_copy(update={"requested_resume_mode": "if_possible"}) + elif mutation == "effective_resume_mode": + client_result = client_result.model_copy(update={"effective_resume_mode": "always"}) + elif mutation == "client_dataset_path": + client_result = client_result.model_copy(update={"dataset_path": "/workspace/other/dataset"}) + elif mutation == "missing_candidate_reference": + client_result = client_result.model_copy(update={"candidate_output_manifest": None}) + elif mutation == "candidate_reference_path": + client_result = client_result.model_copy( + update={ + "candidate_output_manifest": ArtifactReference( + path=posixpath.join(posixpath.dirname(candidate_reference.path), "other-manifest.json"), + sha256=candidate_reference.sha256, + ) + } + ) + elif mutation == "attempt_candidate_reference": + attempt = attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=candidate_reference.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "winner_candidate_reference": + winner = winner.model_copy( + update={ + "candidate_manifest": ArtifactReference( + path=candidate_reference.path, + sha256="a" * 64, + ) + } + ) + elif mutation == "candidate_before_attempt": + candidate = candidate.model_copy(update={"created_at": attempt.created_at - timedelta(seconds=1)}) + candidate_reference = ArtifactReference( + path=candidate_reference.path, + sha256=candidate.compute_sha256(), + ) + client_result = client_result.model_copy(update={"candidate_output_manifest": candidate_reference}) + attempt = attempt.model_copy(update={"candidate_output": candidate_reference}) + winner = winner.model_copy(update={"candidate_manifest": candidate_reference}) + else: + attempt = attempt.model_copy(update={"updated_at": client_result.completed_at - timedelta(seconds=1)}) + + with pytest.raises(IntegrationContractError, match=message): + records.validator.validate_finalization_chain( + records.plan.shards[0], + attempt, + client_result, + candidate, + winner, + ) + + +@pytest.mark.parametrize("mutation", ["path", "digest"]) +def test_planned_attempt_rejects_plan_reference_drift(records: IntegrationRecords, mutation: str) -> None: + reference = records.attempt.resolved_plan + if mutation == "path": + reference = ArtifactReference( + path=posixpath.join(posixpath.dirname(reference.path), "other-plan.json"), + sha256=reference.sha256, + ) + else: + reference = ArtifactReference(path=reference.path, sha256="a" * 64) + drifted = records.attempt.model_copy(update={"resolved_plan": reference}) + + with pytest.raises(IntegrationContractError, match=mutation): + records.validator.validate_planned_attempt(records.plan.shards[0], drifted) + + +def test_finalization_accepts_effective_resume_always(records: IntegrationRecords) -> None: + plan_payload = records.plan.model_dump(mode="python") + plan_payload["invocation"]["authored"]["resume"] = "always" + plan = ResolvedSlurmRunPlan.model_validate(plan_payload) + planned_shard = plan.shards[0] + candidate = CandidateOutputManifest.model_validate( + records.candidate.model_dump(mode="python") + | { + "dataset_path": planned_shard.resume_workspace.path, + "provenance_digest": plan.compute_sha256(), + } + ) + original_candidate_reference = records.client_result.candidate_output_manifest + assert original_candidate_reference is not None + candidate_reference = ArtifactReference( + path=original_candidate_reference.path, + sha256=candidate.compute_sha256(), + ) + attempt = AttemptManifest.model_validate( + records.attempt.model_dump(mode="python") + | {"candidate_output": candidate_reference, "resolved_plan": _plan_reference(plan)} + ) + client_result = ClientResult.model_validate( + records.client_result.model_dump(mode="python") + | { + "requested_resume_mode": "always", + "effective_resume_mode": "always", + "dataset_path": planned_shard.resume_workspace.path, + "candidate_output_manifest": candidate_reference, + } + ) + winner = ShardWinner.model_validate( + records.winner.model_dump(mode="python") | {"candidate_manifest": candidate_reference} + ) + validator = PlanStateValidator(plan) + + assert ( + validator.validate_finalization_chain( + planned_shard, + attempt, + client_result, + candidate, + winner, + ) + is winner + ) + + +def test_integration_golden_contains_no_environment_specific_values() -> None: + payload = (INTEGRATION_GOLDEN_DIR / "finalization_chain.json").read_text().casefold() + + for forbidden in ('"token":', '"password":', '"account":', '"partition":', "/users/", "/home/"): + assert forbidden not in payload + + +def _load_plan(name: str) -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIR / name).read_text()) + + +def _load_record(record_type: type[_RecordT], payload: object) -> _RecordT: + return record_type.model_validate_json(json.dumps(payload)) + + +def _plan_reference(plan: ResolvedSlurmRunPlan) -> ArtifactReference: + return ArtifactReference( + path=posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json"), + sha256=plan.compute_sha256(), + ) + + +def _state_shards_for_plan(plan: ResolvedSlurmRunPlan) -> tuple[RunManifest, tuple[ShardManifest, ...]]: + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=CREATED_AT, + authored_config=plan.authored_config, + resolved_plan=_plan_reference(plan), + shard_count=len(plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=planned.shard_id, + shard_index=planned.shard_index, + record_range=planned.record_range, + input_partition=planned.input_partition, + resume_workspace=planned.resume_workspace, + created_at=CREATED_AT + timedelta(seconds=index + 1), + ) + for index, planned in enumerate(plan.shards) + ) + return run, shards + + +def _attempt_for_plan(plan: ResolvedSlurmRunPlan) -> AttemptManifest: + return _attempt_for_planned_shard(plan, plan.shards[0]) + + +def _attempt_for_planned_shard(plan: ResolvedSlurmRunPlan, shard: PlannedShard) -> AttemptManifest: + return AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=_plan_reference(plan), + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=shard.array_task_index), + created_at=CREATED_AT, + updated_at=CREATED_AT, + ) + + +def _pending_readiness(plan: ResolvedSlurmRunPlan, attempt: AttemptManifest) -> AttemptReadiness: + return AttemptReadiness( + schema_version=1, + run_id=attempt.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=attempt.created_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 plan.deployments + ), + ) diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py new file mode 100644 index 000000000..49f124289 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -0,0 +1,857 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone + +import pytest +from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner + +from data_designer.slurm.launcher.client import SlurmCommandClient, SlurmExecutables +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + + +def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + submission = client.submit("/workspace/run.sbatch") + queue = client.query_queue((submission.job_id,)) + + assert submission.job_id == 4101 + assert tuple(record.state for record in queue) == (SchedulerState.PENDING, SchedulerState.RUNNING) + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "--export=ALL", "/workspace/run.sbatch"), + ("squeue", "--noheader", "--array", "--format=%i|%T", "--jobs=4101"), + ] + + +def test_client_submits_verified_script_text_through_standard_input() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + client = SlurmCommandClient(runner) + script = "#!/usr/bin/env bash\nexit 0\n" + + submission = client.submit_script(script) + + assert submission.job_id == 5101 + assert runner.calls == [("sbatch", "--parsable", "--export=ALL")] + assert runner.inputs == [script] + + +def test_client_holds_submission_until_initial_state_is_recorded() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + runner.script_next("scontrol", FakeCommandResponse()) + client = SlurmCommandClient(runner) + + submission = client.submit_script("#!/usr/bin/env bash\n", hold=True) + client.release(submission.job_id) + + assert runner.calls == [ + ("sbatch", "--parsable", "--hold", "--export=ALL"), + ("scontrol", "release", "5101"), + ] + + +def test_client_preserves_explicit_deployment_type_without_values_in_argv(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NEMO_DEPLOYMENT_TYPE", "nvidia-internal") + runner = _EnvironmentRunner() + client = SlurmCommandClient(runner) + + client.submit_script( + "#!/usr/bin/env bash\n", + export_environment={"HF_TOKEN": "secret-value", "NEMO_DEPLOYMENT_TYPE": "library", "SLURM_EXPORT_ENV": "ALL"}, + ) + + assert runner.command == ( + "sbatch", + "--parsable", + "--export=ALL", + ) + assert runner.environment == { + "HF_TOKEN": "secret-value", + "NEMO_DEPLOYMENT_TYPE": "library", + "SLURM_EXPORT_ENV": "ALL", + } + assert "secret-value" not in " ".join(runner.command) + + +def test_client_sets_slurm_deployment_type_when_unset(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("NEMO_DEPLOYMENT_TYPE", raising=False) + runner = _EnvironmentRunner() + + SlurmCommandClient(runner).submit_script("#!/usr/bin/env bash\n") + + assert runner.command == ("sbatch", "--parsable", "--export=ALL") + assert runner.environment == {"NEMO_DEPLOYMENT_TYPE": "slurm"} + + +def test_client_preserves_deployment_type_from_submitter_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NEMO_DEPLOYMENT_TYPE", "slurm-nvidia-internal") + runner = _EnvironmentRunner() + + SlurmCommandClient(runner).submit_script("#!/usr/bin/env bash\n") + + assert runner.command == ("sbatch", "--parsable", "--export=ALL") + assert runner.environment == {"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"} + + +def test_script_submission_classifies_scheduler_rejection_as_definite() -> None: + runner = FakeSlurmRunner() + runner.script_next("sbatch", FakeCommandResponse(stderr="rejected", returncode=2)) + + with pytest.raises(SlurmSubmissionError, match="rejected") as error: + SlurmCommandClient(runner).submit_script("#!/bin/sh\n") + + assert not error.value.may_have_succeeded + + +def test_script_submission_classifies_timeout_as_ambiguous() -> None: + with pytest.raises(SlurmSubmissionError, match="timed out") as error: + SlurmCommandClient(_TimeoutRunner()).submit_script("#!/bin/sh\n") + + assert error.value.may_have_succeeded + + +def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + client.submit("run.sbatch") + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=1) + + client.cancel(scheduler) + accounting = client.query_accounting((scheduler,)) + + assert len(accounting) == 1 + assert accounting[0].job_identity == scheduler + assert accounting[0].state is SchedulerState.CANCELLED + assert fake_slurm_runner.calls[-2:] == [ + ("scancel", "4101_1"), + ( + "sacct", + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", + "--jobs=4101_1", + ), + ] + + +def test_client_finds_one_exact_named_array_across_queue_and_accounting() -> None: + runner = FakeSlurmRunner() + job_name = f"dd-retry-{'a' * 32}" + runner.script_next("squeue", FakeCommandResponse(stdout=f"4201_0|{job_name}\n")) + runner.script_next("sacct", FakeCommandResponse(stdout=f"4201_0|{job_name}\n4201_1|{job_name}\n")) + submitted_after = datetime(2026, 9, 2, 18, tzinfo=timezone.utc) + + matches = SlurmCommandClient(runner).query_submissions_by_name(job_name, submitted_after=submitted_after) + + assert len(matches) == 1 + assert matches[0].job_id == 4201 + assert matches[0].array_task_ids == (0, 1) + assert runner.calls[0] == ( + "squeue", + "--noheader", + "--array", + "--format=%i|%.128j", + "--me", + f"--name={job_name}", + ) + assert runner.calls[1][0:6] == ( + "sacct", + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,JobName%128", + ) + assert runner.calls[1][6].startswith("--uid=") + assert runner.calls[1][7].startswith("--starttime=") + assert runner.calls[1][8] == f"--name={job_name}" + + +def test_client_rejects_a_named_lookup_result_with_a_different_name() -> None: + runner = FakeSlurmRunner() + runner.script_next("squeue", FakeCommandResponse(stdout="4201|unrelated\n")) + runner.script_next("sacct", FakeCommandResponse()) + + with pytest.raises(SlurmCommandOutputError, match="outside the requested exact name"): + SlurmCommandClient(runner).query_submissions_by_name( + f"dd-retry-{'a' * 32}", + submitted_after=datetime(2026, 9, 2, 18, tzinfo=timezone.utc), + ) + + +def test_client_deduplicates_explicit_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + client.submit("run.sbatch") + + client.query_queue((4101, 4101, SchedulerIdentity(array_job_id=4101, array_task_id=0))) + + assert fake_slurm_runner.calls[-1][-1] == "--jobs=4101,4101_0" + + +def test_client_rejects_unrequested_scheduler_records(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout="9999_0|RUNNING\n")) + + with pytest.raises(SlurmCommandOutputError, match="unrequested"): + client.query_queue((4101,)) + + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stdout="9999_0|FAILED|1:0\n")) + with pytest.raises(SlurmCommandOutputError, match="unrequested"): + client.query_accounting((SchedulerIdentity(array_job_id=4101, array_task_id=0),)) + + +def test_client_observes_regular_cpu_job() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + client = SlurmCommandClient(runner) + + submission = client.submit("image-build.sbatch") + queue = client.query_queue((5101,)) + runner.set_job_state(5101, queue_state=None, accounting_state="COMPLETED") + accounting = client.query_accounting((5101,)) + + assert submission.job_id == 5101 + assert queue[0].job_identity == 5101 + assert queue[0].state is SchedulerState.PENDING + assert accounting[0].job_identity == 5101 + assert accounting[0].state is SchedulerState.COMPLETED + + +def test_client_ignores_only_array_parent_observation_for_exact_task() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) + client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + + records = client.query_accounting((task,)) + + assert tuple(record.job_identity for record in records) == (task,) + + +def test_client_keeps_explicitly_selected_parent_and_task_observations() -> None: + runner = FakeSlurmRunner() + runner.script_next("sacct", FakeCommandResponse(stdout="4101|RUNNING|0:0\n4101_0|COMPLETED|0:0\n")) + client = SlurmCommandClient(runner) + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + + records = client.query_accounting((4101, task)) + + assert tuple(record.job_identity for record in records) == (4101, task) + + +def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="at least one"): + client.query_queue(()) + with pytest.raises(ValueError, match="positive 32-bit integers"): + client.query_accounting((0,)) + with pytest.raises(ValueError, match="positive 32-bit integers"): + client.cancel(True) + with pytest.raises(ValueError, match="32-bit"): + client.cancel(1 << 32) + with pytest.raises(ValueError, match="array-task IDs"): + client.cancel(SchedulerIdentity(array_job_id=4101, array_task_id=1 << 32)) + + assert fake_slurm_runner.calls == [] + + +def test_client_queries_default_partition_gpu_inventory(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + assert client.query_gpu_counts() == (2,) + assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%P|%G")] + + +def test_client_queries_partition_scoped_gpu_inventory() -> None: + command = ("sinfo", "--noheader", "--format=%G", "--partition=batch") + runner = FakeSlurmRunner(sinfo_responses={command: FakeCommandResponse(stdout="gpu:a100:8\n")}) + + assert SlurmCommandClient(runner).query_gpu_counts(partition="batch") == (8,) + assert runner.calls == [command] + + +def test_client_queries_multiple_partition_gpu_inventory() -> None: + command = ("sinfo", "--noheader", "--format=%P|%G", "--partition=batch,other") + runner = FakeSlurmRunner( + sinfo_responses={command: FakeCommandResponse(stdout="batch|gpu:a100:8\nother|gpu:h100:8\n")} + ) + + assert SlurmCommandClient(runner).query_gpu_counts(partition="batch,other") == (8, 8) + assert runner.calls == [command] + + +@pytest.mark.parametrize( + "output", + ["batch|gpu:a100:8\n", "batch|gpu:a100:8\nother|(null)\n", "batch|gpu:a100:8\nother|gpu:h100:8\nother|(null)\n"], +) +def test_client_rejects_selected_partition_without_gpu_inventory(output: str) -> None: + command = ("sinfo", "--noheader", "--format=%P|%G", "--partition=batch,other") + runner = FakeSlurmRunner(sinfo_responses={command: FakeCommandResponse(stdout=output)}) + + with pytest.raises(SlurmCommandOutputError, match="requested partition|no GPU resources"): + SlurmCommandClient(runner).query_gpu_counts(partition="batch,other") + + +@pytest.mark.parametrize("selection", ["batch,,other", "batch,other;echo", "batch,batch"]) +def test_client_rejects_invalid_gpu_partition_without_running_command( + fake_slurm_runner: FakeSlurmRunner, selection: str +) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="comma-separated list"): + client.query_gpu_counts(partition=selection) + + assert fake_slurm_runner.calls == [] + + +def test_client_normalizes_command_failures(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next( + "sacct", + FakeCommandResponse(stderr="accounting\nservice unavailable\n", returncode=2), + ) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError, match="sacct failed with exit code 2: accounting service unavailable"): + client.query_accounting((4101,)) + + +@pytest.mark.parametrize( + "detail", + ( + "Invalid job id specified", + "slurm_load_jobs error: Invalid job id specified", + ), +) +def test_client_treats_unknown_queue_job_as_an_absent_row( + fake_slurm_runner: FakeSlurmRunner, + detail: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=f"{detail}\n", returncode=1)) + + assert SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) == () + + +@pytest.mark.parametrize( + "detail", + ( + "Invalid job id specified for partition", + "queue unavailable", + ), +) +def test_client_preserves_non_missing_queue_failures( + fake_slurm_runner: FakeSlurmRunner, + detail: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=f"{detail}\n", returncode=1)) + + with pytest.raises(SlurmCommandError, match=detail): + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + +def test_client_removes_terminal_controls_from_command_failures(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next( + "squeue", + FakeCommandResponse(stderr="queue unavailable\x1b[31m\n", returncode=2), + ) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError) as error: + client.query_queue((4101,)) + + assert "\x1b" not in str(error.value) + + +@pytest.mark.parametrize( + ("diagnostic", "expected", "secret"), + ( + ("--api-key\x00plaintext-secret", "--api-key ", "plaintext-secret"), + ("--api_key\x1funderscore-secret", "--api_key ", "underscore-secret"), + ("--access-token\u200bformat-secret", "--access-token ", "format-secret"), + ( + "Authorization: Custom control-secret\x00value\nstatus=failed", + "Authorization: status=failed", + "control-secret", + ), + ("HF_TOKEN=secret\x00suffix status=failed", "HF_TOKEN= status=failed", "suffix"), + ( + "https://user:secret\x00suffix@example.test/path status=failed", + "https://@example.test/path status=failed", + "suffix", + ), + ("HF_TOKEN\x00=boundary-secret", "HF_TOKEN =", "boundary-secret"), + ( + 'HF_TOKEN="first-secret";status=failed --api-key\x00second-secret next=ready', + "HF_TOKEN=;status=failed --api-key next=ready", + "second-secret", + ), + ( + "HF_TOKEN=\x00suffix status=failed", + "HF_TOKEN= status=failed", + "suffix", + ), + ), + ids=( + "nul", + "unit-separator", + "format-control", + "authorization-line-boundary", + "assignment-secret-suffix", + "uri-userinfo-secret-suffix", + "assignment-hidden-boundary", + "raw-and-normalized-pass", + "placeholder-collision", + ), +) +def test_client_redacts_secrets_obscured_by_control_characters( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + expected: str, + secret: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == expected + assert secret not in detail + + +@pytest.mark.parametrize( + "control", + ( + pytest.param("\t", id="tab"), + pytest.param("\v", id="vertical-tab"), + pytest.param("\f", id="form-feed"), + pytest.param("\x1c", id="file-separator"), + pytest.param("\x1d", id="group-separator"), + pytest.param("\x1e", id="record-separator"), + pytest.param("\x1f", id="unit-separator"), + pytest.param("\x85", id="next-line"), + ), +) +def test_client_redacts_whitespace_control_credential_suffixes( + fake_slurm_runner: FakeSlurmRunner, + control: str, +) -> None: + diagnostic = ( + f"HF_TOKEN=assignment-secret{control}suffix status=failed " + f"https://user:uri-secret{control}suffix@example.test/path" + ) + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == "HF_TOKEN= status=failed https://@example.test/path" + assert "secret" not in detail + assert "suffix" not in detail + + +@pytest.mark.parametrize( + "boundary", + ( + pytest.param("\u00a0", id="no-break-space"), + pytest.param("\u1680", id="ogham-space"), + pytest.param("\u2007", id="figure-space"), + pytest.param("\u202f", id="narrow-no-break-space"), + pytest.param("\u205f", id="medium-mathematical-space"), + pytest.param("\u3000", id="ideographic-space"), + ), +) +def test_client_redacts_unicode_whitespace_credential_suffixes( + fake_slurm_runner: FakeSlurmRunner, + boundary: str, +) -> None: + diagnostic = ( + f"HF_TOKEN=assignment-secret{boundary}suffix status=failed " + f"https://user:uri-secret{boundary}suffix@example.test/path" + ) + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == "HF_TOKEN= status=failed https://@example.test/path" + assert "secret" not in detail + assert "suffix" not in detail + + +@pytest.mark.parametrize("line_break", ("\n", "\r", "\r\n"), ids=("line-feed", "carriage-return", "crlf")) +def test_client_retains_unambiguous_diagnostic_line_boundaries( + fake_slurm_runner: FakeSlurmRunner, + line_break: str, +) -> None: + fake_slurm_runner.script_next( + "squeue", + FakeCommandResponse(stderr=f"HF_TOKEN=secret{line_break}status=failed", returncode=2), + ) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + assert str(error.value).endswith("HF_TOKEN= status=failed") + + +@pytest.mark.parametrize( + ("diagnostic", "secret"), + ( + ("HF_TOKEN=super-secret-value", "super-secret-value"), + ('HF_TOKEN="quoted secret value"', "quoted secret value"), + ("--api-key plaintext-secret", "plaintext-secret"), + ("Authorization: Bearer bearer-secret;suffix status=failed", "bearer-secret;suffix"), + ('{"access_token":"json-secret"}', "json-secret"), + ("Authorization: Bearer bearer-secret", "bearer-secret"), + ("https://user:url-secret@example.test/index", "url-secret"), + (f"token github_pat_{'a' * 24}", f"github_pat_{'a' * 24}"), + ), + ids=( + "environment", + "quoted-environment", + "option", + "authorization-punctuation", + "json", + "authorization", + "url", + "known-token", + ), +) +def test_client_redacts_secrets_from_command_failures( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + secret: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError) as error: + client.query_queue((4101,)) + + assert secret not in str(error.value) + assert "" in str(error.value) + + +@pytest.mark.parametrize( + ("diagnostic", "expected", "secret"), + ( + ("Authorization: Token token-secret status=failed", "Authorization: ", "token-secret"), + ("Authorization=ApiKey key-secret", "Authorization=", "key-secret"), + ("authorization: opaque secret value", "authorization: ", "opaque secret value"), + ( + "Authorization: Custom custom-secret\nstatus=failed", + "Authorization: status=failed", + "custom-secret", + ), + ( + '{"Authorization": "Bearer serialized-secret-value", "status": "failed"}', + '{"Authorization": ', + "serialized-secret-value", + ), + ), + ids=("token", "api-key", "bare", "newline-boundary", "serialized-header"), +) +def test_client_redacts_complete_authorization_values( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + expected: str, + secret: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == expected + assert secret not in detail + + +def test_client_redacts_quoted_authorization_value_with_control_obscured_suffix( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + scheme = "Bear" + "er" + diagnostic = f'Authorization: "{scheme} token"\x00secret-suffix\nstatus=failed' + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == "Authorization: status=failed" + assert "token" not in detail + assert "suffix" not in detail + + +@pytest.mark.parametrize( + ("diagnostic", "expected"), + ( + ( + "https://user:first@second@example.test/index", + "https://@example.test/index", + ), + ("ssh://user:secret@example.test/repository", "ssh://@example.test/repository"), + ("git+ssh://user:p@ss@example.test/repository", "git+ssh://@example.test/repository"), + ("https://example.test/path@owner", "https://example.test/path@owner"), + ), + ids=("last-at", "ssh", "compound-scheme", "at-after-path"), +) +def test_client_redacts_general_uri_userinfo( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + expected: str, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + assert str(error.value).partition(": ")[2] == expected + + +@pytest.mark.parametrize( + ("diagnostic", "expected"), + ( + ("HF_TOKEN=secret;status=failed job=4", "HF_TOKEN= job=4"), + ("HF_TOKEN=secret,status=failed job=4", "HF_TOKEN= job=4"), + ("HF_TOKEN=secret status=failed", "HF_TOKEN= status=failed"), + ("status=failed;HF_TOKEN=secret;job=4 next=ready", "status=failed;HF_TOKEN= next=ready"), + ("status=failed,HF_TOKEN=secret,job=4 next=ready", "status=failed,HF_TOKEN= next=ready"), + ("status=failed HF_TOKEN=secret job=4", "status=failed HF_TOKEN= job=4"), + ("status=failed; HF_TOKEN=secret, job=4", "status=failed; HF_TOKEN= job=4"), + ("status=failed;HF_TOKEN=secret", "status=failed;HF_TOKEN="), + ("status=failed,HF_TOKEN=secret", "status=failed,HF_TOKEN="), + ("status=failed HF_TOKEN=secret", "status=failed HF_TOKEN="), + ("status=failed|HF_TOKEN=secret", "status=failed|HF_TOKEN="), + ("status=failed/HF_TOKEN=secret", "status=failed/HF_TOKEN="), + ("HF_TOKEN=secret;suffix status=failed", "HF_TOKEN= status=failed"), + ("HF_TOKEN=secret,suffix status=failed", "HF_TOKEN= status=failed"), + ("HF_TOKEN=secret;part=value status=failed", "HF_TOKEN= status=failed"), + ("HF_TOKEN=secret,part=value status=failed", "HF_TOKEN= status=failed"), + ('HF_TOKEN="secret";status=failed', "HF_TOKEN=;status=failed"), + ), + ids=( + "first-semicolon", + "first-comma", + "first-whitespace", + "middle-semicolon", + "middle-comma", + "middle-whitespace", + "middle-spaced-punctuation", + "last-semicolon", + "last-comma", + "last-whitespace", + "pipe-before-secret", + "slash-before-secret", + "semicolon-inside-value", + "comma-inside-value", + "assignment-looking-semicolon-suffix", + "assignment-looking-comma-suffix", + "quoted-value-boundary", + ), +) +def test_client_redacts_adjacent_assignments_fail_closed( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + expected: str, +) -> None: + fake_slurm_runner.script_next( + "squeue", + FakeCommandResponse(stderr=diagnostic, returncode=2), + ) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + assert expected in str(error.value) + assert "HF_TOKEN=secret" not in str(error.value) + + +@pytest.mark.parametrize( + ("diagnostic", "expected", "secret"), + ( + ("--format compact|--api-key secret", "--format compact|--api-key ", "secret"), + ("--format compact|--api_key secret", "--format compact|--api_key ", "secret"), + ("--api-key plaintext-secret", "--api-key ", "plaintext-secret"), + ("--api_key plaintext-secret", "--api_key ", "plaintext-secret"), + ('--api_Key "quoted secret value" status=failed', "--api_Key status=failed", "quoted secret value"), + ( + "status=failed/--access_token.key 'single quoted secret' next=ready", + "status=failed/--access_token.key next=ready", + "single quoted secret", + ), + ("--api-key secret;part=value status=failed", "--api-key status=failed", "secret;part=value"), + ("--api_key secret,part=value status=failed", "--api_key status=failed", "secret,part=value"), + ("--access-token_key mixed-secret", "--access-token_key ", "mixed-secret"), + ("--output_format compact status=failed", "--output_format compact status=failed", None), + ), + ids=( + "hyphen-after-unknown-separator", + "underscore-after-unknown-separator", + "hyphen-unquoted", + "underscore-unquoted", + "mixed-case-double-quoted", + "mixed-dot-underscore-single-quoted", + "hyphen-punctuation-suffix", + "underscore-punctuation-suffix", + "mixed-hyphen-underscore", + "nonsecret-underscore-preserved", + ), +) +def test_client_redacts_option_values_without_overlap( + fake_slurm_runner: FakeSlurmRunner, + diagnostic: str, + expected: str, + secret: str | None, +) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr=diagnostic, returncode=2)) + + with pytest.raises(SlurmCommandError) as error: + SlurmCommandClient(fake_slurm_runner).query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert detail == expected + if secret is not None: + assert secret not in detail + + +def test_client_bounds_command_failure_detail(fake_slurm_runner: FakeSlurmRunner) -> None: + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stderr="x" * 600, returncode=2)) + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(SlurmCommandError) as error: + client.query_queue((4101,)) + + detail = str(error.value).partition(": ")[2] + assert len(detail) == 512 + assert detail.endswith("...") + + +def test_client_normalizes_execution_errors() -> None: + client = SlurmCommandClient(_FailingRunner()) + + with pytest.raises(SlurmCommandError, match="squeue could not be executed") as error: + client.query_queue((4101,)) + + assert isinstance(error.value.__cause__, FileNotFoundError) + + +def test_client_normalizes_command_timeouts() -> None: + client = SlurmCommandClient(_TimeoutRunner()) + + with pytest.raises(SlurmCommandError, match="command timed out") as error: + client.query_queue((4101,)) + + assert isinstance(error.value.__cause__, subprocess.TimeoutExpired) + + +@pytest.mark.parametrize("returncode", (0, 2)) +def test_client_rejects_non_text_runner_output(returncode: int) -> None: + client = SlurmCommandClient(_NonTextRunner(returncode)) + + with pytest.raises(SlurmCommandError, match="malformed process result"): + client.query_queue((4101,)) + + +def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + client.submit("/workspace/run; touch injected.sbatch") + + assert fake_slurm_runner.calls[0] == ( + "sbatch", + "--parsable", + "--export=ALL", + "/workspace/run; touch injected.sbatch", + ) + + +def test_client_rejects_option_like_script_path(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="must not begin"): + client.submit("--wrap=unexpected") + + assert fake_slurm_runner.calls == [] + + +@pytest.mark.parametrize("script_path", ("", "bad\npath")) +def test_client_rejects_invalid_script_path(fake_slurm_runner: FakeSlurmRunner, script_path: str) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="batch script path"): + client.submit(script_path) + + assert fake_slurm_runner.calls == [] + + +@pytest.mark.parametrize("script", ("", "bad\0script", b"not text")) +def test_client_rejects_invalid_script_text(fake_slurm_runner: FakeSlurmRunner, script: object) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="batch script text"): + client.submit_script(script) # type: ignore[arg-type] + + assert fake_slurm_runner.calls == [] + + +@pytest.mark.parametrize("executable", ("", "sbatch --wait", "sbatch\n")) +def test_executables_reject_invalid_tokens(executable: str) -> None: + with pytest.raises(ValueError, match="Slurm executable"): + SlurmExecutables(sbatch=executable) + + +class _FailingRunner: + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + del command + raise FileNotFoundError("missing executable") + + +class _EnvironmentRunner: + def __init__(self) -> None: + self.command: tuple[str, ...] = () + self.environment: dict[str, str] = {} + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + self.command = tuple(command) + self.environment = dict(environment or {}) + return subprocess.CompletedProcess(command, 0, stdout="42\n", stderr="") + + +class _TimeoutRunner: + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + del input_text, environment + raise subprocess.TimeoutExpired(command, 30.0) + + +class _NonTextRunner: + def __init__(self, returncode: int) -> None: + self._returncode = returncode + + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + completed = subprocess.CompletedProcess(command, self._returncode, stdout="ok", stderr="") + if self._returncode: + completed.stderr = b"not text" # type: ignore[assignment] + else: + completed.stdout = b"not text" # type: ignore[assignment] + return completed diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py new file mode 100644 index 000000000..861bc1731 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import shlex +import subprocess +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import ContainerMount, injected_profile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.launcher.collection import render_collection_script +from data_designer.slurm.launcher.renderer import render_generation_retry_script +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import CollectionPlan, CollectionShard, RetryPlan, RetryShard +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.errors import StateConflictError + + +def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace_mount = ContainerMount(source="/workspace", target="/workspace", read_only=False) + output_mount = ContainerMount( + source="/workspace/primary/runs/run-001", + target="/exports", + read_only=False, + ) + mounts = (workspace_mount, output_mount) + profile = multi_node_plan.selected_profile.profile.model_copy(update={"container_mounts": list(mounts)}) + selection = multi_node_plan.selected_profile.model_copy( + update={ + "profile": profile, + "profile_sha256": compute_canonical_json_sha256(profile.model_dump(mode="json")), + } + ) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps( + multi_node_plan.model_copy(update={"container_mounts": mounts, "selected_profile": selection}).model_dump( + mode="json" + ) + ) + ) + destination = CollectionDestinationResolver().resolve(plan) + collection = CollectionPlan( + schema_version=1, + collection_id="collection-0001", + run_id=plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=plan.compute_sha256(), + ), + planned_shards=( + CollectionShard( + shard_id="shard-00000", + winner_manifest=ArtifactReference( + path="/workspace/primary/runs/run-001/shards/shard-00000/winner.json", + sha256="a" * 64, + ), + ), + ), + python_executable="/shared/data-designer/bin/python", + package_version="0.9.9", + runtime_bundle=ArtifactReference(path="/workspace/primary/runtime/collection.tar.gz", sha256="b" * 64), + host_destination=destination.host_path, + num_partitions=1, + ) + + script = render_collection_script(plan, collection, destination) + + assert f"#SBATCH --job-name={collection.submission_job_name}" in script + assert "#SBATCH --partition=cpu" in script + assert "#SBATCH --partition=batch" not in script + assert f"dd-collect-{plan.run_id}" not in script + assert ( + 'readonly DD_COLLECTION_PLAN="/workspace/primary/runs/run-001/collections/collection-0001/plan.json"' in script + ) + assert 'readonly DD_COLLECTION_PYTHON="/shared/data-designer/bin/python"' in script + assert 'readonly DD_PACKAGE_VERSION="0.9.9"' in script + assert f'readonly DD_RUNTIME_ARCHIVE="{collection.runtime_bundle.path}"' in script + assert f'readonly DD_RUNTIME_SHA256="{collection.runtime_bundle.sha256}"' in script + assert 'export PYTHONPATH="${DD_SCRATCH_ROOT}/runtime"' in script + assert "data_designer.slurm.state.collection_worker" in script + assert '"${DD_COLLECTION_PYTHON}" -m data_designer.slurm.state.collection_worker' in script + assert "enroot start" not in script + assert "DD_CLIENT_IMAGE" not in script + assert "DD_ENROOT_MOUNTS" not in script + assert "#SBATCH --gres=" not in script + assert "#SBATCH --gpus=" not in script + assert script.index("trap 'exit 143' TERM") < script.index("\ndd_initialize_local_scratch\n") + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + legacy = CollectionPlan.model_validate( + collection.model_dump(mode="python", exclude={"package_version", "runtime_bundle"}) + ) + legacy_script = render_collection_script(plan, legacy, destination) + assert 'readonly DD_PACKAGE_VERSION="0.9.2"' in legacy_script + assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"' in legacy_script + assert f'readonly DD_RUNTIME_SHA256="{plan.runtime_bundle.sha256}"' in legacy_script + with pytest.raises(ValidationError, match="package version and runtime bundle must be set together"): + CollectionPlan.model_validate(legacy.model_dump(mode="python") | {"package_version": "0.9.9"}) + + +def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( + multi_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, +) -> None: + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=multi_node_plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=multi_node_plan.compute_sha256(), + ), + planned_shards=( + RetryShard( + shard_id="shard-00001", + attempt_id="attempt-0002", + attempt_ordinal=2, + array_task_index=1, + ), + ), + effective_resume_mode="never", + ) + + script = render_generation_retry_script(multi_node_plan, retry) + + assert f"#SBATCH --job-name={retry.submission_job_name}" in script + assert "#SBATCH --cpus-per-gpu=4\n" in script + assert "#SBATCH --cpus-per-task=" not in script + assert multi_node_plan.submission.job_name not in script + assert "#SBATCH --array=1%2" in script + assert 'DD_ATTEMPT_ORDINAL="0002"' in script + assert 'readonly DD_ATTEMPT_MANIFEST="${DD_ATTEMPT_DIR}/attempt.json"' in script + assert f'readonly DD_RETRY_ID="{retry.retry_id}"' in script + assert f'readonly DD_RETRY_PLAN_SHA256="{retry.compute_sha256()}"' in script + assert 'readonly DD_EFFECTIVE_RESUME_MODE="never"' in script + assert script.index("DD_ATTEMPT_MANIFEST") < script.index("DD_RUNTIME_ROOT") + assert "data_designer.slurm.state.attempt_identity" in script + assert "enroot start" not in script + assert "--container-image=" not in script + assert "--container-mounts=" not in script + assert "--container-env=" not in script + binding = next(line for line in script.splitlines() if line.partition("=")[0].endswith("DD_CLIENT_PYTHON")) + plan_path = tmp_path / "resolved-plan.json" + plan_path.write_text(multi_node_plan.serialize_json()) + reader = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/plan_reader.sh" + completed = subprocess.run( + ( + "bash", + "-c", + f'set -e\n{binding}\nsource {shlex.quote(str(reader))}\ndd_read_control_plan "$1"\n' + 'printf "%s\\n" "$DD_CLIENT_PYTHON"', + "retry", + str(plan_path), + ), + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == multi_node_plan.client.runtime.python_executable + assert 'export PYTHONPATH="${DD_RUNTIME_ROOT}"' in script + assert "--gres=none" in script + assert "--export=ALL" in script + assert '"${DD_CLIENT_PYTHON}" -m data_designer.slurm.state.attempt_identity' in script + assert '--array-job-id "${DD_ARRAY_JOB_ID}" --array-task-id "${DD_ARRAY_TASK_ID}"' in script + assert script.index("dd_stage_allocation_runtime") < script.index("data_designer.slurm.state.attempt_identity") + assert script.index("trap 'exit 143' TERM") < script.index("\ndd_initialize_local_scratch\n") + assert f'#SBATCH --output="/workspace/primary/runs/run-001/retries/{retry.retry_id}/slurm-%A_%a.out"' in script + assert f'#SBATCH --error="/workspace/primary/runs/run-001/retries/{retry.retry_id}/slurm-%A_%a.err"' in script + assert '"${DD_RETRY_PLAN_SHA256}" "${DD_EFFECTIVE_RESUME_MODE}"' in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +def test_retry_renderer_uses_profile_slurm_bin_path( + multi_node_plan: ResolvedSlurmRunPlan, gpu_request_mode: str +) -> None: + scheduler = multi_node_plan.selected_profile.profile.scheduler.model_copy(update={"bin_path": "/opt/slurm/bin"}) + profile = multi_node_plan.selected_profile.profile.model_copy( + update={"scheduler": scheduler, "gpu_request_mode": gpu_request_mode} + ) + plan = multi_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=plan.compute_sha256(), + ), + planned_shards=( + RetryShard( + shard_id="shard-00001", + attempt_id="attempt-0002", + attempt_ordinal=2, + array_task_index=1, + ), + ), + effective_resume_mode="never", + ) + + script = render_generation_retry_script(plan, retry) + + assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in script + assert "#SBATCH --exclusive" in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_destination_reauthorizes_explicit_path_through_workspace_mapping( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace_root = multi_node_plan.selected_profile.profile.workspace_root + requested = (Path(workspace_root) / "explicit" / "collected").as_posix() + plan = multi_node_plan.model_copy(update={"container_mounts": ()}) + resolver = CollectionDestinationResolver() + destination = resolver.resolve(plan, requested) + collection = CollectionPlan( + schema_version=1, + collection_id="collection-0001", + run_id=plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-001/resolved-plan.json", + sha256=plan.compute_sha256(), + ), + planned_shards=( + CollectionShard( + shard_id="shard-00000", + winner_manifest=ArtifactReference( + path="/workspace/primary/runs/run-001/shards/shard-00000/winner.json", + sha256="a" * 64, + ), + ), + ), + python_executable="/shared/data-designer/bin/python", + package_version="0.9.9", + runtime_bundle=ArtifactReference(path="/workspace/primary/runtime/collection.tar.gz", sha256="b" * 64), + host_destination=requested, + num_partitions=1, + ) + + assert destination.authorized_root == workspace_root + assert resolver.validate_persisted(plan, collection) == destination + + +@pytest.mark.parametrize( + ("destination_suffix", "error"), + [ + ("images/collected", "package-managed workspace state"), + ("runtime/collected", "package-managed workspace state"), + ("benchmarks/collected", "package-managed workspace state"), + ("managed-assets/collected", "managed assets"), + ("runs/run-001/collections", "package-managed run state"), + ("runs/run-001/retries", "package-managed run state"), + ("runs/other-run/output", "package-managed run state"), + ("runs/other-run/merge-output", "package-managed run state"), + ], +) +def test_destination_rejects_package_managed_paths( + multi_node_plan: ResolvedSlurmRunPlan, + destination_suffix: str, + error: str, +) -> None: + workspace_root = multi_node_plan.selected_profile.profile.workspace_root + requested = (Path(workspace_root) / destination_suffix).as_posix() + + with pytest.raises(StateConflictError, match=error): + CollectionDestinationResolver().resolve(multi_node_plan, requested) + + +def test_destination_allows_current_run_merge_output_subtree(multi_node_plan: ResolvedSlurmRunPlan) -> None: + requested = ( + Path(multi_node_plan.selected_profile.profile.workspace_root) + / "runs" + / multi_node_plan.run_id + / "merge-output" + / "collected" + ).as_posix() + + assert CollectionDestinationResolver().resolve(multi_node_plan, requested).host_path == requested + + +def test_destination_ignores_duplicate_container_targets_for_one_host_root( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + requested = "/host/exports/run-001" + mounts = ( + ContainerMount(source="/host/exports", target="/container/a", read_only=False), + ContainerMount(source="/host/exports", target="/container/b", read_only=False), + ) + plan = multi_node_plan.model_copy(update={"container_mounts": mounts}) + + assert CollectionDestinationResolver().resolve(plan, requested).authorized_root == "/host/exports" diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py new file mode 100644 index 000000000..c6d1f955e --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +import pytest + +from data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import SlurmQueueEntry +from data_designer.slurm.launcher.parsing import ( + parse_accounting, + parse_default_partition_gpu_counts, + parse_gpu_counts, + parse_queue, + parse_state, + parse_submission, +) +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "slurm" +OVERSIZED_DECIMAL = "9" * 5000 + + +def test_parse_submission_accepts_non_federated_parsable_sbatch_output() -> None: + submission = parse_submission("4101\n") + + assert submission.job_id == 4101 + + +def test_parse_submission_rejects_federated_receipts() -> None: + with pytest.raises(SlurmCommandOutputError, match="federated"): + parse_submission("4101;primary\n") + + +@pytest.mark.parametrize( + "output", + ("", "0", "Submitted batch job 4101", "٤١٠١", "4101;", "4101;bad name", f"{'0' * 5000}1"), +) +def test_parse_submission_rejects_malformed_output(output: str) -> None: + with pytest.raises(SlurmCommandOutputError, match="invalid"): + parse_submission(output) + + +def test_parse_submission_enforces_slurm_job_id_width() -> None: + assert parse_submission(str((1 << 32) - 1)).job_id == (1 << 32) - 1 + + with pytest.raises(SlurmCommandOutputError, match="invalid job ID"): + parse_submission(str(1 << 32)) + + +def test_parse_queue_normalizes_active_array_tasks() -> None: + records = parse_queue((GOLDEN_DIRECTORY / "squeue_active.txt").read_text()) + + assert records == ( + _make_queue_entry(0, SchedulerState.PENDING), + _make_queue_entry(1, SchedulerState.RUNNING), + ) + + +def test_parse_queue_normalizes_regular_jobs() -> None: + records = parse_queue("5101|RUNNING\n") + + assert records == (SlurmQueueEntry(job_identity=5101, state=SchedulerState.RUNNING),) + + +@pytest.mark.parametrize( + ("raw_state", "expected"), + ( + ("CONFIGURING", SchedulerState.PENDING), + ("COMPLETING", SchedulerState.RUNNING), + ("COMPLETED+", SchedulerState.COMPLETED), + ("CANCELLED by 1234", SchedulerState.CANCELLED), + ("TIMEOUT", SchedulerState.TIMED_OUT), + ("NODE_FAIL", SchedulerState.NODE_FAILED), + ("PREEMPTED", SchedulerState.PREEMPTED), + ("REQUEUED", SchedulerState.REQUEUED), + ("REQUEUE_HOLD", SchedulerState.PENDING), + ("RESV_DEL_HOLD", SchedulerState.PENDING), + ("SPECIAL_EXIT", SchedulerState.PENDING), + ("OUT_OF_MEMORY", SchedulerState.OUT_OF_MEMORY), + ("A_NEW_STATE", SchedulerState.UNKNOWN), + ), +) +def test_parse_state_normalizes_long_slurm_spellings(raw_state: str, expected: SchedulerState) -> None: + assert parse_state(raw_state) is expected + + +def test_parse_accounting_normalizes_parent_and_terminal_array_task_rows() -> None: + output = "4101|RUNNING|0:0\n" + (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + + records = parse_accounting(output) + + assert tuple(record.state for record in records) == ( + SchedulerState.RUNNING, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.PREEMPTED, + SchedulerState.REQUEUED, + SchedulerState.OUT_OF_MEMORY, + SchedulerState.CANCELLED, + ) + assert records[1].process_exit_code.exit_status == 0 + assert records[1].process_exit_code.termination_signal == 125 + + +def test_parse_accounting_normalizes_regular_jobs() -> None: + records = parse_accounting("5101|COMPLETED|0:0\n") + + assert len(records) == 1 + assert records[0].job_identity == 5101 + assert records[0].state is SchedulerState.COMPLETED + + +def test_empty_scheduler_output_preserves_absent_evidence_for_reconciliation() -> None: + assert parse_queue("") == () + assert parse_accounting("\n") == () + + +@pytest.mark.parametrize( + ("parser", "output", "message"), + ( + (parse_queue, "malformed scheduler output\n", "two fields"), + (parse_queue, "not-a-job|RUNNING\n", "job or array-task ID"), + (parse_queue, "4101_0|RUNNING\n4101_0|PENDING\n", "duplicates"), + (parse_accounting, "4101_0|FAILED\n", "three fields"), + (parse_accounting, "4101_0|FAILED|not-an-exit-code\n", "exit code"), + (parse_accounting, "4101_0.batch|FAILED|1:0\n", "job or array-task ID"), + (parse_accounting, "garbage.step|FAILED|1:0\n", "job or array-task ID"), + (parse_queue, "4101_0|COMPLETED unexpectedly\n", "unexpected whitespace"), + ), +) +def test_scheduler_parsers_reject_malformed_or_ambiguous_rows( + parser: Callable[[str], object], + output: str, + message: str, +) -> None: + with pytest.raises(SlurmCommandOutputError, match=message): + parser(output) + + +@pytest.mark.parametrize( + ("parser", "output", "message"), + ( + (parse_submission, OVERSIZED_DECIMAL, "invalid job ID"), + (parse_queue, f"4101_{OVERSIZED_DECIMAL}|RUNNING", "array-task ID"), + (parse_accounting, f"4101_0|FAILED|{OVERSIZED_DECIMAL}:0", "exit code"), + (parse_gpu_counts, f"gpu:{OVERSIZED_DECIMAL}", "invalid GPU resource"), + ), + ids=("submission-job-id", "queue-task-id", "accounting-exit-code", "gpu-count"), +) +def test_parsers_normalize_oversized_numeric_fields( + parser: Callable[[str], object], + output: str, + message: str, +) -> None: + with pytest.raises(SlurmCommandOutputError, match=message): + parser(output) + + +@pytest.mark.parametrize( + ("output", "expected"), + ( + ("gpu:2\n", (2,)), + ("gpu:a100:8(S:0-7)\n", (8,)), + ("gpu:a100:4(S:0-1,4-5)\n", (4,)), + ("mps:100,gpu:a100:4\n", (4,)), + ("gpu:a100:4,gpu:h100:4\n(null)\n", (8,)), + ("(null)\nN/A\n", ()), + ), +) +def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tuple[int, ...]) -> None: + assert parse_gpu_counts(output) == expected + + +def test_parse_default_partition_gpu_counts_ignores_other_partitions() -> None: + output = "cpu|(null)\ngpu*|gpu:h100:8\nother|gpu:a100:4\ngpu*|gpu:h100:8\n" + + assert parse_default_partition_gpu_counts(output) == (8, 8) + + +@pytest.mark.parametrize( + "output", + ("gpu*|gpu:8|extra\n", "*|gpu:8\n", "gpu*|gpu:8\nother*|gpu:4\n"), + ids=("field-count", "missing-name", "multiple-defaults"), +) +def test_parse_default_partition_gpu_counts_rejects_malformed_rows(output: str) -> None: + with pytest.raises(SlurmCommandOutputError): + parse_default_partition_gpu_counts(output) + + +@pytest.mark.parametrize( + "output", + ( + "gpu:a100:not-a-count\n", + "gpu:a100:4(S:0-1,4-5\n", + "gpu:a100:4)\n", + "gpu:a100:4((S:0-1))\n", + "gpu:a100:4,\n", + ",gpu:a100:4\n", + "gpu:a100:4,,mps:100\n", + ), +) +def test_parse_gpu_counts_rejects_malformed_gpu_resources(output: str) -> None: + with pytest.raises(SlurmCommandOutputError, match="invalid GPU resource"): + parse_gpu_counts(output) + + +@pytest.mark.parametrize("state", ("", "CANCELLED by root")) +def test_parse_state_rejects_invalid_spellings(state: str) -> None: + with pytest.raises(SlurmCommandOutputError): + parse_state(state) + + +def _make_queue_entry(array_task_id: int, state: SchedulerState) -> SlurmQueueEntry: + return SlurmQueueEntry( + job_identity=SchedulerIdentity(array_job_id=4101, array_task_id=array_task_id), + state=state, + ) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py new file mode 100644 index 000000000..91757b78a --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +import shlex +import subprocess +import tarfile +from pathlib import Path +from typing import Literal + +import pytest + +from data_designer.slurm.config import SchedulerProfile, injected_profile +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +from data_designer.slurm.planning import ResolvedSlurmRunPlan, ResolvedSubmission + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" + + +@pytest.mark.parametrize( + ("fixture_name", "plan_fixture"), + (("single_node.sbatch", "single_node_plan"), ("multi_node.sbatch", "multi_node_plan")), +) +def test_renderer_matches_contract_bound_goldens( + fixture_name: str, + plan_fixture: str, + request: pytest.FixtureRequest, +) -> None: + plan = request.getfixturevalue(plan_fixture) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert script == (GOLDEN_DIRECTORY / fixture_name).read_text() + assert script.index("trap 'exit 143' TERM") < script.index("\ndd_initialize_local_scratch\n") + + +def test_renderer_omits_unspecified_array_throttle(multi_node_plan: ResolvedSlurmRunPlan) -> None: + array_tasks = multi_node_plan.array_tasks.model_copy(update={"max_concurrent": None}) + plan = multi_node_plan.model_copy(update={"array_tasks": array_tasks}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert "#SBATCH --array=0-1\n" in script + assert "#SBATCH --array=0-1%" not in script + + +def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fields( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": SchedulerProfile(account="research", partition="batch"), + } + ) + plan = single_node_plan.model_copy( + update={ + "selected_profile": injected_profile(profile), + "submission": ResolvedSubmission( + job_name="data-designer", + account=None, + partition=None, + time_limit="01:00:00", + comment="safe test run", + ), + } + ) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert "#SBATCH --gres=" not in script + assert "#SBATCH --account=" not in script + assert "#SBATCH --partition=" not in script + assert "#SBATCH --exclusive\n" in script + assert '#SBATCH --comment="safe test run"\n' in script + + +def test_renderer_emits_multiple_gpu_partitions(single_node_plan: ResolvedSlurmRunPlan) -> None: + submission = ResolvedSubmission( + job_name="data-designer", + account="research", + partition=["gpu-a", "gpu-b"], + time_limit="01:00:00", + ) + plan = single_node_plan.model_copy(update={"submission": submission}) + + assert "#SBATCH --partition=gpu-a,gpu-b\n" in render_generation_attempt_script(plan, attempt_ordinal=1) + + +def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlurmRunPlan) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={"scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G")} + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + + assert "#SBATCH --mem-per-gpu=80G\n" in render_generation_attempt_script(plan, attempt_ordinal=1) + + +def test_renderer_uses_profile_slurm_bin_path(single_node_plan: ResolvedSlurmRunPlan) -> None: + scheduler = single_node_plan.selected_profile.profile.scheduler.model_copy(update={"bin_path": "/opt/slurm/bin"}) + profile = single_node_plan.selected_profile.profile.model_copy(update={"scheduler": scheduler}) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in script + assert "#SBATCH --exclusive" not in script + + +def test_renderer_uses_discovered_slurm_bin_path_without_profile_override( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = single_node_plan.model_copy(update={"scheduler_bin_path": "/shared/slurm/bin"}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert plan.selected_profile.profile.scheduler.bin_path is None + assert 'export PATH="/shared/slurm/bin:/usr/local/sbin:' in script + assert ResolvedSlurmRunPlan.model_validate_json(plan.serialize_json()).scheduler_bin_path == "/shared/slurm/bin" + + +@pytest.mark.parametrize( + ("gpu_request_mode", "client_cpus", "expected_directive"), + ( + ("gres", 4, "#SBATCH --cpus-per-gpu=1\n"), + ("gres", 16, "#SBATCH --cpus-per-gpu=2\n"), + ("gres", 17, "#SBATCH --cpus-per-gpu=3\n"), + ("visible", 17, "#SBATCH --cpus-per-task=17\n"), + ), +) +def test_renderer_reserves_client_cpus_without_restricting_gpu_socket_visibility( + single_node_plan: ResolvedSlurmRunPlan, + gpu_request_mode: Literal["gres", "visible"], + client_cpus: int, + expected_directive: str, +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy(update={"gpu_request_mode": gpu_request_mode}) + client = single_node_plan.client.model_copy( + update={"authored": single_node_plan.client.authored.model_copy(update={"cpus": client_cpus})} + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile), "client": client}) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert expected_directive in script + assert script.count("#SBATCH --cpus-") == 1 + + +def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + profile = single_node_plan.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), + } + ) + selected_profile = single_node_plan.selected_profile.model_copy(update={"profile": profile}) + plan = single_node_plan.model_copy(update={"selected_profile": selected_profile}) + + with pytest.raises(SlurmBatchRenderError, match="requires GRES"): + render_generation_attempt_script(plan, attempt_ordinal=1) + + +def test_renderer_escapes_shell_expansion_in_structured_paths( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = single_node_plan.model_copy( + update={ + "runtime_bundle": ArtifactReference( + path='/workspace/runtime/$(touch owned)-`whoami`-"bundle".tar.gz', + sha256="e" * 64, + ) + } + ) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + assert ( + 'readonly DD_RUNTIME_ARCHIVE="/workspace/runtime/\\$(touch owned)-\\`whoami\\`-\\"bundle\\".tar.gz"' in script + ) + completed = subprocess.run(("bash", "-n"), input=script, capture_output=True, text=True, check=False) + assert completed.returncode == 0, completed.stderr + + +def test_renderer_keeps_user_text_on_one_non_executable_directive( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + comment = '$(touch owned)" --output=/tmp/owned; `whoami`' + plan = single_node_plan.model_copy( + update={"submission": single_node_plan.submission.model_copy(update={"comment": comment})} + ) + + script = render_generation_attempt_script(plan, attempt_ordinal=1) + + comment_lines = [line for line in script.splitlines() if line.startswith("#SBATCH --comment=")] + assert len(comment_lines) == 1 + assert shlex.split(comment_lines[0].removeprefix("#SBATCH ")) == [f"--comment={comment}"] + assert "\\$" not in comment_lines[0] + assert "\\`" not in comment_lines[0] + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +@pytest.mark.parametrize( + "attempt_ordinal", + (0, -1, True), +) +def test_renderer_rejects_invalid_attempt_ordinals( + single_node_plan: ResolvedSlurmRunPlan, + attempt_ordinal: object, +) -> None: + with pytest.raises(SlurmBatchRenderError, match="positive integer"): + render_generation_attempt_script(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] + + +def test_renderer_rejects_control_characters_from_unvalidated_plan_copies( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = single_node_plan.model_copy( + update={"submission": single_node_plan.submission.model_copy(update={"comment": "unsafe\ntext"})} + ) + + with pytest.raises(SlurmBatchRenderError, match="control characters"): + render_generation_attempt_script(plan, attempt_ordinal=1) + + +def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) -> None: + script = render_generation_attempt_script(single_node_plan, attempt_ordinal=12) + + assert script.count("dd_slurm_run_allocation") == 1 + assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script + assert len(script.splitlines()) <= 49 + assert script.endswith("\n") + + +def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_runtime( + single_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, +) -> None: + run_root = tmp_path / "run" + run_root.mkdir() + (run_root / "shards/shard-00000/attempts/attempt-0001").mkdir(parents=True) + scratch_parent = tmp_path / "scratch" + scratch_parent.mkdir() + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_srun = fake_bin / "srun" + fake_srun.write_text( + '#!/usr/bin/env bash\nset -Eeuo pipefail\nwhile [[ $# -gt 0 && $1 != -- ]]; do shift; done\nshift\nexec "$@"\n' + ) + fake_srun.chmod(0o700) + captured_plan_path = tmp_path / "captured-plan.json" + entrypoint_path = tmp_path / "entrypoint.sh" + entrypoint_path.write_text( + f'dd_slurm_run_allocation() {{\n cp -- "$1" {shlex.quote(str(captured_plan_path))}\n}}\n' + ) + runtime_archive_path = tmp_path / "runtime.tar.gz" + scratch_path = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/scratch.sh" + with tarfile.open(runtime_archive_path, mode="w:gz") as runtime_archive: + runtime_archive.add(scratch_path, arcname="scratch.sh") + runtime_archive.add(entrypoint_path, arcname="entrypoint.sh") + + scheduler = single_node_plan.selected_profile.profile.scheduler.model_copy(update={"bin_path": fake_bin.as_posix()}) + profile = single_node_plan.selected_profile.profile.model_copy(update={"scheduler": scheduler}) + plan = single_node_plan.model_copy( + update={ + "authored_config": ArtifactReference( + path=str(run_root / "authored-config.json"), + sha256="a" * 64, + ), + "runtime_bundle": ArtifactReference( + path=str(runtime_archive_path), + sha256=hashlib.sha256(runtime_archive_path.read_bytes()).hexdigest(), + ), + "selected_profile": injected_profile(profile), + } + ) + plan_path = run_root / "resolved-plan.json" + serialized_plan = plan.serialize_json() + serialized_plan_bytes = serialized_plan.encode() + plan_path.write_bytes(serialized_plan_bytes) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + environment = { + **os.environ, + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_ID": "5101", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_TMPDIR": scratch_parent.as_posix(), + } + + valid = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert valid.returncode == 0 + assert captured_plan_path.read_bytes() == serialized_plan_bytes + + captured_plan_path.unlink() + plan_path.write_bytes(serialized_plan_bytes + b" ") + tampered = subprocess.run( + ("bash",), + input=script, + env=environment, + text=True, + check=False, + ) + + assert tampered.returncode != 0 + assert not captured_plan_path.exists() diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py new file mode 100644 index 000000000..98c74b636 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import subprocess +from collections.abc import Mapping, Sequence + +import pytest + +from data_designer.slurm.launcher.runner import SubprocessRunner + + +def test_subprocess_runner_uses_argv_and_only_explicit_environment(monkeypatch: pytest.MonkeyPatch) -> None: + observed: dict[str, object] = {} + + def fake_run( + command: Sequence[str], + *, + check: bool, + stdin: int, + capture_output: bool, + text: bool, + encoding: str, + errors: str, + env: Mapping[str, str], + timeout: float, + ) -> subprocess.CompletedProcess[str]: + observed.update( + command=command, + check=check, + stdin=stdin, + capture_output=capture_output, + text=text, + encoding=encoding, + errors=errors, + env=env, + timeout=timeout, + ) + return subprocess.CompletedProcess(command, 0, stdout="ok\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + supplied_environment = {"PATH": "/usr/bin", "LC_ALL": "fr_FR.UTF-8"} + runner = SubprocessRunner(environment=supplied_environment, timeout_seconds=4.0) + supplied_environment["SECRET"] = "must-not-leak" + + completed = runner.run(("squeue", "--noheader")) + + assert completed.stdout == "ok\n" + assert observed == { + "command": ("squeue", "--noheader"), + "check": False, + "stdin": subprocess.DEVNULL, + "capture_output": True, + "text": True, + "encoding": "utf-8", + "errors": "replace", + "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, + "timeout": 4.0, + } + + +def test_subprocess_runner_default_environment_forwards_only_slurm_bootstrap(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "/workspace/slurm/bin:/usr/bin") + monkeypatch.setenv("SLURM_CONF", "/etc/slurm/slurm.conf") + monkeypatch.setenv("SECRET", "must-not-leak") + + runner = SubprocessRunner() + + assert runner.environment == { + "LC_ALL": "C", + "PATH": "/workspace/slurm/bin:/usr/bin", + "SLURM_CONF": "/etc/slurm/slurm.conf", + } + + +def test_subprocess_runner_forwards_explicit_standard_input(monkeypatch: pytest.MonkeyPatch) -> None: + observed: dict[str, object] = {} + + def fake_run( + command: Sequence[str], + *, + check: bool, + input: str, + capture_output: bool, + text: bool, + encoding: str, + errors: str, + env: Mapping[str, str], + timeout: float, + ) -> subprocess.CompletedProcess[str]: + observed.update(command=command, check=check, input=input, env=env, timeout=timeout) + assert capture_output is text is True + assert encoding == "utf-8" + assert errors == "replace" + return subprocess.CompletedProcess(command, 0, stdout="5101\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + completed = SubprocessRunner(environment={"PATH": "/usr/bin"}).run( + ("sbatch", "--parsable"), + input_text="#!/bin/sh\n", + ) + + assert completed.stdout == "5101\n" + assert observed == { + "command": ("sbatch", "--parsable"), + "check": False, + "input": "#!/bin/sh\n", + "env": {"LC_ALL": "C", "PATH": "/usr/bin"}, + "timeout": 30.0, + } + + +def test_subprocess_runner_default_environment_replaces_empty_search_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "") + monkeypatch.delenv("SLURM_CONF", raising=False) + + runner = SubprocessRunner() + + assert runner.environment == {"LC_ALL": "C", "PATH": os.defpath} + + +def test_subprocess_runner_environment_is_immutable() -> None: + runner = SubprocessRunner() + + with pytest.raises(TypeError): + runner.environment["SECRET"] = "value" # type: ignore[index] + + +@pytest.mark.parametrize("timeout_seconds", [0, -1, True, "30", float("nan"), float("inf")]) +def test_subprocess_runner_rejects_invalid_timeout(timeout_seconds: object) -> None: + with pytest.raises(ValueError, match="finite positive"): + SubprocessRunner(timeout_seconds=timeout_seconds) # type: ignore[arg-type] + + +@pytest.mark.parametrize("environment", ({"BAD=NAME": "value"}, {"NAME": "bad\0value"})) +def test_subprocess_runner_rejects_invalid_environment(environment: dict[str, str]) -> None: + with pytest.raises(ValueError, match="environment"): + SubprocessRunner(environment=environment) diff --git a/packages/data-designer-slurm/tests/launcher/test_scheduler_path.py b/packages/data-designer-slurm/tests/launcher/test_scheduler_path.py new file mode 100644 index 000000000..415f66221 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_scheduler_path.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 os +from pathlib import Path + +import pytest + +from data_designer.slurm.launcher.scheduler_path import SchedulerPathError, resolve_scheduler_bin_path + + +def _install_fake_slurm_tools(directory: Path) -> None: + directory.mkdir() + for name in ("srun", "scontrol"): + executable = directory / name + executable.write_text("#!/bin/sh\nexit 0\n") + executable.chmod(0o755) + + +def test_scheduler_directory_is_discovered_from_submit_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / "slurm-bin" + _install_fake_slurm_tools(directory) + monkeypatch.setenv("PATH", str(directory)) + + assert resolve_scheduler_bin_path(None) == str(directory) + + +def test_scheduler_directory_preserves_command_symlinks(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / "slurm-bin" + directory.mkdir() + srun_target = tmp_path / "srun-target" + scontrol_target = tmp_path / "scontrol-target" + srun_target.write_text("#!/bin/sh\nexit 0\n") + scontrol_target.write_text("#!/bin/sh\nexit 0\n") + srun_target.chmod(0o755) + scontrol_target.chmod(0o755) + (directory / "srun").symlink_to(srun_target) + (directory / "scontrol").symlink_to(scontrol_target) + monkeypatch.setenv("PATH", str(directory)) + + assert resolve_scheduler_bin_path(None) == str(directory) + + +def test_relative_submit_path_is_normalized(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / "slurm-bin" + _install_fake_slurm_tools(directory) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("PATH", "./slurm-bin") + + assert resolve_scheduler_bin_path(None) == str(directory) + + +def test_explicit_override_works_without_submit_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / "slurm-bin" + _install_fake_slurm_tools(directory) + monkeypatch.setenv("PATH", os.defpath) + + assert resolve_scheduler_bin_path(str(directory)) == str(directory) + + +def test_missing_srun_reports_actionable_error(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", str(tmp_path)) + + with pytest.raises(SchedulerPathError, match="srun was not found"): + resolve_scheduler_bin_path(None) + + +def test_missing_scontrol_fails_before_submission(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + directory = tmp_path / "slurm-bin" + _install_fake_slurm_tools(directory) + (directory / "scontrol").unlink() + monkeypatch.setenv("PATH", str(directory)) + + with pytest.raises(SchedulerPathError, match="runnable srun and scontrol"): + resolve_scheduler_bin_path(None) diff --git a/packages/data-designer-slurm/tests/planning/test_colocation.py b/packages/data-designer-slurm/tests/planning/test_colocation.py new file mode 100644 index 000000000..515093ede --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_colocation.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from data_designer.config import DataDesignerConfig +from data_designer.slurm.config import ( + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + injected_profile, + load_run_config, +) +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.errors import SlurmConfigResolutionError +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.runtime.ports import resolve_allocation_plan +from data_designer.slurm.serving.resolver import resolve_vllm_server + + +@pytest.mark.parametrize("suffix", ("yaml", "json")) +def test_shared_node_authoring_round_trip( + shared_node_config: DataDesignerSlurmConfig, tmp_path: Path, suffix: str +) -> None: + builder = DataDesignerSlurmConfigBuilder(shared_node_config.builder).with_invocation( + num_records=4, dataset_name="judged", model_concurrency={"writer": 1, "judge": 1} + ) + for deployment in shared_node_config.deployments: + builder.with_deployment(deployment) + path = tmp_path / f"run.{suffix}" + builder.write_config(path) + assert load_run_config(path) == builder.build() == shared_node_config + + +@pytest.mark.parametrize( + ("resources", "topology", "message"), + [ + ({"shared_node": "models"}, {}, "provided together"), + ({"gpus_per_node": 4}, {}, "provided together"), + ({"shared_node": "models", "gpus_per_node": 0}, {}, "allowed minimum"), + ({"shared_node": "models", "gpus_per_node": -1}, {}, "allowed minimum"), + ({"shared_node": "models", "gpus_per_node": 4, "nodes": 2}, {}, "exactly one"), + ({"shared_node": "models", "gpus_per_node": 4}, {"nodes_per_replica": 2}, "single-node"), + ({"shared_node": "models", "gpus_per_node": 3}, {"tensor_parallel": 2}, "GPU quota"), + ({"gpu_start_index": 0}, {}, "requires shared_node"), + ({"shared_node": "models", "gpus_per_node": 1, "gpu_start_index": -1}, {}, "allowed minimum"), + ], +) +def test_shared_node_builder_rejects_invalid_resources( + shared_node_config: DataDesignerSlurmConfig, resources: dict[str, object], topology: dict[str, int], message: str +) -> None: + deployment = shared_node_config.deployments[0].model_dump(mode="python") + deployment["resources"] = resources + deployment["topology"].update(topology) + with pytest.raises(ValueError, match=message): + DataDesignerSlurmConfigBuilder(shared_node_config.builder).with_deployment(deployment) + + +@pytest.mark.parametrize("mismatch", ("singleton", "disjoint", "environment", "queue_backpressure")) +def test_shared_node_rejects_incompatible_groups(shared_node_config: DataDesignerSlurmConfig, mismatch: str) -> None: + payload = shared_node_config.model_dump(mode="python") + if mismatch == "singleton": + payload["deployments"][1]["resources"]["shared_node"] = "other" + elif mismatch == "disjoint": + payload["deployments"][1]["resources"] = {"nodes": 1} + elif mismatch == "environment": + payload["deployments"][1]["server"]["environment"] = {"MODE": {"type": "literal", "value": "other"}} + else: + payload["deployments"][1]["server"]["queue_backpressure"]["max_waiting_requests"] = 2 + with pytest.raises(ValidationError, match="shared_node"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_shared_node_plan_assigns_distinct_models_gpus_ports_and_one_node( + shared_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = ResolvedSlurmRunPlan.model_validate_json(shared_node_plan.serialize_json()) + assert plan.compute_sha256() == shared_node_plan.compute_sha256() + assert [deployment.node_indices for deployment in plan.deployments] == [(0,), (0,)] + assert [deployment.gpu_start_index for deployment in plan.deployments] == [0, 4] + servers = tuple(resolve_vllm_server(plan, deployment.deployment_id) for deployment in plan.deployments) + assert len({server.model for server in servers}) == 2 + assert [process.gpu_indices for server in servers for process in server.processes] == [ + (0, 1), + (2, 3), + (4, 5), + (6, 7), + ] + assert [backend.port for server in servers for backend in server.backend_endpoints] == [18000, 18001, 18002, 18003] + assert [server.logical_endpoint.port for server in servers] == [17000, 17001] + script = render_generation_attempt_script(plan, attempt_ordinal=1) + assert "#SBATCH --nodes=1\n" in script + assert "#SBATCH --gres=gpu:8\n" in script + assert "#SBATCH --exclusive" not in script + + +@pytest.mark.parametrize( + ("mutation", "message"), + ( + ("overlap", "GPU ownership"), + ("oversubscription", "quotas exceed"), + ("node", "placement and GPU ownership"), + ("image", "same resolved serving image"), + ("visible", "requires GRES"), + ("environment", "matching server environment"), + ("queue", "matching server environment"), + ), +) +def test_shared_node_plan_rejects_invalid_ownership( + shared_node_plan: ResolvedSlurmRunPlan, mutation: str, message: str +) -> None: + payload = shared_node_plan.model_dump(mode="python") + if mutation == "overlap": + payload["deployments"][1]["gpu_start_index"] = 0 + elif mutation == "oversubscription": + payload["deployments"][1]["authored"]["resources"]["gpus_per_node"] = 8 + elif mutation == "node": + payload["deployments"][1]["node_indices"] = (1,) + for port in payload["deployments"][1]["ports"]: + port["node_index"] = 1 + elif mutation == "image": + payload["deployments"][1]["image"]["path"] = "/images/other.sqsh" + elif mutation == "visible": + profile = shared_node_plan.selected_profile.profile + visible_profile = profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": profile.scheduler.model_copy(update={"mem_per_gpu": None}), + } + ) + payload["selected_profile"] = injected_profile(visible_profile).model_dump(mode="python") + elif mutation == "environment": + payload["deployments"][1]["authored"]["server"]["environment"] = {"MODE": {"type": "literal", "value": "other"}} + else: + payload["deployments"][1]["authored"]["server"]["queue_backpressure"]["max_waiting_requests"] = 2 + with pytest.raises(ValidationError, match=message): + ResolvedSlurmRunPlan.model_validate(payload) + + +def test_documented_colocated_example_validates_both_public_configs() -> None: + path = Path(__file__).parents[2] / "examples/colocated-models.yaml" + config = load_run_config(path) + builder = DataDesignerConfig.model_validate(config.builder.inline["data_designer"]) + assert [model.alias for model in builder.model_configs] == ["writer", "judge"] + assert [column.name for column in builder.columns] == ["answer", "review"] + assert len({model.model for model in builder.model_configs}) == 2 + + +def test_shared_node_runtime_ports_are_distinct_and_isolated(shared_node_plan: ResolvedSlurmRunPlan) -> None: + first = resolve_allocation_plan(shared_node_plan, {"SLURM_JOB_GPUS": "1,3,5,7"}) + second = resolve_allocation_plan(shared_node_plan, {"SLURM_JOB_GPUS": "0,2,4,6"}) + first_ports = [port.port for port in first.client.ports] + [ + port.port for deployment in first.deployments for port in deployment.ports + ] + second_ports = [port.port for port in second.client.ports] + [ + port.port for deployment in second.deployments for port in deployment.ports + ] + assert len(set(first_ports)) == 6 + assert not set(first_ports).intersection(second_ports) + assert [deployment.node_indices for deployment in first.deployments] == [(0,), (0,)] + assert first.authored_config == shared_node_plan.authored_config + + +@pytest.mark.parametrize("invalid", ("oversubscription", "offset", "visible", "image")) +def test_preparation_rejects_unsupported_shared_placement( + shared_node_config: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, + invalid: str, +) -> None: + payload = shared_node_config.model_dump(mode="python") + profile = single_node_plan.selected_profile.profile + images = (single_node_plan.deployments[0].image,) * 2 + if invalid == "oversubscription": + payload["deployments"][1]["resources"]["gpus_per_node"] = 8 + elif invalid == "offset": + payload["deployments"][1]["resources"]["gpu_start_index"] = 8 + elif invalid == "visible": + profile = profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": profile.scheduler.model_copy(update={"mem_per_gpu": None}), + } + ) + else: + images = (images[0], images[1].model_copy(update={"path": "/images/other.sqsh"})) + with pytest.raises(SlurmConfigResolutionError, match="shared_node"): + resolve_slurm_config( + DataDesignerSlurmConfig.model_validate(payload), + selected_profile=injected_profile(profile), + client_runtime=single_node_plan.client.runtime, + deployment_images=images, + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=single_node_plan.run_id, + package_version=single_node_plan.package_version, + ) + + +@pytest.mark.parametrize( + ("quotas", "starts", "budget", "expected_starts"), + [ + ((1, 2), (None, None), 3, [0, 1]), + ((1, 1), (0, 0), 1, [0, 0]), + ((2, 1), (0, 1), 2, [0, 1]), + ((1, 1), (1, None), 3, [1, 2]), + ], +) +def test_partial_gpu_allocation_uses_per_model_quota_for_replica_count( + shared_node_config: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, + quotas: tuple[int, int], + starts: tuple[int | None, int | None], + budget: int, + expected_starts: list[int], +) -> None: + payload = shared_node_config.model_dump(mode="python") + for deployment, quota, start in zip(payload["deployments"], quotas, starts, strict=True): + deployment["resources"]["gpus_per_node"] = quota + deployment["resources"]["gpu_start_index"] = start + deployment["topology"]["tensor_parallel"] = quota + profile = single_node_plan.selected_profile.profile.model_copy(update={"gpus_per_node": budget}) + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + DataDesignerSlurmConfig.model_validate(payload), + selected_profile=injected_profile(profile), + client_runtime=single_node_plan.client.runtime, + deployment_images=(single_node_plan.deployments[0].image,) * 2, + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=single_node_plan.run_id, + package_version=single_node_plan.package_version, + ) + ) + assert [deployment.topology.replica_count for deployment in plan.deployments] == [1, 1] + assert [deployment.gpu_start_index for deployment in plan.deployments] == expected_starts + assert [deployment.gpus_per_node for deployment in plan.deployments] == list(quotas) + assert ResolvedSlurmRunPlan.model_validate_json(plan.serialize_json()) == plan + script = render_generation_attempt_script(plan, attempt_ordinal=1) + assert "#SBATCH --nodes=1\n" in script and f"#SBATCH --gres=gpu:{budget}\n" in script + assert "#SBATCH --exclusive" not in script diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py new file mode 100644 index 000000000..2d58d9908 --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -0,0 +1,802 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path + +import pytest + +from data_designer.config import ( + DataDesignerConfigBuilder, + DropColumnsProcessorConfig, + JudgeScoreProfilerConfig, + LLMTextColumnConfig, + ModelConfig, +) +from data_designer.slurm.config import ( + BuilderInput, + DataDesignerSlurmConfig, + InputBindings, + SlurmProfileCatalog, + select_profile, +) +from data_designer.slurm.contracts import compute_serialized_json_sha256 +from data_designer.slurm.planning import ( + ArtifactReference, + ResolvedDependencyLock, + ResolvedSlurmRunPlan, +) +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.errors import ( + SlurmConfigResolutionError, + SlurmPlanCompilationError, + SlurmPlanContractError, +) +from data_designer.slurm.planning.resolution import EffectiveDataDesignerSlurmConfig, resolve_slurm_config + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" + + +def _resolve_fixture( + authored: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + expected: ResolvedSlurmRunPlan, + **updates: object, +) -> EffectiveDataDesignerSlurmConfig: + values = { + "selected_profile": expected.selected_profile, + "client_runtime": expected.client.runtime, + "deployment_images": tuple(deployment.image for deployment in expected.deployments), + "dependency_lock": dependency_lock, + "runtime_bundle": expected.runtime_bundle, + "run_id": expected.run_id, + "package_version": expected.package_version, + "resolved_gpus_per_node": expected.resolved_gpus_per_node, + } + values.update(updates) + return resolve_slurm_config(authored, **values) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("authored_fixture", "lock_fixture", "plan_fixture", "golden_name"), + [ + ("authored_run_single", "dependency_lock_single", "single_node_plan", "single_node_plan.json"), + ("authored_run", "dependency_lock", "multi_node_plan", "multi_node_plan.json"), + ], +) +def test_compiler_reproduces_plan_goldens_byte_for_byte( + request: pytest.FixtureRequest, + authored_fixture: str, + lock_fixture: str, + plan_fixture: str, + golden_name: str, +) -> None: + authored = request.getfixturevalue(authored_fixture) + dependency_lock = request.getfixturevalue(lock_fixture) + expected = request.getfixturevalue(plan_fixture) + effective = _resolve_fixture(authored, dependency_lock, expected) + + first = SlurmRunCompiler.compile(effective) + second = SlurmRunCompiler.compile(effective) + + assert first.serialize_json() == (GOLDEN_DIRECTORY / golden_name).read_text() + assert first.serialize_canonical_json() == second.serialize_canonical_json() + assert first.compute_sha256() == second.compute_sha256() == expected.compute_sha256() + + +def test_compiler_resolves_explicit_hostname_and_default_profile_selection( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + profile_catalog: SlurmProfileCatalog, +) -> None: + selections = ( + select_profile(profile_catalog, cluster="primary", hostnames=("lab-login-1",)), + select_profile(profile_catalog, hostnames=("PRIMARY-LOGIN-1",)), + select_profile(profile_catalog, hostnames=("unmatched",)), + ) + + plans = tuple( + SlurmRunCompiler.compile( + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + ) + ) + for selected in selections + ) + + assert {plan.selected_profile.selection_source.value for plan in plans} == { + "explicit", + "hostname", + "default", + } + assert all(plan.deployments[0].node_indices == (0,) for plan in plans) + assert all( + plan.invocation.effective_input_bindings.managed_assets_path == "/workspace/primary/managed-assets" + for plan in plans + ) + assert all(plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 4 for plan in plans) + + +def test_auto_gpu_resolution_is_explicit_and_scheduler_free( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + profile_catalog: SlurmProfileCatalog, +) -> None: + selected = select_profile(profile_catalog, cluster="lab") + runtime_bundle = single_node_plan.runtime_bundle.model_copy( + update={"path": f"/workspace/lab/runtime/{single_node_plan.runtime_bundle.sha256}.tar.gz"} + ) + + with pytest.raises(SlurmConfigResolutionError, match="auto gpus_per_node"): + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + runtime_bundle=runtime_bundle, + resolved_gpus_per_node=None, + ) + + plan = SlurmRunCompiler.compile( + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + selected_profile=selected, + runtime_bundle=runtime_bundle, + resolved_gpus_per_node=8, + ) + ) + + assert plan.resolved_gpus_per_node == 8 + assert plan.submission.account == "lab" + + +def test_compiler_preserves_explicit_compatibility_run_values( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = { + "buffer_size": 2048, + "disable_early_shutdown": False, + "max_conversation_restarts": 3, + "otel_metrics_port": 24000, + "shutdown_error_rate": 0.25, + "non_inference_max_parallel_workers": 7, + } + payload["invocation"]["input_bindings"]["managed_assets_path"] = "/workspace/explicit-assets" + authored = DataDesignerSlurmConfig.model_validate(payload) + + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + assert effective.invocation.effective_run_config["buffer_size"] == 2048 + assert effective.invocation.effective_run_config["disable_early_shutdown"] is False + assert effective.invocation.effective_run_config["max_conversation_restarts"] == 3 + assert effective.invocation.effective_run_config["otel_metrics_port"] == 24000 + assert effective.invocation.effective_run_config["shutdown_error_rate"] == 0.25 + assert effective.invocation.effective_run_config["non_inference_max_parallel_workers"] == 7 + assert effective.invocation.effective_input_bindings.managed_assets_path == "/workspace/explicit-assets" + + +@pytest.mark.parametrize( + ("run_config", "expected_rate", "expected_window"), + [ + ({"shutdown_error_rate": 0.25}, 0.25, 10), + ({"shutdown_error_window": 25}, 0.5, 25), + ({"disable_early_shutdown": False}, 0.5, 10), + ({"disable_early_shutdown": False, "shutdown_error_window": 25}, 0.5, 25), + ], +) +def test_compiler_preserves_partial_early_shutdown_configuration( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + run_config: dict[str, object], + expected_rate: float, + expected_window: int, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = run_config + authored = DataDesignerSlurmConfig.model_validate(payload) + + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + assert effective.invocation.effective_run_config["disable_early_shutdown"] is False + assert effective.invocation.effective_run_config["shutdown_error_rate"] == expected_rate + assert effective.invocation.effective_run_config["shutdown_error_window"] == expected_window + + +def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["deployments"][0]["topology"]["tensor_parallel"] = 3 + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmPlanCompilationError, match="tensor_parallel"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +@pytest.mark.parametrize( + "managed_assets_path", + [ + "/workspace/primary", + "/workspace/primary/images/assets", + "/workspace/primary/runs", + ], +) +def test_resolution_rejects_managed_assets_overlapping_workspace_state( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + managed_assets_path: str, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["input_bindings"]["managed_assets_path"] = managed_assets_path + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match="managed_assets_path"): + _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + +def test_compiler_rejects_direct_effective_input_resolution_bypass( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run, dependency_lock, multi_node_plan) + + with pytest.raises(SlurmPlanCompilationError, match="runtime bundle"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "runtime_bundle": ArtifactReference(path="/tmp/runtime.tar.gz", sha256="e" * 64), + } + ) + ) + + profile = effective.selected_profile.profile.model_copy( + update={ + "gpu_request_mode": "visible", + "scheduler": effective.selected_profile.profile.scheduler.model_copy(update={"mem_per_gpu": "80G"}), + } + ) + with pytest.raises(SlurmPlanCompilationError, match="requires GRES"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"selected_profile": effective.selected_profile.model_copy(update={"profile": profile})} + ) + ) + + +def test_compiler_rejects_direct_effective_record_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + run_config = dict(effective.invocation.effective_run_config) + run_config["jinja_rendering_engine"] = "native" + + with pytest.raises(SlurmPlanCompilationError, match="resolved invocation"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"invocation": effective.invocation.model_copy(update={"effective_run_config": run_config})} + ) + ) + bindings = effective.invocation.effective_input_bindings.model_copy( + update={"managed_assets_path": "/workspace/other-assets"} + ) + with pytest.raises(SlurmPlanCompilationError, match="resolved invocation"): + SlurmRunCompiler.compile( + effective.model_copy( + update={"invocation": effective.invocation.model_copy(update={"effective_input_bindings": bindings})} + ) + ) + + +def test_compiler_rejects_model_alias_missing_from_builder( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"][0]["alias"] = "other" + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmPlanCompilationError, match="exactly cover"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +def test_compiler_rejects_builder_model_alias_without_deployment( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + sourced: bool, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"].append( + {"alias": "undeployed", "model": "example/undeployed", "provider": "openai"} + ) + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmPlanCompilationError, match="exactly cover"): + SlurmRunCompiler.compile( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + +def test_compiler_rejects_inline_builder_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + + with pytest.raises(SlurmPlanCompilationError, match="builder digest"): + SlurmRunCompiler.compile( + effective.model_copy(update={"builder": effective.builder.model_copy(update={"content_sha256": "a" * 64})}) + ) + + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["model_configs"].append( + {"alias": "undeployed", "model": "example/undeployed", "provider": "openai"} + ) + authored = DataDesignerSlurmConfig.model_validate(payload) + effective = _resolve_fixture(authored, dependency_lock_single, single_node_plan) + forged_builder = effective.builder.model_copy(update={"model_aliases": ("generator",)}) + + with pytest.raises(SlurmPlanCompilationError, match="model aliases"): + SlurmRunCompiler.compile(effective.model_copy(update={"builder": forged_builder})) + + +def test_compiler_rejects_resolved_image_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + deployment_images = ( + effective.deployment_images[0].model_copy(update={"sha256": "a" * 64}), + *effective.deployment_images[1:], + ) + with pytest.raises(SlurmPlanCompilationError, match="inspection record"): + SlurmRunCompiler.compile(effective.model_copy(update={"deployment_images": deployment_images})) + + +@pytest.mark.parametrize("invalid_field", ["client_runtime", "runtime_bundle"]) +def test_compiler_revalidates_nested_effective_contracts( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + invalid_field: str, +) -> None: + effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) + if invalid_field == "client_runtime": + invalid_value = effective.client_runtime.model_copy(update={"python_executable": "relative/python"}) + else: + invalid_value = effective.runtime_bundle.model_copy( + update={"path": f"/workspace/primary/runtime/../{effective.runtime_bundle.sha256}.tar.gz"} + ) + + with pytest.raises(SlurmPlanCompilationError, match="failed validation"): + SlurmRunCompiler.compile(effective.model_copy(update={invalid_field: invalid_value})) + + +def test_compiler_rejects_otel_collision_before_runtime( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["run_config"] = {"otel_metrics_port": 17000} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmPlanCompilationError, match="OTEL"): + SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + +def test_sharded_seed_inputs_have_stable_ranges_and_partition_digests( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invocation = authored_run.invocation.model_copy( + update={"input_bindings": InputBindings(seed_path="/datasets/seed.parquet")} + ) + authored = authored_run.model_copy(update={"invocation": invocation}) + + first = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + second = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert [(shard.record_range.start_index, shard.record_range.end_index_exclusive) for shard in first.shards] == [ + (0, 50), + (50, 100), + ] + assert all(shard.input_partition is not None for shard in first.shards) + for shard in first.shards: + assert shard.input_partition is not None + assert shard.input_partition.sha256 == compute_serialized_json_sha256( + { + "record_range": shard.record_range.model_dump(mode="json"), + "seed_path": "/datasets/seed.parquet", + } + ) + assert first.shards == second.shards + + +def test_uneven_shards_assign_remainder_to_the_final_task( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run.model_copy(update={"array_tasks": authored_run.array_tasks.model_copy(update={"count": 3})}) + + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert tuple((shard.record_range.start_index, shard.record_range.end_index_exclusive) for shard in plan.shards) == ( + (0, 33), + (33, 66), + (66, 100), + ) + + +@pytest.mark.parametrize( + ("builder_update", "message"), + [ + ( + { + "seed_config": { + "sampling_strategy": "shuffle", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + "shuffled seed", + ), + ( + { + "seed_config": { + "sampling_strategy": "ordered", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + "seed_path", + ), + ( + { + "seed_config": { + "sampling_strategy": "ordered", + "selection_strategy": {"start": 0, "end": 9}, + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + } + }, + "selection strategies", + ), + ( + { + "columns": [ + { + "column_type": "image", + "model_alias": "generator", + "name": "picture", + "prompt": "an image", + } + ] + }, + "media output", + ), + ], +) +def test_resolution_rejects_unshardable_big_iron_fields( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + builder_update: dict[str, object], + message: str, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"].update(builder_update) + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match=message): + _resolve_fixture(authored, dependency_lock, multi_node_plan) + + +@pytest.mark.parametrize("field", ["processors", "profilers"]) +def test_resolution_rejects_real_global_builder_configs( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + field: str, +) -> None: + builder = DataDesignerConfigBuilder( + model_configs=[ModelConfig(alias="generator", model="example/generator", provider="openai")] + ) + builder.add_column(LLMTextColumnConfig(name="generated", prompt="hello", model_alias="generator")) + if field == "processors": + builder.add_processor(DropColumnsProcessorConfig(name="global", column_names=["generated"])) + else: + builder.add_profiler(JudgeScoreProfilerConfig(model_alias="generator")) + data_designer = builder.get_builder_config().to_dict()["data_designer"] + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"][field] = data_designer[field] + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match=field): + _resolve_fixture(authored, dependency_lock, multi_node_plan) + + +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +@pytest.mark.parametrize( + ("column", "message"), + [ + ({"column_type": "future-column", "name": "future"}, "unknown column semantics"), + ({"column_type": "fake-slurm-column", "name": "plugin"}, "plugin"), + ({"column_type": "custom", "generator_function": "generate", "name": "custom"}, "custom"), + ( + { + "column_type": "validation", + "name": "validate", + "target_columns": ["generated"], + "validator_params": { + "validation_function": "validate", + "validator_type": "local_callable", + }, + "validator_type": "local_callable", + }, + "local callable", + ), + ], +) +def test_resolution_rejects_unportable_multi_shard_columns( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + sourced: bool, + column: dict[str, object], + message: str, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [column] + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match=message): + _resolve_fixture( + authored, + dependency_lock, + multi_node_plan, + builder_payload=builder_payload, + ) + + +def test_sharded_seed_binding_may_override_authored_source( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["seed_config"] = { + "sampling_strategy": "ordered", + "source": {"path": "/datasets/original.parquet", "seed_type": "local"}, + } + payload["invocation"]["input_bindings"]["seed_path"] = "/datasets/override.parquet" + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock, multi_node_plan)) + + assert all(shard.input_partition is not None for shard in plan.shards) + + +def test_single_shard_allows_non_collectable_big_iron_fields( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"].update( + { + "columns": [ + { + "column_type": "image", + "model_alias": "generator", + "name": "picture", + "prompt": "an image", + } + ], + "seed_config": { + "sampling_strategy": "shuffle", + "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, + }, + } + ) + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + assert plan.array_tasks.count == 1 + + +@pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) +def test_single_shard_defers_opaque_plugin_aliases_to_client_preflight( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + sourced: bool, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [ + { + "column_type": "plugin-column", + "fallback_model_alias": None, + "name": "plugin", + } + ] + builder_payload = None + if sourced: + builder_payload = payload["builder"]["inline"] + payload["builder"] = {"source": "builder.json"} + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = SlurmRunCompiler.compile( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + assert plan.builder.model_aliases == ("generator",) + + +def test_sourced_builder_is_resolved_to_one_digest_bound_run_input( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + assert isinstance(builder_payload, dict) + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + + plan = SlurmRunCompiler.compile( + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + ) + + assert plan.builder.authored_source == "builder.json" + assert plan.builder.source is not None + assert plan.builder.source.path == "/workspace/primary/runs/run-single/builder-config.json" + assert plan.builder.content_sha256 == plan.builder.source.sha256 + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("bytes", "builder digest"), + ("aliases", "model aliases"), + ], +) +def test_compiler_rejects_sourced_builder_payload_identity_drift( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + mutation: str, + message: str, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + effective = _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + drifted_payload = deepcopy(builder_payload) + if mutation == "bytes": + drifted_payload["library_version"] = "drifted" + else: + drifted_payload["data_designer"]["model_configs"][0]["alias"] = "drifted" + drifted = effective.model_copy(update={"builder_payload": drifted_payload}) + + with pytest.raises(SlurmPlanContractError, match=message): + SlurmRunCompiler.compile(drifted) + + +def test_resolution_rejects_secret_values_in_sourced_builder_payload( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + builder_payload = authored_run_single.model_dump(mode="json")["builder"]["inline"] + secret = "super-secret-token" + builder_payload["data_designer"]["api_key"] = secret + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + + with pytest.raises(SlurmConfigResolutionError, match="failed validation") as error: + _resolve_fixture( + authored, + dependency_lock_single, + single_node_plan, + builder_payload=builder_payload, + ) + + assert secret not in str(error.value) + assert error.value.__cause__ is None + + +def test_resolution_rejects_dependency_identity_mismatch( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + wrong_lock = dependency_lock_single.model_copy(update={"client_runtime_sha256": "a" * 64}) + + with pytest.raises(SlurmConfigResolutionError, match="client runtime"): + _resolve_fixture(authored_run_single, wrong_lock, single_node_plan) + + +@pytest.mark.parametrize( + "runtime_path", + [ + "/tmp/" + "e" * 64 + ".tar.gz", + "/workspace/primary/runtime/runtime.tar.gz", + "/workspace/primary/runtime/" + "e" * 64 + ".tgz", + ], +) +def test_resolution_rejects_invalid_runtime_bundle_paths( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + runtime_path: str, +) -> None: + wrong_runtime = ArtifactReference(path=runtime_path, sha256="e" * 64) + + with pytest.raises(SlurmConfigResolutionError, match="runtime bundle"): + _resolve_fixture( + authored_run_single, + dependency_lock_single, + single_node_plan, + runtime_bundle=wrong_runtime, + ) + + +def test_plan_contains_secret_references_without_credentials( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = SlurmRunCompiler.compile(_resolve_fixture(authored_run, dependency_lock, multi_node_plan)) + serialized = plan.serialize_json() + + assert "PACKAGE_INDEX_TOKEN" in serialized + assert "HF_TOKEN" in serialized + assert "secret-value" not in serialized diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py new file mode 100644 index 000000000..253f7a771 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Literal, cast + +import pytest + +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStep, RuntimeStepRole +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + SchedulerIdentity, + ShardWinner, + StateNotFoundError, + validate_attempt_transition, + validate_readiness_transition, +) + + +@dataclass(slots=True) +class RuntimeCase: + workspace: Path + context: AllocationContext + created_at: datetime + + +@dataclass(slots=True) +class FakeStateStore: + attempt: AttemptManifest + readiness: list[AttemptReadiness] = field(default_factory=list) + dataset_workspace_lease_active: bool = False + dataset_workspace_modes: list[str] = field(default_factory=list) + winners: list[ShardWinner] = field(default_factory=list) + + def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: + validate_attempt_transition(self.attempt, attempt) + self.attempt = attempt + return attempt + + def publish_attempt_result( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + assert self.dataset_workspace_lease_active + reference = client_result.candidate_output_manifest + assert reference is not None + assert candidate.compute_sha256() == reference.sha256 + assert self.attempt.state is AttemptLifecycleState.RUNNING + bound_attempt = self.attempt.model_copy(update={"candidate_output": reference}) + validate_attempt_transition(self.attempt, bound_attempt) + self.attempt = bound_attempt + return client_result, candidate + + @contextmanager + def acquire_dataset_workspace( + self, + shard_id: str, + attempt_id: str, + resume_mode: Literal["never", "always", "if_possible"], + ) -> Iterator[Path]: + assert shard_id == self.attempt.shard_id + assert attempt_id == self.attempt.attempt_id + assert not self.dataset_workspace_lease_active + self.dataset_workspace_modes.append(resume_mode) + self.dataset_workspace_lease_active = True + try: + yield Path(self.attempt.resolved_plan.path).parent / "dataset" + finally: + self.dataset_workspace_lease_active = False + + def finalize_winner( + self, + shard_id: str, + attempt_id: str, + *, + published_at: datetime, + completed_at: datetime | None = None, + ) -> ShardWinner: + if completed_at is not None: + self.update_attempt( + self.attempt.model_copy( + update={ + "state": AttemptLifecycleState.SUCCEEDED, + "terminal_classification": AttemptTerminalClassification.SUCCEEDED, + "updated_at": completed_at, + } + ) + ) + assert self.attempt.state is AttemptLifecycleState.SUCCEEDED + assert self.attempt.candidate_output is not None + winner = ShardWinner( + schema_version=1, + run_id=self.attempt.run_id, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_ordinal=self.attempt.attempt_ordinal, + candidate_manifest=self.attempt.candidate_output, + published_at=published_at, + ) + self.winners.append(winner) + return winner + + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: + if self.readiness: + validate_readiness_transition(self.readiness[-1], readiness) + else: + assert readiness.revision == 1 + self.readiness.append(readiness) + return readiness + + def load_readiness(self, shard_id: str, attempt_id: str) -> AttemptReadiness: + if not self.readiness: + raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") + readiness = self.readiness[-1] + assert readiness.shard_id == shard_id + assert readiness.attempt_id == attempt_id + return readiness + + def load_attempt(self, shard_id: str, attempt_id: str) -> AttemptManifest: + assert self.attempt.shard_id == shard_id + assert self.attempt.attempt_id == attempt_id + return self.attempt + + +class FakePreflight: + def __init__(self, failure: BaseException | None = None) -> None: + self.failure = failure + self.calls = 0 + + def verify(self, context: AllocationContext, environment: object) -> None: + del context, environment + self.calls += 1 + if self.failure is not None: + raise self.failure + + +class FakeClientStepBuilder: + def __init__(self) -> None: + self.retry_resume_modes: list[Literal["never", "always"] | None] = [] + + def build_preflight_step( + self, + plan: ResolvedSlurmRunPlan, + shard: object, + attempt: object, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: object, + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + del plan, shard, attempt, endpoints, source_environment + self.retry_resume_modes.append(retry_resume_mode) + return _step("client-preflight", RuntimeStepRole.CLIENT_PREFLIGHT, attempt_directory) + + def build_generation_step( + self, + plan: ResolvedSlurmRunPlan, + shard: object, + attempt: object, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: object, + *, + retry_resume_mode: Literal["never", "always"] | None = None, + ) -> RuntimeStep: + del plan, shard, attempt, endpoints, source_environment + self.retry_resume_modes.append(retry_resume_mode) + return _step("client-generation", RuntimeStepRole.CLIENT, attempt_directory) + + +@pytest.fixture +def runtime_case(tmp_path: Path, single_node_plan: ResolvedSlurmRunPlan) -> RuntimeCase: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + plan = relocate_plan(single_node_plan, workspace) + created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) + plan_reference = ArtifactReference( + path=Path(plan.authored_config.path).with_name("resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ) + attempt = AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=plan.shards[0].shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + created_at=created_at, + updated_at=created_at + timedelta(seconds=1), + ) + attempt_directory = ( + Path(plan.authored_config.path).parent / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + ) + attempt_directory.mkdir(parents=True, mode=0o700) + return RuntimeCase( + workspace=workspace, + context=AllocationContext( + plan=plan, + shard=plan.shards[0], + attempt=attempt, + attempt_directory=attempt_directory, + ), + created_at=created_at, + ) + + +def relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: + previous_workspace = plan.selected_profile.profile.workspace_root + payload = cast( + dict[str, object], + json.loads(plan.serialize_json().replace(previous_workspace, workspace.as_posix())), + ) + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected_profile["profile"]) + relocated_mount = { + "source": workspace.as_posix(), + "target": previous_workspace, + "read_only": False, + } + profile_payload["container_mounts"] = [relocated_mount] + payload["container_mounts"] = [relocated_mount] + profile = SlurmProfile.model_validate(profile_payload) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _step(step_id: str, role: RuntimeStepRole, attempt_directory: Path) -> RuntimeStep: + return RuntimeStep( + step_id=step_id, + role=role, + command=("true",), + environment={"PATH": "/usr/bin", "LC_ALL": "C"}, + stdout_path=attempt_directory / "logs" / f"{step_id}.out", + stderr_path=attempt_directory / "logs" / f"{step_id}.err", + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_backpressure.py b/packages/data-designer-slurm/tests/runtime/test_backpressure.py new file mode 100644 index 000000000..e0235d97f --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -0,0 +1,119 @@ +# 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.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + AsgiMessage, + QueueBackpressureController, + QueueBackpressureSettings, + QueueDepthBackpressureMiddleware, +) + + +def test_backpressure_settings_are_strict_and_support_absent_retry_header() -> None: + assert QueueBackpressureSettings.from_environment({}).max_waiting_requests == 128 + settings = QueueBackpressureSettings.from_environment( + { + MAX_WAITING_REQUESTS_ENVIRONMENT: "7", + RETRY_AFTER_SECONDS_ENVIRONMENT: "", + } + ) + assert settings.max_waiting_requests == 7 + assert settings.retry_after_seconds is None + with pytest.raises(ValueError, match="threshold"): + QueueBackpressureSettings.from_environment({MAX_WAITING_REQUESTS_ENVIRONMENT: "-1"}) + + +@pytest.mark.asyncio +async def test_middleware_rejects_fresh_overload_and_preserves_health() -> None: + downstream_calls: list[str] = [] + messages: list[AsgiMessage] = [] + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + del receive, send + downstream_calls.append(str(scope["path"])) + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + messages.append(message) + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=lambda: 4, + start_sampler=False, + ) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + + await middleware({"type": "http", "path": "/v1/chat/completions"}, receive, send) + await middleware({"type": "http", "path": "/health"}, receive, send) + + assert messages[0]["status"] == 429 + assert (b"retry-after", b"3") in messages[0]["headers"] + assert downstream_calls == ["/health"] + + +@pytest.mark.asyncio +async def test_middleware_fails_open_when_metrics_are_unavailable() -> None: + called = False + + async def app(scope: dict[str, object], receive: object, send: object) -> None: + nonlocal called + del scope, receive, send + called = True + + async def receive() -> AsgiMessage: + return {"type": "http.request"} + + async def send(message: AsgiMessage) -> None: + raise AssertionError(f"unexpected middleware response: {message}") + + controller = QueueBackpressureController(reader=lambda: None, start_sampler=False) + controller.sample_once() + middleware = QueueDepthBackpressureMiddleware(app, controller) + await middleware({"type": "http", "path": "/v1/completions"}, receive, send) + + assert called + + +def test_sampler_fails_open_after_reader_error_and_recovers() -> None: + readings: list[Exception | int] = [RuntimeError("metrics unavailable"), 4] + + def read_queue_depth() -> int: + reading = readings.pop(0) + if isinstance(reading, Exception): + raise reading + return reading + + controller = QueueBackpressureController( + QueueBackpressureSettings(2, 3), + reader=read_queue_depth, + start_sampler=False, + ) + + failed_snapshot = controller.sample_once() + failed_reject, _ = controller.should_reject() + recovered_snapshot = controller.sample_once() + recovered_reject, _ = controller.should_reject() + + assert failed_snapshot.depth is None + assert not failed_reject + assert recovered_snapshot.depth == 4 + assert recovered_reject + + +def test_sampler_preserves_process_control_exceptions() -> None: + def interrupt() -> int: + raise KeyboardInterrupt + + controller = QueueBackpressureController(reader=interrupt, start_sampler=False) + + with pytest.raises(KeyboardInterrupt): + controller.sample_once() diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py new file mode 100644 index 000000000..5c92d5a38 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import shlex +import subprocess +from dataclasses import replace +from pathlib import Path + +import pytest +from conftest import RuntimeCase, relocate_plan + +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import PortClaim, ResolvedSlurmRunPlan, ResolvedTopology +from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, RuntimeStepSpec, build_runtime_manifest +from data_designer.slurm.runtime.distributed import build_vllm_process_command +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import AllocationContext, RuntimeStepRole +from data_designer.slurm.runtime.node_spec import decode_node_worker_spec +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.serving.vllm import ResolvedVllmProcess +from data_designer.slurm.state import RetryPlan, RetryShard + + +def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + runtime_root = Path("/tmp/data-designer-slurm-4101-0/runtime") + log_directory = context.attempt_directory / "logs/execution-00000002" + + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0"}, + runtime_root=runtime_root, + log_directory=log_directory, + layout=AllocationLayout(("compute-001",)), + ) + reloaded = RuntimeBootstrapManifest.model_validate_json(manifest.serialize_json()) + + assert reloaded == manifest + assert [step.role for step in manifest.steps] == [ + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER, + RuntimeStepRole.ENDPOINT, + RuntimeStepRole.CLIENT, + ] + assert all(step.command[0] != "srun" for step in manifest.steps) + assert all(step.stdout_path.startswith(context.attempt_directory.as_posix()) for step in manifest.steps) + assert manifest.steps[-1].command[:4] == ( + context.plan.client.runtime.python_executable, + "-m", + "data_designer.slurm.runtime.entrypoint", + "client", + ) + assert "--shard-id" not in manifest.steps[-1].command + assert "--attempt-id" not in manifest.steps[-1].command + assert "--plan" in manifest.steps[-1].command + assert "--attempt-dir" in manifest.steps[-1].command + assert all(step.node_hosts == ("compute-001",) for step in manifest.steps) + assert all(step.role is not RuntimeStepRole.SERVER_PREFLIGHT for step in manifest.steps) + server = next(step for step in manifest.steps if step.role is RuntimeStepRole.SERVER) + assert server.command[server.command.index("--host") + 1] == "127.0.0.1" + endpoint = next(step for step in manifest.steps if step.role is RuntimeStepRole.ENDPOINT) + assert endpoint.literal_environment["PYTHONPATH"] == runtime_root.as_posix() + assert endpoint.container_environment == () + client_steps = tuple( + step for step in manifest.steps if step.role in {RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.CLIENT} + ) + assert all( + step.literal_environment["DATA_DESIGNER_SLURM_SCRATCH_ROOT"] == runtime_root.parent.as_posix() + for step in client_steps + ) + assert "SLURM_TMPDIR" not in manifest.serialize_json() + + +@pytest.mark.parametrize( + ("nodes", "tensor_parallel", "replicas", "stagger"), ((1, 1, 8, 0), (1, 2, 4, 2), (1, 4, 2, 0), (2, 4, 4, 2)) +) +def test_multi_replica_manifest_uses_one_gpu_owned_worker( + runtime_case: RuntimeCase, nodes: int, tensor_parallel: int, replicas: int, stagger: int +) -> None: + plan = runtime_case.context.plan + placement = plan.deployments[0] + authored = placement.authored.model_copy( + update={ + "topology": placement.authored.topology.model_copy(update={"tensor_parallel": tensor_parallel}), + "resources": placement.authored.resources.model_copy(update={"nodes": nodes}), + "server": placement.authored.server.model_copy( + update={ + "startup_timeout": "10s", + "distributed_init_timeout": "1s", + "lead_boot_standoff": "4s", + "rank_launch_stagger": f"{stagger}s", + } + ), + } + ) + topology = ResolvedTopology.derive( + node_count=nodes, + gpus_per_node=plan.resolved_gpus_per_node, + tensor_parallel=tensor_parallel, + nodes_per_replica=1, + ) + ports = tuple( + PortClaim( + name=f"{placement.deployment_id}-http-{index:05d}", + role="http", + node_index=index // (replicas // nodes), + port=18000 + index % (replicas // nodes), + ) + for index in range(replicas) + ) + placement = placement.model_copy( + update={"authored": authored, "topology": topology, "ports": ports, "node_indices": tuple(range(nodes))} + ) + plan = ResolvedSlurmRunPlan.model_validate(plan.model_copy(update={"deployments": (placement,)}).model_dump()) + context = replace(runtime_case.context, plan=plan) + + layout = AllocationLayout(tuple(f"compute-{index + 1:03d}" for index in range(nodes))) + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=layout, + ) + + servers = tuple(step for step in manifest.steps if step.role is RuntimeStepRole.SERVER) + assert len(servers) == 1 + assert servers[0].gpu_indices == tuple(range(8)) + assert len(servers[0].readiness) == replicas + assert {probe.host for probe in servers[0].readiness} == ({"127.0.0.1"} if nodes == 1 else set(layout.node_hosts)) + worker = decode_node_worker_spec(servers[0].command[-1]) + processes = tuple(process for node in worker.nodes for process in node.processes) + assert len(processes) == replicas + assert [process.gpu_indices for process in processes] == [ + tuple( + range((index % (replicas // nodes)) * tensor_parallel, (index % (replicas // nodes) + 1) * tensor_parallel) + ) + for index in range(replicas) + ] + assert all( + ("--host", "127.0.0.1" if nodes == 1 else "0.0.0.0") == process.command[process.command.index("--host") :][:2] + for process in processes + ) + endpoint = next(step for step in manifest.steps if step.role is RuntimeStepRole.ENDPOINT) + assert "0.0.0.0" not in endpoint.command + assert " ".join(endpoint.command).count("http://") == replicas + _assert_launch_readiness(manifest, runtime_case.workspace) + + +@pytest.mark.parametrize("role", (RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.SERVER)) +def test_runtime_step_rejects_execution_mode_that_conflicts_with_role( + runtime_case: RuntimeCase, role: RuntimeStepRole +) -> None: + context = runtime_case.context + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0"}, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), + ) + step = next(item for item in manifest.steps if item.role is role) + payload = step.model_dump(mode="python") + payload["execution"] = "native" if step.execution == "container" else "container" + payload["image_path"] = None if payload["execution"] == "native" else "/workspace/client.sqsh" + if payload["execution"] == "native": + payload["container_environment"] = () + + with pytest.raises(ValueError, match="execution does not match its role"): + RuntimeStepSpec.model_validate(payload) + + +def test_bootstrap_manifest_rejects_runtime_outside_allocation_scratch(runtime_case: RuntimeCase) -> None: + with pytest.raises(SlurmRuntimeError, match="outside allocation-local scratch"): + build_runtime_manifest( + runtime_case.context, + {"SLURM_JOB_GPUS": "0"}, + runtime_root=Path("/shared/runtime"), + log_directory=runtime_case.context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), + ) + + +def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=context.plan.run_id, + created_at=runtime_case.created_at, + resolved_plan=context.attempt.resolved_plan, + planned_shards=( + RetryShard( + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + attempt_ordinal=context.attempt.attempt_ordinal, + array_task_index=context.shard.array_task_index, + ), + ), + effective_resume_mode="never", + ) + retry_context = replace(context, retry_plan=retry) + + manifest = build_runtime_manifest( + retry_context, + {"SLURM_JOB_GPUS": "0"}, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), + ) + + preflight = manifest.steps[0].command + client = manifest.steps[-1].command + assert ("--resume-mode", "never") == preflight[preflight.index("--resume-mode") :][:2] + assert ("--retry-id", retry.retry_id) == client[client.index("--retry-id") :][:2] + assert ("--retry-plan-sha256", retry.compute_sha256()) == client[client.index("--retry-plan-sha256") :][:2] + assert ("--effective-resume-mode", "never") == client[client.index("--effective-resume-mode") :][:2] + assert "--resume-mode" not in client + + +def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( + runtime_case: RuntimeCase, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployments = tuple( + deployment.model_copy( + update={ + "authored": deployment.authored.model_copy( + update={ + "model": f"/workspace/primary/models/model-{deployment_index}", + "served_model_name": deployment.served_model_name, + "server": deployment.authored.server.model_copy( + update={ + "startup_timeout": "10s", + "distributed_init_timeout": "1s", + "lead_boot_standoff": "4s", + "rank_launch_stagger": "2s", + } + ), + } + ), + "model": f"/workspace/primary/models/model-{deployment_index}", + } + ) + for deployment_index, deployment in enumerate(multi_node_plan.deployments) + ) + context = _replace_plan(runtime_case, multi_node_plan.model_copy(update={"deployments": deployments})) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=layout, + ) + + distributed = next(step for step in manifest.steps if step.step_id == "deployment-00000-serve") + preflight = next(step for step in manifest.steps if step.step_id == "deployment-00000-preflight") + worker_spec = decode_node_worker_spec(distributed.command[-1]) + endpoint = next(step for step in manifest.steps if step.step_id == "deployment-00000-endpoint") + remote_server = next(step for step in manifest.steps if step.step_id == "deployment-00001-serve") + remote_preflight = next(step for step in manifest.steps if step.step_id == "deployment-00001-preflight") + + assert distributed.node_hosts == ("compute-001", "compute-002") + assert distributed.kill_on_bad_exit + assert preflight.role is RuntimeStepRole.SERVER_PREFLIGHT + assert tuple(node.host for node in worker_spec.nodes) == distributed.node_hosts + assert "--master-addr" in worker_spec.nodes[0].processes[0].command + assert "compute-001" in worker_spec.nodes[0].processes[0].command + assert worker_spec.nodes[0].processes[0].command[2] == "/workspace/primary/models/model-0" + assert worker_spec.required_model_path == "/workspace/primary/models/model-0" + assert "--headless" in worker_spec.nodes[1].processes[0].command + assert tuple(probe.host for probe in distributed.readiness) == ("compute-001",) + assert "http://compute-001:" in " ".join(endpoint.command) + assert endpoint.node_hosts == ("compute-001",) + assert endpoint.literal_environment["PYTHONPATH"] == "/tmp/data-designer-slurm-4101-0/runtime" + assert endpoint.container_environment == () + assert remote_preflight.node_hosts == ("compute-003",) + remote_worker_spec = decode_node_worker_spec(remote_preflight.command[-1]) + assert remote_worker_spec.required_model_path == "/workspace/primary/models/model-1" + assert remote_server.node_hosts == ("compute-003",) + assert remote_server.gpu_indices == tuple(range(8)) + remote_server_spec = decode_node_worker_spec(remote_server.command[-1]) + assert remote_server_spec.required_model_path == "/workspace/primary/models/model-1" + assert remote_server_spec.nodes[0].host == "compute-003" + assert len(remote_server_spec.nodes[0].processes) == 8 + assert all( + process.command[process.command.index("--distributed-executor-backend") + 1] == "uni" + for process in remote_server_spec.nodes[0].processes + ) + assert tuple(probe.host for probe in remote_server.readiness) == ("compute-003",) * 8 + _assert_launch_readiness(manifest, runtime_case.workspace) + + +def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path) -> None: + manifest_path = workspace / "readiness-manifest.json" + manifest_path.write_text(manifest.serialize_json()) + cases: list[str] = [] + budgets: list[tuple[int, int]] = [] + for step in manifest.steps: + if step.role is not RuntimeStepRole.SERVER: + continue + worker = decode_node_worker_spec(step.command[-1]) + for probe in step.readiness: + process = next( + process + for node in worker.nodes + if node.host == probe.host or probe.host == "127.0.0.1" + for process in node.processes + if "--headless" not in process.command + if str(probe.port) == process.command[process.command.index("--port") + 1] + ) + budgets.append((probe.deadline_seconds, process.launch_delay_seconds + 10)) + cases.append(f"*://{probe.host}:{probe.port}/*) ((SECONDS >= {process.launch_delay_seconds + 1}));;") + command = f""" +DD_RUNTIME_MANIFEST={shlex.quote(str(manifest_path))} +SECONDS=0 +dd_start_step() {{ DD_LAST_PID=1; }} +dd_register_required_pid() {{ :; }} +dd_require_running() {{ :; }} +dd_sleep() {{ if [[ $1 == 0.5 ]]; then SECONDS=$((SECONDS + 1)); else SECONDS=$((SECONDS + $1)); fi; }} +curl() {{ case "${{@: -1}}" in {" ".join(cases)} *) return 99;; esac; }} +dd_start_servers +dd_wait_for_role_readiness server +""" + completed = _run_readiness_command(command) + assert completed.returncode == 0, completed.stderr + assert all(actual == expected for actual, expected in budgets) + + +def _run_readiness_command(body: str) -> subprocess.CompletedProcess[str]: + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + return subprocess.run( + ("bash", "-c", f"source {shlex.quote(str(entrypoint))}\n{body}"), + capture_output=True, + text=True, + timeout=10, + ) + + +def test_server_readiness_uses_one_deadline_for_all_probes() -> None: + command = """ +DD_RUNTIME_MANIFEST=unused +SECONDS=0 +DD_STEP_STARTED_AT[server]=0 +dd_read_step_ids() { DD_STEP_IDS=(server); } +dd_read_step() { DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4); } +dd_require_running() { :; } +dd_sleep() { SECONDS=$((SECONDS + 1)); } +curl() { [[ "${@: -1}" == *:10001/* && $SECONDS -ge 3 ]]; } +if dd_wait_for_role_readiness server; then + echo unexpected-success +else + echo "timeout-at:$SECONDS" +fi +""" + completed = _run_readiness_command(command) + + assert completed.returncode == 0 + assert "timeout-at:4" in completed.stdout + assert "readiness timed out" in completed.stderr + + +@pytest.mark.parametrize( + "probe_fields", + ( + pytest.param("127.0.0.1 10001 /health 4", id="single-probe"), + pytest.param("127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4", id="second-probe"), + ), +) +def test_server_readiness_rejects_late_success(probe_fields: str) -> None: + command = f""" +DD_RUNTIME_MANIFEST=unused +SECONDS=0 +DD_STEP_STARTED_AT[server]=0 +dd_read_step_ids() {{ DD_STEP_IDS=(server); }} +dd_read_step() {{ DD_STEP_PROBE_FIELDS=({probe_fields}); }} +dd_require_running() {{ :; }} +dd_sleep() {{ :; }} +curl() {{ SECONDS=5; return 0; }} +if dd_wait_for_role_readiness server; then + echo unexpected-success +else + echo "timeout-at:$SECONDS" +fi +""" + completed = _run_readiness_command(command) + + assert completed.returncode == 0 + assert "timeout-at:5" in completed.stdout + assert "readiness timed out" in completed.stderr + + +def test_server_readiness_checks_later_steps_before_earlier_steps_finish() -> None: + command = """ +DD_RUNTIME_MANIFEST=unused +SECONDS=0 +DD_STEP_STARTED_AT[first]=0 +DD_STEP_STARTED_AT[second]=0 +dd_read_step_ids() { DD_STEP_IDS=(first second); } +dd_read_step() { + case "$2" in + first) DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 120);; + second) DD_STEP_PROBE_FIELDS=(127.0.0.1 10002 /health 60);; + esac +} +dd_require_running() { :; } +dd_sleep() { SECONDS=$((SECONDS + 1)); } +curl() { + case "${@: -1}" in + *:10001/*) ((SECONDS >= 110));; + *:10002/*) ((SECONDS >= 10));; + esac +} +dd_wait_for_role_readiness server +echo "ready-at:$SECONDS" +""" + completed = _run_readiness_command(command) + + assert completed.returncode == 0, completed.stderr + assert "ready-at:110" in completed.stdout + + +def test_pipeline_parallel_process_uses_multi_process_executor( + runtime_case: RuntimeCase, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + context = _replace_plan(runtime_case, multi_node_plan) + layout = AllocationLayout(("compute-001", "compute-002", "compute-003")) + deployment = resolve_allocation_deployments( + context, + {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, + )[0] + payload = deployment.processes[0].model_dump(mode="python") + payload.update({"gpu_indices": (0,), "tensor_parallel": 1}) + process = ResolvedVllmProcess.model_validate(payload) + + command = build_vllm_process_command(deployment, process, context.plan, layout) + + executor_index = command.index("--distributed-executor-backend") + assert process.tensor_parallel == 1 + assert process.pipeline_parallel == 2 + assert command[executor_index + 1] == "mp" + + +def _replace_plan(runtime_case: RuntimeCase, source_plan: ResolvedSlurmRunPlan) -> AllocationContext: + plan = relocate_plan(source_plan, runtime_case.workspace) + shard = plan.shards[0] + plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + attempt = runtime_case.context.attempt.model_copy( + update={ + "run_id": plan.run_id, + "shard_id": shard.shard_id, + "resolved_plan": ArtifactReference(path=plan_path.as_posix(), sha256=plan.compute_sha256()), + "scheduler": runtime_case.context.attempt.scheduler.model_copy( + update={"array_task_id": shard.array_task_index} + ), + } + ) + attempt_directory = plan_path.parent / "shards" / shard.shard_id / "attempts" / attempt.attempt_id + attempt_directory.mkdir(parents=True, mode=0o700) + return AllocationContext( + plan=plan, + shard=shard, + attempt=attempt, + attempt_directory=attempt_directory, + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py new file mode 100644 index 000000000..ad8361a77 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import io +import os +import stat +import subprocess +import sys +import tarfile +from pathlib import Path + +import pytest + +import data_designer.slurm.runtime.bundle as runtime_bundle +from data_designer.slurm.runtime.bundle import stage_runtime_bundle +from data_designer.slurm.runtime.errors import SlurmRuntimeError + + +def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + + first = stage_runtime_bundle(workspace) + second = stage_runtime_bundle(workspace) + + assert first == second + archive_path = Path(first.path) + content = archive_path.read_bytes() + assert hashlib.sha256(content).hexdigest() == first.sha256 + assert archive_path.name == f"{first.sha256}.tar.gz" + assert stat.S_IMODE(archive_path.stat().st_mode) == 0o600 + assert stat.S_IMODE(archive_path.parent.stat().st_mode) == 0o700 + with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as archive: + names = archive.getnames() + assert names[:2] == ["scratch.sh", "entrypoint.sh"] + assert names[2:5] == ["plan_reader.sh", "step_runner.sh", "cleanup.sh"] + assert names[5] == "data_designer/slurm/runtime/slurm-sources.txt" + assert names[6] == "data_designer/slurm/__init__.py" + assert "data_designer/slurm/runtime/controller.py" in names + assert "data_designer/slurm/runtime/entrypoint.py" in names + assert "data_designer/slurm/state/store.py" in names + assert "data_designer/slurm/contracts.py" in names + assert archive.getmember("entrypoint.sh").mode == 0o500 + assert all(archive.getmember(name).uid == 0 for name in names) + entrypoint = archive.extractfile("entrypoint.sh") + assert entrypoint is not None + entrypoint_content = entrypoint.read() + assert b"SPDX-License-Identifier: Apache-2.0" in entrypoint_content + assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint_content + assert b"--retry-plan-sha256" in entrypoint_content + assert b"--effective-resume-mode" in entrypoint_content + + +def test_runtime_bundle_recursively_collects_and_imports_nested_packages( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = tmp_path / "workspace" + source_root = tmp_path / "source" / "slurm" + runtime_root = source_root / "runtime" + nested_root = runtime_root / "nested" + workspace.mkdir(mode=0o700) + nested_root.mkdir(parents=True) + (source_root / "__init__.py").write_text("") + (runtime_root / "__init__.py").write_text("") + (runtime_root / "bundle.py").write_text("") + for name in ("scratch.sh", "entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh"): + (runtime_root / name).write_text("") + (nested_root / "__init__.py").write_text("") + (nested_root / "worker.py").write_text("VALUE = 42\n") + monkeypatch.setattr(runtime_bundle, "__file__", (runtime_root / "bundle.py").as_posix()) + + reference = stage_runtime_bundle(workspace) + extracted = tmp_path / "extracted" + with tarfile.open(reference.path, mode="r:gz") as archive: + archive.extractall(extracted, filter="data") + manifest_path = extracted / "data_designer/slurm/runtime/slurm-sources.txt" + manifest = manifest_path.read_text().splitlines() + assert "data_designer/slurm/runtime/nested/worker.py" in manifest + environment = dict(os.environ) + environment["PYTHONPATH"] = extracted.as_posix() + + completed = subprocess.run( + ( + sys.executable, + "-c", + "from data_designer.slurm.runtime.nested.worker import VALUE; print(VALUE, end='')", + ), + capture_output=True, + text=True, + check=False, + env=environment, + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout == "42" + + +def test_runtime_bundle_rejects_different_bytes_at_digest_path(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + reference = stage_runtime_bundle(workspace) + path = Path(reference.path) + path.write_bytes(b"different") + path.chmod(0o600) + + with pytest.raises(SlurmRuntimeError, match="cannot stage"): + stage_runtime_bundle(workspace) + + +def test_runtime_bundle_recovers_an_interrupted_immutable_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + original_unlink = runtime_bundle.os.unlink + + def interrupt_temporary_unlink(path: str, *, dir_fd: int | None = None) -> None: + del path, dir_fd + raise KeyboardInterrupt("injected runtime publication interruption") + + monkeypatch.setattr(runtime_bundle.os, "unlink", interrupt_temporary_unlink) + with pytest.raises(KeyboardInterrupt, match="runtime publication interruption"): + stage_runtime_bundle(workspace) + + runtime_root = workspace / "runtime" + archive_paths = tuple(runtime_root.glob("*.tar.gz")) + temporary_paths = tuple(runtime_root.glob(".runtime.*.tmp")) + assert len(archive_paths) == len(temporary_paths) == 1 + assert archive_paths[0].stat().st_ino == temporary_paths[0].stat().st_ino + assert archive_paths[0].stat().st_nlink == 2 + + monkeypatch.setattr(runtime_bundle.os, "unlink", original_unlink) + reference = stage_runtime_bundle(workspace) + assert Path(reference.path).stat().st_nlink == 1 + assert not tuple(runtime_root.glob(".runtime.*.tmp")) + + +def test_runtime_bundle_restats_after_a_peer_recovers_interrupted_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + original_unlink = runtime_bundle.os.unlink + original_listdir = runtime_bundle.os.listdir + + def interrupt_temporary_unlink(path: str, *, dir_fd: int | None = None) -> None: + del path, dir_fd + raise KeyboardInterrupt("injected runtime publication interruption") + + monkeypatch.setattr(runtime_bundle.os, "unlink", interrupt_temporary_unlink) + with pytest.raises(KeyboardInterrupt, match="runtime publication interruption"): + stage_runtime_bundle(workspace) + monkeypatch.setattr(runtime_bundle.os, "unlink", original_unlink) + + def recover_before_candidate_stat(directory_descriptor: int) -> list[str]: + names = original_listdir(directory_descriptor) + temporary_name = next(name for name in names if name.startswith(".runtime.")) + original_unlink(temporary_name, dir_fd=directory_descriptor) + return names + + monkeypatch.setattr(runtime_bundle.os, "listdir", recover_before_candidate_stat) + + reference = stage_runtime_bundle(workspace) + + assert Path(reference.path).stat().st_nlink == 1 + + +def test_runtime_bundle_rejects_symlinked_runtime_root(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + outside = tmp_path / "outside" + workspace.mkdir(mode=0o700) + outside.mkdir(mode=0o700) + (workspace / "runtime").symlink_to(outside, target_is_directory=True) + + with pytest.raises(SlurmRuntimeError, match="cannot stage"): + stage_runtime_bundle(workspace) + + assert tuple(outside.iterdir()) == () + + +def test_runtime_bundle_rejects_a_group_writable_workspace(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o770) + workspace.chmod(0o770) + + with pytest.raises(SlurmRuntimeError, match="cannot stage"): + stage_runtime_bundle(workspace) + + +def test_runtime_bundle_normalizes_source_read_failures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + + def fail_archive() -> bytes: + raise OSError("injected source race") + + monkeypatch.setattr(runtime_bundle, "_build_runtime_archive", fail_archive) + + with pytest.raises(SlurmRuntimeError, match="cannot stage"): + stage_runtime_bundle(workspace) + + +def test_extracted_bundle_runtime_takes_precedence_over_installed_sources(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + extracted = tmp_path / "extracted" + workspace.mkdir(mode=0o700) + extracted.mkdir(mode=0o700) + reference = stage_runtime_bundle(workspace) + with tarfile.open(reference.path, mode="r:gz") as archive: + archive.extractall(extracted, filter="data") + environment = dict(os.environ) + environment["PYTHONPATH"] = extracted.as_posix() + manifest = (extracted / "data_designer/slurm/runtime/slurm-sources.txt").read_text().splitlines() + modules = tuple(_get_module_name(source_name) for source_name in manifest) + import_all = "; ".join(f"importlib.import_module({module_name!r})" for module_name in modules) + + completed = subprocess.run( + ( + sys.executable, + "-c", + f"import importlib; {import_all}; " + "import data_designer.slurm.runtime as runtime; print(runtime.__file__, end='')", + ), + capture_output=True, + text=True, + check=False, + env=environment, + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout == (extracted / "data_designer/slurm/runtime/__init__.py").as_posix() + + +def test_backpressure_module_imports_from_the_bundle_without_site_packages(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + extracted = tmp_path / "extracted" + workspace.mkdir(mode=0o700) + reference = stage_runtime_bundle(workspace) + with tarfile.open(reference.path, mode="r:gz") as archive: + archive.extractall(extracted, filter="data") + + completed = subprocess.run( + ( + sys.executable, + "-S", + "-c", + "import sys; " + f"sys.path.insert(0, {extracted.as_posix()!r}); " + "from data_designer.slurm.runtime.backpressure import QueueDepthBackpressureMiddleware", + ), + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + + +def _get_module_name(source_name: str) -> str: + module_name = source_name.removesuffix(".py").replace("/", ".") + return module_name.removesuffix(".__init__") diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py new file mode 100644 index 000000000..fe9553100 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast +from unittest.mock import Mock + +import pytest + +import data_designer.slurm.runtime.context as runtime_context +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.context import load_allocation_context +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + RetryPlan, + RetryShard, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateWriter, +) + + +def test_runtime_retry_binding_requires_all_arguments() -> None: + with pytest.raises(SlurmRuntimeError, match="binding is incomplete"): + runtime_context._load_retry_plan( + Mock(), + retry_id="retry-0001", + retry_plan_sha256=None, + effective_resume_mode="never", + ) + + +@pytest.mark.parametrize( + ("digest", "resume_mode"), + [ + pytest.param("f" * 64, "never", id="digest"), + pytest.param(None, "always", id="resume-mode"), + ], +) +def test_runtime_retry_binding_rejects_tampered_arguments( + single_node_plan: ResolvedSlurmRunPlan, + digest: str | None, + resume_mode: str, +) -> None: + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=single_node_plan.run_id, + created_at=datetime(2026, 9, 2, tzinfo=timezone.utc), + resolved_plan=ArtifactReference( + path="/workspace/primary/runs/run-single/resolved-plan.json", + sha256=single_node_plan.compute_sha256(), + ), + planned_shards=( + RetryShard( + shard_id="shard-00000", + attempt_id="attempt-0002", + attempt_ordinal=2, + array_task_index=0, + ), + ), + effective_resume_mode="never", + ) + writer = Mock() + writer.load_retry_plan.return_value = retry + + with pytest.raises(SlurmRuntimeError, match="binding differs"): + runtime_context._load_retry_plan( + writer, + retry_id=retry.retry_id, + retry_plan_sha256=retry.compute_sha256() if digest is None else digest, + effective_resume_mode=resume_mode, + ) + + +def test_allocation_context_rejects_container_remapped_workspace( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + physical_workspace = tmp_path / "workspace" + physical_workspace.mkdir() + logical_workspace = single_node_plan.selected_profile.profile.workspace_root + logical_runs = Path(logical_workspace) / "runs" + fast_runs = tmp_path / "fast-runs" + logical_attempts = ( + Path(logical_workspace) + / "runs" + / single_node_plan.run_id + / "shards" + / single_node_plan.shards[0].shard_id + / "attempts" + ) + fast_attempts = tmp_path / "fast-attempts" + payload = cast(dict[str, object], json.loads(single_node_plan.serialize_json())) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mounts = [ + {"source": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False}, + {"source": logical_runs.as_posix(), "target": fast_runs.as_posix(), "read_only": False}, + {"source": logical_attempts.as_posix(), "target": fast_attempts.as_posix(), "read_only": False}, + ] + profile_payload["container_mounts"] = mounts + payload["container_mounts"] = mounts + profile = SlurmProfile.model_validate(profile_payload) + selected["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + created_at = datetime(2026, 9, 9, tzinfo=timezone.utc) + plan_reference = ArtifactReference( + path=f"{logical_workspace}/runs/{plan.run_id}/resolved-plan.json", + sha256=plan.compute_sha256(), + ) + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=created_at, + authored_config=plan.authored_config, + resolved_plan=plan_reference, + shard_count=1, + ) + shard = ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=plan.shards[0].shard_id, + shard_index=0, + record_range=plan.shards[0].record_range, + input_partition=plan.shards[0].input_partition, + resume_workspace=plan.shards[0].resume_workspace, + created_at=created_at, + ) + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=0) + attempt = AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.SUBMITTED, + scheduler=scheduler, + created_at=created_at, + updated_at=created_at, + ) + host_writer = SlurmStateWriter( + physical_workspace, + plan.run_id, + logical_workspace_root=logical_workspace, + ) + host_writer.initialize_run(authored_run_single, plan, run, (shard,)) + host_writer.create_attempt(attempt) + plan_path = fast_runs / plan.run_id / "resolved-plan.json" + plan_path.parent.mkdir(parents=True, mode=0o700) + plan_path.write_text(plan.serialize_json()) + plan_path.chmod(0o600) + attempt_directory = fast_attempts / attempt.attempt_id + attempt_directory.mkdir(parents=True, mode=0o700) + + with pytest.raises(SlurmRuntimeError, match="resolved plan path is invalid"): + load_allocation_context( + plan_path, + attempt_directory, + {"SLURM_ARRAY_TASK_ID": "0", "SLURM_ARRAY_JOB_ID": "4101"}, + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py new file mode 100644 index 000000000..ceb4a1e73 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -0,0 +1,846 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, replace +from datetime import timedelta +from pathlib import Path +from typing import Callable + +import pytest +from conftest import FakeClientStepBuilder, FakePreflight, FakeStateStore, RuntimeCase, relocate_plan +from slurm_test_fakes import FakeClock + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.controller import OneNodeAllocationController +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.signals import TerminationSignalCoordinator +from data_designer.slurm.runtime.supervisor import StepSupervisor +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, + RetryPlan, + RetryShard, + RunManifest, + ShardManifest, + SlurmStateWriter, +) +from data_designer.slurm.state.artifacts import compute_candidate_schema_digest + +_ALLOCATION_ENVIRONMENT = {"SLURM_JOB_GPUS": "0"} + + +def test_independent_controller_rejects_shared_plan_before_startup( + runtime_case: RuntimeCase, + shared_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = relocate_plan(shared_node_plan, runtime_case.workspace) + context = replace(runtime_case.context, plan=plan, shard=plan.shards[0]) + clock = FakeClock(runtime_case.created_at) + runner = _FakeRunner() + with pytest.raises(SlurmRuntimeError, match="coordinated"): + OneNodeAllocationController( + context, + runtime_proxy_path=context.attempt_directory / "runtime/proxy.py", + state=FakeStateStore(context.attempt), + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + assert not runner.steps + + +@dataclass(slots=True) +class _FakeProcess: + pid: int + returncode: int | None + fail_terminate: bool = False + remain_running: bool = False + terminate_calls: int = 0 + kill_calls: int = 0 + + def poll(self) -> int | None: + return self.returncode + + def terminate(self) -> None: + self.terminate_calls += 1 + if self.fail_terminate: + raise OSError("injected termination failure") + if not self.remain_running: + self.returncode = -15 + + def kill(self) -> None: + self.kill_calls += 1 + if not self.remain_running: + self.returncode = -9 + + def wait(self, timeout: float | None = None) -> int: + del timeout + if self.returncode is None: + raise TimeoutError + return self.returncode + + +class _FakeRunner: + def __init__( + self, + *, + generation_hook: Callable[[], None] | None = None, + server_returncode: int | None = None, + failed_role: RuntimeStepRole | None = None, + fail_cleanup: bool = False, + incomplete_cleanup: bool = False, + ) -> None: + self.generation_hook = generation_hook + self.server_returncode = server_returncode + self.failed_role = failed_role + self.fail_cleanup = fail_cleanup + self.incomplete_cleanup = incomplete_cleanup + self.steps: list[RuntimeStep] = [] + self.processes: list[_FakeProcess] = [] + + def start(self, step: RuntimeStep) -> _FakeProcess: + self.steps.append(step) + if step.role is RuntimeStepRole.CLIENT and self.generation_hook is not None: + self.generation_hook() + is_managed_service = step.role in {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} + if step.role is self.failed_role: + returncode = 23 + elif step.role is RuntimeStepRole.SERVER: + returncode = self.server_returncode + else: + returncode = None if is_managed_service else 0 + process = _FakeProcess( + pid=1000 + len(self.processes), + returncode=returncode, + fail_terminate=self.fail_cleanup and is_managed_service, + remain_running=self.incomplete_cleanup and is_managed_service, + ) + self.processes.append(process) + return process + + +class _FakeProber: + def __init__(self, *, ready: bool, clock: FakeClock, advance_on_failure: float = 0) -> None: + self.ready = ready + self.clock = clock + self.advance_on_failure = advance_on_failure + self.calls = 0 + + def is_ready(self, host: str, port: int, path: str, *, timeout_seconds: float) -> bool: + del host, port, path, timeout_seconds + self.calls += 1 + if not self.ready and self.advance_on_failure: + self.clock.advance(self.advance_on_failure) + return self.ready + + +@dataclass(slots=True) +class _RestartAwarePreflight: + state: FakeStateStore + + def verify(self, context: object, environment: object) -> None: + del context, environment + assert self.state.readiness[-1].state is ReadinessState.RESTARTING + + +def _supervisor( + runner: _FakeRunner, + clock: FakeClock, + *, + poll_interval_seconds: float = 0.1, + termination_grace_seconds: float = 10.0, +) -> StepSupervisor: + return StepSupervisor( + runner, + signals=TerminationSignalCoordinator(), + clock=clock, + poll_interval_seconds=poll_interval_seconds, + termination_grace_seconds=termination_grace_seconds, + ) + + +def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + + def write_result_under_dataset_lease() -> None: + assert state.dataset_workspace_lease_active + _write_complete_result(runtime_case, clock) + + runner = _FakeRunner(generation_hook=write_result_under_dataset_lease) + supervisor = _supervisor(runner, clock, poll_interval_seconds=0.1) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=supervisor, + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert result.terminal_classification is AttemptTerminalClassification.SUCCEEDED + assert result.candidate_output is not None + assert state.dataset_workspace_modes == ["never"] + assert not state.dataset_workspace_lease_active + assert state.winners[0].attempt_id == result.attempt_id + assert [step.role for step in runner.steps] == [ + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER, + RuntimeStepRole.ENDPOINT, + RuntimeStepRole.CLIENT, + ] + assert [item.state for item in state.readiness][0] is ReadinessState.PENDING + 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"} + + +def test_controller_publishes_result_before_success_with_real_state_writer( + runtime_case: RuntimeCase, + authored_run_single: DataDesignerSlurmConfig, +) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = _initialize_real_state(runtime_case, authored_run_single) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + persisted = state.load_attempt(result.shard_id, result.attempt_id) + assert persisted == result + assert persisted.state is AttemptLifecycleState.SUCCEEDED + assert persisted.candidate_output is not None + assert state.load_winner(result.shard_id).attempt_id == result.attempt_id + + +def test_winner_finalization_failure_leaves_controller_attempt_retryable(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + + class FailingState(FakeStateStore): + def finalize_winner(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("injected finalization failure") + + state = FailingState(runtime_case.context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="allocation runtime failed"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.attempt.candidate_output is not None + assert state.winners == [] + + +def test_controller_binds_retry_resume_mode_to_client_and_workspace(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=context.plan.run_id, + created_at=runtime_case.created_at, + resolved_plan=context.attempt.resolved_plan, + planned_shards=( + RetryShard( + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + attempt_ordinal=context.attempt.attempt_ordinal, + array_task_index=context.shard.array_task_index, + ), + ), + effective_resume_mode="never", + ) + retry_context = replace(context, retry_plan=retry) + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + client_steps = FakeClientStepBuilder() + controller = OneNodeAllocationController( + retry_context, + runtime_proxy_path=context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=client_steps, + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert client_steps.retry_resume_modes == ["never", "never"] + assert state.dataset_workspace_modes == ["never"] + + +def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner() + failure = SlurmRuntimeError(SlurmRuntimeErrorCode.PREFLIGHT_FAILED, "injected preflight failure") + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(failure), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="injected preflight failure") as raised: + controller.run() + + assert raised.value.__cause__ is None + assert runner.steps == [] + assert state.readiness == [] + assert state.attempt.state is AttemptLifecycleState.FAILED + + +@pytest.mark.parametrize( + ("persisted_state", "ready_backends", "endpoint_publication"), + ( + (ReadinessState.STARTING, 0, EndpointPublicationState.PENDING), + (ReadinessState.READY, 1, EndpointPublicationState.PUBLISHED), + ), +) +def test_requeued_running_attempt_publishes_restart_epoch_and_uses_fresh_logs( + runtime_case: RuntimeCase, + persisted_state: ReadinessState, + ready_backends: int, + endpoint_publication: EndpointPublicationState, +) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + running_attempt = runtime_case.context.attempt.__class__.model_validate( + runtime_case.context.attempt.model_dump(mode="python") + | { + "state": AttemptLifecycleState.RUNNING, + "updated_at": runtime_case.created_at + timedelta(seconds=2), + } + ) + context = replace(runtime_case.context, attempt=running_attempt) + persisted_readiness = AttemptReadiness( + schema_version=1, + run_id=running_attempt.run_id, + shard_id=running_attempt.shard_id, + attempt_id=running_attempt.attempt_id, + revision=7, + updated_at=runtime_case.created_at + timedelta(seconds=3), + state=persisted_state, + deployments=( + DeploymentReadiness( + deployment_id=context.plan.deployments[0].deployment_id, + model_alias=context.plan.deployments[0].authored.model_alias, + state=persisted_state, + expected_backends=context.plan.deployments[0].topology.replica_count, + ready_backends=ready_backends, + endpoint_publication=endpoint_publication, + ), + ), + ) + state = FakeStateStore(running_attempt, readiness=[persisted_readiness]) + restarted_case = replace(runtime_case, context=context) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(restarted_case, clock)) + controller = OneNodeAllocationController( + context, + runtime_proxy_path=context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=_RestartAwarePreflight(state), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert state.readiness[1].revision == 8 + assert state.readiness[1].state is ReadinessState.RESTARTING + assert state.readiness[1].deployments[0].state is ReadinessState.RESTARTING + assert state.readiness[1].deployments[0].ready_backends == 0 + assert state.readiness[1].deployments[0].endpoint_publication is EndpointPublicationState.PENDING + assert state.readiness[-1].state is ReadinessState.STOPPED + assert [step.role for step in runner.steps] == [ + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.SERVER, + RuntimeStepRole.ENDPOINT, + RuntimeStepRole.CLIENT, + ] + assert {step.stdout_path.parent.name for step in runner.steps} == {"execution-00000008"} + + +def test_required_server_exit_fails_and_cleans_partial_start(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(server_returncode=17) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="required runtime step"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert ReadinessState.FAILED in [item.state for item in state.readiness] + assert state.readiness[-1].state is ReadinessState.STOPPED + + +def test_interrupted_failed_readiness_write_cannot_bypass_cleanup(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + + class InterruptingState(FakeStateStore): + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: + if readiness.state is ReadinessState.FAILED: + raise KeyboardInterrupt("injected readiness interruption") + return super().write_readiness(readiness) + + state = InterruptingState(runtime_case.context.attempt) + runner = _FakeRunner(server_returncode=17) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="required runtime step") as raised: + controller.run() + + assert any("failed readiness could not be persisted: KeyboardInterrupt" in note for note in raised.value.__notes__) + assert all(process.poll() is not None for process in runner.processes) + assert state.attempt.state is AttemptLifecycleState.FAILED + + +def test_readiness_timeout_fails_and_terminates_server(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner() + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=False, clock=clock, advance_on_failure=10_000), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="readiness timed out"): + controller.run() + + server = next( + process + for step, process in zip(runner.steps, runner.processes, strict=True) + if step.role is RuntimeStepRole.SERVER + ) + assert server.terminate_calls == 1 + assert server.poll() == -15 + assert state.attempt.state is AttemptLifecycleState.FAILED + + +@pytest.mark.parametrize( + "failed_role", + (RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.ENDPOINT, RuntimeStepRole.CLIENT), +) +def test_managed_step_failure_fails_attempt_and_cleans_started_services( + runtime_case: RuntimeCase, + failed_role: RuntimeStepRole, +) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(failed_role=failed_role) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="status 23"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.readiness[-1].state is ReadinessState.STOPPED + services = ( + process + for step, process in zip(runner.steps, runner.processes, strict=True) + if step.role in {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} + ) + assert all(process.poll() is not None for process in services) + + +def test_cleanup_failure_prevents_false_success(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner( + generation_hook=lambda: _write_complete_result(runtime_case, clock), + fail_cleanup=True, + ) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor( + runner, + clock, + poll_interval_seconds=1, + termination_grace_seconds=1, + ), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="cleanup failed"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.attempt.candidate_output is not None + assert all(process.poll() is not None for process in runner.processes) + + +def test_cleanup_failure_is_retained_as_a_note_on_the_primary_failure(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(failed_role=RuntimeStepRole.CLIENT, fail_cleanup=True) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor( + runner, + clock, + poll_interval_seconds=1, + termination_grace_seconds=1, + ), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="status 23") as raised: + controller.run() + + assert any("cleanup also failed" in note for note in raised.value.__notes__) + assert state.attempt.state is AttemptLifecycleState.FAILED + + +def test_incomplete_cleanup_does_not_publish_stopped_readiness(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner( + generation_hook=lambda: _write_complete_result(runtime_case, clock), + incomplete_cleanup=True, + ) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor( + runner, + clock, + poll_interval_seconds=1, + termination_grace_seconds=1, + ), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="cleanup failed"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + assert state.readiness[-1].state is ReadinessState.FAILED + assert ReadinessState.STOPPED not in {readiness.state for readiness in state.readiness} + + +def test_future_client_timestamp_cannot_push_persisted_state_clock_forward(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + + def write_future_result() -> None: + _write_complete_result(runtime_case, clock) + clock.current_time -= timedelta(milliseconds=500) + + runner = _FakeRunner(generation_hook=write_future_result) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="later than the allocation clock"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + + +def test_partial_client_result_requires_candidate_manifest(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_partial_result(runtime_case, clock)) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError) as raised: + controller.run() + + assert raised.value.code is SlurmRuntimeErrorCode.FINALIZATION_FAILED + assert not (runtime_case.context.attempt_directory / "output-manifest.json").exists() + + +def test_nonempty_partial_client_result_completes(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(runtime_case.context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock, shortfall=1)) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert state.winners[0].attempt_id == result.attempt_id + + +def test_stale_candidate_from_an_earlier_generation_cannot_succeed(runtime_case: RuntimeCase) -> None: + clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) + _write_complete_result(runtime_case, clock) + clock.advance(5) + state = FakeStateStore(runtime_case.context.attempt) + controller = OneNodeAllocationController( + runtime_case.context, + runtime_proxy_path=runtime_case.context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(_FakeRunner(), clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment=_ALLOCATION_ENVIRONMENT, + ) + + with pytest.raises(SlurmRuntimeError, match="predates the current generation"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + + +def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock, *, shortfall: int = 0) -> None: + context = runtime_case.context + requested = context.shard.requested_records + dataset_path = (context.attempt_directory / "dataset").as_posix() + output_path = Path(dataset_path) / "part-00000.parquet" + output_path.parent.mkdir(exist_ok=True) + actual_records = requested - shortfall + table = lazy.pa.table({"record_id": range(actual_records)}) + lazy.pq.write_table(table, output_path) + output_bytes = output_path.read_bytes() + candidate = CandidateOutputManifest( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + attempt_ordinal=context.attempt.attempt_ordinal, + created_at=clock.now(), + dataset_path=dataset_path, + requested_records=requested, + actual_records=actual_records, + outcome=CandidateOutcome.COMPLETE if not shortfall else CandidateOutcome.PARTIAL, + files=( + CandidateOutputFile( + relative_path="part-00000.parquet", + sha256=hashlib.sha256(output_bytes).hexdigest(), + byte_size=len(output_bytes), + record_count=actual_records, + ), + ), + dataset_schema_digest=compute_candidate_schema_digest(table.schema), + provenance_digest=context.plan.compute_sha256(), + ) + candidate_path = context.attempt_directory / "output-manifest.json" + _write_record(candidate_path, candidate.serialize_json()) + clock.advance(1) + result = ClientResult( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + completed_at=clock.now(), + requested_records=requested, + actual_records=actual_records, + outcome=ClientOutcome.COMPLETE if not shortfall else ClientOutcome.PARTIAL, + dataset_path=dataset_path, + early_shutdown=False, + requested_resume_mode=context.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=ArtifactReference( + path=candidate_path.as_posix(), + sha256=candidate.compute_sha256(), + ), + ) + _write_record(context.attempt_directory / "client-result.json", result.serialize_json()) + + +def _initialize_real_state( + runtime_case: RuntimeCase, + authored_config: DataDesignerSlurmConfig, +) -> SlurmStateWriter: + context = runtime_case.context + planned_shard = context.shard + run = RunManifest( + schema_version=1, + run_id=context.plan.run_id, + created_at=runtime_case.created_at, + authored_config=context.plan.authored_config, + resolved_plan=context.attempt.resolved_plan, + shard_count=len(context.plan.shards), + ) + shard = ShardManifest( + schema_version=1, + run_id=context.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=runtime_case.created_at, + ) + state = SlurmStateWriter(runtime_case.workspace, context.plan.run_id) + state.initialize_run(authored_config, context.plan, run, (shard,)) + state.create_attempt(context.attempt) + return state + + +def _write_partial_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: + context = runtime_case.context + candidate_path = context.attempt_directory / "output-manifest.json" + result = ClientResult( + schema_version=1, + run_id=context.plan.run_id, + shard_id=context.shard.shard_id, + attempt_id=context.attempt.attempt_id, + completed_at=clock.now(), + requested_records=context.shard.requested_records, + actual_records=context.shard.requested_records - 1, + outcome=ClientOutcome.PARTIAL, + dataset_path=(context.attempt_directory / "dataset").as_posix(), + early_shutdown=True, + requested_resume_mode=context.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=ArtifactReference( + path=candidate_path.as_posix(), + sha256="a" * 64, + ), + ) + _write_record(context.attempt_directory / "client-result.json", result.serialize_json()) + + +def _write_record(path: Path, content: str) -> None: + path.write_text(content) + path.chmod(0o600) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py new file mode 100644 index 000000000..59c34c1f7 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -0,0 +1,512 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import csv +import hashlib +import io +import json +import os +import sys +import zipfile +from dataclasses import dataclass, replace +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import pyarrow.parquet as pq +import pytest +from conftest import FakeStateStore, RuntimeCase, relocate_plan +from packaging.tags import interpreter_name, interpreter_version + +import data_designer.slurm.runtime.entrypoint as entrypoint +from data_designer.slurm.client.process import ClientWorkerProcess +from data_designer.slurm.client.records import ClientEnvironmentManifest, ClientInstallerOutcome, ClientResult +from data_designer.slurm.contracts import ( + ArtifactReference, + InstalledDistribution, + compute_canonical_json_sha256, + compute_serialized_json_sha256, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.ports import resolve_allocation_plan +from data_designer.slurm.state import AttemptLifecycleState, CandidateOutputManifest, ReadinessState + + +class _InjectedFailure(Exception): + pass + + +@dataclass(frozen=True) +class _PluginRuntimeCase: + context: AllocationContext + plan_path: Path + image_metadata_directory: Path + + +def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: + assert ( + entrypoint.main( + ( + "prepare", + "--plan", + "resolved-plan.json", + "--attempt-dir", + "attempt-0001", + "--runtime-root", + "runtime", + "--manifest", + "runtime-manifest.json", + "--node-host", + "compute-001", + ) + ) + == 64 + ) + captured = capsys.readouterr() + assert "runtime paths must be absolute" in captured.err + assert "Traceback" not in captured.err + + +def test_native_phases_use_the_supplied_host_attempt_directory( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + container_attempt_directory = Path("/container/workspace/runs/run-single/shards/shard-00000/attempts/attempt-0001") + _patch_runtime_context(monkeypatch, runtime_case, state) + + def verify_attempt_directory(path: Path) -> None: + assert path == container_attempt_directory + raise _InjectedFailure + + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", verify_attempt_directory) + prepare = entrypoint._parse_arguments( + _phase_arguments( + "prepare", + runtime_case, + runtime_case.context.attempt_directory / "runtime", + container_attempt_directory / "runtime-manifest.json", + attempt_directory=container_attempt_directory, + ) + ) + with pytest.raises(_InjectedFailure): + entrypoint._prepare(prepare, {}) + + def load_candidate(*args: object, attempt_directory: Path | None = None) -> None: + assert attempt_directory == container_attempt_directory + raise _InjectedFailure + + monkeypatch.setattr(entrypoint, "load_nonempty_client_candidate", load_candidate) + client = entrypoint._parse_arguments( + _phase_arguments("client", runtime_case, attempt_directory=container_attempt_directory) + ) + with pytest.raises(_InjectedFailure): + entrypoint._client(client, {}, client_worker=ClientWorkerProcess(executor=lambda command: 0)) + + +def test_client_phase_installs_and_runs_plugin_in_default_worker_process( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + fake_plugin_overlay: Path, + tmp_path: Path, +) -> None: + case = _prepare_plugin_runtime_case(runtime_case, fake_plugin_overlay, tmp_path) + scratch_root = tmp_path / "scratch" + scratch_root.mkdir() + monkeypatch.setenv("DATA_DESIGNER_SLURM_SCRATCH_ROOT", scratch_root.as_posix()) + state = FakeStateStore(case.context.attempt) + runtime_case.context = case.context + _patch_runtime_context(monkeypatch, runtime_case, state) + assert "data_designer.config.column_types" in sys.modules + _use_controlled_image_inventory(monkeypatch, case.image_metadata_directory) + + allocation = resolve_allocation_plan(case.context.plan, {"SLURM_JOB_GPUS": "0"}) + endpoint = f"generator=http://127.0.0.1:{allocation.client.ports[0].port}/v1" + worker_arguments = ( + "preflight", + "--plan", + case.plan_path.as_posix(), + "--shard-id", + case.context.shard.shard_id, + "--attempt-id", + case.context.attempt.attempt_id, + "--attempt-dir", + case.context.attempt_directory.as_posix(), + "--endpoint", + endpoint, + ) + assert ClientWorkerProcess().run(worker_arguments) == 0 + + client = entrypoint._parse_arguments((*_phase_arguments("client", runtime_case), "--endpoint", endpoint)) + entrypoint._client(client, {}) + + environment = ClientEnvironmentManifest.model_validate_json( + (case.context.attempt_directory / "client-environment.json").read_text() + ) + result = ClientResult.model_validate_json((case.context.attempt_directory / "client-result.json").read_text()) + candidate = CandidateOutputManifest.model_validate_json( + (case.context.attempt_directory / "output-manifest.json").read_text() + ) + dataset = pq.read_table(Path(candidate.dataset_path) / candidate.files[0].relative_path) + + assert environment.installer_outcome is ClientInstallerOutcome.INSTALLED + assert [plugin.plugin_name for plugin in environment.plugins] == ["fake-slurm-column"] + assert result.candidate_output_manifest is not None + assert result.candidate_output_manifest.sha256 == candidate.compute_sha256() + assert state.attempt.candidate_output == result.candidate_output_manifest + assert set(dataset.column_names) == {"record_id", "custom"} + assert dataset.column("custom").to_pylist() == ["plugin-marker"] * case.context.shard.requested_records + + +@pytest.mark.parametrize("plan_fixture", ("single_node_plan", "shared_node_plan")) +def test_control_phases_record_running_ready_and_failed( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + request: pytest.FixtureRequest, + plan_fixture: str, +) -> None: + plan = relocate_plan(request.getfixturevalue(plan_fixture), runtime_case.workspace) + runtime_case = replace(runtime_case, context=replace(runtime_case.context, plan=plan, shard=plan.shards[0])) + state = FakeStateStore(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert state.attempt.state is AttemptLifecycleState.RUNNING + assert [item.state for item in state.readiness] == [ReadinessState.PENDING, ReadinessState.STARTING] + assert manifest_path.read_text() == "{}" + + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + assert state.readiness[-1].state is ReadinessState.READY + assert [ + (deployment.model_alias, deployment.expected_backends, deployment.ready_backends) + for deployment in state.readiness[-1].deployments + ] == [ + (deployment.authored.model_alias, deployment.topology.replica_count, deployment.topology.replica_count) + for deployment in plan.deployments + ] + + assert entrypoint.main(_phase_arguments("fail", runtime_case)) == 0 + assert state.attempt.state is AttemptLifecycleState.FAILED + assert [item.state for item in state.readiness[-2:]] == [ReadinessState.FAILED, ReadinessState.STOPPED] + assert all(deployment.ready_backends == 0 for deployment in state.readiness[-1].deployments) + + +def test_succeed_phase_stops_runtime_and_finalizes_winner( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + state.attempt = state.attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=(runtime_case.context.attempt_directory / "output-manifest.json").as_posix(), + sha256="a" * 64, + ) + } + ) + + assert entrypoint.main(_phase_arguments("succeed", runtime_case)) == 0 + + assert state.attempt.state is AttemptLifecycleState.SUCCEEDED + assert state.readiness[-1].state is ReadinessState.STOPPED + assert state.winners[0].attempt_id == state.attempt.attempt_id + + +def test_prepare_rejects_unverified_artifacts_before_starting_attempt( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + capsys: pytest.CaptureFixture[str], +) -> None: + state = FakeStateStore(runtime_case.context.attempt) + verify_artifacts = entrypoint.SystemAllocationPreflight.verify_artifacts + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_artifacts", verify_artifacts) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 64 + + assert state.attempt.state is AttemptLifecycleState.SUBMITTED + assert not state.readiness + assert not manifest_path.exists() + captured = capsys.readouterr() + assert "preflight_failed" in captured.err + assert "Traceback" not in captured.err + + +def test_succeed_phase_does_not_strand_success_when_winner_finalization_fails( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, +) -> None: + class FailingState(FakeStateStore): + def finalize_winner(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("injected finalization failure") + + state = FailingState(runtime_case.context.attempt) + runtime_root = runtime_case.context.attempt_directory / "runtime" + manifest_path = runtime_case.context.attempt_directory / "runtime-manifest.json" + _patch_runtime_context(monkeypatch, runtime_case, state) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_attempt_directory", lambda path: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_ports", lambda *args: None) + monkeypatch.setattr( + entrypoint, + "build_runtime_manifest", + lambda *args, **kwargs: SimpleNamespace(serialize_json=lambda: "{}"), + ) + assert entrypoint.main(_phase_arguments("prepare", runtime_case, runtime_root, manifest_path)) == 0 + assert entrypoint.main(_phase_arguments("ready", runtime_case)) == 0 + state.attempt = state.attempt.model_copy( + update={ + "candidate_output": ArtifactReference( + path=(runtime_case.context.attempt_directory / "output-manifest.json").as_posix(), + sha256="a" * 64, + ), + } + ) + assert entrypoint.main(_phase_arguments("succeed", runtime_case)) == 70 + assert state.attempt.state is AttemptLifecycleState.RUNNING + + +def _patch_runtime_context( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + state: FakeStateStore, +) -> None: + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_client_runtime", lambda context: None) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_artifacts", lambda context: None) + monkeypatch.setenv("SLURM_JOB_GPUS", "0") + + +def _prepare_plugin_runtime_case( + runtime_case: RuntimeCase, + fake_plugin_overlay: Path, + tmp_path: Path, +) -> _PluginRuntimeCase: + runtime_distributions = runtime_case.context.plan.client.runtime.distributions + wheel_path = Path(runtime_case.context.plan.authored_config.path).parent / ( + "dependencies/fake_data_designer_plugin-1.0.0-py3-none-any.whl" + ) + _build_plugin_wheel(fake_plugin_overlay, wheel_path) + payload = runtime_case.context.plan.model_dump(mode="json") + _configure_identity_mount(payload, runtime_case.workspace) + _configure_plugin_builder(payload) + client = cast(dict[str, object], payload["client"]) + _configure_client_runtime(client, runtime_distributions) + lock = _plugin_dependency_lock(client, runtime_distributions, wheel_path) + _write_dependency_lock(client, lock) + + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + plan_path.write_text(plan.serialize_json()) + Path(plan.invocation.effective_input_bindings.managed_assets_path).mkdir(parents=True, exist_ok=True) + attempt = runtime_case.context.attempt.model_copy( + update={ + "resolved_plan": ArtifactReference(path=plan_path.as_posix(), sha256=plan.compute_sha256()), + "state": AttemptLifecycleState.RUNNING, + } + ) + context = AllocationContext(plan, plan.shards[0], attempt, runtime_case.context.attempt_directory) + image_metadata_directory = _write_inventory_bootstrap(runtime_distributions, tmp_path) + return _PluginRuntimeCase(context, plan_path, image_metadata_directory) + + +def _configure_identity_mount(payload: dict[str, object], workspace: Path) -> None: + mount = {"source": workspace.as_posix(), "target": workspace.as_posix(), "read_only": False} + payload["container_mounts"] = [mount] + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile = cast(dict[str, object], selected_profile["profile"]) + profile["container_mounts"] = [mount] + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile) + + +def _configure_plugin_builder(payload: dict[str, object]) -> None: + resolved_builder = cast(dict[str, object], payload["builder"]) + builder = cast(dict[str, object], resolved_builder["inline"]) + cast(dict[str, object], builder["data_designer"])["columns"] = [ + {"name": "record_id", "column_type": "sampler", "sampler_type": "uuid", "params": {}}, + {"name": "custom", "column_type": "fake-slurm-column"}, + ] + resolved_builder["content_sha256"] = compute_serialized_json_sha256(builder) + + +def _configure_client_runtime( + client: dict[str, object], + runtime_distributions: tuple[InstalledDistribution, ...], +) -> None: + runtime = cast(dict[str, object], client["runtime"]) + runtime.update( + { + "python_abi": f"{interpreter_name()}{interpreter_version()}", + "python_executable": sys.executable, + "python_implementation": sys.implementation.name, + "python_version": sys.version.split()[0], + "distributions": [item.model_dump(mode="json") for item in runtime_distributions], + } + ) + runtime["runtime_sha256"] = compute_canonical_json_sha256( + {key: value for key, value in runtime.items() if key != "runtime_sha256"} + ) + dependencies = cast(dict[str, object], cast(dict[str, object], client["authored"])["dependencies"]) + dependencies["requirements"] = ["fake-data-designer-plugin==1.0.0"] + + +def _plugin_dependency_lock( + client: dict[str, object], + runtime_distributions: tuple[InstalledDistribution, ...], + wheel_path: Path, +) -> ResolvedDependencyLock: + runtime = cast(dict[str, object], client["runtime"]) + return ResolvedDependencyLock.model_validate( + { + "schema_version": 1, + "resolver_version": "resolver-1", + "python_abi": runtime["python_abi"], + "client_runtime_sha256": runtime["runtime_sha256"], + "authored_requirements": ("fake-data-designer-plugin==1.0.0",), + "authored_source": None, + "source": None, + "base_distributions": runtime_distributions, + "overlay_packages": ( + { + "name": "fake-data-designer-plugin", + "version": "1.0.0", + "artifact": { + "path": wheel_path.as_posix(), + "sha256": hashlib.sha256(wheel_path.read_bytes()).hexdigest(), + }, + }, + ), + } + ) + + +def _write_dependency_lock(client: dict[str, object], lock: ResolvedDependencyLock) -> None: + lock_path = Path(cast(dict[str, object], client["dependency_lock"])["path"]) + lock_path.parent.mkdir(parents=True, exist_ok=True) + lock_path.write_text(lock.serialize_json()) + cast(dict[str, object], client["dependency_lock"])["sha256"] = lock.compute_sha256() + + +def _build_plugin_wheel(source: Path, destination: Path) -> None: + """Package the installed-layout fixture as the locked wheel under test.""" + destination.parent.mkdir(parents=True, exist_ok=True) + dist_info = "fake_data_designer_plugin-1.0.0.dist-info" + members = { + path.relative_to(source).as_posix(): path.read_bytes() + for path in source.rglob("*") + if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc" + } + members[f"{dist_info}/WHEEL"] = ( + b"Wheel-Version: 1.0\nGenerator: data-designer-tests\nRoot-Is-Purelib: true\nTag: py3-none-any\n" + ) + record_path = f"{dist_info}/RECORD" + output = io.StringIO() + writer = csv.writer(output, lineterminator="\n") + for name in (*sorted(members), record_path): + writer.writerow((name, "", "")) + members[record_path] = output.getvalue().encode() + with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_DEFLATED) as archive: + for name, content in sorted(members.items()): + archive.writestr(name, content) + + +def _write_inventory_bootstrap( + distributions: tuple[InstalledDistribution, ...], + tmp_path: Path, +) -> Path: + """Present plan-recorded runtime metadata while imports use the editable checkout.""" + bootstrap_directory = tmp_path / "client-bootstrap" + image_metadata_directory = bootstrap_directory / "image-metadata" + image_metadata_directory.mkdir(parents=True) + for distribution in distributions: + metadata = image_metadata_directory / ( + f"{distribution.name.replace('-', '_')}-{distribution.version}.dist-info/METADATA" + ) + metadata.parent.mkdir() + metadata.write_text(f"Metadata-Version: 2.1\nName: {distribution.name}\nVersion: {distribution.version}\n") + (bootstrap_directory / "sitecustomize.py").write_text( + """from __future__ import annotations + +import importlib.metadata +import os +import sys + +_distributions = importlib.metadata.distributions + + +def _controlled_distributions(**kwargs): + paths = kwargs.get("path") + if paths is not None and len(paths) == 1 and paths[0].endswith("/client-env/site-packages"): + return _distributions(**kwargs) + paths = [os.environ["DATA_DESIGNER_TEST_IMAGE_METADATA"]] + paths.extend(path for path in sys.path if path.endswith("/client-env/site-packages")) + return _distributions(path=paths) + + +importlib.metadata.distributions = _controlled_distributions +""" + ) + return image_metadata_directory + + +def _use_controlled_image_inventory(monkeypatch: pytest.MonkeyPatch, image_metadata_directory: Path) -> None: + """Limit child-process inventory discovery to the plan runtime and installed overlay.""" + bootstrap_directory = image_metadata_directory.parent + python_path = os.environ.get("PYTHONPATH") + paths = (bootstrap_directory.as_posix(),) if python_path is None else (bootstrap_directory.as_posix(), python_path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(paths)) + monkeypatch.setenv("DATA_DESIGNER_TEST_IMAGE_METADATA", image_metadata_directory.as_posix()) + + +def _phase_arguments( + operation: str, + runtime_case: RuntimeCase, + runtime_root: Path | None = None, + manifest_path: Path | None = None, + *, + attempt_directory: Path | None = None, +) -> tuple[str, ...]: + arguments = ( + operation, + "--plan", + (runtime_case.workspace / "runs/run-single/resolved-plan.json").as_posix(), + "--attempt-dir", + (attempt_directory or runtime_case.context.attempt_directory).as_posix(), + ) + if operation == "prepare": + assert runtime_root is not None and manifest_path is not None + return ( + *arguments, + "--runtime-root", + runtime_root.as_posix(), + "--manifest", + manifest_path.as_posix(), + "--node-host", + "compute-001", + ) + return arguments diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py new file mode 100644 index 000000000..10c1314e6 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import base64 +import json +import os +import subprocess +import sys +import tarfile +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from pathlib import Path + +import pytest + +from data_designer.slurm.runtime import node_worker as runtime_node_worker +from data_designer.slurm.runtime.bundle import stage_runtime_bundle +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec, + NodeReadinessSpec, + NodeSpec, + NodeWorkerSpec, + decode_node_worker_spec, + encode_node_worker_spec, +) +from data_designer.slurm.runtime.node_worker import NodeProcessSupervisor + + +@pytest.fixture +def shared_gpu_spec() -> NodeWorkerSpec: + processes = tuple( + NodeProcessSpec( + f"lane-{index}", + ("serve",), + (0,), + 0, + NodeReadinessSpec(port=18000 + index, path="/health", deadline_seconds=1), + ) + for index in range(2) + ) + return NodeWorkerSpec(1, 1, None, (NodeSpec(0, "compute-001", (18000, 18001), processes),), allow_gpu_sharing=True) + + +@pytest.mark.parametrize( + "mutation", + ( + None, + "policy", + "boolean", + "missing_probe", + "remote_probe", + "unowned_port", + "bad_path", + "bad_deadline", + "bad_port_type", + "outside", + "unknown", + ), +) +def test_shared_gpu_transport_requires_explicit_policy_and_valid_probes( + shared_gpu_spec: NodeWorkerSpec, + mutation: str | None, +) -> None: + payload = json.loads(base64.urlsafe_b64decode(encode_node_worker_spec(shared_gpu_spec))) + process = payload["nodes"][0]["processes"][0] + if mutation == "policy": + del payload["allow_gpu_sharing"] + elif mutation == "boolean": + payload["allow_gpu_sharing"] = "true" + elif mutation == "missing_probe": + del process["readiness"] + elif mutation == "remote_probe": + process["readiness"]["host"] = "other-host" + elif mutation == "unowned_port": + process["readiness"]["port"] = 18002 + elif mutation == "bad_path": + process["readiness"]["path"] = "health" + elif mutation == "bad_deadline": + process["readiness"]["deadline_seconds"] = 0 + elif mutation == "bad_port_type": + process["readiness"]["port"] = "18000" + elif mutation == "outside": + process["gpu_indices"] = [1] + elif mutation == "unknown": + payload["unknown"] = True + encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + if mutation is None: + assert decode_node_worker_spec(encoded) == shared_gpu_spec + else: + with pytest.raises(ValueError): + decode_node_worker_spec(encoded) + + +@pytest.mark.parametrize("outcome", ("ready", "failed", "cancelled", "timeout")) +def test_shared_gpu_startup_waits_for_readiness_and_cleans_up( + shared_gpu_spec: NodeWorkerSpec, + monkeypatch: pytest.MonkeyPatch, + outcome: str, +) -> None: + children: list[_FakeProcess] = [] + events: list[str] = [] + clock = [0.0] + supervisor = NodeProcessSupervisor(environment={}) + + def start(*args: object, **kwargs: object) -> _FakeProcess: + del args + events.append("start") + environment = kwargs["env"] + assert isinstance(environment, dict) + assert environment["CUDA_VISIBLE_DEVICES"] == "GPU-remapped" + child = _FakeProcess(pid=40 + len(children), returncode=23 if children else None) + children.append(child) + return child + + def probe(*args: object, **kwargs: object) -> bool: + del args, kwargs + events.append("probe") + if outcome == "failed": + children[0].returncode = 23 + elif outcome == "cancelled": + supervisor.request_stop() + elif outcome == "timeout": + clock[0] = 2.0 + return outcome == "ready" + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + monkeypatch.setattr(runtime_node_worker.HttpReadinessProber, "is_ready", probe) + monkeypatch.setattr(runtime_node_worker.time, "monotonic", lambda: clock[0]) + monkeypatch.setattr( + runtime_node_worker, "_signal_process_group", lambda child, selected: setattr(child, "returncode", -15) + ) + if outcome == "timeout": + with pytest.raises(RuntimeError, match="readiness timed out"): + supervisor.run(shared_gpu_spec.nodes[0], ("GPU-remapped",)) + else: + assert supervisor.run(shared_gpu_spec.nodes[0], ("GPU-remapped",)) == (1 if outcome == "cancelled" else 23) + assert events == (["start", "probe", "start"] if outcome == "ready" else ["start", "probe"]) + assert supervisor.cleanup_complete + assert all(child.poll() is not None for child in children) + + +@pytest.mark.parametrize("visible", ("0", "7", "GPU-remapped")) +def test_shared_gpu_real_children_wait_for_health_and_receive_the_same_device( + shared_gpu_spec: NodeWorkerSpec, + tmp_path: Path, + unused_tcp_port: int, + visible: str, +) -> None: + first, second = tmp_path / "first", tmp_path / "second" + program = f""" +import os, time +from http.server import BaseHTTPRequestHandler, HTTPServer +from pathlib import Path + +class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.end_headers() + +Path({str(first)!r}).write_text(os.environ['CUDA_VISIBLE_DEVICES']) +time.sleep(0.2) +HTTPServer(('127.0.0.1', {unused_tcp_port}), Handler).serve_forever() +""" + follower = ( + "import os; from pathlib import Path; from urllib.request import urlopen; " + f"urlopen('http://127.0.0.1:{unused_tcp_port}/health',timeout=1).close(); " + f"Path({str(second)!r}).write_text(os.environ['CUDA_VISIBLE_DEVICES']); raise SystemExit(23)" + ) + head, tail = shared_gpu_spec.nodes[0].processes + assert head.readiness is not None + head = replace( + head, + command=(sys.executable, "-c", program), + readiness=replace(head.readiness, port=unused_tcp_port, deadline_seconds=3), + ) + tail = replace(tail, command=(sys.executable, "-c", follower)) + node = replace(shared_gpu_spec.nodes[0], ports=tuple(sorted((unused_tcp_port, 18001))), processes=(head, tail)) + spec = decode_node_worker_spec(encode_node_worker_spec(replace(shared_gpu_spec, nodes=(node,)))) + supervisor = NodeProcessSupervisor(environment={}) + assert supervisor.run(spec.nodes[0], (visible,)) == 23 + assert supervisor.cleanup_complete + assert first.read_text() == second.read_text() == visible + + +def test_bundled_node_worker_imports_without_client_packages(tmp_path: Path) -> None: + workspace, extracted = tmp_path / "workspace", tmp_path / "extracted" + workspace.mkdir(mode=0o700) + bundle = stage_runtime_bundle(workspace) + with tarfile.open(bundle.path) as archive: + archive.extractall(extracted, filter="data") + program = """ +import importlib.abc +import sys + +class NoClientPackages(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname in {'data_designer.config', 'data_designer.engine', 'data_designer.interface'}: + raise ModuleNotFoundError(fullname) + +sys.meta_path.insert(0, NoClientPackages()) +from data_designer.slurm.runtime import node_worker +assert 'data_designer.slurm.planning' not in sys.modules +print(node_worker.__file__, end='') +""" + result = subprocess.run( + (sys.executable, "-c", program), + env={**os.environ, "PYTHONPATH": str(extracted)}, + text=True, + capture_output=True, + ) + assert result.returncode == 0, result.stderr + assert result.stdout == str(extracted / "data_designer/slurm/runtime/node_worker.py") + + +def test_node_worker_spec_round_trips_and_rejects_untrusted_payload() -> None: + spec = _worker_spec((NodeProcessSpec("lane-0", ("true",), (0,), 0),)) + assert decode_node_worker_spec(encode_node_worker_spec(spec)) == spec + with pytest.raises(ValueError, match="invalid"): + decode_node_worker_spec("not-base64") + + +def test_node_preflight_rejects_missing_required_model_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + model_path = tmp_path / "missing-model" + spec = _worker_spec( + (NodeProcessSpec("lane-0", ("true",), (0,), 0),), + required_model_path=model_path.as_posix(), + ) + _set_node_environment(monkeypatch) + + assert runtime_node_worker.main(("preflight", "--spec", encode_node_worker_spec(spec))) == 70 + assert capsys.readouterr().err == "node worker failed at a validated runtime boundary\n" + + +def test_node_preflight_rejects_unreadable_required_model_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + model_path = tmp_path / "model" + model_path.mkdir() + spec = _worker_spec( + (NodeProcessSpec("lane-0", ("true",), (0,), 0),), + required_model_path=model_path.as_posix(), + ) + _set_node_environment(monkeypatch) + monkeypatch.setattr(runtime_node_worker.os, "access", lambda path, mode: False) + + assert runtime_node_worker.main(("preflight", "--spec", encode_node_worker_spec(spec))) == 70 + assert capsys.readouterr().err == "node worker failed at a validated runtime boundary\n" + + +def test_partial_startup_failure_cleans_already_started_lane(monkeypatch: pytest.MonkeyPatch) -> None: + process = _FakeProcess(pid=41) + calls = 0 + signals: list[int] = [] + cache_paths: list[Path] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + nonlocal calls + del arguments + environment = keywords["env"] + assert isinstance(environment, dict) + cache_paths.append(Path(environment["TIKTOKEN_RS_CACHE_DIR"])) + calls += 1 + if calls == 2: + raise OSError("injected partial startup") + return process + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + + def terminate(child: _FakeProcess, selected: int) -> None: + signals.append(selected) + child.returncode = -15 + + monkeypatch.setattr(runtime_node_worker, "_signal_process_group", terminate) + supervisor = NodeProcessSupervisor(environment={}) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("lane-0", ("first",), (0,), 0), + NodeProcessSpec("lane-1", ("second",), (1,), 0), + ), + ) + + with pytest.raises(OSError, match="partial startup"): + supervisor.run(node, ("0", "1")) + + assert signals + assert supervisor.cleanup_complete + assert len(cache_paths) == 2 + assert all(not path.exists() for path in cache_paths) + + +def test_node_worker_assigns_distinct_visible_gpus_to_two_tp4_replicas(monkeypatch: pytest.MonkeyPatch) -> None: + visible: list[str] = [] + cache_paths: list[Path] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + del arguments + environment = keywords["env"] + assert isinstance(environment, dict) + visible.append(environment["CUDA_VISIBLE_DEVICES"]) + cache_path = Path(environment["TIKTOKEN_RS_CACHE_DIR"]) + assert cache_path.is_dir() + cache_paths.append(cache_path) + return _FakeProcess(pid=40 + len(visible), returncode=1 if len(visible) == 2 else None) + + def terminate(child: _FakeProcess, selected: int) -> None: + del selected + child.returncode = -15 + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + monkeypatch.setattr(runtime_node_worker, "_signal_process_group", terminate) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("replica-0", ("first",), (0, 1, 2, 3), 0), + NodeProcessSpec("replica-1", ("second",), (4, 5, 6, 7), 0), + ), + ) + + assert NodeProcessSupervisor(environment={}).run(node, tuple(str(index) for index in range(8))) == 1 + assert visible == ["0,1,2,3", "4,5,6,7"] + assert len(set(cache_paths)) == 2 + assert all(not path.exists() for path in cache_paths) + + +def test_node_worker_preserves_configured_tokenizer_cache(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + cache_path = tmp_path / "prewarmed-cache" + cache_path.mkdir() + observed: list[str] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + del arguments + environment = keywords["env"] + assert isinstance(environment, dict) + observed.append(environment["TIKTOKEN_RS_CACHE_DIR"]) + return _FakeProcess(pid=41, returncode=1) + + monkeypatch.setattr(runtime_node_worker.subprocess, "Popen", start) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("replica-0", ("serve",), (0,), 0),), + ) + + assert NodeProcessSupervisor(environment={"TIKTOKEN_RS_CACHE_DIR": cache_path.as_posix()}).run(node, ("0",)) == 1 + assert observed == [cache_path.as_posix()] + assert cache_path.is_dir() + + +def test_follower_failure_terminates_sibling_without_orphan(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + long_running = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + failing = "import time; time.sleep(0.2); raise SystemExit(23)" + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=( + NodeProcessSpec("head", (sys.executable, "-c", long_running), (0,), 0), + NodeProcessSpec("follower", (sys.executable, "-c", failing), (1,), 0), + ), + ) + supervisor = NodeProcessSupervisor(environment={}) + + assert supervisor.run(node, ("0", "1")) == 23 + supervisor.cleanup() + + assert ready.exists() + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_request_stops_all_node_lanes(tmp_path: Path) -> None: + ready = tmp_path / "ready" + stopped = tmp_path / "stopped" + child = ( + "import signal,time; from pathlib import Path; " + f"ready=Path({ready.as_posix()!r}); stopped=Path({stopped.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(30)" + ) + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("lane", (sys.executable, "-c", child), (0,), 0),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + _wait_for_file(ready) + supervisor.request_stop() + assert result.result(timeout=5) == 1 + + assert stopped.read_text() == "stopped" + assert supervisor.cleanup_complete + + +def test_cancellation_interrupts_a_pending_launch_delay() -> None: + node = NodeSpec( + node_index=0, + host="compute-001", + ports=(), + processes=(NodeProcessSpec("delayed", ("must-not-start",), (0,), 30),), + ) + supervisor = NodeProcessSupervisor(environment={}) + with ThreadPoolExecutor(max_workers=1) as executor: + result = executor.submit(supervisor.run, node, ("0",)) + time.sleep(0.05) + supervisor.request_stop() + assert result.result(timeout=2) == 1 + + assert supervisor.cleanup_complete + + +@dataclass(slots=True) +class _FakeProcess: + pid: int + returncode: int | None = None + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + del timeout + self.returncode = -15 + return self.returncode + + +def _worker_spec( + processes: tuple[NodeProcessSpec, ...], + *, + required_model_path: str | None = None, +) -> NodeWorkerSpec: + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=8, + required_model_path=required_model_path, + nodes=(NodeSpec(node_index=0, host="compute-001", ports=(18000,), processes=processes),), + ) + + +def _set_node_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("SLURM_NTASKS", "1") + monkeypatch.setenv("SLURMD_NODENAME", "compute-001") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + + +def _wait_for_file(path: Path) -> None: + deadline = time.monotonic() + 3 + while not path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert path.exists() diff --git a/packages/data-designer-slurm/tests/runtime/test_paths.py b/packages/data-designer-slurm/tests/runtime/test_paths.py new file mode 100644 index 000000000..7e55a3806 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_paths.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from conftest import RuntimeCase + +from data_designer.slurm.config import ContainerMount +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.paths import get_container_path, get_host_path + + +def test_container_path_uses_most_specific_mount_and_preserves_relative_path(runtime_case: RuntimeCase) -> None: + workspace = runtime_case.workspace.as_posix() + nested = f"{workspace}/runs" + plan = runtime_case.context.plan.model_copy( + update={ + "container_mounts": ( + ContainerMount(source=workspace, target="/container/workspace"), + ContainerMount(source=nested, target="/container/runs"), + ) + } + ) + + mapped = get_container_path(plan, f"{nested}/run-single/resolved-plan.json") + + assert mapped == "/container/runs/run-single/resolved-plan.json" + assert get_host_path(plan, mapped) == f"{nested}/run-single/resolved-plan.json" + + +def test_container_path_rejects_unmounted_and_read_only_writes(runtime_case: RuntimeCase) -> None: + plan = runtime_case.context.plan.model_copy( + update={ + "container_mounts": ( + ContainerMount( + source=runtime_case.workspace.as_posix(), + target="/container/workspace", + read_only=True, + ), + ) + } + ) + + with pytest.raises(SlurmRuntimeError, match="writable"): + get_container_path(plan, runtime_case.context.attempt_directory.as_posix(), require_writable=True) + with pytest.raises(SlurmRuntimeError, match="not available"): + get_container_path(plan, "/different/root/file.json") diff --git a/packages/data-designer-slurm/tests/runtime/test_ports.py b/packages/data-designer-slurm/tests/runtime/test_ports.py new file mode 100644 index 000000000..337d7cee5 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_ports.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from conftest import RuntimeCase + +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.ports import allocation_ports, resolve_allocation_deployments + + +def test_allocation_ports_are_deterministic_and_isolated_by_gpu(runtime_case: RuntimeCase) -> None: + first = allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) + second = allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "1"}) + + assert first == allocation_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) + assert set(first).isdisjoint(second) + assert all(10000 <= port < 10256 for port in first) + assert all(10256 <= port < 10512 for port in second) + + +def test_allocation_deployment_uses_remapped_ports(runtime_case: RuntimeCase) -> None: + environment = {"SLURM_JOB_GPUS": "0"} + deployment = resolve_allocation_deployments(runtime_case.context, environment)[0] + ports = set(allocation_ports(runtime_case.context, environment)) + + assert deployment.logical_endpoint.port in ports + assert {backend.port for backend in deployment.backend_endpoints} <= ports + assert {probe.port for probe in deployment.readiness_probes} <= ports + assert {process.http_port for process in deployment.processes if process.http_port is not None} <= ports + + +def test_allocation_ports_skip_client_otel_port(runtime_case: RuntimeCase) -> None: + environment = {"SLURM_JOB_GPUS": "0"} + otel_port = allocation_ports(runtime_case.context, environment)[0] + invocation = runtime_case.context.plan.invocation.model_copy( + update={"effective_run_config": {"otel_metrics_port": otel_port}} + ) + plan = runtime_case.context.plan.model_copy(update={"invocation": invocation}) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=runtime_case.context.attempt_directory, + ) + + assert otel_port not in allocation_ports(context, environment) + + +@pytest.mark.parametrize("environment", ({}, {"SLURM_JOB_GPUS": ""}, {"SLURM_JOB_GPUS": "0,0"})) +def test_allocation_ports_reject_invalid_gpu_ids( + runtime_case: RuntimeCase, + environment: dict[str, str], +) -> None: + with pytest.raises(SlurmRuntimeError, match="SLURM_JOB_GPUS"): + allocation_ports(runtime_case.context, environment) + + +def test_visible_gpu_mode_keeps_planned_ports(runtime_case: RuntimeCase) -> None: + profile = runtime_case.context.plan.selected_profile.profile.model_copy(update={"gpu_request_mode": "visible"}) + selected_profile = runtime_case.context.plan.selected_profile.model_copy(update={"profile": profile}) + plan = runtime_case.context.plan.model_copy(update={"selected_profile": selected_profile}) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=runtime_case.context.attempt_directory, + ) + planned = tuple(port.port for port in plan.client.ports) + tuple( + port.port for deployment in plan.deployments for port in deployment.ports + ) + + assert allocation_ports(context, {}) == planned diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py new file mode 100644 index 000000000..b9dcf5bfe --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import pytest +from conftest import RuntimeCase + +from data_designer.slurm.config import ContainerMount +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.runtime import preflight as runtime_preflight +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.preflight import AllocationLayout, SystemAllocationPreflight, _verify_artifact + + +def test_allocation_layout_rejects_invalid_hosts_and_indices() -> None: + layout = AllocationLayout(("compute-001", "compute-002")) + + with pytest.raises(SlurmRuntimeError, match="outside"): + layout.get_host(-1) + with pytest.raises(SlurmRuntimeError, match="identity"): + AllocationLayout(("compute-001", "--invalid")) + + +def test_scheduler_preflight_accepts_complete_multi_node_plan() -> None: + context = cast( + AllocationContext, + SimpleNamespace( + plan=SimpleNamespace( + client=SimpleNamespace(host_node_index=0), + deployments=(SimpleNamespace(node_indices=(0, 1)), SimpleNamespace(node_indices=(2,))), + resolved_gpus_per_node=8, + ), + attempt=SimpleNamespace(scheduler=SimpleNamespace(array_job_id=4101)), + shard=SimpleNamespace(array_task_index=0), + ), + ) + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "3", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", + } + + SystemAllocationPreflight._verify_scheduler(context, environment) + + +def test_scheduler_preflight_accepts_exact_one_node_gpu_shape(runtime_case: RuntimeCase) -> None: + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", + } + + SystemAllocationPreflight._verify_scheduler(runtime_case.context, environment) + + +def test_scheduler_preflight_accepts_exclusive_allocation_gpu_superset(runtime_case: RuntimeCase) -> None: + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7,8", + } + + SystemAllocationPreflight._verify_scheduler(runtime_case.context, environment) + + +@pytest.mark.parametrize( + ("name", "value"), + ( + ("SLURM_ARRAY_TASK_ID", "1"), + ("SLURM_ARRAY_JOB_ID", "4102"), + ("SLURM_JOB_NUM_NODES", "2"), + ("SLURM_NODEID", "1"), + ("CUDA_VISIBLE_DEVICES", "0,1"), + ), +) +def test_scheduler_preflight_rejects_mismatched_allocation_fact( + runtime_case: RuntimeCase, + name: str, + value: str, +) -> None: + environment = { + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_NODEID": "0", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7", + } + environment[name] = value + + with pytest.raises(SlurmRuntimeError, match="does not match|GPU visibility"): + SystemAllocationPreflight._verify_scheduler(runtime_case.context, environment) + + +def test_artifact_verification_checks_exact_bytes_and_rejects_symlink(tmp_path: Path) -> None: + artifact = tmp_path / "artifact.bin" + content = b"reviewed artifact" + artifact.write_bytes(content) + reference = ArtifactReference(path=artifact.as_posix(), sha256=hashlib.sha256(content).hexdigest()) + + _verify_artifact(reference) + artifact.write_bytes(b"changed") + with pytest.raises(OSError, match="digest"): + _verify_artifact(reference) + + target = tmp_path / "target.bin" + target.write_bytes(content) + artifact.unlink() + artifact.symlink_to(target) + with pytest.raises(OSError, match="regular file"): + _verify_artifact(reference) + + +@pytest.fixture +def artifact_context(tmp_path: Path) -> AllocationContext: + artifact = tmp_path / "runtime.tar.gz" + artifact.write_bytes(b"runtime") + reference = ArtifactReference(path=artifact.as_posix(), sha256=hashlib.sha256(b"runtime").hexdigest()) + image = tmp_path / "serving.sqsh" + content = b"serving image" + image.write_bytes(content) + image_reference = ArtifactReference(path=image.as_posix(), sha256=hashlib.sha256(content).hexdigest()) + plan_path = tmp_path / "resolved-plan.json" + plan_path.write_text("{}") + plan = SimpleNamespace( + authored_config=SimpleNamespace(path=(tmp_path / "run.yaml").as_posix()), + compute_sha256=lambda: hashlib.sha256(b"{}").hexdigest(), + runtime_bundle=reference, + client=SimpleNamespace(dependency_lock=reference), + deployments=(SimpleNamespace(image=image_reference), SimpleNamespace(image=image_reference)), + builder=SimpleNamespace(source=reference), + ) + return cast(AllocationContext, SimpleNamespace(plan=plan, shard=SimpleNamespace(input_partition=reference))) + + +def test_artifact_preflight_accepts_shared_artifacts(artifact_context: AllocationContext) -> None: + SystemAllocationPreflight.verify_artifacts(artifact_context) + + +@pytest.mark.parametrize("damage", ("contents", "missing", "symlink")) +def test_artifact_preflight_rejects_changed_serving_image( + artifact_context: AllocationContext, + damage: str, + tmp_path: Path, +) -> None: + image = Path(artifact_context.plan.deployments[0].image.path) + if damage == "contents": + image.write_bytes(b"changed image") + else: + image.unlink() + if damage == "symlink": + target = tmp_path / "target.sqsh" + target.write_bytes(b"serving image") + image.symlink_to(target) + + with pytest.raises(SlurmRuntimeError, match="could not verify launch inputs"): + SystemAllocationPreflight.verify_artifacts(artifact_context) + + +def test_artifact_verification_rejects_path_replacement_during_read( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + artifact = tmp_path / "artifact.bin" + content = b"reviewed artifact" + artifact.write_bytes(content) + reference = ArtifactReference(path=artifact.as_posix(), sha256=hashlib.sha256(content).hexdigest()) + replacement = tmp_path / "replacement.bin" + replacement.write_bytes(b"different artifact") + original_fstat = runtime_preflight.os.fstat + fstat_calls = 0 + replaced = False + + def replace_after_descriptor_validation(descriptor: int) -> os.stat_result: + nonlocal fstat_calls, replaced + status = original_fstat(descriptor) + fstat_calls += 1 + if fstat_calls == 2: + replaced = True + os.replace(replacement, artifact) + return status + + monkeypatch.setattr(runtime_preflight.os, "fstat", replace_after_descriptor_validation) + with pytest.raises(OSError, match="replaced while it was verified"): + _verify_artifact(reference) + assert replaced + + +def test_port_preflight_detects_collision_before_launch( + runtime_case: RuntimeCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _UnavailableSocket: + closed: bool = False + + def __init__(self, *arguments: object) -> None: + del arguments + + def setsockopt(self, *arguments: object) -> None: + del arguments + + def bind(self, address: tuple[str, int]) -> None: + del address + raise OSError("injected port collision") + + def close(self) -> None: + type(self).closed = True + + monkeypatch.setattr("data_designer.slurm.runtime.preflight.socket.socket", _UnavailableSocket) + + with pytest.raises(SlurmRuntimeError, match="ports are unavailable"): + SystemAllocationPreflight.verify_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) + assert _UnavailableSocket.closed + + +def test_attempt_directory_must_be_restrictive(runtime_case: RuntimeCase) -> None: + runtime_case.context.attempt_directory.chmod(0o755) + + with pytest.raises(SlurmRuntimeError, match="restrictive directory"): + SystemAllocationPreflight.verify_attempt_directory(runtime_case.context.attempt_directory) + + +def test_client_runtime_preflight_accepts_exact_resolved_runtime( + runtime_case: RuntimeCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + runtime_preflight.ClientRuntimeInspector, + "inspect", + lambda self: runtime_case.context.plan.client.runtime, + ) + + SystemAllocationPreflight.verify_client_runtime(runtime_case.context) + + +def test_client_runtime_preflight_rejects_compute_node_drift( + runtime_case: RuntimeCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = runtime_case.context.plan.client.runtime.model_copy(update={"python_version": "3.12.1"}) + monkeypatch.setattr(runtime_preflight.ClientRuntimeInspector, "inspect", lambda self: runtime) + + with pytest.raises(SlurmRuntimeError, match="differs from the resolved plan"): + SystemAllocationPreflight.verify_client_runtime(runtime_case.context) + + +def test_preflight_uses_host_paths_without_container_translation( + runtime_case: RuntimeCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + logical_attempt = runtime_case.context.attempt_directory + container_workspace = tmp_path / "container-workspace" + plan = runtime_case.context.plan.model_copy( + update={ + "container_mounts": ( + ContainerMount( + source=runtime_case.workspace.as_posix(), + target=container_workspace.as_posix(), + ), + ) + } + ) + context = AllocationContext( + plan=plan, + shard=runtime_case.context.shard, + attempt=runtime_case.context.attempt, + attempt_directory=logical_attempt, + ) + verified_paths: list[str] = [] + monkeypatch.setattr(SystemAllocationPreflight, "_verify_scheduler", lambda *args: None) + monkeypatch.setattr(SystemAllocationPreflight, "verify_client_runtime", lambda *args: None) + monkeypatch.setattr(SystemAllocationPreflight, "verify_ports", lambda *args: None) + monkeypatch.setattr(runtime_preflight, "_verify_artifact", lambda reference: verified_paths.append(reference.path)) + + SystemAllocationPreflight().verify(context, {}) + + assert verified_paths + assert any(path.startswith(runtime_case.workspace.as_posix()) for path in verified_paths) + assert not any(path.startswith(container_workspace.as_posix()) for path in verified_paths) diff --git a/packages/data-designer-slurm/tests/runtime/test_probes.py b/packages/data-designer-slurm/tests/runtime/test_probes.py new file mode 100644 index 000000000..4b49b5100 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_probes.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from data_designer.slurm.runtime import probes +from data_designer.slurm.runtime.probes import HttpReadinessProber + + +def test_readiness_probe_does_not_buffer_response_body(monkeypatch: pytest.MonkeyPatch) -> None: + connection = _FakeConnection() + monkeypatch.setattr(probes.http.client, "HTTPConnection", lambda *args, **kwargs: connection) + + assert HttpReadinessProber().is_ready("127.0.0.1", 8000, "/health", timeout_seconds=1.0) + + assert connection.requested == ("GET", "/health", {"Connection": "close"}) + assert connection.closed + + +class _FakeConnection: + def __init__(self) -> None: + self.requested: tuple[str, str, dict[str, str]] | None = None + self.closed = False + + def request(self, method: str, path: str, *, headers: dict[str, str]) -> None: + self.requested = (method, path, headers) + + def getresponse(self) -> SimpleNamespace: + return SimpleNamespace(status=200) + + def close(self) -> None: + self.closed = True diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py new file mode 100644 index 000000000..a8afe7af5 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -0,0 +1,532 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import asyncio +import socket +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager + +import pytest +from aiohttp import ClientSession, ClientTimeout, TCPConnector, web +from aiohttp.test_utils import TestServer + +from data_designer.slurm.runtime import proxy as runtime_proxy +from data_designer.slurm.runtime.proxy import _Backend, _BackendPool, _parse_backend, _ProxyApplication + +_Handler = Callable[[web.Request], Awaitable[web.StreamResponse]] + + +@asynccontextmanager +async def _serve(application: web.Application) -> AsyncIterator[TestServer]: + server = TestServer(application) + await server.start_server() + try: + yield server + finally: + await server.close() + + +def _application(handler: _Handler) -> web.Application: + application = web.Application() + application.router.add_route("*", "/{path:.*}", handler) + return application + + +def _backend(server: TestServer) -> _Backend: + assert server.host is not None + assert server.port is not None + return _Backend(server.host, server.port) + + +@asynccontextmanager +async def _serve_production( + backend: _Backend, + monkeypatch: pytest.MonkeyPatch, +) -> AsyncIterator[tuple[str, TCPConnector]]: + runner: web.AppRunner | None = None + connector: TCPConnector | None = None + + def run_app(application: web.Application, **options: object) -> None: + nonlocal runner + for name in ("host", "port", "print"): + options.pop(name) + runner = web.AppRunner(application, **options) + + def create_connector(**options: object) -> TCPConnector: + nonlocal connector + connector = TCPConnector(**options) + return connector + + monkeypatch.setattr(runtime_proxy.web, "run_app", run_app) + monkeypatch.setattr(runtime_proxy, "TCPConnector", create_connector) + assert runtime_proxy.main(["--listen-port", "1", "--backend", backend.origin]) == 0 + assert runner is not None + await runner.setup() + try: + await web.TCPSite(runner, "127.0.0.1", 0).start() + assert connector is not None + yield f"http://127.0.0.1:{runner.addresses[0][1]}", connector + finally: + await runner.cleanup() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", (False, True), ids=("before-headers", "between-sse-chunks")) +async def test_production_proxy_disconnect_releases_upstream( + monkeypatch: pytest.MonkeyPatch, + streaming: bool, +) -> None: + entered, closed, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + transport: asyncio.Transport | None = None + + async def delayed(request: web.Request) -> web.StreamResponse: + nonlocal transport + if request.path == "/complete": + return web.Response(text="complete") + transport = request.transport + response = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) + try: + if streaming: + await response.prepare(request) + await response.write(b"data: first\n\n") + entered.set() + await release.wait() + if streaming: + await response.write(b"data: second\n\n") + return response + finally: + closed.set() + + async with _serve(_application(delayed)) as backend: + async with _serve_production(_backend(backend), monkeypatch) as (origin, connector): + port = int(origin.rsplit(":", 1)[1]) + reader, writer = await asyncio.open_connection("127.0.0.1", port) + try: + writer.write(b"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + await asyncio.wait_for(entered.wait(), 2) + if streaming: + await asyncio.wait_for(reader.readuntil(b"data: first\n\n"), 2) + assert len(connector._acquired) == 1 + writer.close() + await writer.wait_closed() + await asyncio.wait_for(closed.wait(), 2) + assert transport is not None and transport.is_closing() + async with ClientSession() as client: + async with client.get(f"{origin}/metrics") as response: + metrics = await response.json() + assert metrics["backends"][0]["active_requests"] == 0 + assert not connector._acquired + assert metrics["connections"]["idle"] == 0 + async with client.get(f"{origin}/complete") as response: + assert response.status == 200 + assert await response.text() == "complete" + finally: + writer.close() + release.set() + + +@pytest.fixture +def total_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + def timeout(**options: object) -> ClientTimeout: + return ClientTimeout(**(options | {"total": 0.05, "sock_read": 1.0})) + + monkeypatch.setattr(runtime_proxy, "ClientTimeout", timeout) + + +@pytest.mark.asyncio +async def test_proxy_total_timeout_fails_over_and_classifies_timeout(total_timeout: None) -> None: + calls: list[str] = [] + + async def slow(request: web.Request) -> web.Response: + calls.append("slow") + await asyncio.sleep(1) + return web.Response(text="late") + + async def available(request: web.Request) -> web.Response: + calls.append("spare") + return web.Response(text="complete") + + async with _serve(_application(slow)) as first, _serve(_application(available)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 200 + assert await response.text() == "complete" + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + assert metrics["retries"] == {"timeout": 1} + assert metrics["connections"]["failed"] == 1 + assert [item["active_requests"] for item in metrics["backends"]] == [0, 0] + assert calls == ["slow", "spare"] + + +@pytest.mark.asyncio +async def test_proxy_health_aggregates_total_timeout(total_timeout: None) -> None: + async def slow(request: web.Request) -> web.Response: + await asyncio.sleep(1) + return web.Response() + + async def healthy(request: web.Request) -> web.Response: + return web.Response() + + async with _serve(_application(slow)) as first, _serve(_application(healthy)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/health")) as response: + assert response.status == 503 + assert await response.json() == {"status": "unavailable", "backends_ready": 1, "backends": 2} + async with client.get(endpoint.make_url("/metrics")) as response: + assert all(item["active_requests"] == 0 for item in (await response.json())["backends"]) + + +@pytest.mark.asyncio +async def test_proxy_retries_overload_on_another_least_active_backend() -> None: + calls: list[str] = [] + + async def overloaded(request: web.Request) -> web.Response: + calls.append("overloaded") + assert await request.read() == b"{}" + return web.Response(status=429, text="busy") + + async def available(request: web.Request) -> web.Response: + calls.append("available") + assert await request.read() == b"{}" + return web.json_response({"status": "complete"}) + + async with _serve(_application(overloaded)) as first, _serve(_application(available)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), data=b"{}") as response: + assert response.status == 200 + assert await response.json() == {"status": "complete"} + + assert calls == ["overloaded", "available"] + + +@pytest.mark.asyncio +async def test_proxy_limits_failover_to_three_distinct_backends_and_preserves_final_retry_after() -> None: + calls: list[int] = [] + + def build_handler(index: int) -> _Handler: + async def overloaded(request: web.Request) -> web.Response: + del request + calls.append(index) + return web.Response(status=429, text=f"busy-{index}", headers={"Retry-After": "99"}) + + return overloaded + + applications = tuple(_application(build_handler(index)) for index in range(4)) + async with ( + _serve(applications[0]) as first, + _serve(applications[1]) as second, + _serve(applications[2]) as third, + _serve(applications[3]) as fourth, + ): + backends = tuple(_backend(server) for server in (first, second, third, fourth)) + proxy = _ProxyApplication(backends, 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 429 + assert await response.text() == f"busy-{calls[-1]}" + assert response.headers["Retry-After"] == "99" + + assert calls == [0, 1, 2] + + +@pytest.mark.asyncio +async def test_proxy_adds_configured_retry_after_when_final_backend_omits_it() -> None: + async def overloaded(request: web.Request) -> web.Response: + del request + return web.Response(status=429, text="busy") + + async with _serve(_application(overloaded)) as backend: + proxy = _ProxyApplication((_backend(backend),), 3) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 429 + assert response.headers["Retry-After"] == "3" + + +@pytest.mark.asyncio +async def test_proxy_rejects_requests_over_the_configured_limit(monkeypatch: pytest.MonkeyPatch) -> None: + calls = 0 + + async def available(request: web.Request) -> web.Response: + nonlocal calls + calls += 1 + return web.Response(text="complete") + + monkeypatch.setattr(runtime_proxy, "_MAXIMUM_REQUEST_BYTES", 4) + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), data=b"12345") as response: + assert response.status == 413 + + assert calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", (500, 502, 503, 504)) +async def test_proxy_retries_transient_http_failures(status: int) -> None: + async def unavailable(request: web.Request) -> web.Response: + del request + return web.Response(status=status) + + async def available(request: web.Request) -> web.Response: + del request + return web.Response(text="complete") + + async with _serve(_application(unavailable)) as first, _serve(_application(available)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 200 + assert await response.text() == "complete" + + +@pytest.mark.asyncio +async def test_proxy_retries_connection_failure_on_another_backend() -> None: + unavailable = socket.socket() + unavailable.bind(("127.0.0.1", 0)) + unavailable_port = unavailable.getsockname()[1] + unavailable.close() + + async def available(request: web.Request) -> web.Response: + del request + return web.Response(text="complete") + + async with _serve(_application(available)) as server: + proxy = _ProxyApplication((_Backend("127.0.0.1", unavailable_port), _backend(server)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 200 + assert await response.text() == "complete" + + +@pytest.mark.asyncio +async def test_proxy_preserves_overload_after_final_backend_connection_failure() -> None: + unavailable = socket.socket() + unavailable.bind(("127.0.0.1", 0)) + unavailable_port = unavailable.getsockname()[1] + unavailable.close() + + async def overloaded(request: web.Request) -> web.Response: + del request + return web.Response(status=429, headers={"Retry-After": "17"}) + + async with _serve(_application(overloaded)) as server: + proxy = _ProxyApplication((_backend(server), _Backend("127.0.0.1", unavailable_port)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 429 + assert response.headers["Retry-After"] == "17" + + +@pytest.mark.asyncio +async def test_proxy_keeps_earlier_retry_after_when_later_overload_omits_it() -> None: + unavailable = socket.socket() + unavailable.bind(("127.0.0.1", 0)) + unavailable_port = unavailable.getsockname()[1] + unavailable.close() + + async def first_overload(request: web.Request) -> web.Response: + del request + return web.Response(status=429, headers={"Retry-After": "17"}) + + async def second_overload(request: web.Request) -> web.Response: + del request + return web.Response(status=429) + + async with _serve(_application(first_overload)) as first, _serve(_application(second_overload)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second), _Backend("127.0.0.1", unavailable_port)), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.post(endpoint.make_url("/v1/chat/completions"), json={}) as response: + assert response.status == 429 + assert response.headers["Retry-After"] == "17" + + +@pytest.mark.asyncio +async def test_proxy_streams_response_without_waiting_for_completion() -> None: + release_tail = asyncio.Event() + + async def streaming(request: web.Request) -> web.StreamResponse: + response = web.StreamResponse(status=200) + await response.prepare(request) + await response.write(b"first") + await release_tail.wait() + await response.write(b"second") + await response.write_eof() + return response + + async with _serve(_application(streaming)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/v1/models")) as response: + assert await asyncio.wait_for(response.content.readexactly(5), timeout=1) == b"first" + release_tail.set() + assert await response.read() == b"second" + + +@pytest.mark.asyncio +async def test_proxy_reuses_upstream_connections_and_reports_pool_metrics() -> None: + transports: set[object] = set() + + async def available(request: web.Request) -> web.Response: + assert request.transport is not None + transports.add(request.transport) + return web.Response(text="complete") + + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + for _ in range(2): + async with client.get(endpoint.make_url("/v1/models")) as response: + assert await response.text() == "complete" + async with client.get(endpoint.make_url("/metrics")) as response: + metrics = await response.json() + + assert len(transports) == 1 + assert metrics["connections"]["opened"] == 1 + assert metrics["connections"]["reused"] == 1 + assert metrics["connections"]["idle"] == 1 + + +@pytest.mark.asyncio +async def test_proxy_waits_for_a_pooled_connection_without_timing_out(monkeypatch: pytest.MonkeyPatch) -> None: + first_started = asyncio.Event() + release_first = asyncio.Event() + requests = 0 + + async def available(request: web.Request) -> web.Response: + del request + nonlocal requests + requests += 1 + if requests == 1: + first_started.set() + await release_first.wait() + return web.Response(text="complete") + + monkeypatch.setattr(runtime_proxy, "_MAXIMUM_CONNECTIONS", 1) + monkeypatch.setattr(runtime_proxy, "_CONNECT_TIMEOUT_SECONDS", 0.05) + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + first = asyncio.create_task(client.get(endpoint.make_url("/first"))) + await first_started.wait() + second = asyncio.create_task(client.get(endpoint.make_url("/second"))) + await asyncio.sleep(0.1) + release_first.set() + async with await first as response: + assert response.status == 200 + assert await response.text() == "complete" + async with await second as response: + assert response.status == 200 + assert await response.text() == "complete" + + assert requests == 2 + + +@pytest.mark.asyncio +async def test_proxy_releases_backend_when_upstream_request_is_cancelled(monkeypatch: pytest.MonkeyPatch) -> None: + proxy = _ProxyApplication((_Backend("127.0.0.1", 8001),), 1) + + class _Request: + async def read(self) -> bytes: + return b"" + + async def cancelled(*args: object) -> None: + del args + raise asyncio.CancelledError + + monkeypatch.setattr(proxy, "_request_backend", cancelled) + with pytest.raises(asyncio.CancelledError): + await proxy._forward(_Request()) # type: ignore[arg-type] + + assert proxy.pool.active_counts() == (0,) + + +@pytest.mark.asyncio +async def test_proxy_health_requires_every_backend_to_be_ready() -> None: + async def healthy(request: web.Request) -> web.Response: + return web.Response(status=200 if request.path == "/ready" else 404) + + async def unhealthy(request: web.Request) -> web.Response: + return web.Response(status=503 if request.path == "/ready" else 404) + + async with _serve(_application(healthy)) as first, _serve(_application(unhealthy)) as second: + proxy = _ProxyApplication((_backend(first), _backend(second)), 1, health_path="/ready") + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get(endpoint.make_url("/health")) as response: + assert response.status == 503 + assert await response.json() == {"status": "unavailable", "backends_ready": 1, "backends": 2} + + +@pytest.mark.asyncio +async def test_proxy_strips_connection_nominated_headers_in_both_directions() -> None: + forwarded_headers: dict[str, str] = {} + + async def available(request: web.Request) -> web.Response: + forwarded_headers.update(request.headers) + return web.Response( + text="complete", + headers={"Connection": "X-Backend-Hop", "X-Backend-Hop": "private"}, + ) + + async with _serve(_application(available)) as backend: + proxy = _ProxyApplication((_backend(backend),), 1) + async with _serve(proxy.create()) as endpoint, ClientSession() as client: + async with client.get( + endpoint.make_url("/v1/models"), + headers={ + "Authorization": "Bearer reviewed", + "Connection": "X-Client-Hop", + "X-Client-Hop": "private", + }, + ) as response: + assert await response.text() == "complete" + assert "X-Backend-Hop" not in response.headers + + assert forwarded_headers["Authorization"] == "Bearer reviewed" + assert "X-Client-Hop" not in forwarded_headers + + +@pytest.mark.parametrize( + "value", + ( + "https://127.0.0.1:8000", + "http://example.com:8000", + "http://127.0.0.1:70000", + "http://user@127.0.0.1:8000", + "http://127.0.0.1:8000/path", + ), +) +def test_proxy_rejects_non_loopback_or_malformed_backend(value: str) -> None: + with pytest.raises(argparse.ArgumentTypeError, match="backend"): + _parse_backend(value) + + +def test_proxy_matches_allowed_backend_hosts_case_insensitively() -> None: + backend = _parse_backend("http://compute-001:8000", frozenset({"Compute-001"})) + + assert backend == _Backend("compute-001", 8000) + + +def test_pool_selects_least_active_backend_deterministically() -> None: + pool = _BackendPool( + (_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002), _Backend("127.0.0.1", 8003)), + 1, + ) + first = pool.acquire(frozenset()) + second = pool.acquire(frozenset()) + pool.release(first) + third = pool.acquire(frozenset()) + pool.release(second) + pool.release(third) + + assert (first, second, third) == (0, 1, 2) + assert pool.active_counts() == (0, 0, 0) diff --git a/packages/data-designer-slurm/tests/runtime/test_shared_node_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shared_node_runtime.py new file mode 100644 index 000000000..b07850c27 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_shared_node_runtime.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import shlex +import subprocess +import sys +import time +from dataclasses import replace +from pathlib import Path + +import pytest +from conftest import RuntimeCase, relocate_plan + +from data_designer.slurm.config import DataDesignerSlurmConfig, injected_profile +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.runtime.bootstrap import build_runtime_manifest +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec, + NodeSpec, + NodeWorkerSpec, + decode_node_worker_spec, + encode_node_worker_spec, +) +from data_designer.slurm.runtime.node_worker import NodeProcessSupervisor, main +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.steps import build_vllm_steps +from data_designer.slurm.serving.resolver import resolve_vllm_server +from data_designer.slurm.state import RetryPlan, RetryShard + + +@pytest.mark.parametrize(("shared_gpu", "replicas"), ((False, 1), (True, 1), (True, 2))) +def test_shared_worker_has_both_models_readiness_and_gpu_free_clients( + runtime_case: RuntimeCase, + shared_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, + shared_node_config: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + shared_gpu: bool, + replicas: int, +) -> None: + if shared_gpu: + payload = shared_node_config.model_dump(mode="python") + for index, deployment in enumerate(payload["deployments"]): + deployment["resources"].update( + gpus_per_node=replicas if index == 0 else 1, gpu_start_index=0 if index == 0 else replicas - 1 + ) + deployment["topology"]["tensor_parallel"] = 1 + shared_node_plan = SlurmRunCompiler.compile( + resolve_slurm_config( + DataDesignerSlurmConfig.model_validate(payload), + selected_profile=injected_profile( + shared_node_plan.selected_profile.profile.model_copy(update={"gpus_per_node": replicas}) + ), + client_runtime=shared_node_plan.client.runtime, + deployment_images=tuple(member.image for member in shared_node_plan.deployments), + dependency_lock=dependency_lock_single, + runtime_bundle=shared_node_plan.runtime_bundle, + run_id=shared_node_plan.run_id, + package_version=shared_node_plan.package_version, + ) + ) + gpu_budget = replicas if shared_gpu else 8 + plan = relocate_plan(shared_node_plan, runtime_case.workspace) + context = replace(runtime_case.context, plan=plan, shard=plan.shards[0]) + manifest = build_runtime_manifest( + context, + { + "SLURM_JOB_GPUS": ",".join(str(index + 7) for index in range(replicas)) + if shared_gpu + else "1,3,5,7,9,11,13,15" + }, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs", + layout=AllocationLayout(("compute-001",)), + ) + servers = [step for step in manifest.steps if step.role is RuntimeStepRole.SERVER] + assert len(servers) == 1 + server = servers[0] + assert server.gpu_indices == tuple(range(gpu_budget)) + assert server.gpu_binding == "none" + assert len(server.readiness) == (replicas + 1 if shared_gpu else 4) + assert len({probe.port for probe in server.readiness}) == len(server.readiness) + worker = decode_node_worker_spec(server.command[-1]) + processes = worker.nodes[0].processes + expected_gpus = [(0,), (0,)] if replicas == 1 else [(0,), (1,), (1,)] + assert [process.gpu_indices for process in processes] == ( + expected_gpus if shared_gpu else [(0, 1), (4, 5), (2, 3), (6, 7)] + ) + expected_delays = [0, 0] if replicas == 1 else [0, 0, 65] + assert [process.launch_delay_seconds for process in processes] == ( + expected_delays if shared_gpu else [0, 0, 65, 65] + ) + assert worker.allow_gpu_sharing is shared_gpu + assert all((process.readiness is not None) is shared_gpu for process in processes) + if shared_gpu: + assert [probe.port for probe in server.readiness] == [process.readiness.port for process in processes] + assert {process.command[2] for process in processes} == { + deployment.authored.model for deployment in plan.deployments + } + endpoints = [step for step in manifest.steps if step.role is RuntimeStepRole.ENDPOINT] + assert len(endpoints) == 2 + assert all(not step.gpu_indices for step in manifest.steps if step.execution == "native") + assert all(len(step.readiness) == 1 for step in endpoints) + client = manifest.steps[-1] + assert client.command.count("--endpoint") == 2 + assert any(argument.startswith("writer=http://") for argument in client.command) + assert any(argument.startswith("judge=http://") for argument in client.command) + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text(manifest.serialize_json()) + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + script = f""" +set -Eeuo pipefail +source {shlex.quote(str(runtime_root / "plan_reader.sh"))} +source {shlex.quote(str(runtime_root / "step_runner.sh"))} +DD_GPU_REQUEST_MODE=gres +DD_CONTAINER_MOUNTS= +DD_SCRATCH_ROOT=/tmp/data-designer-slurm-4101-0 +DD_SCRATCH_CONTAINER_ROOT=/run/data-designer-slurm +dd_read_step {shlex.quote(str(manifest_path))} {shlex.quote(server.step_id)} +dd_build_srun_command +printf '%s\n' "${{DD_SRUN_COMMAND[@]}}" +""" + result = subprocess.run(("bash", "-c", script), capture_output=True, text=True, check=True) + arguments = result.stdout.splitlines() + assert "--nodes=1" in arguments + assert f"--gpus-per-task={gpu_budget}" in arguments + assert "--gpu-bind=none" in arguments + assert "--overlap" in arguments + assert not any(argument.startswith("--exclusive") for argument in arguments) + + +@pytest.mark.parametrize("visible", (("1", "3", "5", "7"), ("0", "1", "2", "3"), ("GPU-a", "GPU-b", "GPU-c", "GPU-d"))) +def test_distinct_model_children_and_their_children_receive_disjoint_devices( + tmp_path: Path, visible: tuple[str, ...] +) -> None: + files = (tmp_path / "writer", tmp_path / "judge") + programs = [] + for path in files: + child = "import os; print(os.environ['CUDA_VISIBLE_DEVICES'])" + program = ( + "import os,signal,subprocess,sys,time; from pathlib import Path; " + f"value=subprocess.check_output([sys.executable,'-c',{child!r}],text=True); " + f"Path({str(path)!r}).write_text(os.environ['CUDA_VISIBLE_DEVICES']+'\\n'+value); " + "time.sleep(0.4); sys.exit(23)" + ) + programs.append((sys.executable, "-c", program)) + node = NodeSpec( + 0, + "compute-001", + (), + ( + NodeProcessSpec("writer", programs[0], (0, 1), 0), + NodeProcessSpec("judge", programs[1], (2, 3), 0), + ), + ) + supervisor = NodeProcessSupervisor(environment={"PATH": os.defpath}) + assert supervisor.run(node, visible) == 23 + assert supervisor.cleanup_complete + assert files[0].read_text().splitlines() == [",".join(visible[:2])] * 2 + assert files[1].read_text().splitlines() == [",".join(visible[2:])] * 2 + + +def test_shared_preflight_verifies_every_model_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + first = tmp_path / "first" + first.mkdir() + second = tmp_path / "second" + node = NodeSpec( + 0, + "compute-001", + (), + (NodeProcessSpec("writer", ("true",), (0,), 0), NodeProcessSpec("judge", ("true",), (1,), 0)), + ) + spec = NodeWorkerSpec(1, 2, None, (node,), (str(first), str(second))) + monkeypatch.setenv("SLURM_PROCID", "0") + monkeypatch.setenv("SLURM_NTASKS", "1") + monkeypatch.setenv("SLURMD_NODENAME", "compute-001") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + assert decode_node_worker_spec(encode_node_worker_spec(spec)) == spec + assert main(("preflight", "--spec", encode_node_worker_spec(spec))) == 70 + second.mkdir() + assert main(("preflight", "--spec", encode_node_worker_spec(spec))) == 0 + + +def test_model_failure_interrupts_sibling_launch_standoff() -> None: + node = NodeSpec( + 0, + "compute-001", + (), + ( + NodeProcessSpec("writer", (sys.executable, "-c", "raise SystemExit(23)"), (0,), 0), + NodeProcessSpec("judge", ("must-not-start",), (1,), 30), + ), + ) + supervisor = NodeProcessSupervisor(environment={}) + started = time.monotonic() + assert supervisor.run(node, ("0", "1")) == 23 + assert time.monotonic() - started < 2 + assert supervisor.cleanup_complete + + +def test_independent_steps_reject_shared_plan_at_entry( + shared_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, +) -> None: + server = resolve_vllm_server(shared_node_plan, shared_node_plan.deployments[0].deployment_id) + with pytest.raises(SlurmRuntimeError, match="coordinated"): + build_vllm_steps(server, shared_node_plan, tmp_path, {}, tmp_path) + + +def test_fresh_process_retry_preserves_shared_placement( + shared_node_plan: ResolvedSlurmRunPlan, + runtime_case: RuntimeCase, + tmp_path: Path, +) -> None: + plan = shared_node_plan + retry = RetryPlan( + schema_version=1, + retry_id="retry-0001", + run_id=plan.run_id, + created_at=runtime_case.created_at, + resolved_plan=ArtifactReference( + path=str(Path(plan.authored_config.path).with_name("resolved-plan.json")), sha256=plan.compute_sha256() + ), + planned_shards=( + RetryShard( + shard_id=plan.shards[0].shard_id, attempt_id="attempt-0002", attempt_ordinal=2, array_task_index=0 + ), + ), + effective_resume_mode="never", + ) + plan_path, retry_path = tmp_path / "plan.json", tmp_path / "retry.json" + plan_path.write_text(plan.serialize_json()) + retry_path.write_text(retry.serialize_json()) + program = ( + "import json,sys; from pathlib import Path; " + "from data_designer.slurm.planning import ResolvedSlurmRunPlan; " + "from data_designer.slurm.state import RetryPlan; " + "from data_designer.slurm.launcher.renderer import render_generation_retry_script; " + "from data_designer.slurm.serving.resolver import resolve_vllm_server; " + "plan=ResolvedSlurmRunPlan.model_validate_json(Path(sys.argv[1]).read_text()); " + "retry=RetryPlan.model_validate_json(Path(sys.argv[2]).read_text()); " + "print(json.dumps({'script':render_generation_retry_script(plan,retry)," + "'gpus':[list(p.gpu_indices) for d in plan.deployments for p in resolve_vllm_server(plan,d.deployment_id).processes]}))" + ) + result = subprocess.run( + (sys.executable, "-c", program, str(plan_path), str(retry_path)), capture_output=True, text=True, check=True + ) + evidence = json.loads(result.stdout) + assert evidence["gpus"] == [[0, 1], [2, 3], [4, 5], [6, 7]] + assert "#SBATCH --nodes=1\n" in evidence["script"] + assert "#SBATCH --gres=gpu:8\n" in evidence["script"] + assert "#SBATCH --exclusive" not in evidence["script"] diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py new file mode 100644 index 000000000..7fe9e6b28 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -0,0 +1,686 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import shutil +import subprocess +import tarfile +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( + tmp_path: Path, + gpu_request_mode: str, +) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scratch_root = tmp_path / "data-designer-slurm-4101-0" + runtime_copy = scratch_root / "runtime" + shutil.copytree(runtime_root, runtime_copy) + attempt_directory = tmp_path / "runs/run-shell/shards/shard-00000/attempts/attempt-0001" + log_directory = attempt_directory / "logs/execution-00000002" + fake_bin = tmp_path / "bin" + log_directory.mkdir(parents=True) + fake_bin.mkdir() + artifacts = tuple(_artifact(tmp_path, name) for name in ("runtime", "lock", "client", "server")) + plan_path = tmp_path / "runs/run-shell/resolved-plan.json" + plan_path.parent.mkdir(parents=True, exist_ok=True) + plan = _plan( + tmp_path, runtime_copy, artifacts, gpu_request_mode, python_executable=(fake_bin / "python3").as_posix() + ) + plan_path.write_text(json.dumps(plan)) + manifest_path = tmp_path / "manifest-source.json" + manifest_path.write_text(json.dumps(_manifest(attempt_directory, artifacts, "a" * 64))) + marker_path = tmp_path / "host-python-ran" + _write_executable(fake_bin / "python3", f"#!/usr/bin/env bash\nprintf ran > {marker_path}\nexit 99\n") + _write_executable(fake_bin / "curl", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "getent", "#!/usr/bin/env bash\nexit 0\n") + _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nprintf 'compute-001\\n'\n") + _write_executable(fake_bin / "srun", _fake_srun()) + command = f""" + set -Eeuo pipefail +DD_PLAN_SHA256={"a" * 64} +DD_SHARD_ID=shard-00000 +DD_ATTEMPT_ORDINAL=0001 +DD_RUNTIME_ARCHIVE={shlex.quote(artifacts[0][0])} +DD_RUNTIME_SHA256={artifacts[0][1]} +readonly DD_PLAN={shlex.quote(plan_path.as_posix())} +readonly DD_ATTEMPT_DIR={shlex.quote(attempt_directory.as_posix())} +readonly DD_RUNTIME_ROOT={shlex.quote(runtime_copy.as_posix())} +source {shlex.quote((runtime_copy / "entrypoint.sh").as_posix())} +dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" +""" + environment = { + **os.environ, + "PATH": f"{fake_bin}:{os.environ['PATH']}", + "FAKE_MANIFEST_SOURCE": manifest_path.as_posix(), + "FAKE_GPU_MODE": gpu_request_mode, + "SOURCE_TOKEN": "supersecret", + "CUDA_VISIBLE_DEVICES": "0", + "DD_SCRATCH_ROOT": scratch_root.as_posix(), + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_ID": "4101", + "SLURM_JOB_NUM_NODES": "1", + "SLURM_JOB_NODELIST": "compute-001", + "SLURM_JOB_GPUS": "0", + "SLURM_NODEID": "0", + } + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + check=False, + env=environment, + timeout=10, + ) + + assert completed.returncode == 0, completed.stderr + assert not scratch_root.exists() + assert not marker_path.exists() + assert "supersecret" not in completed.stdout + completed.stderr + manifest_path.read_text() + assert "Starting vLLM servers..." in completed.stdout + assert "Waiting for vLLM servers to become ready..." in completed.stdout + assert "vLLM servers are ready." in completed.stdout + assert ( + f"Follow vLLM logs: tail -F {log_directory / 'deployment-00000-replica-00000-rank-00000.out'} " + f"{log_directory / 'deployment-00000-replica-00000-rank-00000.err'}" + ) in completed.stdout + assert ( + f"Generation started. Follow client logs: tail -F {log_directory / 'client-generation.out'} " + f"{log_directory / 'client-generation.err'}" + ) in completed.stdout + assert f"Inspect generation progress: cat {attempt_directory / 'client-progress.json'}" in completed.stdout + assert "Generation finished; finalizing run..." in completed.stdout + assert "Run completed." in completed.stdout + + +def test_staged_shell_modules_parse_as_bash(tmp_path: Path) -> None: + del tmp_path + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scripts = tuple( + runtime_root / name + for name in ("scratch.sh", "entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh") + ) + + completed = subprocess.run(("bash", "-n", *(path.as_posix() for path in scripts)), capture_output=True, text=True) + + assert completed.returncode == 0, completed.stderr + + +def test_scratch_prefers_slurm_tmpdir_and_cleans_private_tree(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + slurm_tmp = tmp_path / "slurm-tmp" + fallback_tmp = tmp_path / "fallback-tmp" + slurm_tmp.mkdir() + fallback_tmp.mkdir() + expected_root = slurm_tmp / "data-designer-slurm-4101-7" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "scratch.sh").as_posix())} +dd_initialize_local_scratch +[[ $DD_SCRATCH_ROOT == {shlex.quote(expected_root.as_posix())} ]] +[[ $HOME == "${{DD_SCRATCH_ROOT}}/home" ]] +[[ $ENROOT_TEMP_PATH == "${{DD_SCRATCH_ROOT}}/enroot/tmp" ]] +dd_remove_local_scratch "${{DD_SCRATCH_ROOT}}" +[[ ! -e $DD_SCRATCH_ROOT ]] +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={ + **os.environ, + "SLURM_ARRAY_TASK_ID": "7", + "SLURM_JOB_ID": "4101", + "SLURM_TMPDIR": slurm_tmp.as_posix(), + "TMPDIR": fallback_tmp.as_posix(), + }, + ) + + assert completed.returncode == 0, completed.stderr + assert not expected_root.exists() + + +def test_scratch_falls_back_to_tmpdir_and_rejects_symlink_root(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + fallback_tmp = tmp_path / "fallback-tmp" + outside = tmp_path / "outside" + fallback_tmp.mkdir() + outside.mkdir() + expected_root = fallback_tmp / "data-designer-slurm-4101-0" + expected_root.symlink_to(outside, target_is_directory=True) + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "scratch.sh").as_posix())} +if dd_initialize_local_scratch; then + exit 1 +else + [[ $? == 73 ]] +fi +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={ + **os.environ, + "SLURM_JOB_ID": "4101", + "SLURM_TMPDIR": (tmp_path / "missing").as_posix(), + "TMPDIR": fallback_tmp.as_posix(), + }, + ) + + assert completed.returncode == 0, completed.stderr + assert expected_root.is_symlink() + assert not tuple(outside.iterdir()) + + +def test_scratch_rejects_inconsistent_allocation_root(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scratch_parent = tmp_path / "scratch" + scratch_parent.mkdir() + expected_root = scratch_parent / "data-designer-slurm-4101-0" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "scratch.sh").as_posix())} +if dd_initialize_local_scratch /different/data-designer-slurm-4101-0; then + exit 1 +else + [[ $? == 73 ]] +fi +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "SLURM_JOB_ID": "4101", "SLURM_TMPDIR": scratch_parent.as_posix()}, + ) + + assert completed.returncode == 0, completed.stderr + assert not expected_root.exists() + + +def test_runtime_bundle_rejects_digest_mismatch_and_existing_target(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scratch_parent = tmp_path / "scratch" + scratch_parent.mkdir() + archive = tmp_path / "runtime.tar.gz" + with tarfile.open(archive, mode="w:gz") as bundle: + bundle.add(runtime_root / "scratch.sh", arcname="scratch.sh") + bundle.add(runtime_root / "entrypoint.sh", arcname="entrypoint.sh") + digest = hashlib.sha256(archive.read_bytes()).hexdigest() + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "scratch.sh").as_posix())} +dd_initialize_local_scratch +if dd_extract_runtime_bundle {shlex.quote(archive.as_posix())} {"0" * 64}; then + exit 1 +else + [[ $? == 65 ]] +fi +mkdir "${{DD_SCRATCH_ROOT}}/runtime" +if dd_extract_runtime_bundle {shlex.quote(archive.as_posix())} {digest}; then + exit 1 +else + [[ $? == 73 ]] +fi +dd_remove_local_scratch "${{DD_SCRATCH_ROOT}}" +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "SLURM_JOB_ID": "4101", "SLURM_TMPDIR": scratch_parent.as_posix()}, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_scratch_stages_runtime_for_fake_multi_node_allocation(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + scratch_parent = tmp_path / "scratch" + fake_bin = tmp_path / "bin" + scratch_parent.mkdir() + fake_bin.mkdir() + entrypoint = tmp_path / "entrypoint.sh" + entrypoint.write_text("true\n") + archive = tmp_path / "runtime.tar.gz" + with tarfile.open(archive, mode="w:gz") as bundle: + bundle.add(runtime_root / "scratch.sh", arcname="scratch.sh") + bundle.add(entrypoint, arcname="entrypoint.sh") + digest = hashlib.sha256(archive.read_bytes()).hexdigest() + srun_log = tmp_path / "srun.log" + _write_executable( + fake_bin / "srun", + "#!/usr/bin/env bash\n" + "set -Eeuo pipefail\n" + 'printf \'%s\\n\' "$*" >> "${FAKE_SRUN_LOG}"\n' + "while [[ $# -gt 0 && $1 != -- ]]; do shift; done\n" + "shift\n" + 'exec "$@"\n', + ) + expected_root = scratch_parent / "data-designer-slurm-4101-0" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "scratch.sh").as_posix())} +DD_RUNTIME_ARCHIVE={shlex.quote(archive.as_posix())} +DD_RUNTIME_SHA256={digest} +dd_initialize_local_scratch +dd_stage_allocation_runtime +[[ -f "${{DD_SCRATCH_ROOT}}/runtime/entrypoint.sh" ]] +dd_cleanup_allocation_scratch +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={ + **os.environ, + "FAKE_SRUN_LOG": srun_log.as_posix(), + "PATH": f"{fake_bin}:{os.environ['PATH']}", + "SLURM_JOB_ID": "4101", + "SLURM_JOB_NUM_NODES": "2", + "SLURM_TMPDIR": scratch_parent.as_posix(), + }, + ) + + assert completed.returncode == 0, completed.stderr + assert "--nodes=2 --ntasks=2" in srun_log.read_text() + assert "--time=00:01:00 --wait=0" in srun_log.read_text() + assert not expected_root.exists() + + +@pytest.mark.parametrize( + ("visible_gpus", "expected_gpus", "expected_status"), + (("0,1,2,3,4,5,6,7", 1, 0), ("0", 2, 65)), +) +def test_bash_gpu_count_requires_planned_minimum( + visible_gpus: str, + expected_gpus: int, + expected_status: int, +) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +DD_EXPECTED_GPUS={expected_gpus} +dd_verify_gpu_count +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "CUDA_VISIBLE_DEVICES": visible_gpus}, + ) + + assert completed.returncode == expected_status + + +def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "plan_reader.sh").as_posix())} +source {shlex.quote((runtime_root / "cleanup.sh").as_posix())} +dd_read_null_values DD_VALUES < <(printf '') +(( ${{#DD_VALUES[@]}} == 0 )) +dd_start_runtime_timer +sleep 30 & +pid=$! +DD_MANAGED_PIDS[2]=${{pid}} +dd_cleanup_steps +! kill -0 "${{pid}}" 2>/dev/null +dd_stop_runtime_timer +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=5) + + assert completed.returncode == 0, completed.stderr + + +def test_allocation_runtime_reports_preparation_phase_and_preserves_status() -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +dd_read_control_plan() {{ return 4; }} +dd_cleanup_allocation_scratch() {{ :; }} +dd_cleanup_steps() {{ :; }} +dd_stop_runtime_timer() {{ :; }} +dd_slurm_run_allocation /workspace/resolved-plan.json /workspace/attempt +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True) + + assert completed.returncode == 4 + assert completed.stderr == "allocation runtime failed: phase=read_control_plan status=4\n" + + +def test_step_runner_reports_failed_step_log_paths(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + attempt_directory = tmp_path / "attempt" + log_directory = attempt_directory / "logs/execution-00000002" + log_directory.mkdir(parents=True) + image = tmp_path / "client.sqsh" + image.touch() + manifest = { + "steps": [ + _step( + attempt_directory, + "client-preflight", + "client_preflight", + image.as_posix(), + "true", + ) + ] + } + manifest_path = tmp_path / "runtime-manifest.json" + manifest_path.write_text(json.dumps(manifest)) + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + _write_executable(fake_bin / "srun", "#!/usr/bin/env bash\nexit 4\n") + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "plan_reader.sh").as_posix())} +source {shlex.quote((runtime_root / "step_runner.sh").as_posix())} +DD_GPU_REQUEST_MODE=gres +DD_CONTAINER_MOUNTS= +DD_SCRATCH_ROOT=/tmp/data-designer-slurm-4101-0 +DD_SCRATCH_CONTAINER_ROOT=/run/data-designer-slurm +DD_ALL_SECRET_NAMES=() +DD_MANAGED_PIDS=() +status=0 +dd_run_step {shlex.quote(manifest_path.as_posix())} client-preflight || status=$? +exit "${{status}}" +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "PATH": f"{fake_bin}:{os.environ['PATH']}"}, + ) + + assert completed.returncode == 4 + assert "runtime step failed: step_id=client-preflight status=4" in completed.stderr + assert f"stdout={log_directory / 'client-preflight.out'}" in completed.stderr + assert f"stderr={log_directory / 'client-preflight.err'}" in completed.stderr + + +def test_step_runner_builds_one_coordinated_srun_across_selected_nodes() -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "step_runner.sh").as_posix())} +DD_STEP_NODE_HOSTS=(compute-001 compute-002) +DD_STEP_GPU_INDICES=(0 1) +DD_STEP_CONTAINER_ENV=() +DD_STEP_EXECUTION=container +DD_STEP_CPUS=4 +DD_STEP_IMAGE=/images/server.sqsh +DD_STEP_KILL_ON_BAD_EXIT=true +DD_GPU_REQUEST_MODE=gres +DD_CONTAINER_MOUNTS= +DD_SCRATCH_ROOT=/tmp/data-designer-slurm-4101-0 +DD_SCRATCH_CONTAINER_ROOT=/run/data-designer-slurm +dd_build_srun_command +command=${{DD_SRUN_COMMAND[*]}} +[[ $command == *--nodes=2* ]] +[[ $command == *--ntasks=2* ]] +[[ $command == *--ntasks-per-node=1* ]] +[[ $command == *--nodelist=compute-001,compute-002* ]] +[[ $command == *--kill-on-bad-exit=1* ]] +[[ $command == *--gpus-per-task=2* ]] +[[ $command == *--container-mounts=/tmp/data-designer-slurm-4101-0:/run/data-designer-slurm* ]] +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True) + + assert completed.returncode == 0, completed.stderr + + +def test_shell_resolves_scheduler_hosts_in_planner_order(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + _write_executable(fake_bin / "scontrol", "#!/usr/bin/env bash\nprintf 'compute-001\\ncompute-002\\n'\n") + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +DD_EXPECTED_NODES=2 +DD_CLIENT_NODE_INDEX=0 +SLURM_JOB_NODELIST=compute-[001-002] +dd_resolve_allocation_hosts +[[ ${{DD_ALLOCATION_HOSTS[*]}} == 'compute-001 compute-002' ]] +[[ $DD_CLIENT_HOST == compute-001 ]] +DD_EXPECTED_NODES=3 +! dd_resolve_allocation_hosts +""" + + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env={**os.environ, "PATH": f"{fake_bin}:{os.environ['PATH']}"}, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_plan_reader_rejects_mount_overlapping_allocation_scratch(tmp_path: Path) -> None: + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + artifacts = tuple(_artifact(tmp_path, name) for name in ("runtime", "lock", "client", "server")) + payload = _plan(tmp_path, tmp_path / "runtime-root", artifacts, "gres") + mounts = payload["container_mounts"] + assert isinstance(mounts, list) + mounts.append({"source": "/source", "target": "/run", "read_only": False}) + plan_path = tmp_path / "plan.json" + plan_path.write_text(json.dumps(payload)) + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "plan_reader.sh").as_posix())} +if dd_read_control_plan {shlex.quote(plan_path.as_posix())}; then + exit 1 +fi +""" + + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True) + + assert completed.returncode == 0, completed.stderr + + +def _artifact(root: Path, name: str) -> tuple[str, str]: + path = root / f"{name}.artifact" + content = name.encode() + path.write_bytes(content) + return path.as_posix(), hashlib.sha256(content).hexdigest() + + +def _plan( + root: Path, + runtime_root: Path, + artifacts: tuple[tuple[str, str], ...], + gpu_request_mode: str, + *, + python_executable: str = "/bin/true", +) -> dict[str, object]: + runtime, lock, client, server = artifacts + mounts = [ + {"source": root.as_posix(), "target": root.as_posix(), "read_only": False}, + {"source": runtime_root.as_posix(), "target": runtime_root.as_posix(), "read_only": False}, + ] + return { + "schema_version": 1, + "runtime_bundle": {"path": runtime[0], "sha256": runtime[1]}, + "client": { + "host_node_index": 0, + "authored": {"cpus": 1}, + "dependency_lock": {"path": lock[0], "sha256": lock[1]}, + "runtime": {"python_executable": python_executable}, + }, + "resolved_gpus_per_node": 1, + "selected_profile": {"profile": {"gpu_request_mode": gpu_request_mode}}, + "container_mounts": mounts, + "deployments": [ + { + "node_indices": [0], + "image": {"path": server[0], "sha256": server[1]}, + "authored": { + "server": {"environment": {"SERVER_TOKEN": {"type": "secret", "environment": "SOURCE_TOKEN"}}} + }, + } + ], + "builder": {"source": None}, + "shards": [{"array_task_index": 0, "input_partition": None}], + } + + +def _manifest( + attempt_directory: Path, + artifacts: tuple[tuple[str, str], ...], + plan_sha256: str, +) -> dict[str, object]: + _, _, _, server = artifacts + return { + "schema_version": 1, + "run_id": "run-shell", + "shard_id": "shard-00000", + "attempt_id": "attempt-0001", + "plan_sha256": plan_sha256, + "all_secret_environment_names": ["SERVER_TOKEN", "SOURCE_TOKEN"], + "steps": [ + _step(attempt_directory, "client-preflight", "client_preflight", None, "fake-preflight"), + _step( + attempt_directory, + "deployment-00000-replica-00000-rank-00000", + "server", + server[0], + "fake-server", + gpu_indices=[0], + secret_environment={"SERVER_TOKEN": "SOURCE_TOKEN"}, + container_environment=["SERVER_TOKEN"], + readiness={"host": "127.0.0.1", "port": 18000, "path": "/health", "deadline_seconds": 2}, + ), + _step( + attempt_directory, + "deployment-00000-endpoint", + "endpoint", + None, + "fake-endpoint", + readiness={"host": "127.0.0.1", "port": 17000, "path": "/health", "deadline_seconds": 2}, + ), + _step( + attempt_directory, + "client-generation", + "client", + None, + "fake-client", + secret_environment={"SOURCE_TOKEN": "SOURCE_TOKEN"}, + ), + ], + } + + +def _step( + attempt_directory: Path, + step_id: str, + role: str, + image_path: str | None, + command: str, + *, + gpu_indices: list[int] | None = None, + secret_environment: dict[str, str] | None = None, + container_environment: list[str] | None = None, + readiness: dict[str, object] | None = None, +) -> dict[str, object]: + log_root = attempt_directory / "logs/execution-00000002" + return { + "step_id": step_id, + "role": role, + "execution": "container" if image_path is not None else "native", + "image_path": image_path, + "command": [command], + "cpus": 1, + "gpu_indices": gpu_indices or [], + "node_hosts": ["compute-001"], + "kill_on_bad_exit": False, + "literal_environment": {"LC_ALL": "C"}, + "secret_environment": secret_environment or {}, + "environment_prefixes": {}, + "container_environment": container_environment or [], + "stdout_path": (log_root / f"{step_id}.out").as_posix(), + "stderr_path": (log_root / f"{step_id}.err").as_posix(), + "launch_delay_seconds": 0, + "readiness": [readiness] if readiness is not None else [], + } + + +def _write_executable(path: Path, content: str) -> None: + path.write_text(content) + path.chmod(0o700) + + +def _fake_srun() -> str: + return """#!/usr/bin/env bash +set -Eeuo pipefail +arguments=("$@") +for ((index = 0; index < ${#arguments[@]}; index++)); do + [[ ${arguments[index]} == -- ]] && break +done +command=${arguments[index + 1]} +if [[ ${command} == */python3 ]]; then + operation=${arguments[index + 4]} + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + [[ ${arguments[*]} != *--container-* ]] + if [[ ${operation} == prepare ]]; then + for ((position = index + 5; position < ${#arguments[@]}; position++)); do + if [[ ${arguments[position]} == --manifest ]]; then + cp "${FAKE_MANIFEST_SOURCE}" "${arguments[position + 1]}" + break + fi + done + fi + exit 0 +fi +case ${command} in + fake-preflight) + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + ;; + fake-client) + [[ ${SOURCE_TOKEN} == supersecret && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + ;; + fake-server) + [[ ${SERVER_TOKEN} == supersecret && ! ${SOURCE_TOKEN+x} ]] + if [[ ${FAKE_GPU_MODE} == visible ]]; then + [[ ${CUDA_VISIBLE_DEVICES} == 0 ]] + else + [[ ! ${CUDA_VISIBLE_DEVICES+x} ]] + fi + trap 'exit 0' TERM INT + while :; do :; done + ;; + fake-endpoint) + [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] + trap 'exit 0' TERM INT + while :; do :; done + ;; + *) + exit 90 + ;; +esac +""" diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py new file mode 100644 index 000000000..7b9cbb3b5 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast + +import pytest +from conftest import RuntimeCase, relocate_plan + +from data_designer.slurm.config import QueueBackpressureConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.backpressure import ( + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, +) +from data_designer.slurm.runtime.distributed import build_vllm_process_command +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.preflight import AllocationLayout +from data_designer.slurm.runtime.steps import ( + DefaultClientStepBuilder, + build_endpoint_steps, + build_vllm_command, + build_vllm_steps, +) +from data_designer.slurm.serving.resolver import resolve_vllm_server +from data_designer.slurm.state import AttemptLifecycleState, AttemptManifest, SchedulerIdentity + + +def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + deployment = resolve_vllm_server(context.plan, context.plan.deployments[0].deployment_id) + source_environment = { + "DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0", + "PATH": "/custom/bin", + "SLURM_JOB_ID": "4101", + "UNREVIEWED_VALUE": "must-not-forward", + } + + server_steps = build_vllm_steps( + deployment, + context.plan, + context.attempt_directory, + source_environment, + context.attempt_directory / "runtime", + ) + endpoint_steps = build_endpoint_steps( + (deployment,), + context.plan, + context.attempt_directory, + source_environment, + context.attempt_directory / "runtime/proxy.py", + ) + endpoints = tuple(endpoint for _, endpoint in endpoint_steps) + client_builder = DefaultClientStepBuilder() + client_steps = ( + client_builder.build_preflight_step( + context.plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + source_environment, + ), + client_builder.build_generation_step( + context.plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + source_environment, + ), + ) + + all_steps = server_steps + tuple(step for step, _ in endpoint_steps) + client_steps + assert all(step.command[0] == "srun" for step in all_steps) + assert all("--" in step.command for step in all_steps) + assert all("UNREVIEWED_VALUE" not in step.environment for step in all_steps) + assert all(step.environment["SLURM_JOB_ID"] == "4101" for step in all_steps) + assert all(step.environment["LC_ALL"] == "C" for step in all_steps) + assert all(step.stdout_path.parent == context.attempt_directory / "logs" for step in all_steps) + assert all(f"--cpus-per-task={context.plan.client.authored.cpus}" in step.command for step in all_steps) + assert all( + any( + argument.endswith(":/run/data-designer-slurm") + for argument in step.command + if argument.startswith("--container-mounts=") + ) + for step in server_steps + ) + native_steps = tuple(step for step, _ in endpoint_steps) + client_steps + assert all(not any(argument.startswith("--container-") for argument in step.command) for step in native_steps) + + server = server_steps[0] + assert server.role is RuntimeStepRole.SERVER + assert f"--gpus-per-task={deployment.topology.tensor_parallel}" in server.command + expected_mask = sum(1 << index for index in deployment.processes[0].gpu_indices) + assert f"--gpu-bind=mask_gpu:{expected_mask:#x}" in server.command + assert "CUDA_VISIBLE_DEVICES" not in server.environment + assert all("CUDA_VISIBLE_DEVICES" not in argument for argument in server.command) + assert "--middleware" in server.command + assert "data_designer.slurm.runtime.backpressure.QueueDepthBackpressureMiddleware" in server.command + assert server.environment[MAX_WAITING_REQUESTS_ENVIRONMENT] == "128" + assert server.environment[RETRY_AFTER_SECONDS_ENVIRONMENT] == "1" + assert server.environment["PYTHONPATH"].endswith("/runtime") + assert any( + argument.startswith("--container-env=") + and all( + name in argument + for name in ("PYTHONPATH", MAX_WAITING_REQUESTS_ENVIRONMENT, RETRY_AFTER_SECONDS_ENVIRONMENT) + ) + for argument in server.command + ) + assert all("--gpus-per-task=" not in argument for step in client_steps for argument in step.command) + assert all("--gres=none" in step.command for step in client_steps) + assert all("CUDA_VISIBLE_DEVICES" not in step.environment for step in client_steps) + endpoint_command = tuple(step for step, _ in endpoint_steps)[0].command + assert endpoint_command[endpoint_command.index("--health-path") + 1] == deployment.launch_policy.readiness_path + + +def test_vllm_command_maps_absolute_model_path(runtime_case: RuntimeCase) -> None: + plan = runtime_case.context.plan + host_model = runtime_case.workspace / "models" / "generator" + payload = cast(dict[str, object], json.loads(plan.serialize_json())) + deployment_payload = cast(dict[str, object], cast(list[object], payload["deployments"])[0]) + authored = cast(dict[str, object], deployment_payload["authored"]) + authored["model"] = host_model.as_posix() + authored["served_model_name"] = deployment_payload["served_model_name"] + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected_profile["profile"]) + profile = SlurmProfile.model_validate(profile_payload) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + deployment = resolve_vllm_server(plan, plan.deployments[0].deployment_id) + + command = build_vllm_command(deployment, deployment.processes[0], plan) + + assert command[2] == "/workspace/primary/models/generator" + + +@pytest.mark.parametrize("plugin", [None, "deberta_gliner_io"]) +@pytest.mark.parametrize("plan_fixture", ["single_node_plan", "multi_node_plan", "shared_node_plan"]) +def test_vllm_commands_preserve_typed_io_processor_plugin( + request: pytest.FixtureRequest, plan_fixture: str, plugin: str | None +) -> None: + plan = request.getfixturevalue(plan_fixture) + deployment = resolve_vllm_server(plan, plan.deployments[0].deployment_id) + policy = deployment.launch_policy.model_copy(update={"io_processor_plugin": plugin}) + deployment = deployment.model_copy(update={"launch_policy": policy}) + node_count = 1 + max(index for member in plan.deployments for index in member.node_indices) + layout = AllocationLayout(tuple(f"compute-{index}" for index in range(node_count))) + + for process in deployment.processes: + commands = (build_vllm_process_command(deployment, process, plan, layout),) + if process.http_port is not None: + commands += (build_vllm_command(deployment, process, plan),) + for command in commands: + assert command.count("--io-processor-plugin") == (plugin is not None) + if plugin is not None: + assert command[command.index("--io-processor-plugin") + 1] == plugin + + +def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + deployment = resolve_vllm_server(context.plan, context.plan.deployments[0].deployment_id) + endpoint_steps = build_endpoint_steps( + (deployment,), + context.plan, + context.attempt_directory, + {"DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0"}, + context.attempt_directory / "runtime/proxy.py", + ) + endpoints = tuple(endpoint for _, endpoint in endpoint_steps) + + step = DefaultClientStepBuilder().build_generation_step( + context.plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + {"DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0"}, + ) + + separator = step.command.index("--") + worker = step.command[separator + 1 :] + assert worker[:5] == ( + context.plan.client.runtime.python_executable, + "-m", + "data_designer.slurm.client.worker", + "run", + "--plan", + ) + assert context.attempt.attempt_id in worker + assert context.shard.shard_id in worker + assert f"generator=http://127.0.0.1:{endpoints[0].port}/v1" in worker + + +@pytest.mark.parametrize("environment", ({}, {"DD_SCRATCH_ROOT": "relative"})) +def test_client_step_rejects_unavailable_allocation_scratch( + runtime_case: RuntimeCase, + environment: dict[str, str], +) -> None: + context = runtime_case.context + + with pytest.raises(SlurmRuntimeError, match="allocation scratch"): + DefaultClientStepBuilder().build_preflight_step( + context.plan, + context.shard, + context.attempt, + context.attempt_directory, + (), + environment, + ) + + +def test_endpoint_step_uses_resolved_retry_policy_and_backends(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + deployment = resolve_vllm_server(context.plan, context.plan.deployments[0].deployment_id) + + ((step, endpoint),) = build_endpoint_steps( + (deployment,), + context.plan, + context.attempt_directory, + {}, + context.attempt_directory / "runtime/proxy.py", + ) + + assert step.role is RuntimeStepRole.ENDPOINT + assert endpoint.port == deployment.logical_endpoint.port + retry_after_seconds = deployment.launch_policy.queue_backpressure.retry_after_seconds + assert retry_after_seconds is not None + assert ("--retry-after-seconds", str(retry_after_seconds)) == step.command[ + step.command.index("--retry-after-seconds") : step.command.index("--retry-after-seconds") + 2 + ] + assert "--max-waiting-requests" not in step.command + assert step.environment["PYTHONPATH"].endswith("/runtime") + assert "--container-env=PYTHONPATH" not in step.command + assert (context.attempt_directory / "runtime/proxy.py").as_posix() in step.command + for backend in deployment.backend_endpoints: + assert f"http://127.0.0.1:{backend.port}" in step.command + + disabled_policy = deployment.launch_policy.model_copy( + update={"queue_backpressure": QueueBackpressureConfig(max_waiting_requests=0, retry_after_seconds=None)} + ) + disabled = deployment.model_copy(update={"launch_policy": disabled_policy}) + ((disabled_step, _),) = build_endpoint_steps( + (disabled,), + context.plan, + context.attempt_directory, + {}, + context.attempt_directory / "runtime/proxy.py", + ) + assert "--retry-after-seconds" not in disabled_step.command + assert "--max-waiting-requests" not in disabled_step.command + + +def test_client_receives_client_secrets_without_server_only_secrets( + tmp_path: Path, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + plan = relocate_plan(multi_node_plan, workspace) + shard = plan.shards[0] + attempt_directory = workspace / "attempt" + attempt = AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=ArtifactReference( + path=Path(plan.authored_config.path).with_name("resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ), + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + created_at=datetime(2026, 9, 1, tzinfo=timezone.utc), + updated_at=datetime(2026, 9, 1, tzinfo=timezone.utc), + ) + endpoints = tuple( + RuntimeEndpoint( + model_alias=deployment.authored.model_alias, + served_model_name=deployment.served_model_name, + host="127.0.0.1", + port=port.port, + ) + for deployment, port in zip(plan.deployments, plan.client.ports, strict=True) + ) + + step = DefaultClientStepBuilder().build_preflight_step( + plan, + shard, + attempt, + attempt_directory, + endpoints, + { + "DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0", + "PACKAGE_INDEX_TOKEN": "client-secret", + "HF_TOKEN": "server-secret", + "SLURM_JOB_GPUS": "0", + }, + ) + + assert step.environment["PACKAGE_INDEX_TOKEN"] == "client-secret" + assert "HF_TOKEN" not in step.environment + assert not any(argument.startswith("--container-") for argument in step.command) + + with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): + DefaultClientStepBuilder().build_preflight_step( + plan, + shard, + attempt, + attempt_directory, + endpoints, + { + "DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0", + "HF_TOKEN": "server-secret", + }, + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_supervisor.py b/packages/data-designer-slurm/tests/runtime/test_supervisor.py new file mode 100644 index 000000000..308449d78 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_supervisor.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + +import pytest +from slurm_test_fakes import FakeClock + +from data_designer.slurm.runtime import signals as runtime_signals +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import RuntimeStep, RuntimeStepRole +from data_designer.slurm.runtime.signals import TerminationSignalCoordinator +from data_designer.slurm.runtime.supervisor import StepSupervisor, SubprocessStepRunner + + +@dataclass(slots=True) +class _Process: + pid: int + returncode: int | None = None + terminated: int = 0 + killed: int = 0 + + def poll(self) -> int | None: + return self.returncode + + def terminate(self) -> None: + self.terminated += 1 + self.returncode = -15 + + def kill(self) -> None: + self.killed += 1 + self.returncode = -9 + + def wait(self, timeout: float | None = None) -> int: + del timeout + assert self.returncode is not None + return self.returncode + + +class _Runner: + def __init__(self, processes: list[_Process]) -> None: + self.processes = processes + self.started = 0 + + def start(self, step: RuntimeStep) -> _Process: + del step + process = self.processes[self.started] + self.started += 1 + return process + + +def test_cleanup_is_reverse_order_and_idempotent(tmp_path: Path, fake_clock: FakeClock) -> None: + first = _Process(1) + second = _Process(2) + runner = _Runner([first, second]) + supervisor = StepSupervisor(runner, signals=TerminationSignalCoordinator(), clock=fake_clock) + supervisor.start(_step(tmp_path, "first")) + supervisor.start(_step(tmp_path, "second")) + + supervisor.cleanup() + supervisor.cleanup() + + assert first.terminated == second.terminated == 1 + assert first.killed == second.killed == 0 + with pytest.raises(SlurmRuntimeError, match="after cleanup"): + supervisor.start(_step(tmp_path, "third")) + + +def test_wait_fails_when_required_service_exits(tmp_path: Path, fake_clock: FakeClock) -> None: + service = _Process(1, returncode=19) + target = _Process(2) + supervisor = StepSupervisor( + _Runner([service, target]), + signals=TerminationSignalCoordinator(), + clock=fake_clock, + ) + managed_service = supervisor.start(_step(tmp_path, "service")) + managed_target = supervisor.start(_step(tmp_path, "target")) + + with pytest.raises(SlurmRuntimeError, match="status 19"): + supervisor.wait(managed_target, required=(managed_service,)) + + +def test_subprocess_runner_uses_restrictive_logs_and_rejects_symlink_root(tmp_path: Path) -> None: + attempt = tmp_path / "attempt" + attempt.mkdir(mode=0o700) + step = RuntimeStep( + step_id="real-step", + role=RuntimeStepRole.CLIENT_PREFLIGHT, + command=(sys.executable, "-c", "import sys; print('out'); print('err', file=sys.stderr)"), + environment={"PATH": "/usr/bin", "LC_ALL": "C"}, + stdout_path=attempt / "logs" / "execution-00000001" / "real-step.out", + stderr_path=attempt / "logs" / "execution-00000001" / "real-step.err", + ) + + process = SubprocessStepRunner().start(step) + assert process.wait(timeout=10) == 0 + assert step.stdout_path.read_text() == "out\n" + assert step.stderr_path.read_text() == "err\n" + assert step.stdout_path.stat().st_mode & 0o077 == 0 + + restarted = RuntimeStep( + step_id=step.step_id, + role=step.role, + command=step.command, + environment=step.environment, + stdout_path=attempt / "logs" / "execution-00000002" / "real-step.out", + stderr_path=attempt / "logs" / "execution-00000002" / "real-step.err", + ) + restarted_process = SubprocessStepRunner().start(restarted) + assert restarted_process.wait(timeout=10) == 0 + assert restarted.stdout_path.read_text() == "out\n" + + outside = tmp_path / "outside" + outside.mkdir() + symlink_attempt = tmp_path / "symlink-attempt" + symlink_attempt.mkdir(mode=0o700) + (symlink_attempt / "logs").symlink_to(outside, target_is_directory=True) + unsafe = RuntimeStep( + step_id="unsafe-step", + role=RuntimeStepRole.CLIENT, + command=("true",), + environment={"PATH": "/usr/bin"}, + stdout_path=symlink_attempt / "logs" / "execution-00000001" / "unsafe.out", + stderr_path=symlink_attempt / "logs" / "execution-00000001" / "unsafe.err", + ) + with pytest.raises(OSError, match="not a directory"): + SubprocessStepRunner().start(unsafe) + assert tuple(outside.iterdir()) == () + + +def test_subprocess_cleanup_terminates_the_complete_step_process_group(tmp_path: Path) -> None: + attempt = tmp_path / "attempt" + attempt.mkdir(mode=0o700) + parent_stopped = attempt / "parent-stopped" + child_stopped = attempt / "child-stopped" + child_ready = attempt / "child-ready" + child_code = ( + "import signal,time; from pathlib import Path; " + f"stopped=Path({child_stopped.as_posix()!r}); ready=Path({child_ready.as_posix()!r}); " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "ready.write_text('ready'); time.sleep(60)" + ) + parent_code = ( + "import signal,subprocess,sys,time; from pathlib import Path; " + f"stopped=Path({parent_stopped.as_posix()!r}); child_code={child_code!r}; " + "signal.signal(signal.SIGTERM, lambda *_: (stopped.write_text('stopped'), exit(0))); " + "subprocess.Popen([sys.executable, '-c', child_code]); time.sleep(60)" + ) + step = RuntimeStep( + step_id="process-tree", + role=RuntimeStepRole.SERVER, + command=(sys.executable, "-c", parent_code), + environment={"PATH": "/usr/bin"}, + stdout_path=attempt / "logs" / "execution-00000001" / "process-tree.out", + stderr_path=attempt / "logs" / "execution-00000001" / "process-tree.err", + ) + supervisor = StepSupervisor( + SubprocessStepRunner(), + signals=TerminationSignalCoordinator(), + poll_interval_seconds=0.01, + termination_grace_seconds=3, + ) + supervisor.start(step) + deadline = time.monotonic() + 3 + while not child_ready.exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert child_ready.exists() + + supervisor.cleanup() + + assert parent_stopped.read_text() == "stopped" + assert child_stopped.read_text() == "stopped" + + +def test_start_normalizes_process_creation_failure(tmp_path: Path, fake_clock: FakeClock) -> None: + class _FailingRunner: + def start(self, step: RuntimeStep) -> _Process: + del step + raise subprocess.SubprocessError("injected") + + with pytest.raises(SlurmRuntimeError, match="cannot start"): + StepSupervisor( + _FailingRunner(), + signals=TerminationSignalCoordinator(), + clock=fake_clock, + ).start(_step(tmp_path, "failed")) + + +def test_start_registers_process_before_replaying_deferred_termination( + tmp_path: Path, + fake_clock: FakeClock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + process = _Process(1) + installed_handlers: dict[object, object] = {} + signals = TerminationSignalCoordinator() + supervisor: StepSupervisor + + def fake_getsignal(selected: object) -> object: + del selected + return runtime_signals.signal.SIG_DFL + + def fake_signal(selected: object, handler: object) -> object: + installed_handlers[selected] = handler + return runtime_signals.signal.SIG_DFL + + def fake_kill(process_id: int, selected: object) -> None: + del process_id + assert len(supervisor.managed_steps) == 1 + handler = installed_handlers[selected] + assert callable(handler) + handler(selected, None) + + class _InterruptingRunner(_Runner): + def start(self, step: RuntimeStep) -> _Process: + started = super().start(step) + handler = installed_handlers[runtime_signals.signal.SIGTERM] + assert callable(handler) + handler(runtime_signals.signal.SIGTERM, None) + return started + + monkeypatch.setattr(runtime_signals.signal, "getsignal", fake_getsignal) + monkeypatch.setattr(runtime_signals.signal, "signal", fake_signal) + monkeypatch.setattr(runtime_signals.os, "kill", fake_kill) + supervisor = StepSupervisor(_InterruptingRunner([process]), signals=signals, clock=fake_clock) + + with ( + pytest.raises(KeyboardInterrupt), + signals.interrupt_on_termination(supervisor.cleanup), + ): + supervisor.start(_step(tmp_path, "server")) + + assert supervisor.managed_steps[0].process is process + assert process.terminated == 1 + + +def test_start_can_register_from_a_worker_thread(tmp_path: Path, fake_clock: FakeClock) -> None: + signals = TerminationSignalCoordinator() + supervisor = StepSupervisor(_Runner([_Process(1)]), signals=signals, clock=fake_clock) + + with ( + signals.interrupt_on_termination(supervisor.cleanup), + ThreadPoolExecutor(max_workers=1) as executor, + ): + managed = executor.submit(supervisor.start, _step(tmp_path, "server")).result() + + assert supervisor.managed_steps == (managed,) + supervisor.cleanup() + + +def test_cleanup_remains_incomplete_while_a_process_is_still_running( + tmp_path: Path, + fake_clock: FakeClock, +) -> None: + class _StubbornProcess(_Process): + def terminate(self) -> None: + self.terminated += 1 + + def kill(self) -> None: + self.killed += 1 + + def wait(self, timeout: float | None = None) -> int: + del timeout + raise TimeoutError("still running") + + process = _StubbornProcess(1) + supervisor = StepSupervisor( + _Runner([process]), + signals=TerminationSignalCoordinator(), + clock=fake_clock, + poll_interval_seconds=1, + termination_grace_seconds=1, + ) + supervisor.start(_step(tmp_path, "server")) + + with pytest.raises(SlurmRuntimeError, match="cleanup failed"): + supervisor.cleanup() + + assert process.terminated == process.killed == 1 + assert not supervisor.cleanup_complete + + +def _step(root: Path, step_id: str) -> RuntimeStep: + return RuntimeStep( + step_id=step_id, + role=RuntimeStepRole.SERVER, + command=("true",), + environment={"PATH": "/usr/bin"}, + stdout_path=(root / f"{step_id}.out").absolute(), + stderr_path=(root / f"{step_id}.err").absolute(), + ) diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py new file mode 100644 index 000000000..9358e6418 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -0,0 +1,771 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from collections.abc import Callable +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from slurm_test_fakes import FakeClock + +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildProfile, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + SchedulerProfile, + SlurmProfile, +) +from data_designer.slurm.images.service import VerifiedImageRegistry +from data_designer.slurm.launcher.errors import SlurmSubmissionError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmProcessExitCode, + SlurmQueueEntry, +) +from data_designer.slurm.services import ( + SlurmImageService, + SlurmServiceError, + SlurmServiceErrorCode, + create_slurm_image_service, +) +from data_designer.slurm.state import SchedulerState + +_JOB_ID = 42 +_LIFECYCLE_ID = "image-test" + + +class _Launcher: + def __init__( + self, + *, + queue_state: SchedulerState | None = None, + accounting_state: SchedulerState | None = SchedulerState.COMPLETED, + exit_status: int = 0, + termination_signal: int = 0, + on_submit: Callable[[], None] | None = None, + submission_error: SlurmSubmissionError | None = None, + ) -> None: + self.queue_state = queue_state + self.accounting_state = accounting_state + self.exit_status = exit_status + self.termination_signal = termination_signal + self.on_submit = on_submit + self.submission_error = submission_error + self.submissions: list[str] = [] + self.cancellations: list[int] = [] + + def submit_script(self, script: str, **_: object) -> SlurmJobSubmissionReceipt: + self.submissions.append(script) + if self.submission_error is not None: + raise self.submission_error + if self.on_submit is not None: + self.on_submit() + return SlurmJobSubmissionReceipt(job_id=_JOB_ID) + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + del selectors + if self.queue_state is None: + return () + return (SlurmQueueEntry(job_identity=_JOB_ID, state=self.queue_state),) + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + del selectors + if self.accounting_state is None: + return () + return ( + SlurmAccountingEntry( + job_identity=_JOB_ID, + state=self.accounting_state, + process_exit_code=SlurmProcessExitCode( + exit_status=self.exit_status, + termination_signal=self.termination_signal, + ), + ), + ) + + def cancel(self, job_id: int) -> None: + self.cancellations.append(job_id) + + +class _DelayedAccountingLauncher(_Launcher): + def __init__(self, *, on_submit: Callable[[], None]) -> None: + super().__init__(queue_state=SchedulerState.COMPLETED, on_submit=on_submit) + self.queue_queries = 0 + self.accounting_queries = 0 + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + self.queue_queries += 1 + if self.queue_queries == 1: + return super().query_queue(selectors) + return () + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + self.accounting_queries += 1 + if self.accounting_queries < 3: + return ( + SlurmAccountingEntry( + job_identity=_JOB_ID, + state=SchedulerState.RUNNING, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + return super().query_accounting(selectors) + + +class _TransitionLauncher(_Launcher): + def __init__(self, *, on_submit: Callable[[], None], running_checks: int = 1) -> None: + super().__init__(queue_state=SchedulerState.PENDING, accounting_state=None, on_submit=on_submit) + self.queue_queries = 0 + self.running_checks = running_checks + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + self.queue_queries += 1 + if self.queue_queries == 2: + self.queue_state = SchedulerState.RUNNING + elif self.queue_queries == 2 + self.running_checks: + self.queue_state = SchedulerState.COMPLETED + self.accounting_state = SchedulerState.COMPLETED + return super().query_queue(selectors) + + +class _DelayedCancellationLauncher(_Launcher): + def __init__(self) -> None: + super().__init__(accounting_state=None) + + def cancel(self, job_id: int) -> None: + super().cancel(job_id) + self.accounting_state = SchedulerState.RUNNING + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + entries = super().query_accounting(selectors) + if self.cancellations: + self.accounting_state = SchedulerState.CANCELLED + return entries + + +class _ImmediateCancellationLauncher(_Launcher): + def cancel(self, job_id: int) -> None: + super().cancel(job_id) + self.accounting_state = SchedulerState.CANCELLED + + +def _interrupt(_: float) -> None: + raise KeyboardInterrupt + + +def test_default_image_add_runs_lifecycle_and_registers_existing_sqsh(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _Launcher(on_submit=lambda: _write_inspection(workspace, content)) + service = _service(workspace, launcher) + + registered = service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert registered.path == source.as_posix() + assert registered.sqsh_sha256 == hashlib.sha256(content).hexdigest() + assert service.get("client") == registered + assert ( + VerifiedImageRegistry(workspace) + .resolve_for_planning( + ImageRef(name="client"), + expected_kind=ImageKind.CLIENT, + ) + .path + == source.as_posix() + ) + assert len(launcher.submissions) == 1 + assert "#SBATCH --partition=cpu" in launcher.submissions[0] + assert not _job_directory(workspace).exists() + + +def test_default_image_add_waits_for_completed_job_accounting_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _DelayedAccountingLauncher(on_submit=lambda: _write_inspection(workspace, content)) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + registered = _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert registered.path == source.as_posix() + assert launcher.queue_queries == 3 + assert launcher.accounting_queries == 3 + assert clock.sleep_calls == [30.0, 30.0] + assert not _job_directory(workspace).exists() + + +def test_image_add_waits_for_accounting_without_repeating_status(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _DelayedAccountingLauncher(on_submit=lambda: _write_inspection(workspace, content)) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + progress: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=clock.now, + sleep=clock.sleep, + progress=progress.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert clock.sleep_calls == [30.0, 30.0] + assert progress == [ + "Submitting image add job...", + "Image add job 42 submitted; checking every 30 seconds.", + "Image add job 42: completed (elapsed 00:00:00).", + "Waiting for successful job exit evidence...", + "Registering image alias...", + "Image add complete: client (total elapsed 00:01:00).", + ] + + +def test_image_add_reports_pending_running_and_completed(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _TransitionLauncher(on_submit=lambda: _write_inspection(workspace, content), running_checks=2) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + progress: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=clock.now, + sleep=clock.sleep, + progress=progress.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert clock.sleep_calls == [30.0, 30.0, 30.0] + assert "Image add job 42: pending (elapsed 00:00:00)." in progress + assert "Image add job 42: running (elapsed 00:00:30)." in progress + assert "Image add job 42: still running (elapsed 00:01:00)." not in progress + assert "Image add job 42: completed (elapsed 00:01:30)." in progress + + +def test_image_add_reports_job_phases_and_quiet_heartbeat(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + + class PhasedLauncher(_TransitionLauncher): + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + entries = super().query_queue(selectors) + if self.queue_queries == 2: + (_job_directory(workspace) / "phase").write_text("inspecting_image\n") + elif self.queue_queries == 7: + (_job_directory(workspace) / "phase").write_text("inspection_complete\n") + return entries + + launcher = PhasedLauncher(on_submit=lambda: _write_inspection(workspace, content), running_checks=6) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + progress: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=clock.now, + sleep=clock.sleep, + progress=progress.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert "Image add job 42: inspecting image." in progress + assert "Image add job 42: inspection complete." in progress + assert "Image add job 42: still running (elapsed 00:02:30)." in progress + assert "Image add job 42: still running (elapsed 00:01:00)." not in progress + assert "Image add complete: client (total elapsed 00:03:30)." in progress + + +def test_image_add_skips_empty_followed_log_lines(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + + def on_submit() -> None: + _write_inspection(workspace, content) + (_job_directory(workspace) / f"slurm-{_JOB_ID}.err").write_bytes(b"first\n\n \r\nsecond\n") + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + logs=logs.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == ["[image stderr] first", "[image stderr] second"] + + +def test_image_add_follows_bounded_job_logs_before_success_cleanup(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + + def on_submit() -> None: + _write_inspection(workspace, content) + (_job_directory(workspace) / f"slurm-{_JOB_ID}.out").write_bytes(b"importing\nready\n") + (_job_directory(workspace) / f"slurm-{_JOB_ID}.err").write_bytes(b"warning\x1b[31m\ntoken=secretvalue\n") + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + logs=logs.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == [ + "[image stdout] importing", + "[image stdout] ready", + "[image stderr] warning [31m", + "[image stderr] token=", + ] + assert not _job_directory(workspace).exists() + + +def test_image_add_follows_only_new_bytes_across_status_checks(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + log_path = _job_directory(workspace) / f"slurm-{_JOB_ID}.out" + + def on_submit() -> None: + _write_inspection(workspace, content) + log_path.write_bytes(b"first\nsecond part \xe2") + + class AppendingLauncher(_TransitionLauncher): + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + entries = super().query_queue(selectors) + if self.queue_queries == 2: + with log_path.open("ab") as log: + log.write(b"\x82") + elif self.queue_queries == 3: + with log_path.open("ab") as log: + log.write(b"\xac done\n") + return entries + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=AppendingLauncher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)).now, + sleep=lambda _: None, + logs=logs.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == ["[image stdout] first", "[image stdout] second part € done"] + + +def test_image_add_shows_carriage_return_progress_while_running(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + log_path = _job_directory(workspace) / f"slurm-{_JOB_ID}.err" + + def on_submit() -> None: + _write_inspection(workspace, content) + log_path.write_bytes(b"10%\r") + + class AppendingLauncher(_TransitionLauncher): + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + entries = super().query_queue(selectors) + if self.queue_queries == 2: + with log_path.open("ab") as log: + log.write(b"\n20%\r") + elif self.queue_queries == 3: + with log_path.open("ab") as log: + log.write(b"done\n") + return entries + + logs: list[str] = [] + sleep_calls = 0 + + def sleep(_: float) -> None: + nonlocal sleep_calls + sleep_calls += 1 + if sleep_calls == 1: + assert logs == ["[image stderr] 10%"] + + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=AppendingLauncher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + sleep=sleep, + logs=logs.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == ["[image stderr] 10%", "[image stderr] 20%", "[image stderr] done"] + + +def test_image_add_follows_failure_logs_before_cleanup(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + + def on_submit() -> None: + (_job_directory(workspace) / f"slurm-{_JOB_ID}.err").write_text("import failed\n") + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(accounting_state=SchedulerState.FAILED, on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + logs=logs.append, + ) + + with pytest.raises(SlurmServiceError): + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == ["[image stderr] import failed"] + assert not _job_directory(workspace).exists() + + +def test_image_add_ignores_symlinked_logs_without_changing_job_outcome(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + outside = tmp_path / "outside.txt" + outside.write_text("private content") + + def on_submit() -> None: + _write_inspection(workspace, content) + (_job_directory(workspace) / f"slurm-{_JOB_ID}.out").symlink_to(outside) + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + logs=logs.append, + ) + + registered = service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert registered.name == "client" + assert logs == ["Image job logs are unavailable."] + assert "private content" not in "".join(logs) + + +def test_image_add_caps_followed_log_output(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + + def on_submit() -> None: + _write_inspection(workspace, content) + (_job_directory(workspace) / f"slurm-{_JOB_ID}.out").write_bytes(b"x\n" * 200_000) + + logs: list[str] = [] + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + logs=logs.append, + ) + + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs[-1] == "Image job log display reached its 1 MiB limit; further output is hidden." + assert sum(len(line.encode("utf-8")) + 1 for line in logs) <= 1024 * 1024 + + +def test_progress_output_failure_cannot_cancel_image_job(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + content = b"client image" + source.write_bytes(content) + launcher = _Launcher(on_submit=lambda: _write_inspection(workspace, content)) + + def broken_progress(_: str) -> None: + raise RuntimeError("progress stream failed") + + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + progress=broken_progress, + ) + + registered = service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert registered.name == "client" + assert launcher.cancellations == [] + + +def test_default_image_add_preserves_collision_then_replaces_explicitly(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + first = tmp_path / "first.sqsh" + second = tmp_path / "second.sqsh" + first.write_bytes(b"first") + second.write_bytes(b"second") + current_content = b"first" + launcher = _Launcher(on_submit=lambda: _write_inspection(workspace, current_content)) + service = _service(workspace, launcher) + original = service.add(ImageBuildRequest(name="client", kind="client", source=first.as_posix())) + current_content = b"second" + + with pytest.raises(SlurmServiceError) as caught: + service.add(ImageBuildRequest(name="client", kind="client", source=second.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert service.get("client") == original + + replacement = service.add( + ImageBuildRequest(name="client", kind="client", source=second.as_posix()), + replace=True, + ) + + assert replacement.path == second.as_posix() + assert service.get("client") == replacement + assert not _job_directory(workspace).exists() + + +@pytest.mark.parametrize( + ("state", "exit_status"), + ((SchedulerState.FAILED, 1), (SchedulerState.COMPLETED, 1)), + ids=("terminal-failure", "nonzero-exit"), +) +def test_default_image_add_fails_closed_and_cleans_terminal_jobs( + tmp_path: Path, + state: SchedulerState, + exit_status: int, +) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(accounting_state=state, exit_status=exit_status) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert not _job_directory(workspace).exists() + assert not (workspace / "images" / "registry.yaml").exists() + + +def test_default_image_add_cleans_definitive_submission_failure(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher( + submission_error=SlurmSubmissionError("rejected", may_have_succeeded=False), + ) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert not _job_directory(workspace).exists() + + +def test_default_image_add_retains_ambiguous_submission_state(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher( + submission_error=SlurmSubmissionError("unknown", may_have_succeeded=True), + ) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher).add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == ("image lifecycle submission outcome is unknown; lifecycle image-test was retained") + assert _job_directory(workspace).is_dir() + + +def test_default_image_add_waits_for_cancelled_job_before_cleanup(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _DelayedCancellationLauncher() + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [30.0] * 12 + assert not _job_directory(workspace).exists() + + +def test_default_image_add_retains_state_without_terminal_cancellation_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(accounting_state=None) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError): + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [30.0] * 21 + assert _job_directory(workspace).is_dir() + + +def test_default_image_add_retains_completed_job_without_accounting_exit_evidence(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _Launcher(queue_state=SchedulerState.COMPLETED, accounting_state=None) + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + + with pytest.raises(SlurmServiceError) as caught: + _service(workspace, launcher, clock=clock).add( + ImageBuildRequest(name="client", kind="client", source=source.as_posix()) + ) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.cancellations == [_JOB_ID] + assert clock.sleep_calls == [30.0] * 21 + assert _job_directory(workspace).is_dir() + + +def test_image_add_displays_partial_log_when_cancellation_is_unconfirmed(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + + def on_submit() -> None: + (_job_directory(workspace) / f"slurm-{_JOB_ID}.err").write_bytes(b"partial diagnostic") + + logs: list[str] = [] + clock = FakeClock(datetime(2026, 9, 10, tzinfo=timezone.utc)) + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=_Launcher(accounting_state=None, on_submit=on_submit), # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=clock.now, + sleep=clock.sleep, + logs=logs.append, + ) + + with pytest.raises(SlurmServiceError): + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert logs == ["[image stderr] partial diagnostic"] + assert _job_directory(workspace).is_dir() + + +def test_default_image_add_cancels_and_cleans_on_interrupt(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + source = tmp_path / "client.sqsh" + source.write_bytes(b"client") + launcher = _ImmediateCancellationLauncher(queue_state=SchedulerState.PENDING, accounting_state=None) + + service = create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + sleep=_interrupt, + ) + + with pytest.raises(KeyboardInterrupt): + service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) + + assert launcher.cancellations == [_JOB_ID] + assert not _job_directory(workspace).exists() + + +def _service(workspace: Path, launcher: _Launcher, *, clock: FakeClock | None = None) -> SlurmImageService: + return create_slurm_image_service( + profile=_profile(workspace), + launcher=launcher, # type: ignore[arg-type] + lifecycle_id_factory=lambda: _LIFECYCLE_ID, + clock=None if clock is None else clock.now, + sleep=None if clock is None else clock.sleep, + ) + + +def _profile(workspace: Path) -> SlurmProfile: + return SlurmProfile( + schema_version=1, + scheduler=SchedulerProfile(account="research", partition="gpu"), + gpus_per_node=8, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile( + partition="cpu", + cpus_per_task=2, + memory="8G", + time_limit="04:00:00", + ), + ) + + +def _write_inspection(workspace: Path, content: bytes) -> None: + inspection = ImageInspectionRecord( + schema_version=1, + inspector_version="inspector-1", + sqsh_sha256=hashlib.sha256(content).hexdigest(), + inspection=ClientImageInspection( + kind="client", + python_implementation="cpython", + python_version="3.13.3", + python_abi="cp313", + distributions=tuple( + InstalledDistribution(name=name, version="0.9.2") + for name in ( + "data-designer", + "data-designer-config", + "data-designer-engine", + "data-designer-slurm", + ) + ) + + ( + InstalledDistribution(name="aiohttp", version="3.14.3"), + InstalledDistribution(name="pip", version="26.1"), + ), + installer_path="/usr/bin/pip", + installer_version="26.1", + ), + ) + output = _job_directory(workspace) / "output" / "inspection.json" + output.write_text(inspection.model_dump_json()) + + +def _job_directory(workspace: Path) -> Path: + return workspace / "images" / ".tmp" / "jobs" / _LIFECYCLE_ID diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py new file mode 100644 index 000000000..8914146d4 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -0,0 +1,447 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import stat +from pathlib import Path + +import pytest +import yaml + +from data_designer.slurm.config import ( + ImageBuildProfile, + ProfileSelectionSource, + SchedulerProfile, + SlurmProfile, + SlurmProfileCatalog, + load_profile_catalog, +) +from data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.services import SlurmServiceError, SlurmServiceErrorCode, create_slurm_profile_service + + +class _Launcher: + def __init__(self, gpu_counts: tuple[int, ...] = ()) -> None: + self.gpu_counts = gpu_counts + self.partitions: list[str | None] = [] + + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + self.partitions.append(partition) + return self.gpu_counts + + +class _MalformedLauncher(_Launcher): + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + del partition + raise SlurmCommandOutputError("malformed") + + +def test_profile_init_creates_deterministic_private_starter_without_side_effects(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + workspace = tmp_path / "workspace" + launcher = _Launcher() + service = create_slurm_profile_service( + profile_file=profile_file, + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("Ignored.EXAMPLE.test",), + ) + + result = service.initialize( + workspace_root=workspace, + image_build_partition="cpu", + cluster="primary", + account="research", + partition="gpu", + host_patterns=("Login.EXAMPLE.test",), + ) + + assert result.profile_file == profile_file.as_posix() + assert result.validation_command == f"data-designer slurm profile validate --profile-file {profile_file.as_posix()}" + assert stat.S_IMODE(profile_file.stat().st_mode) == 0o600 + assert not workspace.exists() + assert launcher.partitions == [] + assert load_profile_catalog(profile_file).model_dump(mode="json") == { + "schema_version": 1, + "default_cluster": "primary", + "clusters": { + "primary": { + "schema_version": 1, + "host_patterns": ["Login.EXAMPLE.test"], + "scheduler": { + "account": "research", + "partition": "gpu", + "mem_per_gpu": None, + "bin_path": None, + }, + "gpus_per_node": "auto", + "workspace_root": workspace.as_posix(), + "image_build": { + "partition": "cpu", + "cpus_per_task": 2, + "memory": "8G", + "time_limit": "04:00:00", + }, + "gpu_request_mode": "gres", + "container_mounts": [], + } + }, + } + assert "container_mounts" not in profile_file.read_text() + + +def test_profile_init_creates_missing_parent_directories(tmp_path: Path) -> None: + profile_file = tmp_path / "nested" / "config" / "profile.yml" + + result = create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert result.profile_file == profile_file.as_posix() + assert load_profile_catalog(profile_file).default_cluster == "default" + for directory in (profile_file.parent, profile_file.parent.parent): + assert stat.S_IMODE(directory.stat().st_mode) == 0o700 + assert stat.S_IMODE(profile_file.stat().st_mode) == 0o600 + + +def test_profile_init_creates_usable_private_directories_with_restrictive_umask(tmp_path: Path) -> None: + profile_file = tmp_path / "nested" / "config" / "profile.yml" + previous_umask = os.umask(0o777) + try: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + finally: + os.umask(previous_umask) + + assert load_profile_catalog(profile_file).default_cluster == "default" + for directory in (profile_file.parent, profile_file.parent.parent): + assert stat.S_IMODE(directory.stat().st_mode) == 0o700 + + +def test_profile_init_handles_unsupported_nofollow_chmod(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + profile_file = tmp_path / "nested" / "profile.yml" + original_chmod = os.chmod + + def reject_nofollow_chmod( + path: str, + mode: int, + *, + dir_fd: int | None = None, + follow_symlinks: bool = True, + ) -> None: + if not follow_symlinks: + raise NotImplementedError("chmod without following symlinks is unavailable") + original_chmod(path, mode, dir_fd=dir_fd, follow_symlinks=follow_symlinks) + + monkeypatch.setattr(os, "chmod", reject_nofollow_chmod) + + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert load_profile_catalog(profile_file).default_cluster == "default" + assert stat.S_IMODE(profile_file.parent.stat().st_mode) == 0o700 + + +def test_profile_init_does_not_follow_replaced_parent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + profile_file = tmp_path / "nested" / "profile.yml" + redirect = tmp_path / "redirect" + redirect.mkdir() + original_mkdir = os.mkdir + + def replace_new_directory(name: str, mode: int = 0o777, *, dir_fd: int | None = None) -> None: + if name == "nested" and dir_fd is not None: + os.symlink(redirect, name, dir_fd=dir_fd) + return + original_mkdir(name, mode, dir_fd=dir_fd) + + monkeypatch.setattr(os, "mkdir", replace_new_directory) + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert tuple(redirect.iterdir()) == () + + +def test_profile_init_and_auto_gpu_resolution_accept_multiple_partitions(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + workspace = tmp_path / "workspace" + workspace.mkdir() + launcher = _Launcher((8, 8)) + service = create_slurm_profile_service( + profile_file=profile_file, + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("login.example.test",), + ) + + service.initialize( + workspace_root=workspace, + image_build_partition="cpu", + partition="gpu-a,gpu-b", + ) + assert load_profile_catalog(profile_file).clusters["default"].scheduler.partition == "gpu-a,gpu-b" + assert load_profile_catalog(profile_file).clusters["default"].host_patterns == [] + + result = service.validate() + + assert result.gpus_per_node == 8 + assert result.selection_source is ProfileSelectionSource.DEFAULT + assert launcher.partitions == ["gpu-a,gpu-b"] + + +def test_profile_init_uses_explicit_environment_then_home_path_precedence(tmp_path: Path) -> None: + environment_file = tmp_path / "environment.yml" + explicit_file = tmp_path / "explicit.json" + environment = {"DATA_DESIGNER_SLURM_PROFILE_FILE": environment_file.as_posix()} + + explicit = create_slurm_profile_service( + profile_file=explicit_file, + environ=environment, + home_directory=tmp_path, + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + from_environment = create_slurm_profile_service( + environ=environment, + home_directory=tmp_path, + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + from_home = create_slurm_profile_service( + environ={}, + home_directory=tmp_path, + ).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + ) + + assert explicit.profile_file == explicit_file.as_posix() + assert from_environment.profile_file == environment_file.as_posix() + assert from_home.profile_file == (tmp_path / ".data-designer-slurm-profile.yml").as_posix() + assert yaml.safe_load(explicit_file.read_text()) == json.loads(explicit_file.read_text()) + assert load_profile_catalog(from_home.profile_file).clusters["default"].host_patterns == [] + + +def test_profile_init_refuses_to_overwrite_or_leave_temporary_files(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + profile_file.write_text("owned by caller\n") + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert profile_file.read_text() == "owned by caller\n" + assert tuple(tmp_path.iterdir()) == (profile_file,) + + +def test_profile_init_refuses_to_follow_dangling_destination_symlink(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + target = tmp_path / "target.yml" + profile_file.symlink_to(target) + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert profile_file.is_symlink() + assert not target.exists() + + +@pytest.mark.parametrize( + "profile_file", + (Path("profile.txt"), Path("missing/profile.txt")), + ids=("unsupported-suffix", "unsupported-suffix-with-missing-parent"), +) +def test_profile_init_rejects_invalid_destination(tmp_path: Path, profile_file: Path) -> None: + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=tmp_path / profile_file).initialize( + workspace_root=tmp_path / "workspace", + image_build_partition="cpu", + host_patterns=("login",), + ) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_profile_validate_uses_strict_loader_selection_and_effective_checks(tmp_path: Path) -> None: + selected_workspace = tmp_path / "primary" + selected_workspace.mkdir() + catalog = _catalog(selected_workspace, tmp_path / "unused") + profile_file = tmp_path / "profile.yml" + profile_file.write_text(yaml.safe_dump(catalog.model_dump(mode="json"), sort_keys=False)) + launcher = _Launcher((8, 8)) + + result = create_slurm_profile_service( + profile_file=profile_file, + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("LOGIN-01.EXAMPLE.TEST", "login-01.example.test"), + ).validate() + + assert result.profile_file == profile_file.as_posix() + assert result.hostnames == ("login-01.example.test",) + assert result.default_cluster == "fallback" + assert result.selected_cluster == "primary" + assert result.selection_source is ProfileSelectionSource.HOSTNAME + assert result.matched_pattern == "login-*.example.test" + assert [(match.cluster, match.patterns) for match in result.matched_clusters] == [ + ("primary", ("login-*.example.test",)) + ] + assert result.workspace_root == selected_workspace.as_posix() + assert result.image_root == (selected_workspace / "images").as_posix() + assert result.registry_file == (selected_workspace / "images" / "registry.yaml").as_posix() + assert result.gpus_per_node == 8 + assert launcher.partitions == ["gpu"] + assert tuple(selected_workspace.iterdir()) == () + assert not (tmp_path / "unused").exists() + + +def test_profile_validate_honors_explicit_cluster_and_fixed_gpu_count(tmp_path: Path) -> None: + workspace = tmp_path / "fallback" + workspace.mkdir() + catalog = _catalog(tmp_path / "unused", workspace) + launcher = _Launcher() + + result = create_slurm_profile_service( + catalog=catalog, + cluster="fallback", + launcher=launcher, # type: ignore[arg-type] + hostname_resolver=lambda: ("login-01.example.test",), + ).validate() + + assert result.selected_cluster == "fallback" + assert result.selection_source is ProfileSelectionSource.EXPLICIT + assert result.gpus_per_node == 4 + assert launcher.partitions == [] + assert [(match.cluster, match.patterns) for match in result.matched_clusters] == [ + ("primary", ("login-*.example.test",)) + ] + + +@pytest.mark.parametrize("gpu_counts", ((), (4, 8)), ids=("missing", "heterogeneous")) +def test_profile_validate_rejects_ambiguous_automatic_gpu_count( + tmp_path: Path, + gpu_counts: tuple[int, ...], +) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service( + profile=_profile(workspace), + launcher=_Launcher(gpu_counts), # type: ignore[arg-type] + ).validate() + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "eligible Slurm nodes do not report one GPU count" + + +def test_profile_validate_rejects_ambiguous_selection_and_unavailable_workspace(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + overlapping = SlurmProfileCatalog( + schema_version=1, + default_cluster="first", + clusters={ + "first": _profile(workspace, host_patterns=["node-*"]), + "second": _profile(workspace, host_patterns=["*.example.test"]), + }, + ) + + with pytest.raises(SlurmServiceError) as ambiguous: + create_slurm_profile_service( + catalog=overlapping, + hostname_resolver=lambda: ("node-01.example.test",), + ).validate() + + missing = create_slurm_profile_service( + profile=_profile(tmp_path / "missing", gpus_per_node=4), + hostname_resolver=lambda: ("login",), + ) + with pytest.raises(SlurmServiceError) as unavailable: + missing.validate() + + assert ambiguous.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert unavailable.value.code is SlurmServiceErrorCode.UNAVAILABLE + + +def test_profile_validate_rejects_duplicate_keys_with_stable_error(tmp_path: Path) -> None: + profile_file = tmp_path / "profile.yml" + profile_file.write_text("schema_version: 1\nschema_version: 1\n") + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service(profile_file=profile_file).validate() + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert str(caught.value) == "profile configuration cannot be resolved" + + +def test_profile_validate_reports_malformed_slurm_output_as_unavailable(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + + with pytest.raises(SlurmServiceError) as caught: + create_slurm_profile_service( + profile=_profile(workspace), + launcher=_MalformedLauncher(), # type: ignore[arg-type] + ).validate() + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "Slurm is unavailable" + + +def _catalog(primary_workspace: Path, fallback_workspace: Path) -> SlurmProfileCatalog: + return SlurmProfileCatalog( + schema_version=1, + default_cluster="fallback", + clusters={ + "primary": _profile(primary_workspace, host_patterns=["login-*.example.test"]), + "fallback": _profile(fallback_workspace, gpus_per_node=4), + }, + ) + + +def _profile( + workspace: Path, + *, + host_patterns: list[str] | None = None, + gpus_per_node: int | str = "auto", +) -> SlurmProfile: + return SlurmProfile( + schema_version=1, + host_patterns=[] if host_patterns is None else host_patterns, + scheduler=SchedulerProfile(partition="gpu"), + gpus_per_node=gpus_per_node, + workspace_root=workspace.as_posix(), + image_build=ImageBuildProfile( + partition="cpu", + cpus_per_task=2, + memory="8G", + time_limit="04:00:00", + ), + ) diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py new file mode 100644 index 000000000..9f6c5e972 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -0,0 +1,735 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pickle +from pathlib import Path +from unittest.mock import Mock + +import pytest +from slurm_test_fakes import ( + FakeBatchScriptRenderer, + FakeBenchmarkBackend, + FakeImageResolver, + FakeRunPlanningBackend, + FakeScriptError, +) + +import data_designer.slurm.services as slurm_services +from data_designer.errors import DataDesignerError +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageInspectionRecord, + ImageKind, + ImageRef, +) +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmBatchScriptRenderer, + SlurmBenchmarkBackend, + SlurmBenchmarkService, + SlurmCollectionExecution, + SlurmImageManager, + SlurmImageResolver, + SlurmImageService, + SlurmRetryExecution, + SlurmRunArtifactPublisher, + SlurmRunBackend, + SlurmRunExecution, + SlurmRunPlanner, + SlurmRunService, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) +from data_designer.slurm.state import CollectionState + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" + + +@pytest.fixture +def single_node_script() -> str: + return (GOLDEN_DIRECTORY / "single_node.sbatch").read_text() + + +@pytest.fixture +def manifest_with_matching_config_digest( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, +) -> BenchmarkManifest: + reference = benchmark_manifest.benchmark_config.model_copy(update={"sha256": benchmark_config.compute_sha256()}) + return benchmark_manifest.model_copy(update={"benchmark_config": reference}) + + +def test_run_service_returns_plan_for_requested_config( + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + planner = FakeRunPlanningBackend(((authored_run_single, single_node_plan),)) + renderer = FakeBatchScriptRenderer(()) + + result = SlurmRunService(planner, renderer).plan(authored_run_single) + + assert result is single_node_plan + assert planner.calls == [authored_run_single] + assert renderer.calls == [] + planner.assert_complete() + renderer.assert_complete() + + +def test_run_service_rejects_invalid_plan_result(authored_run_single: DataDesignerSlurmConfig) -> None: + planner = FakeRunPlanningBackend(((authored_run_single, object()),)) # type: ignore[arg-type] + renderer = FakeBatchScriptRenderer(()) + + with pytest.raises(SlurmServiceError) as caught: + SlurmRunService(planner, renderer).plan(authored_run_single) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + assert str(caught.value) == "plan run failed" + assert planner.calls == [authored_run_single] + assert renderer.calls == [] + planner.assert_complete() + renderer.assert_complete() + + +def test_run_service_renders_attempt_without_replanning( + single_node_plan: ResolvedSlurmRunPlan, + single_node_script: str, +) -> None: + script = single_node_script.replace( + 'readonly DD_ATTEMPT_ORDINAL="0001"', + 'readonly DD_ATTEMPT_ORDINAL="0002"', + ) + planner = FakeRunPlanningBackend(()) + renderer = FakeBatchScriptRenderer((((single_node_plan, 2), script),)) + + result = SlurmRunService(planner, renderer).render_attempt(single_node_plan, attempt_ordinal=2) + + assert result == script + assert planner.calls == [] + assert renderer.calls == [(single_node_plan, 2)] + planner.assert_complete() + renderer.assert_complete() + + +@pytest.mark.parametrize("invalid_script", ["", object()], ids=["empty", "wrong-type"]) +def test_run_service_rejects_invalid_rendered_script( + single_node_plan: ResolvedSlurmRunPlan, + invalid_script: object, +) -> None: + planner = FakeRunPlanningBackend(()) + renderer = FakeBatchScriptRenderer((((single_node_plan, 1), invalid_script),)) # type: ignore[arg-type] + service = SlurmRunService(planner, renderer) + + with pytest.raises(SlurmServiceError) as caught: + service.render_attempt(single_node_plan) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RENDER_ATTEMPT + assert str(caught.value) == "render attempt failed" + assert planner.calls == [] + assert renderer.calls == [(single_node_plan, 1)] + planner.assert_complete() + renderer.assert_complete() + + +def test_run_service_rejects_plan_for_another_config( + authored_run_single: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + planner = FakeRunPlanningBackend(((authored_run_single, multi_node_plan),)) + renderer = FakeBatchScriptRenderer(()) + + with pytest.raises(SlurmServiceError) as caught: + SlurmRunService(planner, renderer).plan(authored_run_single) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + assert planner.calls == [authored_run_single] + assert renderer.calls == [] + planner.assert_complete() + renderer.assert_complete() + + +@pytest.mark.parametrize("attempt_ordinal", [0, -1, True, 1.5]) +def test_run_service_rejects_invalid_attempt_ordinals( + single_node_plan: ResolvedSlurmRunPlan, + attempt_ordinal: object, +) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(())) + + with pytest.raises(SlurmServiceError) as caught: + service.render_attempt(single_node_plan, attempt_ordinal=attempt_ordinal) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RENDER_ATTEMPT + + +def test_run_service_rejects_untyped_resolved_plan() -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(())) + + with pytest.raises(SlurmServiceError) as caught: + service.render_attempt(object()) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RENDER_ATTEMPT + + +def test_run_service_rejects_untyped_config() -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(())) + + with pytest.raises(SlurmServiceError) as caught: + service.plan(object()) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + + +def test_run_service_delegates_execute_actions( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, +) -> None: + class RunBackend: + def __init__(self) -> None: + self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool]] = [] + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + ) -> SlurmRunExecution: + self.calls.append((config, source_root, dry_run)) + return SlurmRunExecution( + run_id="run-0001", + state="dry_run", + plan_sha256="1" * 64, + shard_count=1, + batch_script="script", + ) + + def status(self, run_id): + raise AssertionError(run_id) + + def cancel(self, run_id): + raise AssertionError(run_id) + + backend = RunBackend() + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + result = service.execute(authored_run_single, source_root=tmp_path, dry_run=True) + + assert result.run_id == "run-0001" + assert backend.calls == [(authored_run_single, tmp_path.resolve(), True)] + + +def test_run_service_delegates_retry_with_stable_shard_order() -> None: + backend = Mock(spec=SlurmRunBackend) + expected = SlurmRetryExecution( + run_id="run-0001", + state="submitted", + shard_ids=("shard-00001", "shard-00003"), + attempt_ids=("attempt-0002", "attempt-0004"), + effective_resume_mode="always", + job_id=43, + ) + backend.retry.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + result = service.retry( + "42", + shard_ids=("shard-00003", "shard-00001"), + resume="always", + ) + + assert result is expected + backend.retry.assert_called_once_with( + "42", + shard_ids=("shard-00001", "shard-00003"), + resume="always", + dry_run=False, + ) + + +@pytest.mark.parametrize( + ("shard_ids", "resume", "dry_run"), + [ + ((), "never", False), + (("shard-00001", "shard-00001"), "never", False), + (None, "sometimes", False), + (None, "never", 1), + ], +) +def test_run_service_rejects_invalid_retry_actions( + shard_ids: object, + resume: object, + dry_run: object, +) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), Mock(spec=SlurmRunBackend)) + + with pytest.raises(SlurmServiceError) as caught: + service.retry( # type: ignore[arg-type] + "run-0001", + shard_ids=shard_ids, + resume=resume, + dry_run=dry_run, + ) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RETRY_RUN + + +def test_run_service_normalizes_collection_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.chdir(tmp_path) + backend = Mock(spec=SlurmRunBackend) + expected = SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=44, + output_path=(tmp_path / "collected").as_posix(), + num_partitions=2, + ) + backend.collect.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + result = service.collect("runs/run-0001", destination="collected", num_partitions=2) + + assert result is expected + backend.collect.assert_called_once_with( + tmp_path / "runs/run-0001", + destination=tmp_path / "collected", + num_partitions=2, + ) + + +def test_run_service_defaults_collection_to_one_partition(tmp_path: Path) -> None: + backend = Mock(spec=SlurmRunBackend) + expected = SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=44, + output_path=(tmp_path / "collected").as_posix(), + num_partitions=1, + ) + backend.collect.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + assert service.collect(tmp_path / "runs/run-0001", destination=tmp_path / "collected") is expected + backend.collect.assert_called_once_with( + tmp_path / "runs/run-0001", + destination=tmp_path / "collected", + num_partitions=None, + ) + + +def test_run_service_forwards_run_id_without_destination(tmp_path: Path) -> None: + backend = Mock(spec=SlurmRunBackend) + expected = SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=44, + output_path=(tmp_path / "runs/run-0001/merge-output").as_posix(), + num_partitions=1, + ) + backend.collect.return_value = expected + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) + + assert service.collect("run-0001") is expected + backend.collect.assert_called_once_with("run-0001", destination=None, num_partitions=None) + + +@pytest.mark.parametrize("num_partitions", [0, -1, True, 1.5]) +def test_run_service_rejects_invalid_collection_partitions(num_partitions: object) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), Mock(spec=SlurmRunBackend)) + + with pytest.raises(SlurmServiceError) as caught: + service.collect("runs/run-0001", destination="collected", num_partitions=num_partitions) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.COLLECT_RUN + + +def test_run_service_normalizes_and_redacts_unexpected_backend_errors( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + planner = FakeRunPlanningBackend(((authored_run_single, RuntimeError("secret backend detail")),)) + service = SlurmRunService(planner, FakeBatchScriptRenderer(())) + + with pytest.raises(SlurmServiceError) as caught: + service.plan(authored_run_single) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + assert str(caught.value) == "plan run failed" + assert caught.value.__suppress_context__ + + +def test_run_service_redacts_matching_internal_backend_errors( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + error = SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.PLAN_RUN, + "secret backend detail", + ) + service = SlurmRunService( + FakeRunPlanningBackend(((authored_run_single, error),)), + FakeBatchScriptRenderer(()), + ) + + with pytest.raises(SlurmServiceError) as caught: + service.plan(authored_run_single) + + assert caught.value is not error + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + assert str(caught.value) == "plan run failed" + assert caught.value.__suppress_context__ + + +def test_run_service_preserves_matching_public_safe_errors( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + error = SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.PLAN_RUN, + "planning dependencies are unavailable", + ) + service = SlurmRunService( + FakeRunPlanningBackend(((authored_run_single, error),)), + FakeBatchScriptRenderer(()), + ) + + with pytest.raises(SlurmServiceError) as caught: + service.plan(authored_run_single) + + assert caught.value is error + + +@pytest.mark.parametrize( + ("code", "message", "expected_message"), + [ + (SlurmServiceErrorCode.NOT_FOUND, "image not found", "image not found"), + (SlurmServiceErrorCode.INTERNAL, "image lookup failed internally", "plan run failed"), + ], +) +def test_run_service_reattributes_nested_public_safe_errors( + authored_run_single: DataDesignerSlurmConfig, + code: SlurmServiceErrorCode, + message: str, + expected_message: str, +) -> None: + error = SlurmServiceError( + code, + SlurmServiceOperation.RESOLVE_IMAGE, + message, + ) + service = SlurmRunService( + FakeRunPlanningBackend(((authored_run_single, error),)), + FakeBatchScriptRenderer(()), + ) + + with pytest.raises(SlurmServiceError) as caught: + service.plan(authored_run_single) + + assert caught.value is not error + assert caught.value.code is code + assert caught.value.operation is SlurmServiceOperation.PLAN_RUN + assert str(caught.value) == expected_message + assert caught.value.__suppress_context__ + + +def test_run_service_normalizes_renderer_errors( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + planner = FakeRunPlanningBackend(()) + renderer = FakeBatchScriptRenderer((((single_node_plan, 1), RuntimeError("secret renderer detail")),)) + service = SlurmRunService(planner, renderer) + + with pytest.raises(SlurmServiceError) as caught: + service.render_attempt(single_node_plan) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RENDER_ATTEMPT + assert str(caught.value) == "render attempt failed" + assert caught.value.__suppress_context__ + assert planner.calls == [] + assert renderer.calls == [(single_node_plan, 1)] + planner.assert_complete() + renderer.assert_complete() + + +def test_fake_scripting_errors_are_not_normalized( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(())) + + with pytest.raises(FakeScriptError, match="unexpected run plan"): + service.plan(authored_run_single) + + +def test_service_boundary_does_not_swallow_cancellation_signals( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + service = SlurmRunService( + FakeRunPlanningBackend(((authored_run_single, KeyboardInterrupt()),)), + FakeBatchScriptRenderer(()), + ) + + with pytest.raises(KeyboardInterrupt): + service.plan(authored_run_single) + + +def test_image_service_returns_correlated_resolution( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + image = single_node_plan.deployments[0].image + reference = image.authored_ref + resolver = FakeImageResolver((((reference, ImageKind.SERVING), image),)) + + result = SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.SERVING) + + assert result is image + assert resolver.calls == [(reference, ImageKind.SERVING)] + resolver.assert_complete() + + +def test_image_service_rejects_invalid_resolution(single_node_plan: ResolvedSlurmRunPlan) -> None: + reference = single_node_plan.deployments[0].image.authored_ref + resolver = FakeImageResolver((((reference, ImageKind.SERVING), object()),)) # type: ignore[arg-type] + + with pytest.raises(SlurmServiceError) as caught: + SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.SERVING) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + assert str(caught.value) == "resolve image failed" + assert resolver.calls == [(reference, ImageKind.SERVING)] + resolver.assert_complete() + + +def test_image_service_rejects_untyped_image_kind(single_node_plan: ResolvedSlurmRunPlan) -> None: + service = SlurmImageService(FakeImageResolver(())) + + with pytest.raises(SlurmServiceError) as caught: + service.resolve(single_node_plan.deployments[0].image.authored_ref, expected_kind="serving") # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + + +def test_image_service_rejects_untyped_reference() -> None: + service = SlurmImageService(FakeImageResolver(())) + + with pytest.raises(SlurmServiceError) as caught: + service.resolve(object(), expected_kind=ImageKind.SERVING) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + + +def test_image_service_requires_package_owned_manager() -> None: + service = SlurmImageService(FakeImageResolver(())) + + with pytest.raises(SlurmServiceError) as caught: + service.list() + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert caught.value.operation is SlurmServiceOperation.LIST_IMAGES + + +@pytest.mark.parametrize("mismatch", ["reference", "kind"]) +def test_image_service_rejects_uncorrelated_results( + single_node_plan: ResolvedSlurmRunPlan, + client_image_inspection: ImageInspectionRecord, + mismatch: str, +) -> None: + source = single_node_plan.deployments[0].image + reference = source.authored_ref + if mismatch == "reference": + image = source.model_copy(update={"authored_ref": ImageRef(name="other-serving")}) + else: + image = source.model_copy(update={"inspection": client_image_inspection}) + resolver = FakeImageResolver((((reference, ImageKind.SERVING), image),)) + + with pytest.raises(SlurmServiceError) as caught: + SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.SERVING) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + assert resolver.calls == [(reference, ImageKind.SERVING)] + resolver.assert_complete() + + +def test_benchmark_service_delegates_run_and_analysis( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + manifest_with_matching_config_digest: BenchmarkManifest, + benchmark_report: BenchmarkReport, +) -> None: + benchmark_manifest = manifest_with_matching_config_digest + source_root = Path.cwd().resolve() + backend = FakeBenchmarkBackend( + 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, 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: + 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) + + 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 == [request] + backend.assert_complete() + + +def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkManifest) -> None: + backend = FakeBenchmarkBackend( + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), object()),), # type: ignore[arg-type] + ) + + with pytest.raises(SlurmServiceError) as caught: + SlurmBenchmarkService(backend).analyze(benchmark_manifest.benchmark_id) + + 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, False)] + backend.assert_complete() + + +def test_benchmark_service_rejects_manifest_for_another_config( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + manifest_with_matching_config_digest: BenchmarkManifest, +) -> 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}) + 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 == [request] + backend.assert_complete() + + +def test_benchmark_service_rejects_untyped_config() -> None: + service = SlurmBenchmarkService(FakeBenchmarkBackend()) + + with pytest.raises(SlurmServiceError) as caught: + service.run(object()) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK + + +@pytest.mark.parametrize( + ("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( # 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 + + +def test_benchmark_service_rejects_uncorrelated_report( + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, +) -> None: + report = benchmark_report.model_copy(update={"benchmark_id": "other-benchmark"}) + backend = FakeBenchmarkBackend( + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), report),), + ) + + with pytest.raises(SlurmServiceError) as caught: + SlurmBenchmarkService(backend).analyze(benchmark_manifest.benchmark_id) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False, False)] + backend.assert_complete() + + +def test_service_errors_reject_unstable_messages() -> None: + with pytest.raises(ValueError, match="control characters"): + SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "unsafe\nmessage", + ) + + +def test_service_errors_use_the_data_designer_error_hierarchy() -> None: + assert issubclass(SlurmServiceError, DataDesignerError) + + +@pytest.mark.parametrize( + ("name", "contract"), + [ + ("SlurmBatchScriptRenderer", SlurmBatchScriptRenderer), + ("SlurmBenchmarkBackend", SlurmBenchmarkBackend), + ("SlurmImageManager", SlurmImageManager), + ("SlurmImageResolver", SlurmImageResolver), + ("SlurmRunArtifactPublisher", SlurmRunArtifactPublisher), + ("SlurmRunBackend", SlurmRunBackend), + ("SlurmRunPlanner", SlurmRunPlanner), + ], +) +def test_services_export_dependency_contracts(name: str, contract: type[object]) -> None: + assert name in slurm_services.__all__ + assert getattr(slurm_services, name) is contract + + +def test_service_errors_round_trip_through_pickle() -> None: + error = SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.PLAN_RUN, + "planning dependencies are unavailable", + ) + + restored = pickle.loads(pickle.dumps(error)) + + assert restored.code is error.code + assert restored.operation is error.operation + assert str(restored) == str(error) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py new file mode 100644 index 000000000..5da340517 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -0,0 +1,1735 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import data_designer.slurm.services.retry_collection as retry_collection_module +import data_designer.slurm.services.wiring as wiring_module +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.client.dependencies import ResolvedClientDependencies +from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector +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 +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmProcessExitCode, + SlurmQueueEntry, + SlurmSubmissionMatch, +) +from data_designer.slurm.launcher.scheduler_path import SchedulerPathError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, + create_slurm_benchmark_service, + create_slurm_image_service, + create_slurm_run_service, +) +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + CollectionState, + RunManifest, + SchedulerIdentity, + SchedulerState, + ShardManifest, + SlurmStateWriter, + StateConflictError, + StateNotFoundError, +) + + +@pytest.fixture(autouse=True) +def _use_persisted_client_runtime( + monkeypatch: pytest.MonkeyPatch, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + monkeypatch.setattr(ClientRuntimeInspector, "inspect", lambda self: single_node_plan.client.runtime) + monkeypatch.setattr(wiring_module, "resolve_scheduler_bin_path", lambda configured: configured or "/opt/slurm/bin") + + +class _Launcher: + def __init__( + self, + gpu_counts: tuple[int, ...] = (), + *, + 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] = [] + self.releases: list[int] = [] + self.held_submissions: list[bool] = [] + self.exported_environments: list[dict[str, str]] = [] + self.queue_entries: tuple[SlurmQueueEntry, ...] = () + self.accounting_entries: tuple[SlurmAccountingEntry, ...] = () + self.submission_matches: tuple[SlurmSubmissionMatch, ...] = () + 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, + script: str, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: + 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]) + + def cancel(self, job_id: int) -> None: + self.cancellations.append(job_id) + if self.cancel_error is not None: + raise self.cancel_error + + def release(self, job_id: int) -> None: + self.releases.append(job_id) + if self.release_error is not None: + raise self.release_error + + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + assert partition is not None + return self.gpu_counts + + def query_queue(self, selectors: object) -> tuple[SlurmQueueEntry, ...]: + del selectors + return self.queue_entries + + def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ...]: + del selectors + return self.accounting_entries + + def query_submissions_by_name( + self, + job_name: str, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + del job_name, submitted_after + return self.submission_matches + + +class _Publisher: + def __init__( + self, + *, + initialization_error: BaseException | None = None, + 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, int, datetime]] = [] + self.initialization_error = initialization_error + self.submission_error = submission_error + + def initialize( + self, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + dependencies: ResolvedClientDependencies, + builder_payload: dict[str, object] | None, + *, + force: bool, + ) -> None: + if self.initialization_error is not None: + raise self.initialization_error + 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: + if self.submission_error is not None: + raise self.submission_error + self.submissions.append((plan.run_id, job_id, submitted_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: + return profile_catalog.clusters["primary"].model_copy(update={"workspace_root": tmp_path.as_posix()}) + + +def _register_images( + tmp_path: Path, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> None: + store = ImageRegistryStore(tmp_path) + pairs = tuple( + (deployment.server.image, resolved.image) + for deployment, resolved in zip(authored.deployments, plan.deployments, strict=True) + ) + for index, (reference, resolved) in enumerate(pairs): + path = tmp_path / f"image-{index}.sqsh" + path.write_bytes(f"image-{index}".encode()) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + inspection = resolved.inspection.model_copy(update={"sqsh_sha256": digest}) + assert reference.name is not None + image = RegisteredImage( + schema_version=1, + name=reference.name, + path=path.as_posix(), + sqsh_sha256=digest, + inspection=inspection, + ) + store.register(image, verify_before_publish=lambda _: None) + + +def test_production_wiring_dry_run_resolves_and_renders_without_submission( + 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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + result = service.execute(authored_run_single, source_root=tmp_path, dry_run=True) + + assert result.run_id == "run-wired" + assert result.state == "dry_run" + assert result.batch_script is not None + assert "#SBATCH --array=0" in result.batch_script + assert 'export PATH="/opt/slurm/bin:/usr/local/sbin:' in result.batch_script + plan = service.plan(authored_run_single, source_root=tmp_path) + assert plan.scheduler_bin_path == "/opt/slurm/bin" + assert plan.selected_profile.profile.scheduler.bin_path is None + assert launcher.submissions == [] + + +def test_missing_srun_fails_before_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def missing_scheduler(_configured: str | None) -> str: + raise SchedulerPathError("srun was not found on the submit host PATH") + + monkeypatch.setattr(wiring_module, "resolve_scheduler_bin_path", missing_scheduler) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-missing-srun", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path, dry_run=False) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert "srun was not found" in str(caught.value) + assert launcher.submissions == [] + + +def test_dry_run_reports_missing_native_runtime_dependency( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_inspection(_inspector: ClientRuntimeInspector) -> None: + raise ClientRuntimeInspectionError("client Python requires aiohttp>=3.14.3,<4 for the inference proxy") + + monkeypatch.setattr(ClientRuntimeInspector, "inspect", fail_inspection) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-missing-runtime-dependency", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path, dry_run=True) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert str(caught.value) == "client Python requires aiohttp>=3.14.3,<4 for the inference proxy" + + +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, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.builder.inline + assert payload is not None + (tmp_path / "builder.json").write_bytes(canonical_json(payload)) + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder.json")}) + _register_images(tmp_path, authored, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + plan = service.plan(authored, source_root=tmp_path) + + assert plan.builder.authored_source == "builder.json" + assert plan.builder.source is not None + assert plan.builder.source.path == (tmp_path / "runs/run-wired/builder-config.json").as_posix() + + +def test_execute_reports_missing_builder_source( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy(update={"builder": BuilderInput(source="builder-20b.yaml")}) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.EXECUTE_RUN + assert str(caught.value) == ( + f"run configuration cannot be prepared: builder source file not found: {tmp_path / 'builder-20b.yaml'}" + ) + assert not launcher.submissions + + +@pytest.mark.parametrize("invalid_columns", (1, 100)) +def test_execute_rejects_invalid_builtin_builder_before_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + invalid_columns: int, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [ + {"name": f"invalid-{index}", "column_type": "sampler", "sampler_type": "not-a-sampler"} + for index in range(invalid_columns) + ] + authored = DataDesignerSlurmConfig.model_validate(payload) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="builder configuration failed validation") as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert len(str(caught.value)) <= 512 + assert not launcher.submissions + + +def test_submit_validation_defers_overlay_plugin_components( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["client"]["dependencies"]["requirements"] = ["example-plugin==1.0.0"] + payload["builder"]["inline"]["data_designer"]["columns"] = [ + {"name": "plugin", "column_type": "example-plugin-column"} + ] + payload["builder"]["inline"]["data_designer"]["processors"] = [{"processor_type": "example-plugin-processor"}] + authored = DataDesignerSlurmConfig.model_validate(payload) + + wiring_module._validate_builder_for_submission( # noqa: SLF001 + authored, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + + +def test_submit_validation_rejects_invalid_model_with_overlay_plugin_column( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["client"]["dependencies"]["requirements"] = ["example-plugin==1.0.0"] + builder = payload["builder"]["inline"]["data_designer"] + builder["columns"] = [{"name": "plugin", "column_type": "example-plugin-column"}] + builder["model_configs"] = [{"alias": "generator", "model": 123}] + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmServiceError, match="builder configuration failed validation") as caught: + wiring_module._validate_builder_for_submission( # noqa: SLF001 + authored, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_submit_validation_defers_overlay_plugin_seed_source( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["client"]["dependencies"]["requirements"] = ["example-plugin==1.0.0"] + payload["builder"]["inline"]["data_designer"]["seed_config"] = { + "source": {"seed_type": "plugin-seed-source", "dataset": "demo"}, + "sampling_strategy": "ordered", + } + authored = DataDesignerSlurmConfig.model_validate(payload) + + wiring_module._validate_builder_for_submission( # noqa: SLF001 + authored, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + + payload["builder"]["inline"]["data_designer"]["seed_config"]["sampling_strategy"] = "invalid" + invalid = DataDesignerSlurmConfig.model_validate(payload) + with pytest.raises(SlurmServiceError, match="builder configuration failed validation") as caught: + wiring_module._validate_builder_for_submission( # noqa: SLF001 + invalid, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_submit_validation_defers_compute_only_local_seed_path( + authored_run_single: DataDesignerSlurmConfig, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["seed_config"] = { + "source": {"seed_type": "local", "path": "/compute-only/seed.parquet"}, + "sampling_strategy": "ordered", + } + authored = DataDesignerSlurmConfig.model_validate(payload) + + wiring_module._validate_builder_for_submission( # noqa: SLF001 + authored, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + + payload["builder"]["inline"]["data_designer"]["seed_config"]["source"]["path"] = "/compute-only/seed.txt" + invalid = DataDesignerSlurmConfig.model_validate(payload) + with pytest.raises(SlurmServiceError, match="builder configuration failed validation") as caught: + wiring_module._validate_builder_for_submission( # noqa: SLF001 + invalid, + None, + SlurmServiceOperation.EXECUTE_RUN, + ) + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_production_wiring_submits_after_publisher_initialization( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("NEMO_DEPLOYMENT_TYPE", raising=False) + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + publisher = _Publisher() + submitted_at = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=publisher, # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: submitted_at, + package_version="0.9.2", + ) + + result = service.execute(authored_run_single, source_root=tmp_path) + + assert result.state == "submitted" + assert result.job_id == 42 + assert publisher.initializations == [("run-wired", False)] + assert publisher.submissions == [("run-wired", 42, submitted_at)] + assert len(launcher.submissions) == 1 + assert launcher.held_submissions == [True] + assert launcher.releases == [42] + assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": "slurm"}] + assert (tmp_path / "managed-assets").is_dir() + + +def test_execute_releases_run_reference_lock_before_submission( + 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) + + class _CheckingLauncher(_Launcher): + def submit_script( + self, + script: str, + *, + hold: bool = False, + export_environment: Mapping[str, str] | None = None, + ) -> SlurmJobSubmissionReceipt: + def acquire_deletion_lock() -> None: + with ImageRegistryStore(tmp_path).acquire_run_reference_lock(exclusive=True): + pass + + executor = ThreadPoolExecutor(max_workers=1) + try: + executor.submit(acquire_deletion_lock).result(timeout=2) + finally: + executor.shutdown(wait=False, cancel_futures=True) + return super().submit_script(script, hold=hold, export_environment=export_environment) + + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(), # type: ignore[arg-type] + launcher=_CheckingLauncher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + assert service.execute(authored_run_single, source_root=tmp_path).state == "submitted" + + +def test_production_wiring_exports_referenced_secrets_to_the_allocation( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployment = authored_run_single.deployments[0] + server = deployment.server.model_copy( + update={"environment": {"HF_TOKEN": SecretRef(type="secret", environment="SOURCE_TOKEN")}} + ) + authored = authored_run_single.model_copy( + update={"deployments": [deployment.model_copy(update={"server": server})]} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + source_environment={"SOURCE_TOKEN": "secret-value"}, + ) + + service.execute(authored, source_root=tmp_path) + + assert launcher.exported_environments == [ + {"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": "slurm", "SOURCE_TOKEN": "secret-value"} + ] + + +def test_production_wiring_preserves_selected_deployment_type( + 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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + source_environment={"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"}, + ) + + service.execute(authored_run_single, source_root=tmp_path) + + assert launcher.exported_environments == [ + {"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"} + ] + + +def test_production_wiring_rejects_missing_secret_before_publishing( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + deployment = authored_run_single.deployments[0] + server = deployment.server.model_copy( + update={"environment": {"HF_TOKEN": SecretRef(type="secret", environment="MISSING_TOKEN")}} + ) + authored = authored_run_single.model_copy( + update={"deployments": [deployment.model_copy(update={"server": server})]} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + publisher = _Publisher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=publisher, # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + source_environment={}, + ) + + with pytest.raises(SlurmServiceError, match="MISSING_TOKEN") as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert publisher.initializations == [] + assert launcher.submissions == [] + + +def test_production_wiring_publishes_initial_state_before_releasing_submission( + 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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + result = service.execute(authored_run_single, source_root=tmp_path) + + writer = SlurmStateWriter(tmp_path, result.run_id) + attempts = writer.load_attempts("shard-00000") + assert result.state == "submitted" + assert writer.load_run().run_id == result.run_id + dependency_bytes = (tmp_path / "runs/run-wired/dependency-lock.json").read_bytes() + assert hashlib.sha256(dependency_bytes).hexdigest() == writer.load_resolved_plan().client.dependency_lock.sha256 + assert attempts[0].state is AttemptLifecycleState.SUBMITTED + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=42, array_task_id=0) + assert launcher.held_submissions == [True] + assert launcher.releases == [42] + + +def test_execute_reports_safe_submission_and_log_guidance( + 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) + progress: list[str] = [] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + progress=progress.append, + ) + + result = service.execute(authored_run_single, source_root=tmp_path) + + assert result.state == "submitted" + assert progress[0] == "Rendered batch script; invoking sbatch --parsable --hold (script via stdin)." + assert "data-designer slurm status run-wired" in progress[1] + assert "tail -F" in progress[2] + assert "slurm-attempt-0001-42_0.out" in progress[2] + assert "slurm-attempt-0001-42_0.err" in progress[2] + assert "watch shard 0 vLLM logs (including new executions)" in progress[3] + assert "shards/shard-00000/attempts/attempt-0001/logs" in progress[3] + assert "watch -n 2 -t" in progress[3] + assert "deployment-*.out" in progress[3] + assert "deployment-*.err" in progress[3] + assert "-exec tail -n 10 {} +" in progress[3] + + +def test_execute_distinguishes_shard_index_from_shard_id_in_log_guidance( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + sharded_run = authored_run_single.model_copy( + update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + ) + _register_images(tmp_path, sharded_run, single_node_plan) + progress: list[str] = [] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + progress=progress.append, + ) + + result = service.execute(sharded_run, source_root=tmp_path) + + assert result.shard_count == 2 + assert "_1 and shard-00001" in progress[4] + + +def test_execute_status_guidance_preserves_explicit_profile_selection( + 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) + catalog = profile_catalog.model_copy( + update={"clusters": {**profile_catalog.clusters, "primary": _profile(tmp_path, profile_catalog)}} + ) + profile_file = tmp_path / "profiles.json" + profile_file.write_text(catalog.model_dump_json()) + progress: list[str] = [] + service = create_slurm_run_service( + profile_file=profile_file, + cluster="primary", + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + progress=progress.append, + ) + + result = service.execute(authored_run_single, source_root=tmp_path) + + assert result.state == "submitted" + assert progress[1] == ( + f"Submitted. Check status: data-designer slurm status run-wired --profile-file {profile_file} --cluster primary" + ) + + +@pytest.mark.parametrize( + ("submission_error", "message"), + [ + ( + SlurmSubmissionError( + "sbatch failed with exit code 1: Invalid account or account/partition combination specified", + may_have_succeeded=False, + ), + "Slurm submission failed: sbatch failed with exit code 1: " + "Invalid account or account/partition combination specified", + ), + ( + SlurmSubmissionError("sbatch could not be executed: missing executable", may_have_succeeded=False), + "Slurm submission failed: sbatch could not be executed: missing executable", + ), + ( + SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True), + "Slurm submission outcome is uncertain", + ), + ], +) +def test_execute_reports_bounded_redacted_submission_diagnostic( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + submission_error: SlurmSubmissionError, + message: str, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(submission_error=submission_error), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match=message) as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert SlurmStateWriter(tmp_path, "run-wired").load_attempts("shard-00000") == () + + +def test_execute_redacts_submission_diagnostic( + 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) + error = SlurmSubmissionError( + "sbatch rejected account with token=secretvalue\n" + "extra detail " * 100, + may_have_succeeded=False, + ) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(submission_error=error), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path) + + message = str(caught.value) + assert message.startswith("Slurm submission failed: sbatch rejected account with token=") + assert "secretvalue" not in message + assert "\n" not in message + assert len(message) <= 512 + + +def test_status_reconciles_cancelled_scheduler_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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: datetime(2026, 9, 8, tzinfo=timezone.utc), + package_version="0.9.2", + ) + result = service.execute(authored_run_single, source_root=tmp_path) + service.cancel(result.run_id) + scheduler = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=scheduler, + state=SchedulerState.CANCELLED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=15), + ), + ) + + status = service.status(result.run_id) + + attempt = status.shards[0].attempts[0].attempt + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + + +def test_status_expires_unrequeued_preemption_and_cancel_skips_terminal_job( + 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) + launcher = _Launcher() + current_time = [datetime(2026, 9, 8, tzinfo=timezone.utc)] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: current_time[0], + package_version="0.9.2", + ) + result = service.execute(authored_run_single, source_root=tmp_path) + scheduler = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=scheduler, + state=SchedulerState.PREEMPTED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + + pending = service.status(result.run_id) + current_time[0] += timedelta(minutes=5, seconds=1) + failed = service.status(result.run_id) + cancellation = service.cancel(result.run_id) + + assert pending.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + assert failed.shards[0].attempts[0].attempt.state is AttemptLifecycleState.FAILED + assert cancellation.job_ids == () + assert launcher.cancellations == [] + + +def test_status_does_not_expire_requeue_window_from_another_attempt_clock( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + current_time = [datetime(2026, 9, 8, tzinfo=timezone.utc)] + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: current_time[0], + package_version="0.9.2", + ) + result = service.execute(authored, source_root=tmp_path) + preempted = SchedulerIdentity(array_job_id=42, array_task_id=0) + running = SchedulerIdentity(array_job_id=42, array_task_id=1) + launcher.queue_entries = (SlurmQueueEntry(job_identity=running, state=SchedulerState.RUNNING),) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=preempted, + state=SchedulerState.PREEMPTED, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ), + ) + service.status(result.run_id) + writer = SlurmStateWriter(tmp_path, result.run_id) + other_attempt = writer.load_attempt("shard-00001", "attempt-0001") + writer.update_attempt(other_attempt.model_copy(update={"updated_at": current_time[0] + timedelta(minutes=10)})) + + current_time[0] += timedelta(minutes=1) + still_requeueing = service.status(result.run_id) + current_time[0] += timedelta(minutes=1) + launcher.queue_entries = ( + SlurmQueueEntry(job_identity=preempted, state=SchedulerState.PENDING), + SlurmQueueEntry(job_identity=running, state=SchedulerState.RUNNING), + ) + requeued = service.status(result.run_id) + + assert still_requeueing.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + assert requeued.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING + + +@pytest.mark.parametrize( + ("source_environment", "expected_deployment_type"), + ( + pytest.param({"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"}, "slurm-nvidia-internal", id="selected"), + pytest.param({}, "slurm", id="unset"), + ), +) +def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, + source_environment: Mapping[str, str], + expected_deployment_type: str, +) -> None: + monkeypatch.setenv("NEMO_DEPLOYMENT_TYPE", "api") + authored = authored_run_single.model_copy( + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher(submission_job_ids=(42, 43)) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + source_environment=source_environment, + ) + service.execute(authored, source_root=tmp_path) + assert SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan().invocation.authored.resume == "if_possible" + failed = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=failed, + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + + preview = service.retry("42", shard_ids=("shard-00000",), resume="if_possible", dry_run=True) + assert not (tmp_path / "runs" / "run-wired" / "retries").exists() + assert not (tmp_path / "runs" / "run-wired" / "retry.lock").exists() + assert not ( + tmp_path / "runs" / "run-wired" / "shards" / "shard-00000" / "attempts" / "attempt-0001" / "scheduler.json" + ).exists() + assert len(launcher.submissions) == 1 + submitted = service.retry("42", shard_ids=("shard-00000",), resume="if_possible") + repeated = service.retry("run-wired", shard_ids=("shard-00000",), resume="if_possible") + + assert preview.state == "dry_run" + assert preview.job_id is None + assert preview.shard_ids == ("shard-00000",) + assert preview.effective_resume_mode == "never" + assert preview.batch_script is not None + assert "#SBATCH --array=0" in preview.batch_script + assert submitted.state == "submitted" + assert submitted.job_id == 43 + assert submitted.shard_ids == ("shard-00000",) + assert repeated == submitted + assert len(launcher.submissions) == 2 + assert launcher.exported_environments == [ + {"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": expected_deployment_type}, + {"NEMO_DEPLOYMENT_TYPE": expected_deployment_type}, + ] + + +def test_production_retry_preview_recovers_active_pinned_mode( + 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) + launcher = _Launcher(submission_job_ids=(42, 43)) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=42, array_task_id=0), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + submitted = service.retry("run-wired", resume="never") + + preview = service.retry("run-wired", resume="never", dry_run=True) + + assert preview.state == "dry_run" + assert preview.shard_ids == submitted.shard_ids + assert preview.attempt_ids == submitted.attempt_ids + assert preview.effective_resume_mode == "never" + assert len(launcher.submissions) == 2 + + +def test_production_retry_if_possible_reuses_populated_workspace( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={"invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"})} + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored, source_root=tmp_path) + plan = SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan() + resume_workspace = Path(plan.shards[0].resume_workspace.path) + resume_workspace.mkdir(parents=True) + (resume_workspace / "partial.parquet").touch() + failed = SchedulerIdentity(array_job_id=42, array_task_id=0) + launcher.accounting_entries = ( + SlurmAccountingEntry( + job_identity=failed, + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + + preview = service.retry("run-wired", resume="if_possible", dry_run=True) + + assert preview.effective_resume_mode == "always" + assert preview.shard_ids == ("shard-00000",) + + +def test_production_retry_if_possible_rejects_mixed_workspace_availability( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={ + "array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2}), + "invocation": authored_run_single.invocation.model_copy(update={"resume": "if_possible"}), + } + ) + _register_images(tmp_path, authored, single_node_plan) + launcher = _Launcher() + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: now, + package_version="0.9.2", + ) + service.execute(authored, source_root=tmp_path) + plan = SlurmStateWriter(tmp_path, "run-wired").load_resolved_plan() + resume_workspace = Path(plan.shards[0].resume_workspace.path) + resume_workspace.mkdir(parents=True) + (resume_workspace / "partial.parquet").touch() + launcher.accounting_entries = tuple( + SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=42, array_task_id=task_id), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ) + for task_id in range(2) + ) + + with pytest.raises(SlurmServiceError, match="mixes resumable and fresh shards") as caught: + service.retry("run-wired", resume="if_possible", dry_run=True) + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_production_retry_rejects_no_retryable_shards( + 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) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="no retryable shards") as caught: + service.retry("run-wired", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert caught.value.operation is SlurmServiceOperation.RETRY_RUN + + +def test_production_retry_rejects_ambiguous_scheduler_job_id( + 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) + run_ids = iter(("run-one", "run-two")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: next(run_ids), + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="multiple managed runs") as caught: + service.retry("42", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + + +def test_production_retry_rejects_numeric_run_and_job_ambiguity( + 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) + run_ids = iter(("run-one", "42")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(submission_job_ids=(42, 43)), # type: ignore[arg-type] + run_id_factory=lambda: next(run_ids), + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError, match="numeric reference") as caught: + service.retry("42", resume="never") + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + + +def test_collection_rejects_unmanaged_input_and_destination( + 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) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + + with pytest.raises(SlurmServiceError) as unmanaged: + service.collect(tmp_path / "run-wired", destination=tmp_path / "collected") + with pytest.raises(SlurmServiceError) as destination: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path.parent / "collected") + with pytest.raises(SlurmServiceError) as managed_state: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path / "runs/run-wired/collections") + with pytest.raises(SlurmServiceError) as partitions: + service.collect(tmp_path / "runs/run-wired", destination=tmp_path / "collected", num_partitions=9) + + assert unmanaged.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert destination.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert managed_state.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert partitions.value.code is SlurmServiceErrorCode.INVALID_REQUEST + + +def test_collection_returns_active_and_completed_jobs( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + service.execute(authored_run_single, source_root=tmp_path) + statuses = iter( + ( + SimpleNamespace(collection_id="collection-0001", state=CollectionState.RUNNING, scheduler=43), + SimpleNamespace(collection_id="collection-0001", state=CollectionState.SUCCEEDED, scheduler=43), + ) + ) + + class Coordinator: + def submit(self, *, destination: Path, num_partitions: int | None, submitted_at: datetime) -> SimpleNamespace: + del destination, num_partitions, submitted_at + return next(statuses) + + monkeypatch.setattr(retry_collection_module, "SlurmCollectionCoordinator", lambda *_: Coordinator()) + monkeypatch.setattr( + retry_collection_module, + "CollectionStorage", + lambda *_: SimpleNamespace( + read_plan=lambda _: SimpleNamespace(num_partitions=3, host_destination=output_path.as_posix()) + ), + ) + output_path = tmp_path / "collected" + + active = service.collect("run-wired", destination=output_path, num_partitions=3) + completed = service.collect("run-wired") + + assert active.state is CollectionState.RUNNING + assert completed.state is CollectionState.SUCCEEDED + assert active.job_id == completed.job_id == 43 + assert active.num_partitions == completed.num_partitions == 3 + assert active.output_path == completed.output_path == output_path.as_posix() + + +def test_auto_gpu_resolution_rejects_mixed_node_shapes( + 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_catalog.clusters["lab"].model_copy(update={"workspace_root": tmp_path.as_posix()}) + service = create_slurm_run_service( + profile=profile, + launcher=_Launcher((4, 8)), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path, dry_run=True) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + + +def test_publisher_value_error_is_not_misattributed_to_run_preparation( + 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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(initialization_error=ValueError("backend bug")), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert str(caught.value) == "execute run failed" + assert launcher.submissions == [] + + +def test_recording_conflict_cancels_the_accepted_job( + 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) + launcher = _Launcher() + publisher = _Publisher(submission_error=StateConflictError("conflict")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=publisher, # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.cancellations == [42] + 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, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + authored = authored_run_single.model_copy( + update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + ) + _register_images(tmp_path, authored, single_node_plan) + original_create = SlurmStateWriter.create_attempt + created = 0 + + def fail_second_attempt(writer: SlurmStateWriter, attempt: AttemptManifest) -> AttemptManifest: + nonlocal created + created += 1 + if created == 2: + raise StateConflictError("injected partial submission failure") + return original_create(writer, attempt) + + monkeypatch.setattr(SlurmStateWriter, "create_attempt", fail_second_attempt) + launcher = _Launcher() + failed_at = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: failed_at, + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.cancellations == [42] + writer = SlurmStateWriter(tmp_path, "run-wired") + attempt = writer.load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + with pytest.raises(StateNotFoundError): + writer.load_attempt("shard-00001", "attempt-0001") + + +def test_release_failure_cancels_the_held_job( + 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) + launcher = _Launcher(release_error=SlurmLauncherError("release failed")) + failed_at = datetime(2026, 9, 8, tzinfo=timezone.utc) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + clock=lambda: failed_at, + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="could not be released and was cancelled") as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert launcher.releases == [42] + assert launcher.cancellations == [42] + attempt = SlurmStateWriter(tmp_path, "run-wired").load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.FAILED + assert attempt.terminal_classification is AttemptTerminalClassification.CANCELLED + assert attempt.updated_at == failed_at + + +def test_failed_compensating_cancel_reports_accepted_job_id( + 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) + launcher = _Launcher(cancel_error=RuntimeError("cancel failed")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(submission_error=StateConflictError("conflict")), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="Slurm job 42") as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + + +def test_recording_interrupt_is_preserved_when_compensating_cancel_fails( + 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) + launcher = _Launcher(cancel_error=RuntimeError("cancel failed")) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=_Publisher(submission_error=KeyboardInterrupt()), # type: ignore[arg-type] + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(KeyboardInterrupt): + service.execute(authored_run_single, source_root=tmp_path) + + assert launcher.cancellations == [42] + + +def test_production_status_and_cancel_use_only_persisted_m2_records( + 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) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + plan = service.plan(authored_run_single) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + plan_reference = ArtifactReference( + path=(tmp_path / "runs/run-wired/resolved-plan.json").as_posix(), + sha256=plan.compute_sha256(), + ) + run = RunManifest( + schema_version=1, + run_id=plan.run_id, + created_at=now, + authored_config=plan.authored_config, + resolved_plan=plan_reference, + shard_count=len(plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shard.shard_id, + shard_index=shard.shard_index, + record_range=shard.record_range, + input_partition=shard.input_partition, + resume_workspace=shard.resume_workspace, + created_at=now, + ) + for shard in plan.shards + ) + writer = SlurmStateWriter(tmp_path, plan.run_id) + writer.initialize_run(authored_run_single, plan, run, shards) + created = AttemptManifest( + schema_version=1, + run_id=plan.run_id, + shard_id=shards[0].shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=plan_reference, + state=AttemptLifecycleState.CREATED, + created_at=now, + updated_at=now, + ) + writer.create_attempt(created) + assert service.cancel(plan.run_id).job_ids == () + assert launcher.cancellations == [] + writer.update_attempt( + created.model_copy( + update={ + "state": AttemptLifecycleState.SUBMITTED, + "scheduler": SchedulerIdentity(array_job_id=42, array_task_id=0), + } + ) + ) + + status = service.status(plan.run_id) + cancellation = service.cancel(plan.run_id) + persisted_after_cancel = service.status(plan.run_id) + + assert status.run == run + assert status.shards[0].attempts[0].attempt.state is AttemptLifecycleState.SUBMITTED + assert cancellation.job_ids == (42,) + assert persisted_after_cancel.shards[0].attempts[0].attempt.state is AttemptLifecycleState.SUBMITTED + assert launcher.cancellations == [42] + + +def test_production_status_maps_unknown_run_to_not_found( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, +) -> None: + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), # type: ignore[arg-type] + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.status("run-missing") + + assert caught.value.code is SlurmServiceErrorCode.NOT_FOUND + + +def test_production_image_registry_operations( + 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) + service = create_slurm_image_service(profile=_profile(tmp_path, profile_catalog)) + + images = service.list() + selected = service.get(images[0].name) + removed = service.remove(images[0].name) + + assert selected == removed == images[0] + assert service.list() == () + + +def test_production_image_remove_refuses_to_delete_external_sqsh( + 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) + service = create_slurm_image_service(profile=_profile(tmp_path, profile_catalog)) + image = service.list()[0] + + with pytest.raises(SlurmServiceError) as caught: + service.remove(image.name, delete_artifact=True) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert "package-owned" in str(caught.value) + assert service.get(image.name) == image diff --git a/packages/data-designer-slurm/tests/serving/conftest.py b/packages/data-designer-slurm/tests/serving/conftest.py new file mode 100644 index 000000000..589895c44 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/conftest.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import pytest + +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import resolve_vllm_server + + +@pytest.fixture +def multi_replica_plan(multi_node_plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: + payload = multi_node_plan.model_dump(mode="json") + deployment = payload["deployments"][0] + deployment["authored"]["server"].update( + lead_boot_standoff="30s", + rank_launch_stagger="2s", + ) + deployment["authored"]["topology"]["tensor_parallel"] = 4 + deployment["topology"].update( + tensor_parallel=4, + replicas_per_node_group=2, + replica_count=2, + gpus_per_replica=8, + ) + deployment["ports"] = [ + { + "name": "deployment-00000-http-00000", + "role": "http", + "node_index": 0, + "port": 18000, + }, + { + "name": "deployment-00000-http-00001", + "role": "http", + "node_index": 0, + "port": 18001, + }, + { + "name": "deployment-00000-rendezvous-00000", + "role": "rendezvous", + "node_index": 0, + "port": 19000, + }, + { + "name": "deployment-00000-rendezvous-00001", + "role": "rendezvous", + "node_index": 0, + "port": 19001, + }, + ] + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +@pytest.fixture +def resolved_multi_replica_server(multi_replica_plan: ResolvedSlurmRunPlan) -> ResolvedVllmServerDeployment: + deployment = multi_replica_plan.deployments[0] + return _resolve(multi_replica_plan, deployment.deployment_id) + + +@pytest.fixture +def resolved_single_node_server(single_node_plan: ResolvedSlurmRunPlan) -> ResolvedVllmServerDeployment: + deployment = single_node_plan.deployments[0] + return _resolve(single_node_plan, deployment.deployment_id) + + +def _resolve(plan: ResolvedSlurmRunPlan, deployment_id: str) -> ResolvedVllmServerDeployment: + return resolve_vllm_server(plan, deployment_id) diff --git a/packages/data-designer-slurm/tests/serving/golden/multi_node.json b/packages/data-designer-slurm/tests/serving/golden/multi_node.json new file mode 100644 index 000000000..bf42f0050 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/golden/multi_node.json @@ -0,0 +1,231 @@ +{ + "backend_endpoints": [ + { + "backend_id": "deployment-00000-backend-00000", + "deployment_replica_index": 0, + "node_group_index": 0, + "node_index": 0, + "port": 18000, + "replica_index_in_node_group": 0, + "served_model_name": "example/generator" + }, + { + "backend_id": "deployment-00000-backend-00001", + "deployment_replica_index": 1, + "node_group_index": 0, + "node_index": 0, + "port": 18001, + "replica_index_in_node_group": 1, + "served_model_name": "example/generator" + } + ], + "deployment_id": "deployment-00000", + "executable_path": "/usr/local/bin/vllm", + "failure_policy": "coordinated", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "launch_policy": { + "distributed_init_timeout_seconds": 600, + "enable_expert_parallel": false, + "environment": { + "HF_TOKEN": { + "environment": "HF_TOKEN", + "type": "secret" + } + }, + "extra_args": [ + "--max-model-len", + "32768" + ], + "lead_boot_standoff_seconds": 30, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "rank_launch_stagger_seconds": 2, + "readiness_path": "/health", + "startup_timeout_seconds": 900 + }, + "logical_endpoint": { + "backend_ids": [ + "deployment-00000-backend-00000", + "deployment-00000-backend-00001" + ], + "endpoint_id": "deployment-00000-logical-endpoint", + "load_balancing": "least_connections", + "model_alias": "generator", + "node_index": 0, + "port": 17000, + "preserve_final_overload_response": true, + "require_all_backends_ready": true, + "retry_status_codes": [ + 429 + ], + "served_model_name": "example/generator" + }, + "model": "example/generator", + "model_alias": "generator", + "node_indices": [ + 0, + 1 + ], + "processes": [ + { + "deployment_id": "deployment-00000", + "deployment_replica_index": 0, + "gpu_indices": [ + 0, + 1, + 2, + 3 + ], + "http_port": 18000, + "launch_delay_seconds": 0, + "node_group_index": 0, + "node_index": 0, + "pipeline_parallel": 2, + "pipeline_rank": 0, + "process_id": "deployment-00000-replica-00000-rank-00000", + "rendezvous": { + "master_node_index": 0, + "node_group_index": 0, + "port": 19000, + "replica_index_in_node_group": 0, + "timeout_seconds": 600 + }, + "replica_index_in_node_group": 0, + "role": "api_server", + "tensor_parallel": 4 + }, + { + "deployment_id": "deployment-00000", + "deployment_replica_index": 0, + "gpu_indices": [ + 0, + 1, + 2, + 3 + ], + "http_port": null, + "launch_delay_seconds": 0, + "node_group_index": 0, + "node_index": 1, + "pipeline_parallel": 2, + "pipeline_rank": 1, + "process_id": "deployment-00000-replica-00000-rank-00001", + "rendezvous": { + "master_node_index": 0, + "node_group_index": 0, + "port": 19000, + "replica_index_in_node_group": 0, + "timeout_seconds": 600 + }, + "replica_index_in_node_group": 0, + "role": "follower", + "tensor_parallel": 4 + }, + { + "deployment_id": "deployment-00000", + "deployment_replica_index": 1, + "gpu_indices": [ + 4, + 5, + 6, + 7 + ], + "http_port": 18001, + "launch_delay_seconds": 32, + "node_group_index": 0, + "node_index": 0, + "pipeline_parallel": 2, + "pipeline_rank": 0, + "process_id": "deployment-00000-replica-00001-rank-00000", + "rendezvous": { + "master_node_index": 0, + "node_group_index": 0, + "port": 19001, + "replica_index_in_node_group": 1, + "timeout_seconds": 600 + }, + "replica_index_in_node_group": 1, + "role": "api_server", + "tensor_parallel": 4 + }, + { + "deployment_id": "deployment-00000", + "deployment_replica_index": 1, + "gpu_indices": [ + 4, + 5, + 6, + 7 + ], + "http_port": null, + "launch_delay_seconds": 32, + "node_group_index": 0, + "node_index": 1, + "pipeline_parallel": 2, + "pipeline_rank": 1, + "process_id": "deployment-00000-replica-00001-rank-00001", + "rendezvous": { + "master_node_index": 0, + "node_group_index": 0, + "port": 19001, + "replica_index_in_node_group": 1, + "timeout_seconds": 600 + }, + "replica_index_in_node_group": 1, + "role": "follower", + "tensor_parallel": 4 + } + ], + "readiness_probes": [ + { + "backend_id": "deployment-00000-backend-00000", + "deadline_seconds": 900, + "expected_status_code": 200, + "node_index": 0, + "path": "/health", + "port": 18000, + "probe_id": "deployment-00000-readiness-00000" + }, + { + "backend_id": "deployment-00000-backend-00001", + "deadline_seconds": 900, + "expected_status_code": 200, + "node_index": 0, + "path": "/health", + "port": 18001, + "probe_id": "deployment-00000-readiness-00001" + } + ], + "served_model_name": "example/generator", + "server_type": "vllm", + "topology": { + "gpus_per_replica": 8, + "node_group_count": 1, + "nodes_per_replica": 2, + "pipeline_parallel": 2, + "replica_count": 2, + "replicas_per_node_group": 2, + "tensor_parallel": 4 + } +} diff --git a/packages/data-designer-slurm/tests/serving/golden/single_node.json b/packages/data-designer-slurm/tests/serving/golden/single_node.json new file mode 100644 index 000000000..d4ae8e57b --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/golden/single_node.json @@ -0,0 +1,120 @@ +{ + "backend_endpoints": [ + { + "backend_id": "deployment-00000-backend-00000", + "deployment_replica_index": 0, + "node_group_index": 0, + "node_index": 0, + "port": 18000, + "replica_index_in_node_group": 0, + "served_model_name": "example/generator" + } + ], + "deployment_id": "deployment-00000", + "executable_path": "/usr/local/bin/vllm", + "failure_policy": "coordinated", + "gpus_per_node": 8, + "image": { + "authored_ref": { + "name": "vllm-0-21", + "path": null + }, + "inspection": { + "inspection": { + "executable_path": "/usr/local/bin/vllm", + "kind": "serving", + "runtime_version": "0.21.0", + "server_type": "vllm" + }, + "inspector_version": "inspector-1", + "schema_version": 1, + "sqsh_sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "path": "/images/vllm-0-21.sqsh", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + }, + "launch_policy": { + "distributed_init_timeout_seconds": 600, + "enable_expert_parallel": false, + "environment": {}, + "extra_args": [], + "lead_boot_standoff_seconds": 60, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "rank_launch_stagger_seconds": 5, + "readiness_path": "/health", + "startup_timeout_seconds": 900 + }, + "logical_endpoint": { + "backend_ids": [ + "deployment-00000-backend-00000" + ], + "endpoint_id": "deployment-00000-logical-endpoint", + "load_balancing": "least_connections", + "model_alias": "generator", + "node_index": 0, + "port": 17000, + "preserve_final_overload_response": true, + "require_all_backends_ready": true, + "retry_status_codes": [ + 429 + ], + "served_model_name": "example/generator" + }, + "model": "example/generator", + "model_alias": "generator", + "node_indices": [ + 0 + ], + "processes": [ + { + "deployment_id": "deployment-00000", + "deployment_replica_index": 0, + "gpu_indices": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "http_port": 18000, + "launch_delay_seconds": 0, + "node_group_index": 0, + "node_index": 0, + "pipeline_parallel": 1, + "pipeline_rank": 0, + "process_id": "deployment-00000-replica-00000-rank-00000", + "rendezvous": null, + "replica_index_in_node_group": 0, + "role": "api_server", + "tensor_parallel": 8 + } + ], + "readiness_probes": [ + { + "backend_id": "deployment-00000-backend-00000", + "deadline_seconds": 900, + "expected_status_code": 200, + "node_index": 0, + "path": "/health", + "port": 18000, + "probe_id": "deployment-00000-readiness-00000" + } + ], + "served_model_name": "example/generator", + "server_type": "vllm", + "topology": { + "gpus_per_replica": 8, + "node_group_count": 1, + "nodes_per_replica": 1, + "pipeline_parallel": 1, + "replica_count": 1, + "replicas_per_node_group": 1, + "tensor_parallel": 8 + } +} diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py new file mode 100644 index 000000000..26532ca01 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import ast +import json +from pathlib import Path + +import pytest + +import data_designer.slurm.serving.resolver as resolver_module +from data_designer.slurm.config import ImageInspectionRecord, QueueBackpressureConfig +from data_designer.slurm.contracts import compute_serialized_json_sha256, pretty_json +from data_designer.slurm.planning import ResolvedDeployment, ResolvedSlurmRunPlan +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.resolver import ( + VllmServerResolutionError, + resolve_vllm_server, +) +from data_designer.slurm.serving.vllm import VllmProcessRole + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" + + +def test_single_node_resolution_matches_golden(single_node_plan: ResolvedSlurmRunPlan) -> None: + placement = single_node_plan.deployments[0] + + first = _resolve(single_node_plan, placement.deployment_id) + second = _resolve(single_node_plan, placement.deployment_id) + + assert first == second + assert pretty_json(first.model_dump(mode="json")) == (GOLDEN_DIRECTORY / "single_node.json").read_text() + assert first.processes[0].role is VllmProcessRole.API_SERVER + assert first.processes[0].launch_delay_seconds == 0 + assert first.launch_policy.lead_boot_standoff_seconds == 60 + assert first.launch_policy.rank_launch_stagger_seconds == 5 + assert first.launch_policy.queue_backpressure == QueueBackpressureConfig() + assert first.logical_endpoint.load_balancing == "least_connections" + assert first.failure_policy == "coordinated" + + +def test_multi_node_replica_resolution_matches_golden(multi_replica_plan: ResolvedSlurmRunPlan) -> None: + placement = multi_replica_plan.deployments[0] + + resolved = _resolve(multi_replica_plan, placement.deployment_id) + + assert pretty_json(resolved.model_dump(mode="json")) == (GOLDEN_DIRECTORY / "multi_node.json").read_text() + assert [process.pipeline_rank for process in resolved.processes] == [0, 1, 0, 1] + assert [process.gpu_indices for process in resolved.processes] == [ + (0, 1, 2, 3), + (0, 1, 2, 3), + (4, 5, 6, 7), + (4, 5, 6, 7), + ] + assert [process.launch_delay_seconds for process in resolved.processes] == [0, 0, 32, 32] + assert [process.role for process in resolved.processes] == [ + VllmProcessRole.API_SERVER, + VllmProcessRole.FOLLOWER, + VllmProcessRole.API_SERVER, + VllmProcessRole.FOLLOWER, + ] + assert resolved.logical_endpoint.load_balancing == "least_connections" + + +def test_resolution_preserves_io_processor_plugin(single_node_plan: ResolvedSlurmRunPlan) -> None: + placement = single_node_plan.deployments[0] + server = placement.authored.server.model_copy(update={"io_processor_plugin": "deberta_gliner_io"}) + authored = placement.authored.model_copy(update={"server": server}) + plan = _plan_with_placement(single_node_plan, placement.model_copy(update={"authored": authored})) + + resolved = _resolve(plan, placement.deployment_id) + + assert resolved.launch_policy.io_processor_plugin == "deberta_gliner_io" + assert ResolvedVllmServerDeployment.model_validate_json(resolved.model_dump_json()) == resolved + + +def test_multiple_two_node_groups_preserve_global_replica_order_and_delays( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan = _multi_group_plan(multi_node_plan) + placement = plan.deployments[0] + + resolved = _resolve(plan, placement.deployment_id) + + assert resolved.topology.node_group_count == 2 + assert resolved.topology.replica_count == 4 + assert [endpoint.node_index for endpoint in resolved.backend_endpoints] == [0, 0, 2, 2] + assert [process.node_index for process in resolved.processes] == [0, 1, 0, 1, 2, 3, 2, 3] + assert [process.launch_delay_seconds for process in resolved.processes] == [0, 0, 32, 32, 34, 34, 36, 36] + + +def test_two_deployments_keep_images_and_endpoint_identities_isolated( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + resolved = tuple(_resolve(multi_node_plan, placement.deployment_id) for placement in multi_node_plan.deployments) + + assert resolved[0].image.sha256 != resolved[1].image.sha256 + assert resolved[0].model_alias != resolved[1].model_alias + assert resolved[0].model != resolved[1].model + assert resolved[0].served_model_name != resolved[1].served_model_name + assert resolved[0].image.inspection_facts.runtime_version != resolved[1].image.inspection_facts.runtime_version + node_indices = [node_index for deployment in resolved for node_index in deployment.node_indices] + assert len(node_indices) == len(set(node_indices)) + process_ids = [process.process_id for deployment in resolved for process in deployment.processes] + assert len(process_ids) == len(set(process_ids)) + backend_addresses = [ + (backend.node_index, backend.port) for deployment in resolved for backend in deployment.backend_endpoints + ] + assert len(backend_addresses) == len(set(backend_addresses)) + logical_addresses = [ + (deployment.logical_endpoint.node_index, deployment.logical_endpoint.port) for deployment in resolved + ] + assert len(logical_addresses) == len(set(logical_addresses)) + rendezvous_addresses = [ + (process.rendezvous.master_node_index, process.rendezvous.port) + for deployment in resolved + for process in deployment.processes + if process.pipeline_rank == 0 and process.rendezvous is not None + ] + network_addresses = backend_addresses + logical_addresses + rendezvous_addresses + assert len(network_addresses) == len(set(network_addresses)) + + +def test_resolution_preserves_inspected_runtime_version_as_provenance(single_node_plan: ResolvedSlurmRunPlan) -> None: + payload = single_node_plan.deployments[0].model_dump(mode="json") + payload["image"]["inspection"]["inspection"]["runtime_version"] = "0.22.0+vendor.1" + placement = ResolvedDeployment.model_validate_json(json.dumps(payload)) + plan = _plan_with_placement(single_node_plan, placement) + + resolved = _resolve(plan, placement.deployment_id) + + assert resolved.image.inspection_facts.runtime_version == "0.22.0+vendor.1" + + +@pytest.mark.parametrize( + "runtime_version", + ("vendor-vllm-build", "0.20.9", "0.22.0rc1", "0.23.0"), +) +def test_resolution_rejects_unsupported_runtime_versions( + single_node_plan: ResolvedSlurmRunPlan, + runtime_version: str, +) -> None: + payload = single_node_plan.deployments[0].model_dump(mode="json") + payload["image"]["inspection"]["inspection"]["runtime_version"] = runtime_version + placement = ResolvedDeployment.model_validate_json(json.dumps(payload)) + plan = _plan_with_placement(single_node_plan, placement) + + with pytest.raises(VllmServerResolutionError, match="unsupported vLLM runtime version"): + _resolve(plan, placement.deployment_id) + + +def test_resolution_rejects_image_inspection_mismatch( + single_node_plan: ResolvedSlurmRunPlan, + client_image_inspection: ImageInspectionRecord, +) -> None: + image = single_node_plan.deployments[0].image.model_copy(update={"inspection": client_image_inspection}) + placement = single_node_plan.deployments[0].model_copy(update={"image": image}) + invalid_plan = single_node_plan.model_copy(update={"deployments": (placement,)}) + + with pytest.raises(VllmServerResolutionError, match="serving image"): + _resolve(invalid_plan, placement.deployment_id) + + +def test_resolution_rejects_multi_node_expert_parallel(multi_node_plan: ResolvedSlurmRunPlan) -> None: + placement = multi_node_plan.deployments[0] + server = placement.authored.server.model_copy(update={"enable_expert_parallel": True}) + authored = placement.authored.model_copy(update={"server": server}) + invalid_placement = placement.model_copy(update={"authored": authored}) + invalid_plan = multi_node_plan.model_copy( + update={"deployments": (invalid_placement, *multi_node_plan.deployments[1:])} + ) + + with pytest.raises(VllmServerResolutionError, match="multi-node expert parallel"): + _resolve(invalid_plan, invalid_placement.deployment_id) + + +def test_resolution_allows_independent_single_node_expert_parallel_replicas( + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + placement = _independent_replica_placement(multi_node_plan) + plan = _plan_with_placement(multi_node_plan, placement) + + resolved = _resolve(plan, placement.deployment_id) + + assert resolved.launch_policy.enable_expert_parallel + assert resolved.topology.pipeline_parallel == 1 + assert [process.node_index for process in resolved.processes] == [0, 1] + assert all(process.rendezvous is None for process in resolved.processes) + + +def test_serving_package_has_no_scheduler_shell_or_runtime_dependency() -> None: + imports: set[str] = set() + for source_path in Path(resolver_module.__file__).parent.glob("*.py"): + tree = ast.parse(source_path.read_text()) + imports.update(alias.name for node in ast.walk(tree) if isinstance(node, ast.Import) for alias in node.names) + imports.update( + node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) and node.module is not None + ) + + assert "subprocess" not in imports + assert not any( + module.startswith( + ( + "data_designer.slurm.launcher", + "data_designer.slurm.runtime", + "data_designer.slurm.state", + ) + ) + for module in imports + ) + + +def test_resolution_requires_matching_logical_endpoint(single_node_plan: ResolvedSlurmRunPlan) -> None: + deployment = single_node_plan.deployments[0] + unrelated_port = single_node_plan.client.ports[0].model_copy(update={"name": "other-logical-endpoint"}) + client = single_node_plan.client.model_copy(update={"ports": (unrelated_port,)}) + invalid_plan = single_node_plan.model_copy(update={"client": client}) + + with pytest.raises(VllmServerResolutionError, match="exactly one logical endpoint"): + resolve_vllm_server(invalid_plan, deployment.deployment_id) + + +def test_resolution_rejects_cross_input_network_collision(single_node_plan: ResolvedSlurmRunPlan) -> None: + deployment = single_node_plan.deployments[0] + deployment_port = deployment.ports[0] + colliding_port = single_node_plan.client.ports[0].model_copy( + update={"node_index": deployment_port.node_index, "port": deployment_port.port} + ) + client = single_node_plan.client.model_copy(update={"ports": (colliding_port,)}) + invalid_plan = single_node_plan.model_copy(update={"client": client}) + + with pytest.raises(VllmServerResolutionError, match="must not collide"): + resolve_vllm_server(invalid_plan, deployment.deployment_id) + + +def test_resolution_normalizes_inconsistent_planner_output(single_node_plan: ResolvedSlurmRunPlan) -> None: + deployment = single_node_plan.deployments[0] + inconsistent_topology = deployment.topology.model_copy(update={"replica_count": 2}) + inconsistent_deployment = deployment.model_copy(update={"topology": inconsistent_topology}) + invalid_plan = single_node_plan.model_copy(update={"deployments": (inconsistent_deployment,)}) + + with pytest.raises(VllmServerResolutionError, match="inconsistent vLLM server specification"): + resolve_vllm_server(invalid_plan, deployment.deployment_id) + + +def test_resolution_normalizes_intermediate_record_validation(single_node_plan: ResolvedSlurmRunPlan) -> None: + deployment = single_node_plan.deployments[0] + long_deployment_id = "d" * 128 + renamed_deployment = deployment.model_copy(update={"deployment_id": long_deployment_id}) + endpoint = single_node_plan.client.ports[0].model_copy(update={"name": f"{long_deployment_id}-logical-endpoint"}) + client = single_node_plan.client.model_copy(update={"ports": (endpoint,)}) + invalid_plan = single_node_plan.model_copy(update={"client": client, "deployments": (renamed_deployment,)}) + + with pytest.raises(VllmServerResolutionError, match="inconsistent vLLM server specification"): + resolve_vllm_server(invalid_plan, long_deployment_id) + + +def test_resolution_derives_placement_and_endpoint_from_one_plan( + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = single_node_plan.model_dump(mode="json") + payload["deployments"][0]["ports"][0]["port"] = 18100 + payload["client"]["ports"][0]["port"] = 17100 + alternate_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + resolved = _resolve(alternate_plan, alternate_plan.deployments[0].deployment_id) + + assert resolved.backend_endpoints[0].port == 18100 + assert resolved.logical_endpoint.port == 17100 + + +def _resolve( + plan: ResolvedSlurmRunPlan, + deployment_id: str, +) -> ResolvedVllmServerDeployment: + return resolve_vllm_server(plan, deployment_id) + + +def _plan_with_placement( + plan: ResolvedSlurmRunPlan, + placement: ResolvedDeployment, +) -> ResolvedSlurmRunPlan: + payload = plan.model_dump(mode="json") + matching_indices = [ + index for index, candidate in enumerate(plan.deployments) if candidate.deployment_id == placement.deployment_id + ] + assert len(matching_indices) == 1 + payload["deployments"][matching_indices[0]] = placement.model_dump(mode="json") + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _multi_group_plan(plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: + payload = plan.model_dump(mode="json") + payload["deployments"] = payload["deployments"][:1] + payload["client"]["ports"] = payload["client"]["ports"][:1] + payload["builder"]["inline"]["data_designer"]["model_configs"] = payload["builder"]["inline"]["data_designer"][ + "model_configs" + ][:1] + payload["builder"]["model_aliases"] = ["generator"] + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) + deployment = payload["deployments"][0] + deployment["authored"]["resources"]["nodes"] = 4 + deployment["authored"]["topology"]["tensor_parallel"] = 4 + deployment["authored"]["server"].update( + lead_boot_standoff="30s", + rank_launch_stagger="2s", + ) + deployment["node_indices"] = [0, 1, 2, 3] + deployment["topology"].update( + tensor_parallel=4, + node_group_count=2, + replicas_per_node_group=2, + replica_count=4, + gpus_per_replica=8, + ) + deployment["ports"] = [ + { + "name": f"deployment-00000-http-{deployment_replica_index:05d}", + "role": "http", + "node_index": 0 if deployment_replica_index < 2 else 2, + "port": 18000 + deployment_replica_index % 2, + } + for deployment_replica_index in range(4) + ] + [ + { + "name": f"deployment-00000-rendezvous-{deployment_replica_index:05d}", + "role": "rendezvous", + "node_index": 0 if deployment_replica_index < 2 else 2, + "port": 19000 + deployment_replica_index % 2, + } + for deployment_replica_index in range(4) + ] + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _independent_replica_placement(plan: ResolvedSlurmRunPlan) -> ResolvedDeployment: + payload = plan.deployments[0].model_dump(mode="json") + payload["authored"]["server"]["enable_expert_parallel"] = True + payload["authored"]["topology"]["nodes_per_replica"] = 1 + payload["topology"].update( + nodes_per_replica=1, + pipeline_parallel=1, + node_group_count=2, + replica_count=2, + gpus_per_replica=8, + ) + payload["ports"] = [ + { + "name": "deployment-00000-http-00000", + "role": "http", + "node_index": 0, + "port": 18000, + }, + { + "name": "deployment-00000-http-00001", + "role": "http", + "node_index": 1, + "port": 18001, + }, + ] + return ResolvedDeployment.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/serving/test_serving_records.py b/packages/data-designer-slurm/tests/serving/test_serving_records.py new file mode 100644 index 000000000..3340709d7 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/test_serving_records.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import ImageInspectionRecord +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.endpoints import ResolvedLogicalEndpoint +from data_designer.slurm.serving.vllm import ResolvedVllmLaunchPolicy, ResolvedVllmProcess + + +@pytest.mark.parametrize("plugin", ["", "../plugin", "plugin module", "--plugin", "plugin;exit", "x" * 129]) +def test_resolved_launch_policy_rejects_invalid_io_processor_plugin( + resolved_multi_replica_server: ResolvedVllmServerDeployment, plugin: str +) -> None: + payload = resolved_multi_replica_server.launch_policy.model_dump(mode="json") + payload["io_processor_plugin"] = plugin + + with pytest.raises(ValidationError, match="io_processor_plugin"): + ResolvedVllmLaunchPolicy.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("gpu_order", "sorted and unique"), + ("gpu_count", "equal tensor parallelism"), + ("pipeline_rank", "below pipeline parallelism"), + ("head_role", "rank zero"), + ("head_port", "require an HTTP port"), + ("follower_role", "must be followers"), + ("follower_port", "must not publish"), + ("missing_rendezvous", "require rendezvous"), + ("wrong_rendezvous_lane", "rendezvous identity"), + ], +) +def test_process_specs_reject_one_field_topology_drift( + resolved_multi_replica_server: ResolvedVllmServerDeployment, + mutation: str, + message: str, +) -> None: + process_index = ( + 1 if mutation.startswith("follower") or mutation in {"missing_rendezvous", "wrong_rendezvous_lane"} else 0 + ) + payload = resolved_multi_replica_server.processes[process_index].model_dump(mode="json") + if mutation == "gpu_order": + payload["gpu_indices"] = list(reversed(payload["gpu_indices"])) + elif mutation == "gpu_count": + payload["gpu_indices"].pop() + elif mutation == "pipeline_rank": + payload["pipeline_rank"] = payload["pipeline_parallel"] + elif mutation == "head_role": + payload["role"] = "follower" + elif mutation == "head_port": + payload["http_port"] = None + elif mutation == "follower_role": + payload["role"] = "api_server" + elif mutation == "follower_port": + payload["http_port"] = 18000 + elif mutation == "missing_rendezvous": + payload["rendezvous"] = None + else: + payload["rendezvous"]["replica_index_in_node_group"] = 1 + + with pytest.raises(ValidationError, match=message): + ResolvedVllmProcess.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("image_kind", "server image inspection"), + ("executable", "executable path"), + ("nodes", "sorted and unique"), + ("node_group_divisibility", "must divide the deployment node count"), + ("gpu_divisibility", "must divide resolved GPUs per node"), + ("topology", "node and GPU resources"), + ("expert_parallel", "multi-node expert parallel"), + ("logical_alias", "model alias"), + ("logical_served_name", "served model name"), + ("logical_id", "endpoint ID"), + ("deployment_replica_index", "replica identities"), + ("backend_id", "endpoint IDs"), + ("logical_backends", "resolved backend order"), + ("process_ids", "process IDs"), + ("process_order", "ordered replica and pipeline identities"), + ("readiness_count", "readiness probes"), + ("backend_placement", "replica placement"), + ("readiness_target", "readiness probes"), + ("pipeline_ranks", "ordered replica and pipeline identities"), + ("process_topology", "topology and launch policy"), + ("head_endpoint", "head process"), + ("rendezvous_consistency", "share one rendezvous"), + ("rendezvous_master", "group head and distributed timeout"), + ("rendezvous_timeout", "group head and distributed timeout"), + ("network_address", "network addresses must be unique"), + ("launch_policy_owned_arg", "owned by the compiler or runtime"), + ("launch_policy_literal_secret", "secret-shaped environment names"), + ("process_count", "process count"), + ], +) +def test_resolved_server_rejects_one_field_join_drift( + multi_node_plan: ResolvedSlurmRunPlan, + resolved_multi_replica_server: ResolvedVllmServerDeployment, + client_image_inspection: ImageInspectionRecord, + mutation: str, + message: str, +) -> None: + payload = resolved_multi_replica_server.model_dump(mode="json") + if mutation == "image_kind": + payload["image"]["inspection"] = client_image_inspection.model_copy( + update={"sqsh_sha256": payload["image"]["sha256"]} + ).model_dump(mode="json") + elif mutation == "executable": + payload["executable_path"] = "/usr/local/bin/other" + elif mutation == "nodes": + payload["node_indices"] = [1, 0] + elif mutation == "node_group_divisibility": + payload["node_indices"] = [0, 1, 2] + elif mutation == "gpu_divisibility": + payload["gpus_per_node"] = 7 + elif mutation == "topology": + payload["topology"]["replica_count"] = 3 + elif mutation == "expert_parallel": + payload["launch_policy"]["enable_expert_parallel"] = True + elif mutation == "logical_alias": + payload["logical_endpoint"]["model_alias"] = "other" + elif mutation == "logical_served_name": + payload["logical_endpoint"]["served_model_name"] = "other" + elif mutation == "logical_id": + payload["logical_endpoint"]["endpoint_id"] = "other-logical-endpoint" + elif mutation == "deployment_replica_index": + payload["backend_endpoints"][1]["deployment_replica_index"] = 2 + elif mutation == "backend_id": + payload["backend_endpoints"][1]["backend_id"] = "deployment-00000-backend-99999" + payload["readiness_probes"][1]["backend_id"] = "deployment-00000-backend-99999" + payload["logical_endpoint"]["backend_ids"][1] = "deployment-00000-backend-99999" + elif mutation == "logical_backends": + payload["logical_endpoint"]["backend_ids"] = list(reversed(payload["logical_endpoint"]["backend_ids"])) + elif mutation == "process_ids": + payload["processes"][1]["process_id"] = payload["processes"][0]["process_id"] + elif mutation == "process_order": + payload["processes"][1], payload["processes"][2] = payload["processes"][2], payload["processes"][1] + elif mutation == "readiness_count": + payload["readiness_probes"].append(payload["readiness_probes"][0]) + elif mutation == "backend_placement": + payload["backend_endpoints"][1]["replica_index_in_node_group"] = 0 + elif mutation == "readiness_target": + payload["readiness_probes"][1]["path"] = "/other" + elif mutation == "pipeline_ranks": + payload["processes"][1]["deployment_replica_index"] = 1 + elif mutation == "process_topology": + payload["processes"][2]["node_index"] = 1 + elif mutation == "head_endpoint": + payload["processes"][2]["http_port"] = 18000 + elif mutation == "rendezvous_consistency": + payload["processes"][1]["rendezvous"]["port"] = 19001 + elif mutation == "rendezvous_master": + payload["processes"][0]["rendezvous"]["master_node_index"] = 1 + payload["processes"][1]["rendezvous"]["master_node_index"] = 1 + elif mutation == "rendezvous_timeout": + payload["processes"][0]["rendezvous"]["timeout_seconds"] = 1 + payload["processes"][1]["rendezvous"]["timeout_seconds"] = 1 + elif mutation == "network_address": + payload["logical_endpoint"]["port"] = payload["backend_endpoints"][0]["port"] + elif mutation == "launch_policy_owned_arg": + payload["launch_policy"]["extra_args"] = ["--model", "example/other"] + elif mutation == "launch_policy_literal_secret": + payload["launch_policy"]["environment"] = {"HF_TOKEN": {"type": "literal", "value": "secret"}} + else: + extra = dict(payload["processes"][0]) + extra.update( + process_id="deployment-00000-replica-99999-rank-00000", + deployment_replica_index=99999, + ) + payload["processes"].append(extra) + + with pytest.raises(ValidationError, match=message): + ResolvedVllmServerDeployment.model_validate_json(json.dumps(payload)) + + +def test_single_node_process_rejects_rendezvous( + resolved_single_node_server: ResolvedVllmServerDeployment, +) -> None: + payload = resolved_single_node_server.processes[0].model_dump(mode="json") + payload["rendezvous"] = { + "node_group_index": 0, + "replica_index_in_node_group": 0, + "master_node_index": 0, + "port": 19000, + "timeout_seconds": 600, + } + + with pytest.raises(ValidationError, match="must not carry rendezvous"): + ResolvedVllmProcess.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("duplicate_backends", "must be unique"), + ("missing_retry", "must retry HTTP 429"), + ], +) +def test_logical_endpoint_rejects_invalid_backend_or_retry_contract( + resolved_single_node_server: ResolvedVllmServerDeployment, + mutation: str, + message: str, +) -> None: + payload = resolved_single_node_server.logical_endpoint.model_dump(mode="json") + if mutation == "duplicate_backends": + payload["backend_ids"] *= 2 + else: + payload["retry_status_codes"] = [] + + with pytest.raises(ValidationError, match=message): + ResolvedLogicalEndpoint.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/serving/test_vllm_launch_policy.py b/packages/data-designer-slurm/tests/serving/test_vllm_launch_policy.py new file mode 100644 index 000000000..2f8a28a16 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/test_vllm_launch_policy.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.config import QueueBackpressureConfig +from data_designer.slurm.serving.vllm import ResolvedVllmLaunchPolicy + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ({"readiness_path": "health"}, "absolute URL path"), + ({"readiness_path": "//other-host/health"}, "absolute URL path"), + ({"readiness_path": "/health check"}, "without whitespace"), + ({"startup_timeout_seconds": 10, "distributed_init_timeout_seconds": 11}, "must not exceed"), + ], +) +def test_launch_policy_rejects_invalid_readiness_or_deadlines(update: dict[str, object], message: str) -> None: + payload = { + "startup_timeout_seconds": 900, + "distributed_init_timeout_seconds": 600, + "lead_boot_standoff_seconds": 60, + "rank_launch_stagger_seconds": 5, + "readiness_path": "/health", + "enable_expert_parallel": False, + "queue_backpressure": {"max_waiting_requests": 128, "retry_after_seconds": 1}, + **update, + } + with pytest.raises(ValidationError, match=message): + ResolvedVllmLaunchPolicy.model_validate(payload) + + +@pytest.mark.parametrize( + "argument", + [ + "-a", + "-as", + "-dc", + "-dcp=2", + "-e", + "-n2", + "-n+2", + "-pc", + "-pcp=2", + "-r1", + "-r0", + "-t", + "--kv-transfer-config={}", + "--numa-bind", + "--reasoning-parser-plugin=/tmp/reasoning.py", + "--worker-cls=custom.Worker", + ], +) +def test_launch_policy_rejects_runtime_owned_arguments(argument: str) -> None: + with pytest.raises(ValidationError, match="owned by the compiler or runtime"): + ResolvedVllmLaunchPolicy( + startup_timeout_seconds=900, + distributed_init_timeout_seconds=600, + lead_boot_standoff_seconds=60, + rank_launch_stagger_seconds=5, + readiness_path="/health", + enable_expert_parallel=False, + queue_backpressure=QueueBackpressureConfig(), + extra_args=(argument,), + ) + + +@pytest.mark.parametrize( + "extra_args", + [ + ("--hf-overrides", '{"api_key":"plaintext-secret"}'), + ('--hf-overrides={"nested":{"access_token":"plaintext-secret"}}',), + ], +) +def test_launch_policy_rejects_plaintext_secrets_in_argument_payloads(extra_args: tuple[str, ...]) -> None: + with pytest.raises(ValidationError, match="plaintext values under secret-bearing keys"): + ResolvedVllmLaunchPolicy( + startup_timeout_seconds=900, + distributed_init_timeout_seconds=600, + lead_boot_standoff_seconds=60, + rank_launch_stagger_seconds=5, + readiness_path="/health", + enable_expert_parallel=False, + queue_backpressure=QueueBackpressureConfig(), + extra_args=extra_args, + ) + + +def test_launch_policy_rejects_argument_outer_whitespace() -> None: + with pytest.raises(ValidationError, match="leading or trailing whitespace"): + ResolvedVllmLaunchPolicy( + startup_timeout_seconds=900, + distributed_init_timeout_seconds=600, + lead_boot_standoff_seconds=60, + rank_launch_stagger_seconds=5, + readiness_path="/health", + enable_expert_parallel=False, + queue_backpressure=QueueBackpressureConfig(), + extra_args=(" --max-model-len",), + ) + + +@pytest.mark.parametrize("argument", ["=", "=x"]) +def test_launch_policy_rejects_empty_option_names(argument: str) -> None: + with pytest.raises(ValidationError, match="option name must not be empty"): + ResolvedVllmLaunchPolicy( + startup_timeout_seconds=900, + distributed_init_timeout_seconds=600, + lead_boot_standoff_seconds=60, + rank_launch_stagger_seconds=5, + readiness_path="/health", + enable_expert_parallel=False, + queue_backpressure=QueueBackpressureConfig(), + extra_args=(argument,), + ) + + +@pytest.mark.parametrize( + "environment_name", + ["CUDA_VISIBLE_DEVICES", "VLLM_API_KEY", "VLLM_DP_RANK", "VLLM_HOST_IP"], +) +def test_launch_policy_rejects_runtime_owned_environment_names(environment_name: str) -> None: + with pytest.raises(ValidationError, match="owned by the compiler or runtime"): + ResolvedVllmLaunchPolicy( + startup_timeout_seconds=900, + distributed_init_timeout_seconds=600, + lead_boot_standoff_seconds=60, + rank_launch_stagger_seconds=5, + readiness_path="/health", + enable_expert_parallel=False, + queue_backpressure=QueueBackpressureConfig(), + environment={environment_name: {"type": "secret", "environment": "EXTERNAL_VALUE"}}, + ) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py new file mode 100644 index 000000000..d566e22d6 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic test doubles for the Slurm integration.""" + +from __future__ import annotations + +from slurm_test_fakes.clock import FakeClock +from slurm_test_fakes.dependencies import FakeDependencyInstaller, FakeDependencyResolver +from slurm_test_fakes.images import FakeInspectionEnvironment +from slurm_test_fakes.services import ( + FakeBatchScriptRenderer, + FakeBenchmarkBackend, + FakeImageResolver, + FakeRunPlanningBackend, + FakeScriptError, +) +from slurm_test_fakes.serving import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend +from slurm_test_fakes.slurm import ( + FakeCommandResponse, + FakeSlurmArray, + FakeSlurmJob, + FakeSlurmRunner, + FakeSlurmTask, +) + +__all__ = [ + "FakeBatchScriptRenderer", + "FakeBenchmarkBackend", + "FakeClock", + "FakeCommandResponse", + "FakeDependencyInstaller", + "FakeDependencyResolver", + "FakeImageResolver", + "FakeInspectionEnvironment", + "FakeLogicalEndpoint", + "FakeRunPlanningBackend", + "FakeScriptError", + "FakeServingState", + "FakeSlurmArray", + "FakeSlurmJob", + "FakeSlurmRunner", + "FakeSlurmTask", + "FakeVllmBackend", +] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py new file mode 100644 index 000000000..45485aae0 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timedelta + + +@dataclass +class FakeClock: + """Clock advanced only by explicit test input.""" + + current_time: datetime + monotonic_time: float = 0.0 + sleep_calls: list[float] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.current_time.tzinfo is None or self.current_time.utcoffset() != timedelta(0): + raise ValueError("current_time must be in UTC") + if self.monotonic_time < 0: + raise ValueError("monotonic_time must not be negative") + + def now(self) -> datetime: + """Return the controlled wall-clock time.""" + return self.current_time + + def monotonic(self) -> float: + """Return the controlled monotonic time.""" + return self.monotonic_time + + def sleep(self, seconds: float) -> None: + """Record a sleep and advance both clocks without blocking.""" + self.advance(seconds) + self.sleep_calls.append(seconds) + + def advance(self, seconds: float) -> None: + """Advance both clocks by a non-negative duration.""" + if seconds < 0: + raise ValueError("seconds must not be negative") + self.current_time += timedelta(seconds=seconds) + self.monotonic_time += seconds diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py new file mode 100644 index 000000000..ac0a8a4cb --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from pathlib import Path +from typing import Generic, TypeVar + +RequestT = TypeVar("RequestT") +ResolutionT = TypeVar("ResolutionT") +LockT = TypeVar("LockT") +InstallationT = TypeVar("InstallationT") + + +class FakeDependencyResolver(Generic[RequestT, ResolutionT]): + """Return or raise exact scripted dependency-resolution outcomes.""" + + def __init__(self, responses: Iterable[tuple[RequestT, ResolutionT | BaseException]]) -> None: + self._responses = deque(responses) + self.calls: list[RequestT] = [] + + def resolve(self, request: RequestT) -> ResolutionT: + """Resolve one expected request without consulting package indexes.""" + self.calls.append(request) + if not self._responses: + raise AssertionError("unexpected dependency resolution") + expected, response = self._responses.popleft() + if request != expected: + raise AssertionError(f"expected dependency request {expected!r}, got {request!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted resolution was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency resolutions remain") + + +class FakeDependencyInstaller(Generic[LockT, InstallationT]): + """Return or raise exact scripted dependency-install outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[tuple[LockT, Path], InstallationT | BaseException]], + ) -> None: + self._responses = deque(responses) + self.calls: list[tuple[LockT, Path]] = [] + + def install(self, lock: LockT, target: Path) -> InstallationT: + """Install one expected lock without invoking an installer process.""" + call = (lock, target) + self.calls.append(call) + if not self._responses: + raise AssertionError("unexpected dependency installation") + expected, response = self._responses.popleft() + if call != expected: + raise AssertionError(f"expected dependency installation {expected!r}, got {call!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted installation was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency installations remain") diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch new file mode 100644 index 000000000..7a2687367 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +#SBATCH --job-name=dd-two-model +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=3 +#SBATCH --cpus-per-gpu=4 +#SBATCH --time=03:55:00 +#SBATCH --array=0-1%2 +#SBATCH --chdir="/workspace/primary/runs/run-001" +#SBATCH --output="/workspace/primary/runs/run-001/slurm-attempt-0001-%A_%a.out" +#SBATCH --error="/workspace/primary/runs/run-001/slurm-attempt-0001-%A_%a.err" +#SBATCH --gres=gpu:8 +#SBATCH --exclusive +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +export PYTHONNOUSERSITE=1 + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" +readonly DD_PLAN_SHA256="cc5b1fd8d7f3ad664ee56098a02747fa69888f0979646f67c78796d46674044b" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${actual_sha256%% *}" == "$1" ]] +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +if [[ ! ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]]; then + printf '%s\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +source <(tar -xOf "${DD_RUNTIME_ARCHIVE}" scratch.sh) +trap 'dd_cleanup_allocation_scratch || true' EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +dd_initialize_local_scratch +dd_stage_allocation_runtime +readonly DD_RUNTIME_ROOT="${DD_SCRATCH_ROOT}/runtime" + +source "${DD_RUNTIME_ROOT}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch new file mode 100644 index 000000000..8fe6975ea --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +#SBATCH --job-name=data-designer +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --time=03:55:00 +#SBATCH --array=0 +#SBATCH --chdir="/workspace/primary/runs/run-single" +#SBATCH --output="/workspace/primary/runs/run-single/slurm-attempt-0001-%A_%a.out" +#SBATCH --error="/workspace/primary/runs/run-single/slurm-attempt-0001-%A_%a.err" +#SBATCH --gres=gpu:8 +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +export PYTHONNOUSERSITE=1 + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" +readonly DD_PLAN_SHA256="8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${actual_sha256%% *}" == "$1" ]] +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +if [[ ! ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]]; then + printf '%s\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 + exit 64 +fi +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +source <(tar -xOf "${DD_RUNTIME_ARCHIVE}" scratch.sh) +trap 'dd_cleanup_allocation_scratch || true' EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +dd_initialize_local_scratch +dd_stage_allocation_runtime +readonly DD_RUNTIME_ROOT="${DD_SCRATCH_ROOT}/runtime" + +source "${DD_RUNTIME_ROOT}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt new file mode 100644 index 000000000..7ac0e7937 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt @@ -0,0 +1,6 @@ +4101_0|TIMEOUT|0:125 +4101_1|NODE_FAIL|1:0 +4101_2|PREEMPTED|0:0 +4101_3|REQUEUED|0:0 +4101_4|OUT_OF_MEMORY|0:125 +4101_5|CANCELLED by 1234|0:15 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt new file mode 100644 index 000000000..17686ec3f --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt @@ -0,0 +1,2 @@ +4101_0|COMPLETED|0:0 +4101_1|FAILED|1:0 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt new file mode 100644 index 000000000..3e5fe27b7 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt @@ -0,0 +1 @@ +batch*|gpu:2 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt new file mode 100644 index 000000000..80e935355 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt @@ -0,0 +1,2 @@ +4101_0|PENDING +4101_1|RUNNING diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt new file mode 100644 index 000000000..670c5234d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt @@ -0,0 +1 @@ +malformed scheduler output diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/images.py b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py new file mode 100644 index 000000000..7ab4c084b --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic image-environment facts for package inspector tests.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field + +from data_designer.slurm.config import InstalledDistribution +from data_designer.slurm.images.errors import ImageInspectionError + + +@dataclass(frozen=True, slots=True) +class FakeInspectionEnvironment: + """Expose explicit client and serving facts without entering a container.""" + + python_implementation: str = "cpython" + python_version: str = "3.12.12" + python_abi: str = "cp312" + distributions: tuple[InstalledDistribution, ...] = () + distribution_versions: Mapping[str, str] = field(default_factory=dict) + executables: Mapping[str, str] = field(default_factory=dict) + + def get_python_implementation(self) -> str: + """Return the configured Python implementation.""" + return self.python_implementation + + def get_python_version(self) -> str: + """Return the configured Python version.""" + return self.python_version + + def get_python_abi(self) -> str: + """Return the configured Python ABI.""" + return self.python_abi + + def list_distributions(self) -> tuple[InstalledDistribution, ...]: + """Return the configured distribution inventory.""" + return self.distributions + + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one configured distribution version and console-script path.""" + try: + return (self.distribution_versions[name], self.executables[name]) + except KeyError: + raise ImageInspectionError(f"required distribution console script {name!r} is not installed") from None diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py new file mode 100644 index 000000000..9db2d7af0 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic fakes for public Slurm service boundaries.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from pathlib import Path +from typing import Generic, TypeVar + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageKind, + ImageRef, +) +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.planning import ResolvedImage, ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmBatchScriptRenderer, + SlurmBenchmarkBackend, + SlurmImageResolver, + SlurmRunPlanner, +) + +_RequestT = TypeVar("_RequestT") +_ResultT = TypeVar("_ResultT") + + +class FakeScriptError(BaseException): + """Signal an invalid deterministic-fake interaction.""" + + +class _ScriptedResponses(Generic[_RequestT, _ResultT]): + def __init__(self, responses: Iterable[tuple[_RequestT, _ResultT | BaseException]]) -> None: + self._responses = deque(responses) + self.calls: list[_RequestT] = [] + + def next(self, request: _RequestT, *, operation: str) -> _ResultT: + self.calls.append(request) + if not self._responses: + raise FakeScriptError(f"unexpected {operation}") + expected, response = self._responses.popleft() + if request != expected: + raise FakeScriptError(f"expected {operation} {expected!r}, got {request!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self, *, operation: str) -> None: + if self._responses: + raise FakeScriptError(f"{len(self._responses)} {operation} responses remain") + + +class FakeRunPlanningBackend(SlurmRunPlanner): + """Return or raise exact scripted run-planning outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[DataDesignerSlurmConfig, ResolvedSlurmRunPlan | BaseException]], + ) -> None: + self._script = _ScriptedResponses(responses) + self.calls = self._script.calls + + def plan(self, config: DataDesignerSlurmConfig, *, source_root: Path) -> ResolvedSlurmRunPlan: + """Return the plan scripted for one exact authored config.""" + del source_root + return self._script.next(config, operation="run plan") + + def assert_complete(self) -> None: + """Assert that every scripted plan was consumed.""" + self._script.assert_complete(operation="run plan") + + +class FakeBatchScriptRenderer(SlurmBatchScriptRenderer): + """Return or raise exact scripted batch-rendering outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[tuple[ResolvedSlurmRunPlan, int], str | BaseException]], + ) -> None: + self._script = _ScriptedResponses(responses) + self.calls = self._script.calls + + def __call__(self, plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: + """Return the script for one exact plan and attempt ordinal.""" + return self._script.next((plan, attempt_ordinal), operation="batch render") + + def assert_complete(self) -> None: + """Assert that every scripted render was consumed.""" + self._script.assert_complete(operation="batch render") + + +class FakeImageResolver(SlurmImageResolver): + """Return or raise exact scripted image-resolution outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[tuple[ImageRef, ImageKind], ResolvedImage | BaseException]], + ) -> None: + self._script = _ScriptedResponses(responses) + self.calls = self._script.calls + + def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedImage: + """Return the image scripted for one exact reference and kind.""" + return self._script.next((reference, expected_kind), operation="image resolution") + + def assert_complete(self) -> None: + """Assert that every scripted image resolution was consumed.""" + self._script.assert_complete(operation="image resolution") + + +class FakeBenchmarkBackend(SlurmBenchmarkBackend): + """Return or raise exact scripted benchmark outcomes.""" + + def __init__( + self, + *, + 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, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: + """Return the manifest scripted for one exact benchmark config.""" + return self._run_script.next((config, source_root, force), operation="benchmark run") + + 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, fail_if_incomplete), + operation="benchmark analysis", + ) + + def assert_complete(self) -> None: + """Assert that every scripted benchmark response was consumed.""" + self._run_script.assert_complete(operation="benchmark run") + self._analysis_script.assert_complete(operation="benchmark analysis") diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py new file mode 100644 index 000000000..0bc05aed3 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from enum import Enum + +import httpx + + +class FakeServingState(str, Enum): + """Lifecycle states shared by the serving test doubles.""" + + CREATED = "created" + STARTING = "starting" + READY = "ready" + FAILED = "failed" + STOPPED = "stopped" + + +class FakeVllmBackend: + """In-memory vLLM boundary with explicit lifecycle and responses.""" + + def __init__( + self, + endpoint: str, + *, + rank: int, + responses: Iterable[httpx.Response | BaseException] = (), + ) -> None: + if rank < 0: + raise ValueError("rank must not be negative") + self.endpoint = endpoint + self.rank = rank + self.state = FakeServingState.CREATED + self.failure_reason: str | None = None + self.requests: list[httpx.Request] = [] + self.start_calls = 0 + self.cleanup_calls = 0 + self._responses = deque(responses) + + def start(self) -> None: + """Enter the explicit startup state.""" + self.start_calls += 1 + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"backend cannot start from {self.state.value}") + self.state = FakeServingState.STARTING + + def mark_ready(self) -> None: + """Mark startup complete.""" + if self.state is not FakeServingState.STARTING: + raise RuntimeError("backend must be starting before it becomes ready") + self.state = FakeServingState.READY + + def fail(self, reason: str) -> None: + """Inject a backend or rank failure.""" + if not reason: + raise ValueError("failure reason must not be empty") + self.failure_reason = reason + self.state = FakeServingState.FAILED + + def queue_response(self, response: httpx.Response | BaseException) -> None: + """Queue one explicit generation response or failure.""" + self._responses.append(response) + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle a health or generation request without opening a socket.""" + self.requests.append(request) + if request.url.path == "/health": + status_code = 200 if self.state is FakeServingState.READY else 503 + return httpx.Response(status_code, request=request) + if self.state is not FakeServingState.READY: + return httpx.Response(503, json={"error": "backend_unavailable"}, request=request) + if not self._responses: + return httpx.Response( + 200, + json={"backend": self.endpoint, "rank": self.rank}, + request=request, + ) + response = self._responses.popleft() + if isinstance(response, BaseException): + raise response + response.request = request + return response + + def cleanup(self) -> None: + """Stop the backend once while allowing cleanup re-entry.""" + self.cleanup_calls += 1 + self.state = FakeServingState.STOPPED + + +class FakeLogicalEndpoint: + """In-memory logical endpoint over one or more fake vLLM backends.""" + + def __init__(self, endpoint: str, backends: Iterable[FakeVllmBackend]) -> None: + self.endpoint = endpoint + self.backends = tuple(backends) + if not self.backends: + raise ValueError("logical endpoints require at least one backend") + self.state = FakeServingState.CREATED + self.published_endpoint: str | None = None + self.requests: list[httpx.Request] = [] + self.publish_calls = 0 + self.cleanup_calls = 0 + self.failure_reason: str | None = None + self._next_backend = 0 + self._publication_failure: Exception | None = None + + def start(self) -> None: + """Start every backend and the readiness aggregator.""" + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"logical endpoint cannot start from {self.state.value}") + for backend in self.backends: + backend.start() + self.state = FakeServingState.STARTING + + def refresh_readiness(self) -> FakeServingState: + """Aggregate backend readiness with coordinated failure semantics.""" + if self.state is FakeServingState.STOPPED: + return self.state + if self._publication_failure is not None and self.publish_calls: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + return self.state + backend_states = {backend.state for backend in self.backends} + if FakeServingState.FAILED in backend_states: + self.state = FakeServingState.FAILED + self.failure_reason = "backend_failed" + elif backend_states == {FakeServingState.READY}: + self.state = FakeServingState.READY + else: + self.state = FakeServingState.STARTING + return self.state + + def publish(self) -> str: + """Publish the logical endpoint after every backend is ready.""" + if self.refresh_readiness() is not FakeServingState.READY: + raise RuntimeError("logical endpoint is not ready") + self.publish_calls += 1 + if self._publication_failure is not None: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + raise self._publication_failure + self.published_endpoint = self.endpoint + return self.endpoint + + def script_publication_failure(self, error: Exception) -> None: + """Inject one explicit endpoint-publication failure.""" + self._publication_failure = error + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle an aggregated health request or forward in round-robin order.""" + self.requests.append(request) + state = self.refresh_readiness() + if request.url.path == "/health": + status_code = 200 if state is FakeServingState.READY and self.published_endpoint else 503 + return httpx.Response(status_code, request=request) + if state is FakeServingState.FAILED: + return httpx.Response(503, json={"error": self.failure_reason}, request=request) + if state is not FakeServingState.READY or self.published_endpoint is None: + return httpx.Response(503, json={"error": "endpoint_unavailable"}, request=request) + response: httpx.Response | None = None + for _ in self.backends: + backend = self.backends[self._next_backend] + self._next_backend = (self._next_backend + 1) % len(self.backends) + response = backend.handle(request) + if response.status_code != 429: + return response + assert response is not None + return response + + def cleanup(self) -> None: + """Stop every backend and unpublish the process-local endpoint idempotently.""" + self.cleanup_calls += 1 + if self.state is FakeServingState.STOPPED: + return + for backend in self.backends: + backend.cleanup() + self.published_endpoint = None + self.state = FakeServingState.STOPPED diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py new file mode 100644 index 000000000..2350b5ab2 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import re +import subprocess +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.state import SchedulerIdentity + +_JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +_SACCT_REQUIRED_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", +) + + +@dataclass(frozen=True) +class FakeCommandResponse: + """One deterministic subprocess response.""" + + stdout: str = "" + stderr: str = "" + returncode: int = 0 + + +@dataclass +class FakeSlurmJob: + """Mutable scheduler views for one ordinary batch-job identity.""" + + job_id: int + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + def __post_init__(self) -> None: + if type(self.job_id) is not int or self.job_id <= 0: + raise ValueError("fake Slurm job IDs must be positive integers") + + +@dataclass +class FakeSlurmTask: + """Mutable scheduler views for one canonical array-task identity.""" + + scheduler: SchedulerIdentity + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + +@dataclass +class FakeSlurmArray: + """A deterministic array submission exposed through Slurm command output.""" + + tasks: tuple[FakeSlurmTask, ...] + + def __post_init__(self) -> None: + if not self.tasks: + raise ValueError("fake Slurm arrays require at least one task") + job_ids = {task.scheduler.array_job_id for task in self.tasks} + task_ids = [task.scheduler.array_task_id for task in self.tasks] + if len(job_ids) != 1: + raise ValueError("fake Slurm array tasks must share one array job ID") + if len(task_ids) != len(set(task_ids)): + raise ValueError("fake Slurm array task IDs must be unique") + + @property + def array_job_id(self) -> int: + """Return the canonical job ID shared by the array tasks.""" + return self.tasks[0].scheduler.array_job_id + + +class FakeSlurmRunner: + """Stateful fake that copies configured jobs before exposing Slurm commands.""" + + def __init__( + self, + arrays: Iterable[FakeSlurmArray] = (), + *, + jobs: Iterable[FakeSlurmJob] = (), + sinfo_responses: Mapping[tuple[str, ...], FakeCommandResponse] | None = None, + ) -> None: + submissions = (*arrays, *jobs) + job_ids = tuple( + submission.array_job_id if isinstance(submission, FakeSlurmArray) else submission.job_id + for submission in submissions + ) + if len(job_ids) != len(set(job_ids)): + raise ValueError("fake Slurm submissions must have unique job IDs") + self._pending_submissions = deque(copy.deepcopy(submissions)) + self._submitted_arrays: dict[int, FakeSlurmArray] = {} + self._submitted_jobs: dict[int, FakeSlurmJob] = {} + self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} + self._sinfo_responses = dict(sinfo_responses or {}) + self.calls: list[tuple[str, ...]] = [] + self.inputs: list[str | None] = [] + + def run( + self, + command: Sequence[str], + *, + check: bool = False, + input_text: str | None = None, + environment: Mapping[str, str] | None = None, + ) -> subprocess.CompletedProcess[str]: + """Run one fake Slurm command and optionally raise on failure.""" + if not command: + raise ValueError("command must not be empty") + argv = tuple(command) + self.calls.append(argv) + self.inputs.append(input_text) + del environment + command_name = Path(argv[0]).name + scripted = self._scripted_responses.get(command_name) + if scripted: + response = scripted.popleft() + else: + response = self._dispatch(command_name, argv) + + completed = subprocess.CompletedProcess( + args=list(argv), + returncode=response.returncode, + stdout=response.stdout, + stderr=response.stderr, + ) + if check and response.returncode: + raise subprocess.CalledProcessError( + response.returncode, + list(argv), + output=response.stdout, + stderr=response.stderr, + ) + return completed + + def script_next(self, command_name: str, response: FakeCommandResponse) -> None: + """Inject one explicit response before the command's stateful behavior.""" + self._scripted_responses.setdefault(command_name, deque()).append(response) + + def set_task_state( + self, + scheduler: SchedulerIdentity, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable queue and accounting states.""" + task = self._find_task(scheduler) + task.queue_state = queue_state + task.accounting_state = accounting_state + task.exit_code = exit_code + + def set_job_state( + self, + job_id: int, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable states for one ordinary job.""" + job = self._submitted_jobs[job_id] + job.queue_state = queue_state + job.accounting_state = accounting_state + job.exit_code = exit_code + + def assert_scripts_consumed(self) -> None: + """Assert that no scripted command response remains.""" + remaining = sum(len(responses) for responses in self._scripted_responses.values()) + if remaining: + raise AssertionError(f"{remaining} scripted Slurm responses remain") + + def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResponse: + handlers = { + "sacct": self._run_sacct, + "sbatch": self._run_sbatch, + "scancel": self._run_scancel, + "sinfo": self._run_sinfo, + "squeue": self._run_squeue, + } + try: + handler = handlers[command_name] + except KeyError: + raise AssertionError(f"unexpected command {command_name!r}") from None + return handler(argv) + + def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: + if not self._pending_submissions: + return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) + submission = self._pending_submissions.popleft() + if isinstance(submission, FakeSlurmArray): + job_id = submission.array_job_id + self._submitted_arrays[job_id] = submission + else: + job_id = submission.job_id + self._submitted_jobs[job_id] = submission + if "--parsable" in argv[1:]: + return FakeCommandResponse(stdout=f"{job_id}\n") + return FakeCommandResponse(stdout=f"Submitted batch job {job_id}\n") + + def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SQUEUE_REQUIRED_ARGUMENTS) + rows = [ + f"{job.job_id}|{job.queue_state}" + for job in self._selected_submitted_jobs(argv) + if job.queue_state is not None + ] + [ + f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.queue_state}" + for task in self._selected_submitted_tasks(argv) + if task.queue_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) + + def _run_sacct(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SACCT_REQUIRED_ARGUMENTS) + rows = [ + f"{job.job_id}|{job.accounting_state}|{job.exit_code}" + for job in self._selected_submitted_jobs(argv) + if job.accounting_state is not None + ] + [ + (f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.accounting_state}|{task.exit_code}") + for task in self._selected_submitted_tasks(argv) + if task.accounting_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in sorted(rows))) + + def _run_scancel(self, argv: tuple[str, ...]) -> FakeCommandResponse: + targets = tuple(argument for argument in argv[1:] if not argument.startswith("-")) + if len(targets) != 1: + return FakeCommandResponse(stderr="expected one cancellation target\n", returncode=1) + target = targets[0] + if "_" not in target and target.isdecimal() and int(target) in self._submitted_jobs: + job = self._submitted_jobs[int(target)] + job.queue_state = None + job.accounting_state = "CANCELLED" + job.exit_code = "0:15" + return FakeCommandResponse() + try: + tasks = self._tasks_for_target(target) + except (KeyError, ValueError): + return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) + for task in tasks: + task.queue_state = None + task.accounting_state = "CANCELLED" + task.exit_code = "0:15" + return FakeCommandResponse() + + def _run_sinfo(self, argv: tuple[str, ...]) -> FakeCommandResponse: + key = ("sinfo", *argv[1:]) + try: + return self._sinfo_responses[key] + except KeyError: + raise AssertionError(f"unexpected sinfo query {key!r}") from None + + @staticmethod + def _require_arguments(argv: tuple[str, ...], required: tuple[str, ...]) -> None: + missing = tuple(argument for argument in required if argument not in argv[1:]) + if missing: + raise AssertionError(f"missing required arguments {missing!r} in {argv!r}") + + def _tasks_for_target(self, target: str) -> tuple[FakeSlurmTask, ...]: + if "_" not in target: + return self._submitted_arrays[int(target)].tasks + job_id, task_id = target.split("_", 1) + scheduler = SchedulerIdentity(array_job_id=int(job_id), array_task_id=int(task_id)) + return (self._find_task(scheduler),) + + def _find_task(self, scheduler: SchedulerIdentity) -> FakeSlurmTask: + array = self._submitted_arrays[scheduler.array_job_id] + for task in array.tasks: + if task.scheduler == scheduler: + return task + raise KeyError(scheduler) + + def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: + return sorted( + (task for array in self._submitted_arrays.values() for task in array.tasks), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) + + def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: + selectors = self._job_selectors(argv) + if not selectors: + return self._sorted_submitted_tasks() + + selected: dict[SchedulerIdentity, FakeSlurmTask] = {} + for selector in selectors: + try: + tasks = self._tasks_for_target(selector) + except KeyError: + continue + selected.update((task.scheduler, task) for task in tasks) + return sorted( + selected.values(), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) + + def _selected_submitted_jobs(self, argv: tuple[str, ...]) -> list[FakeSlurmJob]: + selectors = self._job_selectors(argv) + if not selectors: + return sorted(self._submitted_jobs.values(), key=lambda job: job.job_id) + return sorted( + ( + self._submitted_jobs[int(selector)] + for selector in selectors + if "_" not in selector and int(selector) in self._submitted_jobs + ), + key=lambda job: job.job_id, + ) + + @staticmethod + def _job_selectors(argv: tuple[str, ...]) -> tuple[str, ...]: + selectors: list[str] = [] + for index, argument in enumerate(argv[1:]): + if argument in {"-j", "--jobs"}: + try: + selectors.extend(argv[index + 2].split(",")) + except IndexError: + raise AssertionError(f"missing job selector in {argv!r}") from None + elif argument.startswith("--jobs="): + selectors.extend(argument.partition("=")[2].split(",")) + if not selectors: + return () + for selector in selectors: + if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: + raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") + return tuple(selectors) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py new file mode 100644 index 000000000..b160d5c21 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime, timedelta, timezone + +import pytest + +from slurm_test_fakes import FakeClock + + +def test_fake_clock_advances_without_blocking(fake_clock: FakeClock) -> None: + fake_clock.sleep(2.5) + fake_clock.advance(1.5) + + assert fake_clock.now() == datetime(2026, 8, 18, 12, 0, 4, tzinfo=timezone.utc) + assert fake_clock.monotonic() == 104.0 + assert fake_clock.sleep_calls == [2.5] + + +@pytest.mark.parametrize( + "clock", + ( + pytest.param(lambda: FakeClock(datetime(2026, 8, 18, 12)), id="naive-datetime"), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 13, tzinfo=timezone(timedelta(hours=1)))), + id="non-utc-offset", + ), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=-1), + id="negative-monotonic", + ), + ), +) +def test_fake_clock_rejects_ambient_or_invalid_time(clock: Callable[[], FakeClock]) -> None: + with pytest.raises(ValueError): + clock() + + +def test_fake_clock_rejects_negative_advances(fake_clock: FakeClock) -> None: + with pytest.raises(ValueError, match="must not be negative"): + fake_clock.advance(-0.1) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py new file mode 100644 index 000000000..d32aefa5b --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, ImageInspectionRecord +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + + +def test_shared_contract_fixtures_load_canonical_records( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, + client_result: ClientResult, + run_manifest: RunManifest, + shard_manifests: tuple[ShardManifest, ...], + attempt_manifest: AttemptManifest, + attempt_readiness: AttemptReadiness, + finalization_client_result: ClientResult, + candidate_output_manifest: CandidateOutputManifest, + shard_winner: ShardWinner, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, + client_image_inspection: ImageInspectionRecord, + serving_image_inspection: ImageInspectionRecord, + collection_plan: CollectionPlan, + accounting_lag_observation: SchedulerObservation, +) -> None: + assert single_node_plan.run_id == "run-single" + assert multi_node_plan.run_id == "run-001" + assert client_result.run_id == "run-001" + assert run_manifest.run_id == "run-single" + assert len(shard_manifests) == 1 + assert attempt_manifest.attempt_id == "attempt-0001" + assert attempt_readiness.revision == 1 + assert finalization_client_result.outcome.value == "complete" + assert candidate_output_manifest.winner_eligible + assert shard_winner.attempt_id == "attempt-0001" + assert benchmark_config.name == "generator-scaling" + assert benchmark_manifest.benchmark_id == benchmark_report.benchmark_id + assert client_image_inspection.inspection.kind == "client" + assert serving_image_inspection.inspection.kind == "serving" + assert collection_plan.collection_id == "collection-0001" + assert accounting_lag_observation.state.value == "accounting_lag" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py new file mode 100644 index 000000000..594ca9bbb --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +from importlib.metadata import entry_points +from pathlib import Path + +import pytest + +from data_designer.plugins import PluginType +from data_designer.plugins.registry import PluginRegistry +from data_designer.slurm.config import ClientDependencies, DataDesignerSlurmConfig +from data_designer.slurm.planning import ResolvedDependencyLock +from slurm_test_fakes import FakeDependencyInstaller, FakeDependencyResolver + + +def test_fake_dependency_resolver_scripts_compatible_incompatible_and_missing_cases( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, +) -> None: + compatible = authored_run.client.dependencies + incompatible = ClientDependencies(requirements=["fake-data-designer-plugin==2.0.0"]) + missing = ClientDependencies(requirements=["missing-plugin==1.0.0"]) + resolver: FakeDependencyResolver[ClientDependencies, ResolvedDependencyLock] = FakeDependencyResolver( + ( + (compatible, dependency_lock), + (incompatible, ValueError("incompatible Data Designer version")), + (missing, FileNotFoundError("dependency artifact is missing")), + ) + ) + + assert resolver.resolve(compatible) is dependency_lock + with pytest.raises(ValueError, match="incompatible"): + resolver.resolve(incompatible) + with pytest.raises(FileNotFoundError, match="missing"): + resolver.resolve(missing) + + assert resolver.calls == [compatible, incompatible, missing] + resolver.assert_complete() + + +def test_fake_dependency_installer_scripts_success_and_digest_mismatch( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + first_target = tmp_path / "compatible" + second_target = tmp_path / "digest-mismatch" + mismatched_lock = dependency_lock.model_copy(update={"client_runtime_sha256": "a" * 64}) + installer = FakeDependencyInstaller[ResolvedDependencyLock, tuple[object, ...]]( + ( + ((dependency_lock, first_target), dependency_lock.overlay_packages), + ((mismatched_lock, second_target), ValueError("lock digest mismatch")), + ) + ) + + assert installer.install(dependency_lock, first_target) == dependency_lock.overlay_packages + with pytest.raises(ValueError, match="digest mismatch"): + installer.install(mismatched_lock, second_target) + + assert installer.calls == [(dependency_lock, first_target), (mismatched_lock, second_target)] + installer.assert_complete() + + +def test_dependency_fakes_raise_cancellation_signals( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + request = ClientDependencies(requirements=[]) + resolver = FakeDependencyResolver(((request, KeyboardInterrupt()),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path), SystemExit(2)),)) + + with pytest.raises(KeyboardInterrupt): + resolver.resolve(request) + with pytest.raises(SystemExit, match="2"): + installer.install(dependency_lock, tmp_path) + + +def test_dependency_fakes_reject_unexpected_calls( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + expected = ClientDependencies(requirements=[]) + unexpected = ClientDependencies(requirements=["unexpected==1.0.0"]) + resolver = FakeDependencyResolver(((expected, dependency_lock),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path / "expected"), "installed"),)) + + with pytest.raises(AssertionError, match="expected dependency request"): + resolver.resolve(unexpected) + with pytest.raises(AssertionError, match="expected dependency installation"): + installer.install(dependency_lock, tmp_path / "unexpected") + + +def test_fake_plugin_overlay_supports_real_entry_point_discovery( + fake_plugin_overlay: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(str(fake_plugin_overlay)) + importlib.invalidate_caches() + monkeypatch.setattr("data_designer.plugins.registry.PLUGINS_DISABLED", False) + PluginRegistry.reset() + try: + discovered = tuple( + entry_point + for entry_point in entry_points(group="data_designer.plugins") + if entry_point.name == "fake-slurm-column" + ) + assert len(discovered) == 1 + assert discovered[0].dist is not None + assert discovered[0].dist.version == "1.0.0" + + plugin = discovered[0].load() + registry = PluginRegistry() + + assert plugin.name == "fake-slurm-column" + assert plugin.plugin_type is PluginType.COLUMN_GENERATOR + assert registry.get_plugin("fake-slurm-column") == plugin + finally: + PluginRegistry.reset() diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py new file mode 100644 index 000000000..d1bf2461e --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import posixpath +from pathlib import Path + +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "rendered" + + +def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + _assert_script_matches_plan( + single_node_plan, + "single_node.sbatch", + expected_fixture_sha256="cc24e90595495373d47f470afa877333d41dd5ccb85df57da0e9a3d9e3ac1838", + ) + _assert_script_matches_plan( + multi_node_plan, + "multi_node.sbatch", + expected_fixture_sha256="c185ed76ac5e95e778d2ef8f442c046aa0fb982ef318a4404a9f5124c1a6484d", + ) + + +def test_rendered_scripts_contain_no_site_or_user_data() -> None: + contents = "\n".join(path.read_text().casefold() for path in GOLDEN_DIRECTORY.glob("*.sbatch")) + + for forbidden in ( + "nvidia", + "cluster", + "login", + "lustre", + "/home/", + "/users/", + ): + assert forbidden not in contents + + +def _assert_script_matches_plan( + plan: ResolvedSlurmRunPlan, + filename: str, + *, + expected_fixture_sha256: str, +) -> None: + script = (GOLDEN_DIRECTORY / filename).read_text() + node_indices = ( + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + ) + node_count = max(node_indices) + 1 + array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.count > 1 and plan.array_tasks.max_concurrent is not None: + array = f"{array}%{plan.array_tasks.max_concurrent}" + plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") + run_root = posixpath.dirname(plan.authored_config.path) + + assert f"#SBATCH --job-name={plan.submission.job_name}\n" in script + assert f"#SBATCH --account={plan.submission.account}\n" in script + assert f"#SBATCH --partition={plan.submission.partition}\n" in script + assert f"#SBATCH --nodes={node_count}\n" in script + cpus_per_gpu = max(1, (plan.client.authored.cpus + plan.resolved_gpus_per_node - 1) // plan.resolved_gpus_per_node) + assert f"#SBATCH --cpus-per-gpu={cpus_per_gpu}\n" in script + assert f"#SBATCH --time={plan.submission.time_limit}\n" in script + assert f"#SBATCH --array={array}\n" in script + assert f'#SBATCH --chdir="{run_root}"\n' in script + assert f'#SBATCH --output="{run_root}/slurm-attempt-0001-%A_%a.out"\n' in script + assert f'#SBATCH --error="{run_root}/slurm-attempt-0001-%A_%a.err"\n' in script + assert f"#SBATCH --gres=gpu:{plan.resolved_gpus_per_node}\n" in script + assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"\n' in script + assert f'readonly DD_RUNTIME_SHA256="{plan.runtime_bundle.sha256}"\n' in script + assert f'readonly DD_PLAN="{plan_path}"\n' in script + assert f'readonly DD_PLAN_SHA256="{plan.compute_sha256()}"\n' in script + assert f'readonly DD_RUN_ROOT="{run_root}"\n' in script + assert script.count("dd_slurm_run_allocation") == 1 + assert script.endswith("\n") + assert hashlib.sha256(script.encode()).hexdigest() == expected_fixture_sha256 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py new file mode 100644 index 000000000..cf92ca840 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import httpx +import pytest + +from slurm_test_fakes import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend + + +def _client(endpoint: FakeLogicalEndpoint) -> httpx.Client: + return httpx.Client(transport=httpx.MockTransport(endpoint.handle), base_url=endpoint.endpoint) + + +def _start_and_publish(endpoint: FakeLogicalEndpoint) -> None: + endpoint.start() + for backend in endpoint.backends: + backend.mark_ready() + endpoint.publish() + + +def test_single_backend_startup_readiness_and_endpoint_publication() -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + + endpoint.start() + with _client(endpoint) as client: + assert client.get("/health").status_code == 503 + + backend.mark_ready() + assert endpoint.publish() == endpoint.endpoint + assert client.get("/health").status_code == 200 + assert client.post("/v1/chat/completions", json={"prompt": "fixture"}).json() == { + "backend": backend.endpoint, + "rank": 0, + } + + +def test_multiple_backends_receive_requests_in_deterministic_order( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + with _client(fake_logical_endpoint) as client: + responses = [client.post("/v1/chat/completions").json() for _ in range(3)] + + assert [response["rank"] for response in responses] == [0, 1, 0] + assert [len(backend.requests) for backend in fake_logical_endpoint.backends] == [2, 1] + + +def test_partial_multi_backend_startup_is_not_ready_and_cleans_up_once( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + fake_logical_endpoint.backends[0].mark_ready() + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.STARTING + with pytest.raises(RuntimeError, match="not ready"): + fake_logical_endpoint.publish() + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +@pytest.mark.parametrize("retry_after", ("2", None)) +def test_logical_endpoint_preserves_overload_response(retry_after: str | None) -> None: + headers = {"Retry-After": retry_after} if retry_after is not None else {} + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers=headers, json={"error": "overloaded"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers.get("Retry-After") == retry_after + assert response.json() == {"error": "overloaded"} + + +def test_logical_endpoint_retries_overload_against_another_backend() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, json={"error": "first overloaded"}),), + ) + second = FakeVllmBackend("http://127.0.0.1:31002", rank=1) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 200 + assert response.json()["rank"] == 1 + assert [len(backend.requests) for backend in endpoint.backends] == [1, 1] + + +def test_logical_endpoint_returns_the_final_overload_response() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers={"Retry-After": "1"}),), + ) + second = FakeVllmBackend( + "http://127.0.0.1:31002", + rank=1, + responses=(httpx.Response(429, headers={"Retry-After": "3"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers["Retry-After"] == "3" + + +def test_backend_rank_failure_fails_the_logical_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + fake_logical_endpoint.backends[1].fail("rank exited") + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert response.status_code == 503 + assert response.json() == {"error": "backend_failed"} + + +def test_endpoint_publication_failure_is_explicit_and_cleanup_remains_idempotent( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + for backend in fake_logical_endpoint.backends: + backend.mark_ready() + fake_logical_endpoint.script_publication_failure(RuntimeError("publication failed")) + + with pytest.raises(RuntimeError, match="publication failed"): + fake_logical_endpoint.publish() + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert fake_logical_endpoint.failure_reason == "endpoint_publication_failed" + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + assert response.json() == {"error": "endpoint_publication_failed"} + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +def test_backend_failure_injection_is_explicit() -> None: + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(RuntimeError("scripted backend failure"),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client, pytest.raises(RuntimeError, match="scripted backend failure"): + client.post("/v1/chat/completions") + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_backend_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + if terminal_state is FakeServingState.FAILED: + backend.fail("rank exited") + else: + backend.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + backend.start() + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_logical_endpoint_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + endpoint = FakeLogicalEndpoint( + "http://127.0.0.1:31000", + (FakeVllmBackend("http://127.0.0.1:31001", rank=0),), + ) + if terminal_state is FakeServingState.FAILED: + endpoint.start() + endpoint.backends[0].fail("rank exited") + endpoint.refresh_readiness() + else: + endpoint.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + endpoint.start() + + +def test_cleanup_is_idempotent_and_unpublishes_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + + assert fake_logical_endpoint.state is FakeServingState.STOPPED + assert fake_logical_endpoint.published_endpoint is None + assert fake_logical_endpoint.cleanup_calls == 2 + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py new file mode 100644 index 000000000..3051ca5ee --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from data_designer.slurm.state import SchedulerIdentity +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" +SQUEUE_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +SACCT_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", +) + + +def _submit(runner: FakeSlurmRunner) -> None: + completed = runner.run(("sbatch", "--parsable", "run.sbatch"), check=True) + assert completed.stdout == "4101\n" + + +def test_fake_slurm_runner_models_array_submission_and_active_states( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + observed = fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS), check=True) + + assert observed.stdout == (GOLDEN_DIRECTORY / "squeue_active.txt").read_text() + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "run.sbatch"), + ("squeue", *SQUEUE_ARGUMENTS), + ] + + +def test_fake_slurm_runner_exposes_terminal_accounting_precedence( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=0), + queue_state="RUNNING", + accounting_state="COMPLETED", + ) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ) + + assert "4101_1|RUNNING" in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert ( + fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + == (GOLDEN_DIRECTORY / "sacct_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_models_accounting_lag_and_later_terminal_state( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=0) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state=None, + ) + + assert "4101_0" not in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout == "" + + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + ) + assert "4101_0|COMPLETED|0:0" in fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + + +def test_fake_slurm_runner_filters_job_scoped_queries() -> None: + first = FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + second = FakeSlurmArray( + tasks=( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4201, array_task_id=0), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ), + ) + ) + runner = FakeSlurmRunner((first, second)) + runner.run(("sbatch", "first.sbatch"), check=True) + runner.run(("sbatch", "second.sbatch"), check=True) + + first_queue = runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs", "4101")).stdout + second_accounting = runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=4201")).stdout + + assert first_queue == "4101_0|PENDING\n" + assert second_accounting == "4201_0|FAILED|1:0\n" + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", *SQUEUE_ARGUMENTS, "--jobs", "--noheader"), + ("sacct", *SACCT_ARGUMENTS, "--jobs="), + ), +) +def test_fake_slurm_runner_rejects_malformed_job_selectors( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="malformed job selector"): + fake_slurm_runner.run(command) + + +@pytest.mark.parametrize("target", ("4101", "4101_1")) +def test_fake_slurm_runner_models_cancellation( + fake_slurm_runner: FakeSlurmRunner, + target: str, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("scancel", target), check=True).returncode == 0 + accounting = fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + + expected_tasks = (0, 1) if target == "4101" else (1,) + for task_id in expected_tasks: + assert f"4101_{task_id}|CANCELLED|0:15" in accounting + + +def test_fake_slurm_runner_models_ordinary_job_lifecycle_and_cancellation() -> None: + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(job_id=5101),)) + + assert runner.run(("sbatch", "--parsable", "image-build.sbatch"), check=True).stdout == "5101\n" + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|PENDING\n" + + assert runner.run(("scancel", "5101"), check=True).returncode == 0 + assert runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs=5101"), check=True).stdout == "" + assert runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=5101"), check=True).stdout == "5101|CANCELLED|0:15\n" + + +def test_fake_slurm_runner_rejects_duplicate_job_ids() -> None: + ordinary_job = FakeSlurmJob(job_id=4101) + array = FakeSlurmArray(tasks=(FakeSlurmTask(scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + + with pytest.raises(ValueError, match="unique job IDs"): + FakeSlurmRunner((array,), jobs=(ordinary_job,)) + + +def test_fake_slurm_runner_supports_malformed_output_and_command_failures( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + malformed = (GOLDEN_DIRECTORY / "squeue_malformed.txt").read_text() + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout=malformed)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + + assert fake_slurm_runner.run(("squeue",)).stdout == malformed + assert fake_slurm_runner.run(("sacct",)).returncode == 2 + with pytest.raises(subprocess.CalledProcessError) as error: + fake_slurm_runner.run(("sacct",), check=True) + + assert error.value.stderr == "accounting unavailable\n" + fake_slurm_runner.assert_scripts_consumed() + + +@pytest.mark.parametrize( + ("arguments", "expected_stdout"), + ( + (("run.sbatch",), "Submitted batch job 4101\n"), + (("--parsable", "run.sbatch"), "4101\n"), + ), +) +def test_fake_slurm_runner_matches_sbatch_parsable_mode( + arguments: tuple[str, ...], + expected_stdout: str, +) -> None: + runner = FakeSlurmRunner( + (FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + + assert runner.run(("sbatch", *arguments), check=True).stdout == expected_stdout + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", "--noheader"), + ("sacct", "--noheader", "--format=JobID,State,ExitCode"), + ), +) +def test_fake_slurm_runner_rejects_underspecified_state_queries( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="missing required arguments"): + fake_slurm_runner.run(command) + + +def test_fake_slurm_runner_exposes_retry_terminal_spellings() -> None: + states = ( + ("TIMEOUT", "0:125"), + ("NODE_FAIL", "1:0"), + ("PREEMPTED", "0:0"), + ("REQUEUED", "0:0"), + ("OUT_OF_MEMORY", "0:125"), + ("CANCELLED by 1234", "0:15"), + ) + runner = FakeSlurmRunner( + ( + FakeSlurmArray( + tasks=tuple( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=task_id), + queue_state=None, + accounting_state=state, + exit_code=exit_code, + ) + for task_id, (state, exit_code) in enumerate(states) + ) + ), + ) + ) + _submit(runner) + + assert ( + runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + == (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_bounds_sinfo_queries(fake_slurm_runner: FakeSlurmRunner) -> None: + response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%P|%G"), check=True) + + assert response.stdout == (GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + with pytest.raises(AssertionError, match="unexpected sinfo query"): + fake_slurm_runner.run(("sinfo", "--all")) + + +def test_fake_slurm_runner_rejects_unscripted_commands_and_submissions( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("sbatch", "second.sbatch")).returncode == 1 + with pytest.raises(AssertionError, match="unexpected command"): + fake_slurm_runner.run(("srun", "hostname")) diff --git a/packages/data-designer-slurm/tests/state/golden/accounting_lag.json b/packages/data-designer-slurm/tests/state/golden/accounting_lag.json new file mode 100644 index 000000000..9c97822f4 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/accounting_lag.json @@ -0,0 +1,10 @@ +{ + "observed_at": "2026-08-18T12:05:02Z", + "reconciliation_deadline": "2026-08-18T12:10:02Z", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "state": "accounting_lag" +} diff --git a/packages/data-designer-slurm/tests/state/golden/candidate_output.json b/packages/data-designer-slurm/tests/state/golden/candidate_output.json new file mode 100644 index 000000000..81e9a63f4 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/candidate_output.json @@ -0,0 +1,22 @@ +{ + "actual_records": 100, + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "created_at": "2026-08-18T12:05:00Z", + "dataset_path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/dataset", + "dataset_schema_digest": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + "files": [ + { + "byte_size": 4096, + "record_count": 100, + "relative_path": "part-00000.parquet", + "sha256": "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + } + ], + "outcome": "complete", + "provenance_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "requested_records": 100, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/state/golden/collection_plan.json b/packages/data-designer-slurm/tests/state/golden/collection_plan.json new file mode 100644 index 000000000..1c02ee94a --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -0,0 +1,23 @@ +{ + "collection_id": "collection-0001", + "created_at": "2026-08-18T12:06:00Z", + "host_destination": "/workspace/runs/run-0001/collection/run-0001.parquet", + "num_partitions": 1, + "overwrite": false, + "planned_shards": [ + { + "shard_id": "shard-00000", + "winner_manifest": { + "path": "/workspace/runs/run-0001/shards/shard-00000/winner.json", + "sha256": "e6efb5fd97b97c36e7dd1c4e5ab61dc954b611fd68d5e85956e357b36d2a8930" + } + } + ], + "python_executable": "/shared/data-designer/bin/python", + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "schema_version": 1 +} diff --git a/packages/data-designer-slurm/tests/state/golden/failed_attempt.json b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json new file mode 100644 index 000000000..244b7977d --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/failed_attempt.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-0002", + "attempt_ordinal": 2, + "candidate_output": null, + "created_at": "2026-08-18T13:00:00Z", + "resolved_plan": { + "path": "/workspace/runs/run-failed/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-failed", + "scheduler": { + "array_job_id": 4201, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "terminal_classification": "node_failed", + "updated_at": "2026-08-18T13:05:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json new file mode 100644 index 000000000..444351540 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json @@ -0,0 +1,39 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 2, + "last_probe": { + "observed_at": "2026-08-18T14:02:00Z", + "outcome": "failure", + "reason_code": "backend_starting", + "redacted_message": "One of two backends is ready" + }, + "model_alias": "generator", + "ready_backends": 1, + "state": "starting" + }, + { + "deployment_id": "deployment-00001", + "endpoint_publication": "published", + "expected_backends": 2, + "last_probe": { + "observed_at": "2026-08-18T14:02:00Z", + "outcome": "success", + "reason_code": "backends_ready", + "redacted_message": "All backends are ready" + }, + "model_alias": "judge", + "ready_backends": 2, + "state": "ready" + } + ], + "revision": 4, + "run_id": "run-multi", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "starting", + "updated_at": "2026-08-18T14:02:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/run_manifest.json b/packages/data-designer-slurm/tests/state/golden/run_manifest.json new file mode 100644 index 000000000..38825f78a --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/run_manifest.json @@ -0,0 +1,14 @@ +{ + "authored_config": { + "path": "/workspace/runs/run-0001/authored-config.json", + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + }, + "created_at": "2026-08-18T12:00:00Z", + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "schema_version": 1, + "shard_count": 1 +} diff --git a/packages/data-designer-slurm/tests/state/golden/shard_manifest.json b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json new file mode 100644 index 000000000..5e885b143 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/shard_manifest.json @@ -0,0 +1,18 @@ +{ + "created_at": "2026-08-18T12:00:01Z", + "input_partition": { + "path": "/workspace/runs/run-0001/shards/shard-00000/input-partition.json", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "record_range": { + "end_index_exclusive": 100, + "start_index": 0 + }, + "resume_workspace": { + "path": "/workspace/runs/run-0001/shards/shard-00000/dataset" + }, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "shard_index": 0 +} diff --git a/packages/data-designer-slurm/tests/state/golden/shard_winner.json b/packages/data-designer-slurm/tests/state/golden/shard_winner.json new file mode 100644 index 000000000..20a903c35 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/shard_winner.json @@ -0,0 +1,12 @@ +{ + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_manifest": { + "path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "c639099f85b3cdd2759d3b74c792f06dc4a61a8550622b50ba2a3512c10fdb68" + }, + "published_at": "2026-08-18T12:05:02Z", + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json new file mode 100644 index 000000000..2b5f14067 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "failed" + } + ], + "revision": 4, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:02:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json new file mode 100644 index 000000000..ca41b2d6f --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": null, + "model_alias": "primary", + "ready_backends": 0, + "state": "pending" + } + ], + "revision": 1, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "pending", + "updated_at": "2026-08-18T12:00:02Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json new file mode 100644 index 000000000..5001bc9fb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "failed", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:30Z", + "outcome": "failure", + "reason_code": "endpoint_publication_failed", + "redacted_message": "Logical endpoint publication failed" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "failed" + } + ], + "revision": 3, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:01:30Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json new file mode 100644 index 000000000..9337193d5 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:30Z", + "outcome": "success", + "reason_code": "backend_ready", + "redacted_message": "Backend is ready" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "ready" + } + ], + "revision": 3, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "ready", + "updated_at": "2026-08-18T12:01:30Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json new file mode 100644 index 000000000..7929011bb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:00Z", + "outcome": "failure", + "reason_code": "backend_starting", + "redacted_message": "Backend is starting" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "starting" + } + ], + "revision": 2, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "starting", + "updated_at": "2026-08-18T12:01:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json new file mode 100644 index 000000000..894019149 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "stopped" + } + ], + "revision": 5, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "stopped", + "updated_at": "2026-08-18T12:02:01Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/stale_readiness.json b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json new file mode 100644 index 000000000..511337eb8 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/stale_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0002", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T13:04:00Z", + "outcome": "success", + "reason_code": "backend_ready", + "redacted_message": "Backend was ready before the terminal failure" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "ready" + } + ], + "revision": 2, + "run_id": "run-failed", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "ready", + "updated_at": "2026-08-18T13:04:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/successful_attempt.json b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json new file mode 100644 index 000000000..03d28f620 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/successful_attempt.json @@ -0,0 +1,23 @@ +{ + "attempt_id": "attempt-0001", + "attempt_ordinal": 1, + "candidate_output": { + "path": "/workspace/runs/run-0001/shards/shard-00000/attempts/attempt-0001/output-manifest.json", + "sha256": "c639099f85b3cdd2759d3b74c792f06dc4a61a8550622b50ba2a3512c10fdb68" + }, + "created_at": "2026-08-18T12:00:02Z", + "resolved_plan": { + "path": "/workspace/runs/run-0001/resolved-plan.json", + "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + }, + "run_id": "run-0001", + "scheduler": { + "array_job_id": 4101, + "array_task_id": 0 + }, + "schema_version": 1, + "shard_id": "shard-00000", + "state": "succeeded", + "terminal_classification": "succeeded", + "updated_at": "2026-08-18T12:05:01Z" +} diff --git a/packages/data-designer-slurm/tests/state/test_model_edges.py b/packages/data-designer-slurm/tests/state/test_model_edges.py new file mode 100644 index 000000000..148c34f25 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.state import ( + MAXIMUM_CANDIDATE_OUTPUT_FILES, + ArtifactReference, + AttemptManifest, + AttemptReadiness, + CandidateOutputFile, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" + + +def _golden_payload(filename: str) -> dict[str, Any]: + return json.loads((GOLDEN_DIRECTORY / filename).read_text()) + + +@pytest.mark.parametrize( + "path", + ( + "/", + "//", + "//host/path", + "/workspace/control\ncharacter", + "/workspace//unnormalized", + ), +) +def test_absolute_artifact_paths_reject_unsafe_forms(path: str) -> None: + with pytest.raises(ValidationError): + ArtifactReference(path=path, sha256="a" * 64) + + +@pytest.mark.parametrize( + "path", + ( + "", + "/absolute", + "control\ncharacter", + ".", + "directory//unnormalized", + ), +) +def test_relative_output_paths_reject_unsafe_forms(path: str) -> None: + with pytest.raises(ValidationError): + CandidateOutputFile( + relative_path=path, + sha256="a" * 64, + byte_size=0, + record_count=0, + ) + + +@pytest.mark.parametrize( + "created_at", + ( + "2026-08-18T12:00:00", + "2026-08-18T13:00:00+01:00", + ), +) +def test_run_manifest_requires_utc_timestamp(created_at: str) -> None: + payload = _golden_payload("run_manifest.json") + payload["created_at"] = created_at + + with pytest.raises(ValidationError): + RunManifest.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("mutation", "message"), + ( + ({"updated_at": "2026-08-18T12:00:01Z"}, "must not precede"), + ({"terminal_classification": "failed"}, "succeeded terminal classification"), + ({"state": "failed", "terminal_classification": "succeeded"}, "failed attempts"), + ), +) +def test_attempt_manifest_rejects_invalid_terminal_details( + mutation: dict[str, object], + message: str, +) -> None: + payload = _golden_payload("successful_attempt.json") + payload.update(mutation) + + with pytest.raises(ValidationError, match=message): + AttemptManifest.model_validate_json(json.dumps(payload)) + + +def test_candidate_output_rejects_remaining_invalid_inventory_forms() -> None: + payload = _golden_payload("candidate_output.json") + payload["actual_records"] = 101 + with pytest.raises(ValidationError, match="must not exceed"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["outcome"] = "partial" + with pytest.raises(ValidationError, match="outcome must be"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + duplicate_file = dict(payload["files"][0]) + duplicate_file["record_count"] = 0 + payload["files"].append(duplicate_file) + with pytest.raises(ValidationError, match="paths must be unique"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["files"] = [] + with pytest.raises(ValidationError, match="require at least one file"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + prototype = payload["files"][0] + payload["files"] = [ + prototype + | { + "relative_path": f"part-{index:05d}.parquet", + "record_count": prototype["record_count"] if index == 0 else 0, + } + for index in range(MAXIMUM_CANDIDATE_OUTPUT_FILES + 1) + ] + with pytest.raises(ValidationError, match=rf"at most {MAXIMUM_CANDIDATE_OUTPUT_FILES} items"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + +def test_collection_plan_rejects_duplicate_shards_and_winner_paths() -> None: + payload = _golden_payload("collection_plan.json") + duplicate_shard = json.loads(json.dumps(payload["planned_shards"][0])) + duplicate_shard["winner_manifest"]["path"] = "/workspace/runs/run-0001/shards/shard-00001/winner.json" + payload["planned_shards"].append(duplicate_shard) + with pytest.raises(ValidationError, match="shard IDs must be unique"): + CollectionPlan.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("collection_plan.json") + duplicate_winner = json.loads(json.dumps(payload["planned_shards"][0])) + duplicate_winner["shard_id"] = "shard-00001" + payload["planned_shards"].append(duplicate_winner) + with pytest.raises(ValidationError, match="winner manifest paths must be unique"): + CollectionPlan.model_validate_json(json.dumps(payload)) + + +def test_collection_plan_requires_absolute_python_executable() -> None: + payload = _golden_payload("collection_plan.json") + payload["python_executable"] = ".venv/bin/python" + + with pytest.raises(ValidationError, match="path must be absolute"): + CollectionPlan.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("attempt_state", "deployment_updates", "message"), + ( + ("ready", {"ready_backends": 2}, "must not exceed"), + ("ready", {"endpoint_publication": "failed"}, "failed endpoint publication"), + ( + "pending", + {"state": "pending", "ready_backends": 0, "endpoint_publication": "published"}, + "pending endpoint publication", + ), + ("starting", {"state": "starting"}, "must use the ready state"), + ("ready", {"endpoint_publication": "pending"}, "published endpoint"), + ( + "starting", + {"state": "stopped", "ready_backends": 0}, + "cannot contain failed or stopped", + ), + ("starting", {}, "every deployment ready"), + ( + "ready", + {"state": "starting", "ready_backends": 0, "endpoint_publication": "pending"}, + "every deployment to be ready", + ), + ("failed", {}, "at least one failed deployment"), + ), +) +def test_readiness_rejects_remaining_contradictory_states( + attempt_state: str, + deployment_updates: dict[str, object], + message: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = attempt_state + payload["deployments"][0].update(deployment_updates) + + with pytest.raises(ValidationError, match=message): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_duplicate_deployment_ids() -> None: + payload = _golden_payload("multi_node_readiness.json") + payload["deployments"][1]["deployment_id"] = payload["deployments"][0]["deployment_id"] + + with pytest.raises(ValidationError, match="deployment IDs must be unique"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_probe_observation_after_snapshot() -> None: + payload = _golden_payload("single_node_readiness.json") + payload["deployments"][0]["last_probe"]["observed_at"] = "2026-08-18T12:03:00Z" + + with pytest.raises(ValidationError, match="later than the readiness snapshot"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_scheduler_observation_requires_consistent_deadline() -> None: + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = None + with pytest.raises(ValidationError, match="requires a reconciliation deadline"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = "2026-08-18T12:05:01Z" + with pytest.raises(ValidationError, match="must not precede"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["state"] = "running" + with pytest.raises(ValidationError, match="only accounting lag"): + SchedulerObservation.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("accounting_lag.json") + payload["reconciliation_deadline"] = "2026-08-18T13:10:02+01:00" + with pytest.raises(ValidationError, match="must be in UTC"): + SchedulerObservation.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py new file mode 100644 index 000000000..20436a0b9 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -0,0 +1,1273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import stat +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import cast + +import pytest +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 +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.models import SlurmAccountingEntry, SlurmProcessExitCode, SlurmQueueEntry +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + DeploymentReadiness, + EffectiveAttemptState, + EffectiveRunState, + EndpointPublicationState, + GenerationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + RunManifest, + SchedulerIdentity, + SchedulerJobIdentity, + SchedulerObservation, + SchedulerObservationCollector, + SchedulerState, + ShardManifest, + ShardWinner, + SlurmStateError, + SlurmStateReconciler, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, +) + + +@dataclass(frozen=True, slots=True) +class _ReconciliationCase: + workspace: Path + plan: ResolvedSlurmRunPlan + run: RunManifest + shard: ShardManifest + attempt: AttemptManifest + writer: SlurmStateWriter + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class _StaticSchedulerClient: + queue: tuple[SlurmQueueEntry, ...] + accounting: tuple[SlurmAccountingEntry, ...] + + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + return self.queue + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + return self.accounting + + +def test_collector_prefers_terminal_accounting_for_array_and_collection_jobs() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + queue = ( + SlurmQueueEntry(job_identity=task, state=SchedulerState.RUNNING), + SlurmQueueEntry(job_identity=5101, state=SchedulerState.RUNNING), + ) + accounting = ( + _accounting(task, SchedulerState.NODE_FAILED), + _accounting(5101, SchedulerState.COMPLETED), + ) + observed_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + + observations = SchedulerObservationCollector(_StaticSchedulerClient(queue, accounting)).collect( + (task, 5101), + observed_at=observed_at, + ) + + assert tuple(observation.state for observation in observations) == ( + SchedulerState.NODE_FAILED, + SchedulerState.COMPLETED, + ) + assert tuple(observation.scheduler for observation in observations) == (task, 5101) + + +@pytest.mark.parametrize("transient_state", (SchedulerState.PREEMPTED, SchedulerState.REQUEUED)) +def test_collector_allows_transient_scheduler_states_to_advance( + transient_state: SchedulerState, +) -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + previous = SchedulerObservationCollector( + _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=task, state=transient_state),), + (_accounting(task, transient_state),), + ) + ).collect((task,), observed_at=first_time)[0] + + pending = SchedulerObservationCollector( + _StaticSchedulerClient((SlurmQueueEntry(job_identity=task, state=SchedulerState.PENDING),), ()) + ).collect((task,), observed_at=first_time + timedelta(minutes=1), previous={task: previous})[0] + completed = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time + timedelta(minutes=2), previous={task: pending})[0] + + assert previous.state is transient_state + assert pending.state is SchedulerState.PENDING + assert completed.state is SchedulerState.COMPLETED + + +def test_collector_keeps_active_queue_preemption_nonterminal() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + client = _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=task, state=SchedulerState.PREEMPTED),), + (), + ) + + first = SchedulerObservationCollector(client).collect((task,), observed_at=first_time)[0] + second = SchedulerObservationCollector(client).collect( + (task,), observed_at=first_time + timedelta(minutes=1), previous={task: first} + )[0] + later = SchedulerObservationCollector(client).collect( + (task,), observed_at=first_time + timedelta(minutes=6), previous={task: second} + )[0] + + assert first.reconciliation_deadline is None + assert second.reconciliation_deadline is None + assert later.state is SchedulerState.PREEMPTED + assert later.reconciliation_deadline is None + + +def test_collector_reads_accounting_when_queue_no_longer_knows_job() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + runner = FakeSlurmRunner() + runner.script_next( + "squeue", + FakeCommandResponse(stderr="slurm_load_jobs error: Invalid job id specified\n", returncode=1), + ) + runner.script_next("sacct", FakeCommandResponse(stdout="4101_0|COMPLETED|0:0\n")) + + observation = SchedulerObservationCollector(SlurmCommandClient(runner)).collect( + (task,), + observed_at=datetime(2026, 9, 2, 12, tzinfo=timezone.utc), + )[0] + + assert observation.state is SchedulerState.COMPLETED + assert tuple(call[0] for call in runner.calls) == ("squeue", "sacct") + + +def test_fresh_process_refresh_persists_one_fixed_accounting_lag_deadline( + 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) + first_time = case.created_at + timedelta(minutes=3) + + first = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + ) + second = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + timedelta(minutes=1) + ) + expired = SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=first_time + timedelta(minutes=6) + ) + + first_scheduler = first.shards[0].attempts[0].scheduler + second_scheduler = second.shards[0].attempts[0].scheduler + assert first.effective_state is EffectiveRunState.ACCOUNTING_LAG + assert second.effective_state is EffectiveRunState.ACCOUNTING_LAG + assert expired.effective_state is EffectiveRunState.UNKNOWN + assert first_scheduler is not None and second_scheduler is not None + assert first_scheduler.reconciliation_deadline == second_scheduler.reconciliation_deadline + scheduler_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/scheduler.json" + 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, + 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="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=case.created_at + timedelta(minutes=3)) + + attempt = status.shards[0].attempts[0] + assert attempt.scheduler is not None and attempt.scheduler.state is SchedulerState.FAILED + assert attempt.effective_state is EffectiveAttemptState.FAILED + assert attempt.generation_state is GenerationState.FAILED + assert status.effective_state is EffectiveRunState.FAILED + + +def test_refresh_ignores_uncommitted_partial_client_output_from_failed_attempt( + 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) + running = _copy_attempt( + case.attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.attempt.created_at + timedelta(minutes=2), + ) + case.writer.update_attempt(running) + failed = _copy_attempt( + running, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.FAILED, + updated_at=case.attempt.created_at + timedelta(minutes=3), + ) + case.writer.update_attempt(failed) + + attempt_dir = case.writer.run_root / f"shards/{failed.shard_id}/attempts/{failed.attempt_id}" + candidate_path = attempt_dir / "output-manifest.json" + dataset_path = attempt_dir / "dataset" + requested = case.plan.shards[case.shard.shard_index].requested_records + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=failed.shard_id, + attempt_id=failed.attempt_id, + attempt_ordinal=failed.attempt_ordinal, + created_at=failed.updated_at, + dataset_path=dataset_path.as_posix(), + requested_records=requested, + actual_records=requested - 1, + outcome=CandidateOutcome.PARTIAL, + files=( + CandidateOutputFile( + relative_path="part-00000.parquet", + sha256="a" * 64, + byte_size=1, + record_count=requested - 1, + ), + ), + dataset_schema_digest="b" * 64, + provenance_digest=case.plan.compute_sha256(), + ) + result = ClientResult( + schema_version=1, + run_id=case.plan.run_id, + shard_id=failed.shard_id, + attempt_id=failed.attempt_id, + completed_at=failed.updated_at, + requested_records=requested, + actual_records=requested - 1, + outcome=ClientOutcome.PARTIAL, + dataset_path=dataset_path.as_posix(), + early_shutdown=False, + requested_resume_mode=case.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=ArtifactReference( + path=candidate_path.as_posix(), + sha256=candidate.compute_sha256(), + ), + ) + candidate_path.write_text(candidate.serialize_json()) + (attempt_dir / "client-result.json").write_text(result.serialize_json()) + + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, failed.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state="FAILED", exit_code="70:0") + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=case.created_at + timedelta(minutes=5)) + + attempt = status.shards[0].attempts[0] + assert attempt.effective_state is EffectiveAttemptState.FAILED + assert attempt.generation_state is GenerationState.FAILED + assert attempt.client_result is None + assert attempt.candidate_output is None + + +def test_refresh_marks_unrequeued_preemption_as_failure_after_fixed_deadline( + 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="PREEMPTED", + exit_code="0:15", + ) + first_time = case.created_at + timedelta(minutes=3) + + first = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state=None) + second = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=1)) + expired = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=6)) + + first_attempt = first.shards[0].attempts[0] + second_attempt = second.shards[0].attempts[0] + attempt = expired.shards[0].attempts[0] + assert first_attempt.effective_state is EffectiveAttemptState.PENDING + assert second_attempt.effective_state is EffectiveAttemptState.PENDING + assert first_attempt.scheduler is not None and second_attempt.scheduler is not None + assert first_attempt.scheduler.reconciliation_deadline == second_attempt.scheduler.reconciliation_deadline + assert attempt.scheduler is not None and attempt.scheduler.state is SchedulerState.FAILED + assert attempt.effective_state is EffectiveAttemptState.FAILED + assert attempt.generation_state is GenerationState.FAILED + assert expired.effective_state is EffectiveRunState.FAILED + + +def test_refresh_allows_accounting_preemption_to_reappear_in_queue( + 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="PREEMPTED", + exit_code="0:15", + ) + first_time = case.created_at + timedelta(minutes=3) + first = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time) + + fake_slurm_runner.set_task_state( + scheduler, + queue_state="PREEMPTED", + accounting_state="PREEMPTED", + exit_code="0:15", + ) + still_preempted = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(seconds=30)) + + fake_slurm_runner.set_task_state( + scheduler, + queue_state="PENDING", + accounting_state="PREEMPTED", + exit_code="0:15", + ) + requeued = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=first_time + timedelta(minutes=1)) + + first_attempt = first.shards[0].attempts[0] + still_preempted_attempt = still_preempted.shards[0].attempts[0] + requeued_attempt = requeued.shards[0].attempts[0] + assert first_attempt.scheduler is not None and still_preempted_attempt.scheduler is not None + assert first_attempt.scheduler.state is SchedulerState.PREEMPTED + assert first_attempt.effective_state is EffectiveAttemptState.PENDING + assert still_preempted_attempt.scheduler.reconciliation_deadline is None + assert requeued_attempt.scheduler is not None + assert requeued_attempt.scheduler.state is SchedulerState.PENDING + assert requeued_attempt.effective_state is EffectiveAttemptState.PENDING + assert requeued.effective_state is EffectiveRunState.PENDING + + +def test_refresh_rejects_winner_that_conflicts_with_terminal_scheduler_evidence( + 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") + completed, winner = _publish_winner_state(case) + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state="RUNNING", + accounting_state="NODE_FAIL", + exit_code="1:0", + ) + + with pytest.raises(StateCorruptionError, match="winner conflicts"): + SlurmStateReconciler( + case.workspace, + case.plan.run_id, + 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, + 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") + completed, winner = _publish_winner_state(case) + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + exit_code="0:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + + attempt = status.shards[0].attempts[0] + assert attempt.client_result is not None + assert attempt.candidate_output is not None + assert attempt.effective_state is EffectiveAttemptState.SUCCEEDED + assert attempt.generation_state is GenerationState.WON + assert status.shards[0].winner == winner + assert status.effective_state is EffectiveRunState.SUCCEEDED + + +def test_refresh_does_not_succeed_a_winnerless_persisted_candidate( + 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") + completed, winner = _publish_winner_state(case) + (case.writer.run_root / "shards/shard-00000/winner.json").unlink() + scheduler = cast(SchedulerIdentity, completed.scheduler) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + exit_code="0:0", + ) + + status = SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + + attempt = status.shards[0].attempts[0] + assert attempt.effective_state is EffectiveAttemptState.SUCCEEDED + assert attempt.generation_state is GenerationState.CANDIDATE_READY + assert status.shards[0].winner is None + assert status.shards[0].effective_state is EffectiveAttemptState.UNKNOWN + assert status.effective_state is EffectiveRunState.UNKNOWN + + +def test_refresh_rejects_a_concurrent_attempt_change_instead_of_guessing_status( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + observed_at = case.created_at + timedelta(minutes=4) + + class MutatingClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + return (SlurmQueueEntry(job_identity=scheduler, state=SchedulerState.RUNNING),) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + case.writer.update_attempt( + _copy_attempt(case.attempt, state=AttemptLifecycleState.RUNNING, updated_at=observed_at) + ) + return () + + with pytest.raises(StateConflictError, match="changed during reconciliation"): + SlurmStateReconciler(case.workspace, case.plan.run_id, MutatingClient()).refresh(observed_at=observed_at) + + +def test_refresh_accepts_readiness_published_during_scheduler_query( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + observed_at = case.created_at + timedelta(minutes=4) + ready_at = observed_at + timedelta(seconds=1) + deployment = case.plan.deployments[0] + + class MutatingClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + return (SlurmQueueEntry(job_identity=scheduler, state=SchedulerState.RUNNING),) + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + 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, + ), + ), + ) + ) + return () + + status = SlurmStateReconciler(case.workspace, case.plan.run_id, MutatingClient()).refresh(observed_at=observed_at) + + assert status.observed_at == ready_at + assert status.shards[0].attempts[0].effective_state is EffectiveAttemptState.RUNNING + + +def test_observe_does_not_create_a_missing_shard_lock_or_persist_scheduler_evidence( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + shard_root = case.writer.run_root / "shards/shard-00000" + (shard_root / "shard.lock").unlink() + client = _StaticSchedulerClient( + (SlurmQueueEntry(job_identity=scheduler, state=SchedulerState.RUNNING),), + (), + ) + + status = SlurmStateReconciler(case.workspace, case.plan.run_id, client).observe( + observed_at=case.created_at + timedelta(minutes=3) + ) + + assert status.shards[0].attempts[0].effective_state is EffectiveAttemptState.RUNNING + assert not (shard_root / "shard.lock").exists() + assert not (shard_root / "attempts/attempt-0001/scheduler.json").exists() + + +def test_terminal_observation_remains_authoritative_during_later_accounting_gap() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + first = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time)[0] + later = SchedulerObservationCollector( + _StaticSchedulerClient((SlurmQueueEntry(job_identity=task, state=SchedulerState.RUNNING),), ()) + ).collect((task,), observed_at=first_time + timedelta(minutes=1), previous={task: first})[0] + + assert later.state is SchedulerState.COMPLETED + + +def test_collector_rejects_conflicting_terminal_accounting_evidence() -> None: + task = SchedulerIdentity(array_job_id=4101, array_task_id=0) + first_time = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + previous = SchedulerObservationCollector( + _StaticSchedulerClient((), (_accounting(task, SchedulerState.COMPLETED),)) + ).collect((task,), observed_at=first_time)[0] + + with pytest.raises(SlurmStateError, match="violates persisted chronology"): + SchedulerObservationCollector(_StaticSchedulerClient((), (_accounting(task, SchedulerState.FAILED),))).collect( + (task,), observed_at=first_time + timedelta(minutes=1), previous={task: previous} + ) + + +def test_refresh_normalizes_scheduler_query_failures( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + calls: list[str] = [] + + class FailingClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + del selectors + calls.append("queue") + raise RuntimeError("scheduler unavailable") + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + del selectors + calls.append("accounting") + return () + + with pytest.raises(SlurmStateError, match="cannot query normalized scheduler observations"): + SlurmStateReconciler(case.workspace, case.plan.run_id, FailingClient()).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + assert calls == ["queue", "accounting"] + + +def test_refresh_keeps_an_unsubmitted_attempt_pending_without_querying_slurm( + 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) + + class UnexpectedClient: + def query_queue(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmQueueEntry, ...]: + raise AssertionError(f"unexpected queue query for {selectors!r}") + + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: + raise AssertionError(f"unexpected accounting query for {selectors!r}") + + status = SlurmStateReconciler(case.workspace, case.plan.run_id, UnexpectedClient()).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + attempt = status.shards[0].attempts[0] + assert attempt.scheduler is None + assert attempt.effective_state is EffectiveAttemptState.PENDING + assert attempt.generation_state is GenerationState.NOT_STARTED + assert status.effective_state is EffectiveRunState.PENDING + + +def test_fresh_process_refresh_rejects_mismatched_persisted_scheduler_evidence( + 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="RUNNING", accounting_state=None) + observed_at = case.created_at + timedelta(minutes=3) + SlurmStateReconciler(case.workspace, case.plan.run_id, SlurmCommandClient(fake_slurm_runner)).refresh( + observed_at=observed_at + ) + scheduler_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/scheduler.json" + mismatched = SchedulerObservation( + schema_version=1, + scheduler=SchedulerIdentity(array_job_id=9999, array_task_id=0), + observed_at=observed_at, + state=SchedulerState.RUNNING, + ) + scheduler_path.write_text(mismatched.serialize_json()) + + with pytest.raises(StateCorruptionError, match="mismatched scheduler evidence"): + SlurmStateReconciler( + case.workspace, + case.plan.run_id, + SlurmCommandClient(fake_slurm_runner), + ).refresh(observed_at=observed_at + timedelta(minutes=1)) + + +def _accounting(identity: SchedulerJobIdentity, state: SchedulerState) -> SlurmAccountingEntry: + return SlurmAccountingEntry( + job_identity=identity, + state=state, + process_exit_code=SlurmProcessExitCode(exit_status=0, termination_signal=0), + ) + + +def _initialized_case( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + *, + 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) + created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) + run_root = workspace / "runs" / relocated_plan.run_id + run = RunManifest( + schema_version=1, + run_id=relocated_plan.run_id, + created_at=created_at, + authored_config=relocated_plan.authored_config, + resolved_plan=ArtifactReference( + path=(run_root / "resolved-plan.json").as_posix(), + sha256=relocated_plan.compute_sha256(), + ), + shard_count=len(relocated_plan.shards), + ) + 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, 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) + ) + + +def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: + previous_workspace = plan.selected_profile.profile.workspace_root + payload = cast( + dict[str, object], + json.loads(plan.serialize_json().replace(previous_workspace, workspace.as_posix())), + ) + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile = SlurmProfile.model_validate_json(json.dumps(selected_profile["profile"])) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +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( + attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=attempt.created_at + timedelta(minutes=2), + ) + case.writer.update_attempt(running) + 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 = 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=( + 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, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path="part-00000.parquet", + sha256="a" * 64, + byte_size=1, + record_count=requested, + ), + ), + dataset_schema_digest="b" * 64, + provenance_digest=case.plan.compute_sha256(), + ) + candidate_reference = ArtifactReference(path=candidate_path.as_posix(), sha256=candidate.compute_sha256()) + result = ClientResult( + schema_version=1, + run_id=case.plan.run_id, + shard_id=running.shard_id, + attempt_id=running.attempt_id, + completed_at=result_completed_at, + requested_records=requested, + actual_records=requested, + outcome=ClientOutcome.COMPLETE, + dataset_path=dataset_path.as_posix(), + early_shutdown=False, + requested_resume_mode=case.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=candidate_reference, + ) + case.writer.publish_attempt_result(result, candidate) + completed = _copy_attempt( + running, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + candidate_output=candidate_reference, + updated_at=attempt_stopped_at, + ) + case.writer.update_attempt(completed) + winner = ShardWinner( + schema_version=1, + run_id=case.plan.run_id, + shard_id=completed.shard_id, + attempt_id=completed.attempt_id, + attempt_ordinal=completed.attempt_ordinal, + candidate_manifest=candidate_reference, + published_at=attempt_stopped_at + timedelta(minutes=1), + ) + 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) + return AttemptManifest.model_validate_json(json.dumps(payload, default=_json_value)) + + +def _json_value(value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + raise TypeError(f"unsupported test value: {type(value).__name__}") diff --git a/packages/data-designer-slurm/tests/state/test_records.py b/packages/data-designer-slurm/tests/state/test_records.py new file mode 100644 index 000000000..b1c19e851 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_records.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.state import ( + ArtifactReference, + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + RecordRange, + RunManifest, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" + + +def _golden_payload(filename: str) -> dict[str, Any]: + return json.loads((GOLDEN_DIRECTORY / filename).read_text()) + + +def test_direct_construction_and_json_loading_produce_identical_record() -> None: + record = RunManifest( + schema_version=1, + run_id="run-direct", + created_at=datetime(2026, 8, 18, 12, tzinfo=timezone.utc), + authored_config=ArtifactReference(path="/workspace/run/authored.json", sha256="a" * 64), + resolved_plan=ArtifactReference(path="/workspace/run/plan.json", sha256="b" * 64), + shard_count=2, + ) + + assert RunManifest.model_validate_json(record.serialize_json()) == record + assert record.serialize_json() == RunManifest.model_validate_json(record.serialize_json()).serialize_json() + + +@pytest.mark.parametrize( + ("field_name", "invalid_value"), + ( + ("schema_version", 2), + ("run_id", "contains/slash"), + ("run_id", " leading-space"), + ), +) +def test_run_manifest_rejects_invalid_version_and_identity( + field_name: str, + invalid_value: object, +) -> None: + payload = _golden_payload("run_manifest.json") + payload[field_name] = invalid_value + + with pytest.raises(ValidationError): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_records_reject_unknown_fields() -> None: + payload = _golden_payload("run_manifest.json") + payload["unreviewed"] = True + + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_records_require_explicit_schema_version() -> None: + payload = _golden_payload("run_manifest.json") + payload.pop("schema_version") + + with pytest.raises(ValidationError, match="Field required"): + RunManifest.model_validate_json(json.dumps(payload)) + + +def test_artifact_reference_rejects_invalid_digest_and_path() -> None: + with pytest.raises(ValidationError): + ArtifactReference(path="/workspace/plan.json", sha256="A" * 64) + with pytest.raises(ValidationError, match="parent-directory"): + ArtifactReference(path="/workspace/../private/plan.json", sha256="a" * 64) + with pytest.raises(ValidationError, match="absolute"): + ArtifactReference(path="relative/plan.json", sha256="a" * 64) + + +def test_records_are_frozen() -> None: + record = RunManifest.model_validate_json((GOLDEN_DIRECTORY / "run_manifest.json").read_text()) + + with pytest.raises(ValidationError, match="Instance is frozen"): + record.run_id = "another-run" + + +def test_record_range_is_half_open_and_non_empty() -> None: + record_range = RecordRange(start_index=10, end_index_exclusive=15) + + assert record_range.record_count == 5 + with pytest.raises(ValidationError, match="greater than start_index"): + RecordRange(start_index=10, end_index_exclusive=10) + + +@pytest.mark.parametrize( + ("filename", "mutation"), + ( + ( + "successful_attempt.json", + {"candidate_output": None}, + ), + ( + "successful_attempt.json", + {"state": "running", "terminal_classification": "succeeded"}, + ), + ( + "successful_attempt.json", + {"state": "running", "terminal_classification": None, "scheduler": None}, + ), + ), +) +def test_attempt_lifecycle_invariants(filename: str, mutation: dict[str, object]) -> None: + payload = _golden_payload(filename) + payload.update(mutation) + + with pytest.raises(ValidationError): + AttemptManifest.model_validate_json(json.dumps(payload)) + + +def test_ready_deployment_requires_all_backends_and_published_endpoint() -> None: + payload = _golden_payload("single_node_readiness.json") + deployment = payload["deployments"][0] + deployment["ready_backends"] = 0 + + with pytest.raises(ValidationError, match="every expected backend"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_rejects_duplicate_deployment_identity() -> None: + payload = _golden_payload("multi_node_readiness.json") + payload["deployments"][1]["model_alias"] = payload["deployments"][0]["model_alias"] + + with pytest.raises(ValidationError, match="model aliases must be unique"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_readiness_preserves_the_public_model_alias_contract() -> None: + payload = _golden_payload("single_node_readiness.json") + model_alias = "model/模型 alias " + "x" * 140 + payload["deployments"][0]["model_alias"] = model_alias + + readiness = AttemptReadiness.model_validate_json(json.dumps(payload)) + + assert readiness.deployments[0].model_alias == model_alias + + +def test_probe_evidence_is_bounded_and_single_line() -> None: + payload = _golden_payload("single_node_readiness.json") + payload["deployments"][0]["last_probe"]["redacted_message"] = "line one\nline two" + + with pytest.raises(ValidationError, match="control characters"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + payload["deployments"][0]["last_probe"]["redacted_message"] = "x" * 513 + with pytest.raises(ValidationError): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("attempt_state", "deployment_state"), + ( + ("stopped", "ready"), + ("pending", "ready"), + ("starting", "stopped"), + ), +) +def test_attempt_readiness_rejects_contradictory_deployment_states( + attempt_state: str, + deployment_state: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = attempt_state + payload["deployments"][0]["state"] = deployment_state + + with pytest.raises(ValidationError): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +@pytest.mark.parametrize( + ("state", "endpoint_publication"), + (("pending", "pending"), ("stopped", "published")), +) +def test_pending_and_stopped_deployments_reject_ready_backends( + state: str, + endpoint_publication: str, +) -> None: + payload = _golden_payload("single_node_readiness.json") + payload["state"] = state + payload["deployments"][0].update( + { + "state": state, + "endpoint_publication": endpoint_publication, + "ready_backends": 1, + } + ) + + with pytest.raises(ValidationError, match="cannot have ready backends"): + AttemptReadiness.model_validate_json(json.dumps(payload)) + + +def test_nonempty_partial_candidate_is_winner_eligible() -> None: + payload = _golden_payload("candidate_output.json") + payload["actual_records"] = 99 + payload["outcome"] = "partial" + payload["files"][0]["record_count"] = 99 + + partial_candidate = CandidateOutputManifest.model_validate_json(json.dumps(payload)) + assert partial_candidate.winner_eligible + + +def test_candidate_output_rejects_count_and_file_mismatches() -> None: + payload = _golden_payload("candidate_output.json") + payload["files"][0]["record_count"] = 99 + + with pytest.raises(ValidationError, match="record counts"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) + + payload = _golden_payload("candidate_output.json") + payload["files"][0]["relative_path"] = "../escaped.parquet" + with pytest.raises(ValidationError, match="parent-directory"): + CandidateOutputManifest.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py new file mode 100644 index 000000000..555ed2488 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -0,0 +1,2162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import errno +import hashlib +import json +import os +import shutil +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from threading import Event +from typing import cast +from unittest.mock import patch + +import pytest +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask + +import data_designer.lazy_heavy_imports as lazy +import data_designer.slurm.state.collection_filesystem as collection_filesystem +import data_designer.slurm.state.collection_merge as collection_merge +import data_designer.slurm.state.collection_storage as collection_storage_module +import data_designer.slurm.state.collection_worker as collection_worker_module +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 +from data_designer.slurm.launcher.client import SlurmCommandClient +from data_designer.slurm.launcher.errors import SlurmSubmissionError +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + CollectionState, + RetryState, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateError, + SlurmStateReconciler, + SlurmStateWriter, + StateConflictError, + StateContractError, + StateCorruptionError, + StateNotFoundError, + compute_candidate_schema_digest, +) +from data_designer.slurm.state.attempt_identity import require_attempt_scheduler_identity +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +from data_designer.slurm.state.collection_schema import schemas_match_except_struct_order +from data_designer.slurm.state.collection_storage import CollectionStorage +from data_designer.slurm.state.collection_validation import validate_collection_inputs +from data_designer.slurm.state.collection_worker import SlurmCollectionWorker +from data_designer.slurm.state.destinations import CollectionDestinationResolver +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.retry_storage import RetryStorage +from data_designer.slurm.state.storage import StateStorage + + +@dataclass(frozen=True, slots=True) +class _RunCase: + workspace: Path + plan: ResolvedSlurmRunPlan + run: RunManifest + shards: tuple[ShardManifest, ...] + writer: SlurmStateWriter + created_at: datetime + + +def test_retry_refreshes_failed_shard_and_publishes_exact_next_attempt( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert len(attempts) == 1 + assert attempts[0].attempt_id == "attempt-0002" + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + retry_plan = case.writer.load_retry_plan("retry-0001") + assert retry_plan.effective_resume_mode == "never" + assert retry_plan.planned_shards[0].attempt_id == attempts[0].attempt_id + mapped_workspace = tmp_path / "mapped-workspace" + shutil.copytree(case.workspace / "runs", mapped_workspace / "runs") + mapped_writer = SlurmStateWriter( + mapped_workspace, + case.plan.run_id, + logical_workspace_root=case.workspace, + ) + assert mapped_writer.load_retry_plan("retry-0001") == retry_plan + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=4201, array_task_id=0), + ) + with pytest.raises(StateConflictError, match="scheduler identity"): + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=9999, array_task_id=0), + ) + assert case.writer.load_attempts(case.shards[0].shard_id) == (attempt, attempts[0]) + assert "#SBATCH --array=0" in cast(str, runner.inputs[-1]) + assert 'DD_ATTEMPT_ORDINAL="0002"' in cast(str, runner.inputs[-1]) + assert ( + coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + == attempts + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +def test_active_retry_preview_ignores_statusless_journal_tail( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + storage.ensure_retry("retry-0001") + + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + assert coordinator.preview_active(observed_at=case.created_at + timedelta(minutes=5)) is None + preview, _ = coordinator.preview( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert preview.retry_id == "retry-0001" + assert storage.get_retry_root("retry-0001").is_dir() + + +def test_retry_rejects_a_nonterminal_explicit_shard( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case.writer.create_attempt( + _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + ) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + +def test_retry_rejects_a_shard_with_an_immutable_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + _publish_all_winners(case) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=7), + ) + + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_retry_accepts_unknown_after_bounded_accounting_lag( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state=None) + SlurmStateReconciler(case.workspace, case.plan.run_id, scheduler).refresh( + observed_at=case.created_at + timedelta(minutes=3) + ) + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=9), + ) + + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + + +def test_retry_submits_only_the_failed_sparse_array_task( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index)) for shard in case.shards + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=1)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard in case.shards: + case.writer.create_attempt( + _submitted_attempt( + case, + shard, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index), + ) + ) + runner.set_task_state(initial_tasks[0].scheduler, queue_state="RUNNING", accounting_state=None) + runner.set_task_state(initial_tasks[1].scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + assert tuple(attempt.shard_id for attempt in attempts) == ("shard-00001",) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=1) + assert "#SBATCH --array=1%2" in cast(str, runner.inputs[-1]) + assert 'case "${DD_ARRAY_TASK_ID}"' in cast(str, runner.inputs[-1]) + + +def test_retry_ambiguous_submission_waits_for_scheduler_visibility_without_a_duplicate( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + submissions = 0 + + def ambiguous_submit(script: str) -> object: + nonlocal submissions + del script + submissions += 1 + raise SlurmSubmissionError("sbatch could not be executed: command timed out", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", ambiguous_submit) + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit retry"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + + retry_storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert retry_storage.read_status("retry-0001").state is RetryState.PREPARED + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + with pytest.raises(StateConflictError, match="still being reconciled"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + assert submissions == 1 + + +def test_retry_recovers_an_accepted_submission_after_the_receipt_is_lost( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=(initial,)), FakeSlurmArray(tasks=(retried,)))) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit retry"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + retry_plan = storage.read_plan("retry-0001") + runner.script_next( + "squeue", + FakeCommandResponse(stdout=f"4201_0|{retry_plan.submission_job_name}\n"), + ) + runner.script_next("sacct", FakeCommandResponse()) + + recovered = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert recovered[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + assert case.writer.load_attempts(case.shards[0].shard_id) == (first_attempt, recovered[0]) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert ( + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=7), + ) + == recovered + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +@pytest.mark.parametrize( + ("accepted_before_receipt_loss", "replacement_job_id", "submission_count"), + [ + pytest.param(False, 4201, 2, id="unaccepted"), + pytest.param(True, 4301, 3, id="accepted-but-still-invisible"), + ], +) +def test_retry_replaces_or_fences_an_ambiguous_submission_after_its_deadline( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, + accepted_before_receipt_loss: bool, + replacement_job_id: int, + submission_count: int, +) -> None: + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + hidden = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + replacement = FakeSlurmTask(SchedulerIdentity(array_job_id=replacement_job_id, array_task_id=0)) + retry_arrays = (FakeSlurmArray(tasks=(hidden,)),) if accepted_before_receipt_loss else () + runner = FakeSlurmRunner( + arrays=(FakeSlurmArray(tasks=(initial,)), *retry_arrays, FakeSlurmArray(tasks=(replacement,))) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + original_submit = scheduler.submit_script + + def lose_submission_receipt(script: str) -> object: + if accepted_before_receipt_loss: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", lose_submission_receipt) + with pytest.raises(SlurmStateError, match="cannot submit retry"): + SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + + attempts = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=11), + ) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.FAILED + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=replacement_job_id, array_task_id=0) + with pytest.raises(StateConflictError, match="scheduler identity"): + require_attempt_scheduler_identity( + case.workspace, + case.plan.run_id, + attempts[0].shard_id, + attempts[0].attempt_id, + SchedulerIdentity(array_job_id=4201 if accepted_before_receipt_loss else 4301, array_task_id=0), + ) + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == submission_count + + +def test_retry_definite_submission_failure_settles_and_can_be_retried( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + runner.script_next("sbatch", FakeCommandResponse(stderr="submission rejected", returncode=2)) + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit retry"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.FAILED + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4201, array_task_id=0) + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + + +def test_concurrent_retry_requests_converge_on_one_array( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + + def retry() -> tuple[AttemptManifest, ...]: + return SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _: retry(), range(2))) + + assert results[0] == results[1] + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +def test_default_retry_does_not_hide_failures_outside_an_active_explicit_subset( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4301, array_task_id=1)),)), + ) + ) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + + first = coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + second = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert tuple(attempt.shard_id for attempt in first) == (case.shards[0].shard_id,) + assert tuple(attempt.shard_id for attempt in second) == (case.shards[1].shard_id,) + assert second[0].scheduler == SchedulerIdentity(array_job_id=4301, array_task_id=1) + + +def test_retry_discards_trailing_journal_interrupted_before_submission( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + ) + ) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + attempt = _submitted_attempt(case, case.shards[0], scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0)) + case.writer.create_attempt(attempt) + runner.set_task_state(attempt.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_publish = RetryStorage.publish_status + + def interrupt_before_status(self: RetryStorage, status: object) -> None: + del self, status + raise OSError("injected journal interruption") + + monkeypatch.setattr(RetryStorage, "publish_status", interrupt_before_status) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + monkeypatch.setattr(RetryStorage, "publish_status", original_publish) + + attempts = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert attempts[0].attempt_id == "attempt-0002" + assert RetryStorage(StateStorage(case.workspace, case.plan.run_id)).list_retry_ids() == ("retry-0001",) + + +def test_retry_recovers_evolved_attempt_and_its_exact_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run_single, single_node_plan) + initial = FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=(initial,)), FakeSlurmArray(tasks=(retried,)))) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + first_attempt = _submitted_attempt(case, case.shards[0], scheduler=initial.scheduler) + case.writer.create_attempt(first_attempt) + runner.set_task_state(initial.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry(effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=5)) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + retry_attempt = case.writer.load_attempts(case.shards[0].shard_id)[-1] + running = _copy_attempt( + retry_attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=6), + ) + case.writer.update_attempt(running) + with case.writer.acquire_dataset_workspace(case.shards[0].shard_id, running.attempt_id, "never") as dataset_path: + _publish_candidate(case, case.shards[0], running, dataset_path) + succeeded = case.writer.load_attempts(case.shards[0].shard_id)[-1] + case.writer.finalize_winner( + case.shards[0].shard_id, + succeeded.attempt_id, + published_at=case.created_at + timedelta(minutes=10), + ) + runner.set_task_state(retried.scheduler, queue_state=None, accounting_state="COMPLETED") + + recovered = coordinator.retry( + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=11), + ) + + assert recovered == (succeeded,) + assert ( + RetryStorage(StateStorage(case.workspace, case.plan.run_id)).read_status("retry-0001").state + is RetryState.COMPLETED + ) + + +def test_retry_settles_submitted_journal_before_serving_a_disjoint_selection( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=initial_tasks), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4301, array_task_id=1)),)), + ) + ) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + + attempts = coordinator.retry( + shard_ids=(case.shards[1].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=6), + ) + + assert tuple(attempt.shard_id for attempt in attempts) == (case.shards[1].shard_id,) + assert attempts[0].scheduler == SchedulerIdentity(array_job_id=4301, array_task_id=1) + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert storage.read_status("retry-0002").state is RetryState.COMPLETED + + +def test_retry_does_not_return_recovered_attempt_for_a_different_resume_mode( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + initial_tasks = tuple( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=index)) for index in range(2) + ) + retried = FakeSlurmTask(SchedulerIdentity(array_job_id=4201, array_task_id=0)) + runner = FakeSlurmRunner(arrays=(FakeSlurmArray(tasks=initial_tasks), FakeSlurmArray(tasks=(retried,)))) + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + scheduler = SlurmCommandClient(runner) + scheduler.submit_script("initial") + for shard, task in zip(case.shards, initial_tasks, strict=True): + case.writer.create_attempt(_submitted_attempt(case, shard, scheduler=task.scheduler)) + runner.set_task_state(task.scheduler, queue_state=None, accounting_state="FAILED", exit_code="1:0") + coordinator = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler) + original_complete = SlurmRetryCoordinator._complete_retry + + def interrupt_after_attempts(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected completion interruption") + + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", interrupt_after_attempts) + with pytest.raises(SlurmStateError, match="cannot retry persisted run"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="never", + observed_at=case.created_at + timedelta(minutes=5), + ) + monkeypatch.setattr(SlurmRetryCoordinator, "_complete_retry", original_complete) + + with pytest.raises(StateConflictError, match="sealed or nonterminal"): + coordinator.retry( + shard_ids=(case.shards[0].shard_id,), + effective_resume_mode="always", + observed_at=case.created_at + timedelta(minutes=6), + ) + + storage = RetryStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("retry-0001").state is RetryState.COMPLETED + assert storage.list_retry_ids() == ("retry-0001",) + + +def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + scheduler = SlurmCommandClient(runner) + coordinator = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + scheduler, + {"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"}, + ) + + with patch.object(scheduler, "submit_script", wraps=scheduler.submit_script) as submit_script: + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert submitted.state is CollectionState.SUBMITTED + assert submitted.scheduler == 5101 + assert submit_script.call_args.kwargs["export_environment"] == {"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"} + with pytest.raises(StateConflictError, match="requested destination"): + coordinator.submit(destination=_output_root(case.plan).parent / "other") + script = cast(str, runner.inputs[-1]) + collection_plan = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_plan( + submitted.collection_id + ) + expected_python = Path(sys.executable).absolute().as_posix() + assert collection_plan.python_executable == expected_python + assert f'readonly DD_COLLECTION_PYTHON="{expected_python}"' in script + assert collection_plan.package_version != case.plan.package_version + assert f'readonly DD_PACKAGE_VERSION="{collection_plan.package_version}"' in script + assert collection_plan.runtime_bundle != case.plan.runtime_bundle + assert Path(collection_plan.runtime_bundle.path).is_file() + assert f'readonly DD_RUNTIME_ARCHIVE="{collection_plan.runtime_bundle.path}"' in script + assert f'readonly DD_RUNTIME_SHA256="{collection_plan.runtime_bundle.sha256}"' in script + assert "data_designer.slurm.state.collection_worker" in script + assert "enroot start" not in script + assert "DD_CLIENT_IMAGE" not in script + assert "--gpus" not in script + assert "#SBATCH --gres" not in script + stale_stage = _output_root(case.plan).parent / submitted.staging_directory + stale_stage.mkdir(mode=0o700) + (stale_stage / "partial").write_text("incomplete") + unrelated_stage = _output_root(case.plan).parent / f".dd-collection-{'f' * 32}.tmp" + unrelated_stage.mkdir(mode=0o700) + (unrelated_stage / "active").write_text("preserve") + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = _output_root(case.plan) + assert result.actual_records == case.plan.invocation.authored.num_records + assert len(result.files) == 1 + assert lazy.pq.read_table(destination / result.files[0].relative_path).column("record_id").to_pylist() == list( + range(case.plan.invocation.authored.num_records) + ) + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.SUCCEEDED + assert coordinator.submit() == persisted + assert not stale_stage.exists() + assert (unrelated_stage / "active").read_text() == "preserve" + + +def test_collection_merges_nonempty_partial_shard_winners( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case, shortfall=9) + coordinator = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(FakeSlurmRunner(jobs=(FakeSlurmJob(5101),))), + ) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + expected = sum(shard.record_range.record_count - 9 for shard in case.shards) + assert result.actual_records == expected + assert lazy.pq.read_metadata(_output_root(case.plan) / result.files[0].relative_path).num_rows == expected + assert coordinator.submit().state is CollectionState.SUCCEEDED + + +def test_collection_merges_reordered_struct_fields_from_partial_winners( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case, shortfall=9, alternating_struct_order=True) + coordinator = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(FakeSlurmRunner(jobs=(FakeSlurmJob(5101),))), + ) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + expected = sum(shard.record_range.record_count - 9 for shard in case.shards) + table = lazy.pq.read_table(_output_root(case.plan) / result.files[0].relative_path) + assert result.actual_records == expected + assert table.num_rows == expected + assert table.schema.field("editorial_assessment").type == lazy.pa.struct( + [lazy.pa.field("needs_revision", lazy.pa.bool_()), lazy.pa.field("visible_item", lazy.pa.string())] + ) + rows = table.select(["record_id", "editorial_assessment"]).to_pylist() + assert all(row["editorial_assessment"]["visible_item"] == f"item-{row['record_id']}" for row in rows) + + +def test_collection_schema_reordering_does_not_hide_type_or_field_changes() -> None: + field = lazy.pa.field + reference = lazy.pa.schema( + [field("assessment", lazy.pa.struct([field("flag", lazy.pa.bool_()), field("label", lazy.pa.string())]))] + ) + reordered = lazy.pa.schema( + [field("assessment", lazy.pa.struct([field("label", lazy.pa.string()), field("flag", lazy.pa.bool_())]))] + ) + changed_type = lazy.pa.schema( + [field("assessment", lazy.pa.struct([field("flag", lazy.pa.string()), field("label", lazy.pa.string())]))] + ) + added_field = lazy.pa.schema( + [ + field( + "assessment", + lazy.pa.struct( + [field("flag", lazy.pa.bool_()), field("label", lazy.pa.string()), field("extra", lazy.pa.string())] + ), + ) + ] + ) + + assert schemas_match_except_struct_order(reference, reordered) + assert not schemas_match_except_struct_order(reference, changed_type) + assert not schemas_match_except_struct_order(reference, added_field) + + +def test_collection_requires_stable_plan_provenance_across_distinct_shards( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + state = StateStorage(case.workspace, case.plan.run_id) + collection_plan = CollectionStorage(state).read_plan(submitted.collection_id) + _, candidates = CollectionInputResolver(state, StateReader(state, case.plan.run_id)).resolve(collection_plan) + + assert len({candidate.shard_id for candidate in candidates}) == len(candidates) + assert len({candidate.files[0].sha256 for candidate in candidates}) == len(candidates) + assert {candidate.provenance_digest for candidate in candidates} == {case.plan.compute_sha256()} + drifted = tuple(candidate.model_copy(update={"provenance_digest": "f" * 64}) for candidate in candidates) + with pytest.raises(StateContractError, match="does not match the resolved plan"): + validate_collection_inputs(case.plan, collection_plan, drifted) + + +def test_collection_refresh_observes_running_during_bulk_merge( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + scheduler = SlurmCommandClient(runner) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + runner.set_job_state(5101, queue_state="RUNNING", accounting_state=None) + merge_started = Event() + release_merge = Event() + original_merge = collection_merge.CollectionMerger.merge + + def pause_merge(*args: object, **kwargs: object) -> object: + merge_started.set() + assert release_merge.wait(timeout=3) + return original_merge(*args, **kwargs) + + monkeypatch.setattr(collection_merge.CollectionMerger, "merge", pause_merge) + worker = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + worker_result = executor.submit(worker.run, completed_at=case.created_at + timedelta(minutes=11)) + assert merge_started.wait(timeout=3) + try: + refresh_result = executor.submit( + coordinator.refresh, + observed_at=case.created_at + timedelta(minutes=12), + ).result(timeout=2) + finally: + release_merge.set() + result = worker_result.result(timeout=5) + + assert refresh_result.state is CollectionState.RUNNING + assert result.actual_records == case.plan.invocation.authored.num_records + assert ( + CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id).state + is CollectionState.SUCCEEDED + ) + + +def test_collection_requires_every_planned_winner_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + + with pytest.raises(StateNotFoundError, match="has no winner"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + + +def test_collection_prepares_missing_authorized_destination_parents( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + destination = case.workspace / "new" / "nested" / "output" + assert not destination.parent.exists() + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(destination=destination, submitted_at=case.created_at + timedelta(minutes=10)) + + assert destination.parent.is_dir() + assert not destination.exists() + assert ( + SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit() + == submitted + ) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + assert result.actual_records == case.plan.invocation.authored.num_records + assert destination.is_dir() + + +def test_concurrent_collection_submission_converges_on_one_job( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + def submit() -> object: + return SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _: submit(), range(2))) + + assert results[0] == results[1] + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_collection_ambiguous_submission_waits_for_scheduler_visibility_without_a_duplicate( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner() + scheduler = SlurmCommandClient(runner) + submissions = 0 + + def ambiguous_submit(script: str) -> object: + nonlocal submissions + del script + submissions += 1 + raise SlurmSubmissionError("sbatch could not be executed: command timed out", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", ambiguous_submit) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler) + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.PREPARED + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + with pytest.raises(StateConflictError, match="still being reconciled"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=11) + ) + assert submissions == 1 + + +def test_collection_recovers_an_accepted_submission_after_the_receipt_is_lost( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + scheduler = SlurmCommandClient(runner) + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + collection_plan = storage.read_plan("collection-0001") + runner.script_next( + "squeue", + FakeCommandResponse(stdout=f"5101|{collection_plan.submission_job_name}\n"), + ) + runner.script_next("sacct", FakeCommandResponse()) + waiting_for_binding = Event() + binding_published = Event() + + def wait_for_binding(seconds: float) -> None: + assert seconds == 1 + waiting_for_binding.set() + assert binding_published.wait(timeout=2) + + monkeypatch.setattr(collection_worker_module, "sleep", wait_for_binding) + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_result = executor.submit( + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + "collection-0001", + environment={"SLURM_JOB_ID": "5101"}, + ).run, + completed_at=case.created_at + timedelta(minutes=12), + ) + assert waiting_for_binding.wait(timeout=2) + recovered = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).refresh( + observed_at=case.created_at + timedelta(minutes=11) + ) + binding_published.set() + result = worker_result.result(timeout=5) + + assert recovered.collection_id == "collection-0001" + assert recovered.state is CollectionState.PENDING + assert recovered.scheduler == 5101 + assert result.actual_records == case.plan.invocation.authored.num_records + settled = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=13) + ) + assert settled.state is CollectionState.SUCCEEDED + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 1 + + +def test_collection_fences_an_invisible_accepted_submission_before_replacement_writes( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101), FakeSlurmJob(5102))) + scheduler = SlurmCommandClient(runner) + original_submit = scheduler.submit_script + + def accept_then_lose_receipt(script: str) -> object: + original_submit(script) + raise SlurmSubmissionError("sbatch response was lost", may_have_succeeded=True) + + monkeypatch.setattr(scheduler, "submit_script", accept_then_lose_receipt) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=10) + ) + monkeypatch.setattr(scheduler, "submit_script", original_submit) + runner.script_next("squeue", FakeCommandResponse()) + runner.script_next("sacct", FakeCommandResponse()) + + submitted = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, scheduler).submit( + submitted_at=case.created_at + timedelta(minutes=16) + ) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.FAILED + assert submitted.collection_id == "collection-0002" + assert submitted.state is CollectionState.SUBMITTED + assert submitted.scheduler == 5102 + with pytest.raises(StateConflictError, match="ordinary Slurm job identity"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + "collection-0001", + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=17)) + assert not _output_root(case.plan).exists() + assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 + + +def test_collection_definite_submission_failure_settles_and_can_be_retried( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + runner.script_next("sbatch", FakeCommandResponse(stderr="submission rejected", returncode=2)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + destination = case.workspace / "exports" / "custom" + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit( + destination=destination, num_partitions=3, submitted_at=case.created_at + timedelta(minutes=10) + ) + + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert storage.read_status("collection-0001").state is CollectionState.FAILED + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + assert submitted.collection_id == "collection-0002" + assert submitted.state is CollectionState.SUBMITTED + assert storage.read_plan(submitted.collection_id).num_partitions == 3 + assert storage.read_plan(submitted.collection_id).host_destination == destination.as_posix() + + +def test_collection_discards_trailing_journal_interrupted_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + original_publish = CollectionStorage.publish_status + + def interrupt_before_status(self: CollectionStorage, status: object) -> None: + del self, status + raise OSError("injected journal interruption") + + monkeypatch.setattr(CollectionStorage, "publish_status", interrupt_before_status) + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + monkeypatch.setattr(CollectionStorage, "publish_status", original_publish) + + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + + assert submitted.collection_id == "collection-0001" + assert CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).list_collection_ids() == ( + "collection-0001", + ) + + +def test_collection_failure_removes_staging_without_publishing_partial_output( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + def fail_after_writes(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("injected collection failure") + + monkeypatch.setattr( + "data_designer.slurm.state.collection_merge.CollectionMerger._write_result", + fail_after_writes, + ) + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = _output_root(case.plan) + assert not destination.exists() + assert not tuple(destination.parent.glob(".dd-*.tmp")) + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_rejects_output_replacement_before_descriptor_digest( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + original_describe = collection_merge._BoundOutput.describe + injected = False + + def replace_output_before_describe( + output: collection_merge._BoundOutput, + relative_path: str, + record_count: int, + ) -> object: + nonlocal injected + if not injected: + injected = True + output.path.rename(output.path.with_suffix(f"{output.path.suffix}.written")) + output.path.write_bytes(b"replacement") + output.path.chmod(0o600) + return original_describe(output, relative_path, record_count) + + monkeypatch.setattr(collection_merge._BoundOutput, "describe", replace_output_before_describe) + + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + destination = _output_root(case.plan) + assert not destination.exists() + assert not (destination.parent / submitted.staging_directory).exists() + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_refresh_cleans_only_its_exact_stage_after_oom( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + parent = _output_root(case.plan).parent + exact_stage = parent / submitted.staging_directory + exact_stage.mkdir(mode=0o700) + (exact_stage / "partial.parquet").write_text("incomplete") + unrelated_stage = parent / f".dd-collection-{'e' * 32}.tmp" + unrelated_stage.mkdir(mode=0o700) + (unrelated_stage / "active").write_text("preserve") + runner.set_job_state(5101, queue_state=None, accounting_state="OUT_OF_MEMORY", exit_code="0:125") + + refreshed = coordinator.refresh(observed_at=case.created_at + timedelta(minutes=11)) + + assert refreshed.state is CollectionState.FAILED + assert not exact_stage.exists() + assert (unrelated_stage / "active").read_text() == "preserve" + assert not _output_root(case.plan).exists() + + +def test_collection_uses_merge_partition_count( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ) + submitted = coordinator.submit(num_partitions=3, submitted_at=case.created_at + timedelta(minutes=10)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + plan = storage.read_plan(submitted.collection_id) + assert plan.num_partitions == 3 + assert coordinator.submit() == submitted + + with pytest.raises(StateConflictError, match="partition count"): + coordinator.submit(num_partitions=2) + + resolved_plan, candidates = CollectionInputResolver( + StateStorage(case.workspace, case.plan.run_id), + StateReader(StateStorage(case.workspace, case.plan.run_id), case.plan.run_id), + ).resolve(plan) + excessive = plan.model_copy(update={"num_partitions": case.plan.invocation.authored.num_records + 1}) + with pytest.raises(StateConflictError, match="partitions exceed requested records"): + CollectionDestinationResolver().validate_persisted(resolved_plan, excessive) + with pytest.raises(StateContractError, match="partitions exceed available records"): + validate_collection_inputs(resolved_plan, excessive, candidates) + + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert [file.relative_path for file in result.files] == [ + "part-00000.parquet", + "part-00001.parquet", + "part-00002.parquet", + ] + values = [ + value + for file in result.files + for value in lazy.pq.read_table(_output_root(case.plan) / file.relative_path).column("record_id").to_pylist() + ] + assert values == list(range(case.plan.invocation.authored.num_records)) + + +def test_succeeded_collection_rejects_modified_partition_bytes( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + output_path = _output_root(case.plan) / result.files[0].relative_path + output_path.write_bytes(output_path.read_bytes() + b"tampered") + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit() + + +def test_succeeded_collection_rejects_same_size_partition_mutation( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + output_path = _output_root(case.plan) / result.files[0].relative_path + content = bytearray(output_path.read_bytes()) + content[-1] ^= 1 + output_path.write_bytes(content) + + with pytest.raises(SlurmStateError, match="cannot submit collection"): + coordinator.submit() + + +def test_succeeded_collection_status_check_does_not_read_partition_payloads( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + original_open = collection_storage_module.open_verified_regular_file + + def reject_partition_reads(*args: object, **kwargs: object) -> object: + if args[1] != "collection-result.json": + raise AssertionError("login-host validation attempted to read collected partition bytes") + return original_open(*args, **kwargs) + + monkeypatch.setattr(collection_storage_module, "open_verified_regular_file", reject_partition_reads) + + assert coordinator.submit().state is CollectionState.SUCCEEDED + + +def test_succeeded_collection_rejects_status_result_digest_mismatch( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + succeeded = storage.read_status(submitted.collection_id) + assert succeeded.result is not None + storage.replace_status( + succeeded.model_copy( + update={ + "result": succeeded.result.model_copy(update={"sha256": "f" * 64}), + } + ) + ) + + with pytest.raises(StateCorruptionError, match="does not bind its published result"): + coordinator.submit() + + +def test_collection_refuses_destination_collision_before_submission( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + destination = _output_root(case.plan) + destination.mkdir(mode=0o700) + marker = destination / "existing.txt" + marker.write_text("preserve") + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + + with pytest.raises(StateConflictError, match="already exists"): + SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert marker.read_text() == "preserve" + assert not any(Path(call[0]).name == "sbatch" for call in runner.calls) + + +def test_collection_refuses_atomic_publication_collision_without_replacement( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + destination = _output_root(case.plan) + original_rename = collection_filesystem._rename_without_overwrite + + def collide( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + destination.mkdir(mode=0o700) + (destination / "existing.txt").write_text("preserve") + original_rename(source_directory, source_name, destination_directory, destination_name) + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", collide) + + with pytest.raises(StateConflictError, match="already exists"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert (destination / "existing.txt").read_text() == "preserve" + assert not (destination.parent / submitted.staging_directory).exists() + + +def test_collection_publication_falls_back_when_atomic_no_replace_is_unsupported( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent = tmp_path / "publication" + parent.mkdir(mode=0o700) + source = parent / "staged" + source.mkdir(mode=0o700) + (source / "part-00000.parquet").write_text("complete") + + def unsupported(*args: object) -> None: + raise OSError(errno.EINVAL, "unsupported by filesystem") + + monkeypatch.setattr(collection_filesystem, "_atomic_rename_without_overwrite", unsupported) + descriptor = os.open(parent, os.O_RDONLY | os.O_DIRECTORY) + try: + collection_filesystem._rename_without_overwrite(descriptor, source.name, descriptor, "collected") + finally: + os.close(descriptor) + + assert not source.exists() + assert (parent / "collected" / "part-00000.parquet").read_text() == "complete" + assert not any(name.startswith(".dd-collection-reservation-") for name in os.listdir(parent)) + + +def test_collection_publication_fallback_preserves_existing_destination( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent = tmp_path / "publication" + parent.mkdir(mode=0o700) + source = parent / "staged" + source.mkdir(mode=0o700) + destination = parent / "collected" + destination.mkdir(mode=0o700) + marker = destination / "existing.txt" + marker.write_text("preserve") + + def unsupported(*args: object) -> None: + raise OSError(errno.EOPNOTSUPP, "unsupported by filesystem") + + monkeypatch.setattr(collection_filesystem, "_atomic_rename_without_overwrite", unsupported) + descriptor = os.open(parent, os.O_RDONLY | os.O_DIRECTORY) + try: + with pytest.raises(StateConflictError, match="already exists"): + collection_filesystem._rename_without_overwrite(descriptor, source.name, descriptor, destination.name) + finally: + os.close(descriptor) + + assert source.is_dir() + assert marker.read_text() == "preserve" + assert not any(name.startswith(".dd-collection-reservation-") for name in os.listdir(parent)) + + +def test_collection_worker_recovers_interrupted_private_reservation( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + destination = _output_root(case.plan) + marker_name = collection_filesystem._reservation_marker_name(submitted.staging_directory, destination.name) + descriptor = os.open(destination.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.mkdir(destination.name, 0o700, dir_fd=descriptor) + os.fsync(descriptor) + reservation = os.stat(destination.name, dir_fd=descriptor, follow_symlinks=False) + marker_bytes = collection_filesystem._reservation_marker_bytes( + submitted.staging_directory, + destination.name, + (reservation.st_dev, reservation.st_ino, reservation.st_ctime_ns), + ) + collection_filesystem._create_reservation_marker(descriptor, marker_name, marker_bytes) + finally: + os.close(descriptor) + + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert result.actual_records == case.plan.invocation.authored.num_records + collections = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + assert collections.get_result_path(collections.read_plan(submitted.collection_id)).is_file() + assert not (destination.parent / marker_name).exists() + + +def test_collection_reservation_metadata_failure_preserves_unmarked_destination( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + parent = tmp_path / "publication" + parent.mkdir(mode=0o700) + source = parent / "staged" + source.mkdir(mode=0o700) + destination = parent / "collected" + marker_name = collection_filesystem._reservation_marker_name(source.name, destination.name) + + def unsupported(*args: object) -> None: + raise OSError(errno.EINVAL, "unsupported by filesystem") + + original_stat = os.stat + descriptor = os.open(parent, os.O_RDONLY | os.O_DIRECTORY) + try: + with monkeypatch.context() as patcher: + patcher.setattr(collection_filesystem, "_atomic_rename_without_overwrite", unsupported) + + def fail_reservation_stat(path: object, *args: object, **kwargs: object) -> os.stat_result: + if path == destination.name and kwargs.get("dir_fd") == descriptor: + raise OSError(errno.EIO, "reservation metadata unavailable") + return original_stat(path, *args, **kwargs) + + patcher.setattr(os, "stat", fail_reservation_stat) + with pytest.raises(OSError, match="reservation metadata unavailable"): + collection_filesystem._rename_without_overwrite(descriptor, source.name, descriptor, destination.name) + finally: + os.close(descriptor) + + assert destination.is_dir() + assert not (parent / marker_name).exists() + collection_filesystem.remove_collection_stage(destination, source.name, parent) + assert destination.is_dir() + assert not source.exists() + assert not (parent / marker_name).exists() + + +def test_collection_recovery_preserves_unrelated_empty_destination( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + destination = _output_root(case.plan) + unrelated = destination.parent / "unrelated-reservation" + unrelated.mkdir(mode=0o700) + destination.mkdir(mode=0o700) + marker_name = collection_filesystem._reservation_marker_name(submitted.staging_directory, destination.name) + unrelated_stat = unrelated.stat() + marker_bytes = collection_filesystem._reservation_marker_bytes( + submitted.staging_directory, + destination.name, + (unrelated_stat.st_dev, unrelated_stat.st_ino, unrelated_stat.st_ctime_ns), + ) + descriptor = os.open(destination.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + collection_filesystem._create_reservation_marker(descriptor, marker_name, marker_bytes) + finally: + os.close(descriptor) + + with pytest.raises(StateConflictError, match="already exists"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert destination.is_dir() + assert not list(destination.iterdir()) + assert unrelated.is_dir() + assert not (destination.parent / marker_name).exists() + + +def test_collection_recovery_rejects_reused_inode_with_changed_timestamp(tmp_path: Path) -> None: + parent = tmp_path / "publication" + parent.mkdir(mode=0o700) + source = parent / "staged" + source.mkdir(mode=0o700) + destination = parent / "collected" + destination.mkdir(mode=0o700) + current = destination.stat() + marker_name = collection_filesystem._reservation_marker_name(source.name, destination.name) + marker_bytes = collection_filesystem._reservation_marker_bytes( + source.name, + destination.name, + (current.st_dev, current.st_ino, current.st_ctime_ns + 1), + ) + descriptor = os.open(parent, os.O_RDONLY | os.O_DIRECTORY) + try: + collection_filesystem._create_reservation_marker(descriptor, marker_name, marker_bytes) + finally: + os.close(descriptor) + + collection_filesystem.remove_collection_stage(destination, source.name, parent) + + assert destination.is_dir() + assert not source.exists() + assert not (parent / marker_name).exists() + + +def test_collection_failed_publication_preserves_changed_reservation(tmp_path: Path) -> None: + parent = tmp_path / "publication" + parent.mkdir(mode=0o700) + destination = parent / "collected" + destination.mkdir(mode=0o700) + reservation = destination.stat() + destination.chmod(0o500) + assert destination.stat().st_ctime_ns != reservation.st_ctime_ns + + descriptor = os.open(parent, os.O_RDONLY | os.O_DIRECTORY) + try: + collection_filesystem._remove_unchanged_reservation(descriptor, destination.name, reservation) + finally: + os.close(descriptor) + + assert destination.is_dir() + + +def test_collection_detects_destination_parent_replacement_before_success( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(destination=case.workspace / "exports" / "output", submitted_at=case.created_at + timedelta(minutes=10)) + destination = case.workspace / "exports" / "output" + original_parent = destination.parent + moved_parent = original_parent.with_name("moved-exports") + original_rename = collection_filesystem._rename_without_overwrite + injected = False + + def replace_parent( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + nonlocal injected + if not injected: + injected = True + original_parent.rename(moved_parent) + original_parent.mkdir(mode=0o700) + original_rename(source_directory, source_name, destination_directory, destination_name) + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", replace_parent) + + with pytest.raises(SlurmStateError, match="collection .* failed"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert not destination.exists() + assert not (moved_parent / destination.name).exists() + assert not (moved_parent / submitted.staging_directory).exists() + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.FAILED + + +def test_collection_recovers_when_publication_completed_before_reported_error( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + original_rename = collection_filesystem._rename_without_overwrite + + def publish_then_fail( + source_directory: int, + source_name: str, + destination_directory: int, + destination_name: str, + ) -> None: + original_rename(source_directory, source_name, destination_directory, destination_name) + raise OSError("injected post-rename failure") + + monkeypatch.setattr(collection_filesystem, "_rename_without_overwrite", publish_then_fail) + + result = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + assert result.actual_records == case.plan.invocation.authored.num_records + persisted = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)).read_status(submitted.collection_id) + assert persisted.state is CollectionState.SUCCEEDED + assert _output_root(case.plan).is_dir() + + +def test_collection_worker_rejects_login_node_execution( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + + with pytest.raises(StateConflictError, match="inside its recorded Slurm job"): + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + + +def test_collection_worker_rejects_login_node_recovery_before_reading_outputs( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialize_run(tmp_path, authored_run, multi_node_plan) + _publish_all_winners(case) + runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) + submitted = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + SlurmCommandClient(runner), + ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={"SLURM_JOB_ID": "5101"}, + ).run(completed_at=case.created_at + timedelta(minutes=11)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + storage.replace_status(submitted) + + def reject_output_reads(*args: object, **kwargs: object) -> None: + del args, kwargs + raise AssertionError("login-host worker attempted to read collected outputs") + + monkeypatch.setattr(CollectionStorage, "verify_result_files", reject_output_reads) + worker = SlurmCollectionWorker( + case.workspace, + case.plan.run_id, + submitted.collection_id, + environment={}, + ) + + with pytest.raises(StateConflictError, match="inside its recorded Slurm job"): + worker.run(completed_at=case.created_at + timedelta(minutes=12)) + + assert storage.read_status(submitted.collection_id) == submitted + + +def test_collection_worker_validates_location_before_constructing_state(tmp_path: Path) -> None: + with pytest.raises(SlurmStateError, match="invalid persisted collection worker location"): + SlurmCollectionWorker("relative/workspace", "run-001", "collection-0001") + with pytest.raises(SlurmStateError, match="invalid persisted collection worker location"): + SlurmCollectionWorker(tmp_path, "../run", "collection-0001") + + +def _output_root(plan: ResolvedSlurmRunPlan) -> Path: + return Path(plan.selected_profile.profile.workspace_root) / "runs" / plan.run_id / "merge-output" + + +def _initialize_run( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> _RunCase: + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o700) + relocated_plan = _relocate_plan(plan, workspace) + created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) + run_root = workspace / "runs" / relocated_plan.run_id + run = RunManifest( + schema_version=1, + run_id=relocated_plan.run_id, + created_at=created_at, + authored_config=relocated_plan.authored_config, + resolved_plan=ArtifactReference( + path=(run_root / "resolved-plan.json").as_posix(), + sha256=relocated_plan.compute_sha256(), + ), + shard_count=len(relocated_plan.shards), + ) + shards = tuple( + ShardManifest( + schema_version=1, + run_id=run.run_id, + shard_id=planned.shard_id, + shard_index=planned.shard_index, + record_range=planned.record_range, + input_partition=planned.input_partition, + resume_workspace=planned.resume_workspace, + created_at=created_at, + ) + for planned in relocated_plan.shards + ) + writer = SlurmStateWriter(workspace, run.run_id) + writer.initialize_run(authored_config, relocated_plan, run, shards) + return _RunCase(workspace, relocated_plan, run, shards, writer, created_at) + + +def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: + previous_workspace = plan.selected_profile.profile.workspace_root + payload = cast( + dict[str, object], + json.loads(plan.serialize_json().replace(previous_workspace, workspace.as_posix())), + ) + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected_profile["profile"]) + profile_mounts = cast(list[dict[str, object]], profile_payload["container_mounts"]) + resolved_mounts = cast(list[dict[str, object]], payload["container_mounts"]) + for mount in (*profile_mounts, *resolved_mounts): + mount["source"] = workspace.parent.as_posix() + mount["target"] = workspace.parent.as_posix() + profile = SlurmProfile.model_validate(selected_profile["profile"]) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _submitted_attempt( + case: _RunCase, + shard: ShardManifest, + *, + scheduler: SchedulerIdentity, +) -> AttemptManifest: + return AttemptManifest( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=scheduler, + created_at=case.created_at + timedelta(minutes=1), + updated_at=case.created_at + timedelta(minutes=1), + ) + + +def _publish_all_winners( + case: _RunCase, + *, + shortfall: int = 0, + alternating_struct_order: bool = False, +) -> None: + for shard in case.shards: + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index) + submitted = _submitted_attempt(case, shard, scheduler=scheduler) + case.writer.create_attempt(submitted) + running = _copy_attempt( + submitted, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=2), + ) + case.writer.update_attempt(running) + with case.writer.acquire_dataset_workspace(shard.shard_id, submitted.attempt_id, "never") as dataset_path: + _publish_candidate( + case, + shard, + running, + dataset_path, + shortfall=shortfall, + struct_fields_reversed=alternating_struct_order and shard.shard_index % 2 == 1, + with_struct=alternating_struct_order, + ) + case.writer.finalize_winner( + shard.shard_id, + running.attempt_id, + published_at=case.created_at + timedelta(minutes=6), + ) + + +def _publish_candidate( + case: _RunCase, + shard: ShardManifest, + running: AttemptManifest, + dataset_path: Path, + *, + shortfall: int = 0, + with_struct: bool = False, + struct_fields_reversed: bool = False, +) -> None: + output_path = dataset_path / "part-00000.parquet" + actual_records = shard.record_range.record_count - shortfall + values = range(shard.record_range.start_index, shard.record_range.start_index + actual_records) + table = lazy.pa.table({"record_id": values}) + if with_struct: + struct_fields = [ + lazy.pa.field("needs_revision", lazy.pa.bool_()), + lazy.pa.field("visible_item", lazy.pa.string()), + ] + if struct_fields_reversed: + struct_fields.reverse() + assessment = lazy.pa.array( + [{"needs_revision": index % 2 == 0, "visible_item": f"item-{index}"} for index in values], + type=lazy.pa.struct(struct_fields), + ) + table = table.append_column("editorial_assessment", assessment) + lazy.pq.write_table(table, output_path) + output_path.chmod(0o644) + content = output_path.read_bytes() + candidate_path = ( + case.writer.run_root / "shards" / shard.shard_id / "attempts" / running.attempt_id / "output-manifest.json" + ) + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id=running.attempt_id, + attempt_ordinal=running.attempt_ordinal, + created_at=running.updated_at + timedelta(minutes=1), + dataset_path=dataset_path.as_posix(), + requested_records=shard.record_range.record_count, + actual_records=actual_records, + outcome=CandidateOutcome.COMPLETE if not shortfall else CandidateOutcome.PARTIAL, + files=( + CandidateOutputFile( + relative_path=output_path.name, + sha256=hashlib.sha256(content).hexdigest(), + byte_size=len(content), + record_count=actual_records, + ), + ), + dataset_schema_digest=compute_candidate_schema_digest(table.schema), + provenance_digest=case.plan.compute_sha256(), + ) + candidate_reference = ArtifactReference(path=candidate_path.as_posix(), sha256=candidate.compute_sha256()) + result = ClientResult( + schema_version=1, + run_id=case.run.run_id, + shard_id=shard.shard_id, + attempt_id=running.attempt_id, + completed_at=running.updated_at + timedelta(minutes=2), + requested_records=shard.record_range.record_count, + actual_records=actual_records, + outcome=ClientOutcome.COMPLETE if not shortfall else ClientOutcome.PARTIAL, + dataset_path=dataset_path.as_posix(), + early_shutdown=False, + requested_resume_mode=case.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=candidate_reference, + ) + case.writer.publish_attempt_result(result, candidate) + completed = _copy_attempt( + running, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + candidate_output=candidate_reference, + updated_at=running.updated_at + timedelta(minutes=3), + ) + case.writer.update_attempt(completed) + + +def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: + payload = attempt.model_dump(mode="python") + payload.update(updates) + return AttemptManifest.model_validate(payload) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py new file mode 100644 index 000000000..05a8561d3 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, + ShardManifest, + ShardWinner, + StateRecord, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +CONTRACT_GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" +GOLDEN_MODELS: tuple[tuple[str, type[StateRecord]], ...] = ( + ("run_manifest.json", RunManifest), + ("shard_manifest.json", ShardManifest), + ("successful_attempt.json", AttemptManifest), + ("single_node_pending_readiness.json", AttemptReadiness), + ("single_node_starting_readiness.json", AttemptReadiness), + ("single_node_readiness.json", AttemptReadiness), + ("single_node_publication_failed_readiness.json", AttemptReadiness), + ("single_node_failed_readiness.json", AttemptReadiness), + ("single_node_stopped_readiness.json", AttemptReadiness), + ("multi_node_readiness.json", AttemptReadiness), + ("failed_attempt.json", AttemptManifest), + ("stale_readiness.json", AttemptReadiness), + ("accounting_lag.json", SchedulerObservation), + ("candidate_output.json", CandidateOutputManifest), + ("shard_winner.json", ShardWinner), + ("collection_plan.json", CollectionPlan), +) +PERSISTED_INTENT_FIXTURES = ( + CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", + CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", +) +_COMPATIBILITY_FIXTURE_DIGEST = "18dbde1e5a6d883d90d6392487b84750a5dd8133f44d75d425d240ba7feb7b6c" + + +@pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) +def test_state_golden_record_round_trip_is_deterministic(filename: str, model: type[StateRecord]) -> None: + serialized = (GOLDEN_DIRECTORY / filename).read_text() + + record = model.model_validate_json(serialized) + direct_record = model(**record.model_dump(mode="python")) + + assert direct_record == record + assert record.serialize_json() == serialized + assert model.model_validate_json(record.serialize_canonical_json()) == record + assert model.model_validate_json(record.serialize_json()) == record + assert record.compute_sha256() == hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + +def test_golden_artifact_references_hash_exact_persisted_bytes() -> None: + candidate_bytes = (GOLDEN_DIRECTORY / "candidate_output.json").read_bytes() + winner_bytes = (GOLDEN_DIRECTORY / "shard_winner.json").read_bytes() + attempt = AttemptManifest.model_validate_json((GOLDEN_DIRECTORY / "successful_attempt.json").read_text()) + winner = ShardWinner.model_validate_json(winner_bytes) + collection = CollectionPlan.model_validate_json((GOLDEN_DIRECTORY / "collection_plan.json").read_text()) + + assert attempt.candidate_output is not None + assert attempt.candidate_output.sha256 == hashlib.sha256(candidate_bytes).hexdigest() + assert winner.candidate_manifest.sha256 == hashlib.sha256(candidate_bytes).hexdigest() + assert collection.planned_shards[0].winner_manifest.sha256 == hashlib.sha256(winner_bytes).hexdigest() + + +def test_persisted_state_compatibility_fixtures_are_frozen() -> None: + """Make persisted v1 byte changes an explicit compatibility decision.""" + fixture_bytes = b"".join( + path.name.encode("utf-8") + b"\0" + path.read_bytes() for path in PERSISTED_INTENT_FIXTURES + ) + fixture_bytes += b"".join( + filename.encode("utf-8") + b"\0" + (GOLDEN_DIRECTORY / filename).read_bytes() for filename, _ in GOLDEN_MODELS + ) + + assert hashlib.sha256(fixture_bytes).hexdigest() == _COMPATIBILITY_FIXTURE_DIGEST diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py new file mode 100644 index 000000000..9a9d902b9 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -0,0 +1,2534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import os +import stat +import subprocess +import sys +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from threading import Barrier, Event +from typing import TypeVar, cast + +import pytest + +import data_designer.lazy_heavy_imports as lazy +from data_designer.slurm import filesystem as slurm_filesystem +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.config import ContainerMount, DataDesignerSlurmConfig, SlurmProfile +from data_designer.slurm.contracts import ArtifactReference, ContractValue, compute_canonical_json_sha256, pretty_json +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.state import ( + AttemptId, + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerIdentity, + ShardId, + ShardManifest, + ShardWinner, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, + StateNotFoundError, + compute_candidate_schema_digest, +) +from data_designer.slurm.state import artifacts as state_artifacts +from data_designer.slurm.state import filesystem as state_filesystem +from data_designer.slurm.state import finalization as state_finalization +from data_designer.slurm.state import storage as state_storage + +_ValueT = TypeVar("_ValueT", bound=ContractValue) + + +@dataclass(frozen=True, slots=True) +class _StateCase: + workspace: Path + authored_config: DataDesignerSlurmConfig + plan: ResolvedSlurmRunPlan + run: RunManifest + shards: tuple[ShardManifest, ...] + writer: SlurmStateWriter + created_at: datetime + + +@dataclass(frozen=True, slots=True) +class _FinalizationCase: + attempt: AttemptManifest + client_result: ClientResult + candidate: CandidateOutputManifest + published_at: datetime + output_path: Path + + +def test_combined_dataset_and_state_lock_owns_the_canonical_order( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[str] = [] + + @contextmanager + def resume_lock(storage: state_storage.StateStorage, shard_id: ShardId) -> Iterator[None]: + del storage, shard_id + events.append("resume-enter") + try: + yield + finally: + events.append("resume-exit") + + @contextmanager + def shard_lock(storage: state_storage.StateStorage, shard_id: ShardId) -> Iterator[None]: + del storage, shard_id + events.append("shard-enter") + try: + yield + finally: + events.append("shard-exit") + + monkeypatch.setattr(state_storage.StateStorage, "acquire_resume_lock", resume_lock) + monkeypatch.setattr(state_storage.StateStorage, "acquire_shard_lock", shard_lock) + storage = state_storage.StateStorage(tmp_path, "run-lock-order") + + with storage.acquire_resume_and_shard_locks("shard-00000"): + events.append("body") + + assert events == ["resume-enter", "shard-enter", "body", "shard-exit", "resume-exit"] + + +def test_run_initialization_is_idempotent_restrictive_and_reloadable( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + + assert case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) is case.run + assert case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) is case.run + + reader = SlurmStateWriter(case.workspace, case.plan.run_id) + assert reader.load_run() == case.run + assert reader.load_authored_config() == case.authored_config + assert reader.load_resolved_plan() == case.plan + assert reader.load_shards() == case.shards + assert reader.load_attempts("shard-00000") == () + + package_directories = tuple(path for path in (case.workspace / "runs").rglob("*") if path.is_dir()) + package_directories += (case.workspace / "runs",) + package_files = tuple(path for path in (case.workspace / "runs").rglob("*") if path.is_file()) + assert all(stat.S_IMODE(path.stat().st_mode) == 0o700 for path in package_directories) + assert all(stat.S_IMODE(path.stat().st_mode) == 0o600 for path in package_files) + assert not tuple((case.workspace / "runs").rglob(".state.*.tmp")) + + +def test_concurrent_identical_run_initialization_converges( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + barrier = Barrier(2) + + def initialize() -> RunManifest: + barrier.wait() + writer = SlurmStateWriter(case.workspace, case.plan.run_id) + return writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _: initialize(), range(2))) + + assert results == (case.run, case.run) + assert case.writer.load_run() == case.run + + +def test_context_loading_reads_each_immutable_record_once( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + original_read_record = state_storage.StateStorage.read_record + record_names: list[str] = [] + + def track_read( + self: state_storage.StateStorage, + directory_descriptor: int, + name: str, + display_path: Path, + record_type: type[_ValueT], + ) -> _ValueT: + record_names.append(name) + return original_read_record(self, directory_descriptor, name, display_path, record_type) + + monkeypatch.setattr(state_storage.StateStorage, "read_record", track_read) + + assert case.writer.load_attempts("shard-00000") == () + assert record_names.count("run.json") == 1 + assert record_names.count("authored-config.json") == 1 + assert record_names.count("resolved-plan.json") == 1 + assert record_names.count("shard.json") == 1 + + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + runtime_directory = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/runtime" + assert not runtime_directory.exists() + readiness = _readiness(case, attempt) + case.writer.write_readiness(readiness) + record_names.clear() + + assert case.writer.load_readiness(attempt.shard_id, attempt.attempt_id) == readiness + assert record_names.count("run.json") == 1 + assert record_names.count("authored-config.json") == 1 + assert record_names.count("resolved-plan.json") == 1 + assert record_names.count("shard.json") == 1 + + +def test_shard_mutations_do_not_scan_unrelated_attempt_directories( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run, multi_node_plan) + unrelated_attempts = case.writer.run_root / "shards" / "shard-00001" / "attempts" + (unrelated_attempts / "unexpected").mkdir(mode=0o700) + attempt = _submitted_attempt(case) + + assert case.writer.create_attempt(attempt) == attempt + readiness = _readiness(case, attempt) + assert case.writer.write_readiness(readiness) == readiness + assert case.writer.load_attempts(attempt.shard_id) == (attempt,) + with pytest.raises(StateCorruptionError, match="invalid attempt directory"): + case.writer.load_attempts("shard-00001") + + +def test_shard_mutation_rejects_run_count_drift_without_scanning_other_shards( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + run_path = case.writer.run_root / "run.json" + run_path.write_text(case.run.model_copy(update={"shard_count": 2}).serialize_json()) + + with pytest.raises(StateCorruptionError, match="shard .* is invalid"): + case.writer.create_attempt(_submitted_attempt(case)) + + +def test_incomplete_attempt_recognizes_created_state_temporary_file( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + attempt_root = case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + attempt_root.mkdir(mode=0o700) + with state_filesystem.open_verified_directory(attempt_root, require_private=True) as directory_descriptor: + temporary_descriptor, temporary_name = state_filesystem.create_state_temporary_file(directory_descriptor) + os.close(temporary_descriptor) + + assert state_filesystem.is_state_temporary_name(temporary_name) + assert case.writer.load_attempts(attempt.shard_id) == () + + +def test_reader_rejects_unsupported_persisted_schema_version( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + run_path = case.writer.run_root / "run.json" + payload = json.loads(run_path.read_text()) + payload["schema_version"] = 2 + run_path.write_text(pretty_json(payload)) + + with pytest.raises(StateCorruptionError, match="unsupported schema_version"): + case.writer.load_run() + + +def test_reader_rejects_noncanonical_record_that_relies_on_new_defaults( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + authored_config_path = case.writer.run_root / "authored-config.json" + payload = json.loads(authored_config_path.read_text()) + del payload["array_tasks"] + authored_config_path.write_text(pretty_json(payload)) + + with pytest.raises(StateCorruptionError, match="schema migrations are not supported"): + case.writer.load_authored_config() + + +def test_run_commit_marker_is_published_last_and_interrupted_initialization_retries( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + original_publish = state_storage.publish_immutable_text + + def interrupt_before_run_commit( + directory_descriptor: int, + name: str, + content: str, + display_path: Path, + *, + maximum_size: int, + ) -> bool: + if name == "run.json": + assert (case.writer.run_root / "authored-config.json").is_file() + assert (case.writer.run_root / "resolved-plan.json").is_file() + assert (case.writer.run_root / "shards" / "shard-00000" / "shard.json").is_file() + with pytest.raises(StateNotFoundError): + SlurmStateWriter(case.workspace, case.plan.run_id).load_run() + raise OSError("simulated interruption") + return original_publish( + directory_descriptor, + name, + content, + display_path, + maximum_size=maximum_size, + ) + + monkeypatch.setattr(state_storage, "publish_immutable_text", interrupt_before_run_commit) + with pytest.raises(SlurmStateError, match="cannot initialize"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + monkeypatch.setattr(state_storage, "publish_immutable_text", original_publish) + assert case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) == case.run + assert case.writer.load_shards() == case.shards + + +def test_interrupted_immutable_publication_recovers_the_committed_hard_link( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + original_unlink = state_filesystem.os.unlink + + def interrupt_temporary_unlink(path: str, *, dir_fd: int | None = None) -> None: + del path, dir_fd + raise KeyboardInterrupt("injected interruption after immutable publication") + + monkeypatch.setattr(state_filesystem.os, "unlink", interrupt_temporary_unlink) + with pytest.raises(KeyboardInterrupt, match="after immutable publication"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + committed_path = case.writer.run_root / "authored-config.json" + temporary_paths = tuple(case.writer.run_root.glob(".state.*.tmp")) + assert committed_path.stat().st_nlink == 2 + assert len(temporary_paths) == 1 + assert temporary_paths[0].stat().st_ino == committed_path.stat().st_ino + + monkeypatch.setattr(state_filesystem.os, "unlink", original_unlink) + assert case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) == case.run + assert committed_path.stat().st_nlink == 1 + assert not tuple(case.writer.run_root.glob(".state.*.tmp")) + + +def test_concurrent_recovery_does_not_make_the_immutable_publisher_fail( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + original_unlink = state_filesystem.os.unlink + recovered = False + + def recover_before_publisher_unlink(path: str, *, dir_fd: int | None = None) -> None: + nonlocal recovered + assert dir_fd is not None + monkeypatch.setattr(state_filesystem.os, "unlink", original_unlink) + assert ( + state_filesystem.read_regular_text( + dir_fd, + "authored-config.json", + case.writer.run_root / "authored-config.json", + maximum_size=16 * 1024 * 1024, + ) + == authored_run_single.serialize_json() + ) + recovered = True + original_unlink(path, dir_fd=dir_fd) + + monkeypatch.setattr(state_filesystem.os, "unlink", recover_before_publisher_unlink) + assert case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) == case.run + assert recovered + assert not tuple(case.writer.run_root.glob(".state.*.tmp")) + + +def test_interrupted_publication_recovery_accepts_a_concurrent_atomic_replacement( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + directory = tmp_path / "records" + directory.mkdir(mode=0o700) + record_path = directory / "record.json" + record_path.write_text("old") + record_path.chmod(0o600) + temporary_path = directory / ".state.0123456789abcdef.tmp" + os.link(record_path, temporary_path) + replacement_path = directory / "replacement.json" + replacement_path.write_text("new") + replacement_path.chmod(0o600) + original_unlink = state_filesystem.os.unlink + + def replace_after_recovery_unlink(path: str, *, dir_fd: int | None = None) -> None: + assert dir_fd is not None + original_unlink(path, dir_fd=dir_fd) + os.replace( + replacement_path.name, + record_path.name, + src_dir_fd=dir_fd, + dst_dir_fd=dir_fd, + ) + + with state_filesystem.open_verified_directory(directory, require_private=True) as directory_descriptor: + monkeypatch.setattr(state_filesystem.os, "unlink", replace_after_recovery_unlink) + assert ( + state_filesystem.read_regular_text( + directory_descriptor, + record_path.name, + record_path, + maximum_size=16, + ) + == "new" + ) + + +def test_record_reader_retries_an_atomic_replacement_between_stat_and_open( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + directory = tmp_path / "records" + directory.mkdir(mode=0o700) + record_path = directory / "record.json" + record_path.write_text("old") + record_path.chmod(0o600) + replacement_path = directory / "replacement.json" + replacement_path.write_text("new") + replacement_path.chmod(0o600) + original_open = state_filesystem.os.open + replaced = False + + with state_filesystem.open_verified_directory(directory, require_private=True) as directory_descriptor: + + def replace_before_open( + path: str | bytes, + flags: int, + mode: int = 0o777, + *, + dir_fd: int | None = None, + ) -> int: + nonlocal replaced + if path == record_path.name and not replaced: + replaced = True + os.replace( + replacement_path.name, + record_path.name, + src_dir_fd=directory_descriptor, + dst_dir_fd=directory_descriptor, + ) + return original_open(path, flags, mode, dir_fd=dir_fd) + + monkeypatch.setattr(state_filesystem.os, "open", replace_before_open) + assert ( + state_filesystem.read_regular_text( + directory_descriptor, + record_path.name, + record_path, + maximum_size=16, + ) + == "new" + ) + assert replaced + + +def test_record_reader_preserves_on_disk_newlines(tmp_path: Path) -> None: + directory = tmp_path / "records" + directory.mkdir(mode=0o700) + record_path = directory / "record.json" + record_path.write_bytes(b"first\r\nsecond\r\n") + record_path.chmod(0o600) + + with state_filesystem.open_verified_directory(directory, require_private=True) as directory_descriptor: + assert ( + state_filesystem.read_regular_text( + directory_descriptor, + record_path.name, + record_path, + maximum_size=64, + ) + == "first\r\nsecond\r\n" + ) + + +def test_immutable_publication_retry_resyncs_existing_record( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + directory = tmp_path / "records" + directory.mkdir(mode=0o700) + record_path = directory / "record.json" + original_fsync = state_filesystem.os.fsync + directory_syncs = 0 + + def fail_first_directory_fsync(descriptor: int) -> None: + nonlocal directory_syncs + if stat.S_ISDIR(os.fstat(descriptor).st_mode): + directory_syncs += 1 + if directory_syncs == 1: + raise OSError("simulated post-publish fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_first_directory_fsync) + with state_filesystem.open_verified_directory(directory, require_private=True) as directory_descriptor: + with pytest.raises(OSError, match="post-publish"): + state_filesystem.publish_immutable_text( + directory_descriptor, + record_path.name, + "persisted", + record_path, + maximum_size=64, + ) + syncs_before_retry = directory_syncs + assert not state_filesystem.publish_immutable_text( + directory_descriptor, + record_path.name, + "persisted", + record_path, + maximum_size=64, + ) + + assert directory_syncs > syncs_before_retry + + +def test_run_initialization_never_replaces_different_immutable_bytes( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + different_run = _validated_copy(case.run, created_at=case.created_at - timedelta(seconds=1)) + + with pytest.raises(StateConflictError, match="different immutable state"): + case.writer.initialize_run(case.authored_config, case.plan, different_run, case.shards) + + assert case.writer.load_run() == case.run + + +def test_run_reader_rejects_manifest_bound_to_a_different_run( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + mismatched = _validated_copy(case.run, run_id="other-run") + (case.writer.run_root / "run.json").write_text(mismatched.serialize_json()) + + with pytest.raises(StateCorruptionError, match="persisted location"): + case.writer.load_run() + + +@pytest.mark.parametrize("record_kind", ("symlink", "hardlink", "fifo", "permissive")) +def test_run_reader_rejects_unsafe_record_types_and_modes( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + record_kind: str, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + run_path = case.writer.run_root / "run.json" + if record_kind == "symlink": + outside = tmp_path / "outside-run.json" + outside.write_text(case.run.serialize_json()) + outside.chmod(0o600) + run_path.unlink() + run_path.symlink_to(outside) + elif record_kind == "hardlink": + outside = tmp_path / "outside-run.json" + outside.write_text(case.run.serialize_json()) + outside.chmod(0o600) + run_path.unlink() + os.link(outside, run_path) + elif record_kind == "fifo": + run_path.unlink() + os.mkfifo(run_path) + else: + run_path.chmod(0o644) + + with pytest.raises(StateCorruptionError, match="cannot load"): + case.writer.load_run() + + +def test_initialization_rejects_symlinked_package_storage_without_touching_target( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + outside = tmp_path / "outside" + outside.mkdir() + (case.workspace / "runs").symlink_to(outside, target_is_directory=True) + + with pytest.raises(SlurmStateError, match="cannot initialize"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + assert tuple(outside.iterdir()) == () + + +@pytest.mark.parametrize("lock_kind", ("symlink", "fifo")) +def test_initialization_rejects_unsafe_run_lock_without_changing_target( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + lock_kind: str, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + lock_root = case.workspace / "runs" / ".locks" + lock_root.mkdir(parents=True, mode=0o700) + lock_path = lock_root / f"run-{case.plan.run_id}.lock" + if lock_kind == "symlink": + outside = tmp_path / "outside.lock" + outside.write_text("outside") + outside.chmod(0o640) + lock_path.symlink_to(outside) + else: + os.mkfifo(lock_path) + + with pytest.raises(SlurmStateError, match="cannot initialize"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + if lock_kind == "symlink": + assert stat.S_IMODE(outside.stat().st_mode) == 0o640 + + +def test_temporary_file_permission_failure_leaves_no_partial_record( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + original_fchmod = state_filesystem.os.fchmod + regular_file_calls = 0 + + def fail_second_regular_file(descriptor: int, mode: int) -> None: + nonlocal regular_file_calls + if stat.S_ISREG(os.fstat(descriptor).st_mode): + regular_file_calls += 1 + if regular_file_calls == 2: + raise OSError("simulated permission failure") + original_fchmod(descriptor, mode) + + monkeypatch.setattr(state_filesystem.os, "fchmod", fail_second_regular_file) + with pytest.raises(SlurmStateError, match="cannot initialize"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + assert not tuple((case.workspace / "runs").rglob(".state.*.tmp")) + assert not (case.writer.run_root / "authored-config.json").exists() + + +def test_run_lock_rejects_path_replacement_during_acquisition( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + lock_path = case.workspace / "runs" / ".locks" / f"run-{case.plan.run_id}.lock" + original_flock = slurm_filesystem.fcntl.flock + replaced = False + + def replace_acquired_lock(descriptor: int, operation: int) -> None: + nonlocal replaced + original_flock(descriptor, operation) + if operation == slurm_filesystem.fcntl.LOCK_EX and not replaced: + replaced = True + lock_path.unlink() + lock_path.write_text("") + lock_path.chmod(0o600) + + monkeypatch.setattr(slurm_filesystem.fcntl, "flock", replace_acquired_lock) + with pytest.raises(SlurmStateError, match="cannot initialize"): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + assert not (case.writer.run_root / "run.json").exists() + + +def test_run_lock_closes_descriptor_when_acquisition_is_interrupted( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _build_case(tmp_path, authored_run_single, single_node_plan) + lock_descriptor: int | None = None + + def interrupt_flock(descriptor: int, operation: int) -> None: + nonlocal lock_descriptor + del operation + lock_descriptor = descriptor + raise KeyboardInterrupt + + monkeypatch.setattr(slurm_filesystem.fcntl, "flock", interrupt_flock) + with pytest.raises(KeyboardInterrupt): + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + + assert lock_descriptor is not None + with pytest.raises(OSError): + os.fstat(lock_descriptor) + + +def test_reader_rejects_package_directory_with_loosened_permissions( + 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.run_root.chmod(0o755) + + with pytest.raises(StateCorruptionError, match="cannot load"): + case.writer.load_run() + + +def test_attempt_creation_recovers_an_unpublished_next_directory_and_updates_monotonically( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + partial_root = case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + partial_root.mkdir(mode=0o700) + + assert case.writer.load_attempts(attempt.shard_id) == () + assert case.writer.create_attempt(attempt) == attempt + assert case.writer.create_attempt(attempt) == attempt + + submitted = _validated_copy( + attempt, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + updated_at=case.created_at + timedelta(minutes=2), + ) + assert case.writer.update_attempt(submitted) == submitted + assert ( + SlurmStateWriter(case.workspace, case.plan.run_id).load_attempt(attempt.shard_id, attempt.attempt_id) + == submitted + ) + + with pytest.raises(StateConflictError, match="monotonic transition"): + case.writer.update_attempt(attempt) + + +def test_attempt_reader_retries_when_publication_wins_missing_file_race( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + attempt_root = case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + attempt_root.mkdir(mode=0o700) + original_read = state_storage.read_regular_text + published = False + + def publish_before_reporting_missing( + directory_descriptor: int, + name: str, + display_path: Path, + *, + maximum_size: int, + ) -> str: + nonlocal published + if name == "attempt.json" and not published: + published = True + display_path.write_text(attempt.serialize_json()) + display_path.chmod(0o600) + raise FileNotFoundError(name) + return original_read(directory_descriptor, name, display_path, maximum_size=maximum_size) + + monkeypatch.setattr(state_storage, "read_regular_text", publish_before_reporting_missing) + + assert case.writer.load_attempts(attempt.shard_id) == (attempt,) + assert published + + +def test_attempt_recovery_rejects_unexpected_unpublished_records( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + partial_root = case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + partial_root.mkdir(mode=0o700) + (partial_root / "readiness.json").write_text("{}") + (partial_root / "readiness.json").chmod(0o600) + + with pytest.raises(StateCorruptionError, match="unpublished state records"): + case.writer.load_attempts(attempt.shard_id) + + +def test_public_attempt_readers_normalize_invalid_and_missing_identities( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + + with pytest.raises(SlurmStateError, match="invalid shard identity"): + case.writer.load_attempts(cast(ShardId, "invalid")) + with pytest.raises(StateNotFoundError, match="shard"): + case.writer.load_attempts(cast(ShardId, "shard-99999")) + with pytest.raises(StateConflictError, match="next shard ordinal"): + case.writer.create_attempt(_attempt(case, ordinal=2)) + + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with pytest.raises(SlurmStateError, match="invalid attempt identity"): + case.writer.load_attempt(attempt.shard_id, cast(AttemptId, "invalid")) + with pytest.raises(StateNotFoundError, match="attempt"): + case.writer.load_attempt(attempt.shard_id, cast(AttemptId, "attempt-9999")) + with pytest.raises(StateNotFoundError, match="no readiness snapshot"): + case.writer.load_readiness(attempt.shard_id, attempt.attempt_id) + + +@pytest.mark.parametrize("collection_root", ("shards", "attempts")) +def test_collection_readers_reject_hidden_unowned_entries( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + collection_root: str, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + root = case.writer.run_root / "shards" + if collection_root == "attempts": + root = root / "shard-00000" / "attempts" + (root / ".unowned").write_text("unexpected") + (root / ".unowned").chmod(0o600) + + with pytest.raises(StateCorruptionError): + if collection_root == "shards": + case.writer.load_shards() + else: + case.writer.load_attempts("shard-00000") + + +def test_concurrent_attempt_updates_choose_one_scheduler_identity( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + case.writer.create_attempt(attempt) + barrier = Barrier(2) + + def submit(array_job_id: int) -> AttemptManifest: + candidate = _validated_copy( + attempt, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=array_job_id, array_task_id=0), + updated_at=case.created_at + timedelta(minutes=2), + ) + barrier.wait() + return SlurmStateWriter(case.workspace, case.plan.run_id).update_attempt(candidate) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(submit, job_id) for job_id in (4101, 4102)) + + outcomes = tuple(future.result() for future in futures if future.exception() is None) + failures = tuple(future.exception() for future in futures if future.exception() is not None) + assert len(outcomes) == 1 + assert len(failures) == 1 + assert isinstance(failures[0], StateConflictError) + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) == outcomes[0] + + +def test_scheduler_identity_cannot_be_reused_by_a_later_attempt( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + first = _submitted_attempt(case) + case.writer.create_attempt(first) + second = _attempt(case, ordinal=2) + case.writer.create_attempt(second) + reused_scheduler = _validated_copy( + second, + state=AttemptLifecycleState.SUBMITTED, + scheduler=first.scheduler, + updated_at=case.created_at + timedelta(minutes=4), + ) + + with pytest.raises(StateConflictError, match="monotonic transition"): + case.writer.update_attempt(reused_scheduler) + + assert case.writer.load_attempt(second.shard_id, second.attempt_id) == second + + +def test_attempt_reader_revalidates_global_scheduler_ownership( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + first = _submitted_attempt(case) + case.writer.create_attempt(first) + second = _attempt(case, ordinal=2) + case.writer.create_attempt(second) + conflicting = _validated_copy( + second, + state=AttemptLifecycleState.SUBMITTED, + scheduler=first.scheduler, + updated_at=case.created_at + timedelta(minutes=4), + ) + attempt_path = case.writer.run_root / "shards" / second.shard_id / "attempts" / second.attempt_id / "attempt.json" + attempt_path.write_text(conflicting.serialize_json()) + + with pytest.raises(StateCorruptionError, match="invalid persisted attempts"): + SlurmStateWriter(case.workspace, case.plan.run_id).load_attempts(second.shard_id) + + +def test_failed_atomic_attempt_replacement_preserves_previous_manifest( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + case.writer.create_attempt(attempt) + submitted = _validated_copy( + attempt, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + updated_at=case.created_at + timedelta(minutes=2), + ) + + def fail_replace(*args: object, **kwargs: object) -> None: + del args, kwargs + raise OSError("simulated replacement failure") + + monkeypatch.setattr(state_filesystem.os, "replace", fail_replace) + with pytest.raises(SlurmStateError, match="cannot update"): + case.writer.update_attempt(submitted) + + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) == attempt + assert not tuple((case.writer.run_root / "shards").rglob(".state.*.tmp")) + + +def test_attempt_update_retry_converges_after_post_replace_fsync_failure( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _attempt(case) + case.writer.create_attempt(attempt) + submitted = _validated_copy( + attempt, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + updated_at=case.created_at + timedelta(minutes=2), + ) + original_fsync = state_filesystem.os.fsync + directory_syncs = 0 + + def fail_first_directory_fsync(descriptor: int) -> None: + nonlocal directory_syncs + if stat.S_ISDIR(os.fstat(descriptor).st_mode): + directory_syncs += 1 + if directory_syncs == 1: + raise OSError("simulated post-replace fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_first_directory_fsync) + with pytest.raises(SlurmStateError, match="cannot update"): + case.writer.update_attempt(submitted) + + assert case.writer.update_attempt(submitted) == submitted + assert directory_syncs == 2 + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) == submitted + + +def test_concurrent_conflicting_attempt_creation_publishes_one_candidate( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + first = _attempt(case) + second = _validated_copy( + first, + created_at=first.created_at + timedelta(seconds=1), + updated_at=first.updated_at + timedelta(seconds=1), + ) + barrier = Barrier(2) + + def create(candidate: AttemptManifest) -> AttemptManifest: + barrier.wait() + return SlurmStateWriter(case.workspace, case.plan.run_id).create_attempt(candidate) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(create, candidate) for candidate in (first, second)) + + outcomes = tuple(future.result() for future in futures if future.exception() is None) + failures = tuple(future.exception() for future in futures if future.exception() is not None) + assert len(outcomes) == 1 + assert len(failures) == 1 + assert isinstance(failures[0], StateConflictError) + assert case.writer.load_attempt(first.shard_id, first.attempt_id) == outcomes[0] + + +def test_readiness_revisions_are_idempotent_monotonic_and_atomic( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + initial = _readiness(case, attempt) + + assert case.writer.write_readiness(initial) == initial + assert case.writer.write_readiness(initial) == initial + + starting_deployment = _validated_copy( + initial.deployments[0], + state=ReadinessState.STARTING, + ) + revision_two = _validated_copy( + initial, + revision=2, + updated_at=case.created_at + timedelta(minutes=4), + state=ReadinessState.STARTING, + deployments=(starting_deployment,), + ) + assert case.writer.write_readiness(revision_two) == revision_two + assert ( + SlurmStateWriter(case.workspace, case.plan.run_id).load_readiness(attempt.shard_id, attempt.attempt_id) + == revision_two + ) + + skipped = _validated_copy(revision_two, revision=4, updated_at=case.created_at + timedelta(minutes=5)) + with pytest.raises(StateConflictError, match="monotonic transition"): + case.writer.write_readiness(skipped) + assert case.writer.load_readiness(attempt.shard_id, attempt.attempt_id) == revision_two + + +def test_readiness_reader_rebinds_snapshot_to_its_attempt( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + initial = _readiness(case, attempt) + case.writer.write_readiness(initial) + mismatched = _validated_copy(initial, run_id="other-run") + readiness_path = ( + case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id / "readiness.json" + ) + readiness_path.write_text(mismatched.serialize_json()) + + with pytest.raises(StateCorruptionError, match="invalid persisted readiness"): + SlurmStateWriter(case.workspace, case.plan.run_id).load_readiness(attempt.shard_id, attempt.attempt_id) + + +def test_idempotent_readiness_write_revalidates_persisted_snapshot( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + initial = _readiness(case, attempt) + case.writer.write_readiness(initial) + mismatched_deployment = _validated_copy(initial.deployments[0], model_alias="other-model") + mismatched = _validated_copy(initial, deployments=(mismatched_deployment,)) + readiness_path = ( + case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id / "readiness.json" + ) + readiness_path.write_text(mismatched.serialize_json()) + + with pytest.raises(StateCorruptionError, match="invalid persisted readiness"): + case.writer.write_readiness(mismatched) + + +def test_readiness_retry_converges_after_post_replace_fsync_failure( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + initial = _readiness(case, attempt) + case.writer.write_readiness(initial) + starting = _validated_copy(initial.deployments[0], state=ReadinessState.STARTING) + revision_two = _validated_copy( + initial, + revision=2, + updated_at=case.created_at + timedelta(minutes=4), + state=ReadinessState.STARTING, + deployments=(starting,), + ) + original_fsync = state_filesystem.os.fsync + directory_syncs = 0 + + def fail_first_directory_fsync(descriptor: int) -> None: + nonlocal directory_syncs + if stat.S_ISDIR(os.fstat(descriptor).st_mode): + directory_syncs += 1 + if directory_syncs == 1: + raise OSError("simulated post-replace fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_first_directory_fsync) + with pytest.raises(SlurmStateError, match="cannot persist readiness"): + case.writer.write_readiness(revision_two) + + assert case.writer.write_readiness(revision_two) == revision_two + assert directory_syncs == 2 + assert case.writer.load_readiness(attempt.shard_id, attempt.attempt_id) == revision_two + + +def test_concurrent_readiness_writers_cannot_publish_the_same_next_revision( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + initial = _readiness(case, attempt) + case.writer.write_readiness(initial) + barrier = Barrier(2) + + def publish(seconds: int) -> AttemptReadiness: + starting = _validated_copy(initial.deployments[0], state=ReadinessState.STARTING) + candidate = _validated_copy( + initial, + revision=2, + updated_at=case.created_at + timedelta(minutes=4, seconds=seconds), + state=ReadinessState.STARTING, + deployments=(starting,), + ) + barrier.wait() + return SlurmStateWriter(case.workspace, case.plan.run_id).write_readiness(candidate) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(publish, seconds) for seconds in (1, 2)) + + outcomes = tuple(future.result() for future in futures if future.exception() is None) + failures = tuple(future.exception() for future in futures if future.exception() is not None) + assert len(outcomes) == 1 + assert len(failures) == 1 + assert isinstance(failures[0], StateConflictError) + assert case.writer.load_readiness(attempt.shard_id, attempt.attempt_id) == outcomes[0] + + +def test_candidate_finalization_publishes_one_reloadable_winner_and_seals_the_shard( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + + winner = case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + assert winner.candidate_manifest == finalization.client_result.candidate_output_manifest + assert case.writer.load_winner(finalization.attempt.shard_id) == winner + assert stat.S_IMODE((case.writer.run_root / "shards/shard-00000/winner.json").stat().st_mode) == 0o600 + attempt_root = _attempt_root(case, finalization.attempt) + assert stat.S_IMODE((attempt_root / "client-result.json").stat().st_mode) == 0o600 + assert stat.S_IMODE((attempt_root / "output-manifest.json").stat().st_mode) == 0o600 + assert case.writer.publish_attempt_result(finalization.client_result, finalization.candidate) == ( + finalization.client_result, + finalization.candidate, + ) + assert ( + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at + timedelta(minutes=1), + ) + == winner + ) + assert case.writer.create_attempt(finalization.attempt) == finalization.attempt + with pytest.raises(StateConflictError, match="immutable winner"): + case.writer.create_attempt(_attempt(case, ordinal=2)) + with pytest.raises(StateConflictError, match="immutable winner"): + with case.writer.acquire_dataset_workspace( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + "never", + ): + pass + with pytest.raises(StateConflictError, match="timezone-aware UTC"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at.replace(tzinfo=None), + ) + + +def test_result_publication_binds_candidate_before_success_transition( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + + bound_attempt = case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) + assert bound_attempt.state is AttemptLifecycleState.RUNNING + assert bound_attempt.candidate_output == finalization.client_result.candidate_output_manifest + assert bound_attempt.candidate_output is not None + + different_reference = bound_attempt.candidate_output.model_copy(update={"sha256": "d" * 64}) + conflicting_success = _validated_copy( + bound_attempt, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + candidate_output=different_reference, + updated_at=case.created_at + timedelta(minutes=5), + ) + with pytest.raises(StateConflictError, match="valid monotonic transition"): + case.writer.update_attempt(conflicting_success) + + +def test_winner_publication_failure_restores_running_attempt_for_retry( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + original_publish = case.writer._storage.publish_winner + + def fail_winner_publication(winner: ShardWinner) -> None: + raise OSError("injected winner publication failure") + + monkeypatch.setattr(case.writer._storage, "publish_winner", fail_winner_publication) + completed_at = case.created_at + timedelta(minutes=5) + with pytest.raises(SlurmStateError, match="cannot finalize"): + case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=finalization.published_at, + ) + + persisted = case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) + assert persisted.state is AttemptLifecycleState.RUNNING + assert persisted.candidate_output == finalization.client_result.candidate_output_manifest + + monkeypatch.setattr(case.writer._storage, "publish_winner", original_publish) + winner = case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=completed_at, + published_at=finalization.published_at, + ) + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id).state is AttemptLifecycleState.SUCCEEDED + assert case.writer.load_winner(attempt.shard_id) == winner + + +def test_fresh_writer_resumes_finalization_interrupted_after_success_commit( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + original_replace = case.writer._storage.replace_attempt + + def interrupt_after_success_commit(updated: AttemptManifest) -> None: + original_replace(updated) + if updated.state is AttemptLifecycleState.SUCCEEDED: + raise KeyboardInterrupt("injected process interruption") + + monkeypatch.setattr(case.writer._storage, "replace_attempt", interrupt_after_success_commit) + with pytest.raises(KeyboardInterrupt, match="process interruption"): + case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + resumed = SlurmStateWriter(case.workspace, case.plan.run_id) + persisted = resumed.load_attempt(attempt.shard_id, attempt.attempt_id) + assert persisted.state is AttemptLifecycleState.SUCCEEDED + with pytest.raises(StateNotFoundError): + resumed.load_winner(attempt.shard_id) + + winner = resumed.resume_incomplete_finalization( + attempt.shard_id, + published_at=finalization.published_at, + ) + + assert winner is not None + assert resumed.load_winner(attempt.shard_id) == winner + + +def test_interrupted_finalization_resumes_through_nested_attempt_mount( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + base_case = _build_case(tmp_path, authored_run_single, single_node_plan) + logical_attempts = base_case.workspace / "runs" / base_case.plan.run_id / "shards" / "shard-00000" / "attempts" + fast_attempts = tmp_path / "fast-attempts" + mounts = ( + *base_case.plan.container_mounts, + ContainerMount(source=logical_attempts.as_posix(), target=fast_attempts.as_posix()), + ) + profile = SlurmProfile.model_validate( + base_case.plan.selected_profile.profile.model_dump(mode="python") | {"container_mounts": list(mounts)} + ) + selected_profile = base_case.plan.selected_profile.model_copy( + update={ + "profile": profile, + "profile_sha256": compute_canonical_json_sha256(profile.model_dump(mode="json")), + } + ) + plan = ResolvedSlurmRunPlan.model_validate( + base_case.plan.model_dump(mode="python") | {"selected_profile": selected_profile, "container_mounts": mounts} + ) + plan_reference = base_case.run.resolved_plan.model_copy(update={"sha256": plan.compute_sha256()}) + run = base_case.run.model_copy(update={"resolved_plan": plan_reference}) + writer = SlurmStateWriter( + base_case.workspace, + plan.run_id, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), + ) + case = _StateCase( + workspace=base_case.workspace, + authored_config=base_case.authored_config, + plan=plan, + run=run, + shards=base_case.shards, + writer=writer, + created_at=base_case.created_at, + ) + writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + attempt = _submitted_attempt(case).model_copy(update={"resolved_plan": plan_reference}) + writer.create_attempt(attempt) + logical_dataset = logical_attempts / attempt.attempt_id / "dataset" + with writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + assert dataset_path == fast_attempts / attempt.attempt_id / "dataset" + dataset_path.mkdir(parents=True, mode=0o700) + finalization = _persist_complete_result( + case, + attempt, + dataset_path, + manifest_dataset_path=logical_dataset, + complete_attempt=False, + ) + original_replace = writer._storage.replace_attempt + + def interrupt_after_success_commit(updated: AttemptManifest) -> None: + original_replace(updated) + if updated.state is AttemptLifecycleState.SUCCEEDED: + raise KeyboardInterrupt("injected process interruption") + + monkeypatch.setattr(writer._storage, "replace_attempt", interrupt_after_success_commit) + with pytest.raises(KeyboardInterrupt, match="process interruption"): + writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + resumed = SlurmStateWriter( + case.workspace, + plan.run_id, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), + ) + winner = resumed.resume_incomplete_finalization( + attempt.shard_id, + published_at=finalization.published_at, + ) + + assert winner is not None + assert resumed.load_winner(attempt.shard_id) == winner + + +def test_runtime_finalization_converges_after_committed_winner_sync_failure( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + original_fsync = state_filesystem.os.fsync + failed = False + + def fail_after_winner_link(descriptor: int) -> None: + nonlocal failed + if winner_path.exists() and not failed: + failed = True + raise OSError("injected winner directory fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_after_winner_link) + winner = case.writer.finalize_winner( + attempt.shard_id, + attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=5), + published_at=finalization.published_at, + ) + + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id).state is AttemptLifecycleState.SUCCEEDED + assert case.writer.load_winner(attempt.shard_id) == winner + + +def test_attempt_update_cannot_bind_candidate_before_result_publication( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + running = _validated_copy( + attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=2), + ) + case.writer.update_attempt(running) + candidate_reference = ArtifactReference( + path=(_attempt_root(case, attempt) / "output-manifest.json").as_posix(), + sha256="d" * 64, + ) + + with pytest.raises(StateConflictError, match="valid monotonic transition"): + case.writer.update_attempt(_validated_copy(running, candidate_output=candidate_reference)) + + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) == running + + +def test_result_publication_retry_finishes_candidate_binding_after_interruption( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + running = _validated_copy( + attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=2, seconds=30), + ) + case.writer.update_attempt(running) + original_replace = case.writer._storage.replace_attempt + + def interrupt_candidate_binding(_attempt: AttemptManifest) -> None: + raise OSError("injected attempt binding interruption") + + monkeypatch.setattr(case.writer._storage, "replace_attempt", interrupt_candidate_binding) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + with pytest.raises(SlurmStateError, match="cannot publish result"): + _persist_complete_result(case, attempt, dataset_path, complete_attempt=False) + + unbound_attempt = case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) + assert unbound_attempt.state is AttemptLifecycleState.RUNNING + assert unbound_attempt.candidate_output is None + client_result, candidate = case.writer._storage.read_finalization_records(attempt.shard_id, attempt.attempt_id) + + monkeypatch.setattr(case.writer._storage, "replace_attempt", original_replace) + assert case.writer.publish_attempt_result(client_result, candidate) == (client_result, candidate) + assert case.writer.load_attempt(attempt.shard_id, attempt.attempt_id).candidate_output == ( + client_result.candidate_output_manifest + ) + + +def test_result_publication_rejects_a_successful_attempt_bound_to_another_candidate( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + assert finalization.client_result.candidate_output_manifest is not None + different_candidate = finalization.candidate.model_copy(update={"dataset_schema_digest": "d" * 64}) + different_reference = finalization.client_result.candidate_output_manifest.model_copy( + update={"sha256": different_candidate.compute_sha256()} + ) + different_result = finalization.client_result.model_copy(update={"candidate_output_manifest": different_reference}) + + with pytest.raises(StateConflictError, match="persisted run intent"): + case.writer.publish_attempt_result(different_result, different_candidate) + + assert case.writer.publish_attempt_result(finalization.client_result, finalization.candidate) == ( + finalization.client_result, + finalization.candidate, + ) + + +def test_candidate_verifier_leases_manifest_bounded_files_through_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + dataset_path = tmp_path / "dataset" + dataset_path.mkdir(mode=0o700) + table = lazy.pa.table({"record_id": [1]}) + files: list[CandidateOutputFile] = [] + for index in range(16): + output_path = dataset_path / f"part-{index:05d}.parquet" + lazy.pq.write_table(table, output_path) + output_path.chmod(0o644) + content = output_path.read_bytes() + files.append( + CandidateOutputFile( + relative_path=output_path.name, + sha256=hashlib.sha256(content).hexdigest(), + byte_size=len(content), + record_count=1, + ) + ) + candidate = CandidateOutputManifest( + schema_version=1, + run_id="run-bounded", + shard_id="shard-00000", + attempt_id="attempt-0001", + attempt_ordinal=1, + created_at=datetime(2026, 9, 1, tzinfo=timezone.utc), + dataset_path=dataset_path.as_posix(), + requested_records=len(files), + actual_records=len(files), + outcome=CandidateOutcome.COMPLETE, + files=tuple(files), + dataset_schema_digest=compute_candidate_schema_digest(table.schema), + provenance_digest="a" * 64, + ) + original_open = state_artifacts.open_verified_regular_file + active_files = 0 + maximum_active_files = 0 + + @contextmanager + def track_open_file( + directory_descriptor: int, + name: str, + display_path: Path, + *, + expected_size: int, + expected_sha256: str, + require_private: bool = True, + ) -> Iterator[int]: + nonlocal active_files, maximum_active_files + with original_open( + directory_descriptor, + name, + display_path, + expected_size=expected_size, + expected_sha256=expected_sha256, + require_private=require_private, + ) as descriptor: + active_files += 1 + maximum_active_files = max(maximum_active_files, active_files) + try: + yield descriptor + finally: + active_files -= 1 + + monkeypatch.setattr(state_artifacts, "open_verified_regular_file", track_open_file) + with state_artifacts.CandidateArtifactVerifier().verify(candidate) as artifacts: + assert active_files == len(files) + artifacts.rebind() + assert active_files == len(files) + + assert active_files == 0 + assert maximum_active_files == len(files) + + maximum_active_files = 0 + verifier = state_artifacts.CandidateArtifactVerifier() + snapshot = verifier.inspect(candidate) + verifier.rebind(candidate, snapshot) + + assert active_files == 0 + assert maximum_active_files == 1 + + replacement = dataset_path / "replacement.parquet" + replacement.write_bytes((dataset_path / files[0].relative_path).read_bytes()) + replacement.chmod(0o644) + os.replace(replacement, dataset_path / files[0].relative_path) + + with pytest.raises(OSError, match="changed during collection"): + verifier.rebind(candidate, snapshot) + + +def test_candidate_schema_digest_ignores_arrow_metadata() -> None: + schema = lazy.pa.schema([("record_id", lazy.pa.int64())]) + annotated = schema.with_metadata({b"producer": b"worker-a"}) + + assert compute_candidate_schema_digest(annotated) == compute_candidate_schema_digest(schema) + + +def test_finalization_rejects_mismatched_partial_result_and_candidate_file_drift( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + partial = ClientResult.model_validate( + finalization.client_result.model_dump(mode="python") + | { + "actual_records": finalization.client_result.requested_records - 1, + "outcome": ClientOutcome.PARTIAL, + "early_shutdown": True, + } + ) + _write_state_record( + _attempt_root(case, finalization.attempt) / "client-result.json", + partial.serialize_json(), + ) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + _write_state_record( + _attempt_root(case, finalization.attempt) / "client-result.json", + finalization.client_result.serialize_json(), + ) + finalization.output_path.write_bytes(b"changed candidate bytes") + finalization.output_path.chmod(0o600) + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + with pytest.raises(StateNotFoundError, match="no winner"): + case.writer.load_winner(finalization.attempt.shard_id) + + +@pytest.mark.parametrize("file_kind", ("symlink", "hardlink", "fifo", "writable")) +def test_finalization_rejects_unsafe_candidate_files( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + file_kind: str, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + output_path = finalization.output_path + if file_kind == "writable": + output_path.chmod(0o666) + else: + outside = tmp_path / "outside-output" + outside.write_bytes(output_path.read_bytes()) + outside.chmod(0o600) + output_path.unlink() + if file_kind == "symlink": + output_path.symlink_to(outside) + elif file_kind == "hardlink": + os.link(outside, output_path) + else: + os.mkfifo(output_path) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + +def test_finalization_rejects_candidate_path_replacement_after_digest_read( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + replacement = tmp_path / "replacement.parquet" + replacement.write_bytes(finalization.output_path.read_bytes()) + replacement.chmod(0o600) + original_fstat = state_filesystem.os.fstat + output_inode = finalization.output_path.stat().st_ino + candidate_fstat_calls = 0 + replaced = False + + def replace_after_descriptor_validation(descriptor: int) -> os.stat_result: + nonlocal candidate_fstat_calls, replaced + status = original_fstat(descriptor) + if status.st_ino == output_inode: + candidate_fstat_calls += 1 + if candidate_fstat_calls == 2: + os.replace(replacement, finalization.output_path) + replaced = True + return status + + monkeypatch.setattr(state_filesystem.os, "fstat", replace_after_descriptor_validation) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + assert replaced + + +def test_finalization_rebinds_dataset_tree_immediately_before_winner_publication( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + dataset_path = finalization.output_path.parent + moved_path = dataset_path.with_name("moved-dataset") + original_rebind = state_artifacts.VerifiedCandidateArtifacts.rebind + + def replace_before_rebind(artifacts: state_artifacts.VerifiedCandidateArtifacts) -> None: + dataset_path.rename(moved_path) + dataset_path.mkdir(mode=0o700) + original_rebind(artifacts) + + monkeypatch.setattr(state_artifacts.VerifiedCandidateArtifacts, "rebind", replace_before_rebind) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + assert not (case.writer.run_root / "shards/shard-00000/winner.json").exists() + + +def test_finalization_derives_parquet_row_counts_before_publication( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result( + case, + attempt, + dataset_path, + physical_records=case.plan.shards[0].requested_records - 1, + ) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + +def test_finalization_derives_parquet_schema_before_publication( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result( + case, + attempt, + dataset_path, + reported_schema_digest="d" * 64, + ) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + +@pytest.mark.parametrize("replace_parent_with_symlink", (False, True)) +def test_finalization_verifies_nested_candidate_directories( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + replace_parent_with_symlink: bool, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case, relative_path="partition/day-01/part-00000.parquet") + partition_path = finalization.output_path.parents[1] + if replace_parent_with_symlink: + outside = tmp_path / "outside-partition" + partition_path.rename(outside) + partition_path.symlink_to(outside, target_is_directory=True) + + with pytest.raises(StateConflictError, match="not eligible"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + else: + assert ( + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ).attempt_id + == finalization.attempt.attempt_id + ) + + +@pytest.mark.parametrize("record_name", ("client-result.json", "output-manifest.json")) +def test_finalization_classifies_unsafe_result_records_as_corruption( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + record_name: str, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + record_path = _attempt_root(case, finalization.attempt) / record_name + outside = tmp_path / record_name + outside.write_text(record_path.read_text()) + outside.chmod(0o600) + record_path.unlink() + record_path.symlink_to(outside) + + with pytest.raises(StateCorruptionError, match="unsafe or unreadable"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + +def test_concurrent_identical_finalizers_converge_on_the_same_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + barrier = Barrier(2) + + def finalize() -> ShardWinner: + barrier.wait() + writer = SlurmStateWriter(case.workspace, case.plan.run_id) + return writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + winners = tuple(executor.map(lambda _: finalize(), range(2))) + + assert winners[0] == winners[1] == case.writer.load_winner(finalization.attempt.shard_id) + + +def test_finalization_retry_converges_after_post_publish_fsync_failure( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + original_fsync = state_filesystem.os.fsync + failed = False + + def fail_after_winner_link(descriptor: int) -> None: + nonlocal failed + if winner_path.exists() and not failed: + failed = True + raise OSError("injected winner directory fsync failure") + original_fsync(descriptor) + + monkeypatch.setattr(state_filesystem.os, "fsync", fail_after_winner_link) + with pytest.raises(SlurmStateError, match="cannot finalize"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + monkeypatch.setattr(state_filesystem.os, "fsync", original_fsync) + winner = case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at + timedelta(minutes=1), + ) + assert winner == case.writer.load_winner(finalization.attempt.shard_id) + assert not tuple(winner_path.parent.glob(".state.*.tmp")) + + +def test_candidate_file_verification_does_not_hold_the_run_state_lock( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run, multi_node_plan) + finalization = _complete_finalization_case(case) + verification_started = Event() + release_verification = Event() + original_validate = state_finalization.WinnerFinalizer._validate_artifact_metadata + + def block_verification(candidate: CandidateOutputManifest, artifacts: object) -> None: + verification_started.set() + assert release_verification.wait(timeout=5) + original_validate(candidate, artifacts) + + monkeypatch.setattr( + state_finalization.WinnerFinalizer, + "_validate_artifact_metadata", + staticmethod(block_verification), + ) + + with ThreadPoolExecutor(max_workers=1) as executor: + finalizer = executor.submit( + case.writer.finalize_winner, + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + assert verification_started.wait(timeout=5) + other_shard = case.shards[1] + other_attempt = AttemptManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=other_shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.CREATED, + created_at=case.created_at + timedelta(minutes=1), + updated_at=case.created_at + timedelta(minutes=1), + ) + assert case.writer.create_attempt(other_attempt) == other_attempt + release_verification.set() + assert finalizer.result().attempt_id == finalization.attempt.attempt_id + + +def test_finalization_does_not_scan_unrelated_attempt_directories( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run, multi_node_plan) + finalization = _complete_finalization_case(case) + unrelated_attempts = case.writer.run_root / "shards" / "shard-00001" / "attempts" + (unrelated_attempts / "unexpected").mkdir(mode=0o700) + + winner = case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + + assert case.writer.load_winner(finalization.attempt.shard_id) == winner + with pytest.raises(StateCorruptionError, match="invalid attempt directory"): + case.writer.load_attempts("shard-00001") + + +def test_finalization_waits_for_the_dataset_writer_to_release_its_lease( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + finalization_started = Event() + finalization_finished = Event() + + def finalize(finalization: _FinalizationCase) -> ShardWinner: + finalization_started.set() + winner = SlurmStateWriter(case.workspace, case.plan.run_id).finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + finalization_finished.set() + return winner + + with ThreadPoolExecutor(max_workers=1) as executor: + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + finalization = _persist_complete_result(case, attempt, dataset_path) + future = executor.submit(finalize, finalization) + assert finalization_started.wait(timeout=5) + assert not finalization_finished.wait(timeout=0.1) + winner = future.result(timeout=5) + + assert winner == case.writer.load_winner(attempt.shard_id) + + +def test_winner_reader_rejects_a_reference_to_an_unknown_attempt( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + winner = case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + corrupted = winner.model_copy(update={"attempt_id": "attempt-9999"}) + _write_state_record(case.writer.run_root / "shards/shard-00000/winner.json", corrupted.serialize_json()) + + with pytest.raises(StateCorruptionError, match="unknown attempt"): + case.writer.load_winner(finalization.attempt.shard_id) + with pytest.raises(StateCorruptionError, match="unknown attempt"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + with pytest.raises(StateCorruptionError, match="unknown attempt"): + case.writer.create_attempt(_attempt(case, ordinal=2)) + with pytest.raises(StateCorruptionError, match="unknown attempt"): + with case.writer.acquire_dataset_workspace( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + "never", + ): + pass + + +def test_semantically_corrupted_winner_is_corruption_at_every_writer_boundary( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + finalization = _complete_finalization_case(case) + winner = case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + corrupted = winner.model_copy( + update={ + "candidate_manifest": winner.candidate_manifest.model_copy( + update={"sha256": "d" * 64}, + ) + } + ) + _write_state_record(case.writer.run_root / "shards/shard-00000/winner.json", corrupted.serialize_json()) + + with pytest.raises(StateCorruptionError, match="winner chain"): + case.writer.load_winner(finalization.attempt.shard_id) + with pytest.raises(StateCorruptionError, match="winner chain"): + case.writer.finalize_winner( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + published_at=finalization.published_at, + ) + with pytest.raises(StateCorruptionError, match="winner chain"): + case.writer.create_attempt(_attempt(case, ordinal=2)) + with pytest.raises(StateCorruptionError, match="winner chain"): + with case.writer.acquire_dataset_workspace( + finalization.attempt.shard_id, + finalization.attempt.attempt_id, + "never", + ): + pass + + +def test_resumable_workspace_is_shard_owned_and_exclusively_locked( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan_payload = single_node_plan.model_dump(mode="python") + plan_payload["invocation"]["authored"]["resume"] = "always" + resumable_plan = ResolvedSlurmRunPlan.model_validate(plan_payload) + case = _initialized_case(tmp_path, authored_run_single, resumable_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + first_entered = Event() + release_first = Event() + second_entered = Event() + + def first_writer() -> Path: + writer = SlurmStateWriter(case.workspace, case.plan.run_id) + with writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "always") as path: + first_entered.set() + assert release_first.wait(timeout=5) + return path + + def second_writer() -> Path: + assert first_entered.wait(timeout=5) + writer = SlurmStateWriter(case.workspace, case.plan.run_id) + with writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "always") as path: + second_entered.set() + return path + + with ThreadPoolExecutor(max_workers=2) as executor: + first_future = executor.submit(first_writer) + second_future = executor.submit(second_writer) + assert first_entered.wait(timeout=5) + assert not second_entered.wait(timeout=0.1) + release_first.set() + assert first_future.result() == Path(case.shards[0].resume_workspace.path) + assert second_future.result() == Path(case.shards[0].resume_workspace.path) + + dataset_path = Path(case.shards[0].resume_workspace.path) + assert stat.S_IMODE(dataset_path.stat().st_mode) == 0o700 + with pytest.raises(StateConflictError, match="requested dataset workspace"): + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never"): + pass + + +def test_if_possible_lease_does_not_create_resume_workspace( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + plan_payload = single_node_plan.model_dump(mode="python") + plan_payload["invocation"]["authored"]["resume"] = "if_possible" + case = _initialized_case(tmp_path, authored_run_single, ResolvedSlurmRunPlan.model_validate(plan_payload)) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + resume_path = Path(case.shards[0].resume_workspace.path) + + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "if_possible") as path: + assert path == resume_path + assert not path.exists() + + assert not resume_path.exists() + + +def test_dataset_lock_rejects_unsafe_files_without_reclassifying_body_errors( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + lock_path = case.writer.run_root / "shards/shard-00000/resume.lock" + lock_path.unlink() + outside = tmp_path / "outside.lock" + outside.write_text("unchanged") + lock_path.symlink_to(outside) + + with pytest.raises(SlurmStateError, match="cannot lock dataset workspace"): + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never"): + pass + assert outside.read_text() == "unchanged" + + lock_path.unlink() + with pytest.raises(FileNotFoundError, match="body failure"): + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never"): + raise FileNotFoundError("body failure") + + +def test_dataset_lock_interruption_closes_an_open_shard_context( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + shard_closed = False + + @contextmanager + def open_shard(shard_id: ShardId) -> Iterator[int]: + nonlocal shard_closed + del shard_id + try: + yield 17 + finally: + shard_closed = True + + @contextmanager + def interrupt_lock(directory_descriptor: int, name: str, display_path: Path) -> Iterator[None]: + del directory_descriptor, name, display_path + raise KeyboardInterrupt("injected lock interruption") + yield + + monkeypatch.setattr(case.writer._storage, "open_shard_directory", open_shard) + monkeypatch.setattr(state_storage, "acquire_file_lock", interrupt_lock) + + with pytest.raises(KeyboardInterrupt, match="lock interruption"): + with case.writer.acquire_dataset_workspace("shard-00000", "attempt-0001", "never"): + pass + + assert shard_closed + + +def test_fresh_process_loads_only_persisted_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + command = ( + "import sys; " + "from data_designer.slurm.state import SlurmStateWriter; " + "record = SlurmStateWriter(sys.argv[1], sys.argv[2]).load_run(); " + "print(record.serialize_json(), end='')" + ) + + completed = subprocess.run( + (sys.executable, "-c", command, case.workspace.as_posix(), case.plan.run_id), + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout == case.run.serialize_json() + + +def test_state_package_defers_writer_dependencies_until_access() -> None: + command = ( + "import sys; " + "import data_designer.slurm.state as state; " + "assert 'data_designer.slurm.state.store' not in sys.modules; " + "assert 'data_designer.slurm.config' not in sys.modules; " + "assert 'SlurmStateWriter' in dir(state)" + ) + + completed = subprocess.run( + (sys.executable, "-c", command), + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_state_persistence_does_not_import_integration_layer() -> None: + command = ( + "import sys; " + "import data_designer.slurm.state.store; " + "assert 'data_designer.slurm.integration' not in sys.modules" + ) + + completed = subprocess.run( + (sys.executable, "-c", command), + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + + +def _build_case( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> _StateCase: + workspace = tmp_path / "workspace" + workspace.mkdir() + relocated_plan = _relocate_plan(plan, workspace) + created_at = datetime(2026, 9, 1, 12, tzinfo=timezone.utc) + run_root = workspace / "runs" / relocated_plan.run_id + resolved_plan_reference = ArtifactReference( + path=(run_root / "resolved-plan.json").as_posix(), + sha256=relocated_plan.compute_sha256(), + ) + run = RunManifest( + schema_version=1, + run_id=relocated_plan.run_id, + created_at=created_at, + authored_config=relocated_plan.authored_config, + resolved_plan=resolved_plan_reference, + shard_count=len(relocated_plan.shards), + ) + 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 + ) + return _StateCase( + workspace=workspace, + authored_config=authored_config, + plan=relocated_plan, + run=run, + shards=shards, + writer=SlurmStateWriter(workspace, relocated_plan.run_id), + created_at=created_at, + ) + + +def _initialized_case( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> _StateCase: + case = _build_case(tmp_path, authored_config, plan) + case.writer.initialize_run(case.authored_config, case.plan, case.run, case.shards) + return case + + +def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: + previous_workspace = plan.selected_profile.profile.workspace_root + payload = cast( + dict[str, object], + json.loads(plan.serialize_json().replace(previous_workspace, workspace.as_posix())), + ) + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile = SlurmProfile.model_validate_json(json.dumps(selected_profile["profile"])) + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile.model_dump(mode="json")) + return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def _attempt(case: _StateCase, *, ordinal: int = 1) -> AttemptManifest: + return AttemptManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id="shard-00000", + attempt_id=f"attempt-{ordinal:04d}", + attempt_ordinal=ordinal, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.CREATED, + created_at=case.created_at + timedelta(minutes=ordinal), + updated_at=case.created_at + timedelta(minutes=ordinal), + ) + + +def _submitted_attempt(case: _StateCase) -> AttemptManifest: + attempt = _attempt(case) + return _validated_copy( + attempt, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0), + updated_at=case.created_at + timedelta(minutes=2), + ) + + +def _readiness(case: _StateCase, attempt: AttemptManifest) -> AttemptReadiness: + 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 + ) + return AttemptReadiness( + schema_version=1, + run_id=attempt.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=case.created_at + timedelta(minutes=3), + state=ReadinessState.PENDING, + deployments=deployments, + ) + + +def _complete_finalization_case( + case: _StateCase, + *, + relative_path: str = "part-00000.parquet", +) -> _FinalizationCase: + attempt = _submitted_attempt(case) + case.writer.create_attempt(attempt) + with case.writer.acquire_dataset_workspace(attempt.shard_id, attempt.attempt_id, "never") as dataset_path: + return _persist_complete_result(case, attempt, dataset_path, relative_path=relative_path) + + +def _persist_complete_result( + case: _StateCase, + attempt: AttemptManifest, + dataset_path: Path, + *, + relative_path: str = "part-00000.parquet", + physical_records: int | None = None, + reported_schema_digest: str | None = None, + manifest_dataset_path: Path | None = None, + complete_attempt: bool = True, +) -> _FinalizationCase: + running_attempt = _validated_copy( + attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=2, seconds=30), + ) + case.writer.update_attempt(running_attempt) + output_path = dataset_path / relative_path + output_path.parent.mkdir(mode=0o755, parents=True, exist_ok=True) + for parent in output_path.parents: + if parent == dataset_path: + break + parent.chmod(0o755) + requested_records = case.plan.shards[0].requested_records + table = lazy.pa.table({"record_id": range(requested_records if physical_records is None else physical_records)}) + lazy.pq.write_table(table, output_path) + output_path.chmod(0o644) + content = output_path.read_bytes() + persisted_dataset_path = manifest_dataset_path or dataset_path + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + attempt_ordinal=attempt.attempt_ordinal, + created_at=case.created_at + timedelta(minutes=3), + dataset_path=persisted_dataset_path.as_posix(), + requested_records=requested_records, + actual_records=requested_records, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path=relative_path, + sha256=hashlib.sha256(content).hexdigest(), + byte_size=len(content), + record_count=requested_records, + ), + ), + dataset_schema_digest=reported_schema_digest or compute_candidate_schema_digest(table.schema), + provenance_digest=case.plan.compute_sha256(), + ) + candidate_path = _attempt_root(case, attempt) / "output-manifest.json" + candidate_reference = ArtifactReference( + path=candidate_path.as_posix(), + sha256=candidate.compute_sha256(), + ) + client_result = ClientResult( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + completed_at=case.created_at + timedelta(minutes=4), + requested_records=requested_records, + actual_records=requested_records, + outcome=ClientOutcome.COMPLETE, + dataset_path=persisted_dataset_path.as_posix(), + early_shutdown=False, + requested_resume_mode=case.plan.invocation.authored.resume, + effective_resume_mode="never", + candidate_output_manifest=candidate_reference, + ) + case.writer.publish_attempt_result(client_result, candidate) + bound_attempt = case.writer.load_attempt(attempt.shard_id, attempt.attempt_id) + assert bound_attempt.candidate_output == candidate_reference + if complete_attempt: + final_attempt = _validated_copy( + bound_attempt, + state=AttemptLifecycleState.SUCCEEDED, + terminal_classification=AttemptTerminalClassification.SUCCEEDED, + updated_at=case.created_at + timedelta(minutes=5), + ) + case.writer.update_attempt(final_attempt) + else: + final_attempt = bound_attempt + return _FinalizationCase( + attempt=final_attempt, + client_result=client_result, + candidate=candidate, + published_at=case.created_at + timedelta(minutes=6), + output_path=output_path, + ) + + +def _attempt_root(case: _StateCase, attempt: AttemptManifest) -> Path: + return case.writer.run_root / "shards" / attempt.shard_id / "attempts" / attempt.attempt_id + + +def _write_state_record(path: Path, content: str) -> None: + path.write_text(content) + path.chmod(0o600) + + +def _validated_copy(record: _ValueT, **updates: object) -> _ValueT: + payload = record.model_dump(mode="json") + payload.update(updates) + return type(record).model_validate_json(json.dumps(payload, default=_json_value)) + + +def _json_value(value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if isinstance(value, ContractValue): + return value.model_dump(mode="json") + raise TypeError(f"unsupported test JSON value: {type(value).__name__}") diff --git a/packages/data-designer-slurm/tests/state/test_submission_recovery.py b/packages/data-designer-slurm/tests/state/test_submission_recovery.py new file mode 100644 index 000000000..5e29db0d0 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_submission_recovery.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone + +import pytest + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.models import SlurmSubmissionMatch +from data_designer.slurm.state import StateConflictError +from data_designer.slurm.state.submission_recovery import PreparedSubmission, resolve_prepared_submission + + +@dataclass(frozen=True) +class _SubmissionLookup: + matches: tuple[SlurmSubmissionMatch, ...] + + def query_submissions_by_name( + self, + job_name: Identifier, + *, + submitted_after: datetime, + ) -> tuple[SlurmSubmissionMatch, ...]: + assert job_name == "dd-retry-0123456789abcdef0123456789abcdef" + assert submitted_after == datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + return self.matches + + +def test_prepared_submission_recovery_rejects_multiple_exact_matches() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + + with pytest.raises(StateConflictError, match="multiple scheduler jobs"): + resolve_prepared_submission( + _SubmissionLookup( + ( + SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + SlurmSubmissionMatch( + job_id=4301, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + ) + ), + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=submitted_at + timedelta(minutes=5), + expected_array_task_ids=(0,), + ), + observed_at=submitted_at + timedelta(minutes=1), + ) + + +def test_prepared_submission_recovery_returns_definitive_absence_only_after_deadline() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + deadline = submitted_at + timedelta(minutes=5) + lookup = _SubmissionLookup(()) + + with pytest.raises(StateConflictError, match="still being reconciled"): + resolve_prepared_submission( + lookup, + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0,), + ), + observed_at=deadline, + ) + assert ( + resolve_prepared_submission( + lookup, + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0,), + ), + observed_at=deadline + timedelta(microseconds=1), + ) + is None + ) + + +@pytest.mark.parametrize("actual_shape", [None, (0, 2)]) +def test_prepared_submission_recovery_rejects_the_wrong_scheduler_shape( + actual_shape: tuple[int, ...] | None, +) -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + match = SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=actual_shape, + ) + + with pytest.raises(StateConflictError, match="shape"): + resolve_prepared_submission( + _SubmissionLookup((match,)), + PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=submitted_at + timedelta(minutes=5), + expected_array_task_ids=(0, 1), + ), + observed_at=submitted_at + timedelta(minutes=1), + ) + + +def test_prepared_submission_recovery_bounds_a_partial_array_view() -> None: + submitted_at = datetime(2026, 9, 2, 12, tzinfo=timezone.utc) + deadline = submitted_at + timedelta(minutes=5) + lookup = _SubmissionLookup( + ( + SlurmSubmissionMatch( + job_id=4201, + job_name="dd-retry-0123456789abcdef0123456789abcdef", + array_task_ids=(0,), + ), + ) + ) + prepared = PreparedSubmission( + job_name="dd-retry-0123456789abcdef0123456789abcdef", + submitted_after=submitted_at, + reconciliation_deadline=deadline, + expected_array_task_ids=(0, 1), + ) + + with pytest.raises(StateConflictError, match="still being reconciled"): + resolve_prepared_submission(lookup, prepared, observed_at=deadline) + assert ( + resolve_prepared_submission( + lookup, + prepared, + observed_at=deadline + timedelta(microseconds=1), + ) + is None + ) diff --git a/packages/data-designer-slurm/tests/state/test_transition_fixtures.py b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py new file mode 100644 index 000000000..4bc1c7188 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from itertools import pairwise +from pathlib import Path + +from data_designer.slurm.state import AttemptReadiness, validate_readiness_transition + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +TRANSITION_FILES = ( + "single_node_pending_readiness.json", + "single_node_starting_readiness.json", + "single_node_readiness.json", + "single_node_failed_readiness.json", + "single_node_stopped_readiness.json", +) + + +def test_single_node_readiness_transition_fixture_is_canonical() -> None: + records = tuple( + AttemptReadiness.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) for filename in TRANSITION_FILES + ) + + assert tuple(record.revision for record in records) == (1, 2, 3, 4, 5) + for previous, current in pairwise(records): + assert validate_readiness_transition(previous, current) is current + + +def test_endpoint_publication_failure_transition_fixture_is_canonical() -> None: + starting = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_starting_readiness.json").read_text() + ) + failed = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_publication_failed_readiness.json").read_text() + ) + + assert validate_readiness_transition(starting, failed) is failed diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py new file mode 100644 index 000000000..6a254aedb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -0,0 +1,678 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import TypeVar + +import pytest + +from data_designer.slurm.state import ( + ArtifactReference, + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + CollectionPlan, + EffectiveAttemptState, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerObservation, + SchedulerState, + ShardManifest, + ShardWinner, + StateContractError, + StateRecord, + StateValue, + reconcile_attempt_observation, + validate_attempt_manifest, + validate_attempt_set, + validate_attempt_transition, + validate_collection_plan, + validate_readiness_transition, + validate_scheduler_observation_transition, + validate_shard_manifest, + validate_shard_set, + validate_shard_winner, +) + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +RecordT = TypeVar("RecordT", bound=StateRecord) +ValueT = TypeVar("ValueT", bound=StateValue) + + +def _load(model: type[RecordT], filename: str) -> RecordT: + return model.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) + + +def _json_value(value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if isinstance(value, StateValue): + return value.model_dump(mode="json") + if isinstance(value, tuple): + return [_json_value(item) for item in value] + return value + + +def _validated_copy(record: ValueT, **updates: object) -> ValueT: + payload = record.model_dump(mode="json") + payload.update({key: _json_value(value) for key, value in updates.items()}) + return type(record).model_validate_json(json.dumps(payload)) + + +def test_run_shard_attempt_and_winner_contracts_match() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + assert validate_shard_manifest(run, shard) is shard + assert validate_attempt_manifest(run, shard, attempt) is attempt + assert validate_shard_winner(run, shard, attempt, candidate, winner) is winner + + +def test_cross_record_identity_and_digest_mismatches_fail() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + with pytest.raises(StateContractError, match="run_id"): + validate_shard_manifest(run, _validated_copy(shard, run_id="another-run")) + + wrong_reference = ArtifactReference( + path=winner.candidate_manifest.path, + sha256="0" * 64, + ) + wrong_winner = _validated_copy(winner, candidate_manifest=wrong_reference) + with pytest.raises(StateContractError, match="candidate digest"): + validate_shard_winner(run, shard, attempt, candidate, wrong_winner) + + wrong_range = _validated_copy( + shard.record_range, + end_index_exclusive=shard.record_range.end_index_exclusive + 1, + ) + wrong_shard = _validated_copy(shard, record_range=wrong_range) + with pytest.raises(StateContractError, match="record count"): + validate_shard_winner(run, wrong_shard, attempt, candidate, winner) + + +def test_shard_set_rejects_overlapping_ranges_and_shared_resume_workspace() -> None: + run = _validated_copy(_load(RunManifest, "run_manifest.json"), shard_count=2) + first = _load(ShardManifest, "shard_manifest.json") + second_range = _validated_copy( + first.record_range, + start_index=first.record_range.end_index_exclusive, + end_index_exclusive=first.record_range.end_index_exclusive + first.record_range.record_count, + ) + second = _validated_copy( + first, + shard_id="shard-00001", + shard_index=1, + record_range=second_range, + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, + ) + + assert validate_shard_set(run, (first, second)) == (first, second) + + overlapping_range = _validated_copy( + second.record_range, + start_index=first.record_range.end_index_exclusive - 1, + ) + overlapping = _validated_copy(second, record_range=overlapping_range) + with pytest.raises(StateContractError, match="must not overlap"): + validate_shard_set(run, (first, overlapping)) + + shared_workspace = _validated_copy(second, resume_workspace=first.resume_workspace) + with pytest.raises(StateContractError, match="workspace paths must be unique"): + validate_shard_set(run, (first, shared_workspace)) + + +def test_attempt_set_rejects_scheduler_identity_collisions() -> None: + run = _validated_copy(_load(RunManifest, "run_manifest.json"), shard_count=2) + first_shard = _load(ShardManifest, "shard_manifest.json") + second_shard = _validated_copy( + first_shard, + shard_id="shard-00001", + shard_index=1, + record_range=_validated_copy(first_shard.record_range, start_index=100, end_index_exclusive=200), + resume_workspace={"path": "/workspace/runs/run-0001/shards/shard-00001/dataset"}, + ) + first_attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + second_attempt = _validated_copy( + first_attempt, + shard_id=second_shard.shard_id, + ) + assert first_attempt.scheduler is not None + + with pytest.raises(StateContractError, match="scheduler identities must be unique"): + validate_attempt_set(run, (first_shard, second_shard), (first_attempt, second_attempt)) + + unique_scheduler = _validated_copy(first_attempt.scheduler, array_task_id=1) + unique_second_attempt = _validated_copy(second_attempt, scheduler=unique_scheduler) + assert validate_attempt_set( + run, + (first_shard, second_shard), + (first_attempt, unique_second_attempt), + ) == (first_attempt, unique_second_attempt) + + +def test_attempt_transitions_preserve_identity_and_terminal_state() -> None: + succeeded = _load(AttemptManifest, "successful_attempt.json") + running = _validated_copy( + succeeded, + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + updated_at=datetime(2026, 8, 18, 12, 4, tzinfo=timezone.utc), + ) + + assert validate_attempt_transition(running, succeeded) is succeeded + assert validate_attempt_transition(succeeded, succeeded) is succeeded + with pytest.raises(StateContractError, match="state cannot move"): + validate_attempt_transition(running, _validated_copy(running, state=AttemptLifecycleState.PENDING)) + + changed_scheduler = _validated_copy(running.scheduler, array_job_id=9999) + with pytest.raises(StateContractError, match="scheduler identity cannot change"): + validate_attempt_transition(running, _validated_copy(running, scheduler=changed_scheduler)) + + candidate_running = _validated_copy(running, candidate_output=succeeded.candidate_output) + with pytest.raises(StateContractError, match="candidate output cannot change"): + validate_attempt_transition(candidate_running, _validated_copy(candidate_running, candidate_output=None)) + + with pytest.raises(StateContractError, match="terminal attempt manifests are immutable"): + validate_attempt_transition( + succeeded, + _validated_copy(succeeded, updated_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc)), + ) + + +def test_failed_attempts_and_existing_winners_are_rejected_while_partial_winners_are_allowed() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + attempt = _load(AttemptManifest, "successful_attempt.json") + candidate = _load(CandidateOutputManifest, "candidate_output.json") + winner = _load(ShardWinner, "shard_winner.json") + + failed_attempt = _validated_copy( + attempt, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.NODE_FAILED, + ) + with pytest.raises(StateContractError, match="successful attempts"): + validate_shard_winner(run, shard, failed_attempt, candidate, winner) + + partial_payload = json.loads((GOLDEN_DIRECTORY / "candidate_output.json").read_text()) + partial_payload.update( + { + "actual_records": 99, + "outcome": "partial", + } + ) + partial_payload["files"][0]["record_count"] = 99 + partial = CandidateOutputManifest.model_validate_json(json.dumps(partial_payload)) + partial_reference = ArtifactReference( + path=winner.candidate_manifest.path, + sha256=partial.compute_sha256(), + ) + partial_attempt = _validated_copy(attempt, candidate_output=partial_reference) + partial_winner = _validated_copy(winner, candidate_manifest=partial_reference) + assert validate_shard_winner(run, shard, partial_attempt, partial, partial_winner) == partial_winner + + with pytest.raises(StateContractError, match="immutable"): + validate_shard_winner( + run, + shard, + attempt, + candidate, + winner, + existing_winner=winner, + ) + + +def test_readiness_revisions_are_monotonic_and_preserve_authored_order() -> None: + current = _load(AttemptReadiness, "single_node_readiness.json") + starting_deployment = _validated_copy( + current.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, + ) + previous = _validated_copy( + current, + revision=2, + updated_at=datetime(2026, 8, 18, 12, 1, tzinfo=timezone.utc), + state=ReadinessState.STARTING, + deployments=(starting_deployment,), + ) + + assert validate_readiness_transition(previous, current) is current + with pytest.raises(StateContractError, match="revision"): + validate_readiness_transition(previous, _validated_copy(current, revision=2)) + + multi = _load(AttemptReadiness, "multi_node_readiness.json") + reordered = _validated_copy( + multi, + revision=5, + deployments=tuple(reversed(multi.deployments)), + ) + with pytest.raises(StateContractError, match="order or ID"): + validate_readiness_transition(multi, reordered) + + +def test_readiness_cannot_move_backward_or_change_backend_count() -> None: + ready = _load(AttemptReadiness, "single_node_readiness.json") + starting_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + backward = _validated_copy( + ready, + revision=4, + state=ReadinessState.STARTING, + deployments=(starting_deployment,), + ) + with pytest.raises(StateContractError, match="cannot move"): + validate_readiness_transition(ready, backward) + + changed_deployment = _validated_copy( + ready.deployments[0], + expected_backends=2, + ready_backends=2, + ) + changed_count = _validated_copy( + ready, + revision=4, + deployments=(changed_deployment,), + ) + with pytest.raises(StateContractError, match="backend count"): + validate_readiness_transition(ready, changed_count) + + +def test_readiness_can_reset_only_through_an_explicit_restart_transition() -> None: + ready = _load(AttemptReadiness, "single_node_readiness.json") + restarting_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.RESTARTING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, + ) + restarting = _validated_copy( + ready, + revision=ready.revision + 1, + updated_at=ready.updated_at + timedelta(seconds=1), + state=ReadinessState.RESTARTING, + deployments=(restarting_deployment,), + ) + + assert validate_readiness_transition(ready, restarting) is restarting + + starting = _validated_copy(restarting_deployment, state=ReadinessState.STARTING) + with pytest.raises(StateContractError, match="cannot move"): + validate_readiness_transition( + ready, + _validated_copy(restarting, state=ReadinessState.STARTING, deployments=(starting,)), + ) + + stopped_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STOPPED, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + unpublished = _validated_copy( + ready, + revision=ready.revision + 1, + state=ReadinessState.STOPPED, + deployments=(stopped_deployment,), + ) + with pytest.raises(StateContractError, match="endpoint publication"): + validate_readiness_transition(ready, unpublished) + + +def test_readiness_probe_evidence_cannot_regress_or_disappear() -> None: + previous = _load(AttemptReadiness, "single_node_readiness.json") + previous_probe = previous.deployments[0].last_probe + assert previous_probe is not None + + regressed_probe = _validated_copy( + previous_probe, + observed_at=datetime(2026, 8, 18, 12, 1, tzinfo=timezone.utc), + ) + regressed_deployment = _validated_copy(previous.deployments[0], last_probe=regressed_probe) + regressed = _validated_copy( + previous, + revision=previous.revision + 1, + updated_at=datetime(2026, 8, 18, 12, 3, tzinfo=timezone.utc), + deployments=(regressed_deployment,), + ) + with pytest.raises(StateContractError, match="probe observation cannot move backward"): + validate_readiness_transition(previous, regressed) + + removed_deployment = _validated_copy(previous.deployments[0], last_probe=None) + removed = _validated_copy( + previous, + revision=previous.revision + 1, + updated_at=datetime(2026, 8, 18, 12, 3, tzinfo=timezone.utc), + deployments=(removed_deployment,), + ) + with pytest.raises(StateContractError, match="probe evidence cannot be removed"): + validate_readiness_transition(previous, removed) + + +def test_collection_requires_exact_winner_set_and_digests() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + plan = _load(CollectionPlan, "collection_plan.json") + winner = _load(ShardWinner, "shard_winner.json") + + assert validate_shard_set(run, (shard,)) == (shard,) + assert validate_collection_plan(run, plan, (shard,), (winner,)) is plan + with pytest.raises(StateContractError, match="winner set"): + validate_collection_plan(run, plan, (shard,), ()) + + wrong_reference = ArtifactReference( + path=plan.planned_shards[0].winner_manifest.path, + sha256="0" * 64, + ) + wrong_planned_shard = _validated_copy(plan.planned_shards[0], winner_manifest=wrong_reference) + wrong_plan = _validated_copy(plan, planned_shards=(wrong_planned_shard,)) + with pytest.raises(StateContractError, match="digest mismatch"): + validate_collection_plan(run, wrong_plan, (shard,), (winner,)) + + +def test_collection_rejects_an_invented_shard_even_when_its_digest_matches() -> None: + run = _load(RunManifest, "run_manifest.json") + shard = _load(ShardManifest, "shard_manifest.json") + plan = _load(CollectionPlan, "collection_plan.json") + winner = _validated_copy(_load(ShardWinner, "shard_winner.json"), shard_id="shard-99999") + winner_reference = ArtifactReference( + path=plan.planned_shards[0].winner_manifest.path, + sha256=winner.compute_sha256(), + ) + planned_shard = _validated_copy( + plan.planned_shards[0], + shard_id="shard-99999", + winner_manifest=winner_reference, + ) + altered_plan = _validated_copy(plan, planned_shards=(planned_shard,)) + + with pytest.raises(StateContractError, match="planned shards"): + validate_collection_plan(run, altered_plan, (shard,), (winner,)) + + +def test_terminal_attempt_evidence_overrides_stale_readiness() -> None: + attempt = _load(AttemptManifest, "failed_attempt.json") + readiness = _load(AttemptReadiness, "stale_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 13, 5, tzinfo=timezone.utc), + state=SchedulerState.RUNNING, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 13, 5, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +def test_terminal_scheduler_failure_overrides_successful_attempt() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json") + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=SchedulerState.NODE_FAILED, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +@pytest.mark.parametrize("scheduler_state", (SchedulerState.PENDING, SchedulerState.RUNNING)) +def test_failed_readiness_overrides_nonterminal_scheduler_state(scheduler_state: SchedulerState) -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + failed_deployment = _validated_copy( + readiness.deployments[0], + state=ReadinessState.FAILED, + ready_backends=0, + endpoint_publication=EndpointPublicationState.FAILED, + ) + failed_readiness = _validated_copy( + readiness, + state=ReadinessState.FAILED, + deployments=(failed_deployment,), + ) + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=scheduler_state, + ) + + assert ( + reconcile_attempt_observation( + attempt, + failed_readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.FAILED + ) + + +def test_reconciliation_rejects_stale_or_future_evidence() -> None: + attempt = _load(AttemptManifest, "successful_attempt.json") + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + + stale_scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 0, 1, tzinfo=timezone.utc), + state=SchedulerState.NODE_FAILED, + ) + with pytest.raises(StateContractError, match="cannot precede attempt creation"): + reconcile_attempt_observation( + attempt, + readiness, + stale_scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + + current_scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + future_readiness = _validated_copy( + readiness, + updated_at=datetime(2026, 8, 18, 12, 7, tzinfo=timezone.utc), + ) + with pytest.raises(StateContractError, match="cannot precede readiness update"): + reconcile_attempt_observation( + attempt, + future_readiness, + current_scheduler, + current_time=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + + +def test_reconciliation_covers_nonterminal_and_fallback_states() -> None: + successful_attempt = _load(AttemptManifest, "successful_attempt.json") + attempt = _validated_copy( + successful_attempt, + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + ready = _load(AttemptReadiness, "single_node_readiness.json") + pending_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.PENDING, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + last_probe=None, + ) + pending = _validated_copy(ready, state=ReadinessState.PENDING, deployments=(pending_deployment,)) + stopped_deployment = _validated_copy( + ready.deployments[0], + state=ReadinessState.STOPPED, + ready_backends=0, + ) + stopped = _validated_copy(ready, state=ReadinessState.STOPPED, deployments=(stopped_deployment,)) + assert attempt.scheduler is not None + observed_at = datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc) + current_time = observed_at + + cases = ( + (attempt, ready, SchedulerState.PENDING, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.PREEMPTED, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.REQUEUED, EffectiveAttemptState.PENDING), + (attempt, ready, SchedulerState.RUNNING, EffectiveAttemptState.RUNNING), + (attempt, ready, SchedulerState.COMPLETED, EffectiveAttemptState.FAILED), + (attempt, pending, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), + (attempt, stopped, SchedulerState.UNKNOWN, EffectiveAttemptState.UNKNOWN), + (successful_attempt, ready, SchedulerState.UNKNOWN, EffectiveAttemptState.SUCCEEDED), + ) + for case_attempt, case_readiness, scheduler_state, expected in cases: + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=observed_at, + state=scheduler_state, + ) + assert ( + reconcile_attempt_observation( + case_attempt, + case_readiness, + scheduler, + current_time=current_time, + ) + is expected + ) + + +def test_accounting_lag_is_nonterminal_until_its_deadline() -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + scheduler = _load(SchedulerObservation, "accounting_lag.json") + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 7, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.ACCOUNTING_LAG + ) + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.UNKNOWN + ) + + +def test_accounting_lag_deadline_is_fixed_across_observations() -> None: + previous = _load(SchedulerObservation, "accounting_lag.json") + continued = _validated_copy( + previous, + observed_at=datetime(2026, 8, 18, 12, 6, tzinfo=timezone.utc), + ) + refreshed = _validated_copy( + continued, + reconciliation_deadline=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + expired = SchedulerObservation( + schema_version=1, + scheduler=previous.scheduler, + observed_at=datetime(2026, 8, 18, 12, 10, 3, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + + assert validate_scheduler_observation_transition(previous, continued) is continued + assert validate_scheduler_observation_transition(previous, expired) is expired + with pytest.raises(StateContractError, match="deadline cannot change"): + validate_scheduler_observation_transition(previous, refreshed) + + early_unknown = _validated_copy(expired, observed_at=datetime(2026, 8, 18, 12, 9, tzinfo=timezone.utc)) + with pytest.raises(StateContractError, match="before its reconciliation deadline expires"): + validate_scheduler_observation_transition(previous, early_unknown) + + +def test_unknown_scheduler_state_overrides_stale_ready_readiness() -> None: + attempt = _validated_copy( + _load(AttemptManifest, "successful_attempt.json"), + state=AttemptLifecycleState.RUNNING, + terminal_classification=None, + candidate_output=None, + ) + readiness = _load(AttemptReadiness, "single_node_readiness.json") + assert attempt.scheduler is not None + scheduler = SchedulerObservation( + schema_version=1, + scheduler=attempt.scheduler, + observed_at=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + state=SchedulerState.UNKNOWN, + ) + + assert ( + reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=datetime(2026, 8, 18, 12, 11, tzinfo=timezone.utc), + ) + is EffectiveAttemptState.UNKNOWN + ) diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py new file mode 100644 index 000000000..9b491f6fb --- /dev/null +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -0,0 +1,998 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner +from pydantic import BaseModel + +import data_designer.slurm.cli as cli_module +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, + ImageBuildRequest, + ImageInspectionRecord, +) +from data_designer.slurm.images.records import RegisteredImage +from data_designer.slurm.services import ( + SlurmCollectionExecution, + SlurmRetryExecution, + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) +from data_designer.slurm.state import CollectionState + + +def _timestamped_lines(output: str) -> list[str]: + lines: list[str] = [] + for line in output.splitlines(): + match = re.fullmatch(r"\[\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z\] (.*)", line) + assert match is not None, line + lines.append(match.group(1)) + return lines + + +class _RunService: + def __init__(self) -> None: + self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool]] = [] + self.retry_calls: list[tuple[str, tuple[str, ...] | None, str, bool]] = [] + self.collection_calls: list[tuple[str | Path, Path | None, int | None]] = [] + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + ) -> SlurmRunExecution: + self.calls.append((config, source_root, dry_run)) + return SlurmRunExecution( + run_id="run-0001", + state="dry_run" if dry_run else "submitted", + plan_sha256="1" * 64, + shard_count=1, + batch_script="#!/bin/bash\n" if dry_run else None, + job_id=None if dry_run else 43, + ) + + def retry( + self, + run_or_job_id: str, + *, + shard_ids: tuple[str, ...] | None, + resume: str, + dry_run: bool, + ) -> SlurmRetryExecution: + self.retry_calls.append((run_or_job_id, shard_ids, resume, dry_run)) + selected = ("shard-00000",) if shard_ids is None else shard_ids + return SlurmRetryExecution( + run_id="run-0001", + state="dry_run" if dry_run else "submitted", + shard_ids=selected, + attempt_ids=tuple("attempt-0002" for _ in selected), + effective_resume_mode="always", + job_id=None if dry_run else 43, + batch_script="#!/bin/bash\n" if dry_run else None, + ) + + def collect( + self, + input_path: str | Path, + *, + destination: Path | None, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + self.collection_calls.append((input_path, destination, num_partitions)) + return SlurmCollectionExecution( + run_id="run-0001", + collection_id="collection-0001", + state=CollectionState.SUBMITTED, + job_id=43, + output_path=(destination or Path("/workspace/runs/run-0001/merge-output")).resolve().as_posix(), + num_partitions=1 if num_partitions is None else num_partitions, + ) + + +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 + + +@pytest.fixture(autouse=True) +def json_output_for_existing_contract_tests(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep existing JSON contract assertions explicit while the CLI default changes.""" + monkeypatch.setenv("DATA_DESIGNER_SLURM_OUTPUT", "json") + + +def test_default_output_is_readable_and_json_is_explicit( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + run_file = tmp_path / "run.json" + run_file.write_text(authored_run_single.serialize_json()) + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: _RunService()) + + human = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file), "--dry-run"]) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "execute", str(run_file), "--dry-run"]) + + assert human.exit_code == machine.exit_code == 0 + assert "Run id: run-0001" in human.stdout + assert "State: dry_run" in human.stdout + assert "Batch script:\n #!/bin/bash" in human.stdout + assert json.loads(machine.stdout)["run_id"] == "run-0001" + + +def test_submitted_run_hides_absent_batch_script_in_human_output( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + run_file = tmp_path / "run.json" + run_file.write_text(authored_run_single.serialize_json()) + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: _RunService()) + + human = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file)]) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "execute", str(run_file)]) + + assert human.exit_code == machine.exit_code == 0 + assert "Job id: 43" in human.stdout + assert "Batch script:" not in human.stdout + assert json.loads(machine.stdout)["batch_script"] is None + + +def test_default_error_is_readable_and_json_error_is_explicit(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + + def fail(**_: object) -> None: + raise SlurmServiceError(SlurmServiceErrorCode.UNAVAILABLE, SlurmServiceOperation.STATUS_RUN, "try later") + + monkeypatch.setattr(cli_module, "create_slurm_run_service", fail) + + human = CliRunner().invoke(cli_module.create_cli(), ["status", "run-0001"]) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "status", "run-0001"]) + + assert human.exit_code == machine.exit_code == 5 + assert human.stderr == "Error: try later\nOperation: status run\n" + assert json.loads(machine.stderr)["error"]["code"] == "unavailable" + + +def test_default_empty_image_list_is_readable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + + class ImageService: + def list(self) -> tuple[()]: + return () + + monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_: ImageService()) + result = CliRunner().invoke(cli_module.create_cli(), ["image", "ls"]) + + assert result.exit_code == 0 + assert result.stdout == "No registered images.\n" + + +def test_image_add_list_and_info_show_image_details( + monkeypatch: pytest.MonkeyPatch, + serving_image_inspection: ImageInspectionRecord, +) -> None: + image = RegisteredImage( + schema_version=1, + name="vllm", + path="/workspace/images/vllm.sqsh", + sqsh_sha256=serving_image_inspection.sqsh_sha256, + source_oci_digest="a" * 64, + inspection=serving_image_inspection, + ) + + class ImageService: + def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: + assert request.name == image.name + assert not replace + return image + + def list(self) -> tuple[RegisteredImage, ...]: + return (image,) + + def get(self, name: str) -> RegisteredImage: + assert name == image.name + return image + + monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_: ImageService()) + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + source = f"registry.example/vllm@sha256:{'a' * 64}" + + added = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--name", image.name]) + listed = CliRunner().invoke(cli_module.create_cli(), ["image", "ls"]) + info = CliRunner().invoke(cli_module.create_cli(), ["image", "info", image.name]) + machine_added = CliRunner().invoke( + cli_module.create_cli(), ["--output", "json", "image", "add", source, "--name", image.name] + ) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "image", "info", image.name]) + machine_list = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "image", "ls"]) + + expected = ( + "Image alias: vllm\n" + "SQSH path: /workspace/images/vllm.sqsh\n" + f"SQSH sha256: {serving_image_inspection.sqsh_sha256}\n" + f"Source OCI digest: {'a' * 64}\n" + "Image details:\n" + " Kind: serving\n" + " Server type: vllm\n" + " Runtime version: 0.21.0\n" + " Executable path: /usr/local/bin/vllm\n" + ) + assert added.exit_code == listed.exit_code == info.exit_code == 0 + assert machine_added.exit_code == machine.exit_code == machine_list.exit_code == 0 + assert added.stdout == expected.replace("Image alias:", "Registered image alias:", 1) + assert listed.stdout == "vllm (serving, vLLM 0.21.0)\n SQSH: /workspace/images/vllm.sqsh\n" + assert info.stdout == expected + assert json.loads(machine_added.stdout) == image.model_dump(mode="json") + assert json.loads(machine.stdout) == image.model_dump(mode="json") + assert json.loads(machine_list.stdout) == [image.model_dump(mode="json")] + + +def test_existing_client_image_summary_omits_absent_oci_digest( + monkeypatch: pytest.MonkeyPatch, + client_image_inspection: ImageInspectionRecord, +) -> None: + image = RegisteredImage( + schema_version=1, + name="client", + path="/workspace/images/client.sqsh", + sqsh_sha256=client_image_inspection.sqsh_sha256, + inspection=client_image_inspection, + ) + + class ImageService: + def get(self, name: str) -> RegisteredImage: + assert name == image.name + return image + + def list(self) -> tuple[RegisteredImage, ...]: + return (image,) + + monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_: ImageService()) + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "info", image.name]) + listed = CliRunner().invoke(cli_module.create_cli(), ["image", "ls"]) + + assert result.exit_code == listed.exit_code == 0 + assert "Source OCI digest:" not in result.stdout + assert result.stdout.count("Image details:") == 1 + assert " Kind: client\n" in result.stdout + assert " Python version:" in result.stdout + assert listed.stdout == "client (client, Python 3.12.12)\n SQSH: /workspace/images/client.sqsh\n" + + +def test_image_list_sanitizes_runtime_without_changing_json( + monkeypatch: pytest.MonkeyPatch, + serving_image_inspection: ImageInspectionRecord, + client_image_inspection: ImageInspectionRecord, +) -> None: + payload = serving_image_inspection.model_dump(mode="json") + payload["inspection"]["runtime_version"] = "0.21.0\nForged alias: admin\x1b[31m token=secretvalue" + altered_inspection = ImageInspectionRecord.model_validate(payload) + serving = RegisteredImage( + schema_version=1, + name="vllm", + path="/workspace/images/vllm.sqsh", + sqsh_sha256=altered_inspection.sqsh_sha256, + inspection=altered_inspection, + ) + client = RegisteredImage( + schema_version=1, + name="client", + path="/workspace/images/client.sqsh", + sqsh_sha256=client_image_inspection.sqsh_sha256, + inspection=client_image_inspection, + ) + + class ImageService: + def list(self) -> tuple[RegisteredImage, ...]: + return (serving, client) + + monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_: ImageService()) + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + + human = CliRunner().invoke(cli_module.create_cli(), ["image", "ls"]) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "image", "ls"]) + + assert human.exit_code == machine.exit_code == 0 + lines = human.stdout.splitlines() + assert len(lines) == 5 + assert lines[0].startswith("vllm (serving, vLLM 0.21.0 Forged alias: admin") + assert "token=" in lines[0] + assert lines[1] == " SQSH: /workspace/images/vllm.sqsh" + assert lines[2] == "" + assert lines[3:] == ["client (client, Python 3.12.12)", " SQSH: /workspace/images/client.sqsh"] + assert "\x1b" not in human.stdout + assert "secretvalue" not in human.stdout + assert json.loads(machine.stdout) == [serving.model_dump(mode="json"), client.model_dump(mode="json")] + + +def test_image_remove_confirms_alias_removal_and_preserves_json(monkeypatch: pytest.MonkeyPatch) -> None: + class RemovedImage(BaseModel): + name: str + path: str + + removed = RemovedImage(name="vllm-gpt-oss", path="/workspace/images/vllm-gpt-oss.sqsh") + calls: list[bool] = [] + + class ImageService: + def remove(self, name: str, *, delete_artifact: bool) -> RemovedImage: + assert name == removed.name + calls.append(delete_artifact) + return removed + + monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_: ImageService()) + monkeypatch.delenv("DATA_DESIGNER_SLURM_OUTPUT") + + human = CliRunner().invoke(cli_module.create_cli(), ["image", "rm", removed.name]) + deleted = CliRunner().invoke(cli_module.create_cli(), ["image", "rm", removed.name, "--delete-artifact"]) + machine = CliRunner().invoke(cli_module.create_cli(), ["--output", "json", "image", "rm", removed.name]) + + assert human.exit_code == deleted.exit_code == machine.exit_code == 0 + assert human.stdout == ( + "Removed image alias: vllm-gpt-oss\nSQSH artifact retained: /workspace/images/vllm-gpt-oss.sqsh\n" + ) + assert deleted.stdout == ( + "Removed image alias: vllm-gpt-oss\nSQSH artifact deleted: /workspace/images/vllm-gpt-oss.sqsh\n" + ) + assert json.loads(machine.stdout) == removed.model_dump() + assert calls == [False, True, False] + + +def test_json_output_works_without_click_context(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cli_module.click, "get_current_context", lambda silent=True: None) + token = cli_module._ACTIVE_OUTPUT_FORMAT.set(cli_module._OutputFormat.JSON) + try: + assert cli_module._output_format() is cli_module._OutputFormat.JSON + finally: + cli_module._ACTIVE_OUTPUT_FORMAT.reset(token) + assert cli_module._output_format() is cli_module._OutputFormat.HUMAN + + +def test_human_output_redacts_credentials_and_terminal_controls() -> None: + rendered = "\n".join( + cli_module._human_lines( + {"detail": "token=secretvalue\x1b[31m", "script": "first line\napi_key=anothersecret\nlast line"} + ) + ) + + assert "secretvalue" not in rendered + assert "anothersecret" not in rendered + assert "\x1b" not in rendered + assert "token=" in rendered + assert "api_key=" in rendered + assert " first line\n api_key=\n last line" in rendered + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + (["status", "--help"], "show persisted run status"), + (["image", "add", "--help"], "Import or inspect an image and register its alias"), + ], +) +def test_help_uses_public_terminology(arguments: list[str], expected: str) -> None: + result = CliRunner().invoke(cli_module.create_cli(), arguments) + + assert result.exit_code == 0 + assert expected in result.output + + +def test_execute_emits_deterministic_json_and_forwards_actions( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch, +) -> None: + run_file = tmp_path / "run.json" + run_file.write_text(authored_run_single.serialize_json()) + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file), "--dry-run"]) + + assert result.exit_code == 0 + assert json.loads(result.stdout) == { + "batch_script": "#!/bin/bash\n", + "job_id": None, + "plan_sha256": "1" * 64, + "run_id": "run-0001", + "shard_count": 1, + "state": "dry_run", + } + assert service.calls == [(authored_run_single, tmp_path, True)] + assert result.stderr == "" + + +def test_execute_reports_submission_in_interactive_terminal( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + run_file = tmp_path / "run.json" + run_file.write_text(authored_run_single.serialize_json()) + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: _RunService()) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + + result = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file)]) + + assert result.exit_code == 0 + assert json.loads(result.stdout)["state"] == "submitted" + assert _timestamped_lines(result.stderr) == ["Preparing and submitting run..."] + + +def test_progress_stream_failure_does_not_stop_submissions( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + run_file = tmp_path / "run.json" + run_file.write_text(authored_run_single.serialize_json()) + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + original_echo = cli_module.typer.echo + + def echo_with_broken_progress(message: str, *, err: bool = False) -> None: + if "Preparing and submitting" in message: + raise BrokenPipeError("progress stream closed") + original_echo(message, err=err) + + monkeypatch.setattr(cli_module.typer, "echo", echo_with_broken_progress) + + execute = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file)]) + retry = CliRunner().invoke(cli_module.create_cli(), ["retry", "42", "--force"]) + merge = CliRunner().invoke( + cli_module.create_cli(), + ["merge", "--input-path", str(tmp_path), "--output-path", str(tmp_path / "collected")], + ) + + assert [result.exit_code for result in (execute, retry, merge)] == [0, 0, 0] + assert [json.loads(result.stdout)["state"] for result in (execute, retry, merge)] == [ + "submitted", + "submitted", + "submitted", + ] + assert len(service.calls) == 1 + assert len(service.retry_calls) == 1 + assert len(service.collection_calls) == 1 + + +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) + + def create_benchmark_service(**kwargs: object) -> _BenchmarkService: + if "progress" in kwargs: + progress = kwargs["progress"] + assert callable(progress) + progress("Benchmark case 1/1: submitted.") + return service + + monkeypatch.setattr(benchmark_cli_module, "create_slurm_benchmark_service", create_benchmark_service) + monkeypatch.setattr( + benchmark_cli_module, + "resolve_profile", + lambda **_: SimpleNamespace(profile=SimpleNamespace(workspace_root="/workspace")), + ) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + + 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)] + assert _timestamped_lines(run_result.stderr) == ["Preparing benchmark launch...", "Benchmark case 1/1: submitted."] + + +def test_benchmark_analyze_rejects_directory_outside_selected_workspace( + benchmark_manifest: BenchmarkManifest, benchmark_report: BenchmarkReport, monkeypatch: pytest.MonkeyPatch +) -> None: + service = _BenchmarkService(benchmark_manifest, benchmark_report) + monkeypatch.setattr(benchmark_cli_module, "create_slurm_benchmark_service", lambda **_: service) + monkeypatch.setattr( + benchmark_cli_module, + "resolve_profile", + lambda **_: SimpleNamespace(profile=SimpleNamespace(workspace_root="/workspace")), + ) + + result = CliRunner().invoke( + cli_module.create_cli(), + ["benchmark", "analyze", "/another-workspace/benchmarks/benchmark-001"], + ) + + assert result.exit_code == 2 + assert "outside the selected profile" in result.output + assert service.analysis_calls == [] + + +def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke( + cli_module.create_cli(), + ["retry", "42", "--task-id", "1", "--task-id", "3", "--resume", "always", "--dry-run"], + ) + + assert result.exit_code == 0 + assert json.loads(result.stdout) == { + "attempt_ids": ["attempt-0002", "attempt-0002"], + "batch_script": "#!/bin/bash\n", + "effective_resume_mode": "always", + "job_id": None, + "run_id": "run-0001", + "shard_ids": ["shard-00001", "shard-00003"], + "state": "dry_run", + } + assert service.retry_calls == [("42", ("shard-00001", "shard-00003"), "always", True)] + + +def test_retry_auto_selects_tasks_and_confirms_submission(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="y\n") + + assert result.exit_code == 0 + assert json.loads(result.stdout.splitlines()[-1]) == { + "attempt_ids": ["attempt-0002"], + "batch_script": None, + "effective_resume_mode": "always", + "job_id": 43, + "run_id": "run-0001", + "shard_ids": ["shard-00000"], + "state": "submitted", + } + assert "Submit this retry?" in result.stderr + assert "Retry shard-00000 with resume=always" in result.stderr + assert service.retry_calls == [ + ("42", None, "if_possible", True), + ("42", ("shard-00000",), "always", False), + ] + + +def test_retry_force_skips_confirmation(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42", "--force"]) + + assert result.exit_code == 0 + assert service.retry_calls == [("42", None, "if_possible", False)] + assert _timestamped_lines(result.stderr) == ["Preparing and submitting retry..."] + + +def test_retry_decline_emits_stable_result(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="n\n") + + assert result.exit_code == 0 + assert json.loads(result.stdout.splitlines()[-1]) == {"operation": "retry_run", "state": "declined"} + assert service.retry_calls == [("42", None, "if_possible", True)] + + +def test_retry_noninteractive_confirmation_is_invalid_request(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42"], input="") + + assert result.exit_code == 2 + assert json.loads(result.stderr.splitlines()[-1]) == { + "error": { + "code": "invalid_request", + "message": "interactive confirmation is unavailable; pass --force or --dry-run", + "operation": "retry_run", + } + } + assert service.retry_calls == [("42", None, "if_possible", True)] + + +def test_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + input_path = tmp_path / "runs/run-0001" + output_path = tmp_path / "collected" + + result = CliRunner().invoke( + cli_module.create_cli(), + [ + "merge", + "--input-path", + str(input_path), + "--output-path", + str(output_path), + "--num-partitions", + "2", + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.stdout) == { + "collection_id": "collection-0001", + "job_id": 43, + "num_partitions": 2, + "output_path": output_path.as_posix(), + "run_id": "run-0001", + "state": "submitted", + } + assert service.collection_calls == [(input_path, output_path, 2)] + assert _timestamped_lines(result.stderr) == ["Preparing and submitting collection job..."] + + +def test_merge_accepts_run_id_with_default_destination(monkeypatch: pytest.MonkeyPatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + result = CliRunner().invoke(cli_module.create_cli(), ["merge", "run-0001"]) + + assert result.exit_code == 0 + assert json.loads(result.stdout)["output_path"] == "/workspace/runs/run-0001/merge-output" + assert service.collection_calls == [("run-0001", None, None)] + + +def test_merge_accepts_run_id_with_explicit_destination(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + destination = tmp_path / "collected" + + result = CliRunner().invoke(cli_module.create_cli(), ["merge", "run-0001", "--output-path", str(destination)]) + + assert result.exit_code == 0 + assert service.collection_calls == [("run-0001", destination, None)] + + +@pytest.mark.parametrize("arguments", [[], ["run-0001", "--input-path", "/workspace/runs/run-0001"]]) +def test_merge_requires_one_run_reference(arguments: list[str], monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: _RunService()) + + result = CliRunner().invoke(cli_module.create_cli(), ["merge", *arguments]) + + assert result.exit_code == 2 + assert json.loads(result.stderr.splitlines()[-1])["error"]["message"] == "specify a run ID or --input-path" + + +@pytest.mark.parametrize( + ("code", "exit_code"), + [ + (SlurmServiceErrorCode.INVALID_REQUEST, 2), + (SlurmServiceErrorCode.NOT_FOUND, 3), + (SlurmServiceErrorCode.CONFLICT, 4), + (SlurmServiceErrorCode.UNAVAILABLE, 5), + ], +) +def test_cli_emits_stable_service_error(monkeypatch, code: SlurmServiceErrorCode, exit_code: int) -> None: + error = SlurmServiceError( + code, + SlurmServiceOperation.STATUS_RUN, + "stable failure", + ) + + def fail(**_): + raise error + + monkeypatch.setattr(cli_module, "create_slurm_run_service", fail) + + result = CliRunner().invoke(cli_module.create_cli(), ["status", "run-0001"]) + + assert result.exit_code == exit_code + assert json.loads(result.stderr) == { + "error": { + "code": code.value, + "message": "stable failure", + "operation": "status_run", + } + } + + +def test_execute_preserves_sanitized_config_diagnostic(tmp_path: Path) -> None: + run_file = tmp_path / "run.json" + run_file.write_text('{"schema_version":1}') + + result = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file)]) + + assert result.exit_code == 2 + error = json.loads(result.stderr)["error"] + assert error["code"] == "invalid_request" + assert "failed validation" in error["message"] + + +def test_execute_bounds_large_config_diagnostic(tmp_path: Path) -> None: + run_file = tmp_path / f"{'x' * 240}.json" + run_file.write_text(json.dumps({f"unknown_field_{index}": index for index in range(100)})) + + result = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file)]) + + assert result.exit_code == 2 + message = json.loads(result.stderr)["error"]["message"] + assert len(message) == 512 + assert message.endswith("...") + + +@pytest.mark.parametrize( + "source", + ["nvcr.io/nvidia/pytorch:24.01-py3", "docker://ubuntu:22.04"], +) +def test_image_add_rejects_mutable_oci_source(source: str) -> None: + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--kind", "client"]) + + assert result.exit_code == 2 + assert json.loads(result.stderr) == { + "error": { + "code": "invalid_request", + "message": "OCI image source must be digest-qualified as name@sha256:", + "operation": "add_image", + } + } + + +@pytest.mark.parametrize("interactive", (False, True)) +def test_image_add_resolves_versioned_vllm_before_import( + monkeypatch: pytest.MonkeyPatch, + interactive: bool, +) -> None: + source = "vllm/vllm-openai:v0.22.0" + pinned = f"{source}@sha256:{'a' * 64}" + requests: list[ImageBuildRequest] = [] + + class ImageService: + def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest: + assert not replace + requests.append(request) + return request + + monkeypatch.setattr(cli_module, "resolve_versioned_vllm_source", lambda _source: pinned) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: interactive) + + def create_service(**kwargs: object) -> ImageService: + progress = kwargs["progress"] + assert callable(progress) + progress("Image add job 42: pending (elapsed 00:00:00).") + return ImageService() + + monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source]) + + assert result.exit_code == 0, result.output + assert len(requests) == 1 + assert requests[0].kind == "serving" + assert requests[0].source == pinned + assert requests[0].name == "vllm-openai-v0.22.0" + assert json.loads(result.stdout)["name"] == "vllm-openai-v0.22.0" + assert _timestamped_lines(result.stderr) == ( + [ + "Resolving vLLM image tag...", + "Preparing image registration...", + "Image add job 42: pending (elapsed 00:00:00).", + ] + if interactive + else [] + ) + + +def test_image_add_follow_logs_is_explicit_and_keeps_json_on_stdout(monkeypatch: pytest.MonkeyPatch) -> None: + source = f"vllm/vllm-openai@sha256:{'a' * 64}" + + class ImageService: + def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest: + assert not replace + return request + + def create_service(**kwargs: object) -> ImageService: + logs = kwargs["logs"] + assert callable(logs) + logs("[image stderr] inspection started") + return ImageService() + + monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: False) + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--follow-logs"]) + + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["source"] == source + assert _timestamped_lines(result.stderr) == ["[image stderr] inspection started"] + + +@pytest.mark.parametrize("interactive", [False, True]) +def test_image_add_follows_logs_by_default_in_terminals( + monkeypatch: pytest.MonkeyPatch, + interactive: bool, +) -> None: + source = f"vllm/vllm-openai@sha256:{'a' * 64}" + + class ImageService: + def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest: + return request + + def create_service(**kwargs: object) -> ImageService: + logs = kwargs["logs"] + assert callable(logs) is interactive + if callable(logs): + logs("[image stderr] inspection started") + return ImageService() + + monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: interactive) + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source]) + + assert result.exit_code == 0, result.output + assert json.loads(result.stdout)["source"] == source + assert ("[image stderr] inspection started" in result.stderr) is interactive + + +def test_image_add_can_disable_default_log_following(monkeypatch: pytest.MonkeyPatch) -> None: + source = f"vllm/vllm-openai@sha256:{'a' * 64}" + + class ImageService: + def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest: + return request + + def create_service(**kwargs: object) -> ImageService: + assert kwargs["logs"] is None + return ImageService() + + monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: True) + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--no-follow-logs"]) + + assert result.exit_code == 0, result.output + assert "[image stderr]" not in result.stderr + + +def test_image_add_rejects_versioned_vllm_as_client() -> None: + result = CliRunner().invoke( + cli_module.create_cli(), ["image", "add", "vllm/vllm-openai:v0.22.0", "--kind", "client"] + ) + + assert result.exit_code == 2 + assert json.loads(result.stderr)["error"]["message"] == "versioned vLLM images require --kind serving" + + +def test_profile_init_creates_starter_and_emits_validation_command(tmp_path: Path) -> None: + profile_file = tmp_path / "config" / "profile.yml" + workspace = tmp_path / "workspace" + + result = CliRunner().invoke( + cli_module.create_cli(), + [ + "profile", + "init", + "--workspace-root", + str(workspace), + "--image-build-partition", + "cpu", + "--profile-file", + str(profile_file), + "--host-pattern", + "login.example.test", + ], + ) + + assert result.exit_code == 0, result.output + assert json.loads(result.stdout) == { + "profile_file": profile_file.as_posix(), + "validation_command": f"data-designer slurm profile validate --profile-file {profile_file.as_posix()}", + } + assert profile_file.is_file() + assert profile_file.parent.is_dir() + assert not workspace.exists() + + +def test_profile_validate_emits_selected_effective_paths(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + profile_file = tmp_path / "profile.json" + profile_file.write_text( + json.dumps( + { + "schema_version": 1, + "default_cluster": "local", + "clusters": { + "local": { + "schema_version": 1, + "gpus_per_node": 4, + "workspace_root": workspace.as_posix(), + "image_build": { + "partition": "cpu", + "cpus_per_task": 2, + "memory": "8G", + "time_limit": "04:00:00", + }, + } + }, + } + ) + ) + + result = CliRunner().invoke( + cli_module.create_cli(), + ["profile", "validate", "--profile-file", str(profile_file), "--cluster", "local"], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["profile_file"] == profile_file.as_posix() + assert payload["selected_cluster"] == "local" + assert payload["selection_source"] == "explicit" + assert payload["workspace_root"] == workspace.as_posix() + assert payload["image_root"] == (workspace / "images").as_posix() + assert payload["registry_file"] == (workspace / "images" / "registry.yaml").as_posix() + assert payload["gpus_per_node"] == 4 + assert tuple(workspace.iterdir()) == () + + +def test_image_add_rejects_credential_bearing_oci_source() -> None: + source = f"https://user:secret@registry.example/image@sha256:{'a' * 64}" + + result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source, "--kind", "client"]) + + assert result.exit_code == 2 + error = json.loads(result.stderr)["error"] + assert error == { + "code": "invalid_request", + "message": "OCI image source must be a credential-free registry reference without a scheme", + "operation": "add_image", + } + assert "secret" not in result.stderr + + +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", "benchmark") + ) diff --git a/packages/data-designer-slurm/tests/test_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py new file mode 100644 index 000000000..8911689b5 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib + +import pytest +from pydantic import ValidationError + +from data_designer.slurm import contracts as contracts_module +from data_designer.slurm import types as types_module +from data_designer.slurm.contracts import ( + ArtifactReference as ContractArtifactReference, +) +from data_designer.slurm.contracts import ( + ContractRecord, + ContractValue, + ResumeWorkspace, + canonical_json, + compute_canonical_json_sha256, + pretty_json, +) +from data_designer.slurm.contracts import ( + RecordRange as ContractRecordRange, +) +from data_designer.slurm.state import ( + ArtifactReference as StateArtifactReference, +) +from data_designer.slurm.state import ( + ContractRecord as StateContractRecord, +) +from data_designer.slurm.state import ( + ContractValue as StateContractValue, +) +from data_designer.slurm.state import ( + RecordRange as StateRecordRange, +) +from data_designer.slurm.state import ( + ResumeWorkspace as StateResumeWorkspace, +) +from data_designer.slurm.state import ( + StateRecord, + StateValue, +) + + +@pytest.mark.parametrize("type_name", types_module.__all__) +def test_temporary_contracts_compatibility_exports_match_shared_scalar_types(type_name: str) -> None: + assert getattr(contracts_module, type_name) is getattr(types_module, type_name) + + +def test_state_exports_exact_shared_contract_types() -> None: + assert StateArtifactReference is ContractArtifactReference + assert StateRecordRange is ContractRecordRange + assert StateResumeWorkspace is ResumeWorkspace + assert StateContractValue is ContractValue + assert StateContractRecord is ContractRecord + assert StateValue is ContractValue + assert StateRecord is ContractRecord + + +def test_shared_json_helpers_are_deterministic() -> None: + value = {"unicode": "模型", "number": 1} + + assert canonical_json(value) == b'{"number":1,"unicode":"\xe6\xa8\xa1\xe5\x9e\x8b"}' + assert pretty_json(value) == '{\n "number": 1,\n "unicode": "模型"\n}\n' + assert compute_canonical_json_sha256(value) == hashlib.sha256(canonical_json(value)).hexdigest() + + +def test_resume_workspace_requires_a_safe_absolute_path() -> None: + workspace = ResumeWorkspace(path="/workspace/runs/run-0001/shards/shard-00000/dataset") + + assert workspace.path.endswith("/dataset") + with pytest.raises(ValidationError, match="exactly one leading slash"): + ResumeWorkspace(path="//workspace/dataset") diff --git a/packages/data-designer-slurm/tests/test_package.py b/packages/data-designer-slurm/tests/test_package.py new file mode 100644 index 000000000..24a4b4955 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_package.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.metadata +from pathlib import Path + +import data_designer +import data_designer.slurm + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def test_slurm_uses_shared_namespace() -> None: + assert data_designer.__file__ is None + assert data_designer.slurm.__name__ == "data_designer.slurm" + + +def test_slurm_license_matches_repository() -> None: + package_license = REPO_ROOT / "packages" / "data-designer-slurm" / "LICENSE" + + assert package_license.read_bytes() == (REPO_ROOT / "LICENSE").read_bytes() + + +def test_slurm_is_published_before_base_extra() -> None: + publish_script = (REPO_ROOT / "scripts" / "publish.sh").read_text() + + assert publish_script.index('"packages/data-designer-slurm"') < publish_script.index('"packages/data-designer"') + + +def test_slurm_registers_lazy_cli_extension() -> None: + entry_points = importlib.metadata.entry_points(group="data_designer.cli") + slurm_entry_point = next( + entry_point + for entry_point in entry_points + if entry_point.name == "slurm" and entry_point.dist.name == "data-designer-slurm" + ) + + assert slurm_entry_point.value == "data_designer.slurm.cli:create_cli" diff --git a/packages/data-designer-slurm/tests/test_public_artifacts.py b/packages/data-designer-slurm/tests/test_public_artifacts.py new file mode 100644 index 000000000..73b9f1361 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_public_artifacts.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import subprocess +import sys +import tarfile +import zipfile +from io import BytesIO +from pathlib import Path + +import pytest + +REPOSITORY_ROOT = Path(__file__).parents[3] +AUDIT_SCRIPT = REPOSITORY_ROOT / "scripts" / "audit_slurm_public_artifacts.py" +PACKAGE_LICENSE = REPOSITORY_ROOT / "packages" / "data-designer-slurm" / "LICENSE" +SPDX_HEADER = """# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" + + +def test_repository_slurm_artifacts_pass_public_audit() -> None: + result = _run_audit() + + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize( + ("content", "rule"), + ( + (f"token=github_pat_{'a' * 24}\n", "GitHub token"), + ("workspace=/home/specific-user/run\n", "user-specific absolute path"), + ("endpoint=10.23.45.67\n", "private infrastructure address"), + ("host=service.internal.nvidia.com\n", "internal NVIDIA hostname"), + ), + ids=("credential", "user-path", "private-address", "internal-host"), +) +def test_public_audit_reports_rule_without_echoing_sensitive_content( + tmp_path: Path, + content: str, + rule: str, +) -> None: + artifact = tmp_path / "runtime.log" + artifact.write_text(content) + + result = _run_audit(artifact) + + assert result.returncode == 1 + assert rule in result.stderr + assert content.strip() not in result.stderr + assert str(tmp_path) not in result.stderr + + +def test_public_audit_detects_quoted_authorization_header(tmp_path: Path) -> None: + artifact = tmp_path / "runtime.log" + serialized_header = '{"Author' + 'ization": "Bearer ' + "serialized-secret-value" + '"}\n' + artifact.write_text(serialized_header) + + result = _run_audit(artifact) + + assert result.returncode == 1 + assert "authorization credential" in result.stderr + assert "serialized-secret-value" not in result.stderr + + +def test_public_audit_checks_wheel_members_and_license_text(tmp_path: Path) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr("data_designer/slurm/runtime.py", f"{SPDX_HEADER}\nfrom __future__ import annotations\n") + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/licenses/LICENSE", + PACKAGE_LICENSE.read_bytes(), + ) + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/METADATA", + "Metadata-Version: 2.5\nLicense-Expression: Apache-2.0\n", + ) + + result = _run_audit(wheel) + + assert result.returncode == 0, result.stderr + + +def test_public_audit_rejects_wheel_with_truncated_license(tmp_path: Path) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/licenses/LICENSE", + "Apache License\nVersion 2.0\n", + ) + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/METADATA", + "Metadata-Version: 2.5\nLicense-Expression: Apache-2.0\n", + ) + + result = _run_audit(wheel) + + assert result.returncode == 1 + assert "wheel does not contain exactly one canonical license text" in result.stderr + + +def test_public_audit_rejects_multiple_wheel_distribution_roots(tmp_path: Path) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/licenses/LICENSE", + PACKAGE_LICENSE.read_bytes(), + ) + archive.writestr( + "decoy-1.0.0.dist-info/METADATA", + "Metadata-Version: 2.5\nLicense-Expression: Apache-2.0\n", + ) + + result = _run_audit(wheel) + + assert result.returncode == 1 + assert "wheel must contain exactly one .dist-info root" in result.stderr + + +@pytest.mark.parametrize( + ("extra_member", "expected"), + ( + ( + "data_designer_slurm-1.0.0.dist-info/licenses/LICENSE.txt", + "wheel does not contain exactly one canonical license text", + ), + ( + "data_designer_slurm-1.0.0.dist-info/metadata", + "wheel does not contain exactly one Apache-2.0 metadata record", + ), + ), + ids=("license", "metadata"), +) +def test_public_audit_rejects_duplicate_wheel_distribution_records( + tmp_path: Path, + extra_member: str, + expected: str, +) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + metadata = "Metadata-Version: 2.5\nLicense-Expression: Apache-2.0\n" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr("data_designer_slurm-1.0.0.dist-info/licenses/LICENSE", PACKAGE_LICENSE.read_bytes()) + archive.writestr("data_designer_slurm-1.0.0.dist-info/METADATA", metadata) + archive.writestr(extra_member, metadata if extra_member.endswith("metadata") else PACKAGE_LICENSE.read_bytes()) + + result = _run_audit(wheel) + + assert result.returncode == 1 + assert expected in result.stderr + + +def test_public_audit_rejects_unsafe_or_unlicensed_wheel_members(tmp_path: Path) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr("../escaped.py", "pass\n") + archive.writestr("data_designer/slurm/runtime.py", "from __future__ import annotations\n") + + result = _run_audit(wheel) + + assert result.returncode == 1 + assert "archive member path is unsafe" in result.stderr + assert "missing the NVIDIA Apache-2.0 SPDX header" in result.stderr + assert "wheel must contain exactly one .dist-info root" in result.stderr + + +def test_public_audit_scans_archive_member_names_without_echoing_them(tmp_path: Path) -> None: + wheel = tmp_path / "data_designer_slurm-1.0.0-py3-none-any.whl" + sensitive_member = "data_designer/slurm/10.23.45.67.py" + with zipfile.ZipFile(wheel, mode="w") as archive: + archive.writestr(sensitive_member, SPDX_HEADER) + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/licenses/LICENSE", + PACKAGE_LICENSE.read_bytes(), + ) + archive.writestr( + "data_designer_slurm-1.0.0.dist-info/METADATA", + "Metadata-Version: 2.5\nLicense-Expression: Apache-2.0\n", + ) + + result = _run_audit(wheel) + + assert result.returncode == 1 + assert "private infrastructure address" in result.stderr + assert sensitive_member not in result.stderr + + +def test_public_audit_sanitizes_encrypted_zip_member_failures(tmp_path: Path) -> None: + archive_path = tmp_path / "artifacts.zip" + sensitive_member = "sensitive-member-name.txt" + with zipfile.ZipFile(archive_path, mode="w") as archive: + archive.writestr(sensitive_member, "safe\n") + content = bytearray(archive_path.read_bytes()) + local_header = content.index(b"PK\x03\x04") + central_header = content.index(b"PK\x01\x02") + content[local_header + 6] |= 1 + content[central_header + 8] |= 1 + archive_path.write_bytes(content) + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "archive member cannot be read" in result.stderr + assert sensitive_member not in result.stderr + + +def test_public_audit_scans_unknown_suffix_archive_members(tmp_path: Path) -> None: + archive_path = tmp_path / "artifacts.zip" + with zipfile.ZipFile(archive_path, mode="w") as archive: + archive.writestr("credentials.env", f"token=github_pat_{'a' * 24}\n") + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "GitHub token" in result.stderr + + +def test_public_audit_detects_zip_archive_by_content(tmp_path: Path) -> None: + archive_path = tmp_path / "opaque-artifact" + with zipfile.ZipFile(archive_path, mode="w") as archive: + archive.writestr("credentials.env", f"token=github_pat_{'a' * 24}\n") + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "GitHub token" in result.stderr + + +def test_public_audit_detects_tar_archive_by_content(tmp_path: Path) -> None: + archive_path = tmp_path / "opaque-artifact" + secret = f"nvapi-{'a' * 24}".encode() + with tarfile.open(archive_path, mode="w:gz") as archive: + member = tarfile.TarInfo("credentials.env") + member.size = len(secret) + archive.addfile(member, BytesIO(secret)) + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "NGC API key" in result.stderr + + +def test_public_audit_checks_runtime_tar_content_and_entrypoint_license(tmp_path: Path) -> None: + archive_path = tmp_path / "runtime.tar.gz" + secret = f"nvapi-{'a' * 24}" + with tarfile.open(archive_path, mode="w:gz") as archive: + entrypoint = tarfile.TarInfo("entrypoint.sh") + entrypoint_content = f"#!/bin/sh\necho {secret}\n".encode() + entrypoint.size = len(entrypoint_content) + archive.addfile(entrypoint, BytesIO(entrypoint_content)) + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "NGC API key" in result.stderr + assert "missing the NVIDIA Apache-2.0 SPDX header" in result.stderr + assert secret not in result.stderr + + +@pytest.mark.parametrize( + "member_type", + (tarfile.FIFOTYPE, tarfile.CHRTYPE, tarfile.BLKTYPE), + ids=("fifo", "character-device", "block-device"), +) +def test_public_audit_rejects_tar_special_members(tmp_path: Path, member_type: bytes) -> None: + archive_path = tmp_path / "runtime.tar" + with tarfile.open(archive_path, mode="w") as archive: + member = tarfile.TarInfo("unsupported-member") + member.type = member_type + archive.addfile(member) + + result = _run_audit(archive_path) + + assert result.returncode == 1 + assert "archive member type is unsupported" in result.stderr + + +def test_public_audit_allows_tar_directories(tmp_path: Path) -> None: + archive_path = tmp_path / "runtime.tar" + with tarfile.open(archive_path, mode="w") as archive: + directory = tarfile.TarInfo("runtime") + directory.type = tarfile.DIRTYPE + archive.addfile(directory) + + result = _run_audit(archive_path) + + assert result.returncode == 0, result.stderr + + +def test_public_audit_rejects_explicit_symbolic_link_without_disclosing_its_parent(tmp_path: Path) -> None: + target = tmp_path / "target.log" + target.write_text("safe\n") + link = tmp_path / "linked.log" + link.symlink_to(target) + + result = _run_audit(link) + + assert result.returncode == 1 + assert "symbolic-link artifact requires explicit review" in result.stderr + assert str(tmp_path) not in result.stderr + + +def test_public_audit_rejects_symbolic_link_named_as_generated_cache(tmp_path: Path) -> None: + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + target = tmp_path / "target" + target.mkdir() + (artifacts / "__pycache__").symlink_to(target, target_is_directory=True) + + result = _run_audit(artifacts) + + assert result.returncode == 1 + assert "symbolic-link artifact requires explicit review" in result.stderr + + +@pytest.mark.parametrize("scan_directory", (False, True), ids=("explicit", "directory")) +def test_public_audit_rejects_fifo_without_opening_it(tmp_path: Path, scan_directory: bool) -> None: + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + fifo = artifacts / "named-pipe" + os.mkfifo(fifo) + + result = _run_audit(artifacts if scan_directory else fifo) + + assert result.returncode == 1 + assert "artifact is not a regular file" in result.stderr + + +@pytest.mark.parametrize("filename", ("credentials.env", "records.csv", "METADATA", "LICENSE")) +def test_public_audit_scans_content_with_unknown_or_empty_suffixes(tmp_path: Path, filename: str) -> None: + artifact = tmp_path / filename + artifact.write_text(f"token=github_pat_{'a' * 24}\n") + + result = _run_audit(artifact) + + assert result.returncode == 1 + assert "GitHub token" in result.stderr + + +def test_public_audit_does_not_allow_test_sentinels_outside_exact_sources(tmp_path: Path) -> None: + artifact = tmp_path / "test_credentials.py" + artifact.write_text('secret = "super-secret-token"\n') + + result = _run_audit(artifact) + + assert result.returncode == 1 + assert "plaintext secret assignment" in result.stderr + + +def test_public_audit_does_not_mask_sentinel_prefixes(tmp_path: Path) -> None: + audit_script = tmp_path / "scripts" / "audit_slurm_public_artifacts.py" + audit_script.parent.mkdir() + audit_script.write_bytes(AUDIT_SCRIPT.read_bytes()) + artifact = tmp_path / "packages" / "data-designer-slurm" / "tests" / "config" / "test_loading_builder.py" + artifact.parent.mkdir(parents=True) + extended_sentinel = "super-secret-token" + "-production-value" + artifact.write_text(f'secret = "{extended_sentinel}"\n') + + result = _run_audit(artifact, audit_script=audit_script) + + assert result.returncode == 1 + assert "plaintext secret assignment" in result.stderr + + +def _run_audit(*paths: Path, audit_script: Path = AUDIT_SCRIPT) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(audit_script), *(str(path) for path in paths)], + cwd=REPOSITORY_ROOT, + check=False, + capture_output=True, + text=True, + timeout=10, + ) diff --git a/packages/data-designer/pyproject.toml b/packages/data-designer/pyproject.toml index 2da5276bd..7caca322f 100644 --- a/packages/data-designer/pyproject.toml +++ b/packages/data-designer/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "data-designer" -dynamic = ["version", "dependencies"] +dynamic = ["version", "dependencies", "optional-dependencies"] description = "General framework for synthetic data generation" readme = "README.md" requires-python = ">=3.10" @@ -56,6 +56,7 @@ dependencies = [ "rich>=13.7.1,<15", "typer>=0.12.0,<1", ] +optional-dependencies = { slurm = ["data-designer-slurm=={{ version }}"] } [tool.hatch.build.targets.wheel] packages = ["src/data_designer"] @@ -74,3 +75,4 @@ package = true [tool.uv.sources] data-designer-config = { workspace = true } data-designer-engine = { workspace = true } +data-designer-slurm = { workspace = true } diff --git a/packages/data-designer/src/data_designer/cli/lazy_group.py b/packages/data-designer/src/data_designer/cli/lazy_group.py index 6b6f055aa..f0f0a49bf 100644 --- a/packages/data-designer/src/data_designer/cli/lazy_group.py +++ b/packages/data-designer/src/data_designer/cli/lazy_group.py @@ -4,11 +4,205 @@ from __future__ import annotations import importlib +import importlib.metadata +from collections import defaultdict from typing import Any import click import typer -from typer.core import TyperGroup +from typer.core import TyperCommand, TyperGroup + +CLI_EXTENSION_ENTRY_POINT_GROUP = "data_designer.cli" + + +def create_lazy_typer_group( + lazy_subcommands: dict[str, dict[str, str]], + *, + entry_point_group: str | None = None, +) -> type[TyperGroup]: + """Factory that returns a ``TyperGroup`` subclass with lazy-loaded commands. + + ``list_commands`` includes lazy command names so that ``--help`` works + without importing any command module. ``get_command`` returns a lightweight + ``_LazyCommand`` stub for lazy entries; the real Typer/Click command is only + built when the stub is invoked. + + Args: + lazy_subcommands: Mapping of command names to metadata dicts with keys: + - ``module``: Dotted module path (e.g. ``data_designer.cli.commands.preview``) + - ``attr``: Function attribute name in the module (e.g. ``preview_command``) + - ``help``: (optional) Short help text for group listing + - ``rich_help_panel``: (optional) Rich help panel name + entry_point_group: Optional entry-point group for lazy top-level command + extensions. Entry-point targets must be zero-argument callables that + return a ``click.Command``. + + Returns: + A ``TyperGroup`` subclass. + """ + + class LazyTyperGroup(TyperGroup): + _extension_entry_points: dict[str, list[importlib.metadata.EntryPoint]] | None = None + + def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]: + if not args and self.no_args_is_help and not ctx.resilient_parsing: + click.echo(ctx.get_help(), color=ctx.color) + ctx.exit(0) + return super().parse_args(ctx, args) + + def list_commands(self, ctx: click.Context) -> list[str]: + eager = super().list_commands(ctx) + lazy_names = [name for name in lazy_subcommands if name not in eager] + built_in_names = set(eager) | set(lazy_names) + extension_names = [name for name in self._discover_extension_entry_points() if name not in built_in_names] + return eager + sorted(lazy_names) + sorted(extension_names) + + def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: + cmd = super().get_command(ctx, cmd_name) + extension_entry_points = self._discover_extension_entry_points().get(cmd_name, []) + if extension_entry_points and (cmd is not None or cmd_name in lazy_subcommands): + providers = ", ".join(_entry_point_label(entry_point) for entry_point in extension_entry_points) + return _UnavailableCommand( + cmd_name, + f"CLI extension command {cmd_name!r} from {providers} conflicts with a built-in command.", + ) + if len(extension_entry_points) > 1: + providers = ", ".join(_entry_point_label(entry_point) for entry_point in extension_entry_points) + return _UnavailableCommand( + cmd_name, + f"CLI command {cmd_name!r} is provided by multiple extensions: {providers}.", + ) + if cmd is not None: + return cmd + if cmd_name in lazy_subcommands: + info = lazy_subcommands[cmd_name] + return _LazyCommand( + name=cmd_name, + module_path=info["module"], + attr_name=info["attr"], + help=info.get("help"), + rich_help_panel=info.get("rich_help_panel"), + ) + if extension_entry_points: + return _LazyEntryPointCommand(extension_entry_points[0]) + return None + + def _discover_extension_entry_points(self) -> dict[str, list[importlib.metadata.EntryPoint]]: + if self._extension_entry_points is not None: + return self._extension_entry_points + if entry_point_group is None: + self._extension_entry_points = {} + return self._extension_entry_points + + try: + entry_points = sorted( + importlib.metadata.entry_points(group=entry_point_group), + key=_entry_point_sort_key, + ) + except Exception as e: + click.echo(f"Warning: Failed to discover CLI extensions from {entry_point_group!r}: {e}.", err=True) + self._extension_entry_points = {} + return self._extension_entry_points + + discovered: defaultdict[str, list[importlib.metadata.EntryPoint]] = defaultdict(list) + for entry_point in entry_points: + if ( + not entry_point.name + or entry_point.name.startswith("-") + or any(character.isspace() for character in entry_point.name) + ): + click.echo( + f"Warning: Ignoring invalid CLI extension command name {entry_point.name!r} " + f"from {_entry_point_label(entry_point)}.", + err=True, + ) + continue + discovered[entry_point.name].append(entry_point) + self._extension_entry_points = dict(discovered) + return self._extension_entry_points + + return LazyTyperGroup + + +_DATA_DESIGNER_DISTRIBUTION = "data-designer" + + +def _distribution_name(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is None: + return "unknown-distribution" + return distribution.metadata.get("Name") or "unknown-distribution" + + +def _distribution_version(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is None: + return "unknown-version" + return distribution.version + + +def _entry_point_label(entry_point: importlib.metadata.EntryPoint) -> str: + return f"{_distribution_name(entry_point)} {_distribution_version(entry_point)} ({entry_point.value})" + + +def _entry_point_help(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is not None: + summary = distribution.metadata.get("Summary") + if summary: + return summary + return f"CLI extension provided by {_distribution_name(entry_point)}" + + +def _entry_point_sort_key(entry_point: importlib.metadata.EntryPoint) -> tuple[str, str, str, str, str]: + distribution_name = _distribution_name(entry_point) + return ( + entry_point.name, + distribution_name.casefold(), + distribution_name, + _distribution_version(entry_point), + entry_point.value, + ) + + +def _validate_entry_point_compatibility(entry_point: importlib.metadata.EntryPoint) -> None: + packaging_requirements = importlib.import_module("packaging.requirements") + packaging_utils = importlib.import_module("packaging.utils") + + distribution = getattr(entry_point, "dist", None) + label = _entry_point_label(entry_point) + if distribution is None: + raise click.ClickException(f"CLI extension {entry_point.name!r} has no owning distribution: {label}.") + + try: + requirements = [packaging_requirements.Requirement(value) for value in distribution.requires or []] + except packaging_requirements.InvalidRequirement as e: + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} has invalid dependency metadata: {e}." + ) from None + + data_designer_requirements = [ + requirement + for requirement in requirements + if packaging_utils.canonicalize_name(requirement.name) == _DATA_DESIGNER_DISTRIBUTION + and (requirement.marker is None or requirement.marker.evaluate()) + ] + if not data_designer_requirements: + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} must declare a dependency on data-designer." + ) + + try: + installed_version = importlib.metadata.version(_DATA_DESIGNER_DISTRIBUTION) + except importlib.metadata.PackageNotFoundError: + raise click.ClickException("Unable to resolve the installed data-designer version.") from None + + if not all(installed_version in requirement.specifier for requirement in data_designer_requirements): + expected = " and ".join(str(requirement) for requirement in data_designer_requirements) + raise click.ClickException( + f"CLI extension {entry_point.name!r} from {label} is incompatible with " + f"data-designer {installed_version}; requires {expected}." + ) class _LazyCommand(click.Command): @@ -35,6 +229,15 @@ def __init__( self._resolved: click.Command | None = None self.rich_help_panel = rich_help_panel + def make_context( + self, + info_name: str, + args: list[str], + parent: click.Context | None = None, + **extra: Any, + ) -> click.Context: + return self._resolve().make_context(info_name, args, parent, **extra) + def _resolve(self) -> click.Command: if self._resolved is not None: return self._resolved @@ -51,6 +254,14 @@ def _resolve(self) -> click.Command: self._resolved = click_cmd return self._resolved + +class _LazyEntryPointCommand(click.Command): + def __init__(self, entry_point: importlib.metadata.EntryPoint) -> None: + super().__init__(name=entry_point.name, help=_entry_point_help(entry_point)) + self._entry_point = entry_point + self._resolved: click.Command | None = None + self.rich_help_panel = "Extensions" + def make_context( self, info_name: str, @@ -60,53 +271,44 @@ def make_context( ) -> click.Context: return self._resolve().make_context(info_name, args, parent, **extra) + def _resolve(self) -> click.Command: + if self._resolved is not None: + return self._resolved -def create_lazy_typer_group( - lazy_subcommands: dict[str, dict[str, str]], -) -> type[TyperGroup]: - """Factory that returns a ``TyperGroup`` subclass with lazy-loaded commands. + _validate_entry_point_compatibility(self._entry_point) + label = _entry_point_label(self._entry_point) + try: + factory = self._entry_point.load() + except Exception as e: + raise click.ClickException(f"Failed to load CLI extension {self.name!r} from {label}: {e}.") from None + if not callable(factory): + raise click.ClickException(f"CLI extension {self.name!r} from {label} must load a zero-argument callable.") - ``list_commands`` includes lazy command names so that ``--help`` works - without importing any command module. ``get_command`` returns a lightweight - ``_LazyCommand`` stub for lazy entries; the real Typer/Click command is only - built when the stub is invoked. + try: + command = factory() + except Exception as e: + raise click.ClickException(f"Failed to create CLI extension {self.name!r} from {label}: {e}.") from None + if not isinstance(command, (click.Command, TyperCommand, TyperGroup)): + raise click.ClickException( + f"CLI extension {self.name!r} from {label} returned {type(command).__name__}, expected click.Command." + ) - Args: - lazy_subcommands: Mapping of command names to metadata dicts with keys: - - ``module``: Dotted module path (e.g. ``data_designer.cli.commands.preview``) - - ``attr``: Function attribute name in the module (e.g. ``preview_command``) - - ``help``: (optional) Short help text for group listing - - ``rich_help_panel``: (optional) Rich help panel name + command.name = self.name + self._resolved = command + return command - Returns: - A ``TyperGroup`` subclass. - """ - - class LazyTyperGroup(TyperGroup): - def parse_args(self, ctx: click.Context, args: list[str]) -> list[str]: - if not args and self.no_args_is_help and not ctx.resilient_parsing: - click.echo(ctx.get_help(), color=ctx.color) - ctx.exit(0) - return super().parse_args(ctx, args) - def list_commands(self, ctx: click.Context) -> list[str]: - eager = super().list_commands(ctx) - lazy_names = [name for name in lazy_subcommands if name not in eager] - return eager + sorted(lazy_names) +class _UnavailableCommand(click.Command): + def __init__(self, name: str, message: str) -> None: + super().__init__(name=name, help=f"Unavailable: {message}") + self._message = message + self.rich_help_panel = "Extensions" - def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: - cmd = super().get_command(ctx, cmd_name) - if cmd is not None: - return cmd - if cmd_name in lazy_subcommands: - info = lazy_subcommands[cmd_name] - return _LazyCommand( - name=cmd_name, - module_path=info["module"], - attr_name=info["attr"], - help=info.get("help"), - rich_help_panel=info.get("rich_help_panel"), - ) - return None - - return LazyTyperGroup + def make_context( + self, + info_name: str, + args: list[str], + parent: click.Context | None = None, + **extra: Any, + ) -> click.Context: + raise click.ClickException(self._message) diff --git a/packages/data-designer/src/data_designer/cli/main.py b/packages/data-designer/src/data_designer/cli/main.py index 502631048..26bf335f3 100644 --- a/packages/data-designer/src/data_designer/cli/main.py +++ b/packages/data-designer/src/data_designer/cli/main.py @@ -10,7 +10,7 @@ import typer from data_designer.cli.agent_command_defs import AGENT_COMMANDS -from data_designer.cli.lazy_group import create_lazy_typer_group +from data_designer.cli.lazy_group import CLI_EXTENSION_ENTRY_POINT_GROUP, create_lazy_typer_group from data_designer.cli.runtime import ensure_cli_default_model_settings from data_designer.config.utils.constants import DATA_DESIGNER_PACKAGE_NAME @@ -84,7 +84,8 @@ def _is_version_request(args: list[str]) -> bool: "help": "Check that every referenced model and MCP tool is reachable", "rich_help_panel": "Generation", }, - } + }, + entry_point_group=CLI_EXTENSION_ENTRY_POINT_GROUP, ), add_completion=False, no_args_is_help=True, diff --git a/packages/data-designer/src/data_designer/interface/__init__.py b/packages/data-designer/src/data_designer/interface/__init__.py index 7a1b0ed69..c9677ea14 100644 --- a/packages/data-designer/src/data_designer/interface/__init__.py +++ b/packages/data-designer/src/data_designer/interface/__init__.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: from data_designer.config.run_config import ResumeMode # noqa: F401 + from data_designer.engine.dataset_builders.errors import ArtifactStorageError # noqa: F401 from data_designer.interface.composite_workflow import ( # noqa: F401 CompositeWorkflow, CompositeWorkflowResults, @@ -33,6 +34,7 @@ ) _LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "ArtifactStorageError": ("data_designer.engine.dataset_builders.errors", "ArtifactStorageError"), "CompositeWorkflow": ("data_designer.interface.composite_workflow", "CompositeWorkflow"), "CompositeWorkflowResults": ("data_designer.interface.composite_workflow", "CompositeWorkflowResults"), "CompletedWorkflowStageMetadata": ( diff --git a/packages/data-designer/src/data_designer/interface/composite_workflow.py b/packages/data-designer/src/data_designer/interface/composite_workflow.py index 021374a3f..dcb3fd7a5 100644 --- a/packages/data-designer/src/data_designer/interface/composite_workflow.py +++ b/packages/data-designer/src/data_designer/interface/composite_workflow.py @@ -382,6 +382,9 @@ def run( stage=stage, stage_dir_name=stage_dir_name, stage_builder=stage_builder, + requested_num_records=stage_metadata[ + "num_records_actual" if stage.output_processors else "num_records_requested" + ], ) stage_results[stage.name] = output_result stage_output_paths[stage.name] = output_seed_path @@ -663,6 +666,9 @@ def _get_prior_stage_metadata( def _can_skip_prior_stage(stage: _WorkflowStage, prior_stage_metadata: dict[str, Any], workflow_path: Path) -> bool: if prior_stage_metadata.get("status") not in COMPLETED_STAGE_STATUSES: return False + requested_records_key = "num_records_actual" if stage.output_processors else "num_records_requested" + if not isinstance(prior_stage_metadata.get(requested_records_key), int): + return False if stage.on_success is not None and stage.on_success_version is None: return False output_seed_path = prior_stage_metadata.get("output_seed_path") @@ -704,6 +710,7 @@ def _stage_result_from_metadata( stage: _WorkflowStage, stage_dir_name: str, stage_builder: DataDesignerConfigBuilder, + requested_num_records: int, ) -> DatasetCreationResults: main_storage = ArtifactStorage(artifact_path=workflow_path, dataset_name=stage_dir_name, resume=ResumeMode.ALWAYS) result_storage = main_storage @@ -724,6 +731,7 @@ def _stage_result_from_metadata( analysis=_load_stage_analysis(result_storage), config_builder=result_builder, dataset_metadata=DatasetMetadata(), + requested_num_records=requested_num_records, ) diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index f773686d0..75e576edf 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -267,10 +267,12 @@ def create( Returns: DatasetCreationResults object with methods for loading the generated dataset, - analysis results, and displaying sample records for inspection. + analysis results, displaying sample records, and inspecting completion, + artifact, and resume metadata. Raises: DataDesignerGenerationError: If an error occurs during dataset generation. + DataDesignerEarlyShutdownError: If early shutdown is triggered before any records are produced. DataDesignerProfilingError: If an error occurs during dataset profiling. """ logger.info("🎨 Creating Data Designer dataset") @@ -378,6 +380,10 @@ def create( config_builder=config_builder, dataset_metadata=dataset_metadata, task_traces=task_traces, + requested_num_records=num_records, + early_shutdown=builder.early_shutdown, + requested_resume_mode=resume, + effective_resume_mode=builder.artifact_storage.resume, ) async def acreate( diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py index 6c5b076a3..1b0f3589e 100644 --- a/packages/data-designer/src/data_designer/interface/results.py +++ b/packages/data-designer/src/data_designer/interface/results.py @@ -11,6 +11,7 @@ from data_designer.config.config_builder import DataDesignerConfigBuilder from data_designer.config.dataset_metadata import DatasetMetadata from data_designer.config.errors import InvalidFileFormatError +from data_designer.config.run_config import ResumeMode from data_designer.config.seed_source_dataframe import DataFrameSeedSource from data_designer.config.utils.visualization import WithRecordSamplerMixin from data_designer.engine.dataset_builders.errors import ArtifactStorageError @@ -49,7 +50,11 @@ def __init__( analysis: DatasetProfilerResults, config_builder: DataDesignerConfigBuilder, dataset_metadata: DatasetMetadata, + requested_num_records: int, task_traces: list[TaskTrace] | None = None, + early_shutdown: bool | None = None, + requested_resume_mode: ResumeMode | None = None, + effective_resume_mode: ResumeMode | None = None, ): """Creates a new instance with results based on a dataset creation run. @@ -58,16 +63,42 @@ def __init__( analysis: Profiling results for the generated dataset. config_builder: Configuration builder used to create the dataset. dataset_metadata: Metadata about the generated dataset (e.g., seed column names). + requested_num_records: Number of records requested for this invocation. task_traces: Optional list of TaskTrace objects from the async scheduler. Resume note: only contains traces for the current invocation; traces from earlier ``create()`` calls that this run resumed are not retained. + early_shutdown: Whether generation stopped at the early-shutdown gate. + ``None`` when no generation invocation produced this result object. + requested_resume_mode: Resume mode requested for this invocation, or + ``None`` when no generation invocation produced this result object. + effective_resume_mode: Resume mode selected after compatibility checks, + or ``None`` when no generation invocation produced this result object. """ self.artifact_storage = artifact_storage self._analysis = analysis self._config_builder = config_builder self.dataset_metadata = dataset_metadata self.task_traces: list[TaskTrace] = task_traces or [] + self.requested_num_records = requested_num_records + self.early_shutdown = early_shutdown + self.requested_resume_mode = requested_resume_mode + self.effective_resume_mode = effective_resume_mode + + @property + def dataset_path(self) -> Path: + """Return the resolved dataset directory for this result.""" + return self.artifact_storage.base_dataset_path + + @property + def actual_num_records(self) -> int: + """Return the total number of records in the result dataset.""" + return self.count_records() + + @property + def is_partial(self) -> bool: + """Return whether the result contains fewer records than requested.""" + return self.actual_num_records < self.requested_num_records def load_analysis(self) -> DatasetProfilerResults: """Load the profiling analysis results for the generated dataset. @@ -83,8 +114,17 @@ def load_dataset(self) -> pd.DataFrame: Returns: A pandas DataFrame containing the full generated dataset. + + Raises: + ArtifactStorageError: If the dataset artifacts are missing or unreadable. """ - return self.artifact_storage.load_dataset() + try: + dataset = self.artifact_storage.load_dataset() + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to load dataset artifacts: {e}") from e + if dataset is not None and dataset.empty: + self._get_batch_files() + return dataset def to_config_builder(self, columns: list[str] | None = None) -> DataDesignerConfigBuilder: """Create a new config builder seeded from this result dataset. @@ -108,9 +148,21 @@ def count_records(self) -> int: Returns: Total row count across all batch parquet files. + + Raises: + ArtifactStorageError: If the dataset artifacts are missing or unreadable. """ + batch_files = self._get_batch_files() + try: + return sum(lazy.pq.read_metadata(f).num_rows for f in batch_files) + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to read dataset artifacts: {e}") from e + + def _get_batch_files(self) -> list[Path]: batch_files = sorted(self.artifact_storage.final_dataset_path.glob("batch_*.parquet")) - return sum(lazy.pq.read_metadata(f).num_rows for f in batch_files) + if not batch_files: + raise ArtifactStorageError("No batch parquet files found.") + return batch_files def load_processor_dataset(self, processor_name: str) -> pd.DataFrame: """Load the dataset generated by a processor. @@ -122,8 +174,14 @@ def load_processor_dataset(self, processor_name: str) -> pd.DataFrame: Returns: A pandas DataFrame containing the dataset generated by the processor. + + Raises: + ArtifactStorageError: If the processor artifacts are missing or unreadable. """ - return self.artifact_storage.load_processor_dataset(processor_name) + try: + return self.artifact_storage.load_processor_dataset(processor_name) + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to load processor dataset artifacts: {e}") from e def get_path_to_processor_artifacts(self, processor_name: str) -> Path: """Get the path to the artifacts generated by a processor. @@ -161,7 +219,7 @@ def export(self, path: Path | str, *, format: ExportFormat | None = None) -> Pat Raises: InvalidFileFormatError: If the format cannot be determined or is not one of the supported values. - ArtifactStorageError: If no batch parquet files are found. + ArtifactStorageError: If batch parquet files are missing or unreadable. Example: >>> results = data_designer.create(config, num_records=1000) @@ -178,9 +236,7 @@ def export(self, path: Path | str, *, format: ExportFormat | None = None) -> Pat raise InvalidFileFormatError( f"Unsupported export format: {resolved_format!r}. Choose one of: {', '.join(SUPPORTED_EXPORT_FORMATS)}." ) - batch_files = sorted(self.artifact_storage.final_dataset_path.glob("batch_*.parquet")) - if not batch_files: - raise ArtifactStorageError("No batch parquet files found to export.") + batch_files = self._get_batch_files() if resolved_format == "jsonl": _export_jsonl(batch_files, path) elif resolved_format == "csv": @@ -243,7 +299,10 @@ def _export_jsonl(batch_files: list[Path], output: Path) -> None: """ with output.open("w", encoding="utf-8") as f: for batch_file in batch_files: - chunk = lazy.pd.read_parquet(batch_file) + try: + chunk = lazy.pd.read_parquet(batch_file) + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to read dataset artifact {batch_file}: {e}") from e content = chunk.to_json(orient="records", lines=True, force_ascii=False, date_format="iso") if content: f.write(content) @@ -252,7 +311,10 @@ def _export_jsonl(batch_files: list[Path], output: Path) -> None: def _export_csv(batch_files: list[Path], output: Path) -> None: """Write *batch_files* to *output* as CSV with a single header row.""" for i, batch_file in enumerate(batch_files): - chunk = lazy.pd.read_parquet(batch_file) + try: + chunk = lazy.pd.read_parquet(batch_file) + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to read dataset artifact {batch_file}: {e}") from e chunk.to_csv(output, mode="a" if i > 0 else "w", header=(i == 0), index=False) @@ -267,7 +329,10 @@ def _export_parquet(batch_files: list[Path], output: Path) -> None: InvalidFileFormatError: If batch schemas have incompatible column names or types that cannot be unified or cast. """ - schemas = [lazy.pq.read_schema(f) for f in batch_files] + try: + schemas = [lazy.pq.read_schema(f) for f in batch_files] + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to read dataset artifacts: {e}") from e try: # promote_options="permissive" allows minor numeric type drift (e.g. int64 → double) unified_schema = lazy.pa.unify_schemas(schemas, promote_options="permissive") @@ -275,7 +340,10 @@ def _export_parquet(batch_files: list[Path], output: Path) -> None: raise InvalidFileFormatError(f"Cannot unify batch schemas for parquet export: {e}") from e with lazy.pq.ParquetWriter(output, unified_schema) as writer: for batch_file in batch_files: - table = lazy.pq.read_table(batch_file) + try: + table = lazy.pq.read_table(batch_file) + except (OSError, lazy.pa.ArrowException) as e: + raise ArtifactStorageError(f"Failed to read dataset artifact {batch_file}: {e}") from e try: writer.write_table(table.cast(unified_schema)) except (lazy.pa.ArrowInvalid, ValueError) as e: diff --git a/packages/data-designer/tests/cli/test_lazy_group.py b/packages/data-designer/tests/cli/test_lazy_group.py new file mode 100644 index 000000000..26203f014 --- /dev/null +++ b/packages/data-designer/tests/cli/test_lazy_group.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.metadata +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import click +import pytest +import typer +from typer.testing import CliRunner + +from data_designer.cli.lazy_group import create_lazy_typer_group + +ENTRY_POINT_GROUP = "test.data_designer.cli" +runner = CliRunner() + + +def _normalized_output(output: str) -> str: + return " ".join(line.strip(" │") for line in click.unstyle(output).splitlines()) + + +def _app() -> typer.Typer: + app = typer.Typer(cls=create_lazy_typer_group({}, entry_point_group=ENTRY_POINT_GROUP)) + + @app.callback() + def callback() -> None: + pass + + @app.command() + def base() -> None: + click.echo("base") + + return app + + +def _command(output: str = "extension") -> click.Command: + @click.command() + def run() -> None: + click.echo(output) + + return click.Group(name="extension", commands={"run": run}) + + +def _nested_command() -> click.Command: + @click.command() + @click.option("--value", required=True) + @click.pass_context + def init(ctx: click.Context, value: str) -> None: + click.echo(value) + ctx.exit(3) + + profile = click.Group(name="profile", commands={"init": init}) + return click.Group(name="extension", commands={"profile": profile}) + + +def _entry_point( + name: str = "slurm", + *, + distribution_name: str = "data-designer-slurm", + requirements: list[str] | None = None, + loaded: object | None = None, + load_error: Exception | None = None, +) -> SimpleNamespace: + distribution = SimpleNamespace( + metadata={"Name": distribution_name, "Summary": f"{distribution_name} summary"}, + requires=requirements if requirements is not None else ["data-designer>=0"], + version="1.0.0", + ) + load = ( + Mock(side_effect=load_error) + if load_error is not None + else Mock(return_value=loaded or Mock(return_value=_command())) + ) + return SimpleNamespace( + name=name, + value=f"{distribution_name}.cli:create_cli", + group=ENTRY_POINT_GROUP, + dist=distribution, + load=load, + ) + + +def test_no_extensions_preserves_help_and_built_in_commands() -> None: + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[]): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert "base" in help_result.output + assert "slurm" not in help_result.output + assert command_result.exit_code == 0 + assert command_result.output == "base\n" + + +def test_root_help_lists_extension_without_loading_it() -> None: + entry_point = _entry_point() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["--help"]) + + assert result.exit_code == 0 + assert "slurm" in result.output + assert "data-designer-slurm summary" in result.output + entry_point.load.assert_not_called() + + +def test_root_help_uses_fallback_when_distribution_summary_is_missing() -> None: + entry_point = _entry_point() + entry_point.dist.metadata.pop("Summary") + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["--help"]) + + assert result.exit_code == 0 + assert "CLI extension provided by data-designer-slurm" in result.output + entry_point.load.assert_not_called() + + +def test_built_in_command_does_not_load_extension() -> None: + entry_point = _entry_point() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["base"]) + + assert result.exit_code == 0 + assert result.output == "base\n" + entry_point.load.assert_not_called() + + +def test_selected_extension_loads_and_dispatches() -> None: + factory = Mock(return_value=_command("selected")) + entry_point = _entry_point(loaded=factory) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm", "run"]) + + assert result.exit_code == 0 + assert result.output == "selected\n" + entry_point.load.assert_called_once_with() + factory.assert_called_once_with() + + +def test_selected_extension_dispatches_nested_command_and_preserves_exit_code() -> None: + factory = Mock(return_value=_nested_command()) + entry_point = _entry_point(loaded=factory) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm", "profile", "init", "--value", "created"]) + + assert result.exit_code == 3 + assert result.output == "created\n" + entry_point.load.assert_called_once_with() + factory.assert_called_once_with() + + +def test_selecting_one_extension_does_not_load_another() -> None: + alpha = _entry_point("alpha", distribution_name="alpha-extension") + beta = _entry_point("beta", distribution_name="beta-extension") + with patch.object(importlib.metadata, "entry_points", return_value=[beta, alpha]): + result = runner.invoke(_app(), ["alpha", "run"]) + + assert result.exit_code == 0 + alpha.load.assert_called_once_with() + beta.load.assert_not_called() + + +def test_broken_extension_fails_only_when_selected() -> None: + entry_point = _entry_point(load_error=ImportError("missing dependency")) + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(app, ["--help"]) + built_in_result = runner.invoke(app, ["base"]) + command_result = runner.invoke(app, ["slurm"]) + + assert help_result.exit_code == 0 + assert built_in_result.exit_code == 0 + assert built_in_result.output == "base\n" + assert command_result.exit_code == 1 + assert "Failed to load CLI extension 'slurm'" in command_result.output + assert "data-designer-slurm 1.0.0" in command_result.output + assert "missing dependency" in command_result.output + + +def test_extension_target_must_be_callable() -> None: + entry_point = _entry_point(loaded=object()) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "must load a zero-argument callable" in result.output + + +def test_extension_factory_must_return_click_command() -> None: + entry_point = _entry_point(loaded=Mock(return_value=object())) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "returned object" in result.output + assert "click.Command" in result.output + + +def test_extension_factory_error_has_distribution_context() -> None: + entry_point = _entry_point(loaded=Mock(side_effect=RuntimeError("factory failed"))) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "Failed to create CLI extension 'slurm'" in result.output + assert "data-designer-slurm 1.0.0" in result.output + assert "factory failed" in result.output + + +def test_extension_without_owning_distribution_fails_before_import() -> None: + entry_point = _entry_point() + entry_point.dist = None + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "has no owning distribution" in result.output + entry_point.load.assert_not_called() + + +def test_invalid_dependency_metadata_fails_before_import() -> None: + entry_point = _entry_point(requirements=["data-designer=>1"]) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "has invalid dependency metadata" in result.output + entry_point.load.assert_not_called() + + +def test_missing_base_distribution_version_fails_before_import() -> None: + entry_point = _entry_point() + with ( + patch.object(importlib.metadata, "entry_points", return_value=[entry_point]), + patch.object( + importlib.metadata, + "version", + side_effect=importlib.metadata.PackageNotFoundError("data-designer"), + ), + ): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert "Unable to resolve the installed data-designer version" in result.output + entry_point.load.assert_not_called() + + +def test_duplicate_extensions_fail_deterministically_without_loading() -> None: + alpha = _entry_point(distribution_name="alpha-extension") + beta = _entry_point(distribution_name="beta-extension") + outputs = [] + + for entry_points in ([beta, alpha], [alpha, beta]): + with patch.object(importlib.metadata, "entry_points", return_value=entry_points): + result = runner.invoke(_app(), ["slurm"]) + assert result.exit_code == 1 + outputs.append(result.output) + + assert outputs[0] == outputs[1] + assert "CLI command 'slurm' is provided by multiple extensions" in outputs[0] + assert outputs[0].index("alpha-extension") < outputs[0].index("beta-extension") + alpha.load.assert_not_called() + beta.load.assert_not_called() + + +def test_extension_cannot_replace_built_in_command() -> None: + entry_point = _entry_point("base") + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(_app(), ["--help"]) + command_result = runner.invoke(_app(), ["base"]) + + assert help_result.exit_code == 0 + assert "Unavailable" in help_result.output + assert command_result.exit_code == 1 + assert "conflicts with a built-in command" in command_result.output + entry_point.load.assert_not_called() + + +@pytest.mark.parametrize( + ("requirements", "message"), + [ + ([], "must declare a dependency on data-designer"), + (["data-designer<0"], "is incompatible with data-designer"), + (["data-designer>=0", "data-designer<0"], "is incompatible with data-designer"), + ], +) +def test_incompatible_extension_fails_before_import(requirements: list[str], message: str) -> None: + entry_point = _entry_point(requirements=requirements) + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + result = runner.invoke(_app(), ["slurm"]) + + assert result.exit_code == 1 + assert message in _normalized_output(result.output) + entry_point.load.assert_not_called() + + +def test_invalid_extension_name_is_ignored_without_breaking_built_in_commands() -> None: + entry_point = _entry_point("bad name") + app = _app() + with patch.object(importlib.metadata, "entry_points", return_value=[entry_point]): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert "Ignoring invalid CLI extension command name 'bad name'" in help_result.output + assert "data-designer-slurm 1.0.0" in help_result.output + assert command_result.exit_code == 0 + assert command_result.output.endswith("base\n") + entry_point.load.assert_not_called() + + +def test_discovery_error_warns_without_breaking_built_in_commands() -> None: + app = _app() + with patch.object(importlib.metadata, "entry_points", side_effect=RuntimeError("corrupt metadata")): + help_result = runner.invoke(app, ["--help"]) + command_result = runner.invoke(app, ["base"]) + + assert help_result.exit_code == 0 + assert f"Failed to discover CLI extensions from {ENTRY_POINT_GROUP!r}: corrupt metadata" in help_result.output + assert command_result.exit_code == 0 + assert command_result.output.endswith("base\n") diff --git a/packages/data-designer/tests/interface/test_acreate.py b/packages/data-designer/tests/interface/test_acreate.py index 30c452fb2..061e21530 100644 --- a/packages/data-designer/tests/interface/test_acreate.py +++ b/packages/data-designer/tests/interface/test_acreate.py @@ -38,6 +38,7 @@ def _creation_result( analysis=stub_dataset_profiler_results, config_builder=config_builder, dataset_metadata=DatasetMetadata(), + requested_num_records=1, ) diff --git a/packages/data-designer/tests/interface/test_composite_workflow.py b/packages/data-designer/tests/interface/test_composite_workflow.py index 48e0c749b..73b7fb989 100644 --- a/packages/data-designer/tests/interface/test_composite_workflow.py +++ b/packages/data-designer/tests/interface/test_composite_workflow.py @@ -62,6 +62,7 @@ def _result_from_df( analysis=stub_dataset_profiler_results, config_builder=config_builder, dataset_metadata=DatasetMetadata(), + requested_num_records=len(df), ) @@ -633,6 +634,65 @@ def test_composite_workflow_resume_if_possible_skips_legacy_completed_stages( assert resumed_metadata["workflow_extension"] is True assert [stage["stage_extension"] for stage in resumed_metadata["stages"]] == [{"value": 0}, {"value": 1}] assert all("stage_output_override_path" not in stage for stage in resumed_metadata["stages"]) + base_result = results["base"] + assert isinstance(base_result, DatasetCreationResults) + assert base_result.requested_num_records == 3 + assert base_result.actual_num_records == 3 + assert base_result.is_partial is False + assert base_result.early_shutdown is None + assert base_result.requested_resume_mode is None + assert base_result.effective_resume_mode is None + + +def test_composite_workflow_resume_if_possible_missing_requested_count_reruns_stages( + stub_artifact_path: Path, + stub_model_providers: list[ModelProvider], + stub_model_configs: list[ModelConfig], + stub_dataset_profiler_results, +) -> None: + data_designer = _data_designer(stub_artifact_path, stub_model_providers) + create_mock = _patch_create(data_designer, stub_dataset_profiler_results) + workflow = data_designer.compose_workflow(name="resume-missing-count") + workflow.add_stage("base", _category_builder(stub_model_configs), num_records=3) + workflow.add_stage("copy", _copy_builder(stub_model_configs)) + workflow.run() + metadata_path = stub_artifact_path / "resume-missing-count" / "workflow-metadata.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + metadata["stages"][0].pop("num_records_requested") + metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + create_mock.reset_mock() + + resumed = data_designer.compose_workflow(name="resume-missing-count") + resumed.add_stage("base", _category_builder(stub_model_configs), num_records=3) + resumed.add_stage("copy", _copy_builder(stub_model_configs)) + resumed.run(resume=ResumeMode.IF_POSSIBLE) + + assert [call.kwargs["dataset_name"] for call in create_mock.call_args_list] == ["stage-0-base", "stage-1-copy"] + + +def test_composite_workflow_reused_stage_reports_partial_result( + stub_artifact_path: Path, + stub_model_providers: list[ModelProvider], + stub_model_configs: list[ModelConfig], + stub_dataset_profiler_results, +) -> None: + data_designer = _data_designer(stub_artifact_path, stub_model_providers) + create_mock = _patch_create(data_designer, stub_dataset_profiler_results) + workflow = data_designer.compose_workflow(name="resume-partial") + workflow.add_stage("base", _category_builder(stub_model_configs), num_records=3) + first = workflow.run() + batch_path = first["base"].artifact_storage.final_dataset_path / "batch_00000.parquet" + lazy.pd.read_parquet(batch_path).head(2).to_parquet(batch_path, index=False) + create_mock.reset_mock() + + resumed = data_designer.compose_workflow(name="resume-partial") + resumed.add_stage("base", _category_builder(stub_model_configs), num_records=3) + results = resumed.run(resume=ResumeMode.IF_POSSIBLE) + + assert create_mock.call_count == 0 + assert results["base"].requested_num_records == 3 + assert results["base"].actual_num_records == 2 + assert results["base"].is_partial is True def test_composite_workflow_resume_if_possible_skips_stage_with_output_processors( diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 68990ff5c..e510841ec 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -44,6 +44,7 @@ HuggingFaceSeedSource, ) from data_designer.config.seed_source_dataframe import DataFrameSeedSource +from data_designer.engine.dataset_builders.dataset_builder import DatasetBuilder from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode from data_designer.engine.models.errors import ( RETRYABLE_MODEL_ERRORS, @@ -681,6 +682,13 @@ def test_create_dataset_e2e_using_only_sampler_columns( analysis = results.load_analysis() assert analysis.target_num_records == num_records + assert results.dataset_path == stub_artifact_path / "dataset" + assert results.requested_num_records == num_records + assert results.actual_num_records == num_records + assert results.is_partial is False + assert results.early_shutdown is False + assert results.requested_resume_mode == ResumeMode.NEVER + assert results.effective_resume_mode == ResumeMode.NEVER # display report with no errors analysis.to_report() @@ -722,6 +730,104 @@ def test_rounded_sampler_constraints_and_conditions_e2e( assert (frame["label"] == "high").all() +def test_create_reports_partial_early_shutdown( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + original_build = DatasetBuilder.build + + def build_partial(builder: DatasetBuilder, **kwargs: Any) -> Path: + kwargs["num_records"] = 3 + return original_build(builder, **kwargs) + + with ( + patch.object(DatasetBuilder, "build", new=build_partial), + patch.object(DatasetBuilder, "early_shutdown", new_callable=PropertyMock, return_value=True), + ): + results = data_designer.create( + stub_sampler_only_config_builder, + num_records=10, + dataset_name="partial", + ) + + assert results.requested_num_records == 10 + assert results.actual_num_records == 3 + assert results.is_partial is True + assert results.early_shutdown is True + + +def test_create_reports_resolved_colliding_dataset_path( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + + first = data_designer.create(stub_sampler_only_config_builder, num_records=1, dataset_name="collision") + second = data_designer.create(stub_sampler_only_config_builder, num_records=1, dataset_name="collision") + + assert first.dataset_path == stub_artifact_path / "collision" + assert second.dataset_path.parent == stub_artifact_path + assert second.dataset_path.name.startswith("collision_") + assert second.dataset_path != first.dataset_path + + +def test_create_reports_effective_resume_mode( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + + first = data_designer.create(stub_sampler_only_config_builder, num_records=2, dataset_name="resumable") + resumed = data_designer.create( + stub_sampler_only_config_builder, + num_records=3, + dataset_name="resumable", + resume=ResumeMode.IF_POSSIBLE, + ) + + assert resumed.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert resumed.effective_resume_mode == ResumeMode.ALWAYS + assert resumed.dataset_path == first.dataset_path + assert resumed.actual_num_records == 3 + + _add_irrelevant_sampler_column(stub_sampler_only_config_builder) + restarted = data_designer.create( + stub_sampler_only_config_builder, + num_records=3, + dataset_name="resumable", + resume=ResumeMode.IF_POSSIBLE, + ) + + assert restarted.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert restarted.effective_resume_mode == ResumeMode.NEVER + assert restarted.dataset_path != first.dataset_path + + def test_create_with_drop_true_can_skip_dropped_column_artifacts( stub_artifact_path, stub_model_providers, diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py index 200d43810..597133762 100644 --- a/packages/data-designer/tests/interface/test_results.py +++ b/packages/data-designer/tests/interface/test_results.py @@ -15,10 +15,11 @@ from data_designer.config.dataset_metadata import DatasetMetadata from data_designer.config.errors import InvalidFileFormatError from data_designer.config.preview_results import PreviewResults +from data_designer.config.run_config import ResumeMode from data_designer.config.utils.errors import DatasetSampleDisplayError from data_designer.config.utils.visualization import display_sample_record as display_fn -from data_designer.engine.dataset_builders.errors import ArtifactStorageError from data_designer.engine.storage.artifact_storage import ArtifactStorage +from data_designer.interface import ArtifactStorageError from data_designer.interface.results import DatasetCreationResults @@ -46,6 +47,7 @@ def stub_dataset_creation_results( analysis=stub_dataset_profiler_results, config_builder=stub_complete_builder, dataset_metadata=stub_dataset_metadata, + requested_num_records=len(stub_artifact_storage.load_dataset.return_value), ) @@ -56,11 +58,25 @@ def test_init(stub_artifact_storage, stub_dataset_profiler_results, stub_complet analysis=stub_dataset_profiler_results, config_builder=stub_complete_builder, dataset_metadata=stub_dataset_metadata, + requested_num_records=1, ) assert results.artifact_storage == stub_artifact_storage assert results._analysis == stub_dataset_profiler_results assert results._config_builder == stub_complete_builder assert results.dataset_metadata == stub_dataset_metadata + assert results.requested_num_records == 1 + + +def test_init_requires_requested_num_records( + stub_artifact_storage, stub_dataset_profiler_results, stub_complete_builder, stub_dataset_metadata +) -> None: + with pytest.raises(TypeError, match="requested_num_records"): + DatasetCreationResults( + artifact_storage=stub_artifact_storage, + analysis=stub_dataset_profiler_results, + config_builder=stub_complete_builder, + dataset_metadata=stub_dataset_metadata, + ) def test_load_dataset(stub_dataset_creation_results, stub_artifact_storage, stub_dataframe): @@ -199,12 +215,10 @@ def test_display_sample_record_with_empty_dataset(): analysis=MagicMock(spec=DatasetProfilerResults), config_builder=MagicMock(spec=DataDesignerConfigBuilder), dataset_metadata=DatasetMetadata(), + requested_num_records=1, ) - # Empty DataFrame is still a valid DataFrame, so accessing _record_sampler_dataset succeeds - # but display_sample_record fails when trying to access index 0 - # Note: Currently raises UnboundLocalError due to bug in error handling, but tests that it fails - with pytest.raises((DatasetSampleDisplayError, UnboundLocalError)): + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): results.display_sample_record() @@ -218,6 +232,7 @@ def test_display_sample_record_with_none_dataset(): analysis=MagicMock(spec=DatasetProfilerResults), config_builder=MagicMock(spec=DataDesignerConfigBuilder), dataset_metadata=DatasetMetadata(), + requested_num_records=1, ) # Mixin raises DatasetSampleDisplayError when dataset is None @@ -389,6 +404,112 @@ def test_count_records(stub_dataset_creation_results, stub_dataframe, stub_batch assert stub_dataset_creation_results.count_records() == len(stub_dataframe) +@pytest.mark.parametrize("error_type", [PermissionError, lazy.pa.ArrowInvalid]) +def test_count_records_normalizes_unreadable_artifact_error( + stub_dataset_creation_results, stub_batch_dir, error_type +) -> None: + stub_dataset_creation_results.artifact_storage.final_dataset_path = stub_batch_dir + + with ( + patch("data_designer.interface.results.lazy.pq.read_metadata", side_effect=error_type("unreadable")), + pytest.raises(ArtifactStorageError, match="Failed to read dataset artifacts") as exc_info, + ): + stub_dataset_creation_results.count_records() + + assert isinstance(exc_info.value.__cause__, error_type) + + +@pytest.mark.parametrize( + ("format", "reader"), + [ + ("jsonl", "data_designer.interface.results.lazy.pd.read_parquet"), + ("csv", "data_designer.interface.results.lazy.pd.read_parquet"), + ("parquet", "data_designer.interface.results.lazy.pq.read_schema"), + ], +) +@pytest.mark.parametrize("error_type", [PermissionError, lazy.pa.ArrowInvalid]) +def test_export_normalizes_unreadable_artifact_error( + stub_dataset_creation_results, stub_batch_dir, tmp_path, format: str, reader: str, error_type +) -> None: + stub_dataset_creation_results.artifact_storage.final_dataset_path = stub_batch_dir + + with ( + patch(reader, side_effect=error_type("unreadable")), + pytest.raises(ArtifactStorageError, match="Failed to read dataset artifact") as exc_info, + ): + stub_dataset_creation_results.export(tmp_path / f"out.{format}") + + assert isinstance(exc_info.value.__cause__, error_type) + + +def test_missing_dataset_artifacts_raise_public_error(stub_dataset_creation_results, tmp_path) -> None: + empty_dir = tmp_path / "parquet-files" + empty_dir.mkdir() + stub_dataset_creation_results.artifact_storage.final_dataset_path = empty_dir + stub_dataset_creation_results.artifact_storage.load_dataset.return_value = lazy.pd.DataFrame() + stub_dataset_creation_results.requested_num_records = 1 + + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + stub_dataset_creation_results.load_dataset() + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + stub_dataset_creation_results.count_records() + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + _ = stub_dataset_creation_results.actual_num_records + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + _ = stub_dataset_creation_results.is_partial + + +@pytest.mark.parametrize("error_type", [PermissionError, lazy.pa.ArrowInvalid]) +def test_load_dataset_normalizes_artifact_error(stub_dataset_creation_results, error_type) -> None: + stub_dataset_creation_results.artifact_storage.load_dataset.side_effect = error_type("unreadable") + + with pytest.raises(ArtifactStorageError, match="Failed to load dataset artifacts") as exc_info: + stub_dataset_creation_results.load_dataset() + + assert isinstance(exc_info.value.__cause__, error_type) + + +@pytest.mark.parametrize("error_type", [PermissionError, lazy.pa.ArrowInvalid]) +def test_load_processor_dataset_normalizes_artifact_error(stub_dataset_creation_results, error_type) -> None: + stub_dataset_creation_results.artifact_storage.load_processor_dataset.side_effect = error_type("unreadable") + + with pytest.raises(ArtifactStorageError, match="Failed to load processor dataset artifacts") as exc_info: + stub_dataset_creation_results.load_processor_dataset("processor") + + assert isinstance(exc_info.value.__cause__, error_type) + + +def test_public_creation_outcome_fields( + stub_artifact_storage, + stub_dataset_profiler_results, + stub_complete_builder, + stub_dataset_metadata, + stub_dataframe, + stub_batch_dir, + tmp_path, +) -> None: + dataset_path = tmp_path / "dataset" + stub_artifact_storage.base_dataset_path = dataset_path + stub_artifact_storage.final_dataset_path = stub_batch_dir + results = DatasetCreationResults( + artifact_storage=stub_artifact_storage, + analysis=stub_dataset_profiler_results, + config_builder=stub_complete_builder, + dataset_metadata=stub_dataset_metadata, + requested_num_records=len(stub_dataframe) + 1, + early_shutdown=True, + requested_resume_mode=ResumeMode.IF_POSSIBLE, + effective_resume_mode=ResumeMode.NEVER, + ) + + assert results.dataset_path == dataset_path + assert results.actual_num_records == len(stub_dataframe) + assert results.is_partial is True + assert results.early_shutdown is True + assert results.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert results.effective_resume_mode == ResumeMode.NEVER + + def test_export_uppercase_extension_is_recognised(stub_dataset_creation_results, stub_batch_dir, tmp_path) -> None: """export() treats file extensions case-insensitively (e.g. .JSONL → jsonl).""" stub_dataset_creation_results.artifact_storage.final_dataset_path = stub_batch_dir diff --git a/packages/data-designer/tests/test_dependency_audit.py b/packages/data-designer/tests/test_dependency_audit.py index 509433f94..8393a281b 100644 --- a/packages/data-designer/tests/test_dependency_audit.py +++ b/packages/data-designer/tests/test_dependency_audit.py @@ -45,10 +45,11 @@ def audit( workspace: Path, distributions: dict[str, list[str]], requirements: dict[str, list[str]] | None = None, + selected_extras: dict[str, set[str]] | None = None, ) -> dict: module = DEPENDENCY_AUDIT assert isinstance(module, ModuleType) - return module.audit_repository(workspace, distributions, requirements) + return module.audit_repository(workspace, distributions, requirements, selected_extras) def test_marks_transitively_guaranteed_gap_low(workspace: Path) -> None: @@ -144,3 +145,107 @@ def test_applies_module_distribution_override(tmp_path: Path) -> None: result = audit(tmp_path, {}) assert result["packages"][0]["missing"][0]["dependency"] == "pyyaml" + + +def test_includes_selected_dynamic_optional_dependencies(tmp_path: Path) -> None: + base = tmp_path / "packages" / "data-designer" + leaf = tmp_path / "packages" / "data-designer-slurm" + (base / "src").mkdir(parents=True) + (leaf / "src").mkdir(parents=True) + (base / "pyproject.toml").write_text( + """ +[project] +name = "data-designer" +dynamic = ["optional-dependencies"] + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +optional-dependencies = { slurm = ["data-designer-slurm=={{ version }}"] } +""".lstrip() + ) + (leaf / "pyproject.toml").write_text( + """ +[project] +name = "data-designer-slurm" +dynamic = ["dependencies"] + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +dependencies = ["data-designer=={{ version }}"] +""".lstrip() + ) + + result = audit(tmp_path, {}, selected_extras={"data-designer": {"slurm"}}) + + packages = {package["package"]: package for package in result["packages"]} + assert packages["data-designer"]["declared"] == ["data-designer-slurm"] + assert packages["data-designer-slurm"]["declared"] == ["data-designer"] + + +@pytest.mark.parametrize( + ("consumer_dependency", "expected_guaranteed_by"), + [("base", []), ("base[leaf]", ["base"])], +) +def test_scopes_selected_workspace_extra_to_dependency_edge( + tmp_path: Path, + consumer_dependency: str, + expected_guaranteed_by: list[str], +) -> None: + base = tmp_path / "packages" / "base" + leaf = tmp_path / "packages" / "leaf" + consumer = tmp_path / "packages" / "consumer" + for package in (base, leaf, consumer): + (package / "src").mkdir(parents=True) + (base / "pyproject.toml").write_text( + """ +[project] +name = "base" +dynamic = ["optional-dependencies"] + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +optional-dependencies = { leaf = ["leaf=={{ version }}"] } +""".lstrip() + ) + (leaf / "pyproject.toml").write_text('[project]\nname = "leaf"\ndependencies = []\n') + (consumer / "pyproject.toml").write_text( + f'[project]\nname = "consumer"\ndependencies = ["{consumer_dependency}"]\n' + ) + (consumer / "src" / "consumer.py").write_text("import leaf_module\n") + + result = audit( + tmp_path, + {"leaf_module": ["leaf"]}, + selected_extras={"base": {"leaf"}}, + ) + + consumer_result = next(package for package in result["packages"] if package["package"] == "consumer") + assert consumer_result["missing"][0]["guaranteed_by"] == expected_guaranteed_by + assert consumer_result["missing"][0]["severity"] == ("low" if expected_guaranteed_by else "high") + + +def test_does_not_follow_workspace_cycle_back_through_audited_project(tmp_path: Path) -> None: + base = tmp_path / "packages" / "base" + leaf = tmp_path / "packages" / "leaf" + for package in (base, leaf): + (package / "src").mkdir(parents=True) + (base / "pyproject.toml").write_text( + """ +[project] +name = "base" +dependencies = ["support"] +dynamic = ["optional-dependencies"] + +[tool.hatch.metadata.hooks.uv-dynamic-versioning] +optional-dependencies = { leaf = ["leaf=={{ version }}"] } +""".lstrip() + ) + (leaf / "pyproject.toml").write_text('[project]\nname = "leaf"\ndependencies = ["base"]\n') + (base / "src" / "base.py").write_text("import transitive_module\n") + + result = audit( + tmp_path, + {"transitive_module": ["transitive"]}, + {"support": ["transitive"]}, + selected_extras={"base": {"leaf"}}, + ) + + base_result = next(package for package in result["packages"] if package["package"] == "base") + assert base_result["missing"][0]["guaranteed_by"] == ["support"] diff --git a/plans/850/data-designer-contract.md b/plans/850/data-designer-contract.md new file mode 100644 index 000000000..bf41457a0 --- /dev/null +++ b/plans/850/data-designer-contract.md @@ -0,0 +1,330 @@ +--- +date: 2026-08-10 +authors: + - andreatnvidia +issue: https://github.com/NVIDIA-NeMo/DataDesigner/issues/851 +epic: https://github.com/NVIDIA-NeMo/DataDesigner/issues/850 +status: proposal +--- + +# Public Data Designer invocation contract + +## Summary + +An optional execution package can configure, validate, and invoke Data Designer without importing +`data_designer.engine`. The supported boundary consists of serialized builder input, public config models, +`DataDesigner`, `DatasetCreationResults`, and public errors. + +Plugin-independent envelope validation may run before optional plugin packages are installed. Full builder validation +must run in a fresh client process after those packages are installed and available on `sys.path`. + +## Public imports + +The execution package may import these public symbols: + +```python +from data_designer.config import ( + DataDesignerConfigBuilder, + InvalidConfigError, + InvalidFileFormatError, + InvalidFilePathError, + LocalStdioMCPProvider, + MCPProvider, + ModelConfig, + ModelProvider, + PartitionBlock, + ResumeMode, + RunConfig, + SamplingStrategy, +) +from data_designer.interface import ( + ArtifactStorageError, + DataDesigner, + DataDesignerEarlyShutdownError, + DataDesignerGenerationError, + DataDesignerProfilingError, + DatasetCreationResults, +) +from data_designer.plugins import Plugin +``` + +The package must not depend on: + +- Any `data_designer.engine` import. +- `data_designer.config.config_builder.BuilderConfig`. +- `data_designer.config.mcp.MCPProviderT`. +- `DatasetCreationResults.artifact_storage` or `task_traces`. +- The plugin registry implementation. +- Typed errors documented by `DataDesigner.check_models()` outside the public interface error module. + +`DataDesigner` may use engine modules internally. The restriction applies to imports and types crossing the package +boundary. + +## Serialized builder boundary + +Authored input contains exactly one of: + +- A path to a complete JSON or YAML builder config. +- The same builder config as an inline mapping. + +The payload is the format written by `DataDesignerConfigBuilder.write_config()`. The public loader is: + +```python +builder = DataDesignerConfigBuilder.from_config(builder_payload_or_path) +``` + +`BuilderConfig` is an implementation type used by the serializer. It is not part of the cross-package API. +Invocation controls such as record count, dataset name, resume, runtime settings, model endpoints, and output format +remain outside the builder payload. + +## Validation ownership + +### Submission environment + +The submission environment validates only information that does not require plugin-specific config classes: + +- Outer execution schema and version. +- Builder source location, format, digest, and top-level mapping shape. +- Unique aliases declared in the raw `data_designer.model_configs` list. +- Deployment coverage for those declared aliases. +- `RunConfig`, output, image, dependency, and resource fields owned by the execution package. + +It must not instantiate the complete builder when referenced plugins are unavailable. Plugin columns may declare +additional model references, so referenced-alias validation is deferred with full builder validation. + +### Client environment + +The package overlay must be installed before the client process imports Data Designer config unions. Plugin discovery +is process-global and happens on first use, so installing a package after config import is not a supported refresh +path. + +Client preflight runs in a fresh process and performs these steps: + +1. Enumerate the expected `data_designer.plugins` entry points with `importlib.metadata`. +2. Load each entry point and verify that it returns a public `Plugin` object. +3. Resolve `Plugin.config_cls` and `Plugin.impl_cls` so import failures stop preflight. +4. Load the builder with `DataDesignerConfigBuilder.from_config()`. +5. Materialize model bindings, per-model concurrency, seed input bindings, and MCP providers. +6. Verify referenced model aliases through each column config's public `get_model_aliases()` method and each profiler + config's public `model_alias` field when present. +7. Construct `DataDesigner`, apply `RunConfig`, and call `DataDesigner.validate(builder)`. + +The explicit entry-point load is required because normal plugin discovery logs and skips a failing entry point. +`DataDesigner.validate()` validates config structure and seed-dependent compilation without contacting model endpoints. +It does not verify that every referenced column or profiler alias has a matching model config, so client preflight owns +that check. Endpoint readiness belongs after services are available; generation performs its normal readiness check. + +## Model binding + +`ModelConfig.alias` is the workload identity. The original `model` string is not an alias and must not be used to +match deployments. + +For every declared alias, the client creates a stable `ModelProvider` for the resolved logical endpoint and replaces +the matching `ModelConfig` while preserving unrelated inference parameters: + +```python +originals = list(builder.model_configs) +for model_config in originals: + builder.delete_model_config(model_config.alias) + +for model_config in originals: + binding = bindings_by_alias[model_config.alias] + inference_parameters = model_config.inference_parameters.model_copy( + update={"max_parallel_requests": binding.max_parallel_requests} + ) + builder.add_model_config( + model_config.model_copy( + update={ + "model": binding.served_model, + "provider": binding.provider_name, + "inference_parameters": inference_parameters, + } + ) + ) +``` + +Deleting all originals before adding replacements preserves alias uniqueness and original ordering. Full preflight +rejects missing or duplicate declared aliases, missing referenced aliases, missing providers, and deployment aliases +that do not match `ModelConfig.alias`. + +Per-model concurrency is `ModelConfig.inference_parameters.max_parallel_requests`. It is not a `RunConfig` field. + +## Runtime configuration + +The execution package merges its compatibility defaults with the raw authored `run_config` mapping, then validates +the effective mapping once: + +```python +effective_run_config = RunConfig.model_validate(compatibility_defaults | authored_run_config) +data_designer.set_run_config(effective_run_config) +``` + +Merging must occur before constructing the authored `RunConfig`. Model validators normalize related fields, so a +validated model's field-set metadata is not the authoritative record of which keys the user wrote. + +The effective default for every public field is: + +| Field | Data Designer default | Execution-package default | +| --- | ---: | ---: | +| `disable_early_shutdown` | `False` | `True` when no early-shutdown control is authored | +| `shutdown_error_rate` | `0.5` | `1.0` after disabled-shutdown normalization | +| `shutdown_error_window` | `10` | `10` | +| `buffer_size` | `1000` | `16384` | +| `max_concurrent_row_groups` | `3` | `3` | +| `max_in_flight_tasks` | `1024` | `1024` | +| `non_inference_max_parallel_workers` | `4` | `4` | +| `max_conversation_restarts` | `5` | `0` | +| `max_conversation_correction_steps` | `0` | `0` | +| `async_trace` | `False` | `False` | +| `write_scheduler_events` | `False` | `False` | +| `display_tui` | `False` | `False` | +| `progress_interval` | `5.0` | `5.0` | +| `otel_metrics_port` | `9464` | `None` | +| `preserve_dropped_columns` | `True` | `True` | +| `jinja_rendering_engine` | `secure` | `secure` | +| `request_admission` | `None` | `None` | + +Every authored key overrides the corresponding execution-package default. Early-shutdown controls form one related +group: if the user authors `shutdown_error_rate` or `shutdown_error_window` without `disable_early_shutdown`, the +package retains Data Designer's enabled-shutdown default instead of injecting `disable_early_shutdown=True` and +discarding the authored threshold. + +`non_inference_max_parallel_workers` is a public field but is not currently consumed by generation. The execution +package must not claim CPU-derived worker control until Data Designer implements that behavior. + +## Invocation + +The client materializes provider and MCP secrets immediately before constructing public config objects. Secret +values do not enter authored or persisted config. + +```python +data_designer = DataDesigner( + artifact_path=dataset_workspace, + model_providers=model_providers, + managed_assets_path=managed_assets_path, + mcp_providers=mcp_providers, + auto_configure_logging=False, +) +data_designer.set_run_config(effective_run_config) +data_designer.validate(builder) + +results = data_designer.create( + builder, + num_records=requested_num_records, + dataset_name=dataset_name, + resume=resume_mode, +) +``` + +`DataDesigner.acreate()` is the equivalent non-blocking entry point. It delegates generation to a worker thread and +returns the same `DatasetCreationResults` type. + +Remote MCP connections use `MCPProvider`; local subprocess connections use `LocalStdioMCPProvider`. The execution +package owns any secret-reference schema and passes only resolved strings to Data Designer. + +## Results and errors + +A normal return means generation and profiling completed without a public exception. It does not guarantee that the +dataset contains exactly the requested number of records. + +`DatasetCreationResults` exposes: + +| Field or method | Meaning | +| --- | --- | +| `dataset_path` | Resolved dataset directory, including collision or resume resolution. | +| `requested_num_records` | Target passed to `create()`, or the persisted target for a reconstructed workflow result. | +| `actual_num_records` | Current total rows in the final dataset, including rows from an earlier resumed invocation. | +| `is_partial` | `True` when actual records are fewer than requested records. | +| `early_shutdown` | Whether the current invocation stopped through the early-shutdown gate, or `None` when no generation invocation produced the result object. | +| `requested_resume_mode` | Resume mode passed to `create()`, or `None` when no generation invocation produced the result object. | +| `effective_resume_mode` | `always` when the invocation resumed, `never` when it started fresh, or `None` when no generation invocation produced the result object. | +| `count_records()` | Metadata-only row count, equivalent to `actual_num_records`. | +| `export(path, format=...)` | Stream the result to one JSONL, CSV, or Parquet file. | + +The caller derives its own exact-count policy from `actual_num_records == requested_num_records`. A partial result may +be caused by dropped rows or early shutdown; `early_shutdown` distinguishes those cases. + +Public failure behavior is: + +| Condition | Public behavior | +| --- | --- | +| Missing or unreadable local builder file | `InvalidFilePathError` from `data_designer.config`. | +| Malformed local or inline builder data | `InvalidFileFormatError` from `data_designer.config`. | +| Invalid serialized builder shape or remote source | `pydantic.ValidationError` or `ValueError` from `DataDesignerConfigBuilder.from_config()`. | +| Invalid compiled config | `InvalidConfigError` from `data_designer.config`. | +| Generation failure | `DataDesignerGenerationError`. | +| Profiling failure after generation | `DataDesignerProfilingError`. No successful result is returned. | +| Early shutdown with zero records | `DataDesignerEarlyShutdownError`, a `DataDesignerGenerationError` subclass. | +| Early shutdown with some records | A partial `DatasetCreationResults` with `early_shutdown=True`. | +| Invalid export format or incompatible Parquet schemas | `InvalidFileFormatError` from `data_designer.config`. | +| Missing or unreadable dataset or processor artifacts | `ArtifactStorageError` from `data_designer.interface`. | + +The caller must not inspect engine storage or task-trace types to classify an outcome. + +Failure exceptions do not carry a `DatasetCreationResults` object. The resolved dataset path, actual record count, +early-shutdown state, and effective resume mode may therefore be unavailable after a failed invocation. A semantic +failure record owned by an embedding package must make those facts optional rather than inspect engine storage. + +## Output, resume, seed, and telemetry + +- `artifact_path`, `dataset_name`, and `ResumeMode` define the resumable workspace. +- `dataset_path` is the resolved public location. This is important when `ResumeMode.IF_POSSIBLE` starts fresh or a + non-resumable name collides. +- Export format is explicit. Data Designer does not publish a default export-format constant. +- Seed datasets are applied with `with_seed_dataset()` and can be partitioned with public `PartitionBlock`. +- `PartitionBlock` does not define a deterministic generation random seed. No public generation-seed API exists. +- OpenTelemetry metrics are configured with `RunConfig.otel_metrics_port`; `None` disables metrics for the invocation. +- Embedded callers should use `auto_configure_logging=False` when they own process logging. + +`DatasetCreationResults` does not expose per-invocation usage or phase timing. The public OpenTelemetry endpoint +provides create duration, generated/dropped record counters, and request-duration histograms, but not token totals or +separate generation and profiling durations. Callers must treat unavailable values as optional unless a future public +result summary provides them; logs and engine usage types are not a supported substitute. + +## Sharded execution limits + +Public builder methods expose column, processor, and profiler configs through `get_column_configs()`, +`get_processor_configs()`, and `get_profilers()`. Data Designer does not publish capability metadata that identifies +whether a processor, profiler, media output, or plugin is safe to merge across independently generated partitions. +An embedding package must use a conservative policy and reject multi-partition execution when those semantics are +present or unknown. Ordered seed input can use `PartitionBlock`; `SamplingStrategy.SHUFFLE` has no deterministic +partition-then-shuffle contract. + +## Callable and plugin limits + +Entry-point plugins with serializable config models are supported after installation in the client environment. + +`CustomColumnConfig.generator_function` and `LocalCallableValidatorParams.validation_function` are not portable +serialized references. Their serializers write only a function name while validation requires an in-memory callable. +Installing a package later does not resolve that string. A workload requiring these callables needs a separately +approved installed builder factory or a future qualified-callable-reference feature. + +## Import baseline + +The base import budget remains the existing `make perf-import CLEAN=1` measurement: + +```text +import data_designer.config as dd +from data_designer.interface import DataDesigner +``` + +The average of one cold and four warm runs must remain below three seconds. Optional command discovery may read +distribution metadata, but must not import the optional package until its command is selected. CLI extension discovery +is tracked separately by #853. + +## Contract tests + +The owning packages must cover: + +- Public root imports and absence of external `data_designer.engine` imports. +- Builder file and inline mapping loading after plugin installation. +- Plugin entry-point load failure before model services start. +- Model binding by alias, order preservation, per-model concurrency, profiler aliases, and plugin secondary aliases. +- Compatibility-default merging from raw authored keys. +- Complete, partial, partial early-shutdown, zero-record early-shutdown, generation-error, and profiling-error outcomes. +- Failure records that allow result-only path, count, early-shutdown, and resume facts to be unavailable. +- Resolved dataset paths for fresh, colliding, resumed, and incompatible `if_possible` invocations. +- Explicit JSONL, CSV, and Parquet exports. +- Remote and stdio MCP providers without persisted secret values. +- Conservative multi-partition rejection for unsafe or unknown processor, profiler, media, and plugin semantics. +- The existing import-performance threshold. diff --git a/plans/870/model-colocation.md b/plans/870/model-colocation.md new file mode 100644 index 000000000..f162e33dc --- /dev/null +++ b/plans/870/model-colocation.md @@ -0,0 +1,25 @@ +# Model co-location + +## Context + +Different model deployments currently receive disjoint nodes and the full profile GPU budget. + +## Goal + +Allow explicit single-node groups with disjoint per-deployment GPU quotas in one managed run. +Deliver as one PR against `feat/slurm-execution`. + +## Design + +Use `resources.shared_node` and `resources.gpus_per_node`. One group must cover every deployment in the run. Members require one node, single-node replicas, GRES allocation, the same resolved serving image, and matching server environment and queue settings. Tensor parallelism divides each quota, and quotas fit the profile allocation budget. The compiler assigns contiguous GPU slices and unique ports. One coordinated worker launches and supervises all model processes using Slurm's step-local CUDA visibility. Client/proxy steps remain GPU-free. Default fields are omitted from serialization to preserve legacy bytes. + +## Tasks + +- [x] Add authored/preparation validation and deterministic placement/port/GPU contracts, including compatibility tests. +- [x] Compose shared node workers and verify readiness, CUDA mapping, failures, cleanup, and retry resource preservation. +- [x] Document two distinct models and run the package/docs gates and one bounded end-to-end smoke. +- [x] Obtain independent Claude review, address verified findings, and review the final branch. + +## Open questions + +Runtime semantics and shared-worker restrictions require maintainer review before merge. Model memory fit remains the caller's responsibility. diff --git a/plans/870/slurm-early-security-review.md b/plans/870/slurm-early-security-review.md new file mode 100644 index 000000000..c135f1b38 --- /dev/null +++ b/plans/870/slurm-early-security-review.md @@ -0,0 +1,60 @@ +# Slurm early security and provenance review + +This document records the dependency-ready portion of #870. It reviews the public Slurm implementation through +the one-node runtime merge and the hardening changes developed with this review. It is not final release acceptance: +the sealed commit, complete wheel set, runtime checksum, real-cluster scenarios, sanitized profile rerun, and all +dependent implementation and documentation must still be frozen and validated together. + +## Threat model boundaries + +| Boundary | Threats reviewed | Existing or added controls | Evidence | +| --- | --- | --- | --- | +| Slurm submission | Shell injection, option confusion, inherited secrets, unbounded caller-visible diagnostics | Commands use argument vectors without a shell; batch options reject controls; script paths reject option-like names; the launcher forwards only an explicit environment; diagnostic text is normalized, redacted, and limited to 512 characters | `tests/launcher/test_client.py`, `tests/launcher/test_renderer.py`, `tests/launcher/test_runner.py` | +| Batch entrypoint | Directive splitting, shell expansion, plan or runtime substitution, unsafe task identity | Directive names and values are validated; shell values are escaped; the script fixes `PATH`, checks both SHA-256 identities before extraction, validates the array-task ID, and creates a private attempt-local runtime directory | `tests/launcher/test_renderer.py`, `tests/slurm_test_fakes/test_rendered_scripts.py` | +| Runtime commands and environment | Shell fragments, ambient environment leakage, persisted plaintext credentials | Runtime steps are immutable argument vectors; `Popen` uses `shell=False`; only package-owned scheduler variables and explicitly resolved bindings are forwarded; secret-shaped values require environment references and are not persisted | `tests/contracts/test_config_records.py`, `tests/runtime/test_steps.py`, `tests/runtime/test_supervisor.py` | +| Host and container paths | Parent traversal, ambiguous paths, mount escape, read-only mount writes | Persisted paths are normalized absolute POSIX paths below `/`; container translation selects the most-specific resolved mount and separately enforces write access; state and logs use descriptor-bound, no-follow operations | `tests/runtime/test_paths.py`, `tests/runtime/test_preflight.py`, `tests/state/test_store.py` | +| Images and runtime archives | Credential-bearing image references, archive traversal, replacement races, altered runtime source | OCI sources reject credentials and ambiguous schemes; image/state publication is restrictive and atomic; the runtime archive is package-built with fixed member names and metadata, content-addressed, and verified before extraction | `tests/images/test_lifecycle.py`, `tests/images/test_registry_store.py`, `tests/runtime/test_bundle.py` | +| Logs and public evidence | Secret or site-specific data copied into public artifacts | Runtime logs are private `0600` files below private execution directories. The public-artifact audit reports only a display path and rule name, scans explicit log paths without echoing matches, and rejects high-confidence credentials and environment-specific infrastructure values | `tests/runtime/test_supervisor.py`, `tests/test_public_artifacts.py` | +| Cleanup and signals | Orphaned process groups, repeated cleanup, partial publication | Runtime children start in owned sessions; cleanup is idempotent, terminates in reverse order, escalates after a bounded grace period, and surfaces normalized failure state | `tests/runtime/test_supervisor.py`, `tests/runtime/test_controller.py` | + +## Findings resolved in this slice + +- The Slurm wheel declared Apache-2.0 but did not carry the license text. The package now includes a canonical copy of + the repository license, verifies the copies are byte-identical, and fails the wheel audit unless the expanded wheel + contains Apache License 2.0 text. +- Slurm command failures bounded and control-normalized scheduler stderr but did not explicitly redact recognizable + credentials. A shared redaction helper now covers secret-shaped assignments and options, authorization headers, URL + user information, and high-confidence provider token formats before diagnostic truncation. +- Public-artifact checks were limited to individual golden tests. The new scanner covers deployable source, package + metadata, public fixtures, examples/documentation, release scripts, explicit logs, ZIP/wheel members, and tar members. + Archive traversal, links, excessive member counts, excessive expanded content, missing packaged-source SPDX headers, + and missing wheel license text fail closed. + +## Provenance and dependency review + +- All reviewed package Python and shell resources carry NVIDIA Apache-2.0 SPDX headers. Package history identifies the + resources as repository contributions; no copied or adapted third-party source was identified in this scope. +- Direct runtime dependencies remain `data-designer`, `packaging`, `pydantic`, and `pyyaml`. The dependency inventory + reports only `click` and `typer` as transitive imports guaranteed by the exact-version `data-designer` dependency; + there are no unresolved imported modules. +- The scanner deliberately permits only generic test representations that its high-confidence rules do not classify: + `example.test` hosts, loopback addresses, and `/workspace` paths. Its default scope includes the complete Python test + tree and masks only exact, path-scoped synthetic sentinel values before applying the same high-confidence rules. + +## Commands for this review slice + +```bash +python scripts/audit_slurm_public_artifacts.py +python scripts/audit_slurm_public_artifacts.py path/to/wheel.whl path/to/sanitized.log +make test-slurm +make test-slurm-wheel-install +make check-slurm +make check-dependency-licenses +``` + +## Remaining final acceptance + +After all #850 implementation and documentation dependencies are merged, #870 must select one public commit and build +the complete wheel set once. The source, wheels, runtime archive, sanitized scenario evidence, and sanitized profile +rerun must all be tied to those exact digests. Any code or wheel change after sealing invalidates the affected evidence. +Environment-specific profiles and raw scheduler or allocation logs remain outside public artifacts and issue comments. diff --git a/plans/874/slurm-documentation-coverage.md b/plans/874/slurm-documentation-coverage.md new file mode 100644 index 000000000..29356b4cd --- /dev/null +++ b/plans/874/slurm-documentation-coverage.md @@ -0,0 +1,36 @@ +# Slurm documentation coverage map + +This map records the behavioral migration from legacy Slurm guides to the public Fern Slurm guide. It tracks user-visible concepts rather than source wording or environment-specific examples. + +| Legacy area | Behavior retained in Fern | Fern destination | Disposition | +| --- | --- | --- | --- | +| Installation | Python environment, optional Slurm extra, CLI discovery | `overview.mdx`, `getting-started.mdx` | Retained and updated for the same-version optional wheel | +| Environment variables | Profile-file override, secret references, runtime environment boundaries | `profiles.mdx`, `dependencies-and-tools.mdx` | Relocated into the owning workflows | +| Cluster prerequisites | Submit commands, shared storage, Enroot, scheduler and GPU facts | `overview.mdx` | Expanded with workstation, login-host, and compute-node personas | +| Quick start | Install, profile, image import, builder, run config, dry-run, submit, status, cancel, output | `getting-started.mdx` | Replaced with a clean 13-step public flow | +| Profile configuration | Account, partitions, GPU discovery, GPU request mode, mounts | `profiles.mdx` | Expanded to a named multi-cluster catalog | +| Image management | OCI import, existing SQSH inspection, list, inspect, replace, remove | `images.mdx` | Expanded with digest and role verification rules | +| Client and serving images | Separate image roles and compatibility requirements | `images.mdx` | Retained and made explicit | +| Run configuration | Builder input, invocation, client, deployment, submission, output | `run-configuration.mdx` | Replaced with the shipped strict public schema | +| Server options | vLLM timeouts, readiness, backpressure, environment, safe arguments | `run-configuration.mdx` | Retained through supported public fields | +| Multi-node serving | Nodes, tensor parallelism, nodes per replica | `run-configuration.mdx` | Retained and expressed as deployment topology | +| Multiple models | Per-alias deployment, concurrency, and image selection | `run-configuration.mdx` | Retained and expanded to separate serving images | +| Job arrays | Shard count, concurrency throttle, persisted shard identity | `run-configuration.mdx` | Retained with public array constraints | +| Execute and observe | Dry-run, submit, JSON results, scheduler reconciliation | `operations.mdx` | Expanded with durable state behavior | +| Cancellation | Managed job ownership and idempotent cancellation | `operations.mdx` | Retained with run-ID ownership checks | +| Logs and results | Durable logs, attempt records, winners, default output | `operations.mdx` | Expanded with the persistent storage boundary | +| Runtime scratch | Runtime, dependency, cache, and Enroot placement and cleanup | `overview.mdx`, `operations.mdx`, `troubleshooting.mdx` | Intentional change to allocation-local scratch with stable container mount | +| Retry | Sparse task selection, resume policy, confirmation, dry-run | `retry-and-collection.mdx` | Expanded with immutable attempt history | +| Merge and collection | Winner-driven collection to an explicit destination | `retry-and-collection.mdx` | Retained as the public `merge` command | +| Dependencies | Pure-wheel overlay, immutable locks, client-image compatibility | `dependencies-and-tools.mdx` | Expanded with strict accepted sources | +| Plugins | Client image or dependency-overlay installation | `dependencies-and-tools.mdx` | Relocated and linked to public plugin contracts | +| Private packages and secrets | Built client images, immutable wheel or lock inputs, and external secret references | `dependencies-and-tools.mdx` | Replaced with the shipped public-index resolver and typed environment references | +| MCP providers | Remote and local stdio providers in client allocations | `dependencies-and-tools.mdx` | Added to match the shipped public surface | +| Code sandbox | Auxiliary isolated execution service | `dependencies-and-tools.mdx` | Intentionally unsupported in Slurm v1; alternatives documented | +| Benchmarks | Case expansion, child runs, status, point-in-time analysis | `benchmarks.mdx` | Retained and expanded with durable benchmark state | +| FAQ and recovery | Installation, profile, image, scheduler, readiness, retry, output, scratch | `troubleshooting.mdx` | Reorganized by observed failure | +| Command summary | Commands, common options, JSON output, exit codes | `cli-reference.mdx` | Replaced with a concise shipped-CLI reference | + +## Acceptance boundary + +The Fern guide and checked examples are validated in repository CI and against locally built wheels. Formal sealed-artifact acceptance, including the complete real-cluster scenario matrix and sanitized evidence bundle, remains issue #870 and must run after this documentation and release-integration slice merges and any required fixes land. diff --git a/pyproject.toml b/pyproject.toml index d9db9401a..435fa0f09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ build-backend = "hatchling.build" members = [ "packages/data-designer-config", "packages/data-designer-engine", + "packages/data-designer-slurm", "packages/data-designer", ] diff --git a/scripts/audit_package_dependencies.py b/scripts/audit_package_dependencies.py index f044646d6..e4eac95b7 100644 --- a/scripts/audit_package_dependencies.py +++ b/scripts/audit_package_dependencies.py @@ -66,22 +66,29 @@ def requirement_dependencies( return dict(dependencies) -def declared_dependencies(pyproject_path: Path) -> dict[str, set[str]]: +def declared_dependencies(pyproject_path: Path, selected_extras: set[str] | None = None) -> dict[str, set[str]]: with pyproject_path.open("rb") as file: config = tomllib.load(file) - project_dependencies = config.get("project", {}).get("dependencies", []) - dynamic_dependencies = ( - config.get("tool", {}) - .get("hatch", {}) - .get("metadata", {}) - .get("hooks", {}) - .get("uv-dynamic-versioning", {}) - .get("dependencies", []) + project = config.get("project", {}) + metadata_hook = ( + config.get("tool", {}).get("hatch", {}).get("metadata", {}).get("hooks", {}).get("uv-dynamic-versioning", {}) ) + project_dependencies = project.get("dependencies", []) + dynamic_dependencies = metadata_hook.get("dependencies", []) dependencies = requirement_dependencies(project_dependencies) for name, extras in requirement_dependencies(dynamic_dependencies, allow_version_template=True).items(): dependencies.setdefault(name, set()).update(extras) + + for extra in selected_extras or set(): + optional_dependencies = project.get("optional-dependencies", {}).get(extra, []) + dynamic_optional_dependencies = metadata_hook.get("optional-dependencies", {}).get(extra, []) + for name, extras in requirement_dependencies(optional_dependencies).items(): + dependencies.setdefault(name, set()).update(extras) + for name, extras in requirement_dependencies( + dynamic_optional_dependencies, allow_version_template=True + ).items(): + dependencies.setdefault(name, set()).update(extras) return dependencies @@ -134,6 +141,7 @@ def audit_repository( repository_root: Path, module_distributions: dict[str, list[str]] | None = None, requirement_map: dict[str, list[str]] | None = None, + selected_extras_by_project: dict[str, set[str]] | None = None, ) -> dict[str, Any]: repository_root = repository_root.resolve() package_dirs = sorted(path.parent for path in repository_root.glob("packages/*/pyproject.toml")) @@ -144,7 +152,10 @@ def audit_repository( project_name = normalize_name(tomllib.load(file)["project"]["name"]) projects[project_name] = { "path": str(package_dir.relative_to(repository_root)), - "declarations": declared_dependencies(package_dir / "pyproject.toml"), + "declarations": declared_dependencies( + package_dir / "pyproject.toml", + (selected_extras_by_project or {}).get(project_name), + ), "imports": imported_modules(package_dir / "src", repository_root), } @@ -159,16 +170,20 @@ def audit_repository( requirement_cache: dict[tuple[str, frozenset[str]], dict[str, set[str]]] = {} def requirements_for(distribution_name: str, selected_extras: set[str]) -> dict[str, set[str]]: - if distribution_name in projects: - return projects[distribution_name]["declarations"] - if requirement_map is not None: - return requirement_dependencies(requirement_map.get(distribution_name, []), selected_extras) cache_key = (distribution_name, frozenset(selected_extras)) if cache_key not in requirement_cache: - requirement_cache[cache_key] = installed_requirements(distribution_name, selected_extras) + if distribution_name in projects: + package_path = repository_root / projects[distribution_name]["path"] / "pyproject.toml" + requirement_cache[cache_key] = declared_dependencies(package_path, selected_extras) + elif requirement_map is not None: + requirement_cache[cache_key] = requirement_dependencies( + requirement_map.get(distribution_name, []), selected_extras + ) + else: + requirement_cache[cache_key] = installed_requirements(distribution_name, selected_extras) return requirement_cache[cache_key] - def dependency_closure(distribution_name: str, selected_extras: set[str]) -> set[str]: + def dependency_closure(distribution_name: str, selected_extras: set[str], project_name: str) -> set[str]: closure = set() pending = [(distribution_name, selected_extras)] visited: set[tuple[str, frozenset[str]]] = set() @@ -178,6 +193,8 @@ def dependency_closure(distribution_name: str, selected_extras: set[str]) -> set if current_key in visited: continue visited.add(current_key) + if current == project_name: + continue for dependency, dependency_extras in requirements_for(current, extras).items(): closure.add(dependency) pending.append((dependency, dependency_extras)) @@ -188,7 +205,8 @@ def dependency_closure(distribution_name: str, selected_extras: set[str]) -> set declarations = project["declarations"] declared = set(declarations) dependency_closures = { - dependency: dependency_closure(dependency, extras) for dependency, extras in declarations.items() + dependency: dependency_closure(dependency, extras, project_name) + for dependency, extras in declarations.items() } resolved_imports: dict[str, dict[str, set[str]]] = defaultdict(lambda: {"modules": set(), "files": set()}) unresolved = [] @@ -243,9 +261,17 @@ def main() -> None: parser = argparse.ArgumentParser(description="Inventory package import/dependency gaps") parser.add_argument("--root", type=Path, default=Path.cwd()) parser.add_argument("--output", type=Path) + parser.add_argument("--extra", action="append", default=[], metavar="PACKAGE:EXTRA") args = parser.parse_args() - result = audit_repository(args.root) + selected_extras_by_project: dict[str, set[str]] = defaultdict(set) + for value in args.extra: + package, separator, extra = value.partition(":") + if not separator or not package or not extra: + parser.error("--extra must use PACKAGE:EXTRA format") + selected_extras_by_project[normalize_name(package)].add(extra) + + result = audit_repository(args.root, selected_extras_by_project=dict(selected_extras_by_project)) payload = json.dumps(result, indent=2) + "\n" if args.output: args.output.write_text(payload) diff --git a/scripts/audit_slurm_public_artifacts.py b/scripts/audit_slurm_public_artifacts.py new file mode 100644 index 000000000..41bdcbc85 --- /dev/null +++ b/scripts/audit_slurm_public_artifacts.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Scan public Slurm artifacts without rendering matched sensitive content. + +The default scope covers deployable source, test source and fixtures, package +metadata, public documentation and examples, and release scripts. A small path-scoped allowlist masks only exact +synthetic test sentinels. Generic example.test hosts, loopback addresses, and +/workspace paths are the only implicit test-data allowances. +""" + +from __future__ import annotations + +import argparse +import os +import re +import stat +import sys +import tarfile +import zipfile +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import BinaryIO, Iterable + +MAXIMUM_MEMBER_SIZE = 16 * 1024 * 1024 +MAXIMUM_ARCHIVE_CONTENT_SIZE = 128 * 1024 * 1024 +MAXIMUM_ARCHIVE_MEMBERS = 10_000 +REPOSITORY_ROOT = Path(__file__).parents[1] +CANONICAL_PACKAGE_LICENSE = REPOSITORY_ROOT / "packages" / "data-designer-slurm" / "LICENSE" +DEFAULT_ARTIFACTS = ( + "packages/data-designer-slurm/src", + "packages/data-designer-slurm/README.md", + "packages/data-designer-slurm/pyproject.toml", + "packages/data-designer-slurm/LICENSE", + "packages/data-designer-slurm/examples", + "packages/data-designer-slurm/tests", + "fern/versions/latest/pages/slurm", + "fern/versions/latest.yml", + "plans/850/data-designer-contract.md", + "plans/870/slurm-early-security-review.md", + "plans/874/slurm-documentation-coverage.md", + "scripts/publish.sh", + "scripts/test_slurm_package_install.py", +) +_LICENSED_SOURCE_SUFFIXES = frozenset({".py", ".rc", ".sh"}) +_SPDX_COPYRIGHT = "SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved." +_SPDX_LICENSE = "SPDX-License-Identifier: Apache-2.0" +_SLURM_DIST_INFO_PATTERN = re.compile(r"^data_designer_slurm-[^/]+\.dist-info$") +_SENTINEL_VALUE_CHARACTERS = r"A-Za-z0-9+/_.:@-" +_TEST_SOURCE_SENTINELS = { + "packages/data-designer-slurm/tests/config/test_loading_builder.py": ( + '"clusters": {secret: profile_catalog.clusters["primary"]}', + "super-secret-token", + "super-secret-token: x", + "sk_live_ABC123XYZ", + ), + "packages/data-designer-slurm/tests/contracts/test_config_records.py": ("plaintext-secret",), + "packages/data-designer-slurm/tests/launcher/test_client.py": ( + "Authorization: Bearer bearer-secret;suffix status=failed", + "Authorization: Bearer bearer-secret", + '{"Authorization": "Bearer serialized-secret-value", "status": "failed"}', + ), + "packages/data-designer-slurm/tests/planning/test_compiler.py": ("super-secret-token",), + "packages/data-designer-slurm/tests/test_public_artifacts.py": ( + "/home/specific-user/run", + "10.23.45.67", + "data_designer/slurm/10.23.45.67.py", + "service.internal.nvidia.com", + "super-secret-token", + ), +} + + +@dataclass(frozen=True, slots=True) +class AuditRule: + """One high-confidence public-artifact rejection rule.""" + + name: str + pattern: re.Pattern[str] + + +@dataclass(frozen=True, slots=True) +class AuditFinding: + """One path-scoped finding that intentionally omits matched content.""" + + location: str + rule: str + + +@dataclass(frozen=True, slots=True) +class _ZipMemberAudit: + findings: tuple[AuditFinding, ...] + dist_info_root: str | None = None + is_license: bool = False + license_matches: bool = False + is_metadata: bool = False + metadata_declares_license: bool = False + + +_CONTENT_RULES = ( + AuditRule("private key material", re.compile(r"-----BEGIN (?:[A-Z0-9 ]+ )?PRIVATE KEY-----")), + AuditRule("AWS access key", re.compile(r"\bAKIA[0-9A-Z]{16}\b")), + AuditRule("GitHub token", re.compile(r"\b(?:github_pat_|gh[pousr]_)[A-Za-z0-9_]{20,}\b")), + AuditRule("NGC API key", re.compile(r"\bnvapi-[A-Za-z0-9_-]{20,}\b")), + AuditRule("OpenAI API key", re.compile(r"\bsk-(?:proj-)?[A-Za-z0-9_-]{20,}\b")), + AuditRule( + "plaintext secret assignment", + re.compile( + r"(?i)\b(?:api[_-]?key|access[_-]?token|auth[_-]?token|password|secret|credential)" + r"\b\s*[:=]\s*[\"']?(?!<|\$|\{|\[)[A-Za-z0-9+/_.:@-]{16,}" + ), + ), + AuditRule( + "authorization credential", + re.compile(r"(?i)(?P[\"']?)\bauthorization(?P=quote)\s*:\s*[\"']?(?:basic|bearer)\s+\S{16,}"), + ), + AuditRule( + "internal NVIDIA hostname", + re.compile(r"(?i)\b(?:(?:[a-z0-9-]+\.)*(?:corp|internal)|gitlab-master|urm)\.nvidia\.com\b"), + ), + AuditRule("user-specific absolute path", re.compile(r"/(?:Users|home|users)/[A-Za-z0-9._-]+(?:/|\b)")), + AuditRule("site-specific shared path", re.compile(r"/(?:lustre|gpfs|fsx|net)/[^\s\"']+", re.IGNORECASE)), + AuditRule( + "private infrastructure address", + re.compile(r"\b(?:10(?:\.\d{1,3}){3}|172\.(?:1[6-9]|2\d|3[01])(?:\.\d{1,3}){2}|192\.168(?:\.\d{1,3}){2})\b"), + ), +) + + +def audit_public_artifacts(paths: Iterable[Path]) -> tuple[AuditFinding, ...]: + """Return sanitized findings from files, directories, wheels, and tar archives.""" + findings: list[AuditFinding] = [] + for path in sorted((Path(os.path.abspath(candidate)) for candidate in paths), key=str): + if path.is_symlink(): + findings.append(AuditFinding(_display_path(path), "symbolic-link artifact requires explicit review")) + elif not path.exists(): + findings.append(AuditFinding(_display_path(path), "artifact does not exist")) + elif path.is_dir(): + for child in sorted(path.rglob("*")): + if child.is_symlink(): + findings.append( + AuditFinding(_display_path(child), "symbolic-link artifact requires explicit review") + ) + elif "__pycache__" in child.parts: + continue + elif child.is_dir(): + continue + else: + findings.extend(_audit_file(child)) + else: + findings.extend(_audit_file(path)) + return tuple(findings) + + +def _audit_file(path: Path) -> list[AuditFinding]: + try: + mode = path.lstat().st_mode + except OSError: + return [AuditFinding(_display_path(path), "artifact cannot be inspected")] + if not stat.S_ISREG(mode): + return [AuditFinding(_display_path(path), "artifact is not a regular file")] + if _is_zip_archive(path): + return _audit_zip(path) + if _is_tar_archive(path): + return _audit_tar(path) + try: + with path.open("rb") as stream: + content = _read_bounded(stream, expected_size=path.stat().st_size) + except OSError: + return [AuditFinding(_display_path(path), "artifact cannot be read")] + except ValueError as error: + return [AuditFinding(_display_path(path), str(error))] + location = _display_path(path) + return _audit_content( + location, + location, + content, + allowed_sentinels=_TEST_SOURCE_SENTINELS.get(location, ()), + ) + + +def _audit_zip(path: Path) -> list[AuditFinding]: + findings: list[AuditFinding] = [] + member_audits: list[_ZipMemberAudit] = [] + try: + canonical_license = CANONICAL_PACKAGE_LICENSE.read_bytes() + except OSError: + return [AuditFinding(_display_path(path), "canonical package license cannot be read")] + try: + with zipfile.ZipFile(path) as archive: + members = archive.infolist() + limit_finding = _get_archive_limit_finding( + path, + member_count=len(members), + content_size=sum(max(member.file_size, 0) for member in members), + ) + if limit_finding is not None: + return [limit_finding] + for index, member in enumerate(sorted(members, key=lambda item: item.filename), start=1): + location = f"{_display_path(path)}!member-{index}" + member_audit = _audit_zip_member( + archive, + member, + location, + canonical_license, + ) + findings.extend(member_audit.findings) + member_audits.append(member_audit) + except (OSError, zipfile.BadZipFile): + return [AuditFinding(_display_path(path), "artifact is not a readable ZIP archive")] + if path.suffix.casefold() == ".whl": + findings.extend(_audit_wheel_distribution(path, member_audits)) + return findings + + +def _audit_tar(path: Path) -> list[AuditFinding]: + findings: list[AuditFinding] = [] + content_size = 0 + try: + with tarfile.open(path, mode="r:*") as archive: + for member_count, member in enumerate(archive, start=1): + content_size += max(member.size, 0) + limit_finding = _get_archive_limit_finding( + path, + member_count=member_count, + content_size=content_size, + ) + if limit_finding is not None: + return [limit_finding] + location = f"{_display_path(path)}!member-{member_count}" + findings.extend(_audit_tar_member(archive, member, location)) + except (OSError, tarfile.TarError): + return [AuditFinding(_display_path(path), "artifact is not a readable tar archive")] + return findings + + +def _audit_zip_member( + archive: zipfile.ZipFile, + member: zipfile.ZipInfo, + location: str, + canonical_license: bytes, +) -> _ZipMemberAudit: + dist_info_root = _get_dist_info_root(member.filename) + if _is_unsafe_archive_name(member.filename): + return _ZipMemberAudit((AuditFinding(location, "archive member path is unsafe"),)) + if stat.S_ISLNK(member.external_attr >> 16): + return _ZipMemberAudit((AuditFinding(location, "archive member is a symbolic link"),)) + if member.is_dir(): + return _ZipMemberAudit((), dist_info_root=dist_info_root) + is_license = _is_distribution_license(member.filename, dist_info_root) + is_metadata = _is_distribution_metadata(member.filename, dist_info_root) + try: + with archive.open(member) as stream: + content = _read_bounded(stream, expected_size=member.file_size) + except ValueError as error: + return _ZipMemberAudit((AuditFinding(location, str(error)),), dist_info_root=dist_info_root) + except (OSError, RuntimeError, NotImplementedError, zipfile.BadZipFile): + return _ZipMemberAudit( + (AuditFinding(location, "archive member cannot be read"),), + dist_info_root=dist_info_root, + ) + return _ZipMemberAudit( + findings=tuple(_audit_content(location, member.filename, content)), + dist_info_root=dist_info_root, + is_license=is_license, + license_matches=is_license and content == canonical_license, + is_metadata=is_metadata, + metadata_declares_license=is_metadata and b"\nLicense-Expression: Apache-2.0\n" in b"\n" + content, + ) + + +def _audit_wheel_distribution(path: Path, members: Iterable[_ZipMemberAudit]) -> list[AuditFinding]: + member_audits = tuple(members) + roots = frozenset(member.dist_info_root for member in member_audits if member.dist_info_root is not None) + location = _display_path(path) + if len(roots) != 1: + return [AuditFinding(location, "wheel must contain exactly one .dist-info root")] + root = next(iter(roots)) + if _SLURM_DIST_INFO_PATTERN.fullmatch(root.casefold()) is None: + return [AuditFinding(location, "wheel .dist-info root does not identify data-designer-slurm")] + licenses = tuple(member for member in member_audits if member.dist_info_root == root and member.is_license) + metadata = tuple(member for member in member_audits if member.dist_info_root == root and member.is_metadata) + findings: list[AuditFinding] = [] + if len(licenses) != 1 or not licenses[0].license_matches: + findings.append(AuditFinding(location, "wheel does not contain exactly one canonical license text")) + if len(metadata) != 1 or not metadata[0].metadata_declares_license: + findings.append(AuditFinding(location, "wheel does not contain exactly one Apache-2.0 metadata record")) + return findings + + +def _audit_tar_member( + archive: tarfile.TarFile, + member: tarfile.TarInfo, + location: str, +) -> list[AuditFinding]: + if _is_unsafe_archive_name(member.name): + return [AuditFinding(location, "archive member path is unsafe")] + if member.issym() or member.islnk(): + return [AuditFinding(location, "archive member is a link")] + if member.isdir(): + return [] + if not member.isfile(): + return [AuditFinding(location, "archive member type is unsupported")] + stream = archive.extractfile(member) + if stream is None: + return [AuditFinding(location, "archive member cannot be read")] + try: + with stream: + content = _read_bounded(stream, expected_size=member.size) + except (OSError, ValueError) as error: + return [AuditFinding(location, str(error))] + return _audit_content(location, member.name, content) + + +def _get_archive_limit_finding(path: Path, *, member_count: int, content_size: int) -> AuditFinding | None: + if member_count > MAXIMUM_ARCHIVE_MEMBERS: + return AuditFinding(_display_path(path), "archive exceeds the member-count scan limit") + if content_size > MAXIMUM_ARCHIVE_CONTENT_SIZE: + return AuditFinding(_display_path(path), "archive exceeds the expanded-content scan limit") + return None + + +def _audit_content( + location: str, + logical_name: str, + content: bytes, + *, + allowed_sentinels: Iterable[str] = (), +) -> list[AuditFinding]: + text = content.decode("utf-8", errors="replace") + for sentinel in sorted(allowed_sentinels, key=len, reverse=True): + text = re.sub( + rf"(?", + text, + ) + findings = [ + AuditFinding(location, rule.name) + for rule in _CONTENT_RULES + if rule.pattern.search(logical_name) or rule.pattern.search(text) + ] + if _requires_spdx(logical_name) and not _has_spdx_header(text): + findings.append(AuditFinding(location, "packaged source is missing the NVIDIA Apache-2.0 SPDX header")) + return findings + + +def _read_bounded(stream: BinaryIO, *, expected_size: int) -> bytes: + if expected_size > MAXIMUM_MEMBER_SIZE: + raise ValueError("artifact exceeds the per-file scan limit") + content = stream.read(MAXIMUM_MEMBER_SIZE + 1) + if len(content) > MAXIMUM_MEMBER_SIZE: + raise ValueError("artifact exceeds the per-file scan limit") + return content + + +def _requires_spdx(logical_name: str) -> bool: + normalized = logical_name.replace("\\", "/") + suffix = PurePosixPath(normalized).suffix.casefold() + return suffix in _LICENSED_SOURCE_SUFFIXES and ( + normalized == "entrypoint.sh" + or normalized.startswith("data_designer/slurm/") + or "/src/data_designer/slurm/" in normalized + ) + + +def _has_spdx_header(text: str) -> bool: + header = "\n".join(text.splitlines()[:5]) + return _SPDX_COPYRIGHT in header and _SPDX_LICENSE in header + + +def _get_dist_info_root(name: str) -> str | None: + parts = PurePosixPath(name).parts + if not parts or not parts[0].casefold().endswith(".dist-info"): + return None + return parts[0] + + +def _is_distribution_license(name: str, root: str | None) -> bool: + if root is None: + return False + parts = tuple(part.casefold() for part in PurePosixPath(name).parts) + license_names = {"license", "license.md", "license.txt"} + return (len(parts) == 2 and parts[1] in license_names) or ( + len(parts) == 3 and parts[1] == "licenses" and parts[2] in license_names + ) + + +def _is_distribution_metadata(name: str, root: str | None) -> bool: + return root is not None and name.casefold() == f"{root.casefold()}/metadata" + + +def _is_unsafe_archive_name(name: str) -> bool: + path = PurePosixPath(name) + return path.is_absolute() or ".." in path.parts or "\\" in name + + +def _is_zip_archive(path: Path) -> bool: + if path.suffix.casefold() in {".whl", ".zip"}: + return True + try: + return zipfile.is_zipfile(path) + except OSError: + return False + + +def _is_tar_archive(path: Path) -> bool: + normalized = path.name.casefold() + if normalized.endswith((".tar", ".tar.gz", ".tgz")): + return True + try: + return tarfile.is_tarfile(path) + except OSError: + return False + + +def _display_path(path: Path) -> str: + try: + return path.relative_to(REPOSITORY_ROOT).as_posix() + except ValueError: + return "" + + +def main(arguments: list[str] | None = None) -> int: + """Run the public-artifact audit and print only path-scoped rule names.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="*", type=Path, help="Files or directories to scan") + parsed = parser.parse_args(arguments) + paths = parsed.paths or [REPOSITORY_ROOT / relative for relative in DEFAULT_ARTIFACTS] + findings = audit_public_artifacts(paths) + if findings: + for finding in findings: + print(f"{finding.location}: {finding.rule}", file=sys.stderr) + return 1 + print(f"Slurm public-artifact audit passed for {len(paths)} target(s).") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/check_dependency_licenses.py b/scripts/check_dependency_licenses.py index da7d50ae5..c8bb0afac 100644 --- a/scripts/check_dependency_licenses.py +++ b/scripts/check_dependency_licenses.py @@ -21,6 +21,7 @@ POLICY_PATH = Path(__file__).resolve().parents[1] / "dependency-license-policy.toml" LICENSE_ALIASES = { + "Apache License 2.0": {"Apache-2.0"}, "Apache Software License": {"Apache-2.0"}, "BSD License": {"BSD-3-Clause"}, "ISC License (ISCL)": {"ISC"}, diff --git a/scripts/publish.sh b/scripts/publish.sh index 70e29fd8b..aec30b0c0 100755 --- a/scripts/publish.sh +++ b/scripts/publish.sh @@ -4,7 +4,7 @@ # # Publish script for DataDesigner -# Publishes all three subpackages to PyPI with the same version. +# Publishes all four subpackages to PyPI with the same version. # # Usage: # ./scripts/publish.sh 0.3.9rc1 # Full publish @@ -31,6 +31,8 @@ NC='\033[0m' # No Color PACKAGE_DIRS=( "packages/data-designer-config" "packages/data-designer-engine" + # Publish the provider before the package that advertises its exact-version extra. + "packages/data-designer-slurm" "packages/data-designer" ) @@ -544,6 +546,7 @@ main() { echo " https://test.pypi.org/project/data-designer-config/$VERSION/" echo " https://test.pypi.org/project/data-designer-engine/$VERSION/" echo " https://test.pypi.org/project/data-designer/$VERSION/" + echo " https://test.pypi.org/project/data-designer-slurm/$VERSION/" echo "" echo "Test installation with:" echo " pip install --index-url $TEST_PYPI_URL data-designer==$VERSION" @@ -564,6 +567,7 @@ main() { echo " https://pypi.org/project/data-designer-config/$VERSION/" echo " https://pypi.org/project/data-designer-engine/$VERSION/" echo " https://pypi.org/project/data-designer/$VERSION/" + echo " https://pypi.org/project/data-designer-slurm/$VERSION/" fi } diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py new file mode 100644 index 000000000..ce2c8b178 --- /dev/null +++ b/scripts/test_slurm_package_install.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +import shutil +import statistics +import subprocess +import sys +import tempfile +import time +from contextlib import contextmanager +from email.message import Message +from email.parser import BytesParser +from functools import partial +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Thread +from typing import Iterator +from zipfile import ZipFile + +from packaging.requirements import Requirement +from packaging.utils import canonicalize_name + +REPOSITORY_ROOT = Path(__file__).parents[1] +PACKAGE_PATHS = ( + "packages/data-designer-config", + "packages/data-designer-engine", + "packages/data-designer", + "packages/data-designer-slurm", +) +EXAMPLE_PATH = REPOSITORY_ROOT / "packages" / "data-designer-slurm" / "examples" +SLURM_DOCS_PATH = REPOSITORY_ROOT / "fern" / "versions" / "latest" / "pages" / "slurm" +FIRST_PARTY_PACKAGES = frozenset( + {"data-designer", "data-designer-config", "data-designer-engine", "data-designer-slurm"} +) +CLI_HELP_SAMPLES = 9 +MAX_BASE_CLI_HELP_SECONDS = 1.0 +MAX_EXTENSION_CLI_HELP_OVERHEAD_SECONDS = 0.1 + + +class _QuietIndexHandler(SimpleHTTPRequestHandler): + def log_message(self, _format: str, *_args: object) -> None: + pass + + +def run( + command: list[str], + *, + cwd: Path, + check: bool = True, + isolate_pip: bool = False, +) -> subprocess.CompletedProcess[str]: + environment = os.environ.copy() + environment.pop("PYTHONPATH", None) + environment.pop("VIRTUAL_ENV", None) + if isolate_pip: + environment = {name: value for name, value in environment.items() if not name.startswith("PIP_")} + environment["PIP_CONFIG_FILE"] = os.devnull + result = subprocess.run(command, cwd=cwd, env=environment, capture_output=True, text=True, check=False) + if check and result.returncode: + raise RuntimeError(result.stdout + result.stderr) + return result + + +def wheel_metadata(path: Path) -> Message: + with ZipFile(path) as wheel: + metadata_path = next(name for name in wheel.namelist() if name.endswith(".dist-info/METADATA")) + return BytesParser().parsebytes(wheel.read(metadata_path)) + + +def audit_public_artifacts(*paths: Path) -> None: + result = run( + [sys.executable, str(REPOSITORY_ROOT / "scripts" / "audit_slurm_public_artifacts.py"), *map(str, paths)], + cwd=REPOSITORY_ROOT, + check=False, + ) + if result.returncode: + raise RuntimeError(result.stdout + result.stderr) + + +def verify_documented_examples() -> None: + documented_examples = { + "builder.yaml": "getting-started.mdx", + "run.yaml": "getting-started.mdx", + "benchmark.yaml": "benchmarks.mdx", + "profile-catalog.yaml": "profiles.mdx", + } + for name, document_name in documented_examples.items(): + example = (EXAMPLE_PATH / name).read_text().rstrip() + document = SLURM_DOCS_PATH.joinpath(document_name).read_text() + assert f"```yaml\n{example}\n```" in document + + +def build_wheels(uv: str, wheel_directory: Path) -> dict[str, Path]: + for package_path in PACKAGE_PATHS: + run( + [uv, "build", "--wheel", "--out-dir", str(wheel_directory), package_path], + cwd=REPOSITORY_ROOT, + ) + + wheels = {} + for wheel_path in wheel_directory.glob("*.whl"): + name = canonicalize_name(wheel_metadata(wheel_path)["Name"]) + wheels[name] = wheel_path + return wheels + + +def requirement(metadata: Message, dependency_name: str) -> Requirement: + requirements = [Requirement(value) for value in metadata.get_all("Requires-Dist", [])] + return next(item for item in requirements if canonicalize_name(item.name) == dependency_name) + + +def python_path(environment_path: Path) -> Path: + if os.name == "nt": + return environment_path / "Scripts" / "python.exe" + return environment_path / "bin" / "python" + + +def create_environment(uv: str, environment_path: Path, *, cwd: Path) -> Path: + run([uv, "venv", "--python", sys.executable, str(environment_path)], cwd=cwd) + return python_path(environment_path) + + +def install(uv: str, python: Path, wheel_directory: Path, package: str, *, cwd: Path) -> None: + run( + [ + uv, + "pip", + "install", + "--python", + str(python), + "--prerelease=allow", + "--no-index", + "--find-links", + str(wheel_directory), + package, + ], + cwd=cwd, + ) + + +def download_third_party_wheels(wheel_directory: Path, metadata: dict[str, Message]) -> None: + requirements = { + str(parsed) + for distribution in metadata.values() + for value in distribution.get_all("Requires-Dist", []) + if canonicalize_name((parsed := Requirement(value)).name) not in FIRST_PARTY_PACKAGES + and (parsed.marker is None or parsed.marker.evaluate()) + } + run( + [ + sys.executable, + "-m", + "pip", + "download", + "--disable-pip-version-check", + "--only-binary=:all:", + "--pre", + "--dest", + str(wheel_directory), + *sorted(requirements), + ], + cwd=REPOSITORY_ROOT, + isolate_pip=True, + ) + + +def create_simple_index(root: Path, wheel_directory: Path) -> None: + wheels_by_package: dict[str, list[Path]] = {} + for wheel in wheel_directory.glob("*.whl"): + name = canonicalize_name(wheel_metadata(wheel)["Name"]) + wheels_by_package.setdefault(name, []).append(wheel) + for name, wheels in wheels_by_package.items(): + package_index = root / "simple" / name + package_index.mkdir(parents=True) + links = [] + for wheel in sorted(wheels): + digest = hashlib.sha256(wheel.read_bytes()).hexdigest() + links.append(f'{wheel.name}') + package_index.joinpath("index.html").write_text("".join(f"{link}\n" for link in links)) + + +@contextmanager +def serve_index(root: Path) -> Iterator[str]: + handler = partial(_QuietIndexHandler, directory=str(root)) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/simple" + finally: + server.shutdown() + thread.join() + server.server_close() + + +def install_from_index(python: Path, index_url: str, package: str, *, cwd: Path) -> None: + run( + [ + str(python), + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--pre", + "--index-url", + index_url, + package, + ], + cwd=cwd, + isolate_pip=True, + ) + + +def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> None: + statement = f""" +import sys +from importlib.metadata import version +from importlib.util import find_spec + +import data_designer +import data_designer.config +import data_designer.engine +import data_designer.interface +from pathlib import Path +from tempfile import TemporaryDirectory +from typer.testing import CliRunner + +from data_designer.cli.main import app + +assert data_designer.__file__ is None +assert version("data-designer") == {version!r} +assert (find_spec("data_designer.slurm") is not None) is {slurm!r} +assert "data_designer.slurm" not in sys.modules +assert "packaging.requirements" not in sys.modules + +help_result = CliRunner().invoke(app, ["--help"]) +assert help_result.exit_code == 0, help_result.output +assert ("slurm" in help_result.output) is {slurm!r} +assert "data_designer.slurm" not in sys.modules +assert "packaging.requirements" not in sys.modules +""" + if slurm: + 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", "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.config import load_benchmark_config, load_builder_payload, load_profile_catalog, load_run_config +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 +from data_designer.slurm.integration import PlanStateValidator +from data_designer.slurm.images.registry import ImageRegistryStore +from data_designer.slurm.planning import ArtifactReference as PlanningArtifactReference +from data_designer.slurm.planning import RecordRange as PlanningRecordRange +from data_designer.slurm.planning import ResumeWorkspace as PlanningResumeWorkspace +from data_designer.slurm.runtime import proxy as runtime_proxy +from data_designer.slurm.state import ArtifactReference as StateArtifactReference +from data_designer.slurm.state import RecordRange as StateRecordRange +from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace +from data_designer.slurm.state import ( + CollectionResult, + RetryPlan, + RunManifest, + RunStatus, + SchedulerObservationCollector, + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateReconciler, +) +from data_designer.slurm.services import create_slurm_benchmark_service +assert BenchmarkCompiler.__name__ == "BenchmarkCompiler" +assert callable(runtime_proxy.main) +assert callable(create_slurm_benchmark_service) +assert CollectionResult.__name__ == "CollectionResult" +assert RetryPlan.__name__ == "RetryPlan" +assert RunManifest.__name__ == "RunManifest" +assert RunStatus.__name__ == "RunStatus" +assert SchedulerObservationCollector.__name__ == "SchedulerObservationCollector" +assert SlurmCollectionCoordinator.__name__ == "SlurmCollectionCoordinator" +assert SlurmRetryCoordinator.__name__ == "SlurmRetryCoordinator" +assert SlurmStateReconciler.__name__ == "SlurmStateReconciler" +assert ImageRegistryStore.__name__ == "ImageRegistryStore" +assert PlanningArtifactReference is ContractArtifactReference +assert PlanningRecordRange is ContractRecordRange +assert PlanningResumeWorkspace is ContractResumeWorkspace +assert StateArtifactReference is ContractArtifactReference +assert StateRecordRange is ContractRecordRange +assert StateResumeWorkspace is ContractResumeWorkspace +examples = Path("examples") +builder_example = load_builder_payload(examples / "builder.yaml") +assert [column["name"] for column in builder_example["data_designer"]["columns"]] == ["tone", "greeting"] +run_example = load_run_config(examples / "run.yaml") +benchmark_example = load_benchmark_config(examples / "benchmark.yaml") +assert run_example.name == "greeting-run" +assert benchmark_example.name == "generator-scaling" +assert len(BenchmarkCompiler.compile(benchmark_example, run_example).cases) == 4 +assert load_profile_catalog(examples / "profile-catalog.yaml").default_cluster == "primary" +profile_help_result = CliRunner().invoke(app, ["slurm", "profile", "--help"]) +assert profile_help_result.exit_code == 0, profile_help_result.output +with TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + profile_file = root / "profile.yml" + profile_init_result = CliRunner().invoke( + app, + [ + "slurm", + "profile", + "init", + "--workspace-root", + str(root / "workspace"), + "--image-build-partition", + "cpu", + "--profile-file", + str(profile_file), + "--host-pattern", + "login.example.test", + ], + ) + assert profile_init_result.exit_code == 0, profile_init_result.output + assert profile_file.is_file() +""" + run([str(python), "-c", statement], cwd=cwd) + + +def cli_help_medians(base_python: Path, extension_python: Path, *, cwd: Path) -> tuple[float, float]: + statement = """ +from typer.testing import CliRunner +from data_designer.cli.main import app + +result = CliRunner().invoke(app, ["--help"]) +assert result.exit_code == 0, result.output +""" + for python in (base_python, extension_python): + run([str(python), "-c", statement], cwd=cwd) + + samples: dict[Path, list[float]] = {base_python: [], extension_python: []} + for index in range(CLI_HELP_SAMPLES): + environments = (base_python, extension_python) if index % 2 == 0 else (extension_python, base_python) + for python in environments: + start = time.perf_counter() + run([str(python), "-c", statement], cwd=cwd) + samples[python].append(time.perf_counter() - start) + return statistics.median(samples[base_python]), statistics.median(samples[extension_python]) + + +def main() -> None: + uv = shutil.which("uv") + if uv is None: + raise RuntimeError("uv is required") + if run([sys.executable, "-m", "pip", "--version"], cwd=REPOSITORY_ROOT, check=False).returncode: + raise RuntimeError("pip is required") + verify_documented_examples() + + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + shutil.copytree(EXAMPLE_PATH, root / "examples") + wheel_directory = root / "wheels" + wheel_directory.mkdir() + wheels = build_wheels(uv, wheel_directory) + + assert set(wheels) == FIRST_PARTY_PACKAGES + metadata = {name: wheel_metadata(path) for name, path in wheels.items()} + download_third_party_wheels(wheel_directory, metadata) + create_simple_index(root, wheel_directory) + versions = {item["Version"] for item in metadata.values()} + assert len(versions) == 1 + leaf_wheel = wheels["data-designer-slurm"] + audit_public_artifacts(leaf_wheel) + base_metadata = metadata["data-designer"] + leaf_metadata = metadata["data-designer-slurm"] + version = base_metadata["Version"] + assert leaf_metadata["Version"] == version + + base_config_requirement = requirement(base_metadata, "data-designer-config") + base_engine_requirement = requirement(base_metadata, "data-designer-engine") + engine_config_requirement = requirement(metadata["data-designer-engine"], "data-designer-config") + base_leaf_requirement = requirement(base_metadata, "data-designer-slurm") + leaf_base_requirement = requirement(leaf_metadata, "data-designer") + leaf_packaging_requirement = requirement(leaf_metadata, "packaging") + leaf_pip_requirement = requirement(leaf_metadata, "pip") + leaf_pydantic_requirement = requirement(leaf_metadata, "pydantic") + leaf_pyyaml_requirement = requirement(leaf_metadata, "pyyaml") + leaf_aiohttp_requirement = requirement(leaf_metadata, "aiohttp") + assert str(base_config_requirement.specifier) == f"=={version}" + assert str(base_engine_requirement.specifier) == f"=={version}" + assert str(engine_config_requirement.specifier) == f"=={version}" + assert str(base_leaf_requirement.specifier) == f"=={version}" + 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.12,<3").specifier + assert leaf_pyyaml_requirement.specifier == Requirement("pyyaml>=6.0.1,<7").specifier + assert leaf_aiohttp_requirement.specifier == Requirement("aiohttp>=3.14.3,<4").specifier + assert base_leaf_requirement.marker is not None + assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) + assert not base_leaf_requirement.marker.evaluate({"extra": ""}) + assert "slurm" in base_metadata.get_all("Provides-Extra", []) + with ZipFile(leaf_wheel) as wheel: + assert "data_designer/__init__.py" not in wheel.namelist() + + base_python = create_environment(uv, root / "base", cwd=root) + install(uv, base_python, wheel_directory, f"data-designer=={version}", cwd=root) + verify_install(base_python, version, slurm=False, cwd=root) + + extra_python = create_environment(uv, root / "extra", cwd=root) + install(uv, extra_python, wheel_directory, "pip", cwd=root) + with serve_index(root) as index_url: + install_from_index(extra_python, index_url, f"data-designer[slurm]=={version}", cwd=root) + verify_install(extra_python, version, slurm=True, cwd=root) + base_cli_help, extension_cli_help = cli_help_medians(base_python, extra_python, cwd=root) + extension_overhead = extension_cli_help - base_cli_help + assert base_cli_help <= MAX_BASE_CLI_HELP_SECONDS, ( + f"Base CLI root help took {base_cli_help:.3f}s; budget is {MAX_BASE_CLI_HELP_SECONDS:.3f}s" + ) + assert extension_overhead <= MAX_EXTENSION_CLI_HELP_OVERHEAD_SECONDS, ( + f"CLI extension added {extension_overhead:.3f}s to root help " + f"(base={base_cli_help:.3f}s, extension={extension_cli_help:.3f}s)" + ) + print( + f"CLI help median: base={base_cli_help:.3f}s, extension={extension_cli_help:.3f}s, " + f"overhead={extension_overhead:.3f}s" + ) + + leaf_python = create_environment(uv, root / "leaf", cwd=root) + install(uv, leaf_python, wheel_directory, f"data-designer-slurm=={version}", cwd=root) + verify_install(leaf_python, version, slurm=True, cwd=root) + + leaf_only_directory = root / "leaf-only" + leaf_only_directory.mkdir() + shutil.copy2(leaf_wheel, leaf_only_directory) + missing_python = create_environment(uv, root / "missing", cwd=root) + result = run( + [ + uv, + "pip", + "install", + "--python", + str(missing_python), + "--no-index", + "--find-links", + str(leaf_only_directory), + str(leaf_wheel), + ], + cwd=root, + check=False, + ) + assert result.returncode != 0 + assert f"data-designer=={version}" in (result.stdout + result.stderr) + + +if __name__ == "__main__": + main() diff --git a/tests_e2e/test_dependency_license_check.py b/tests_e2e/test_dependency_license_check.py index 081e71595..0cccf79af 100644 --- a/tests_e2e/test_dependency_license_check.py +++ b/tests_e2e/test_dependency_license_check.py @@ -51,6 +51,7 @@ def test_evaluate_report_accepts_permissive_and_composite_licenses() -> None: ) packages = [ module.PackageLicense("apache", "1", "Apache-2.0", ""), + module.PackageLicense("multidict", "6.9.1", "Apache License 2.0", ""), module.PackageLicense("dual", "1", "Apache-2.0 OR BSD-3-Clause", ""), module.PackageLicense("mit", "1", "MIT License", ""), ] @@ -61,6 +62,25 @@ def test_evaluate_report_accepts_permissive_and_composite_licenses() -> None: assert reviewed == [] +def test_locked_psf_dependency_requires_exact_reviewed_version_and_license() -> None: + module = load_script() + approved = module.load_policy().exceptions["aiohappyeyeballs"] + policy = module.LicensePolicy(allowed_licenses=frozenset(), exceptions={"aiohappyeyeballs": approved}) + + accepted, reviewed = module.evaluate_report( + [ + module.PackageLicense("aiohappyeyeballs", "2.6.1", "PSF-2.0", ""), + module.PackageLicense("aiohappyeyeballs", "2.6.1", "Python Software Foundation License", ""), + ], + policy, + ) + changed, _ = module.evaluate_report([module.PackageLicense("aiohappyeyeballs", "2.7.1", "PSF-2.0", "")], policy) + + assert accepted == [] + assert len(reviewed) == 2 + assert len(changed) == 1 + + def test_evaluate_report_accepts_allowed_or_alternative() -> None: module = load_script() policy = module.LicensePolicy(allowed_licenses=frozenset({"Apache-2.0", "MIT"}), exceptions={}) diff --git a/uv.lock b/uv.lock index 7db38303f..3e6d228f9 100644 --- a/uv.lock +++ b/uv.lock @@ -19,6 +19,7 @@ members = [ "data-designer", "data-designer-config", "data-designer-engine", + "data-designer-slurm", "data-designer-workspace", ] constraints = [{ name = "certifi", specifier = ">=2026.10.19" }] @@ -813,10 +814,16 @@ dependencies = [ { name = "typer" }, ] +[package.optional-dependencies] +slurm = [ + { name = "data-designer-slurm" }, +] + [package.metadata] requires-dist = [ { name = "data-designer-config", editable = "packages/data-designer-config" }, { name = "data-designer-engine", editable = "packages/data-designer-engine" }, + { name = "data-designer-slurm", marker = "extra == 'slurm'", editable = "packages/data-designer-slurm" }, { name = "huggingface-hub", specifier = ">=1.0.1,<2" }, { name = "opentelemetry-api", specifier = ">=1.43,<1.44" }, { name = "opentelemetry-exporter-prometheus", specifier = ">=0.64b0,<0.65" }, @@ -831,6 +838,7 @@ requires-dist = [ { name = "rich", specifier = ">=13.7.1,<15" }, { name = "typer", specifier = ">=0.12.0,<1" }, ] +provides-extras = ["slurm"] [[package]] name = "data-designer-config" @@ -939,6 +947,28 @@ requires-dist = [ { name = "wcwidth", specifier = ">=0.2.13,<1" }, ] +[[package]] +name = "data-designer-slurm" +source = { editable = "packages/data-designer-slurm" } +dependencies = [ + { name = "aiohttp" }, + { name = "data-designer" }, + { name = "packaging" }, + { name = "pip" }, + { name = "pydantic" }, + { name = "pyyaml" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", specifier = ">=3.14.3,<4" }, + { name = "data-designer", editable = "packages/data-designer" }, + { name = "packaging", specifier = ">=25,<27" }, + { name = "pip", specifier = ">=25,<27" }, + { name = "pydantic", specifier = ">=2.12,<3" }, + { name = "pyyaml", specifier = ">=6.0.1,<7" }, +] + [[package]] name = "data-designer-workspace" version = "0.0.0" @@ -3380,6 +3410,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/37/66b12dc476041fe523bca29c9d58fbf1ea92a4eaf73acfb9e993948f1191/pillow-13.0.0-pp312-pypy312_pp80-win_amd64.whl", hash = "sha256:567f20f1c75897d7663830da6ea8af719e03d712b1d9ee7f813319f9ece42e5a", size = 7366030, upload-time = "2026-10-16T22:18:06.638Z" }, ] +[[package]] +name = "pip" +version = "26.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/15/4500e320e6b101ec3b719ae85b697d9940b6cda672bc555bd6016fc60c6f/pip-26.2.1.tar.gz", hash = "sha256:f6ad667e89a1fe78046c8f13232b247200f5258d7828f3f7883d660878e0813f", size = 1848877, upload-time = "2026-08-04T22:51:14.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/6e/1736e5b4ae2b778ef2f81c47d797de9f891d4d8acb047a24ca37a60294dd/pip-26.2.1-py3-none-any.whl", hash = "sha256:71138adf1f4ca900cdb7d289c21b7494329f2332b6d85f0e1c42108c0384ed3e", size = 1816632, upload-time = "2026-08-04T22:51:12.472Z" }, +] + [[package]] name = "platformdirs" version = "4.9.4"