From 1111b90a38df65e979d540179b800a31f11e371b Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:52:28 -0300 Subject: [PATCH 01/62] feat: add optional Slurm package (#856) * feat: add optional Slurm package Add the shared-namespace leaf package and publish it through the same-version data-designer[slurm] extra. Cover resolver, namespace, base-only isolation, and built-wheel installation behavior.\n\nCloses #852 Signed-off-by: Andre Manoel * fix: address Slurm package review findings Signed-off-by: Andre Manoel * fix: publish Slurm package Signed-off-by: Andre Manoel * fix: scope workspace dependency extras Signed-off-by: Andre Manoel * fix: publish Slurm before base package Upload the Slurm distribution before data-designer so the exact-version extra is resolvable as soon as the base package is published. Signed-off-by: Andre Manoel * docs: generalize Slurm package wording Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- .agents/recipes/dependencies/recipe.md | 11 +- .agents/recipes/structure/recipe.md | 13 +- .agents/recipes/test-health/recipe.md | 2 +- .github/workflows/agentic-ci-daily.yml | 1 + .github/workflows/ci.yml | 31 +++- AGENTS.md | 7 +- DEVELOPMENT.md | 5 +- Makefile | 79 ++++++-- architecture/overview.md | 8 +- packages/data-designer-slurm/README.md | 9 + packages/data-designer-slurm/pyproject.toml | 50 +++++ .../src/data_designer/slurm/__init__.py | 6 + .../data-designer-slurm/tests/test_package.py | 22 +++ packages/data-designer/pyproject.toml | 4 +- .../tests/test_dependency_audit.py | 107 ++++++++++- pyproject.toml | 1 + scripts/audit_package_dependencies.py | 62 +++++-- scripts/publish.sh | 6 +- scripts/test_slurm_package_install.py | 175 ++++++++++++++++++ uv.lock | 18 ++ 20 files changed, 555 insertions(+), 62 deletions(-) create mode 100644 packages/data-designer-slurm/README.md create mode 100644 packages/data-designer-slurm/pyproject.toml create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/__init__.py create mode 100644 packages/data-designer-slurm/tests/test_package.py create mode 100644 scripts/test_slurm_package_install.py 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..b36fb1411 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -129,6 +129,31 @@ jobs: uv run --with pytest --with pytest-asyncio --with pytest-httpx --with pytest-env \ pytest packages/data-designer/tests + test-slurm-package: + name: Test Slurm package wheels + needs: validate-dispatch + runs-on: ubuntu-latest + + 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: "3.11" + enable-cache: true + + - name: Install development dependencies + run: make install-dev + + - name: Run package tests + run: .venv/bin/pytest packages/data-designer-slurm/tests + + - name: Run built-wheel installation tests + run: make test-slurm-wheel-install + # =========================================================================== # Combined Coverage Check # Runs all tests together to verify overall coverage threshold @@ -284,7 +309,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 +322,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/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/packages/data-designer-slurm/README.md b/packages/data-designer-slurm/README.md new file mode 100644 index 000000000..04d29ccbd --- /dev/null +++ b/packages/data-designer-slurm/README.md @@ -0,0 +1,9 @@ +# data-designer-slurm + +Optional Slurm execution support for Data Designer. + +Install it through the Data Designer extra: + +```bash +pip install "data-designer[slurm]" +``` diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml new file mode 100644 index 000000000..17d2906df --- /dev/null +++ b/packages/data-designer-slurm/pyproject.toml @@ -0,0 +1,50 @@ +[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" + +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", +] + +[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 = ["data-designer=={{ version }}"] + +[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/tests/test_package.py b/packages/data-designer-slurm/tests/test_package.py new file mode 100644 index 000000000..b9ed63874 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_package.py @@ -0,0 +1,22 @@ +# 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 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_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"') 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/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/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/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..c42bdf477 --- /dev/null +++ b/scripts/test_slurm_package_install.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import tempfile +from email.message import Message +from email.parser import BytesParser +from pathlib import Path +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", +) + + +def run(command: list[str], *, cwd: Path, check: bool = True) -> subprocess.CompletedProcess[str]: + environment = os.environ.copy() + environment.pop("PYTHONPATH", None) + environment.pop("VIRTUAL_ENV", None) + 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 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", + "--find-links", + str(wheel_directory), + package, + ], + cwd=cwd, + ) + + +def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> None: + statement = f""" +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 + +assert data_designer.__file__ is None +assert version("data-designer") == {version!r} +assert (find_spec("data_designer.slurm") is not None) is {slurm!r} +""" + if slurm: + statement += f'\nimport data_designer.slurm\nassert version("data-designer-slurm") == {version!r}\n' + run([str(python), "-c", statement], cwd=cwd) + + +def main() -> None: + uv = shutil.which("uv") + if uv is None: + raise RuntimeError("uv is required") + + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + wheel_directory = root / "wheels" + wheel_directory.mkdir() + wheels = build_wheels(uv, wheel_directory) + + base_wheel = wheels["data-designer"] + leaf_wheel = wheels["data-designer-slurm"] + base_metadata = wheel_metadata(base_wheel) + leaf_metadata = wheel_metadata(leaf_wheel) + version = base_metadata["Version"] + assert leaf_metadata["Version"] == version + + base_leaf_requirement = requirement(base_metadata, "data-designer-slurm") + leaf_base_requirement = requirement(leaf_metadata, "data-designer") + assert str(base_leaf_requirement.specifier) == f"=={version}" + assert str(leaf_base_requirement.specifier) == f"=={version}" + 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, f"data-designer[slurm]=={version}", cwd=root) + verify_install(extra_python, version, slurm=True, cwd=root) + + 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/uv.lock b/uv.lock index 7db38303f..52329761e 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,16 @@ requires-dist = [ { name = "wcwidth", specifier = ">=0.2.13,<1" }, ] +[[package]] +name = "data-designer-slurm" +source = { editable = "packages/data-designer-slurm" } +dependencies = [ + { name = "data-designer" }, +] + +[package.metadata] +requires-dist = [{ name = "data-designer", editable = "packages/data-designer" }] + [[package]] name = "data-designer-workspace" version = "0.0.0" From c63fff4a0b81aef79758f6f0cfd87b0786906294 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:29:02 -0300 Subject: [PATCH 02/62] feat: define public generation invocation contract (#854) * feat: define public generation contract * fix: complete public generation contract * docs: clarify public invocation limits * fix: handle mixed-case config extensions Signed-off-by: Andre Manoel * fix: complete generation contract edge cases Signed-off-by: Andre Manoel * test: make inline config regression self-contained Signed-off-by: Andre Manoel * fix: classify malformed config paths Signed-off-by: Andre Manoel * fix: normalize generation artifact errors Classify malformed inline configs independently of filename suffixes. Wrap unreadable result artifacts in the public storage error across counting and export paths. Signed-off-by: Andre Manoel * fix: complete artifact error boundaries Signed-off-by: Andre Manoel * fix: tighten generation result contract Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- .../concepts/architecture-and-performance.mdx | 16 +- .../src/data_designer/config/__init__.py | 9 + .../data_designer/config/config_builder.py | 19 +- .../data_designer/config/utils/io_helpers.py | 11 +- .../tests/config/test_public_contract.py | 71 ++++ .../src/data_designer/interface/__init__.py | 2 + .../interface/composite_workflow.py | 8 + .../data_designer/interface/data_designer.py | 8 +- .../src/data_designer/interface/results.py | 90 ++++- .../tests/interface/test_acreate.py | 1 + .../interface/test_composite_workflow.py | 60 ++++ .../tests/interface/test_data_designer.py | 106 ++++++ .../tests/interface/test_results.py | 131 ++++++- plans/850/data-designer-contract.md | 330 ++++++++++++++++++ 14 files changed, 836 insertions(+), 26 deletions(-) create mode 100644 packages/data-designer-config/tests/config/test_public_contract.py create mode 100644 plans/850/data-designer-contract.md 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/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/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/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/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. From 969a961f2b776b42f3a6a763d5fc98707284273d Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:29:12 -0300 Subject: [PATCH 03/62] feat: add lazy CLI extensions (#863) * feat: add lazy CLI extensions Signed-off-by: Andre Manoel * fix: isolate CLI extension failures Signed-off-by: Andre Manoel * refactor: order public CLI members first Signed-off-by: Nabin Mulepati --------- Signed-off-by: Andre Manoel Signed-off-by: Nabin Mulepati Co-authored-by: Nabin Mulepati --- architecture/cli.md | 27 ++ packages/data-designer-slurm/pyproject.toml | 3 + .../src/data_designer/slurm/cli.py | 23 ++ .../data-designer-slurm/tests/test_package.py | 12 + .../src/data_designer/cli/lazy_group.py | 292 +++++++++++++--- .../src/data_designer/cli/main.py | 5 +- .../tests/cli/test_lazy_group.py | 324 ++++++++++++++++++ scripts/test_slurm_package_install.py | 58 +++- 8 files changed, 696 insertions(+), 48 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/cli.py create mode 100644 packages/data-designer/tests/cli/test_lazy_group.py 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/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 17d2906df..70866a129 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -19,6 +19,9 @@ classifiers = [ "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" 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..a77cc462b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click +import typer + +app = typer.Typer( + name="slurm", + help="Run Data Designer workloads on Slurm", + no_args_is_help=True, +) + + +@app.callback() +def slurm_callback() -> None: + pass + + +def create_cli() -> click.Command: + """Create the Slurm CLI group.""" + return typer.main.get_command(app) diff --git a/packages/data-designer-slurm/tests/test_package.py b/packages/data-designer-slurm/tests/test_package.py index b9ed63874..2544dd973 100644 --- a/packages/data-designer-slurm/tests/test_package.py +++ b/packages/data-designer-slurm/tests/test_package.py @@ -3,6 +3,7 @@ from __future__ import annotations +import importlib.metadata from pathlib import Path import data_designer @@ -20,3 +21,14 @@ 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/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/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/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index c42bdf477..ff6fd6b6a 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -5,9 +5,11 @@ import os import shutil +import statistics import subprocess import sys import tempfile +import time from email.message import Message from email.parser import BytesParser from pathlib import Path @@ -23,6 +25,9 @@ "packages/data-designer", "packages/data-designer-slurm", ) +CLI_HELP_SAMPLES = 9 +MAX_BASE_CLI_HELP_SECONDS = 1.0 +MAX_EXTENSION_CLI_HELP_OVERHEAD_SECONDS = 0.1 def run(command: list[str], *, cwd: Path, check: bool = True) -> subprocess.CompletedProcess[str]: @@ -90,6 +95,7 @@ def install(uv: str, python: Path, wheel_directory: Path, package: str, *, cwd: 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 @@ -97,16 +103,53 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non import data_designer.config import data_designer.engine import data_designer.interface +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'\nimport data_designer.slurm\nassert version("data-designer-slurm") == {version!r}\n' + 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 "data_designer.slurm.cli" in sys.modules +assert version("data-designer-slurm") == {version!r} +""" 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: @@ -143,6 +186,19 @@ def main() -> None: extra_python = create_environment(uv, root / "extra", cwd=root) install(uv, extra_python, wheel_directory, 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) From 5d8a22db56362c9eaf08045e66dadee42c3cf538 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 20 Aug 2026 11:31:46 -0600 Subject: [PATCH 04/62] feat: define Slurm runtime and state record contracts (#878) * feat: add Slurm state record contracts Define strict, versioned execution, readiness, output, reconciliation, and scheduler records. Enforce cross-record identity, digest, shard-set, and terminal-state invariants with deterministic golden fixtures and built-wheel coverage. Closes #865 * fix: harden Slurm state contracts Reject partial winners and conflicting shard ownership. Require explicit versions, enforce evidence chronology and readiness precedence, and cover every state-contract path with validated fixtures. Refs #865 * fix(slurm): harden state contract invariants Hash persisted record bytes, align state identities with the resolved plan vocabulary, and preserve public model aliases. Add monotonic attempt and scheduler transition checks, including fixed accounting-lag deadlines and UNKNOWN precedence. Refs #865 * refactor(slurm): share contract primitives Make planning and runtime state consume one public immutable model family. Align shard, attempt, range, artifact, and resume workspace shapes so the configuration-plan branch can rebase without preserving duplicate Pydantic types.\n\nRefs #865 --- packages/data-designer-slurm/pyproject.toml | 5 +- .../src/data_designer/slurm/contracts.py | 180 +++++ .../src/data_designer/slurm/state/__init__.py | 116 ++++ .../src/data_designer/slurm/state/base.py | 58 ++ .../data_designer/slurm/state/execution.py | 113 +++ .../src/data_designer/slurm/state/outputs.py | 138 ++++ .../data_designer/slurm/state/readiness.py | 148 ++++ .../slurm/state/reconciliation.py | 170 +++++ .../data_designer/slurm/state/scheduler.py | 63 ++ .../data_designer/slurm/state/validation.py | 267 +++++++ .../tests/state/golden/accounting_lag.json | 10 + .../tests/state/golden/candidate_output.json | 22 + .../tests/state/golden/collection_plan.json | 23 + .../tests/state/golden/failed_attempt.json | 20 + .../state/golden/multi_node_readiness.json | 39 ++ .../tests/state/golden/run_manifest.json | 14 + .../tests/state/golden/shard_manifest.json | 18 + .../tests/state/golden/shard_winner.json | 12 + .../state/golden/single_node_readiness.json | 25 + .../tests/state/golden/stale_readiness.json | 25 + .../state/golden/successful_attempt.json | 23 + .../tests/state/test_golden_records.py | 63 ++ .../tests/state/test_model_edges.py | 214 ++++++ .../tests/state/test_records.py | 229 ++++++ .../tests/state/test_validation.py | 650 ++++++++++++++++++ .../tests/test_contracts.py | 65 ++ scripts/test_slurm_package_install.py | 10 + uv.lock | 6 +- 28 files changed, 2724 insertions(+), 2 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/contracts.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/base.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/execution.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/validation.py create mode 100644 packages/data-designer-slurm/tests/state/golden/accounting_lag.json create mode 100644 packages/data-designer-slurm/tests/state/golden/candidate_output.json create mode 100644 packages/data-designer-slurm/tests/state/golden/collection_plan.json create mode 100644 packages/data-designer-slurm/tests/state/golden/failed_attempt.json create mode 100644 packages/data-designer-slurm/tests/state/golden/multi_node_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/run_manifest.json create mode 100644 packages/data-designer-slurm/tests/state/golden/shard_manifest.json create mode 100644 packages/data-designer-slurm/tests/state/golden/shard_winner.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/stale_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/successful_attempt.json create mode 100644 packages/data-designer-slurm/tests/state/test_golden_records.py create mode 100644 packages/data-designer-slurm/tests/state/test_model_edges.py create mode 100644 packages/data-designer-slurm/tests/state/test_records.py create mode 100644 packages/data-designer-slurm/tests/state/test_validation.py create mode 100644 packages/data-designer-slurm/tests/test_contracts.py diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 70866a129..0fd4ec5ff 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -35,7 +35,10 @@ style = "pep440" bump = true [tool.hatch.metadata.hooks.uv-dynamic-versioning] -dependencies = ["data-designer=={{ version }}"] +dependencies = [ + "data-designer=={{ version }}", + "pydantic>=2.9.2,<3", +] [tool.hatch.build.targets.wheel] packages = ["src/data_designer"] 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..087b2a678 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -0,0 +1,180 @@ +# 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 typing import Annotated, Literal + +from pydantic import ( + BaseModel, + ConfigDict, + NonNegativeInt, + PositiveInt, + StringConstraints, + field_validator, + model_validator, +) + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +ModelAlias = str +ShardId = Annotated[str, StringConstraints(pattern=r"^shard-[0-9]{5,}$")] +AttemptId = Annotated[str, StringConstraints(pattern=r"^attempt-[0-9]{4,}$")] +SchemaVersion = Literal[1] +Sha256Digest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + + +class ContractValue(BaseModel): + """Base for strict immutable values shared across Slurm boundaries.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + allow_inf_nan=False, + protected_namespaces=(), + strict=True, + validate_default=True, + ) + + +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_sha256(value: object) -> Sha256Digest: + """Compute the canonical JSON digest of a JSON-compatible value.""" + return hashlib.sha256(canonical_json(value)).hexdigest() + + +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") + 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: + 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 + + +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", + "ContractRecord", + "ContractValue", + "Identifier", + "ModelAlias", + "RecordRange", + "ResumeWorkspace", + "SchemaVersion", + "Sha256Digest", + "ShardId", + "canonical_json", + "compute_sha256", + "pretty_json", + "validate_absolute_path", + "validate_relative_path", +] 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..2d7eb81d8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -0,0 +1,116 @@ +# 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 + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + ContractValue, + Identifier, + ModelAlias, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, +) +from data_designer.slurm.state.base import ( + SchedulerIdentity, + StateRecord, + StateValue, +) +from data_designer.slurm.state.execution import ( + AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + RunManifest, + ShardManifest, +) +from data_designer.slurm.state.outputs import ( + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + CollectionPlan, + CollectionShard, + 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.scheduler import ( + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, +) +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_manifest, + validate_shard_set, + validate_shard_winner, +) + +__all__ = [ + "ArtifactReference", + "AttemptLifecycleState", + "AttemptManifest", + "AttemptId", + "AttemptReadiness", + "AttemptTerminalClassification", + "CandidateOutcome", + "CandidateOutputFile", + "CandidateOutputManifest", + "CollectionPlan", + "CollectionShard", + "ContractRecord", + "ContractValue", + "DeploymentReadiness", + "EffectiveAttemptState", + "EndpointPublicationState", + "Identifier", + "ModelAlias", + "ProbeEvidence", + "ProbeOutcome", + "ReadinessState", + "ReasonCode", + "RecordRange", + "RunManifest", + "ResumeWorkspace", + "SchedulerIdentity", + "SchedulerObservation", + "SchedulerState", + "Sha256Digest", + "ShardManifest", + "ShardId", + "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", +] 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..e9825bd95 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -0,0 +1,58 @@ +# 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 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 + + +__all__ = [ + "ArtifactReference", + "Identifier", + "SchedulerIdentity", + "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/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/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py new file mode 100644 index 000000000..756e00b44 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -0,0 +1,138 @@ +# 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, +) + + +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, ...] + 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 complete candidate may be published as the shard winner.""" + return self.outcome is CandidateOutcome.COMPLETE + + @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 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) + host_destination: str + container_destination: str + num_partitions: PositiveInt + overwrite: Literal[False] = False + + _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) + _destinations_are_safe = field_validator("host_destination", "container_destination")(validate_absolute_path) + + @model_validator(mode="after") + def validate_shards(self) -> CollectionPlan: + 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/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py new file mode 100644 index 000000000..a7f538165 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.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 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, ModelAlias, 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" + 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: ModelAlias + 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: + if self.ready_backends > self.expected_backends: + raise ValueError("ready_backends must not exceed expected_backends") + if self.endpoint_publication is EndpointPublicationState.FAILED and self.state not in { + ReadinessState.FAILED, + ReadinessState.STOPPED, + }: + raise ValueError("failed endpoint publication requires a failed or stopped deployment") + + if self.state is ReadinessState.PENDING: + if self.ready_backends != 0: + raise ValueError("pending deployments cannot have ready backends") + if self.endpoint_publication is not EndpointPublicationState.PENDING: + raise ValueError("pending deployments require pending endpoint publication") + elif self.state is ReadinessState.STARTING: + if ( + self.ready_backends == self.expected_backends + and self.endpoint_publication is EndpointPublicationState.PUBLISHED + ): + raise ValueError("fully ready published deployments must use the ready state") + elif self.state is ReadinessState.READY: + if self.ready_backends != self.expected_backends: + raise ValueError("ready deployments require every expected backend") + if self.endpoint_publication is not EndpointPublicationState.PUBLISHED: + raise ValueError("ready deployments require a published endpoint") + elif self.state is ReadinessState.STOPPED and self.ready_backends != 0: + raise ValueError("stopped deployments cannot have ready backends") + 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 + state: ReadinessState + deployments: tuple[DeploymentReadiness, ...] = Field(min_length=1) + + _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + + @model_validator(mode="after") + def validate_deployments(self) -> AttemptReadiness: + deployment_ids = [deployment.deployment_id for deployment in self.deployments] + model_aliases = [deployment.model_alias for deployment in self.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") + if any( + deployment.last_probe is not None and deployment.last_probe.observed_at > self.updated_at + for deployment in self.deployments + ): + raise ValueError("probe observations must not be later than the readiness snapshot") + + deployment_states = tuple(deployment.state for deployment in self.deployments) + if self.state is ReadinessState.PENDING: + if any(state is not ReadinessState.PENDING for state in deployment_states): + raise ValueError("pending attempts require every deployment to be pending") + elif self.state is ReadinessState.STARTING: + if any(state in {ReadinessState.FAILED, ReadinessState.STOPPED} for state in deployment_states): + raise ValueError("starting attempts cannot contain failed or stopped deployments") + if all(state is ReadinessState.READY for state in deployment_states): + raise ValueError("attempts with every deployment ready must use the ready state") + elif self.state is ReadinessState.READY: + if any(state is not ReadinessState.READY for state in deployment_states): + raise ValueError("ready attempts require every deployment to be ready") + elif self.state is ReadinessState.FAILED: + if not any(state is ReadinessState.FAILED for state in deployment_states): + raise ValueError("failed attempts require at least one failed deployment") + elif any(state is not ReadinessState.STOPPED for state in deployment_states): + raise ValueError("stopped attempts require every deployment to be stopped") + return self 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..aacfacc65 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -0,0 +1,170 @@ +# 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, +) +from data_designer.slurm.state.validation import StateContractError + +_ALLOWED_READINESS_TRANSITIONS: dict[ReadinessState, frozenset[ReadinessState]] = { + ReadinessState.PENDING: frozenset( + {ReadinessState.PENDING, ReadinessState.STARTING, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.STARTING: frozenset( + {ReadinessState.STARTING, ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED} + ), + ReadinessState.READY: frozenset({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}), +} + +_SCHEDULER_FAILURE_STATES = frozenset( + { + SchedulerState.FAILED, + SchedulerState.CANCELLED, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.PREEMPTED, + SchedulerState.REQUEUED, + SchedulerState.OUT_OF_MEMORY, + } +) + + +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.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", + ) + + 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", + ) + _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: + 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, + 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") + _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") + _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") + + if scheduler.state in _SCHEDULER_FAILURE_STATES: + 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.state is ReadinessState.FAILED: + return EffectiveAttemptState.FAILED + if scheduler.state is SchedulerState.PENDING: + 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/scheduler.py b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py new file mode 100644 index 000000000..3245c75fa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.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 +from enum import Enum + +from pydantic import field_validator, model_validator + +from data_designer.slurm.state.base import ( + SchedulerIdentity, + 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: SchedulerIdentity + 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") + if self.reconciliation_deadline < self.observed_at: + raise ValueError("reconciliation deadline must not precede the observation") + elif self.reconciliation_deadline is not None: + raise ValueError("only accounting lag may have a reconciliation deadline") + return self + + +class EffectiveAttemptState(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ACCOUNTING_LAG = "accounting_lag" + UNKNOWN = "unknown" 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..30fdbc066 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -0,0 +1,267 @@ +# 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 + +_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 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}") + validate_attempt_manifest(run, shard, attempt) + + 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_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 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", + ) + 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/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..1d528f0d1 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -0,0 +1,23 @@ +{ + "collection_id": "collection-0001", + "container_destination": "/mnt/data/run-0001.parquet", + "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" + } + } + ], + "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_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/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_golden_records.py b/packages/data-designer-slurm/tests/state/test_golden_records.py new file mode 100644 index 000000000..fa31859db --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_golden_records.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 + +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" +GOLDEN_MODELS: tuple[tuple[str, type[StateRecord]], ...] = ( + ("run_manifest.json", RunManifest), + ("shard_manifest.json", ShardManifest), + ("successful_attempt.json", AttemptManifest), + ("single_node_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), +) + + +@pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) +def test_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() 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..0babd1f76 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_model_edges.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 json +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from data_designer.slurm.state import ( + 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)) + + +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)) + + +@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_records.py b/packages/data-designer-slurm/tests/state/test_records.py new file mode 100644 index 000000000..2eac315bf --- /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_partial_candidate_is_never_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 not 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_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py new file mode 100644 index 000000000..037785899 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -0,0 +1,650 @@ +# 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 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_partial_and_existing_winners_are_rejected() -> 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) + with pytest.raises(StateContractError, match="winner policy"): + validate_shard_winner(run, shard, partial_attempt, partial, 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) + + 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.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_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py new file mode 100644 index 000000000..3319d16e2 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -0,0 +1,65 @@ +# 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.contracts import ( + ArtifactReference as ContractArtifactReference, +) +from data_designer.slurm.contracts import ( + ContractRecord, + ContractValue, + ResumeWorkspace, + canonical_json, + compute_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 ( + StateRecord, + StateValue, +) + + +def test_state_exports_exact_shared_contract_types() -> None: + assert StateArtifactReference is ContractArtifactReference + assert StateRecordRange is ContractRecordRange + 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_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/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index ff6fd6b6a..d244200d8 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -125,6 +125,14 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} +from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference +from data_designer.slurm.contracts import RecordRange as ContractRecordRange +from data_designer.slurm.state import ArtifactReference as StateArtifactReference +from data_designer.slurm.state import RecordRange as StateRecordRange +from data_designer.slurm.state import RunManifest +assert RunManifest.__name__ == "RunManifest" +assert StateArtifactReference is ContractArtifactReference +assert StateRecordRange is ContractRecordRange """ run([str(python), "-c", statement], cwd=cwd) @@ -170,8 +178,10 @@ def main() -> None: base_leaf_requirement = requirement(base_metadata, "data-designer-slurm") leaf_base_requirement = requirement(leaf_metadata, "data-designer") + leaf_pydantic_requirement = requirement(leaf_metadata, "pydantic") assert str(base_leaf_requirement.specifier) == f"=={version}" assert str(leaf_base_requirement.specifier) == f"=={version}" + assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").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": ""}) diff --git a/uv.lock b/uv.lock index 52329761e..136ce5661 100644 --- a/uv.lock +++ b/uv.lock @@ -952,10 +952,14 @@ name = "data-designer-slurm" source = { editable = "packages/data-designer-slurm" } dependencies = [ { name = "data-designer" }, + { name = "pydantic" }, ] [package.metadata] -requires-dist = [{ name = "data-designer", editable = "packages/data-designer" }] +requires-dist = [ + { name = "data-designer", editable = "packages/data-designer" }, + { name = "pydantic", specifier = ">=2.9.2,<3" }, +] [[package]] name = "data-designer-workspace" From 810daa2006e21fe3e23fe72130fd5de097ffe31e Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:19:16 -0300 Subject: [PATCH 05/62] feat: define Slurm config and plan contracts (#879) * feat: define Slurm config and plan contracts Add strict authored configuration, profile, image, client, benchmark, dependency-lock, and resolved-plan records. Validate cross-record identities and digests with sanitized single-node and multi-node golden fixtures. Closes #873 Signed-off-by: Andre Manoel * fix: tighten Slurm contract validation Signed-off-by: Andre Manoel * fix: align Slurm F3 contract seam * fix: bind client dataset to effective resume mode * fix: harden Slurm shared contracts Prevent secret material from entering persisted configuration and make nested contract collections immutable. Tighten client and benchmark semantic invariants with focused negative coverage. Signed-off-by: Nabin Mulepati * fix: address Slurm contract review findings Signed-off-by: Andre Manoel * fix: close Slurm argument validation bypasses Signed-off-by: Andre Manoel * fix: close remaining Slurm contract gaps Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel Signed-off-by: Nabin Mulepati Co-authored-by: Nabin Mulepati --- packages/data-designer-slurm/pyproject.toml | 1 + .../data_designer/slurm/benchmark/__init__.py | 26 + .../data_designer/slurm/benchmark/records.py | 162 ++++++ .../data_designer/slurm/client/__init__.py | 10 + .../src/data_designer/slurm/client/records.py | 108 ++++ .../data_designer/slurm/config/__init__.py | 106 ++++ .../data_designer/slurm/config/benchmark.py | 110 ++++ .../src/data_designer/slurm/config/images.py | 119 ++++ .../data_designer/slurm/config/profiles.py | 275 +++++++++ .../src/data_designer/slurm/config/run.py | 480 ++++++++++++++++ .../src/data_designer/slurm/contracts.py | 120 +++- .../data_designer/slurm/planning/__init__.py | 47 ++ .../data_designer/slurm/planning/models.py | 522 ++++++++++++++++++ .../slurm/planning/validation.py | 139 +++++ .../tests/contracts/conftest.py | 48 ++ .../tests/contracts/golden/authored_run.json | 148 +++++ .../contracts/golden/authored_run_single.json | 96 ++++ .../contracts/golden/benchmark_config.json | 46 ++ .../contracts/golden/benchmark_manifest.json | 26 + .../contracts/golden/benchmark_report.json | 58 ++ .../golden/client_image_inspection.json | 23 + .../tests/contracts/golden/client_result.json | 20 + .../contracts/golden/dependency_lock.json | 31 ++ .../golden/dependency_lock_single.json | 16 + .../contracts/golden/multi_node_plan.json | 445 +++++++++++++++ .../contracts/golden/profile_catalog.json | 49 ++ .../golden/serving_image_inspection.json | 11 + .../contracts/golden/single_node_plan.json | 276 +++++++++ .../tests/contracts/test_config_records.py | 440 +++++++++++++++ .../tests/contracts/test_golden_records.py | 69 +++ .../tests/contracts/test_planning_records.py | 490 ++++++++++++++++ .../tests/contracts/test_profiles.py | 126 +++++ .../tests/contracts/test_shared_records.py | 283 ++++++++++ ...ecords.py => test_state_golden_records.py} | 2 +- .../tests/test_contracts.py | 4 + scripts/test_slurm_package_install.py | 11 + uv.lock | 2 + 37 files changed, 4941 insertions(+), 4 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/images.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/run.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/models.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py create mode 100644 packages/data-designer-slurm/tests/contracts/conftest.py create mode 100644 packages/data-designer-slurm/tests/contracts/golden/authored_run.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_config.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_manifest.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/benchmark_report.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/client_result.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/serving_image_inspection.json create mode 100644 packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json create mode 100644 packages/data-designer-slurm/tests/contracts/test_config_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_golden_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_planning_records.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_profiles.py create mode 100644 packages/data-designer-slurm/tests/contracts/test_shared_records.py rename packages/data-designer-slurm/tests/state/{test_golden_records.py => test_state_golden_records.py} (95%) diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 0fd4ec5ff..340c8f289 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -37,6 +37,7 @@ bump = true [tool.hatch.metadata.hooks.uv-dynamic-versioning] dependencies = [ "data-designer=={{ version }}", + "packaging>=25,<27", "pydantic>=2.9.2,<3", ] 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..40acb8f5b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py @@ -0,0 +1,26 @@ +# 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.records import ( + BenchmarkCaseResult, + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkRecommendation, + BenchmarkRecommendationKind, + BenchmarkReport, +) + +__all__ = [ + "BenchmarkCaseResult", + "BenchmarkChildRun", + "BenchmarkManifest", + "BenchmarkOutcome", + "BenchmarkRecommendation", + "BenchmarkRecommendationKind", + "BenchmarkReport", +] 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..11c797527 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -0,0 +1,162 @@ +# 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" + + +class BenchmarkCaseResult(ContractValue): + case_id: Identifier + child_run_id: Identifier + outcome: BenchmarkOutcome + topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + 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 + nodes_per_job: PositiveInt + 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.gpu_hours_per_job, + self.total_gpu_hours, + self.target_jobs, + self.feasible, + ) + if self.outcome is BenchmarkOutcome.SUCCEEDED and any(value is None for value in required): + raise ValueError("successful benchmark cases require complete timing and feasibility metrics") + if self.outcome is BenchmarkOutcome.SUCCEEDED: + if self.actual_records != self.requested_records: + raise ValueError("successful benchmark cases require the requested record count") + if self.generation_seconds == 0 or self.wall_seconds == 0 or self.rows_per_second == 0: + raise ValueError("successful benchmark generation, wall time, and throughput must be positive") + 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/client/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py new file mode 100644 index 000000000..96fe5effa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py @@ -0,0 +1,10 @@ +# 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 ClientOutcome, ClientResult + +__all__ = ["ClientOutcome", "ClientResult"] 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..a76290cfc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -0,0 +1,108 @@ +# 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, Literal + +from pydantic import NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator + +from data_designer.slurm.contracts import ( + ArtifactReference, + AttemptId, + ContractRecord, + Identifier, + ShardId, + validate_absolute_path, +) + + +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 != "if_possible" and self.effective_resume_mode not in { + None, + self.requested_resume_mode, + }: + raise ValueError("effective resume mode must match a fixed requested mode") + 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/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py new file mode 100644 index 000000000..5cda08684 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -0,0 +1,106 @@ +# 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.images import ( + ClientImageInspection, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, + ServingImageInspection, +) +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, + InvocationDiagnostics, + LiteralEnvironmentBinding, + LocalStdioMCPProviderConfig, + OutputConfig, + QueueBackpressureConfig, + RemoteMCPProviderConfig, + SecretRef, + ServerDeploymentConfig, + SubmissionConfig, + VllmServerConfig, +) + +__all__ = [ + "AdaptiveRecordPolicy", + "ArrayTasksConfig", + "BenchmarkAnalysisTargets", + "BenchmarkBaseRun", + "BenchmarkDeploymentCase", + "BenchmarkDeploymentOverride", + "BuilderInput", + "ClientConfig", + "ClientDependencies", + "ClientImageInspection", + "ContainerMount", + "DataDesignerSlurmBenchmarkConfig", + "DataDesignerSlurmConfig", + "DeploymentResources", + "DeploymentTopology", + "FixedRecordPolicy", + "GpuRequestMode", + "ImageBuildProfile", + "ImageBuildRequest", + "ImageInspectionRecord", + "ImageKind", + "ImageRef", + "InputBindings", + "InstalledDistribution", + "InvocationConfig", + "InvocationDiagnostics", + "LiteralEnvironmentBinding", + "LocalStdioMCPProviderConfig", + "OutputConfig", + "ProfileSelectionSource", + "QueueBackpressureConfig", + "RemoteMCPProviderConfig", + "SchedulerProfile", + "SecretRef", + "SelectedSlurmProfile", + "ServerDeploymentConfig", + "ServingImageInspection", + "SlurmProfile", + "SlurmProfileCatalog", + "SubmissionConfig", + "VllmServerConfig", + "injected_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..e7c465c2e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -0,0 +1,110 @@ +# 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, + ModelAlias, + 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[ModelAlias, BenchmarkDeploymentOverride] = Field(min_length=1) + + +class FixedRecordPolicy(AuthoredConfig): + type: Literal["fixed"] + records: PositiveInt + + +class AdaptiveRecordPolicy(AuthoredConfig): + type: Literal["adaptive"] + base_records: PositiveInt + max_records: PositiveInt + records_per_concurrency: PositiveFloat + + @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[ModelAlias] + concurrency_values: list[PositiveInt] = Field(min_length=1) + deployment_cases: list[BenchmarkDeploymentCase] = Field(min_length=1) + record_policy: BenchmarkRecordPolicy + 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/images.py b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py new file mode 100644 index 000000000..0b4e983de --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/images.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 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, + Sha256Digest, + validate_absolute_path, + validate_plain_text, +) + +DistributionName = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[a-z0-9]+(?:-[a-z0-9]+)*$", + ), +] + + +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 InstalledDistribution(ContractValue): + name: DistributionName + version: Annotated[str, StringConstraints(min_length=1, max_length=128)] + + +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/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py new file mode 100644 index 000000000..87412d84d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -0,0 +1,275 @@ +# 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_sha256, + validate_absolute_path, + validate_plain_text, +) + + +class GpuRequestMode(str, Enum): + GRES = "gres" + VISIBLE = "visible" + + +class SchedulerProfile(AuthoredConfig): + account: Identifier | None = None + partition: Identifier | None = None + mem_per_gpu: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] | None = None + + +class ImageBuildProfile(AuthoredConfig): + partition: Identifier + + +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") + 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_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_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_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..c6f7d69d4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath +import re +from collections.abc import Mapping +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.images import ImageRef +from data_designer.slurm.contracts import ( + AuthoredConfig, + Duration, + EnvironmentName, + Identifier, + ModelAlias, + SchemaVersion, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, + validate_url, +) + +_OWNED_VLLM_FLAGS = { + "--api-key", + "--distributed-executor-backend", + "--distributed-init-address", + "--enable-expert-parallel", + "--headless", + "--host", + "--middleware", + "--model", + "--pipeline-parallel-size", + "--port", + "--served-model-name", + "--tensor-parallel-size", +} +_DURATION_FACTORS = {"s": 1, "m": 60, "h": 3600, "d": 86400} +_NON_SECRET_PAYLOAD_KEYS = frozenset({"idempotency_key", "partition_key", "primary_key", "sort_key"}) +_SECRET_NAME_PARTS = frozenset({"credential", "credentials", "password", "secret", "token"}) + + +def _duration_seconds(value: Duration) -> int: + return int(value[:-1]) * _DURATION_FACTORS[value[-1]] + + +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 _is_secret_name(value: str) -> bool: + segments = _secret_name_segments(value) + return bool( + _SECRET_NAME_PARTS.intersection(segments) + or {"access", "key"}.issubset(segments) + or segments[-1] in {"auth", "key"} + ) + + +def _is_secret_payload_name(value: str) -> bool: + segments = _secret_name_segments(value) + normalized = "_".join(segments) + return normalized not in _NON_SECRET_PAYLOAD_KEYS and _is_secret_name(value) + + +def _option_flag(value: str) -> str: + return re.split(r"[=\s]", value.lstrip(), maxsplit=1)[0] + + +def _contains_secret_key(value: object) -> bool: + if isinstance(value, Mapping): + return any( + (_is_secret_payload_name(str(key)) and item is not None) or _contains_secret_key(item) + for key, item in value.items() + ) + if isinstance(value, list | tuple): + return any(_contains_secret_key(item) for item in value) + return False + + +def _validate_environment_bindings( + values: dict[EnvironmentName, EnvironmentBinding], +) -> dict[EnvironmentName, EnvironmentBinding]: + literal_secrets = [ + name + for name, binding in values.items() + if _is_secret_name(name) and isinstance(binding, LiteralEnvironmentBinding) + ] + if literal_secrets: + raise ValueError("secret-shaped environment names require external secret references") + return values + + +class LiteralEnvironmentBinding(AuthoredConfig): + 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): + type: Literal["secret"] + environment: EnvironmentName + + +EnvironmentBinding = Annotated[ + LiteralEnvironmentBinding | SecretRef, + Field(discriminator="type"), +] + + +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") + retired = {"dependencies", "sandbox_config", "server_configs"}.intersection(self.inline) + if retired: + raise ValueError(f"builder input contains retired Big Iron 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") + if _contains_secret_key(self.inline): + raise ValueError("inline builder input must not contain secret values") + 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 = _option_flag(value).lstrip("-") + if _is_secret_name(option): + raise ValueError("secret-shaped MCP arguments must use an environment secret reference") + return values + + _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) + + +MCPProviderConfig = Annotated[ + RemoteMCPProviderConfig | LocalStdioMCPProviderConfig, + Field(discriminator="provider_type"), +] + + +class InvocationDiagnostics(AuthoredConfig): + log_requests: bool = False + + +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[ModelAlias, PositiveInt] = Field(default_factory=dict) + diagnostics: InvocationDiagnostics = Field(default_factory=InvocationDiagnostics) + + @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) + 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(f"dependency requirement must identify a package or immutable wheel: {value!r}") + try: + requirement = Requirement(value) + except InvalidRequirement as error: + raise ValueError(f"invalid dependency requirement: {value!r}") 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 = 32 + image: ImageRef + dependencies: ClientDependencies = Field(default_factory=ClientDependencies) + + +class QueueBackpressureConfig(AuthoredConfig): + max_waiting_requests: NonNegativeInt = 128 + retry_after_seconds: NonNegativeInt | None = 1 + + +class VllmServerConfig(AuthoredConfig): + type: Literal["vllm"] + image: ImageRef + startup_timeout: Duration = "15m" + distributed_init_timeout: Duration = "10m" + readiness_path: str = "/health" + enable_expert_parallel: bool = False + 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: + validate_plain_text(value, field_name="readiness path") + if not value.startswith("/") or "?" in value or "#" in value: + raise ValueError("readiness_path must be an absolute URL path without query or fragment") + return value + + @field_validator("extra_args") + @classmethod + def validate_extra_args(cls, values: list[str]) -> list[str]: + for value in values: + validate_plain_text(value, field_name="vLLM argument") + flag = _option_flag(value) + if flag in _OWNED_VLLM_FLAGS: + raise ValueError(f"vLLM argument {flag!r} is owned by the compiler or runtime") + if _is_secret_name(flag.lstrip("-")): + raise ValueError("secret-shaped vLLM arguments must use an environment secret reference") + return values + + _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) + + @model_validator(mode="after") + def validate_timeouts(self) -> VllmServerConfig: + if _duration_seconds(self.distributed_init_timeout) > _duration_seconds(self.startup_timeout): + raise ValueError("distributed_init_timeout must not exceed startup_timeout") + return self + + +class DeploymentResources(AuthoredConfig): + nodes: PositiveInt = 1 + + +class DeploymentTopology(AuthoredConfig): + tensor_parallel: PositiveInt = 1 + nodes_per_replica: PositiveInt = 1 + + +class ServerDeploymentConfig(AuthoredConfig): + model_alias: ModelAlias + 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.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 = 1 + + @model_validator(mode="after") + def validate_concurrency(self) -> ArrayTasksConfig: + if 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: Identifier | 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 OutputConfig(AuthoredConfig): + root: str | None = None + format: Literal["parquet", "jsonl", "csv"] = "parquet" + partitions: PositiveInt = 1 + require_exact_record_count: bool = False + + @field_validator("root") + @classmethod + def validate_root(cls, value: str | None) -> str | None: + return None if value is None else validate_absolute_path(value) + + +class DataDesignerSlurmConfig(AuthoredConfig): + """Complete portable intent for one Data Designer Slurm run.""" + + schema_version: SchemaVersion + name: Identifier + builder: BuilderInput + invocation: InvocationConfig + client: ClientConfig + deployments: list[ServerDeploymentConfig] = Field(min_length=1) + array_tasks: ArrayTasksConfig = Field(default_factory=ArrayTasksConfig) + submission: SubmissionConfig = Field(default_factory=SubmissionConfig) + output: OutputConfig = Field(default_factory=OutputConfig) + + @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") + 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/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index 087b2a678..f07063417 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -8,7 +8,9 @@ import hashlib import json import posixpath -from typing import Annotated, Literal +from collections.abc import Mapping +from typing import Annotated, Literal, TypeVar +from urllib.parse import urlsplit from pydantic import ( BaseModel, @@ -32,7 +34,60 @@ 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)$")] + +_Key = TypeVar("_Key") +_Value = TypeVar("_Value") + + +class _FrozenList(list[_Value]): + """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[_Key, _Value]): + """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): @@ -47,6 +102,24 @@ class ContractValue(BaseModel): validate_default=True, ) + @field_validator("*", mode="after") + @classmethod + def freeze_collections(cls, value: object) -> object: + return _freeze_collections(value) + + +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: + return hashlib.sha256(self.serialize_json().encode("utf-8")).hexdigest() + class ContractRecord(ContractValue): """Base for immutable, explicitly versioned Slurm records.""" @@ -104,8 +177,7 @@ def validate_absolute_path(value: str) -> str: raise ValueError("path must have exactly one leading slash") if value == "/": raise ValueError("path must not be the filesystem root") - if any(ord(character) < 32 or ord(character) == 127 for character in value): - raise ValueError("path must not contain control characters") + 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: @@ -126,6 +198,42 @@ def validate_relative_path(value: str) -> str: return value +def validate_local_config_path(value: str) -> str: + 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: + 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_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.""" @@ -163,8 +271,11 @@ class ResumeWorkspace(ContractValue): __all__ = [ "ArtifactReference", "AttemptId", + "AuthoredConfig", "ContractRecord", "ContractValue", + "Duration", + "EnvironmentName", "Identifier", "ModelAlias", "RecordRange", @@ -176,5 +287,8 @@ class ResumeWorkspace(ContractValue): "compute_sha256", "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/planning/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py new file mode 100644 index 000000000..160b7ef77 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -0,0 +1,47 @@ +# 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, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSlurmRunPlan, + ResolvedSubmission, + ResolvedTopology, + ResumeWorkspace, +) +from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan + +__all__ = [ + "ArtifactReference", + "LockedPackage", + "PlanContractError", + "PlannedShard", + "PortClaim", + "RecordRange", + "ResolvedBuilderInput", + "ResolvedClient", + "ResolvedDependencyLock", + "ResolvedDeployment", + "ResolvedImage", + "ResolvedInvocation", + "ResolvedOutput", + "ResolvedSlurmRunPlan", + "ResolvedSubmission", + "ResolvedTopology", + "ResumeWorkspace", + "validate_resolved_plan", +] 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..d1eff7ab6 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -0,0 +1,522 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import posixpath +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.images import ( + DistributionName, + ImageInspectionRecord, + ImageKind, + ImageRef, + InstalledDistribution, +) +from data_designer.slurm.config.profiles import ContainerMount, SelectedSlurmProfile +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + ClientConfig, + ClientDependencies, + InvocationConfig, + ServerDeploymentConfig, + SubmissionConfig, +) +from data_designer.slurm.contracts import ( + ArtifactReference, + ContractRecord, + ContractValue, + Identifier, + ModelAlias, + RecordRange, + ResumeWorkspace, + Sha256Digest, + ShardId, + compute_sha256, + validate_absolute_path, + validate_local_config_path, + validate_plain_text, +) + + +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 + + +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 image inventory.""" + + resolver_version: Identifier + python_abi: Identifier + client_image_sha256: Sha256Digest + authored_requirements: tuple[str, ...] + authored_source: str | None = None + source: ArtifactReference | None = None + image_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)) + image_names = tuple(distribution.name for distribution in self.image_distributions) + overlay_names = tuple(package.name for package in self.overlay_packages) + if image_names != tuple(sorted(image_names)) or overlay_names != tuple(sorted(overlay_names)): + raise ValueError("dependency lock distributions must be sorted by normalized name") + if len(image_names) != len(set(image_names)) or len(overlay_names) != len(set(overlay_names)): + raise ValueError("dependency lock distribution names must be unique") + overlap = set(image_names).intersection(overlay_names) + if overlap: + raise ValueError(f"overlay packages overlap image-owned distributions: {', '.join(sorted(overlap))}") + image_packages = {distribution.name: distribution for distribution in self.image_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 image_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 + content_sha256: Sha256Digest + model_aliases: tuple[ModelAlias, ...] + referenced_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") + expected_digest = compute_sha256(self.inline) + model_aliases, referenced_aliases = _extract_builder_aliases(self.inline) + if self.model_aliases != model_aliases: + raise ValueError("resolved model aliases do not match the inline builder") + if self.referenced_model_aliases != referenced_aliases: + raise ValueError("resolved referenced 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 len(self.referenced_model_aliases) != len(set(self.referenced_model_aliases)): + raise ValueError("resolved builder referenced 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_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"] + node_index: NonNegativeInt + port: Annotated[int, Field(ge=1024, le=65535)] + + +class ResolvedTopology(ContractValue): + 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 + + +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 + 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") + if self.gpus_per_node % self.authored.topology.tensor_parallel: + raise ValueError("tensor_parallel must divide resolved GPUs per node") + expected = ResolvedTopology( + tensor_parallel=self.authored.topology.tensor_parallel, + nodes_per_replica=self.authored.topology.nodes_per_replica, + pipeline_parallel=self.authored.topology.nodes_per_replica, + node_group_count=self.authored.resources.nodes // self.authored.topology.nodes_per_replica, + replicas_per_node_group=self.gpus_per_node // self.authored.topology.tensor_parallel, + replica_count=(self.authored.resources.nodes // self.authored.topology.nodes_per_replica) + * (self.gpus_per_node // self.authored.topology.tensor_parallel), + gpus_per_replica=self.authored.topology.tensor_parallel * 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 + image: ResolvedImage + dependency_lock: ArtifactReference + host_node_index: NonNegativeInt + gpu_count: Literal[0] + ports: tuple[PortClaim, ...] = () + + @model_validator(mode="after") + def validate_client(self) -> ResolvedClient: + if self.image.kind is not ImageKind.CLIENT: + raise ValueError("Data Designer client requires a client image") + if self.image.authored_ref != self.authored.image: + raise ValueError("resolved client image does not match the authored image reference") + 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: Identifier | 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 ResolvedOutput(ContractValue): + root: str + format: Literal["parquet", "jsonl", "csv"] + partitions: PositiveInt + require_exact_record_count: bool + + _root_is_absolute = field_validator("root")(validate_absolute_path) + + +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 + 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 + output: ResolvedOutput + container_mounts: tuple[ContainerMount, ...] = () + runtime_bundle: ArtifactReference + + @model_validator(mode="after") + def validate_plan(self) -> ResolvedSlurmRunPlan: + profile = self.selected_profile.profile + 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 any(deployment.gpus_per_node != self.resolved_gpus_per_node for deployment in self.deployments): + raise ValueError("every deployment must use the resolved profile GPU count") + 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 not set(aliases).issubset(self.builder.model_aliases): + raise ValueError("each deployment alias must match a resolved Data Designer model alias") + if not set(self.builder.referenced_model_aliases).issubset(aliases): + raise ValueError("each referenced Data Designer model alias requires a deployment") + + node_indices = tuple(index for deployment in self.deployments for index in deployment.node_indices) + if node_indices != tuple(range(len(node_indices))): + raise ValueError("deployment nodes must be disjoint and contiguous in authored order") + 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") + 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 != RunConfig().non_inference_max_parallel_workers: + raise ValueError("default non-inference worker count must match the Data Designer RunConfig 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") + self._validate_shards(run_root) + if not _is_below(self.output.root, profile.workspace_root): + raise ValueError("resolved output root must be below the selected workspace_root") + shards_root = posixpath.join(run_root, "shards") + if ( + self.output.root == shards_root + or _is_below(self.output.root, shards_root) + or _is_below(shards_root, self.output.root) + ): + raise ValueError("resolved output root must not overlap the run shard workspace") + 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.authored.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 _is_below(path: str, root: str) -> bool: + return path != root and posixpath.commonpath((path, root)) == root + + +def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...]]: + 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"]) + + referenced_aliases: list[ModelAlias] = [] + + def collect(value: JsonValue, *, key: str | None = None) -> None: + if key == "model_configs": + return + if key == "model_alias" or (key is not None and key.endswith("_model_alias")): + if not isinstance(value, str): + raise ValueError(f"builder {key} must be a string") + referenced_aliases.append(value) + return + if key == "model_aliases": + if not isinstance(value, list) or any(not isinstance(alias, str) for alias in value): + raise ValueError("builder model_aliases must be a list of strings") + referenced_aliases.extend(value) + return + if isinstance(value, dict): + for child_key, child in value.items(): + collect(child, key=child_key) + elif isinstance(value, list): + for child in value: + collect(child) + + collect(data_designer) + return tuple(model_aliases), tuple(dict.fromkeys(referenced_aliases)) 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..c47ebc512 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pydantic import JsonValue + +from data_designer.config import RunConfig +from data_designer.slurm.config.images import ClientImageInspection +from data_designer.slurm.config.run import DataDesignerSlurmConfig +from data_designer.slurm.planning.models import ( + ResolvedDependencyLock, + ResolvedSlurmRunPlan, + _extract_builder_aliases, +) + + +class PlanContractError(ValueError): + """Raised when a resolved plan does not match its authored inputs.""" + + +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") + 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( + plan.builder.inline == authored.builder.inline, "resolved inline builder does not match authored input" + ) + else: + _require( + plan.builder.authored_source == authored.builder.source, + "resolved builder source does not match authored input", + ) + if builder_payload is None: + raise PlanContractError("sourced builder validation requires its resolved payload") + model_aliases, referenced_aliases = _extract_builder_aliases(builder_payload) + _require(plan.builder.model_aliases == model_aliases, "resolved model aliases do not match builder source") + _require( + plan.builder.referenced_model_aliases == referenced_aliases, + "resolved referenced aliases do 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.output.format == authored.output.format, "resolved output format does not match authored input") + _require( + plan.output.partitions == authored.output.partitions, "resolved output partitions do not match authored input" + ) + _require( + plan.output.require_exact_record_count == authored.output.require_exact_record_count, + "resolved exact-record policy does not match authored input", + ) + if authored.output.root is not None: + _require(plan.output.root == authored.output.root, "resolved output root 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_image_sha256 == plan.client.image.sha256, + "dependency lock client image digest does not match the resolved client image", + ) + inspection = plan.client.image.inspection.inspection + _require(isinstance(inspection, ClientImageInspection), "resolved client image lacks client inspection facts") + _require( + dependency_lock.python_abi == inspection.python_abi, "dependency lock Python ABI does not match client image" + ) + _require( + dependency_lock.image_distributions == inspection.distributions, + "dependency lock image inventory does not match client image 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 PlanContractError(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/tests/contracts/conftest.py b/packages/data-designer-slurm/tests/contracts/conftest.py new file mode 100644 index 000000000..44a74d99c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/conftest.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 + +from pathlib import Path + +import pytest + +from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan + +GOLDEN_DIR = Path(__file__).parent / "golden" + + +@pytest.fixture +def authored_run() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run.json").read_text()) + + +@pytest.fixture +def authored_run_single() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run_single.json").read_text()) + + +@pytest.fixture +def profile_catalog() -> SlurmProfileCatalog: + return SlurmProfileCatalog.model_validate_json((GOLDEN_DIR / "profile_catalog.json").read_text()) + + +@pytest.fixture +def dependency_lock() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock.json").read_text()) + + +@pytest.fixture +def dependency_lock_single() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock_single.json").read_text()) + + +@pytest.fixture +def single_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "single_node_plan.json").read_text()) + + +@pytest.fixture +def multi_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "multi_node_plan.json").read_text()) diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json new file mode 100644 index 000000000..28f65680a --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -0,0 +1,148 @@ +{ + "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" + ] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "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 + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "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" + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 1 + } + } + ], + "invocation": { + "dataset_name": "training-data", + "diagnostics": { + "log_requests": false + }, + "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", + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": null + }, + "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..c6a46897a --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -0,0 +1,96 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, + "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": [] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "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 + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "readiness_path": "/health", + "startup_timeout": "15m", + "type": "vllm" + }, + "topology": { + "nodes_per_replica": 1, + "tensor_parallel": 8 + } + } + ], + "invocation": { + "dataset_name": "single-node", + "diagnostics": { + "log_requests": false + }, + "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", + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": null + }, + "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..6466182d7 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -0,0 +1,23 @@ +{ + "inspection": { + "distributions": [ + { + "name": "data-designer", + "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..87190b04c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -0,0 +1,31 @@ +{ + "authored_requirements": [ + "data-designer-speech==0.2.0" + ], + "authored_source": null, + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "image_distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + }, + { + "name": "pip", + "version": "26.1" + } + ], + "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..b86f29d6c --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -0,0 +1,16 @@ +{ + "authored_requirements": [], + "authored_source": null, + "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", + "image_distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + } + ], + "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..0f2fd69b1 --- /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": "9bfdbcb1ddd374a5d43eea44890d13c08798ea9e18c30ff5d5258fa0d27f539d" + }, + "builder": { + "authored_source": null, + "content_sha256": "f37227ca7c67e203abe881c0228b4308a8e741364296d293159a1201949732f2", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + }, + { + "alias": "judge", + "model": "/models/judge", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator", + "judge" + ], + "referenced_model_aliases": [], + "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" + ] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-001/dependency-lock.json", + "sha256": "a86032b310aa6bdb95fc7f35ef606fec2e56b977ba60e36b8b9293341ac43e00" + }, + "gpu_count": 0, + "host_node_index": 0, + "image": { + "authored_ref": { + "name": "dd-client-0.9", + "path": null + }, + "inspection": { + "inspection": { + "distributions": [ + { + "name": "data-designer", + "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" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "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" + } + ] + }, + "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 + }, + "queue_backpressure": { + "max_waiting_requests": 1024, + "retry_after_seconds": 2 + }, + "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" + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "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", + "diagnostics": { + "log_requests": false + }, + "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_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 + } + }, + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": "/workspace/primary/runs/run-001/output" + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-001", + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "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": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "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..8d500c664 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -0,0 +1,49 @@ +{ + "clusters": { + "lab": { + "container_mounts": [], + "gpu_request_mode": "visible", + "gpus_per_node": "auto", + "host_patterns": [ + "lab-login-*" + ], + "image_build": { + "partition": "cpu" + }, + "scheduler": { + "account": "lab", + "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": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "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..f0794cc75 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -0,0 +1,276 @@ +{ + "array_tasks": { + "count": 1, + "max_concurrent": 1 + }, + "authored_config": { + "path": "/workspace/primary/runs/run-single/authored-config.json", + "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + }, + "builder": { + "authored_source": null, + "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", + "inline": { + "data_designer": { + "columns": [], + "model_configs": [ + { + "alias": "generator", + "model": "example/generator", + "provider": "openai" + } + ] + } + }, + "model_aliases": [ + "generator" + ], + "referenced_model_aliases": [], + "source": null + }, + "client": { + "authored": { + "cpus": 32, + "dependencies": { + "index_credentials": {}, + "lock_file": null, + "requirements": [] + }, + "image": { + "name": "dd-client-0.9", + "path": null + } + }, + "dependency_lock": { + "path": "/workspace/primary/runs/run-single/dependency-lock.json", + "sha256": "2600c1944897a8e3c9a9410fe912bdee0e169e09769e58180cadebe94f3da52c" + }, + "gpu_count": 0, + "host_node_index": 0, + "image": { + "authored_ref": { + "name": "dd-client-0.9", + "path": null + }, + "inspection": { + "inspection": { + "distributions": [ + { + "name": "data-designer", + "version": "0.9.2" + } + ], + "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" + }, + "path": "/images/dd-client-0.9.sqsh", + "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + }, + "ports": [ + { + "name": "deployment-00000-logical-endpoint", + "node_index": 0, + "port": 17000, + "role": "logical_endpoint" + } + ] + }, + "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 + }, + "queue_backpressure": { + "max_waiting_requests": 128, + "retry_after_seconds": 1 + }, + "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", + "diagnostics": { + "log_requests": false + }, + "input_bindings": { + "managed_assets_path": null, + "seed_path": null + }, + "mcp_providers": [], + "model_concurrency": { + "generator": 8 + }, + "num_records": 8, + "resume": "never", + "run_config": {} + }, + "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 + } + }, + "output": { + "format": "parquet", + "partitions": 1, + "require_exact_record_count": false, + "root": "/workspace/primary/runs/run-single/output" + }, + "package_version": "0.9.2", + "resolved_gpus_per_node": 8, + "run_id": "run-single", + "runtime_bundle": { + "path": "/workspace/primary/runtime/runtime.tar.gz", + "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" + }, + "schema_version": 1, + "selected_profile": { + "catalog_path": "/workspace/profile.json", + "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "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": { + "partition": "cpu" + }, + "scheduler": { + "account": "research", + "mem_per_gpu": null, + "partition": "batch" + }, + "schema_version": 1, + "workspace_root": "/workspace/primary" + }, + "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "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..33959d5c1 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -0,0 +1,440 @@ +# 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_secret_reference_serializes_only_external_binding() -> None: + secret = SecretRef(type="secret", environment="HF_TOKEN") + + assert secret.model_dump(mode="json") == {"type": "secret", "environment": "HF_TOKEN"} + with pytest.raises(ValidationError): + SecretRef.model_validate({"type": "secret", "environment": "HF_TOKEN", "value": "secret-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_vllm_defaults_and_backpressure_override() -> None: + server = VllmServerConfig(type="vllm", image=ImageRef(name="vllm")) + overridden = VllmServerConfig( + type="vllm", + image=ImageRef(name="vllm"), + queue_backpressure=QueueBackpressureConfig(max_waiting_requests=0, retry_after_seconds=None), + ) + + assert server.queue_backpressure.model_dump() == {"max_waiting_requests": 128, "retry_after_seconds": 1} + assert overridden.queue_backpressure.model_dump() == {"max_waiting_requests": 0, "retry_after_seconds": None} + + +@pytest.mark.parametrize( + "argument", + [ + "--api-key=plaintext-secret", + "--api-key plaintext-secret", + "--distributed-init-address", + "--host=0.0.0.0", + "--middleware", + "--model", + "--port", + "--port 9000", + " --port 9000", + "--tensor-parallel-size", + ], +) +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( + "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_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", + ) + + +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["server"]["enable_expert_parallel"] = True + 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"): + DataDesignerSlurmConfig.model_validate(payload) + + +@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_secret_values(secret_key: str) -> None: + with pytest.raises(ValidationError, match="secret values"): + BuilderInput(inline={"columns": [], secret_key: "plaintext-secret"}) + + +@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["array_tasks"]["count"] = 101 + payload["array_tasks"]["max_concurrent"] = 1 + with pytest.raises(ValidationError, match="requested records"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_small_config_values_validate_at_boundary() -> 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="readiness_path"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), readiness_path="health") + + +@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..5c64a5c24 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -0,0 +1,490 @@ +# 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_sha256 +from data_designer.slurm.planning import ( + ArtifactReference, + PlanContractError, + ResolvedDependencyLock, + ResolvedDeployment, + ResolvedSlurmRunPlan, + ResolvedSubmission, + validate_resolved_plan, +) + + +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({}) + + +@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["output"].update(root="/outside/output"), + lambda payload: payload.update(container_mounts=[]), + lambda payload: payload["deployments"][0]["topology"].update(replica_count=2), + ], +) +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_preserves_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="RunConfig 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"] = 32 + payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + assert plan.invocation.effective_run_config["non_inference_max_parallel_workers"] == 32 + + +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)) + + +@pytest.mark.parametrize( + "output_root", + [ + "/workspace/primary/runs/run-001/shards/shard-00000/dataset", + "/workspace/primary/runs/run-001", + "/workspace/primary/runs", + ], +) +def test_plan_rejects_output_overlapping_shard_workspace( + multi_node_plan: ResolvedSlurmRunPlan, + output_root: str, +) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["output"]["root"] = output_root + + with pytest.raises(ValidationError, match="must not overlap"): + 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_sha256(payload["builder"]["inline"]) + + with pytest.raises(ValidationError, match="deployment alias"): + ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + +def test_plan_rejects_referenced_alias_without_deployment(multi_node_plan: ResolvedSlurmRunPlan) -> None: + payload = multi_node_plan.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"]["columns"] = [{"model_alias": "missing"}] + payload["builder"]["referenced_model_aliases"] = ["missing"] + payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + + with pytest.raises(ValidationError, match="referenced Data Designer model 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")}) + 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.json", "sha256": "a" * 64}, + "inline": None, + "content_sha256": "a" * 64, + "model_aliases": ["generator", "judge"], + "referenced_model_aliases": [], + } + sourced_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + + with pytest.raises(PlanContractError, 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( + "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["client"]["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(image_distributions=list(reversed(payload["image_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(PlanContractError, 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(PlanContractError, 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(PlanContractError, 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(PlanContractError, 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(PlanContractError, match="Python ABI"): + validate_resolved_plan(authored_run, invalid_lock, invalid_plan) + + +def test_cross_record_validation_rejects_image_inventory( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + invalid_lock = dependency_lock.model_copy(update={"image_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(PlanContractError, match="image inventory"): + 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(PlanContractError, 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(PlanContractError, 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..1e4086f5e --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -0,0 +1,126 @@ +# 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"].update(host_patterns=["login[broken"]), + lambda payload: payload["clusters"]["primary"].update( + container_mounts=[ + {"source": "/one", "target": "/same"}, + {"source": "/two", "target": "/same"}, + ] + ), + ], +) +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) 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..5b511c4c4 --- /dev/null +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -0,0 +1,283 @@ +# 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_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", 999), + ("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_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/state/test_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py similarity index 95% rename from packages/data-designer-slurm/tests/state/test_golden_records.py rename to packages/data-designer-slurm/tests/state/test_state_golden_records.py index fa31859db..e029b99b4 100644 --- a/packages/data-designer-slurm/tests/state/test_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -37,7 +37,7 @@ @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) -def test_golden_record_round_trip_is_deterministic(filename: str, model: type[StateRecord]) -> None: +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) diff --git a/packages/data-designer-slurm/tests/test_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py index 3319d16e2..d1983b915 100644 --- a/packages/data-designer-slurm/tests/test_contracts.py +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -34,6 +34,9 @@ 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, @@ -43,6 +46,7 @@ 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 diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index d244200d8..bf2cffd28 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -127,12 +127,21 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert version("data-designer-slurm") == {version!r} 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.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.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 RunManifest assert RunManifest.__name__ == "RunManifest" +assert PlanningArtifactReference is ContractArtifactReference +assert PlanningRecordRange is ContractRecordRange +assert PlanningResumeWorkspace is ContractResumeWorkspace assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange +assert StateResumeWorkspace is ContractResumeWorkspace """ run([str(python), "-c", statement], cwd=cwd) @@ -178,9 +187,11 @@ def main() -> None: 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_pydantic_requirement = requirement(leaf_metadata, "pydantic") 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_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) diff --git a/uv.lock b/uv.lock index 136ce5661..2e6000e20 100644 --- a/uv.lock +++ b/uv.lock @@ -952,12 +952,14 @@ name = "data-designer-slurm" source = { editable = "packages/data-designer-slurm" } dependencies = [ { name = "data-designer" }, + { name = "packaging" }, { name = "pydantic" }, ] [package.metadata] requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, + { name = "packaging", specifier = ">=25,<27" }, { name = "pydantic", specifier = ">=2.9.2,<3" }, ] From 876a77e98b8fc1c759c746ebb00ce06ac3a9ced0 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 24 Aug 2026 15:33:10 -0600 Subject: [PATCH 06/62] feat: integrate Slurm plan and runtime state (#889) * feat: integrate Slurm plan and state Validate plan-aware shards, readiness, attempts, and finalization through a reusable immutable context. Add deterministic goldens and negative coverage for mismatch and failure paths. Closes #880 Signed-off-by: Nabin Mulepati * fix: reject unsubmitted planned attempts Require scheduler-bound plan validation to start only after an attempt leaves the created state. Cover created attempts that already carry scheduler metadata and verify the integration validator imports from an isolated Slurm wheel. Signed-off-by: Nabin Mulepati * test: harden Slurm integration join coverage Signed-off-by: Nabin Mulepati * refactor: simplify plan state validator Signed-off-by: Nabin Mulepati * refactor: remove validator wrapper functions Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/contracts.py | 3 + .../src/data_designer/slurm/integration.py | 280 +++++++ .../golden/finalization_chain.json | 128 ++++ .../test_integration_validation.py | 708 ++++++++++++++++++ scripts/test_slurm_package_install.py | 1 + 5 files changed, 1120 insertions(+) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/integration.py create mode 100644 packages/data-designer-slurm/tests/integration/golden/finalization_chain.json create mode 100644 packages/data-designer-slurm/tests/integration/test_integration_validation.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index f07063417..fc47926b1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -199,6 +199,7 @@ def validate_relative_path(value: str) -> str: 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") @@ -211,6 +212,7 @@ def validate_local_config_path(value: str) -> str: 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): @@ -219,6 +221,7 @@ def validate_plain_text(value: str, *, field_name: str) -> str: 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) 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..5f968fe73 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure validation across Slurm execution-plan and runtime-state contracts.""" + +from __future__ import annotations + +import posixpath +from collections.abc import Mapping +from types import MappingProxyType + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference, ShardId +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + CandidateOutputManifest, + ReadinessState, + RunManifest, + ShardManifest, + ShardWinner, + StateContractError, + validate_shard_set, +) + + +class IntegrationContractError(ValueError): + """Raised when reviewed plan and state records disagree.""" + + +class PlanStateValidator: + """Validate records against one resolved plan with reusable derived state. + + This is an in-process validation service, not a persisted contract record. + """ + + 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} + ) + + @property + def plan(self) -> ResolvedSlurmRunPlan: + """The resolved plan used to build this validation context.""" + 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.""" + _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") + try: + validate_shard_set(run, shards) + except StateContractError as exc: + raise IntegrationContractError(str(exc)) from exc + + _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(persisted.shard_id == planned.shard_id, "state shard identity does not match planned order") + _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 shards + + def validate_initial_readiness( + self, + attempt: AttemptManifest, + readiness: AttemptReadiness, + ) -> AttemptReadiness: + """Anchor the first readiness snapshot to a fully validated planned attempt.""" + 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.revision == 1, "initial readiness must use revision 1") + _require(readiness.state is ReadinessState.PENDING, "initial readiness must be pending") + _require(readiness.updated_at >= attempt.created_at, "initial 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, "initial 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") + expected_attempt_id = f"attempt-{attempt.attempt_ordinal:04d}" + _require(attempt.attempt_id == expected_attempt_id, "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) + _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", + ) + _require(client_result.outcome is ClientOutcome.COMPLETE, "only complete client results may be finalized") + _require(candidate.winner_eligible, "only complete candidate outputs may be finalized") + + for record_name, run_id in ( + ("client result", client_result.run_id), + ("candidate", candidate.run_id), + ("winner", winner.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), + ("winner", winner.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), + ("winner", winner.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 == winner.attempt_ordinal, + "candidate and winner attempt ordinals must match the attempt", + ) + _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 == planned_shard.requested_records, + "client and candidate actual records must complete the planned shard", + ) + _require( + client_result.requested_resume_mode == self.plan.invocation.authored.resume, + "client requested resume mode does not match the resolved plan", + ) + + expected_dataset_path = self._get_expected_dataset_path(planned_shard, attempt, client_result) + _require( + client_result.dataset_path == expected_dataset_path, "client dataset path does not match planned intent" + ) + _require( + candidate.dataset_path == expected_dataset_path, "candidate dataset path does not match planned intent" + ) + + expected_manifest_path = posixpath.join( + posixpath.dirname(planned_shard.resume_workspace.path), + "attempts", + attempt.attempt_id, + "output-manifest.json", + ) + candidate_reference = client_result.candidate_output_manifest + _require(candidate_reference is not None, "complete client result has no candidate manifest reference") + _require( + candidate_reference.path == expected_manifest_path, + "candidate manifest path does not match planned intent", + ) + _require( + candidate_reference.sha256 == candidate.compute_sha256(), + "client candidate digest does not match the manifest", + ) + _require( + attempt.candidate_output == candidate_reference, + "attempt candidate reference does not match client result", + ) + _require( + winner.candidate_manifest == candidate_reference, + "winner candidate reference does not match client result", + ) + + _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", + ) + _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") + return winner + + def _get_planned_shard(self, shard_id: ShardId) -> PlannedShard: + planned_shard = self._shards_by_id.get(shard_id) + if planned_shard is None: + raise IntegrationContractError("attempt shard_id does not identify a planned shard") + return planned_shard + + 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", + ) + + @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 IntegrationContractError(message) + + +__all__ = [ + "IntegrationContractError", + "PlanStateValidator", +] 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..517a45ac5 --- /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": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "created_at": "2026-08-19T12:00:02Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + }, + "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": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "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": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "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": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + }, + "created_at": "2026-08-19T12:00:00Z", + "resolved_plan": { + "path": "/workspace/primary/runs/run-single/resolved-plan.json", + "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + }, + "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": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + }, + "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..17e99ea2c --- /dev/null +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -0,0 +1,708 @@ +# 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 + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ("partial_result", "complete client results"), + ("failed_result", "complete client results"), + ("failed_attempt", "successful attempts"), + ("wrong_terminal_classification", "successfully classified"), + ("ineligible_candidate", "complete candidate outputs"), + ("stale_winner", "winner attempt_id"), + ("digest_mismatch", "digest"), + ("dataset_path_mismatch", "dataset path"), + ("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 == "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 ordinals"), + ("winner_ordinal", "attempt ordinals"), + ("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} + ) + 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/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index bf2cffd28..f6ba85e19 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -128,6 +128,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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.planning import ArtifactReference as PlanningArtifactReference from data_designer.slurm.planning import RecordRange as PlanningRecordRange from data_designer.slurm.planning import ResumeWorkspace as PlanningResumeWorkspace From dc03e753d806f3bf52a4770c2c713917b95700f6 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Mon, 24 Aug 2026 21:03:09 -0300 Subject: [PATCH 07/62] test: add deterministic Slurm fakes (#888) * test: add deterministic Slurm fakes Add reusable scheduler, dependency, plugin, serving, clock, and state-transition fixtures for local and CI Slurm integration tests. Keep the harness deterministic and isolated from production lane implementations. Signed-off-by: Andre Manoel * test: bind Slurm fakes to contracts Promote the merged planning and finalization goldens to package-wide fixtures, exercise dependency doubles with canonical records, and add site-neutral rendered script goldens for downstream lanes. Signed-off-by: Andre Manoel * test: complete post-F3 Slurm fake coverage Bind rendered scripts and shared fixtures to the merged contracts, tighten scheduler command behavior, and cover retry terminal states and invalid serving restarts. Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- .../data-designer-slurm/tests/conftest.py | 222 ++++++++++++++++ .../tests/contracts/conftest.py | 48 ---- .../METADATA | 4 + .../entry_points.txt | 2 + .../top_level.txt | 1 + .../fake_data_designer_plugin/__init__.py | 6 + .../fake_data_designer_plugin/plugin.py | 30 +++ .../tests/slurm_test_fakes/__init__.py | 29 +++ .../tests/slurm_test_fakes/clock.py | 42 +++ .../tests/slurm_test_fakes/dependencies.py | 68 +++++ .../golden/rendered/multi_node.sbatch | 33 +++ .../golden/rendered/single_node.sbatch | 33 +++ .../golden/slurm/sacct_retry_terminal.txt | 6 + .../golden/slurm/sacct_terminal.txt | 2 + .../golden/slurm/sinfo_gres.txt | 1 + .../golden/slurm/squeue_active.txt | 2 + .../golden/slurm/squeue_malformed.txt | 1 + .../tests/slurm_test_fakes/serving.py | 182 +++++++++++++ .../tests/slurm_test_fakes/slurm.py | 242 ++++++++++++++++++ .../tests/slurm_test_fakes/test_clock.py | 44 ++++ .../test_contract_fixtures.py | 56 ++++ .../slurm_test_fakes/test_dependencies.py | 120 +++++++++ .../slurm_test_fakes/test_rendered_scripts.py | 75 ++++++ .../tests/slurm_test_fakes/test_serving.py | 214 ++++++++++++++++ .../tests/slurm_test_fakes/test_slurm.py | 237 +++++++++++++++++ .../golden/single_node_failed_readiness.json | 25 ++ .../golden/single_node_pending_readiness.json | 20 ++ ...gle_node_publication_failed_readiness.json | 25 ++ .../single_node_starting_readiness.json | 25 ++ .../golden/single_node_stopped_readiness.json | 25 ++ .../tests/state/test_state_golden_records.py | 5 + .../tests/state/test_transition_fixtures.py | 39 +++ 32 files changed, 1816 insertions(+), 48 deletions(-) create mode 100644 packages/data-designer-slurm/tests/conftest.py delete mode 100644 packages/data-designer-slurm/tests/contracts/conftest.py create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/clock.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/serving.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json create mode 100644 packages/data-designer-slurm/tests/state/test_transition_fixtures.py diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py new file mode 100644 index 000000000..a46fcb160 --- /dev/null +++ b/packages/data-designer-slurm/tests/conftest.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import cast + +import pytest +from slurm_test_fakes import ( + FakeClock, + FakeCommandResponse, + FakeLogicalEndpoint, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, + FakeVllmBackend, +) + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + ImageInspectionRecord, + SlurmProfileCatalog, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerIdentity, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + +TEST_DIRECTORY = Path(__file__).parent +CONTRACT_GOLDEN_DIRECTORY = TEST_DIRECTORY / "contracts" / "golden" +INTEGRATION_GOLDEN_PATH = TEST_DIRECTORY / "integration" / "golden" / "finalization_chain.json" +SLURM_GOLDEN_DIRECTORY = TEST_DIRECTORY / "slurm_test_fakes" / "golden" / "slurm" +STATE_GOLDEN_DIRECTORY = TEST_DIRECTORY / "state" / "golden" + + +@pytest.fixture +def authored_run() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "authored_run.json").read_text()) + + +@pytest.fixture +def authored_run_single() -> DataDesignerSlurmConfig: + return DataDesignerSlurmConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json").read_text() + ) + + +@pytest.fixture +def profile_catalog() -> SlurmProfileCatalog: + return SlurmProfileCatalog.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "profile_catalog.json").read_text()) + + +@pytest.fixture +def dependency_lock() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "dependency_lock.json").read_text()) + + +@pytest.fixture +def dependency_lock_single() -> ResolvedDependencyLock: + return ResolvedDependencyLock.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text() + ) + + +@pytest.fixture +def single_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text()) + + +@pytest.fixture +def multi_node_plan() -> ResolvedSlurmRunPlan: + return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "multi_node_plan.json").read_text()) + + +@pytest.fixture +def client_result() -> ClientResult: + return ClientResult.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "client_result.json").read_text()) + + +@pytest.fixture +def benchmark_config() -> DataDesignerSlurmBenchmarkConfig: + return DataDesignerSlurmBenchmarkConfig.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "benchmark_config.json").read_text() + ) + + +@pytest.fixture +def benchmark_manifest() -> BenchmarkManifest: + return BenchmarkManifest.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_manifest.json").read_text()) + + +@pytest.fixture +def benchmark_report() -> BenchmarkReport: + return BenchmarkReport.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "benchmark_report.json").read_text()) + + +@pytest.fixture +def client_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "client_image_inspection.json").read_text() + ) + + +@pytest.fixture +def serving_image_inspection() -> ImageInspectionRecord: + return ImageInspectionRecord.model_validate_json( + (CONTRACT_GOLDEN_DIRECTORY / "serving_image_inspection.json").read_text() + ) + + +@pytest.fixture +def collection_plan() -> CollectionPlan: + return CollectionPlan.model_validate_json((STATE_GOLDEN_DIRECTORY / "collection_plan.json").read_text()) + + +@pytest.fixture +def accounting_lag_observation() -> SchedulerObservation: + return SchedulerObservation.model_validate_json((STATE_GOLDEN_DIRECTORY / "accounting_lag.json").read_text()) + + +@pytest.fixture +def finalization_chain_payload() -> dict[str, object]: + return cast(dict[str, object], json.loads(INTEGRATION_GOLDEN_PATH.read_text())) + + +@pytest.fixture +def run_manifest(finalization_chain_payload: dict[str, object]) -> RunManifest: + return RunManifest.model_validate_json(json.dumps(finalization_chain_payload["run"])) + + +@pytest.fixture +def shard_manifests(finalization_chain_payload: dict[str, object]) -> tuple[ShardManifest, ...]: + return tuple( + ShardManifest.model_validate_json(json.dumps(payload)) + for payload in cast(list[object], finalization_chain_payload["shards"]) + ) + + +@pytest.fixture +def attempt_manifest(finalization_chain_payload: dict[str, object]) -> AttemptManifest: + return AttemptManifest.model_validate_json(json.dumps(finalization_chain_payload["attempt"])) + + +@pytest.fixture +def attempt_readiness(finalization_chain_payload: dict[str, object]) -> AttemptReadiness: + return AttemptReadiness.model_validate_json(json.dumps(finalization_chain_payload["readiness"])) + + +@pytest.fixture +def finalization_client_result(finalization_chain_payload: dict[str, object]) -> ClientResult: + return ClientResult.model_validate_json(json.dumps(finalization_chain_payload["client_result"])) + + +@pytest.fixture +def candidate_output_manifest(finalization_chain_payload: dict[str, object]) -> CandidateOutputManifest: + return CandidateOutputManifest.model_validate_json(json.dumps(finalization_chain_payload["candidate"])) + + +@pytest.fixture +def shard_winner(finalization_chain_payload: dict[str, object]) -> ShardWinner: + return ShardWinner.model_validate_json(json.dumps(finalization_chain_payload["winner"])) + + +@pytest.fixture +def fake_clock() -> FakeClock: + """Return an isolated explicitly controlled clock.""" + return FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=100.0) + + +@pytest.fixture +def fake_slurm_runner() -> FakeSlurmRunner: + """Return an isolated Slurm runner with one array and one bounded sinfo query.""" + return FakeSlurmRunner( + arrays=( + FakeSlurmArray( + tasks=( + FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)), + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + ), + ) + ), + ), + sinfo_responses={ + ("sinfo", "--noheader", "--format=%G"): FakeCommandResponse( + stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + ) + }, + ) + + +@pytest.fixture +def fake_plugin_overlay() -> Path: + """Return the installed-layout fake Data Designer plugin overlay.""" + return TEST_DIRECTORY / "fixtures" / "fake_plugin_overlay" + + +@pytest.fixture +def fake_logical_endpoint() -> FakeLogicalEndpoint: + """Return an isolated two-backend logical endpoint.""" + return FakeLogicalEndpoint( + "http://127.0.0.1:31000", + ( + FakeVllmBackend("http://127.0.0.1:31001", rank=0), + FakeVllmBackend("http://127.0.0.1:31002", rank=1), + ), + ) diff --git a/packages/data-designer-slurm/tests/contracts/conftest.py b/packages/data-designer-slurm/tests/contracts/conftest.py deleted file mode 100644 index 44a74d99c..000000000 --- a/packages/data-designer-slurm/tests/contracts/conftest.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from pathlib import Path - -import pytest - -from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfileCatalog -from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan - -GOLDEN_DIR = Path(__file__).parent / "golden" - - -@pytest.fixture -def authored_run() -> DataDesignerSlurmConfig: - return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run.json").read_text()) - - -@pytest.fixture -def authored_run_single() -> DataDesignerSlurmConfig: - return DataDesignerSlurmConfig.model_validate_json((GOLDEN_DIR / "authored_run_single.json").read_text()) - - -@pytest.fixture -def profile_catalog() -> SlurmProfileCatalog: - return SlurmProfileCatalog.model_validate_json((GOLDEN_DIR / "profile_catalog.json").read_text()) - - -@pytest.fixture -def dependency_lock() -> ResolvedDependencyLock: - return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock.json").read_text()) - - -@pytest.fixture -def dependency_lock_single() -> ResolvedDependencyLock: - return ResolvedDependencyLock.model_validate_json((GOLDEN_DIR / "dependency_lock_single.json").read_text()) - - -@pytest.fixture -def single_node_plan() -> ResolvedSlurmRunPlan: - return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "single_node_plan.json").read_text()) - - -@pytest.fixture -def multi_node_plan() -> ResolvedSlurmRunPlan: - return ResolvedSlurmRunPlan.model_validate_json((GOLDEN_DIR / "multi_node_plan.json").read_text()) diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA new file mode 100644 index 000000000..1e52ff779 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/METADATA @@ -0,0 +1,4 @@ +Metadata-Version: 2.3 +Name: fake-data-designer-plugin +Version: 1.0.0 +Requires-Dist: data-designer diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt new file mode 100644 index 000000000..19e752123 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[data_designer.plugins] +fake-slurm-column = fake_data_designer_plugin.plugin:plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt new file mode 100644 index 000000000..7ca87bde3 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin-1.0.0.dist-info/top_level.txt @@ -0,0 +1 @@ +fake_data_designer_plugin diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py new file mode 100644 index 000000000..aa059a7a1 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal Data Designer plugin used by Slurm client-environment tests.""" + +from __future__ import annotations diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py new file mode 100644 index 000000000..2dc7726b3 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Literal + +from data_designer.config.base import SingleColumnConfig +from data_designer.plugins import Plugin, PluginType + + +class FakePluginConfig(SingleColumnConfig): + """Minimal custom column configuration.""" + + column_type: Literal["fake-slurm-column"] = "fake-slurm-column" + + +class FakePluginImplementation: + """Minimal loadable implementation for entry-point verification.""" + + def generate(self, data: dict[str, object]) -> dict[str, object]: + """Return the provided record unchanged.""" + return data + + +plugin = Plugin( + config_qualified_name="fake_data_designer_plugin.plugin.FakePluginConfig", + impl_qualified_name="fake_data_designer_plugin.plugin.FakePluginImplementation", + plugin_type=PluginType.COLUMN_GENERATOR, +) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py new file mode 100644 index 000000000..a73536ea8 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic test doubles for the Slurm integration.""" + +from __future__ import annotations + +from slurm_test_fakes.clock import FakeClock +from slurm_test_fakes.dependencies import FakeDependencyInstaller, FakeDependencyResolver +from slurm_test_fakes.serving import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend +from slurm_test_fakes.slurm import ( + FakeCommandResponse, + FakeSlurmArray, + FakeSlurmRunner, + FakeSlurmTask, +) + +__all__ = [ + "FakeClock", + "FakeCommandResponse", + "FakeDependencyInstaller", + "FakeDependencyResolver", + "FakeLogicalEndpoint", + "FakeServingState", + "FakeSlurmArray", + "FakeSlurmRunner", + "FakeSlurmTask", + "FakeVllmBackend", +] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py new file mode 100644 index 000000000..45485aae0 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/clock.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timedelta + + +@dataclass +class FakeClock: + """Clock advanced only by explicit test input.""" + + current_time: datetime + monotonic_time: float = 0.0 + sleep_calls: list[float] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.current_time.tzinfo is None or self.current_time.utcoffset() != timedelta(0): + raise ValueError("current_time must be in UTC") + if self.monotonic_time < 0: + raise ValueError("monotonic_time must not be negative") + + def now(self) -> datetime: + """Return the controlled wall-clock time.""" + return self.current_time + + def monotonic(self) -> float: + """Return the controlled monotonic time.""" + return self.monotonic_time + + def sleep(self, seconds: float) -> None: + """Record a sleep and advance both clocks without blocking.""" + self.advance(seconds) + self.sleep_calls.append(seconds) + + def advance(self, seconds: float) -> None: + """Advance both clocks by a non-negative duration.""" + if seconds < 0: + raise ValueError("seconds must not be negative") + self.current_time += timedelta(seconds=seconds) + self.monotonic_time += seconds diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py new file mode 100644 index 000000000..ac0a8a4cb --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/dependencies.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from pathlib import Path +from typing import Generic, TypeVar + +RequestT = TypeVar("RequestT") +ResolutionT = TypeVar("ResolutionT") +LockT = TypeVar("LockT") +InstallationT = TypeVar("InstallationT") + + +class FakeDependencyResolver(Generic[RequestT, ResolutionT]): + """Return or raise exact scripted dependency-resolution outcomes.""" + + def __init__(self, responses: Iterable[tuple[RequestT, ResolutionT | BaseException]]) -> None: + self._responses = deque(responses) + self.calls: list[RequestT] = [] + + def resolve(self, request: RequestT) -> ResolutionT: + """Resolve one expected request without consulting package indexes.""" + self.calls.append(request) + if not self._responses: + raise AssertionError("unexpected dependency resolution") + expected, response = self._responses.popleft() + if request != expected: + raise AssertionError(f"expected dependency request {expected!r}, got {request!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted resolution was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency resolutions remain") + + +class FakeDependencyInstaller(Generic[LockT, InstallationT]): + """Return or raise exact scripted dependency-install outcomes.""" + + def __init__( + self, + responses: Iterable[tuple[tuple[LockT, Path], InstallationT | BaseException]], + ) -> None: + self._responses = deque(responses) + self.calls: list[tuple[LockT, Path]] = [] + + def install(self, lock: LockT, target: Path) -> InstallationT: + """Install one expected lock without invoking an installer process.""" + call = (lock, target) + self.calls.append(call) + if not self._responses: + raise AssertionError("unexpected dependency installation") + expected, response = self._responses.popleft() + if call != expected: + raise AssertionError(f"expected dependency installation {expected!r}, got {call!r}") + if isinstance(response, BaseException): + raise response + return response + + def assert_complete(self) -> None: + """Assert that every scripted installation was consumed.""" + if self._responses: + raise AssertionError(f"{len(self._responses)} dependency installations remain") diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch new file mode 100644 index 000000000..5a78c97b7 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +#SBATCH --job-name=dd-two-model +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=3 +#SBATCH --time=03:55:00 +#SBATCH --array=0-1%2 +#SBATCH --gres=gpu:8 +set -Eeuo pipefail + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" +readonly DD_PLAN_SHA256="dbcae1da6ce8ef6799faf2add8e08b1b093390ae08a37dd13e336da522b9a3fd" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + printf '%s %s\n' "$1" "$2" | sha256sum --check --status - +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +install -d -m 0700 "${DD_ATTEMPT_DIR}" +readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" + +source "${DD_RUNTIME_DIR}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch new file mode 100644 index 000000000..e10478f7d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +#SBATCH --job-name=data-designer +#SBATCH --account=research +#SBATCH --partition=batch +#SBATCH --nodes=1 +#SBATCH --time=03:55:00 +#SBATCH --array=0 +#SBATCH --gres=gpu:8 +set -Eeuo pipefail + +readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" +readonly DD_PLAN_SHA256="d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" +readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" +readonly DD_ATTEMPT_ORDINAL="0001" + +verify_sha256() { + printf '%s %s\n' "$1" "$2" | sha256sum --check --status - +} + +verify_sha256 "${DD_RUNTIME_SHA256}" "${DD_RUNTIME_ARCHIVE}" +verify_sha256 "${DD_PLAN_SHA256}" "${DD_PLAN}" +readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +printf -v DD_SHARD_ID 'shard-%05d' "${DD_ARRAY_TASK_ID}" +readonly DD_SHARD_ID +readonly DD_ATTEMPT_DIR="${DD_RUN_ROOT}/shards/${DD_SHARD_ID}/attempts/attempt-${DD_ATTEMPT_ORDINAL}" +install -d -m 0700 "${DD_ATTEMPT_DIR}" +readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" + +source "${DD_RUNTIME_DIR}/entrypoint.sh" +dd_slurm_run_allocation "${DD_PLAN}" "${DD_ATTEMPT_DIR}" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt new file mode 100644 index 000000000..7ac0e7937 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_retry_terminal.txt @@ -0,0 +1,6 @@ +4101_0|TIMEOUT|0:125 +4101_1|NODE_FAIL|1:0 +4101_2|PREEMPTED|0:0 +4101_3|REQUEUED|0:0 +4101_4|OUT_OF_MEMORY|0:125 +4101_5|CANCELLED by 1234|0:15 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt new file mode 100644 index 000000000..17686ec3f --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sacct_terminal.txt @@ -0,0 +1,2 @@ +4101_0|COMPLETED|0:0 +4101_1|FAILED|1:0 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt new file mode 100644 index 000000000..25175f090 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/sinfo_gres.txt @@ -0,0 +1 @@ +gpu:2 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt new file mode 100644 index 000000000..80e935355 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_active.txt @@ -0,0 +1,2 @@ +4101_0|PENDING +4101_1|RUNNING diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt new file mode 100644 index 000000000..670c5234d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/slurm/squeue_malformed.txt @@ -0,0 +1 @@ +malformed scheduler output diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py new file mode 100644 index 000000000..0bc05aed3 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/serving.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from enum import Enum + +import httpx + + +class FakeServingState(str, Enum): + """Lifecycle states shared by the serving test doubles.""" + + CREATED = "created" + STARTING = "starting" + READY = "ready" + FAILED = "failed" + STOPPED = "stopped" + + +class FakeVllmBackend: + """In-memory vLLM boundary with explicit lifecycle and responses.""" + + def __init__( + self, + endpoint: str, + *, + rank: int, + responses: Iterable[httpx.Response | BaseException] = (), + ) -> None: + if rank < 0: + raise ValueError("rank must not be negative") + self.endpoint = endpoint + self.rank = rank + self.state = FakeServingState.CREATED + self.failure_reason: str | None = None + self.requests: list[httpx.Request] = [] + self.start_calls = 0 + self.cleanup_calls = 0 + self._responses = deque(responses) + + def start(self) -> None: + """Enter the explicit startup state.""" + self.start_calls += 1 + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"backend cannot start from {self.state.value}") + self.state = FakeServingState.STARTING + + def mark_ready(self) -> None: + """Mark startup complete.""" + if self.state is not FakeServingState.STARTING: + raise RuntimeError("backend must be starting before it becomes ready") + self.state = FakeServingState.READY + + def fail(self, reason: str) -> None: + """Inject a backend or rank failure.""" + if not reason: + raise ValueError("failure reason must not be empty") + self.failure_reason = reason + self.state = FakeServingState.FAILED + + def queue_response(self, response: httpx.Response | BaseException) -> None: + """Queue one explicit generation response or failure.""" + self._responses.append(response) + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle a health or generation request without opening a socket.""" + self.requests.append(request) + if request.url.path == "/health": + status_code = 200 if self.state is FakeServingState.READY else 503 + return httpx.Response(status_code, request=request) + if self.state is not FakeServingState.READY: + return httpx.Response(503, json={"error": "backend_unavailable"}, request=request) + if not self._responses: + return httpx.Response( + 200, + json={"backend": self.endpoint, "rank": self.rank}, + request=request, + ) + response = self._responses.popleft() + if isinstance(response, BaseException): + raise response + response.request = request + return response + + def cleanup(self) -> None: + """Stop the backend once while allowing cleanup re-entry.""" + self.cleanup_calls += 1 + self.state = FakeServingState.STOPPED + + +class FakeLogicalEndpoint: + """In-memory logical endpoint over one or more fake vLLM backends.""" + + def __init__(self, endpoint: str, backends: Iterable[FakeVllmBackend]) -> None: + self.endpoint = endpoint + self.backends = tuple(backends) + if not self.backends: + raise ValueError("logical endpoints require at least one backend") + self.state = FakeServingState.CREATED + self.published_endpoint: str | None = None + self.requests: list[httpx.Request] = [] + self.publish_calls = 0 + self.cleanup_calls = 0 + self.failure_reason: str | None = None + self._next_backend = 0 + self._publication_failure: Exception | None = None + + def start(self) -> None: + """Start every backend and the readiness aggregator.""" + if self.state is not FakeServingState.CREATED: + raise RuntimeError(f"logical endpoint cannot start from {self.state.value}") + for backend in self.backends: + backend.start() + self.state = FakeServingState.STARTING + + def refresh_readiness(self) -> FakeServingState: + """Aggregate backend readiness with coordinated failure semantics.""" + if self.state is FakeServingState.STOPPED: + return self.state + if self._publication_failure is not None and self.publish_calls: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + return self.state + backend_states = {backend.state for backend in self.backends} + if FakeServingState.FAILED in backend_states: + self.state = FakeServingState.FAILED + self.failure_reason = "backend_failed" + elif backend_states == {FakeServingState.READY}: + self.state = FakeServingState.READY + else: + self.state = FakeServingState.STARTING + return self.state + + def publish(self) -> str: + """Publish the logical endpoint after every backend is ready.""" + if self.refresh_readiness() is not FakeServingState.READY: + raise RuntimeError("logical endpoint is not ready") + self.publish_calls += 1 + if self._publication_failure is not None: + self.state = FakeServingState.FAILED + self.failure_reason = "endpoint_publication_failed" + raise self._publication_failure + self.published_endpoint = self.endpoint + return self.endpoint + + def script_publication_failure(self, error: Exception) -> None: + """Inject one explicit endpoint-publication failure.""" + self._publication_failure = error + + def handle(self, request: httpx.Request) -> httpx.Response: + """Handle an aggregated health request or forward in round-robin order.""" + self.requests.append(request) + state = self.refresh_readiness() + if request.url.path == "/health": + status_code = 200 if state is FakeServingState.READY and self.published_endpoint else 503 + return httpx.Response(status_code, request=request) + if state is FakeServingState.FAILED: + return httpx.Response(503, json={"error": self.failure_reason}, request=request) + if state is not FakeServingState.READY or self.published_endpoint is None: + return httpx.Response(503, json={"error": "endpoint_unavailable"}, request=request) + response: httpx.Response | None = None + for _ in self.backends: + backend = self.backends[self._next_backend] + self._next_backend = (self._next_backend + 1) % len(self.backends) + response = backend.handle(request) + if response.status_code != 429: + return response + assert response is not None + return response + + def cleanup(self) -> None: + """Stop every backend and unpublish the process-local endpoint idempotently.""" + self.cleanup_calls += 1 + if self.state is FakeServingState.STOPPED: + return + for backend in self.backends: + backend.cleanup() + self.published_endpoint = None + self.state = FakeServingState.STOPPED diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py new file mode 100644 index 000000000..785af0630 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import re +import subprocess +from collections import deque +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.state import SchedulerIdentity + +_JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--format=%i|%T") +_SACCT_REQUIRED_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") + + +@dataclass(frozen=True) +class FakeCommandResponse: + """One deterministic subprocess response.""" + + stdout: str = "" + stderr: str = "" + returncode: int = 0 + + +@dataclass +class FakeSlurmTask: + """Mutable scheduler views for one canonical array-task identity.""" + + scheduler: SchedulerIdentity + queue_state: str | None = "PENDING" + accounting_state: str | None = None + exit_code: str = "0:0" + + +@dataclass +class FakeSlurmArray: + """A deterministic array submission exposed through Slurm command output.""" + + tasks: tuple[FakeSlurmTask, ...] + + def __post_init__(self) -> None: + if not self.tasks: + raise ValueError("fake Slurm arrays require at least one task") + job_ids = {task.scheduler.array_job_id for task in self.tasks} + task_ids = [task.scheduler.array_task_id for task in self.tasks] + if len(job_ids) != 1: + raise ValueError("fake Slurm array tasks must share one array job ID") + if len(task_ids) != len(set(task_ids)): + raise ValueError("fake Slurm array task IDs must be unique") + + @property + def array_job_id(self) -> int: + """Return the canonical job ID shared by the array tasks.""" + return self.tasks[0].scheduler.array_job_id + + +class FakeSlurmRunner: + """Stateful fake that copies its configured arrays before exposing Slurm commands.""" + + def __init__( + self, + arrays: Iterable[FakeSlurmArray] = (), + *, + sinfo_responses: Mapping[tuple[str, ...], FakeCommandResponse] | None = None, + ) -> None: + self._pending_arrays = deque(copy.deepcopy(tuple(arrays))) + self._submitted_arrays: dict[int, FakeSlurmArray] = {} + self._scripted_responses: dict[str, deque[FakeCommandResponse]] = {} + self._sinfo_responses = dict(sinfo_responses or {}) + self.calls: list[tuple[str, ...]] = [] + + def run(self, command: Sequence[str], *, check: bool = False) -> subprocess.CompletedProcess[str]: + """Run one fake Slurm command and optionally raise on failure.""" + if not command: + raise ValueError("command must not be empty") + argv = tuple(command) + self.calls.append(argv) + command_name = Path(argv[0]).name + scripted = self._scripted_responses.get(command_name) + if scripted: + response = scripted.popleft() + else: + response = self._dispatch(command_name, argv) + + completed = subprocess.CompletedProcess( + args=list(argv), + returncode=response.returncode, + stdout=response.stdout, + stderr=response.stderr, + ) + if check and response.returncode: + raise subprocess.CalledProcessError( + response.returncode, + list(argv), + output=response.stdout, + stderr=response.stderr, + ) + return completed + + def script_next(self, command_name: str, response: FakeCommandResponse) -> None: + """Inject one explicit response before the command's stateful behavior.""" + self._scripted_responses.setdefault(command_name, deque()).append(response) + + def set_task_state( + self, + scheduler: SchedulerIdentity, + *, + queue_state: str | None, + accounting_state: str | None, + exit_code: str = "0:0", + ) -> None: + """Set the independently observable queue and accounting states.""" + task = self._find_task(scheduler) + task.queue_state = queue_state + task.accounting_state = accounting_state + task.exit_code = exit_code + + def assert_scripts_consumed(self) -> None: + """Assert that no scripted command response remains.""" + remaining = sum(len(responses) for responses in self._scripted_responses.values()) + if remaining: + raise AssertionError(f"{remaining} scripted Slurm responses remain") + + def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResponse: + handlers = { + "sacct": self._run_sacct, + "sbatch": self._run_sbatch, + "scancel": self._run_scancel, + "sinfo": self._run_sinfo, + "squeue": self._run_squeue, + } + try: + handler = handlers[command_name] + except KeyError: + raise AssertionError(f"unexpected command {command_name!r}") from None + return handler(argv) + + def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: + if not self._pending_arrays: + return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) + array = self._pending_arrays.popleft() + self._submitted_arrays[array.array_job_id] = array + if "--parsable" in argv[1:]: + return FakeCommandResponse(stdout=f"{array.array_job_id}\n") + return FakeCommandResponse(stdout=f"Submitted batch job {array.array_job_id}\n") + + def _run_squeue(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SQUEUE_REQUIRED_ARGUMENTS) + rows = [ + f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.queue_state}" + for task in self._selected_submitted_tasks(argv) + if task.queue_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + + def _run_sacct(self, argv: tuple[str, ...]) -> FakeCommandResponse: + self._require_arguments(argv, _SACCT_REQUIRED_ARGUMENTS) + rows = [ + (f"{task.scheduler.array_job_id}_{task.scheduler.array_task_id}|{task.accounting_state}|{task.exit_code}") + for task in self._selected_submitted_tasks(argv) + if task.accounting_state is not None + ] + return FakeCommandResponse(stdout="".join(f"{row}\n" for row in rows)) + + def _run_scancel(self, argv: tuple[str, ...]) -> FakeCommandResponse: + targets = tuple(argument for argument in argv[1:] if not argument.startswith("-")) + if len(targets) != 1: + return FakeCommandResponse(stderr="expected one cancellation target\n", returncode=1) + try: + tasks = self._tasks_for_target(targets[0]) + except (KeyError, ValueError): + return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) + for task in tasks: + task.queue_state = None + task.accounting_state = "CANCELLED" + task.exit_code = "0:15" + return FakeCommandResponse() + + def _run_sinfo(self, argv: tuple[str, ...]) -> FakeCommandResponse: + key = ("sinfo", *argv[1:]) + try: + return self._sinfo_responses[key] + except KeyError: + raise AssertionError(f"unexpected sinfo query {key!r}") from None + + @staticmethod + def _require_arguments(argv: tuple[str, ...], required: tuple[str, ...]) -> None: + missing = tuple(argument for argument in required if argument not in argv[1:]) + if missing: + raise AssertionError(f"missing required arguments {missing!r} in {argv!r}") + + def _tasks_for_target(self, target: str) -> tuple[FakeSlurmTask, ...]: + if "_" not in target: + return self._submitted_arrays[int(target)].tasks + job_id, task_id = target.split("_", 1) + scheduler = SchedulerIdentity(array_job_id=int(job_id), array_task_id=int(task_id)) + return (self._find_task(scheduler),) + + def _find_task(self, scheduler: SchedulerIdentity) -> FakeSlurmTask: + array = self._submitted_arrays[scheduler.array_job_id] + for task in array.tasks: + if task.scheduler == scheduler: + return task + raise KeyError(scheduler) + + def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: + return sorted( + (task for array in self._submitted_arrays.values() for task in array.tasks), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) + + def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: + selectors: list[str] = [] + for index, argument in enumerate(argv[1:]): + if argument in {"-j", "--jobs"}: + try: + selectors.extend(argv[index + 2].split(",")) + except IndexError: + raise AssertionError(f"missing job selector in {argv!r}") from None + elif argument.startswith("--jobs="): + selectors.extend(argument.partition("=")[2].split(",")) + if not selectors: + return self._sorted_submitted_tasks() + + selected: dict[SchedulerIdentity, FakeSlurmTask] = {} + for selector in selectors: + if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: + raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") + try: + tasks = self._tasks_for_target(selector) + except KeyError: + continue + selected.update((task.scheduler, task) for task in tasks) + return sorted( + selected.values(), + key=lambda task: (task.scheduler.array_job_id, task.scheduler.array_task_id), + ) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py new file mode 100644 index 000000000..b160d5c21 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_clock.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime, timedelta, timezone + +import pytest + +from slurm_test_fakes import FakeClock + + +def test_fake_clock_advances_without_blocking(fake_clock: FakeClock) -> None: + fake_clock.sleep(2.5) + fake_clock.advance(1.5) + + assert fake_clock.now() == datetime(2026, 8, 18, 12, 0, 4, tzinfo=timezone.utc) + assert fake_clock.monotonic() == 104.0 + assert fake_clock.sleep_calls == [2.5] + + +@pytest.mark.parametrize( + "clock", + ( + pytest.param(lambda: FakeClock(datetime(2026, 8, 18, 12)), id="naive-datetime"), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 13, tzinfo=timezone(timedelta(hours=1)))), + id="non-utc-offset", + ), + pytest.param( + lambda: FakeClock(datetime(2026, 8, 18, 12, tzinfo=timezone.utc), monotonic_time=-1), + id="negative-monotonic", + ), + ), +) +def test_fake_clock_rejects_ambient_or_invalid_time(clock: Callable[[], FakeClock]) -> None: + with pytest.raises(ValueError): + clock() + + +def test_fake_clock_rejects_negative_advances(fake_clock: FakeClock) -> None: + with pytest.raises(ValueError, match="must not be negative"): + fake_clock.advance(-0.1) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py new file mode 100644 index 000000000..d32aefa5b --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_contract_fixtures.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.client import ClientResult +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, ImageInspectionRecord +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + CollectionPlan, + RunManifest, + SchedulerObservation, + ShardManifest, + ShardWinner, +) + + +def test_shared_contract_fixtures_load_canonical_records( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, + client_result: ClientResult, + run_manifest: RunManifest, + shard_manifests: tuple[ShardManifest, ...], + attempt_manifest: AttemptManifest, + attempt_readiness: AttemptReadiness, + finalization_client_result: ClientResult, + candidate_output_manifest: CandidateOutputManifest, + shard_winner: ShardWinner, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, + client_image_inspection: ImageInspectionRecord, + serving_image_inspection: ImageInspectionRecord, + collection_plan: CollectionPlan, + accounting_lag_observation: SchedulerObservation, +) -> None: + assert single_node_plan.run_id == "run-single" + assert multi_node_plan.run_id == "run-001" + assert client_result.run_id == "run-001" + assert run_manifest.run_id == "run-single" + assert len(shard_manifests) == 1 + assert attempt_manifest.attempt_id == "attempt-0001" + assert attempt_readiness.revision == 1 + assert finalization_client_result.outcome.value == "complete" + assert candidate_output_manifest.winner_eligible + assert shard_winner.attempt_id == "attempt-0001" + assert benchmark_config.name == "generator-scaling" + assert benchmark_manifest.benchmark_id == benchmark_report.benchmark_id + assert client_image_inspection.inspection.kind == "client" + assert serving_image_inspection.inspection.kind == "serving" + assert collection_plan.collection_id == "collection-0001" + assert accounting_lag_observation.state.value == "accounting_lag" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py new file mode 100644 index 000000000..324a8948d --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +from importlib.metadata import entry_points +from pathlib import Path + +import pytest + +from data_designer.plugins import PluginType +from data_designer.plugins.registry import PluginRegistry +from data_designer.slurm.config import ClientDependencies, DataDesignerSlurmConfig +from data_designer.slurm.planning import ResolvedDependencyLock +from slurm_test_fakes import FakeDependencyInstaller, FakeDependencyResolver + + +def test_fake_dependency_resolver_scripts_compatible_incompatible_and_missing_cases( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, +) -> None: + compatible = authored_run.client.dependencies + incompatible = ClientDependencies(requirements=["fake-data-designer-plugin==2.0.0"]) + missing = ClientDependencies(requirements=["missing-plugin==1.0.0"]) + resolver: FakeDependencyResolver[ClientDependencies, ResolvedDependencyLock] = FakeDependencyResolver( + ( + (compatible, dependency_lock), + (incompatible, ValueError("incompatible Data Designer version")), + (missing, FileNotFoundError("dependency artifact is missing")), + ) + ) + + assert resolver.resolve(compatible) is dependency_lock + with pytest.raises(ValueError, match="incompatible"): + resolver.resolve(incompatible) + with pytest.raises(FileNotFoundError, match="missing"): + resolver.resolve(missing) + + assert resolver.calls == [compatible, incompatible, missing] + resolver.assert_complete() + + +def test_fake_dependency_installer_scripts_success_and_digest_mismatch( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + first_target = tmp_path / "compatible" + second_target = tmp_path / "digest-mismatch" + mismatched_lock = dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + installer = FakeDependencyInstaller[ResolvedDependencyLock, tuple[object, ...]]( + ( + ((dependency_lock, first_target), dependency_lock.overlay_packages), + ((mismatched_lock, second_target), ValueError("lock digest mismatch")), + ) + ) + + assert installer.install(dependency_lock, first_target) == dependency_lock.overlay_packages + with pytest.raises(ValueError, match="digest mismatch"): + installer.install(mismatched_lock, second_target) + + assert installer.calls == [(dependency_lock, first_target), (mismatched_lock, second_target)] + installer.assert_complete() + + +def test_dependency_fakes_raise_cancellation_signals( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + request = ClientDependencies(requirements=[]) + resolver = FakeDependencyResolver(((request, KeyboardInterrupt()),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path), SystemExit(2)),)) + + with pytest.raises(KeyboardInterrupt): + resolver.resolve(request) + with pytest.raises(SystemExit, match="2"): + installer.install(dependency_lock, tmp_path) + + +def test_dependency_fakes_reject_unexpected_calls( + tmp_path: Path, + dependency_lock: ResolvedDependencyLock, +) -> None: + expected = ClientDependencies(requirements=[]) + unexpected = ClientDependencies(requirements=["unexpected==1.0.0"]) + resolver = FakeDependencyResolver(((expected, dependency_lock),)) + installer = FakeDependencyInstaller((((dependency_lock, tmp_path / "expected"), "installed"),)) + + with pytest.raises(AssertionError, match="expected dependency request"): + resolver.resolve(unexpected) + with pytest.raises(AssertionError, match="expected dependency installation"): + installer.install(dependency_lock, tmp_path / "unexpected") + + +def test_fake_plugin_overlay_supports_real_entry_point_discovery( + fake_plugin_overlay: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(str(fake_plugin_overlay)) + importlib.invalidate_caches() + monkeypatch.setattr("data_designer.plugins.registry.PLUGINS_DISABLED", False) + PluginRegistry.reset() + try: + discovered = tuple( + entry_point + for entry_point in entry_points(group="data_designer.plugins") + if entry_point.name == "fake-slurm-column" + ) + assert len(discovered) == 1 + assert discovered[0].dist is not None + assert discovered[0].dist.version == "1.0.0" + + plugin = discovered[0].load() + registry = PluginRegistry() + + assert plugin.name == "fake-slurm-column" + assert plugin.plugin_type is PluginType.COLUMN_GENERATOR + assert registry.get_plugin("fake-slurm-column") == plugin + finally: + PluginRegistry.reset() diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py new file mode 100644 index 000000000..c7d119389 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import posixpath +from pathlib import Path + +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "rendered" + + +def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( + single_node_plan: ResolvedSlurmRunPlan, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + _assert_script_matches_plan( + single_node_plan, + "single_node.sbatch", + expected_fixture_sha256="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", + ) + _assert_script_matches_plan( + multi_node_plan, + "multi_node.sbatch", + expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + ) + + +def test_rendered_scripts_contain_no_site_or_user_data() -> None: + contents = "\n".join(path.read_text().casefold() for path in GOLDEN_DIRECTORY.glob("*.sbatch")) + + for forbidden in ( + "nvidia", + "cluster", + "login", + "lustre", + "/home/", + "/users/", + ): + assert forbidden not in contents + + +def _assert_script_matches_plan( + plan: ResolvedSlurmRunPlan, + filename: str, + *, + expected_fixture_sha256: str, +) -> None: + script = (GOLDEN_DIRECTORY / filename).read_text() + node_indices = ( + plan.client.host_node_index, + *(index for deployment in plan.deployments for index in deployment.node_indices), + ) + node_count = max(node_indices) + 1 + array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + plan_path = posixpath.join(posixpath.dirname(plan.authored_config.path), "resolved-plan.json") + run_root = posixpath.dirname(plan.authored_config.path) + + assert f"#SBATCH --job-name={plan.submission.job_name}\n" in script + assert f"#SBATCH --account={plan.submission.account}\n" in script + assert f"#SBATCH --partition={plan.submission.partition}\n" in script + assert f"#SBATCH --nodes={node_count}\n" in script + assert f"#SBATCH --time={plan.submission.time_limit}\n" in script + assert f"#SBATCH --array={array}\n" in script + assert f"#SBATCH --gres=gpu:{plan.resolved_gpus_per_node}\n" in script + assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"\n' in script + assert f'readonly DD_RUNTIME_SHA256="{plan.runtime_bundle.sha256}"\n' in script + assert f'readonly DD_PLAN="{plan_path}"\n' in script + assert f'readonly DD_PLAN_SHA256="{plan.compute_sha256()}"\n' in script + assert f'readonly DD_RUN_ROOT="{run_root}"\n' in script + assert script.count("dd_slurm_run_allocation") == 1 + assert script.endswith("\n") + assert hashlib.sha256(script.encode()).hexdigest() == expected_fixture_sha256 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py new file mode 100644 index 000000000..cf92ca840 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_serving.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import httpx +import pytest + +from slurm_test_fakes import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend + + +def _client(endpoint: FakeLogicalEndpoint) -> httpx.Client: + return httpx.Client(transport=httpx.MockTransport(endpoint.handle), base_url=endpoint.endpoint) + + +def _start_and_publish(endpoint: FakeLogicalEndpoint) -> None: + endpoint.start() + for backend in endpoint.backends: + backend.mark_ready() + endpoint.publish() + + +def test_single_backend_startup_readiness_and_endpoint_publication() -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + + endpoint.start() + with _client(endpoint) as client: + assert client.get("/health").status_code == 503 + + backend.mark_ready() + assert endpoint.publish() == endpoint.endpoint + assert client.get("/health").status_code == 200 + assert client.post("/v1/chat/completions", json={"prompt": "fixture"}).json() == { + "backend": backend.endpoint, + "rank": 0, + } + + +def test_multiple_backends_receive_requests_in_deterministic_order( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + with _client(fake_logical_endpoint) as client: + responses = [client.post("/v1/chat/completions").json() for _ in range(3)] + + assert [response["rank"] for response in responses] == [0, 1, 0] + assert [len(backend.requests) for backend in fake_logical_endpoint.backends] == [2, 1] + + +def test_partial_multi_backend_startup_is_not_ready_and_cleans_up_once( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + fake_logical_endpoint.backends[0].mark_ready() + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.STARTING + with pytest.raises(RuntimeError, match="not ready"): + fake_logical_endpoint.publish() + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +@pytest.mark.parametrize("retry_after", ("2", None)) +def test_logical_endpoint_preserves_overload_response(retry_after: str | None) -> None: + headers = {"Retry-After": retry_after} if retry_after is not None else {} + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers=headers, json={"error": "overloaded"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers.get("Retry-After") == retry_after + assert response.json() == {"error": "overloaded"} + + +def test_logical_endpoint_retries_overload_against_another_backend() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, json={"error": "first overloaded"}),), + ) + second = FakeVllmBackend("http://127.0.0.1:31002", rank=1) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 200 + assert response.json()["rank"] == 1 + assert [len(backend.requests) for backend in endpoint.backends] == [1, 1] + + +def test_logical_endpoint_returns_the_final_overload_response() -> None: + first = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(httpx.Response(429, headers={"Retry-After": "1"}),), + ) + second = FakeVllmBackend( + "http://127.0.0.1:31002", + rank=1, + responses=(httpx.Response(429, headers={"Retry-After": "3"}),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (first, second)) + _start_and_publish(endpoint) + + with _client(endpoint) as client: + response = client.post("/v1/chat/completions") + + assert response.status_code == 429 + assert response.headers["Retry-After"] == "3" + + +def test_backend_rank_failure_fails_the_logical_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + fake_logical_endpoint.backends[1].fail("rank exited") + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert response.status_code == 503 + assert response.json() == {"error": "backend_failed"} + + +def test_endpoint_publication_failure_is_explicit_and_cleanup_remains_idempotent( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + fake_logical_endpoint.start() + for backend in fake_logical_endpoint.backends: + backend.mark_ready() + fake_logical_endpoint.script_publication_failure(RuntimeError("publication failed")) + + with pytest.raises(RuntimeError, match="publication failed"): + fake_logical_endpoint.publish() + assert fake_logical_endpoint.refresh_readiness() is FakeServingState.FAILED + assert fake_logical_endpoint.failure_reason == "endpoint_publication_failed" + + with _client(fake_logical_endpoint) as client: + response = client.post("/v1/chat/completions") + assert response.json() == {"error": "endpoint_publication_failed"} + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] + + +def test_backend_failure_injection_is_explicit() -> None: + backend = FakeVllmBackend( + "http://127.0.0.1:31001", + rank=0, + responses=(RuntimeError("scripted backend failure"),), + ) + endpoint = FakeLogicalEndpoint("http://127.0.0.1:31000", (backend,)) + _start_and_publish(endpoint) + + with _client(endpoint) as client, pytest.raises(RuntimeError, match="scripted backend failure"): + client.post("/v1/chat/completions") + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_backend_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + backend = FakeVllmBackend("http://127.0.0.1:31001", rank=0) + if terminal_state is FakeServingState.FAILED: + backend.fail("rank exited") + else: + backend.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + backend.start() + + +@pytest.mark.parametrize("terminal_state", (FakeServingState.FAILED, FakeServingState.STOPPED)) +def test_logical_endpoint_rejects_restart_from_terminal_state(terminal_state: FakeServingState) -> None: + endpoint = FakeLogicalEndpoint( + "http://127.0.0.1:31000", + (FakeVllmBackend("http://127.0.0.1:31001", rank=0),), + ) + if terminal_state is FakeServingState.FAILED: + endpoint.start() + endpoint.backends[0].fail("rank exited") + endpoint.refresh_readiness() + else: + endpoint.cleanup() + + with pytest.raises(RuntimeError, match=f"cannot start from {terminal_state.value}"): + endpoint.start() + + +def test_cleanup_is_idempotent_and_unpublishes_endpoint( + fake_logical_endpoint: FakeLogicalEndpoint, +) -> None: + _start_and_publish(fake_logical_endpoint) + + fake_logical_endpoint.cleanup() + fake_logical_endpoint.cleanup() + + assert fake_logical_endpoint.state is FakeServingState.STOPPED + assert fake_logical_endpoint.published_endpoint is None + assert fake_logical_endpoint.cleanup_calls == 2 + assert [backend.cleanup_calls for backend in fake_logical_endpoint.backends] == [1, 1] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py new file mode 100644 index 000000000..8bfb3de0e --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from data_designer.slurm.state import SchedulerIdentity +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" +SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") +SACCT_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") + + +def _submit(runner: FakeSlurmRunner) -> None: + completed = runner.run(("sbatch", "--parsable", "run.sbatch"), check=True) + assert completed.stdout == "4101\n" + + +def test_fake_slurm_runner_models_array_submission_and_active_states( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + observed = fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS), check=True) + + assert observed.stdout == (GOLDEN_DIRECTORY / "squeue_active.txt").read_text() + assert fake_slurm_runner.calls == [ + ("sbatch", "--parsable", "run.sbatch"), + ("squeue", *SQUEUE_ARGUMENTS), + ] + + +def test_fake_slurm_runner_exposes_terminal_accounting_precedence( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=0), + queue_state="RUNNING", + accounting_state="COMPLETED", + ) + fake_slurm_runner.set_task_state( + SchedulerIdentity(array_job_id=4101, array_task_id=1), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ) + + assert "4101_1|RUNNING" in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert ( + fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + == (GOLDEN_DIRECTORY / "sacct_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_models_accounting_lag_and_later_terminal_state( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + scheduler = SchedulerIdentity(array_job_id=4101, array_task_id=0) + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state=None, + ) + + assert "4101_0" not in fake_slurm_runner.run(("squeue", *SQUEUE_ARGUMENTS)).stdout + assert fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout == "" + + fake_slurm_runner.set_task_state( + scheduler, + queue_state=None, + accounting_state="COMPLETED", + ) + assert "4101_0|COMPLETED|0:0" in fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS)).stdout + + +def test_fake_slurm_runner_filters_job_scoped_queries() -> None: + first = FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)) + second = FakeSlurmArray( + tasks=( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4201, array_task_id=0), + queue_state="RUNNING", + accounting_state="FAILED", + exit_code="1:0", + ), + ) + ) + runner = FakeSlurmRunner((first, second)) + runner.run(("sbatch", "first.sbatch"), check=True) + runner.run(("sbatch", "second.sbatch"), check=True) + + first_queue = runner.run(("squeue", *SQUEUE_ARGUMENTS, "--jobs", "4101")).stdout + second_accounting = runner.run(("sacct", *SACCT_ARGUMENTS, "--jobs=4201")).stdout + + assert first_queue == "4101_0|PENDING\n" + assert second_accounting == "4201_0|FAILED|1:0\n" + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", *SQUEUE_ARGUMENTS, "--jobs", "--noheader"), + ("sacct", *SACCT_ARGUMENTS, "--jobs="), + ), +) +def test_fake_slurm_runner_rejects_malformed_job_selectors( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="malformed job selector"): + fake_slurm_runner.run(command) + + +@pytest.mark.parametrize("target", ("4101", "4101_1")) +def test_fake_slurm_runner_models_cancellation( + fake_slurm_runner: FakeSlurmRunner, + target: str, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("scancel", target), check=True).returncode == 0 + accounting = fake_slurm_runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + + expected_tasks = (0, 1) if target == "4101" else (1,) + for task_id in expected_tasks: + assert f"4101_{task_id}|CANCELLED|0:15" in accounting + + +def test_fake_slurm_runner_supports_malformed_output_and_command_failures( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + malformed = (GOLDEN_DIRECTORY / "squeue_malformed.txt").read_text() + fake_slurm_runner.script_next("squeue", FakeCommandResponse(stdout=malformed)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + fake_slurm_runner.script_next("sacct", FakeCommandResponse(stderr="accounting unavailable\n", returncode=2)) + + assert fake_slurm_runner.run(("squeue",)).stdout == malformed + assert fake_slurm_runner.run(("sacct",)).returncode == 2 + with pytest.raises(subprocess.CalledProcessError) as error: + fake_slurm_runner.run(("sacct",), check=True) + + assert error.value.stderr == "accounting unavailable\n" + fake_slurm_runner.assert_scripts_consumed() + + +@pytest.mark.parametrize( + ("arguments", "expected_stdout"), + ( + (("run.sbatch",), "Submitted batch job 4101\n"), + (("--parsable", "run.sbatch"), "4101\n"), + ), +) +def test_fake_slurm_runner_matches_sbatch_parsable_mode( + arguments: tuple[str, ...], + expected_stdout: str, +) -> None: + runner = FakeSlurmRunner( + (FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)),) + ) + + assert runner.run(("sbatch", *arguments), check=True).stdout == expected_stdout + + +@pytest.mark.parametrize( + "command", + ( + ("squeue", "--noheader"), + ("sacct", "--noheader", "--format=%i|%State|%ExitCode"), + ), +) +def test_fake_slurm_runner_rejects_underspecified_state_queries( + fake_slurm_runner: FakeSlurmRunner, + command: tuple[str, ...], +) -> None: + _submit(fake_slurm_runner) + + with pytest.raises(AssertionError, match="missing required arguments"): + fake_slurm_runner.run(command) + + +def test_fake_slurm_runner_exposes_retry_terminal_spellings() -> None: + states = ( + ("TIMEOUT", "0:125"), + ("NODE_FAIL", "1:0"), + ("PREEMPTED", "0:0"), + ("REQUEUED", "0:0"), + ("OUT_OF_MEMORY", "0:125"), + ("CANCELLED by 1234", "0:15"), + ) + runner = FakeSlurmRunner( + ( + FakeSlurmArray( + tasks=tuple( + FakeSlurmTask( + SchedulerIdentity(array_job_id=4101, array_task_id=task_id), + queue_state=None, + accounting_state=state, + exit_code=exit_code, + ) + for task_id, (state, exit_code) in enumerate(states) + ) + ), + ) + ) + _submit(runner) + + assert ( + runner.run(("sacct", *SACCT_ARGUMENTS), check=True).stdout + == (GOLDEN_DIRECTORY / "sacct_retry_terminal.txt").read_text() + ) + + +def test_fake_slurm_runner_bounds_sinfo_queries(fake_slurm_runner: FakeSlurmRunner) -> None: + response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%G"), check=True) + + assert response.stdout == (GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() + with pytest.raises(AssertionError, match="unexpected sinfo query"): + fake_slurm_runner.run(("sinfo", "--all")) + + +def test_fake_slurm_runner_rejects_unscripted_commands_and_submissions( + fake_slurm_runner: FakeSlurmRunner, +) -> None: + _submit(fake_slurm_runner) + + assert fake_slurm_runner.run(("sbatch", "second.sbatch")).returncode == 1 + with pytest.raises(AssertionError, match="unexpected command"): + fake_slurm_runner.run(("srun", "hostname")) diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json new file mode 100644 index 000000000..2b5f14067 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "failed" + } + ], + "revision": 4, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:02:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json new file mode 100644 index 000000000..ca41b2d6f --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_pending_readiness.json @@ -0,0 +1,20 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": null, + "model_alias": "primary", + "ready_backends": 0, + "state": "pending" + } + ], + "revision": 1, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "pending", + "updated_at": "2026-08-18T12:00:02Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json new file mode 100644 index 000000000..5001bc9fb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_publication_failed_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "failed", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:30Z", + "outcome": "failure", + "reason_code": "endpoint_publication_failed", + "redacted_message": "Logical endpoint publication failed" + }, + "model_alias": "primary", + "ready_backends": 1, + "state": "failed" + } + ], + "revision": 3, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "failed", + "updated_at": "2026-08-18T12:01:30Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json new file mode 100644 index 000000000..7929011bb --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_starting_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "pending", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:01:00Z", + "outcome": "failure", + "reason_code": "backend_starting", + "redacted_message": "Backend is starting" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "starting" + } + ], + "revision": 2, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "starting", + "updated_at": "2026-08-18T12:01:00Z" +} diff --git a/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json new file mode 100644 index 000000000..894019149 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/golden/single_node_stopped_readiness.json @@ -0,0 +1,25 @@ +{ + "attempt_id": "attempt-0001", + "deployments": [ + { + "deployment_id": "deployment-00000", + "endpoint_publication": "published", + "expected_backends": 1, + "last_probe": { + "observed_at": "2026-08-18T12:02:00Z", + "outcome": "failure", + "reason_code": "backend_failed", + "redacted_message": "Backend rank failed" + }, + "model_alias": "primary", + "ready_backends": 0, + "state": "stopped" + } + ], + "revision": 5, + "run_id": "run-0001", + "schema_version": 1, + "shard_id": "shard-00000", + "state": "stopped", + "updated_at": "2026-08-18T12:02:01Z" +} diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index e029b99b4..d7b5d1adf 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -25,7 +25,12 @@ ("run_manifest.json", RunManifest), ("shard_manifest.json", ShardManifest), ("successful_attempt.json", AttemptManifest), + ("single_node_pending_readiness.json", AttemptReadiness), + ("single_node_starting_readiness.json", AttemptReadiness), ("single_node_readiness.json", AttemptReadiness), + ("single_node_publication_failed_readiness.json", AttemptReadiness), + ("single_node_failed_readiness.json", AttemptReadiness), + ("single_node_stopped_readiness.json", AttemptReadiness), ("multi_node_readiness.json", AttemptReadiness), ("failed_attempt.json", AttemptManifest), ("stale_readiness.json", AttemptReadiness), diff --git a/packages/data-designer-slurm/tests/state/test_transition_fixtures.py b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py new file mode 100644 index 000000000..4bc1c7188 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_transition_fixtures.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from itertools import pairwise +from pathlib import Path + +from data_designer.slurm.state import AttemptReadiness, validate_readiness_transition + +GOLDEN_DIRECTORY = Path(__file__).parent / "golden" +TRANSITION_FILES = ( + "single_node_pending_readiness.json", + "single_node_starting_readiness.json", + "single_node_readiness.json", + "single_node_failed_readiness.json", + "single_node_stopped_readiness.json", +) + + +def test_single_node_readiness_transition_fixture_is_canonical() -> None: + records = tuple( + AttemptReadiness.model_validate_json((GOLDEN_DIRECTORY / filename).read_text()) for filename in TRANSITION_FILES + ) + + assert tuple(record.revision for record in records) == (1, 2, 3, 4, 5) + for previous, current in pairwise(records): + assert validate_readiness_transition(previous, current) is current + + +def test_endpoint_publication_failure_transition_fixture_is_canonical() -> None: + starting = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_starting_readiness.json").read_text() + ) + failed = AttemptReadiness.model_validate_json( + (GOLDEN_DIRECTORY / "single_node_publication_failed_readiness.json").read_text() + ) + + assert validate_readiness_transition(starting, failed) is failed From 41d544450699ed1e55653337806f1b896a7eaa60 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 14:51:15 -0600 Subject: [PATCH 08/62] feat: add Slurm image registry foundation (#891) * feat: add Slurm image registry foundation Add digest-bound client and serving inspection, atomic workspace-derived registry persistence, and verified alias/path resolution.\n\nPart of #867 Signed-off-by: Nabin Mulepati * test: cover image inspection boundaries Exercise the production environment reader and verify restrictive atomic registry persistence discovered during review.\n\nPart of #867 Signed-off-by: Nabin Mulepati * fix: align image registry persistence Persist the package-owned image registry as deterministic safe YAML at the migration plan's canonical path and declare the runtime YAML dependency explicitly.\n\nPart of #867. Signed-off-by: Nabin Mulepati * fix: harden image registry file boundaries Reject symlinked registry and lock files, use explicit UTF-8 persistence, and normalize malformed byte content into the package image error hierarchy.\n\nPart of #867. Signed-off-by: Nabin Mulepati * test: verify Slurm image wheel integration Exercise the public image service from an isolated installed wheel and assert that the wheel metadata carries the runtime PyYAML dependency.\n\nPart of #867. Signed-off-by: Nabin Mulepati * fix: verify registry file types before use Treat only a genuinely missing registry as empty and reject non-regular registry or lock nodes before reading, locking, or changing permissions.\n\nPart of #867. Signed-off-by: Nabin Mulepati * fix: rollback stale image registrations Reverify image bytes while registry mutation locks remain held and restore the previous snapshot if final verification fails. Part of #867 Signed-off-by: Nabin Mulepati * fix image registry review findings Require Data Designer in client image inspections, verify artifacts before publishing registry updates, harden storage directory boundaries, and normalize recursive YAML failures. * harden image registry contracts - Keep the image-registry foundation internal until the public facade lands\n- Require a complete client package inventory and unambiguous registry schema\n- Close SQSH identity races and exercise publication behavior through owned APIs * strengthen image registry test coverage Exercise persisted semantic invariants and exact-byte hashes. Split record, store, and service coverage by ownership. Signed-off-by: Nabin Mulepati * address image registry review feedback Bind registry transactions to verified directory descriptors, centralize immutable image fact comparisons, and verify internal image modules from the installed wheel. Signed-off-by: Nabin Mulepati * clarify image test helper names Use explicit get prefixes for value-producing test helpers while preserving action verbs for creation, inspection, and filesystem operations. Signed-off-by: Nabin Mulepati * split registry filesystem helpers Move descriptor-bound filesystem operations into a focused internal module while keeping registry transactions and YAML semantics together. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- packages/data-designer-slurm/pyproject.toml | 1 + .../data_designer/slurm/images/__init__.py | 6 + .../src/data_designer/slurm/images/errors.py | 30 ++ .../data_designer/slurm/images/filesystem.py | 163 ++++++++ .../data_designer/slurm/images/inspection.py | 179 +++++++++ .../src/data_designer/slurm/images/records.py | 75 ++++ .../data_designer/slurm/images/registry.py | 302 ++++++++++++++ .../src/data_designer/slurm/images/service.py | 153 +++++++ .../golden/client_image_inspection.json | 12 + .../contracts/golden/dependency_lock.json | 12 + .../golden/dependency_lock_single.json | 16 + .../contracts/golden/multi_node_plan.json | 14 +- .../contracts/golden/single_node_plan.json | 18 +- .../images/test_image_registry_records.py | 49 +++ .../tests/images/test_inspection.py | 234 +++++++++++ .../tests/images/test_registry.py | 333 +++++++++++++++ .../tests/images/test_registry_store.py | 380 ++++++++++++++++++ .../golden/finalization_chain.json | 4 +- .../tests/slurm_test_fakes/__init__.py | 2 + .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../tests/slurm_test_fakes/images.py | 54 +++ .../slurm_test_fakes/test_rendered_scripts.py | 4 +- scripts/test_slurm_package_install.py | 4 + uv.lock | 2 + 25 files changed, 2043 insertions(+), 8 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/registry.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/service.py create mode 100644 packages/data-designer-slurm/tests/images/test_image_registry_records.py create mode 100644 packages/data-designer-slurm/tests/images/test_inspection.py create mode 100644 packages/data-designer-slurm/tests/images/test_registry.py create mode 100644 packages/data-designer-slurm/tests/images/test_registry_store.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/images.py diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 340c8f289..9e4610eff 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "data-designer=={{ version }}", "packaging>=25,<27", "pydantic>=2.9.2,<3", + "pyyaml>=6.0.1,<7", ] [tool.hatch.build.targets.wheel] 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..7a6d426bf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py @@ -0,0 +1,30 @@ +# 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 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.""" 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..36441227f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py @@ -0,0 +1,163 @@ +# 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 fcntl +import os +import secrets +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +from data_designer.slurm.images.errors import ImageRegistryError + +_DIRECTORY_MODE = 0o700 +_FILE_MODE = 0o600 + + +@contextmanager +def acquire_file_lock(directory_descriptor: int, name: str, display_path: Path) -> Iterator[None]: + """Acquire one exclusive advisory lock without reopening its parent path.""" + descriptor: int | None = None + try: + # macOS can report ENOENT to the losing openat(O_CREAT) caller when two + # processes create the same lock file concurrently. Retrying then opens + # the winner's regular file with the same no-follow boundary. + for _ in range(10): + try: + descriptor = os.open( + name, + os.O_CREAT | os.O_RDWR | getattr(os, "O_NOFOLLOW", 0), + _FILE_MODE, + dir_fd=directory_descriptor, + ) + except FileNotFoundError: + continue + break + if descriptor is None: + raise OSError(f"lock path {display_path} could not be opened") + if not stat.S_ISREG(os.fstat(descriptor).st_mode): + raise OSError(f"lock path {display_path} is not a regular file") + os.fchmod(descriptor, _FILE_MODE) + fcntl.flock(descriptor, fcntl.LOCK_EX) + except OSError as error: + if descriptor is not None: + os.close(descriptor) + raise ImageRegistryError(f"cannot lock image registry target {display_path}") from error + + try: + yield + finally: + assert descriptor is not None + try: + fcntl.flock(descriptor, fcntl.LOCK_UN) + finally: + os.close(descriptor) + + +def ensure_private_directory(path: Path, *, parents: bool) -> None: + """Create or restrict one real directory to its owning user.""" + path.mkdir(mode=_DIRECTORY_MODE, parents=parents, exist_ok=True) + with open_verified_directory(path) as descriptor: + os.fchmod(descriptor, _DIRECTORY_MODE) + + +@contextmanager +def open_verified_directory(path: Path) -> Iterator[int]: + """Open one real directory and reject path 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"registry 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) + if (before_open.st_dev, before_open.st_ino) != (after_open.st_dev, after_open.st_ino): + raise OSError(f"registry directory {path} changed while it was being opened") + if not stat.S_ISDIR(after_open.st_mode): + raise OSError(f"registry directory {path} is not a directory") + 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, +) -> Iterator[int]: + """Open one real child directory relative to an already 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"registry 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) + if (before_open.st_dev, before_open.st_ino) != (after_open.st_dev, after_open.st_ino): + raise OSError(f"registry directory {display_path} changed while it was being opened") + if not stat.S_ISDIR(after_open.st_mode): + raise OSError(f"registry directory {display_path} is not a directory") + yield descriptor + finally: + if descriptor is not None: + os.close(descriptor) + + +def create_temporary_file( + directory_descriptor: int, + *, + prefix: str, + suffix: str, +) -> tuple[int, str]: + """Create one restrictive 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), + _FILE_MODE, + dir_fd=directory_descriptor, + ) + except FileExistsError: + continue + try: + os.fchmod(descriptor, _FILE_MODE) + except OSError: + os.close(descriptor) + raise + return descriptor, name + raise OSError("cannot allocate a unique temporary file name") + + +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), + 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: + return registry_file.read() + 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..d8e1fbcfc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py @@ -0,0 +1,179 @@ +# 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 shutil +import sys +from pathlib import Path +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 + +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_version(self, name: str) -> str: + """Return one installed distribution version.""" + + def find_executable(self, name: str) -> str: + """Return one absolute executable 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_version(self, name: str) -> str: + """Return one installed distribution version with a canonical missing-package error.""" + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError as error: + raise ImageInspectionError(f"required distribution {name!r} is not installed") from error + + def find_executable(self, name: str) -> str: + """Return one resolved executable path with a canonical missing-tool error.""" + path = shutil.which(name) + if path is None: + raise ImageInspectionError(f"required executable {name!r} is not installed") + executable = Path(path) + if not executable.is_absolute(): + raise ImageInspectionError(f"required executable {name!r} did not resolve to an absolute path") + return executable.as_posix() + + +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}") + 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=environment.find_executable("pip"), + installer_version=versions_by_name["pip"], + ) + 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: + inspection = ServingImageInspection( + kind="serving", + server_type="vllm", + runtime_version=environment.get_distribution_version("vllm"), + executable_path=environment.find_executable("vllm"), + ) + 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/records.py b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py new file mode 100644 index 000000000..4e6589ec1 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable records persisted by the Slurm image registry.""" + +from __future__ import annotations + +from pydantic import field_validator, model_validator + +from data_designer.slurm.config import ImageInspectionRecord, ImageKind +from data_designer.slurm.contracts import ( + ContractRecord, + ContractValue, + Identifier, + SchemaVersion, + Sha256Digest, + validate_absolute_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..dd3e3ba38 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py @@ -0,0 +1,302 @@ +# 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 json +import os +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.images.errors import ( + 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" +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class ImageRegistryStore: + """Persist image aliases beneath one explicitly selected shared workspace.""" + + 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 + + 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], + replace: bool = False, + ) -> RegisteredImage: + """Atomically add or explicitly replace one verified image alias.""" + self._ensure_storage() + alias_lock_name = f"alias-{image.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, + ), + ): + snapshot = self._load_from_directory(image_root_descriptor) + 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") + self._validate_path_facts(snapshot, image, replacing=existing) + images = tuple(entry for entry in snapshot.images if entry.name != image.name) + (image,) + updated = ImageRegistryDocument( + schema_version=1, + images=tuple(sorted(images, key=lambda entry: entry.name)), + ) + self._save( + image_root_descriptor, + updated, + verify_before_publish=lambda: verify_before_publish(image), + ) + return image + + def unregister(self, name: Identifier) -> RegisteredImage: + """Atomically remove an alias without deleting its underlying SQSH artifact.""" + name = _validate_image_alias(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, + ), + ): + 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, updated) + return removed + + def _load(self) -> ImageRegistryDocument: + try: + with open_verified_directory(self._image_root) as image_root_descriptor: + return self._load_from_directory(image_root_descriptor) + except FileNotFoundError: + return ImageRegistryDocument(schema_version=1) + 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, + snapshot: ImageRegistryDocument, + *, + verify_before_publish: Callable[[], None] | None = None, + ) -> None: + temporary_name: str | None = None + descriptor: int | None = None + try: + descriptor, temporary_name = create_temporary_file( + image_root_descriptor, + prefix=".registry.", + 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()) + 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, + ) + os.fsync(image_root_descriptor) + except (OSError, TypeError, yaml.YAMLError) as error: + raise ImageRegistryError(f"cannot persist image registry {self._registry_path}") from error + finally: + 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 + + 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 _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 _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 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..d04819708 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -0,0 +1,153 @@ +# 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 pathlib import Path + +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, ImageRef +from data_designer.slurm.contracts import Sha256Digest +from data_designer.slurm.images.errors import ImageVerificationError +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.planning import ResolvedImage + +_HASH_BLOCK_SIZE = 1024 * 1024 + + +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) + _validate_inspection(inspection, expected_kind=ImageKind(request.kind)) + image = RegisteredImage( + schema_version=1, + name=request.name, + path=path.as_posix(), + sqsh_sha256=inspection.sqsh_sha256, + inspection=inspection, + ) + return self._registry.register(image, verify_before_publish=_verify_registered_image, 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, + ) + return ResolvedImage( + authored_ref=reference, + path=image.path, + sha256=image.sqsh_sha256, + inspection=image.inspection, + ) + + def unregister(self, name: str) -> RegisteredImage: + """Remove an alias while leaving its underlying artifact untouched.""" + return self._registry.unregister(name) + + +def compute_sqsh_file_sha256(path: Path) -> Sha256Digest: + """Compute the SHA-256 of one regular, non-symlink SQSH file.""" + descriptor: int | None = None + try: + before_open = path.lstat() + if not stat.S_ISREG(before_open.st_mode): + raise ImageVerificationError(f"image path {path} must be a regular non-symlink file") + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + image_file = os.fdopen(descriptor, "rb") + descriptor = None + digest = hashlib.sha256() + with image_file: + before_read = os.fstat(image_file.fileno()) + if (before_open.st_dev, before_open.st_ino) != (before_read.st_dev, before_read.st_ino): + raise ImageVerificationError(f"image path {path} changed while it was being opened") + while block := image_file.read(_HASH_BLOCK_SIZE): + digest.update(block) + after_read = os.fstat(image_file.fileno()) + after_path = path.lstat() + before_facts = ( + before_read.st_dev, + before_read.st_ino, + before_read.st_size, + before_read.st_mtime_ns, + before_read.st_ctime_ns, + ) + after_facts = ( + after_read.st_dev, + after_read.st_ino, + after_read.st_size, + after_read.st_mtime_ns, + after_read.st_ctime_ns, + ) + if before_facts != after_facts: + raise ImageVerificationError(f"image path {path} changed while it was being verified") + path_facts = ( + after_path.st_dev, + after_path.st_ino, + after_path.st_size, + after_path.st_mtime_ns, + after_path.st_ctime_ns, + ) + if not stat.S_ISREG(after_path.st_mode) or path_facts != after_facts: + raise ImageVerificationError(f"image path {path} changed while it was being verified") + return digest.hexdigest() + except FileNotFoundError as error: + raise ImageVerificationError(f"image path {path} does not exist") from error + except OSError as error: + raise ImageVerificationError(f"cannot read image path {path}") from error + finally: + if descriptor is not None: + os.close(descriptor) + + +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 _verify_registered_image(image: RegisteredImage) -> None: + actual_sha256 = compute_sqsh_file_sha256(Path(image.path)) + if actual_sha256 != image.sqsh_sha256: + raise ImageVerificationError(f"registered image {image.name!r} no longer matches its SQSH digest") 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 index 6466182d7..5e83b994b 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -5,6 +5,18 @@ "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" diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json index 87190b04c..d776429b0 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -9,6 +9,18 @@ "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" 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 index b86f29d6c..d7d0836c3 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -6,6 +6,22 @@ { "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" } ], "overlay_packages": [], 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 index 0f2fd69b1..baf366891 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -56,7 +56,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "a86032b310aa6bdb95fc7f35ef606fec2e56b977ba60e36b8b9293341ac43e00" + "sha256": "fb536259ea0c206eca90ee8a4ad0a95345f97eb869e349732d1d19db857ecbc3" }, "gpu_count": 0, "host_node_index": 0, @@ -72,6 +72,18 @@ "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" 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 index f0794cc75..cfc0b5824 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -43,7 +43,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "2600c1944897a8e3c9a9410fe912bdee0e169e09769e58180cadebe94f3da52c" + "sha256": "ce671fcd812110bc09bd39a21c2c294754c8255d42030c1b90e250a1b513a311" }, "gpu_count": 0, "host_node_index": 0, @@ -58,6 +58,22 @@ { "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", 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..5d12cad12 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_inspection.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 importlib.metadata +import re +import shutil +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, +) + + +def _get_client_distributions() -> tuple[InstalledDistribution, ...]: + return ( + 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_uses_installer_version_from_distribution_inventory() -> None: + environment = FakeInspectionEnvironment( + distributions=_get_client_distributions(), + distribution_versions={"pip": "unexpected"}, + executables={"pip": "/usr/bin/pip"}, + ) + + inspection = ClientImageInspector(environment).inspect("a" * 64).inspection + + assert inspection.installer_version == "26.1" + + +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, + "version", + Mock(side_effect=importlib.metadata.PackageNotFoundError("missing")), + ) + + with pytest.raises(ImageInspectionError, match="not installed"): + SystemInspectionEnvironment().get_distribution_version("missing") + + +@pytest.mark.parametrize("path", (None, "relative/pip"), ids=("missing", "relative")) +def test_system_inspection_environment_requires_absolute_executable( + monkeypatch: pytest.MonkeyPatch, + path: str | None, +) -> None: + monkeypatch.setattr(shutil, "which", Mock(return_value=path)) + + with pytest.raises(ImageInspectionError): + SystemInspectionEnvironment().find_executable("pip") + + +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_registry.py b/packages/data-designer-slurm/tests/images/test_registry.py new file mode 100644 index 000000000..662466ebf --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_registry.py @@ -0,0 +1,333 @@ +# 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" + + +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.images.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.images.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() -> FakeInspectionEnvironment: + return 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"}, + ) + + +def _inspect_client(path: Path) -> ImageInspectionRecord: + return ClientImageInspector(_get_client_environment()).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..55f92b6c2 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import stat +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier, Event + +import pytest +import yaml +from slurm_test_fakes import FakeInspectionEnvironment + +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, InstalledDistribution +from data_designer.slurm.images.errors import ImageConflictError, ImageRegistryError +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_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_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) + + +@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")) + + +def test_registration_remains_unpublished_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=1) as executor: + future = executor.submit( + store.register, + image, + verify_before_publish=verify_before_publish, + ) + assert verification_started.wait(timeout=5) + assert ImageRegistryStore(workspace).list_images() == () + finish_verification.set() + assert future.result().name == "client" + + assert tuple(registered.name for registered in ImageRegistryStore(workspace).list_images()) == ("client",) + + +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 _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/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 517a45ac5..e3b564f14 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + "sha256": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" + "sha256": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py index a73536ea8..0b85d89f3 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -7,6 +7,7 @@ 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.serving import FakeLogicalEndpoint, FakeServingState, FakeVllmBackend from slurm_test_fakes.slurm import ( FakeCommandResponse, @@ -20,6 +21,7 @@ "FakeCommandResponse", "FakeDependencyInstaller", "FakeDependencyResolver", + "FakeInspectionEnvironment", "FakeLogicalEndpoint", "FakeServingState", "FakeSlurmArray", diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 5a78c97b7..0cf7a8e85 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="dbcae1da6ce8ef6799faf2add8e08b1b093390ae08a37dd13e336da522b9a3fd" +readonly DD_PLAN_SHA256="fff7573a8b6eb2f71ae12f27900a85e1e121debfe58b0e94fefc3fb9362970ce" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index e10478f7d..60751e07c 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -11,7 +11,7 @@ set -Eeuo pipefail readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="d9155dd04ad439b4cd091eaec9b7b4d938bbdf133a03b4a751defb8ea9329425" +readonly DD_PLAN_SHA256="8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/images.py b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py new file mode 100644 index 000000000..df0d8d278 --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py @@ -0,0 +1,54 @@ +# 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_version(self, name: str) -> str: + """Return one configured distribution version.""" + try: + return self.distribution_versions[name] + except KeyError: + raise ImageInspectionError(f"required distribution {name!r} is not installed") from None + + def find_executable(self, name: str) -> str: + """Return one configured executable path.""" + try: + return self.executables[name] + except KeyError: + raise ImageInspectionError(f"required executable {name!r} is not installed") from None diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index c7d119389..1d793f500 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="faa1dac9b9b0423423c9d06a13d71bee339932ffa45d1e02a8a95012a7934520", + expected_fixture_sha256="cdb5204e8ce7a7affda99ea9c121c6b60a16bf42c510c8eb6bc57661969620b9", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="a4542e56b124c2346d0a92ebadc4e89b45d94c9355d7b86a4a1b05180d331a48", + expected_fixture_sha256="c6708cdbc0a03e095062c0642cfd141f066153b958b7ad3dec779afd6414fa34", ) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index f6ba85e19..f80247355 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -129,6 +129,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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 @@ -137,6 +138,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non from data_designer.slurm.state import ResumeWorkspace as StateResumeWorkspace from data_designer.slurm.state import RunManifest assert RunManifest.__name__ == "RunManifest" +assert ImageRegistryStore.__name__ == "ImageRegistryStore" assert PlanningArtifactReference is ContractArtifactReference assert PlanningRecordRange is ContractRecordRange assert PlanningResumeWorkspace is ContractResumeWorkspace @@ -190,10 +192,12 @@ def main() -> None: leaf_base_requirement = requirement(leaf_metadata, "data-designer") leaf_packaging_requirement = requirement(leaf_metadata, "packaging") leaf_pydantic_requirement = requirement(leaf_metadata, "pydantic") + leaf_pyyaml_requirement = requirement(leaf_metadata, "pyyaml") 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_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier + assert leaf_pyyaml_requirement.specifier == Requirement("pyyaml>=6.0.1,<7").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) assert not base_leaf_requirement.marker.evaluate({"extra": ""}) diff --git a/uv.lock b/uv.lock index 2e6000e20..a995c306b 100644 --- a/uv.lock +++ b/uv.lock @@ -954,6 +954,7 @@ dependencies = [ { name = "data-designer" }, { name = "packaging" }, { name = "pydantic" }, + { name = "pyyaml" }, ] [package.metadata] @@ -961,6 +962,7 @@ requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, { name = "packaging", specifier = ">=25,<27" }, { name = "pydantic", specifier = ">=2.9.2,<3" }, + { name = "pyyaml", specifier = ">=6.0.1,<7" }, ] [[package]] From 7fbf347c85e196df344b8fc46a0fc465b3e6242b Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 27 Aug 2026 14:58:08 -0600 Subject: [PATCH 09/62] feat: add Slurm command client and renderer (#892) * feat: add Slurm command client and renderer Part of #868 Signed-off-by: Nabin Mulepati * fix: tighten Slurm launcher boundaries Part of #868 Signed-off-by: Nabin Mulepati * fix: use native sacct output fields Part of #868 Signed-off-by: Nabin Mulepati * test: require expanded Slurm array queries Part of #868 Signed-off-by: Nabin Mulepati * fix: sanitize Slurm command diagnostics Part of #868 Signed-off-by: Nabin Mulepati * fix: reject invalid visible GPU memory Part of #868 Signed-off-by: Nabin Mulepati * fix: hash rendered inputs without manifests Part of #868 Signed-off-by: Nabin Mulepati * fix: make Slurm output decoding stable Part of #868 Signed-off-by: Nabin Mulepati * fix: validate finite command timeouts Part of #868 Signed-off-by: Nabin Mulepati * fix: harden Slurm boundary inputs Part of #868 Signed-off-by: Nabin Mulepati * fix: reject malformed Slurm GRES lists Part of #868 Signed-off-by: Nabin Mulepati * fix: isolate submitted Slurm environments Part of #868 Signed-off-by: Nabin Mulepati * test: verify launcher wheel import Part of #868 Signed-off-by: Nabin Mulepati * fix: pin batch host tool path Part of #868 Signed-off-by: Nabin Mulepati * test: cover Slurm launcher boundaries Part of #868 Signed-off-by: Nabin Mulepati * fix: correlate Slurm query responses Part of #868 Signed-off-by: Nabin Mulepati * fix: harden Slurm command boundaries Preserve the caller's PATH as the only ambient lookup input for default Slurm commands while continuing to isolate all other environment variables. Normalize oversized numeric scheduler fields into the launcher parse-error boundary. Part of #868 Signed-off-by: Nabin Mulepati * fix: fall back from empty command path Use the platform default search path when the ambient PATH is absent or empty so bare Slurm executables remain resolvable. Part of #868 Signed-off-by: Nabin Mulepati * fix: align Slurm launcher semantics * fix Slurm job observation contracts Support both ordinary jobs and array-task observations while preserving accounting-lag semantics. Bound scheduler numeric fields and diagnostics, and honor unthrottled plan arrays when concurrency is omitted. * refine Slurm launcher boundaries - Internalize implementation APIs and clarify transient types - Harden parsing, rendering, and process error normalization - Extend fakes and executable checksum coverage Part of #868 * fix Slurm selector filtering ownership * simplify Slurm launcher structure --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/launcher/__init__.py | 6 + .../data_designer/slurm/launcher/client.py | 215 ++++++++++++++ .../data_designer/slurm/launcher/errors.py | 22 ++ .../data_designer/slurm/launcher/models.py | 45 +++ .../data_designer/slurm/launcher/parsing.py | 229 +++++++++++++++ .../data_designer/slurm/launcher/renderer.py | 132 +++++++++ .../data_designer/slurm/launcher/runner.py | 66 +++++ .../tests/launcher/test_client.py | 277 ++++++++++++++++++ .../tests/launcher/test_parsing.py | 204 +++++++++++++ .../tests/launcher/test_renderer.py | 234 +++++++++++++++ .../tests/launcher/test_runner.py | 97 ++++++ .../tests/slurm_test_fakes/__init__.py | 2 + .../golden/rendered/multi_node.sbatch | 15 +- .../golden/rendered/single_node.sbatch | 15 +- .../tests/slurm_test_fakes/slurm.py | 130 ++++++-- .../slurm_test_fakes/test_rendered_scripts.py | 9 +- .../tests/slurm_test_fakes/test_slurm.py | 33 ++- 17 files changed, 1695 insertions(+), 36 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_client.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_parsing.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_renderer.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_runner.py 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/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py new file mode 100644 index 000000000..3c0d3e858 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -0,0 +1,215 @@ +# 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 re +import subprocess +import unicodedata +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TypeAlias + +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmQueueEntry, +) +from data_designer.slurm.launcher.parsing import ( + parse_accounting, + parse_gpu_counts, + parse_queue, + parse_submission, +) +from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner +from data_designer.slurm.state import SchedulerIdentity + +_JobSelector: TypeAlias = int | SchedulerIdentity +_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_MAX_SLURM_INTEGER = (1 << 32) - 1 + + +@dataclass(frozen=True) +class SlurmExecutables: + """Executable paths used for bounded Slurm operations.""" + + sbatch: str = "sbatch" + squeue: str = "squeue" + sacct: str = "sacct" + scancel: str = "scancel" + sinfo: str = "sinfo" + + def __post_init__(self) -> None: + for executable in (self.sbatch, self.squeue, self.sacct, self.scancel, 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) -> 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 './'") + output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) + return parse_submission(output) + + def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: + """Return normalized active-queue rows for explicit managed jobs.""" + requested = tuple(selectors) + jobs = _format_selectors(requested) + output = self._run( + ( + self._executables.squeue, + "--noheader", + "--array", + "--format=%i|%T", + f"--jobs={jobs}", + ) + ) + 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[_JobSelector]) -> 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 cancel(self, selector: _JobSelector) -> None: + """Cancel one managed Slurm job, array, or array task.""" + self._run((self._executables.scancel, _format_selector(selector))) + + def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, ...]: + """Return configured GPU counts reported for eligible node groups.""" + command = [self._executables.sinfo, "--noheader", "--format=%G"] + if partition is not None: + if type(partition) is not str or _IDENTIFIER_PATTERN.fullmatch(partition) is None: + raise ValueError("Slurm partition must be a valid identifier") + command.append(f"--partition={partition}") + return parse_gpu_counts(self._run(command)) + + def _run(self, command: Sequence[str]) -> str: + command_name = Path(command[0]).name + try: + completed = self._runner.run(command) + except (OSError, subprocess.SubprocessError) as error: + raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") 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") + 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_selectors(selectors: Sequence[_JobSelector]) -> 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: _JobSelector) -> 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[SlurmObservedJobIdentity], + selectors: Sequence[_JobSelector], + *, + command: str, +) -> frozenset[SlurmObservedJobIdentity]: + """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[SlurmObservedJobIdentity] = 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: + sanitized = "".join(" " if unicodedata.category(character).startswith("C") else character for character in value) + normalized = " ".join(sanitized.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__ 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..4a611152b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -0,0 +1,22 @@ +# 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.""" + + +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..7d1168255 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -0,0 +1,45 @@ +# 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 typing import TypeAlias + +from data_designer.slurm.state import SchedulerIdentity, SchedulerState + +SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity + + +@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: SlurmObservedJobIdentity + state: SchedulerState + + +@dataclass(frozen=True) +class SlurmAccountingEntry: + """One transient normalized accounting entry.""" + + job_identity: SlurmObservedJobIdentity + state: SchedulerState + process_exit_code: SlurmProcessExitCode 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..adf0e13c2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -0,0 +1,229 @@ +# 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 data_designer.slurm.launcher.errors import SlurmCommandOutputError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmObservedJobIdentity, + SlurmProcessExitCode, + SlurmQueueEntry, +) +from data_designer.slurm.state import SchedulerIdentity, 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[SlurmObservedJobIdentity] = 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[SlurmObservedJobIdentity] = 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_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_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) -> SlurmObservedJobIdentity: + 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: SlurmObservedJobIdentity, + identities: set[SlurmObservedJobIdentity], + *, + 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..203cfe14a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -0,0 +1,132 @@ +# 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 +import re +from dataclasses import dataclass + +from data_designer.slurm.launcher.errors import SlurmBatchRenderError +from data_designer.slurm.planning import ResolvedSlurmRunPlan + +_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_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") + directives = _build_generation_directives(plan) + directive_text = "\n".join(directive.render() for directive in directives) + attempt = f"{attempt_ordinal:04d}" + + return f"""#!/usr/bin/env bash +{directive_text} +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +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}}" +install -d -m 0700 "${{DD_ATTEMPT_DIR}}" +DD_RUNTIME_DIR="$(mktemp -d "${{DD_ATTEMPT_DIR}}/runtime.${{DD_RUNTIME_SHA256}}.XXXXXX")" +readonly DD_RUNTIME_DIR +tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_DIR}}" + +source "${{DD_RUNTIME_DIR}}/entrypoint.sh" +dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" +""" + + +def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDirective, ...]: + 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: + array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + + values: list[tuple[str, str | None]] = [ + ("job-name", plan.submission.job_name), + ("account", plan.submission.account), + ("partition", plan.submission.partition), + ("nodes", str(node_count)), + ("cpus-per-task", str(plan.client.authored.cpus)), + ("time", plan.submission.time_limit), + ("array", array), + ] + profile = plan.selected_profile.profile + if profile.gpu_request_mode == "gres": + values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) + elif profile.scheduler.mem_per_gpu is not None: + # TODO(#875): Remove this defense once config and plan validation make this state unrepresentable. + raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") + 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(_BatchDirective(name=name, value=value) for name, value in values if value is not None) + + +def _quote_sbatch_option_value(value: str) -> str: + _reject_control_characters(value, field_name="batch directive value") + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' + + +def _quote_shell_value(value: str) -> str: + _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: + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SlurmBatchRenderError(f"{field_name} must not contain control characters") 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..3679d3252 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -0,0 +1,66 @@ +# 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]) -> 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") + explicit_environment = ( + dict(environment) if environment is not None else {"PATH": os.environ.get("PATH") or os.defpath} + ) + 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]) -> subprocess.CompletedProcess[str]: + """Execute an argument vector with captured text output.""" + return subprocess.run( + tuple(command), + check=False, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + env=dict(self._environment), + timeout=self._timeout_seconds, + ) 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..5155ddb02 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -0,0 +1,277 @@ +# 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 Sequence + +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 +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=NIL", "/workspace/run.sbatch"), + ("squeue", "--noheader", "--array", "--format=%i|%T", "--jobs=4101"), + ] + + +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_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_bounded_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=%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_rejects_invalid_gpu_partition_without_running_command(fake_slurm_runner: FakeSlurmRunner) -> None: + client = SlurmCommandClient(fake_slurm_runner) + + with pytest.raises(ValueError, match="valid identifier"): + client.query_gpu_counts(partition="batch,other") + + 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,)) + + +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) + + +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=NIL", + "/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("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 _TimeoutRunner: + def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + 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_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py new file mode 100644 index 000000000..7509f93c2 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -0,0 +1,204 @@ +# 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_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 + + +@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..c80c95f5f --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.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 +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) + + assert render_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() + + +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 --comment="safe test run"\n' in script + + +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) + + +@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( + single_node_plan: ResolvedSlurmRunPlan, + gpu_request_mode: Literal["gres", "visible"], +) -> 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": 17})} + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile), "client": client}) + + assert "#SBATCH --cpus-per-task=17\n" in render_generation_attempt_script(plan, attempt_ordinal=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"), + } + ) + plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(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()) <= 42 + 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() + 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" + with tarfile.open(runtime_archive_path, mode="w:gz") as runtime_archive: + runtime_archive.add(entrypoint_path, arcname="entrypoint.sh") + + 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(), + ), + } + ) + 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"} + + 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..6bab64a83 --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -0,0 +1,97 @@ +# 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_search_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "/workspace/slurm/bin:/usr/bin") + monkeypatch.setenv("SECRET", "must-not-leak") + + runner = SubprocessRunner() + + assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} + + +def test_subprocess_runner_default_environment_replaces_empty_search_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PATH", "") + + 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/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py index 0b85d89f3..60d45aa06 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -12,6 +12,7 @@ from slurm_test_fakes.slurm import ( FakeCommandResponse, FakeSlurmArray, + FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask, ) @@ -25,6 +26,7 @@ "FakeLogicalEndpoint", "FakeServingState", "FakeSlurmArray", + "FakeSlurmJob", "FakeSlurmRunner", "FakeSlurmTask", "FakeVllmBackend", diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 0cf7a8e85..d934fd637 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -3,10 +3,12 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=3 +#SBATCH --cpus-per-task=32 #SBATCH --time=03:55:00 #SBATCH --array=0-1%2 #SBATCH --gres=gpu:8 set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" @@ -16,17 +18,24 @@ readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" verify_sha256() { - printf '%s %s\n' "$1" "$2" | sha256sum --check --status - + 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}" -readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +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}" install -d -m 0700 "${DD_ATTEMPT_DIR}" -readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +readonly DD_RUNTIME_DIR tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" source "${DD_RUNTIME_DIR}/entrypoint.sh" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 60751e07c..93156136d 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -3,10 +3,12 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=1 +#SBATCH --cpus-per-task=32 #SBATCH --time=03:55:00 #SBATCH --array=0 #SBATCH --gres=gpu:8 set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" @@ -16,17 +18,24 @@ readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" verify_sha256() { - printf '%s %s\n' "$1" "$2" | sha256sum --check --status - + 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}" -readonly DD_ARRAY_TASK_ID="${SLURM_ARRAY_TASK_ID:?SLURM_ARRAY_TASK_ID is required}" +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}" install -d -m 0700 "${DD_ATTEMPT_DIR}" -readonly DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" +readonly DD_RUNTIME_DIR tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" source "${DD_RUNTIME_DIR}/entrypoint.sh" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 785af0630..22da3c51f 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -14,8 +14,14 @@ from data_designer.slurm.state import SchedulerIdentity _JOB_SELECTOR_PATTERN = re.compile(r"^[0-9]+(?:_[0-9]+)?$") -_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--format=%i|%T") -_SACCT_REQUIRED_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") +_SQUEUE_REQUIRED_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +_SACCT_REQUIRED_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", +) @dataclass(frozen=True) @@ -27,6 +33,20 @@ class FakeCommandResponse: 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.""" @@ -60,16 +80,25 @@ def array_job_id(self) -> int: class FakeSlurmRunner: - """Stateful fake that copies its configured arrays before exposing Slurm commands.""" + """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: - self._pending_arrays = deque(copy.deepcopy(tuple(arrays))) + 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, ...]] = [] @@ -120,6 +149,20 @@ def set_task_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()) @@ -141,38 +184,58 @@ def _dispatch(self, command_name: str, argv: tuple[str, ...]) -> FakeCommandResp return handler(argv) def _run_sbatch(self, argv: tuple[str, ...]) -> FakeCommandResponse: - if not self._pending_arrays: + if not self._pending_submissions: return FakeCommandResponse(stderr="no scripted submission\n", returncode=1) - array = self._pending_arrays.popleft() - self._submitted_arrays[array.array_job_id] = array + 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"{array.array_job_id}\n") - return FakeCommandResponse(stdout=f"Submitted batch job {array.array_job_id}\n") + 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 rows)) + 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 rows)) + 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(targets[0]) + tasks = self._tasks_for_target(target) except (KeyError, ValueError): return FakeCommandResponse(stderr="unknown cancellation target\n", returncode=1) for task in tasks: @@ -215,22 +278,12 @@ def _sorted_submitted_tasks(self) -> list[FakeSlurmTask]: ) def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask]: - selectors: list[str] = [] - for index, argument in enumerate(argv[1:]): - if argument in {"-j", "--jobs"}: - try: - selectors.extend(argv[index + 2].split(",")) - except IndexError: - raise AssertionError(f"missing job selector in {argv!r}") from None - elif argument.startswith("--jobs="): - selectors.extend(argument.partition("=")[2].split(",")) + selectors = self._job_selectors(argv) if not selectors: return self._sorted_submitted_tasks() selected: dict[SchedulerIdentity, FakeSlurmTask] = {} for selector in selectors: - if _JOB_SELECTOR_PATTERN.fullmatch(selector) is None: - raise AssertionError(f"malformed job selector {selector!r} in {argv!r}") try: tasks = self._tasks_for_target(selector) except KeyError: @@ -240,3 +293,34 @@ def _selected_submitted_tasks(self, argv: tuple[str, ...]) -> list[FakeSlurmTask 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_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 1d793f500..a8c29e5ef 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="cdb5204e8ce7a7affda99ea9c121c6b60a16bf42c510c8eb6bc57661969620b9", + expected_fixture_sha256="8ddf07c38a825a7487c487fddfe051f0a1940f63063725b54b32bfe4c03fd9ca", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="c6708cdbc0a03e095062c0642cfd141f066153b958b7ad3dec779afd6414fa34", + expected_fixture_sha256="17a4c2e16189d22dfdb6885bf76264844ad3168dea0cf94aef70948d5ab2e6b7", ) @@ -54,7 +54,9 @@ def _assert_script_matches_plan( *(index for deployment in plan.deployments for index in deployment.node_indices), ) node_count = max(node_indices) + 1 - array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + array = "0" if plan.array_tasks.count == 1 else f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.count > 1: + 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) @@ -62,6 +64,7 @@ def _assert_script_matches_plan( assert f"#SBATCH --account={plan.submission.account}\n" in script assert f"#SBATCH --partition={plan.submission.partition}\n" in script assert f"#SBATCH --nodes={node_count}\n" in script + assert f"#SBATCH --cpus-per-task={plan.client.authored.cpus}\n" in script assert f"#SBATCH --time={plan.submission.time_limit}\n" in script assert f"#SBATCH --array={array}\n" in script assert f"#SBATCH --gres=gpu:{plan.resolved_gpus_per_node}\n" in script diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index 8bfb3de0e..c47b21432 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -9,11 +9,17 @@ import pytest from data_designer.slurm.state import SchedulerIdentity -from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmJob, FakeSlurmRunner, FakeSlurmTask GOLDEN_DIRECTORY = Path(__file__).parent / "golden" / "slurm" -SQUEUE_ARGUMENTS = ("--noheader", "--format=%i|%T") -SACCT_ARGUMENTS = ("--noheader", "--parsable2", "--format=%i|%State|%ExitCode") +SQUEUE_ARGUMENTS = ("--noheader", "--array", "--format=%i|%T") +SACCT_ARGUMENTS = ( + "--noheader", + "--array", + "--allocations", + "--parsable2", + "--format=JobID,State,ExitCode", +) def _submit(runner: FakeSlurmRunner) -> None: @@ -135,6 +141,25 @@ def test_fake_slurm_runner_models_cancellation( 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: @@ -174,7 +199,7 @@ def test_fake_slurm_runner_matches_sbatch_parsable_mode( "command", ( ("squeue", "--noheader"), - ("sacct", "--noheader", "--format=%i|%State|%ExitCode"), + ("sacct", "--noheader", "--format=JobID,State,ExitCode"), ), ) def test_fake_slurm_runner_rejects_underspecified_state_queries( From 8377f1813c1733644c5a8f04587769c3f333ce1c Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 28 Aug 2026 09:52:18 -0600 Subject: [PATCH 10/62] feat: resolve Slurm serving deployments (#890) * feat: resolve Slurm serving deployments Add pure, package-owned vLLM compatibility and deterministic process, readiness, backend, and logical endpoint resolution. Validate topology, launch inputs, and network identities with single-node and multi-node golden coverage. Closes #866 Signed-off-by: Nabin Mulepati * fix: harden Slurm serving contracts Revalidate persisted launch policy inputs, require canonical process ordering, and normalize resolver errors. Reject vLLM CLI aliases and transports that could override compiler-owned topology or endpoints. Refs #866 Signed-off-by: Nabin Mulepati * fix: close Slurm serving validation gaps Preserve independent single-node expert-parallel replicas while rejecting cross-node EP. Block attached compiler-owned vLLM topology aliases, require canonical runtime versions, and cover the boundaries with regression tests. Refs #866 Signed-off-by: Nabin Mulepati * fix: harden vLLM launch ownership Reject unmodeled topology, plugin, LoRA, speculative-model, and environment overrides at authored and persisted boundaries. Validate JSON arguments, secret references, readiness paths, and compatibility version lengths. Refs #866 Signed-off-by: Nabin Mulepati * fix: bind serving context to plan Derive serving placement and client endpoint records from the same resolved plan so callers cannot mix contracts across plans. Reject vLLM API authentication until the logical endpoint models credential forwarding. Refs #866 Signed-off-by: Nabin Mulepati * fix: align serving with migration contracts - enforce plan-owned environment and endpoint behavior - limit compatibility to tested vLLM release series - split shared and vLLM configuration from run records - cover multi-group ordering and launch delays Refs #866 Signed-off-by: Nabin Mulepati * docs: clarify option flag extraction Signed-off-by: Nabin Mulepati * refactor: share argument and duration helpers Use action-named shared helpers for option extraction and duration conversion while keeping vLLM flag canonicalization backend-specific. Refs #866 Signed-off-by: Nabin Mulepati * refactor: simplify vLLM compatibility contract Track the tested runtime range and contract revision without repeating baseline vLLM or package-owned behaviors as capability booleans. Refs #866 Signed-off-by: Nabin Mulepati * refactor: make vLLM resolution version agnostic Preserve inspected runtime versions as digest-bound provenance without hard-coded support gates or redundant compatibility records. Continue to validate persisted text and compiler-owned vLLM arguments. Refs #866 Signed-off-by: Nabin Mulepati * refactor: treat runtime version as provenance Signed-off-by: Nabin Mulepati * refactor: name vLLM serving module explicitly Signed-off-by: Nabin Mulepati * refactor: remove serving context wrapper Signed-off-by: Nabin Mulepati * refactor: centralize Slurm scalar types Signed-off-by: Nabin Mulepati * refactor: use str for model aliases Signed-off-by: Nabin Mulepati * refactor: clarify serving replica indices * fix: reject abbreviated vLLM flags * refactor: simplify Slurm serving contracts - make vLLM resolution consume placement inputs before plan assembly - keep serving internals backend-specific and remove speculative exports - clarify validation ownership, secret redaction, and digest semantics * fix: bind Slurm serving resolution inputs - Join placement and endpoint claims through one planner-owned context.\n- Normalize inconsistent resolved records to the serving error boundary.\n- Split resolver, record, and launch-policy tests by ownership. Signed-off-by: Nabin Mulepati * remove empty serving package initializer * fix serving contract validation boundaries Derive deployment and endpoint inputs from one resolved plan and centralize topology arithmetic. Reject plaintext secrets in persisted builder and vLLM argument payloads, normalize resolver errors, and require canonical argv tokens. * fix option argument validation gaps --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/config/__init__.py | 6 +- .../data_designer/slurm/config/benchmark.py | 5 +- .../data_designer/slurm/config/environment.py | 114 ++++++ .../data_designer/slurm/config/profiles.py | 8 +- .../src/data_designer/slurm/config/run.py | 201 +++------- .../src/data_designer/slurm/config/utils.py | 12 + .../src/data_designer/slurm/config/vllm.py | 98 +++++ .../slurm/config/vllm_validation.py | 230 +++++++++++ .../src/data_designer/slurm/contracts.py | 43 ++- .../data_designer/slurm/planning/models.py | 70 +++- .../data_designer/slurm/serving/deployment.py | 192 +++++++++ .../data_designer/slurm/serving/endpoints.py | 60 +++ .../data_designer/slurm/serving/resolver.py | 210 ++++++++++ .../src/data_designer/slurm/serving/vllm.py | 123 ++++++ .../src/data_designer/slurm/state/__init__.py | 2 - .../data_designer/slurm/state/readiness.py | 4 +- .../src/data_designer/slurm/types.py | 39 ++ .../tests/contracts/golden/authored_run.json | 4 + .../contracts/golden/authored_run_single.json | 2 + .../contracts/golden/multi_node_plan.json | 6 +- .../contracts/golden/single_node_plan.json | 4 +- .../tests/contracts/test_config_records.py | 364 +++++++++++++++++- .../tests/contracts/test_planning_records.py | 53 ++- .../golden/finalization_chain.json | 6 +- .../tests/serving/conftest.py | 72 ++++ .../tests/serving/golden/multi_node.json | 231 +++++++++++ .../tests/serving/golden/single_node.json | 120 ++++++ .../tests/serving/test_resolver.py | 324 ++++++++++++++++ .../tests/serving/test_serving_records.py | 208 ++++++++++ .../tests/serving/test_vllm_launch_policy.py | 137 +++++++ .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/test_contracts.py | 11 +- 34 files changed, 2740 insertions(+), 227 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/environment.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/utils.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/vllm_validation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/serving/endpoints.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/types.py create mode 100644 packages/data-designer-slurm/tests/serving/conftest.py create mode 100644 packages/data-designer-slurm/tests/serving/golden/multi_node.json create mode 100644 packages/data-designer-slurm/tests/serving/golden/single_node.json create mode 100644 packages/data-designer-slurm/tests/serving/test_resolver.py create mode 100644 packages/data-designer-slurm/tests/serving/test_serving_records.py create mode 100644 packages/data-designer-slurm/tests/serving/test_vllm_launch_policy.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 5cda08684..8107c1f8a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -14,6 +14,7 @@ DataDesignerSlurmBenchmarkConfig, FixedRecordPolicy, ) +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef from data_designer.slurm.config.images import ( ClientImageInspection, ImageBuildRequest, @@ -47,16 +48,13 @@ InputBindings, InvocationConfig, InvocationDiagnostics, - LiteralEnvironmentBinding, LocalStdioMCPProviderConfig, OutputConfig, - QueueBackpressureConfig, RemoteMCPProviderConfig, - SecretRef, ServerDeploymentConfig, SubmissionConfig, - VllmServerConfig, ) +from data_designer.slurm.config.vllm import QueueBackpressureConfig, VllmServerConfig __all__ = [ "AdaptiveRecordPolicy", 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 index e7c465c2e..8fb92a574 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -12,7 +12,6 @@ AuthoredConfig, Duration, Identifier, - ModelAlias, SchemaVersion, validate_local_config_path, ) @@ -54,7 +53,7 @@ def validate_topology(self) -> BenchmarkDeploymentOverride: class BenchmarkDeploymentCase(AuthoredConfig): name: Identifier - deployments: dict[ModelAlias, BenchmarkDeploymentOverride] = Field(min_length=1) + deployments: dict[str, BenchmarkDeploymentOverride] = Field(min_length=1) class FixedRecordPolicy(AuthoredConfig): @@ -89,7 +88,7 @@ class DataDesignerSlurmBenchmarkConfig(AuthoredConfig): schema_version: SchemaVersion name: Identifier base_run: BenchmarkBaseRun - model_aliases: Literal["all"] | list[ModelAlias] + 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 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..c835c3983 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py @@ -0,0 +1,114 @@ +# 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 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", +] + + +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 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("_") 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 index 87412d84d..a6a0b98c3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -15,7 +15,7 @@ Identifier, SchemaVersion, Sha256Digest, - compute_sha256, + compute_canonical_json_sha256, validate_absolute_path, validate_plain_text, ) @@ -155,7 +155,7 @@ def select_profile( catalog_path: str | None = None, ) -> SelectedSlurmProfile: """Select a profile with explicit, hostname, then default precedence.""" - catalog_sha256 = compute_sha256(catalog.model_dump(mode="json")) + 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}") @@ -216,7 +216,7 @@ def validate_selected_profile( """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_sha256(catalog.model_dump(mode="json")): + 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") @@ -255,7 +255,7 @@ def _selection( def _profile_digest(profile: SlurmProfile) -> Sha256Digest: - return compute_sha256(profile.model_dump(mode="json")) + return compute_canonical_json_sha256(profile.model_dump(mode="json")) def _validate_hostname_glob(value: str) -> None: 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 index c6f7d69d4..cd323b763 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -1,11 +1,12 @@ # 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 collections.abc import Mapping from typing import Annotated, Literal from urllib.parse import urlsplit @@ -14,7 +15,6 @@ from pydantic import ( Field, JsonValue, - NonNegativeInt, PositiveInt, StringConstraints, field_validator, @@ -22,110 +22,47 @@ ) 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.images import ImageRef +from data_designer.slurm.config.vllm import QueueBackpressureConfig, VllmServerConfig from data_designer.slurm.contracts import ( AuthoredConfig, - Duration, - EnvironmentName, - Identifier, - ModelAlias, - SchemaVersion, validate_absolute_path, validate_local_config_path, validate_plain_text, validate_url, ) - -_OWNED_VLLM_FLAGS = { - "--api-key", - "--distributed-executor-backend", - "--distributed-init-address", - "--enable-expert-parallel", - "--headless", - "--host", - "--middleware", - "--model", - "--pipeline-parallel-size", - "--port", - "--served-model-name", - "--tensor-parallel-size", -} -_DURATION_FACTORS = {"s": 1, "m": 60, "h": 3600, "d": 86400} -_NON_SECRET_PAYLOAD_KEYS = frozenset({"idempotency_key", "partition_key", "primary_key", "sort_key"}) -_SECRET_NAME_PARTS = frozenset({"credential", "credentials", "password", "secret", "token"}) - - -def _duration_seconds(value: Duration) -> int: - return int(value[:-1]) * _DURATION_FACTORS[value[-1]] - - -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 _is_secret_name(value: str) -> bool: - segments = _secret_name_segments(value) - return bool( - _SECRET_NAME_PARTS.intersection(segments) - or {"access", "key"}.issubset(segments) - or segments[-1] in {"auth", "key"} - ) - - -def _is_secret_payload_name(value: str) -> bool: - segments = _secret_name_segments(value) - normalized = "_".join(segments) - return normalized not in _NON_SECRET_PAYLOAD_KEYS and _is_secret_name(value) - - -def _option_flag(value: str) -> str: - return re.split(r"[=\s]", value.lstrip(), maxsplit=1)[0] - - -def _contains_secret_key(value: object) -> bool: - if isinstance(value, Mapping): - return any( - (_is_secret_payload_name(str(key)) and item is not None) or _contains_secret_key(item) - for key, item in value.items() - ) - if isinstance(value, list | tuple): - return any(_contains_secret_key(item) for item in value) - return False - - -def _validate_environment_bindings( - values: dict[EnvironmentName, EnvironmentBinding], -) -> dict[EnvironmentName, EnvironmentBinding]: - literal_secrets = [ - name - for name, binding in values.items() - if _is_secret_name(name) and isinstance(binding, LiteralEnvironmentBinding) - ] - if literal_secrets: - raise ValueError("secret-shaped environment names require external secret references") - return values - - -class LiteralEnvironmentBinding(AuthoredConfig): - 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): - type: Literal["secret"] - environment: EnvironmentName - - -EnvironmentBinding = Annotated[ - LiteralEnvironmentBinding | SecretRef, - Field(discriminator="type"), +from data_designer.slurm.types import EnvironmentName, Identifier, SchemaVersion + +__all__ = [ + "ArrayTasksConfig", + "BuilderInput", + "ClientConfig", + "ClientDependencies", + "DataDesignerSlurmConfig", + "DeploymentResources", + "DeploymentTopology", + "EnvironmentBinding", + "InputBindings", + "InvocationConfig", + "InvocationDiagnostics", + "LiteralEnvironmentBinding", + "LocalStdioMCPProviderConfig", + "MCPProviderConfig", + "OutputConfig", + "QueueBackpressureConfig", + "RemoteMCPProviderConfig", + "SecretRef", + "ServerDeploymentConfig", + "SubmissionConfig", + "VllmServerConfig", ] @@ -145,6 +82,7 @@ def validate_input(self) -> BuilderInput: 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 Big Iron fields: {', '.join(sorted(retired))}") @@ -157,8 +95,6 @@ def validate_input(self) -> BuilderInput: valid = isinstance(self.inline.get("columns"), list) if not valid: raise ValueError("inline builder input must be one complete serialized Data Designer config") - if _contains_secret_key(self.inline): - raise ValueError("inline builder input must not contain secret values") return self @@ -208,12 +144,18 @@ def validate_command(cls, value: str) -> str: def validate_args(cls, values: list[str]) -> list[str]: for value in values: validate_plain_text(value, field_name="MCP argument") - option = _option_flag(value).lstrip("-") - if _is_secret_name(option): + 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 - _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) + @field_validator("environment") + @classmethod + def validate_environment( + cls, values: dict[EnvironmentName, EnvironmentBinding] + ) -> dict[EnvironmentName, EnvironmentBinding]: + validate_environment_bindings(values) + return values MCPProviderConfig = Annotated[ @@ -233,7 +175,7 @@ class InvocationConfig(AuthoredConfig): 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[ModelAlias, PositiveInt] = Field(default_factory=dict) + model_concurrency: dict[str, PositiveInt] = Field(default_factory=dict) diagnostics: InvocationDiagnostics = Field(default_factory=InvocationDiagnostics) @field_validator("run_config") @@ -267,11 +209,11 @@ def validate_requirements(cls, values: list[str] | None) -> list[str] | None: 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(f"dependency requirement must identify a package or immutable wheel: {value!r}") + raise ValueError("dependency requirement must identify a package or immutable wheel") try: requirement = Requirement(value) except InvalidRequirement as error: - raise ValueError(f"invalid dependency requirement: {value!r}") from 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: @@ -321,51 +263,6 @@ class ClientConfig(AuthoredConfig): dependencies: ClientDependencies = Field(default_factory=ClientDependencies) -class QueueBackpressureConfig(AuthoredConfig): - max_waiting_requests: NonNegativeInt = 128 - retry_after_seconds: NonNegativeInt | None = 1 - - -class VllmServerConfig(AuthoredConfig): - type: Literal["vllm"] - image: ImageRef - startup_timeout: Duration = "15m" - distributed_init_timeout: Duration = "10m" - readiness_path: str = "/health" - enable_expert_parallel: bool = False - 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: - validate_plain_text(value, field_name="readiness path") - if not value.startswith("/") or "?" in value or "#" in value: - raise ValueError("readiness_path must be an absolute URL path without query or fragment") - return value - - @field_validator("extra_args") - @classmethod - def validate_extra_args(cls, values: list[str]) -> list[str]: - for value in values: - validate_plain_text(value, field_name="vLLM argument") - flag = _option_flag(value) - if flag in _OWNED_VLLM_FLAGS: - raise ValueError(f"vLLM argument {flag!r} is owned by the compiler or runtime") - if _is_secret_name(flag.lstrip("-")): - raise ValueError("secret-shaped vLLM arguments must use an environment secret reference") - return values - - _environment_uses_secret_references = field_validator("environment")(_validate_environment_bindings) - - @model_validator(mode="after") - def validate_timeouts(self) -> VllmServerConfig: - if _duration_seconds(self.distributed_init_timeout) > _duration_seconds(self.startup_timeout): - raise ValueError("distributed_init_timeout must not exceed startup_timeout") - return self - - class DeploymentResources(AuthoredConfig): nodes: PositiveInt = 1 @@ -376,7 +273,7 @@ class DeploymentTopology(AuthoredConfig): class ServerDeploymentConfig(AuthoredConfig): - model_alias: ModelAlias + model_alias: str served_model_name: str | None = None model: str server: VllmServerConfig 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..d440881d5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py @@ -0,0 +1,98 @@ +# 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, 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. + 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 + 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 index fc47926b1..b8c0b200e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -9,7 +9,7 @@ import json import posixpath from collections.abc import Mapping -from typing import Annotated, Literal, TypeVar +from typing import TypeVar from urllib.parse import urlsplit from pydantic import ( @@ -17,26 +17,24 @@ ConfigDict, NonNegativeInt, PositiveInt, - StringConstraints, field_validator, model_validator, ) -Identifier = Annotated[ - str, - StringConstraints( - min_length=1, - max_length=128, - pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", - ), -] +from data_designer.slurm.types import ( + AttemptId, + Duration, + EnvironmentName, + Identifier, + NetworkPort, + NonNegativeDuration, + SchemaVersion, + Sha256Digest, + ShardId, +) + +# TODO: Remove these compatibility exports after Stage 2 branches import shared scalars from slurm.types. ModelAlias = str -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)$")] _Key = TypeVar("_Key") _Value = TypeVar("_Value") @@ -97,6 +95,7 @@ class ContractValue(BaseModel): extra="forbid", frozen=True, allow_inf_nan=False, + hide_input_in_errors=True, protected_namespaces=(), strict=True, validate_default=True, @@ -118,6 +117,7 @@ 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() @@ -164,11 +164,15 @@ def pretty_json(value: object) -> str: ) -def compute_sha256(value: object) -> Sha256Digest: +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() +# 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("/"): @@ -280,14 +284,15 @@ class ResumeWorkspace(ContractValue): "Duration", "EnvironmentName", "Identifier", - "ModelAlias", + "NetworkPort", + "NonNegativeDuration", "RecordRange", "ResumeWorkspace", "SchemaVersion", "Sha256Digest", "ShardId", "canonical_json", - "compute_sha256", + "compute_canonical_json_sha256", "pretty_json", "validate_absolute_path", "validate_local_config_path", 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 index d1eff7ab6..9b71cc31f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -13,12 +13,15 @@ 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 ( + ClientImageInspection, DistributionName, ImageInspectionRecord, ImageKind, ImageRef, InstalledDistribution, + ServingImageInspection, ) from data_designer.slurm.config.profiles import ContainerMount, SelectedSlurmProfile from data_designer.slurm.config.run import ( @@ -34,16 +37,16 @@ ContractRecord, ContractValue, Identifier, - ModelAlias, RecordRange, ResumeWorkspace, Sha256Digest, ShardId, - compute_sha256, + compute_canonical_json_sha256, validate_absolute_path, validate_local_config_path, validate_plain_text, ) +from data_designer.slurm.types import NetworkPort class ResolvedImage(ContractValue): @@ -74,6 +77,11 @@ def validate_image(self) -> ResolvedImage: def kind(self) -> ImageKind: return self.inspection.inspection.kind + @property + def inspection_facts(self) -> ClientImageInspection | ServingImageInspection: + """Return factual image inspection data without exposing record nesting to consumers.""" + return self.inspection.inspection + class LockedPackage(ContractValue): name: DistributionName @@ -151,8 +159,8 @@ class ResolvedBuilderInput(ContractValue): source: ArtifactReference | None = None inline: dict[str, JsonValue] | None = None content_sha256: Sha256Digest - model_aliases: tuple[ModelAlias, ...] - referenced_model_aliases: tuple[ModelAlias, ...] = () + model_aliases: tuple[str, ...] + referenced_model_aliases: tuple[str, ...] = () @model_validator(mode="after") def validate_input(self) -> ResolvedBuilderInput: @@ -161,7 +169,8 @@ def validate_input(self) -> ResolvedBuilderInput: if self.source is None: if self.authored_source is not None: raise ValueError("inline builder input cannot contain authored_source") - expected_digest = compute_sha256(self.inline) + validate_no_plaintext_secrets(self.inline, field_name="resolved inline builder input") + expected_digest = compute_canonical_json_sha256(self.inline) model_aliases, referenced_aliases = _extract_builder_aliases(self.inline) if self.model_aliases != model_aliases: raise ValueError("resolved model aliases do not match the inline builder") @@ -196,11 +205,14 @@ def validate_run_config_is_materialized(cls, value: dict[str, JsonValue]) -> dic 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: Annotated[int, Field(ge=1024, le=65535)] + port: NetworkPort class ResolvedTopology(ContractValue): + """Resource-derived serving topology shared by planning and serving.""" + tensor_parallel: PositiveInt nodes_per_replica: PositiveInt pipeline_parallel: PositiveInt @@ -209,6 +221,34 @@ class ResolvedTopology(ContractValue): 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 @@ -233,17 +273,11 @@ def validate_deployment(self) -> ResolvedDeployment: 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") - if self.gpus_per_node % self.authored.topology.tensor_parallel: - raise ValueError("tensor_parallel must divide resolved GPUs per node") - expected = ResolvedTopology( + 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, - pipeline_parallel=self.authored.topology.nodes_per_replica, - node_group_count=self.authored.resources.nodes // self.authored.topology.nodes_per_replica, - replicas_per_node_group=self.gpus_per_node // self.authored.topology.tensor_parallel, - replica_count=(self.authored.resources.nodes // self.authored.topology.nodes_per_replica) - * (self.gpus_per_node // self.authored.topology.tensor_parallel), - gpus_per_replica=self.authored.topology.tensor_parallel * self.authored.topology.nodes_per_replica, ) if self.topology != expected: raise ValueError("resolved topology does not match deployment resources") @@ -482,7 +516,7 @@ def _is_below(path: str, root: str) -> bool: return path != root and posixpath.commonpath((path, root)) == root -def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[ModelAlias, ...], tuple[ModelAlias, ...]]: +def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[str, ...], tuple[str, ...]]: data_designer = builder.get("data_designer", builder) if not isinstance(data_designer, dict): raise ValueError("builder data_designer value must be an object") @@ -490,13 +524,13 @@ def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[Model if not isinstance(model_configs, list): raise ValueError("builder model_configs must be a list") - model_aliases: list[ModelAlias] = [] + model_aliases: list[str] = [] 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"]) - referenced_aliases: list[ModelAlias] = [] + referenced_aliases: list[str] = [] def collect(value: JsonValue, *, key: str | None = None) -> None: if key == "model_configs": 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..51dc57f5a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py @@ -0,0 +1,192 @@ +# 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 + 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 = 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..ac2bdce27 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py @@ -0,0 +1,210 @@ +# 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 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 + + +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 and endpoint derived from the same run plan. + + Inspected runtime versions remain provenance. This resolver emits only the + baseline behavior required of every admitted vLLM image and therefore does not + maintain a package-version compatibility 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") + 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 = 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, + 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, + 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) 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..73f3e085a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py @@ -0,0 +1,123 @@ +# 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 + 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 index 2d7eb81d8..07885b69c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -11,7 +11,6 @@ ContractRecord, ContractValue, Identifier, - ModelAlias, RecordRange, ResumeWorkspace, Sha256Digest, @@ -85,7 +84,6 @@ "EffectiveAttemptState", "EndpointPublicationState", "Identifier", - "ModelAlias", "ProbeEvidence", "ProbeOutcome", "ReadinessState", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py index a7f538165..fb88e1965 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -9,7 +9,7 @@ from pydantic import Field, NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator -from data_designer.slurm.contracts import AttemptId, Identifier, ModelAlias, ShardId +from data_designer.slurm.contracts import AttemptId, Identifier, ShardId from data_designer.slurm.state.base import StateRecord, StateValue, validate_utc_timestamp ReasonCode = Annotated[ @@ -63,7 +63,7 @@ class DeploymentReadiness(StateValue): """Readiness state for one authored-order model deployment.""" deployment_id: Identifier - model_alias: ModelAlias + model_alias: str state: ReadinessState expected_backends: PositiveInt ready_backends: NonNegativeInt 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..aba456062 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/types.py @@ -0,0 +1,39 @@ +# 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 Field, StringConstraints + +Identifier = Annotated[ + str, + StringConstraints( + min_length=1, + max_length=128, + pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$", + ), +] +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", + "SchemaVersion", + "Sha256Digest", + "ShardId", +] diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index 28f65680a..c1bb0d43f 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -67,10 +67,12 @@ "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" @@ -96,10 +98,12 @@ "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" 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 index c6a46897a..346b4500c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -47,10 +47,12 @@ "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" 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 index baf366891..48427e849 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-001/authored-config.json", - "sha256": "9bfdbcb1ddd374a5d43eea44890d13c08798ea9e18c30ff5d5258fa0d27f539d" + "sha256": "74f30c766e96181f72642dff1da72ae10786d5e404b4f65b927bdf27fa0c6a61" }, "builder": { "authored_source": null, @@ -151,10 +151,12 @@ "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" @@ -231,10 +233,12 @@ "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" 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 index cfc0b5824..4b86c8657 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" }, "builder": { "authored_source": null, @@ -124,10 +124,12 @@ "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" diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index 33959d5c1..dcf9c876d 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -94,12 +94,36 @@ def test_client_dependencies_accept_digest_bound_wheel() -> None: 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): - SecretRef.model_validate({"type": "secret", "environment": "HF_TOKEN", "value": "secret-value"}) + 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: @@ -164,31 +188,147 @@ def test_stdio_mcp_rejects_secret_shaped_arguments(argument: str) -> None: ) +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( "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", - " --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: @@ -196,6 +336,148 @@ def test_vllm_rejects_runtime_owned_arguments(argument: str) -> None: 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", [ @@ -209,6 +491,45 @@ def test_vllm_rejects_secret_shaped_arguments(extra_args: list[str]) -> None: 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( @@ -219,6 +540,12 @@ def test_vllm_rejects_distributed_timeout_beyond_startup_timeout() -> None: ) +@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", @@ -232,7 +559,14 @@ def test_deployment_rejects_invalid_topology() -> None: 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) @@ -287,9 +621,20 @@ def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfi "subscription_key", ], ) -def test_builder_input_rejects_secret_values(secret_key: str) -> None: - with pytest.raises(ValidationError, match="secret values"): - BuilderInput(inline={"columns": [], secret_key: "plaintext-secret"}) +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"]) @@ -348,13 +693,14 @@ def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesig DataDesignerSlurmConfig.model_validate(payload) -def test_small_config_values_validate_at_boundary() -> None: +@pytest.mark.parametrize("readiness_path", ["health", "//other-host/health", "/health check"]) +def test_small_config_values_validate_at_boundary(readiness_path: str) -> 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="readiness_path"): - VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), readiness_path="health") + with pytest.raises(ValidationError, match="absolute URL path"): + VllmServerConfig(type="vllm", image=ImageRef(name="vllm"), readiness_path=readiness_path) @pytest.mark.parametrize( diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index 5c64a5c24..cb907c010 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,18 +10,52 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_sha256 +from data_designer.slurm.contracts import compute_canonical_json_sha256 from data_designer.slurm.planning import ( ArtifactReference, PlanContractError, + ResolvedBuilderInput, ResolvedDependencyLock, ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedSubmission, + ResolvedTopology, 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, @@ -66,6 +100,19 @@ def test_resolved_plan_is_deeply_immutable(multi_node_plan: ResolvedSlurmRunPlan 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_canonical_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", [ @@ -160,7 +207,7 @@ def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: Res "model_configs" ][:1] payload["builder"]["model_aliases"] = ["generator"] - payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + payload["builder"]["content_sha256"] = compute_canonical_json_sha256(payload["builder"]["inline"]) with pytest.raises(ValidationError, match="deployment alias"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -170,7 +217,7 @@ def test_plan_rejects_referenced_alias_without_deployment(multi_node_plan: Resol payload = multi_node_plan.model_dump(mode="json") payload["builder"]["inline"]["data_designer"]["columns"] = [{"model_alias": "missing"}] payload["builder"]["referenced_model_aliases"] = ["missing"] - payload["builder"]["content_sha256"] = compute_sha256(payload["builder"]["inline"]) + payload["builder"]["content_sha256"] = compute_canonical_json_sha256(payload["builder"]["inline"]) with pytest.raises(ValidationError, match="referenced Data Designer model alias"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index e3b564f14..7e993b116 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" + "sha256": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" }, "run_id": "run-single", "scheduler": { @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "fa6ca55eac5075455193628e481b09789566d8f4c54926f45bbf837c20e5ba47" + "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" + "sha256": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" }, "run_id": "run-single", "schema_version": 1, 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..f19973cbd --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/test_resolver.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 ast +import json +from pathlib import Path + +import pytest + +import data_designer.slurm.serving.resolver as resolver_module +from data_designer.slurm.config import QueueBackpressureConfig +from data_designer.slurm.contracts import 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_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"] = "vendor-vllm-build" + 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 == "vendor-vllm-build" + + +def test_resolution_rejects_image_inspection_mismatch(single_node_plan: ResolvedSlurmRunPlan) -> None: + placement = single_node_plan.deployments[0].model_copy(update={"image": single_node_plan.client.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] + 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..3e8a59cc3 --- /dev/null +++ b/packages/data-designer-slurm/tests/serving/test_serving_records.py @@ -0,0 +1,208 @@ +# 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.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 ResolvedVllmProcess + + +@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, + mutation: str, + message: str, +) -> None: + payload = resolved_multi_replica_server.model_dump(mode="json") + if mutation == "image_kind": + payload["image"] = multi_node_plan.client.image.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/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index d934fd637..b3c1e40c8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-001/resolved-plan.json" -readonly DD_PLAN_SHA256="fff7573a8b6eb2f71ae12f27900a85e1e121debfe58b0e94fefc3fb9362970ce" +readonly DD_PLAN_SHA256="00e206f4b759ca06336c02d3694f7289551d1a0b9ba50c05f37cdbf7a213756d" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 93156136d..115d998da 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" readonly DD_RUNTIME_SHA256="eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" readonly DD_PLAN="/workspace/primary/runs/run-single/resolved-plan.json" -readonly DD_PLAN_SHA256="8cf7c2ec3eb398bcddd75e200094fffa8ce4403a655e3340b35e5098d2f308a0" +readonly DD_PLAN_SHA256="897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index a8c29e5ef..34d947543 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="8ddf07c38a825a7487c487fddfe051f0a1940f63063725b54b32bfe4c03fd9ca", + expected_fixture_sha256="244fe9b160b361645e3f9be13d40d5d6157563a69fa37037fa974d2a3acead2f", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="17a4c2e16189d22dfdb6885bf76264844ad3168dea0cf94aef70948d5ab2e6b7", + expected_fixture_sha256="37681e8805eb81ca9c4e1fc2da88b14031d7a150c9fd729104ed875ed2503224", ) diff --git a/packages/data-designer-slurm/tests/test_contracts.py b/packages/data-designer-slurm/tests/test_contracts.py index d1983b915..8911689b5 100644 --- a/packages/data-designer-slurm/tests/test_contracts.py +++ b/packages/data-designer-slurm/tests/test_contracts.py @@ -8,6 +8,8 @@ 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, ) @@ -16,7 +18,7 @@ ContractValue, ResumeWorkspace, canonical_json, - compute_sha256, + compute_canonical_json_sha256, pretty_json, ) from data_designer.slurm.contracts import ( @@ -43,6 +45,11 @@ ) +@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 @@ -58,7 +65,7 @@ def test_shared_json_helpers_are_deterministic() -> None: 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_sha256(value) == hashlib.sha256(canonical_json(value)).hexdigest() + assert compute_canonical_json_sha256(value) == hashlib.sha256(canonical_json(value)).hexdigest() def test_resume_workspace_requires_a_safe_absolute_path() -> None: From 561f139161ed05b6218bbb9ac9c00842d94e7e5d Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:48:10 -0300 Subject: [PATCH 11/62] feat: compile deterministic Slurm run plans (#893) * feat: compile deterministic Slurm run plans Signed-off-by: Andre Manoel * fix: enforce Slurm planning invariants Reject unsupported multi-shard semantics and managed output collisions at the effective configuration boundary. Normalize config errors and cover the corrected Big Iron field disposition. Signed-off-by: Andre Manoel * fix: harden Slurm plan resolution Preserve partial early-shutdown intent and hide secret values in validation errors. Reject unsafe sharding semantics, require exact deployment coverage, and bind sourced payload identity. Signed-off-by: Andre Manoel * fix: bind Slurm effective plan inputs Preserve all authored early-shutdown controls and reject direct effective invocation drift. Bind shard partition digests to deterministic persisted JSON bytes. Signed-off-by: Andre Manoel * fix: align Slurm planning boundaries * fix: unify Slurm builder artifact identity * fix: validate Slurm compiler inputs * fix: complete Slurm plan resolution defaults * fix: address plan compiler review feedback * fix: preserve safe validation paths Signed-off-by: Andre Manoel * fix: preserve union validation paths Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- .../src/data_designer/slurm/_errors.py | 163 ++++ .../data_designer/slurm/config/__init__.py | 17 + .../src/data_designer/slurm/config/builder.py | 200 ++++ .../src/data_designer/slurm/config/errors.py | 18 + .../src/data_designer/slurm/config/loading.py | 227 +++++ .../data_designer/slurm/config/profiles.py | 2 + .../src/data_designer/slurm/config/run.py | 4 +- .../src/data_designer/slurm/contracts.py | 36 +- .../data_designer/slurm/launcher/renderer.py | 5 +- .../data_designer/slurm/planning/__init__.py | 3 - .../slurm/planning/builder_identity.py | 34 + .../data_designer/slurm/planning/compiler.py | 254 +++++ .../data_designer/slurm/planning/errors.py | 22 + .../data_designer/slurm/planning/models.py | 112 +-- .../slurm/planning/resolution.py | 428 +++++++++ .../slurm/planning/validation.py | 38 +- .../tests/config/test_loading_builder.py | 362 +++++++ .../contracts/golden/authored_run_single.json | 2 +- .../contracts/golden/multi_node_plan.json | 9 +- .../contracts/golden/single_node_plan.json | 13 +- .../tests/contracts/test_config_records.py | 1 + .../tests/contracts/test_planning_records.py | 99 +- .../tests/contracts/test_profiles.py | 9 + .../golden/finalization_chain.json | 6 +- .../tests/launcher/test_renderer.py | 13 +- .../tests/planning/test_compiler.py | 904 ++++++++++++++++++ .../tests/serving/test_resolver.py | 7 +- .../golden/rendered/multi_node.sbatch | 4 +- .../golden/rendered/single_node.sbatch | 4 +- .../slurm_test_fakes/test_rendered_scripts.py | 6 +- 30 files changed, 2858 insertions(+), 144 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/_errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/builder.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/config/loading.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py create mode 100644 packages/data-designer-slurm/tests/config/test_loading_builder.py create mode 100644 packages/data-designer-slurm/tests/planning/test_compiler.py 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..584bdf62b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -0,0 +1,163 @@ +# 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 = { + "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/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 8107c1f8a..6d329c23d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -14,7 +14,9 @@ DataDesignerSlurmBenchmarkConfig, FixedRecordPolicy, ) +from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.errors import SlurmConfigBuilderError, SlurmConfigLoadError from data_designer.slurm.config.images import ( ClientImageInspection, ImageBuildRequest, @@ -24,6 +26,13 @@ InstalledDistribution, ServingImageInspection, ) +from data_designer.slurm.config.loading import ( + DEFAULT_PROFILE_FILE_NAME, + PROFILE_FILE_ENVIRONMENT, + load_profile_catalog, + load_run_config, + resolve_profile, +) from data_designer.slurm.config.profiles import ( ContainerMount, GpuRequestMode, @@ -70,6 +79,8 @@ "ContainerMount", "DataDesignerSlurmBenchmarkConfig", "DataDesignerSlurmConfig", + "DataDesignerSlurmConfigBuilder", + "DEFAULT_PROFILE_FILE_NAME", "DeploymentResources", "DeploymentTopology", "FixedRecordPolicy", @@ -87,6 +98,7 @@ "LocalStdioMCPProviderConfig", "OutputConfig", "ProfileSelectionSource", + "PROFILE_FILE_ENVIRONMENT", "QueueBackpressureConfig", "RemoteMCPProviderConfig", "SchedulerProfile", @@ -94,11 +106,16 @@ "SelectedSlurmProfile", "ServerDeploymentConfig", "ServingImageInspection", + "SlurmConfigBuilderError", + "SlurmConfigLoadError", "SlurmProfile", "SlurmProfileCatalog", "SubmissionConfig", "VllmServerConfig", "injected_profile", + "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/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py new file mode 100644 index 000000000..679d818b3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -0,0 +1,200 @@ +# 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.images import ImageRef +from data_designer.slurm.config.run import ( + ArrayTasksConfig, + BuilderInput, + ClientConfig, + ClientDependencies, + DataDesignerSlurmConfig, + InputBindings, + InvocationConfig, + InvocationDiagnostics, + MCPProviderConfig, + OutputConfig, + 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() + self._output = OutputConfig() + + @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, + diagnostics: InvocationDiagnostics | Mapping[str, object] | 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), + "diagnostics": {} if diagnostics is None else diagnostics, + }, + ) + return self + + def with_client( + self, + *, + image: ImageRef | Mapping[str, object], + cpus: int = 32, + dependencies: ClientDependencies | Mapping[str, object] | None = None, + ) -> DataDesignerSlurmConfigBuilder: + """Set the separate zero-GPU Data Designer client declaration.""" + self._client = _validate_model( + ClientConfig, + { + "cpus": cpus, + "image": image, + "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 with_output(self, **values: object) -> DataDesignerSlurmConfigBuilder: + """Set typed dataset output intent.""" + self._output = _validate_model(OutputConfig, 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 self._client is None: + missing.append("client") + if not self._deployments: + missing.append("deployment") + if missing: + raise SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") + assert self._invocation is not None + assert self._client is not None + return _validate_model( + DataDesignerSlurmConfig, + { + "schema_version": 1, + "name": self._name, + "builder": self._builder, + "invocation": self._invocation, + "client": self._client, + "deployments": self._deployments, + "array_tasks": self._array_tasks, + "submission": self._submission, + "output": self._output, + }, + ) + + 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/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/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py new file mode 100644 index 000000000..684465871 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -0,0 +1,227 @@ +# 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 ValidationError +from yaml.nodes import MappingNode + +from data_designer.slurm._errors import format_parse_error, format_validation_error +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 DataDesignerSlurmConfig + +PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" +DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" + +_ConfigT = TypeVar("_ConfigT", 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_profile_catalog(path: str | Path) -> SlurmProfileCatalog: + """Load one strict local YAML or JSON cluster-profile catalog.""" + return _load_config(path, SlurmProfileCatalog) + + +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) + 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 _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 index a6a0b98c3..a5c8c91c0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -76,6 +76,8 @@ 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 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 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 index cd323b763..0d9380a92 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -306,11 +306,11 @@ def validate_topology(self) -> ServerDeploymentConfig: class ArrayTasksConfig(AuthoredConfig): count: PositiveInt = 1 - max_concurrent: PositiveInt = 1 + max_concurrent: PositiveInt | None = None @model_validator(mode="after") def validate_concurrency(self) -> ArrayTasksConfig: - if self.max_concurrent > self.count: + 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index b8c0b200e..226b9899e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -36,11 +36,11 @@ # TODO: Remove these compatibility exports after Stage 2 branches import shared scalars from slurm.types. ModelAlias = str -_Key = TypeVar("_Key") -_Value = TypeVar("_Value") +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") -class _FrozenList(list[_Value]): +class _FrozenList(list[_ValueT]): """List that retains JSON compatibility without exposing mutation.""" def _immutable(self, *args: object, **kwargs: object) -> None: @@ -61,7 +61,7 @@ def _immutable(self, *args: object, **kwargs: object) -> None: sort = _immutable -class _FrozenDict(dict[_Key, _Value]): +class _FrozenDict(dict[_KeyT, _ValueT]): """Dictionary that retains JSON compatibility without exposing mutation.""" def _immutable(self, *args: object, **kwargs: object) -> None: @@ -94,8 +94,8 @@ class ContractValue(BaseModel): model_config = ConfigDict( extra="forbid", frozen=True, - allow_inf_nan=False, hide_input_in_errors=True, + allow_inf_nan=False, protected_namespaces=(), strict=True, validate_default=True, @@ -169,6 +169,26 @@ def compute_canonical_json_sha256(value: object) -> Sha256Digest: 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 @@ -284,6 +304,7 @@ class ResumeWorkspace(ContractValue): "Duration", "EnvironmentName", "Identifier", + "ModelAlias", "NetworkPort", "NonNegativeDuration", "RecordRange", @@ -293,6 +314,11 @@ class ResumeWorkspace(ContractValue): "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", 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 index 203cfe14a..e8941df39 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -91,7 +91,9 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire node_count = max(node_indices) + 1 array = "0" if plan.array_tasks.count > 1: - array = f"0-{plan.array_tasks.count - 1}%{plan.array_tasks.max_concurrent}" + array = f"0-{plan.array_tasks.count - 1}" + if plan.array_tasks.max_concurrent is not None: + array = f"{array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ ("job-name", plan.submission.job_name), @@ -106,7 +108,6 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire if profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) elif profile.scheduler.mem_per_gpu is not None: - # TODO(#875): Remove this defense once config and plan validation make this state unrepresentable. raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") if profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) 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 index 160b7ef77..46dfcbcea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -23,12 +23,10 @@ ResolvedTopology, ResumeWorkspace, ) -from data_designer.slurm.planning.validation import PlanContractError, validate_resolved_plan __all__ = [ "ArtifactReference", "LockedPackage", - "PlanContractError", "PlannedShard", "PortClaim", "RecordRange", @@ -43,5 +41,4 @@ "ResolvedSubmission", "ResolvedTopology", "ResumeWorkspace", - "validate_resolved_plan", ] 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..964a282c3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -0,0 +1,254 @@ +# 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, +) +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, + 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, + output=effective.output, + 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] = [] + next_node_index = 0 + for index, (authored, image) in enumerate( + zip(effective.authored.deployments, effective.deployment_images, strict=True) + ): + topology = ResolvedTopology.derive( + node_count=authored.resources.nodes, + gpus_per_node=effective.resolved_gpus_per_node, + 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}" + node_indices = tuple(range(next_node_index, next_node_index + authored.resources.nodes)) + ports = _compile_deployment_ports(deployment_id, node_indices, topology) + 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=effective.resolved_gpus_per_node, + topology=topology, + ports=ports, + ) + ) + next_node_index += authored.resources.nodes + return tuple(resolved) + + +def _compile_deployment_ports( + deployment_id: str, + node_indices: tuple[int, ...], + topology: ResolvedTopology, +) -> 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 + 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, + image=effective.client_image, + 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 index 9b71cc31f..ca2a57c80 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -28,6 +28,7 @@ ArrayTasksConfig, ClientConfig, ClientDependencies, + InputBindings, InvocationConfig, ServerDeploymentConfig, SubmissionConfig, @@ -37,15 +38,20 @@ ContractRecord, ContractValue, Identifier, + 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 @@ -158,9 +164,9 @@ 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[str, ...] - referenced_model_aliases: tuple[str, ...] = () + model_aliases: tuple[ModelAlias, ...] @model_validator(mode="after") def validate_input(self) -> ResolvedBuilderInput: @@ -170,20 +176,15 @@ def validate_input(self) -> ResolvedBuilderInput: 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_canonical_json_sha256(self.inline) - model_aliases, referenced_aliases = _extract_builder_aliases(self.inline) - if self.model_aliases != model_aliases: + 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") - if self.referenced_model_aliases != referenced_aliases: - raise ValueError("resolved referenced 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 len(self.referenced_model_aliases) != len(set(self.referenced_model_aliases)): - raise ValueError("resolved builder referenced aliases must be unique") if self.content_sha256 != expected_digest: raise ValueError("builder content digest does not match the resolved input") return self @@ -191,6 +192,7 @@ def validate_input(self) -> ResolvedBuilderInput: class ResolvedInvocation(ContractValue): authored: InvocationConfig + effective_input_bindings: InputBindings effective_run_config: dict[str, JsonValue] @field_validator("effective_run_config") @@ -401,6 +403,8 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: profile = self.selected_profile.profile 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") if any(deployment.gpus_per_node != self.resolved_gpus_per_node for deployment in self.deployments): raise ValueError("every deployment must use the resolved profile GPU count") if tuple(profile.container_mounts) != self.container_mounts: @@ -413,20 +417,33 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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 not set(aliases).issubset(self.builder.model_aliases): - raise ValueError("each deployment alias must match a resolved Data Designer model alias") - if not set(self.builder.referenced_model_aliases).issubset(aliases): - raise ValueError("each referenced Data Designer model alias requires a deployment") + if set(aliases) != set(self.builder.model_aliases): + raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") node_indices = tuple(index for deployment in self.deployments for index in deployment.node_indices) if node_indices != tuple(range(len(node_indices))): raise ValueError("deployment nodes must be disjoint and contiguous in authored order") 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 != RunConfig().non_inference_max_parallel_workers: - raise ValueError("default non-inference worker count must match the Data Designer RunConfig default") + 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 @@ -453,15 +470,20 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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) - if not _is_below(self.output.root, profile.workspace_root): + if not is_path_below(self.output.root, profile.workspace_root): raise ValueError("resolved output root must be below the selected workspace_root") + if paths_overlap(self.output.root, managed_assets_path): + raise ValueError("resolved output root must not overlap managed assets") shards_root = posixpath.join(run_root, "shards") - if ( - self.output.root == shards_root - or _is_below(self.output.root, shards_root) - or _is_below(shards_root, self.output.root) - ): + if paths_overlap(self.output.root, shards_root): raise ValueError("resolved output root must not overlap the run shard workspace") return self @@ -474,7 +496,7 @@ def _validate_shards(self, run_root: str) -> None: shard_ids: list[ShardId] = [] workspace_paths: list[str] = [] partition_paths: list[str] = [] - requires_partition = self.invocation.authored.input_bindings.seed_path is not None + 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") @@ -510,47 +532,3 @@ def _validate_shards(self, run_root: str) -> None: 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 _is_below(path: str, root: str) -> bool: - return path != root and posixpath.commonpath((path, root)) == root - - -def _extract_builder_aliases(builder: dict[str, JsonValue]) -> tuple[tuple[str, ...], tuple[str, ...]]: - 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[str] = [] - 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"]) - - referenced_aliases: list[str] = [] - - def collect(value: JsonValue, *, key: str | None = None) -> None: - if key == "model_configs": - return - if key == "model_alias" or (key is not None and key.endswith("_model_alias")): - if not isinstance(value, str): - raise ValueError(f"builder {key} must be a string") - referenced_aliases.append(value) - return - if key == "model_aliases": - if not isinstance(value, list) or any(not isinstance(alias, str) for alias in value): - raise ValueError("builder model_aliases must be a list of strings") - referenced_aliases.extend(value) - return - if isinstance(value, dict): - for child_key, child in value.items(): - collect(child, key=child_key) - elif isinstance(value, list): - for child in value: - collect(child) - - collect(data_designer) - return tuple(model_aliases), tuple(dict.fromkeys(referenced_aliases)) 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..3fc47b05f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -0,0 +1,428 @@ +# 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 ClientImageInspection, 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, + ResolvedDependencyLock, + ResolvedImage, + ResolvedInvocation, + ResolvedOutput, + ResolvedSubmission, +) + +__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, + ResolvedOutput, + 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 + resolved_gpus_per_node: PositiveInt + builder: ResolvedBuilderInput + builder_payload: dict[str, JsonValue] | None = None + invocation: ResolvedInvocation + client_image: ResolvedImage + deployment_images: tuple[ResolvedImage, ...] + dependency_lock: ResolvedDependencyLock + submission: ResolvedSubmission + output: ResolvedOutput + runtime_bundle: ArtifactReference + + +def resolve_slurm_config( + authored: DataDesignerSlurmConfig, + *, + selected_profile: SelectedSlurmProfile, + client_image: ResolvedImage, + deployment_images: tuple[ResolvedImage, ...], + dependency_lock: ResolvedDependencyLock, + runtime_bundle: ArtifactReference, + run_id: str, + package_version: str, + 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, + resolved_gpus_per_node=gpus_per_node, + builder=builder, + builder_payload=resolved_builder_payload, + invocation=_materialize_invocation(authored, workspace_root), + client_image=client_image, + deployment_images=deployment_images, + dependency_lock=dependency_lock, + submission=_materialize_submission(authored, selected_profile), + output=_materialize_output(authored, run_root), + 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 + run_root = posixpath.join(workspace_root, "runs", effective.run_id) + 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") + expected_output = _materialize_output(authored, run_root) + if effective.output != expected_output: + raise SlurmConfigResolutionError("resolved output does not match the authored output") + + _validate_dependency_resolution(authored, effective.client_image, effective.dependency_lock) + _validate_resolved_images(authored, effective.client_image, effective.deployment_images) + _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) + _validate_output_destination( + effective.output.root, + workspace_root, + run_root, + managed_assets_path=managed_assets_path, + ) + if effective.output.partitions > authored.invocation.num_records: + raise SlurmConfigResolutionError("output partitions must not exceed requested records") + 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_output(authored: DataDesignerSlurmConfig, run_root: str) -> ResolvedOutput: + return ResolvedOutput( + root=authored.output.root or posixpath.join(run_root, "output"), + format=authored.output.format, + partitions=authored.output.partitions, + require_exact_record_count=authored.output.require_exact_record_count, + ) + + +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 + if authored.output.format != "parquet": + raise SlurmConfigResolutionError("multi-shard runs require parquet output") + + 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_image: ResolvedImage, + dependency_lock: ResolvedDependencyLock, +) -> None: + inspection = client_image.inspection.inspection + if not isinstance(inspection, ClientImageInspection): + raise SlurmConfigResolutionError("resolved client image lacks dependency inspection facts") + if dependency_lock.client_image_sha256 != client_image.sha256: + raise SlurmConfigResolutionError("dependency lock does not match the resolved client image") + if dependency_lock.python_abi != inspection.python_abi: + raise SlurmConfigResolutionError("dependency lock Python ABI does not match the client image") + if dependency_lock.image_distributions != inspection.distributions: + raise SlurmConfigResolutionError("dependency lock inventory does not match the client image") + 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, + client_image: ResolvedImage, + deployment_images: tuple[ResolvedImage, ...], +) -> None: + _validate_resolved_image_identity(client_image) + if client_image.kind is not ImageKind.CLIENT: + raise SlurmConfigResolutionError("resolved client image must contain client inspection facts") + if client_image.authored_ref != authored.client.image: + raise SlurmConfigResolutionError("resolved client image does not match the authored reference") + 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_output_destination( + output_root: str, + workspace_root: str, + run_root: str, + *, + managed_assets_path: str, +) -> None: + if not is_path_below(output_root, workspace_root): + raise SlurmConfigResolutionError("output root must be below the selected workspace_root") + reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) + if any(paths_overlap(output_root, path) for path in reserved): + raise SlurmConfigResolutionError("output root must not overlap package-managed workspace state") + if paths_overlap(output_root, managed_assets_path): + raise SlurmConfigResolutionError("output root must not overlap managed assets") + runs_root = posixpath.join(workspace_root, "runs") + run_output_root = posixpath.join(run_root, "output") + if paths_overlap(output_root, runs_root) and not ( + output_root == run_output_root or is_path_below(output_root, run_output_root) + ): + raise SlurmConfigResolutionError("output root must not overlap another package-managed run") + + +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 index c47ebc512..5cf7ee9ba 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -3,22 +3,22 @@ from __future__ import annotations +import posixpath + from pydantic import JsonValue from data_designer.config import RunConfig from data_designer.slurm.config.images import ClientImageInspection 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, - _extract_builder_aliases, ) -class PlanContractError(ValueError): - """Raised when a resolved plan does not match its authored inputs.""" - - def validate_resolved_plan( authored: DataDesignerSlurmConfig, dependency_lock: ResolvedDependencyLock, @@ -32,6 +32,14 @@ def validate_resolved_plan( "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, @@ -45,22 +53,28 @@ def validate_resolved_plan( _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 PlanContractError("sourced builder validation requires its resolved payload") - model_aliases, referenced_aliases = _extract_builder_aliases(builder_payload) + 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.referenced_model_aliases == referenced_aliases, - "resolved referenced 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 @@ -126,7 +140,7 @@ def validate_resolved_plan( def _require(condition: bool, message: str) -> None: if not condition: - raise PlanContractError(message) + raise SlurmPlanContractError(message) def _require_json_subset(actual: JsonValue, expected: JsonValue, *, path: str) -> None: 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..e1cd641a2 --- /dev/null +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.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 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, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ProfileSelectionSource, + SlurmConfigBuilderError, + SlurmConfigLoadError, + SlurmProfileCatalog, + 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(image={"name": "dd-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 not tuple(tmp_path.iterdir()) + + +def test_builder_requires_complete_authored_intent() -> None: + builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") + + with pytest.raises(SlurmConfigBuilderError, match="invocation, client, 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_builder_rejects_unsupported_output_format(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"}), + ("with_output", {"partitions": 0}), + ], +) +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(image={"name": "dd-client", "api_key": secret}) + + 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( + image={"name": "dd-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_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/contracts/golden/authored_run_single.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json index 346b4500c..7a852dbcd 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -1,7 +1,7 @@ { "array_tasks": { "count": 1, - "max_concurrent": 1 + "max_concurrent": null }, "builder": { "inline": { 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 index 48427e849..1468aac49 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -9,7 +9,7 @@ }, "builder": { "authored_source": null, - "content_sha256": "f37227ca7c67e203abe881c0228b4308a8e741364296d293159a1201949732f2", + "content_sha256": "2d4536fc8ff5a32e06269769bb563cb354c53cee95f45da31c9c6781130b4222", "inline": { "data_designer": { "columns": [], @@ -31,7 +31,6 @@ "generator", "judge" ], - "referenced_model_aliases": [], "source": null }, "client": { @@ -355,6 +354,10 @@ "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, @@ -385,7 +388,7 @@ "resolved_gpus_per_node": 8, "run_id": "run-001", "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" }, "schema_version": 1, 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 index 4b86c8657..b2747d4c0 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -1,15 +1,15 @@ { "array_tasks": { "count": 1, - "max_concurrent": 1 + "max_concurrent": null }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" + "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" }, "builder": { "authored_source": null, - "content_sha256": "b3ef5fc1fe675a8e004633f84842ac60cf82d5ba3dc68b4d50ee4438448b0570", + "content_sha256": "0bac3a88498774a30b64f0a62617511c8e75601ca54e0b775d86484bb13788d8", "inline": { "data_designer": { "columns": [], @@ -25,7 +25,6 @@ "model_aliases": [ "generator" ], - "referenced_model_aliases": [], "source": null }, "client": { @@ -201,6 +200,10 @@ "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, @@ -231,7 +234,7 @@ "resolved_gpus_per_node": 8, "run_id": "run-single", "runtime_bundle": { - "path": "/workspace/primary/runtime/runtime.tar.gz", + "path": "/workspace/primary/runtime/eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee.tar.gz", "sha256": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" }, "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index dcf9c876d..58651662a 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -695,6 +695,7 @@ def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesig @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"): diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index cb907c010..fa45b8291 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -10,18 +10,18 @@ from pydantic import ValidationError from data_designer.slurm.config import BuilderInput, ClientDependencies, DataDesignerSlurmConfig -from data_designer.slurm.contracts import compute_canonical_json_sha256 +from data_designer.slurm.contracts import compute_serialized_json_sha256 from data_designer.slurm.planning import ( ArtifactReference, - PlanContractError, ResolvedBuilderInput, ResolvedDependencyLock, ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedSubmission, ResolvedTopology, - validate_resolved_plan, ) +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: @@ -105,7 +105,7 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu inline = payload["inline"] assert isinstance(inline, dict) inline["plugin_config"] = {"api_key": "VERY_SECRET_VALUE"} - payload["content_sha256"] = compute_canonical_json_sha256(inline) + 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)) @@ -134,6 +134,12 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu lambda payload: payload["output"].update(root="/outside/output"), 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"), + lambda payload: payload["output"].update(root="/workspace/primary/managed-assets"), ], ) def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, mutator: object) -> None: @@ -152,13 +158,13 @@ def test_plan_rejects_unmaterialized_run_config(multi_node_plan: ResolvedSlurmRu ResolvedSlurmRunPlan.model_validate(payload) -def test_plan_preserves_default_non_inference_worker_count( +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="RunConfig default"): + with pytest.raises(ValidationError, match="Data Designer default"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -166,12 +172,31 @@ 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"] = 32 - payload["invocation"]["effective_run_config"]["non_inference_max_parallel_workers"] = 32 + 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"] == 32 + 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: @@ -207,22 +232,12 @@ def test_plan_rejects_deployment_alias_missing_from_builder(multi_node_plan: Res "model_configs" ][:1] payload["builder"]["model_aliases"] = ["generator"] - payload["builder"]["content_sha256"] = compute_canonical_json_sha256(payload["builder"]["inline"]) + 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_plan_rejects_referenced_alias_without_deployment(multi_node_plan: ResolvedSlurmRunPlan) -> None: - payload = multi_node_plan.model_dump(mode="json") - payload["builder"]["inline"]["data_designer"]["columns"] = [{"model_alias": "missing"}] - payload["builder"]["referenced_model_aliases"] = ["missing"] - payload["builder"]["content_sha256"] = compute_canonical_json_sha256(payload["builder"]["inline"]) - - with pytest.raises(ValidationError, match="referenced Data Designer model 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 @@ -273,19 +288,19 @@ def test_sourced_builder_validation_requires_resolved_payload( 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.json", "sha256": "a" * 64}, + "source": {"path": "/workspace/primary/runs/run-001/builder-config.json", "sha256": builder_digest}, "inline": None, - "content_sha256": "a" * 64, + "content_sha256": builder_digest, "model_aliases": ["generator", "judge"], - "referenced_model_aliases": [], } sourced_plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - with pytest.raises(PlanContractError, match="resolved payload"): + with pytest.raises(SlurmPlanContractError, match="resolved payload"): validate_resolved_plan(sourced_authored, dependency_lock, sourced_plan) assert ( @@ -299,6 +314,26 @@ def test_sourced_builder_validation_requires_resolved_payload( ) +@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", [ @@ -394,7 +429,7 @@ def test_cross_record_validation_rejects_authored_digest( } ) - with pytest.raises(PlanContractError, match="authored config digest"): + with pytest.raises(SlurmPlanContractError, match="authored config digest"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -413,7 +448,7 @@ def test_cross_record_validation_rejects_dependency_lock_digest( ) invalid = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="dependency lock digest"): + with pytest.raises(SlurmPlanContractError, match="dependency lock digest"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -431,7 +466,7 @@ def test_cross_record_validation_binds_authored_lock_source( plan_payload["client"]["authored"] = authored.client.model_dump(mode="json") plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(plan_payload)) - with pytest.raises(PlanContractError, match="authored lock file"): + with pytest.raises(SlurmPlanContractError, match="authored lock file"): validate_resolved_plan(authored, dependency_lock, plan) lock_payload = dependency_lock.model_dump(mode="json") @@ -454,7 +489,7 @@ def test_cross_record_validation_binds_authored_lock_source( ) } ) - with pytest.raises(PlanContractError, match="present for authored requirements"): + with pytest.raises(SlurmPlanContractError, match="present for authored requirements"): validate_resolved_plan(authored_run, matching_lock, unexpected_source_plan) matching_plan = plan.model_copy( @@ -487,7 +522,7 @@ def test_cross_record_validation_rejects_python_abi( ) invalid_plan = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="Python ABI"): + with pytest.raises(SlurmPlanContractError, match="Python ABI"): validate_resolved_plan(authored_run, invalid_lock, invalid_plan) @@ -506,7 +541,7 @@ def test_cross_record_validation_rejects_image_inventory( ) invalid_plan = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(PlanContractError, match="image inventory"): + with pytest.raises(SlurmPlanContractError, match="image inventory"): validate_resolved_plan(authored_run, invalid_lock, invalid_plan) @@ -520,7 +555,7 @@ def test_cross_record_validation_rejects_changed_invocation( ) invalid = multi_node_plan.model_copy(update={"invocation": invocation}) - with pytest.raises(PlanContractError, match="invocation"): + with pytest.raises(SlurmPlanContractError, match="invocation"): validate_resolved_plan(authored_run, dependency_lock, invalid) @@ -533,5 +568,5 @@ def test_cross_record_validation_preserves_explicit_run_config( payload["invocation"]["effective_run_config"]["buffer_size"] = 16384 invalid = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) - with pytest.raises(PlanContractError, match="run_config.buffer_size"): + 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 index 1e4086f5e..dcf7addda 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -124,3 +124,12 @@ def test_profile_requires_explicit_version(profile_catalog: SlurmProfileCatalog) 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/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 7e993b116..4568b5566 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" + "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" }, "run_id": "run-single", "scheduler": { @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "3d092f4a6d3731d048c519b5e5a41b5f0c2973a4b3aff6459710dfbe64f66e18" + "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" + "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index c80c95f5f..08d93d747 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -36,6 +36,16 @@ def test_renderer_matches_contract_bound_goldens( assert render_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() +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: @@ -98,7 +108,8 @@ def test_renderer_rejects_mem_per_gpu_without_a_slurm_gpu_request( "scheduler": SchedulerProfile(account="research", partition="batch", mem_per_gpu="80G"), } ) - plan = single_node_plan.model_copy(update={"selected_profile": injected_profile(profile)}) + 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) 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..f9a2389fa --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -0,0 +1,904 @@ +# 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_image": expected.client.image, + "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 {plan.output.root for plan in plans} == {"/workspace/primary/runs/run-single/output"} + 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.output.root == "/workspace/lab/runs/run-single/output" + 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( + "output_update", + [ + {"root": "/outside/output"}, + {"root": "/workspace/primary/images/output"}, + {"root": "/workspace/primary/runtime/output"}, + {"root": "/workspace/primary/managed-assets/output"}, + {"root": "/workspace/primary/runs/other-run/output"}, + {"root": "/workspace/primary/runs/run-single/shards"}, + {"partitions": 9}, + ], +) +def test_resolution_rejects_invalid_output_destinations( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, + output_update: dict[str, object], +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["output"].update(output_update) + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match="output"): + _resolve_fixture(authored, dependency_lock_single, single_node_plan) + + +def test_resolution_rejects_output_overlapping_explicit_managed_assets( + authored_run_single: DataDesignerSlurmConfig, + dependency_lock_single: ResolvedDependencyLock, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + payload = authored_run_single.model_dump(mode="json") + payload["invocation"]["input_bindings"]["managed_assets_path"] = "/workspace/primary/custom-assets" + payload["output"]["root"] = "/workspace/primary/custom-assets/output" + authored = DataDesignerSlurmConfig.model_validate(payload) + + with pytest.raises(SlurmConfigResolutionError, match="managed assets"): + _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) + + jsonl_output = authored_run.output.model_copy(update={"format": "jsonl"}) + jsonl_authored = authored_run.model_copy(update={"output": jsonl_output}) + with pytest.raises(SlurmPlanCompilationError, match="parquet output"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": jsonl_authored, + "output": effective.output.model_copy(update={"format": "jsonl"}), + } + ) + ) + + other_output = "/workspace/primary/runs/other-run/output" + other_authored = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"root": other_output})} + ) + with pytest.raises(SlurmPlanCompilationError, match="another package-managed run"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": other_authored, + "output": effective.output.model_copy(update={"root": other_output}), + } + ) + ) + + partitioned_authored = authored_run.model_copy( + update={"output": authored_run.output.model_copy(update={"partitions": 101})} + ) + with pytest.raises(SlurmPlanCompilationError, match="requested records"): + SlurmRunCompiler.compile( + effective.model_copy( + update={ + "authored": partitioned_authored, + "output": effective.output.model_copy(update={"partitions": 101}), + } + ) + ) + + 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})} + ) + ) + with pytest.raises(SlurmPlanCompilationError, match="resolved output"): + SlurmRunCompiler.compile( + effective.model_copy(update={"output": effective.output.model_copy(update={"format": "jsonl"})}) + ) + 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) + client_image = effective.client_image.model_copy(update={"sha256": "a" * 64}) + dependency_lock = effective.dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + + with pytest.raises(SlurmPlanCompilationError, match="inspection record"): + SlurmRunCompiler.compile( + effective.model_copy(update={"client_image": client_image, "dependency_lock": dependency_lock}) + ) + + 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_image", "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_image": + invalid_value = effective.client_image.model_copy(update={"path": "/workspace/images/client.txt"}) + 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", "output_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", + ), + ({}, {"format": "jsonl"}, "parquet output"), + ], +) +def test_resolution_rejects_unshardable_big_iron_fields( + authored_run: DataDesignerSlurmConfig, + dependency_lock: ResolvedDependencyLock, + multi_node_plan: ResolvedSlurmRunPlan, + builder_update: dict[str, object], + output_update: dict[str, object], + message: str, +) -> None: + payload = authored_run.model_dump(mode="json") + payload["builder"]["inline"]["data_designer"].update(builder_update) + payload["output"].update(output_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"}, + }, + } + ) + payload["output"]["format"] = "jsonl" + authored = DataDesignerSlurmConfig.model_validate(payload) + + plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) + + assert plan.array_tasks.count == 1 + assert plan.output.format == "jsonl" + + +@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_image_sha256": "a" * 64}) + + with pytest.raises(SlurmConfigResolutionError, match="client image"): + _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/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index f19973cbd..fdd83ee99 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -11,7 +11,7 @@ import data_designer.slurm.serving.resolver as resolver_module from data_designer.slurm.config import QueueBackpressureConfig -from data_designer.slurm.contracts import pretty_json +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 ( @@ -261,6 +261,11 @@ 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 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index b3c1e40c8..fc36170e8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -10,10 +10,10 @@ set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +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="00e206f4b759ca06336c02d3694f7289551d1a0b9ba50c05f37cdbf7a213756d" +readonly DD_PLAN_SHA256="cd95eb995dd9705d12336b9760ca1f354cdd7614dce992cc4f92a6faa47dbfdc" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 115d998da..e3f8529d5 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -10,10 +10,10 @@ set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -readonly DD_RUNTIME_ARCHIVE="/workspace/primary/runtime/runtime.tar.gz" +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="897d35e4145200a2e1d82748948cc281e23ef0e2c389c12c946b80c93ce9d2b5" +readonly DD_PLAN_SHA256="14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 34d947543..9880d08e8 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="244fe9b160b361645e3f9be13d40d5d6157563a69fa37037fa974d2a3acead2f", + expected_fixture_sha256="65032718bfc8fbd6c60700add08ebc098ba2e64858c23808decbefe5a8153d91", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="37681e8805eb81ca9c4e1fc2da88b14031d7a150c9fd729104ed875ed2503224", + expected_fixture_sha256="58c3a2eddc9a0d377eb0a90655e70a6b2484cade0c52b2a01b9b388dcd4667e2", ) @@ -55,7 +55,7 @@ def _assert_script_matches_plan( ) 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: + 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) From 7d9e92091d652cf53ef995f94cb433179f8f6178 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:56:05 -0300 Subject: [PATCH 12/62] feat: add Slurm service foundation (#894) * feat: add Slurm service foundation Signed-off-by: Andre Manoel * fix: harden Slurm service contracts Signed-off-by: Andre Manoel * fix: bind Slurm service results Signed-off-by: Andre Manoel * fix: clarify Slurm service boundaries Separate run planning from attempt rendering so retries can reuse an immutable plan. Keep injected implementation seams private and align public service failures with the Data Designer error hierarchy. Signed-off-by: Andre Manoel * fix: tighten Slurm service contracts * test: fix Slurm seam export guard * fix: define Slurm service boundary contracts * fix: expose Slurm service dependency contracts --------- Signed-off-by: Andre Manoel --- .../data_designer/slurm/services/__init__.py | 28 + .../data_designer/slurm/services/benchmark.py | 100 ++++ .../data_designer/slurm/services/errors.py | 100 ++++ .../data_designer/slurm/services/images.py | 66 +++ .../src/data_designer/slurm/services/run.py | 102 ++++ .../tests/services/test_services.py | 534 ++++++++++++++++++ .../tests/slurm_test_fakes/__init__.py | 12 + .../tests/slurm_test_fakes/services.py | 139 +++++ 8 files changed, 1081 insertions(+) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/images.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/run.py create mode 100644 packages/data-designer-slurm/tests/services/test_services.py create mode 100644 packages/data-designer-slurm/tests/slurm_test_fakes/services.py 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..07c832814 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -0,0 +1,28 @@ +# 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 + +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 SlurmImageResolver, SlurmImageService +from data_designer.slurm.services.run import SlurmBatchScriptRenderer, SlurmRunPlanner, SlurmRunService + +__all__ = [ + "SlurmBatchScriptRenderer", + "SlurmBenchmarkBackend", + "SlurmBenchmarkService", + "SlurmImageResolver", + "SlurmImageService", + "SlurmRunPlanner", + "SlurmRunService", + "SlurmServiceError", + "SlurmServiceErrorCode", + "SlurmServiceOperation", +] 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..4c4fd2192 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py @@ -0,0 +1,100 @@ +# 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 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) -> BenchmarkManifest: + """Persist and start the ordinary child runs for one benchmark. + + Any non-``INTERNAL`` service error must contain a caller-safe message. + """ + + def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + """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) -> 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") + + def run_benchmark() -> BenchmarkManifest: + manifest = self._backend.run(config) + 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) -> 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. + + 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: + raise _make_invalid_request_error(operation, "refresh_state must be a boolean") + + def analyze_benchmark() -> BenchmarkReport: + report = self._backend.analyze(validated_id, refresh_state=refresh_state) + 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..1a24686ae --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -0,0 +1,100 @@ +# 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" + RESOLVE_IMAGE = "resolve_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/images.py b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py new file mode 100644 index 000000000..848eaeced --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -0,0 +1,66 @@ +# 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 data_designer.slurm.config import ImageKind, ImageRef +from data_designer.slurm.planning import ResolvedImage +from data_designer.slurm.services.errors import ( + SlurmServiceOperation, + _invoke_service_backend, + _make_invalid_request_error, +) + + +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 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``. + """ + + def __init__(self, resolver: SlurmImageResolver) -> None: + self._resolver = resolver + + 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) 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..b2aaaf04e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -0,0 +1,102 @@ +# 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 typing import Protocol + +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.errors import ( + SlurmServiceOperation, + _invoke_service_backend, + _make_invalid_request_error, +) + + +class SlurmRunPlanner(Protocol): + """Resolve authored run configs through a supported service dependency.""" + + def plan(self, config: DataDesignerSlurmConfig) -> 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 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``. + """ + + def __init__(self, planner: SlurmRunPlanner, renderer: SlurmBatchScriptRenderer) -> None: + self._planner = planner + self._renderer = renderer + + def plan(self, config: DataDesignerSlurmConfig) -> 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") + + def build_plan() -> ResolvedSlurmRunPlan: + plan = self._planner.plan(config) + 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) 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..b937b5b06 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -0,0 +1,534 @@ +# 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 + +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, + ImageKind, + ImageRef, +) +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmBatchScriptRenderer, + SlurmBenchmarkBackend, + SlurmBenchmarkService, + SlurmImageResolver, + SlurmImageService, + SlurmRunPlanner, + SlurmRunService, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) + +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_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: + reference = single_node_plan.client.image.authored_ref + image = single_node_plan.client.image + resolver = FakeImageResolver((((reference, ImageKind.CLIENT), image),)) + + result = SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.CLIENT) + + assert result is image + assert resolver.calls == [(reference, ImageKind.CLIENT)] + resolver.assert_complete() + + +def test_image_service_rejects_invalid_resolution(single_node_plan: ResolvedSlurmRunPlan) -> None: + reference = single_node_plan.client.image.authored_ref + resolver = FakeImageResolver((((reference, ImageKind.CLIENT), object()),)) # type: ignore[arg-type] + + with pytest.raises(SlurmServiceError) as caught: + SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.CLIENT) + + 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.CLIENT)] + 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.client.image.authored_ref, expected_kind="client") # 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.CLIENT) # type: ignore[arg-type] + + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + + +@pytest.mark.parametrize("mismatch", ["reference", "kind"]) +def test_image_service_rejects_uncorrelated_results( + single_node_plan: ResolvedSlurmRunPlan, + mismatch: str, +) -> None: + reference = single_node_plan.client.image.authored_ref + if mismatch == "reference": + source = single_node_plan.client.image + image = source.model_copy(update={"authored_ref": ImageRef(name="other-client")}) + else: + source = single_node_plan.deployments[0].image + image = source.model_copy(update={"authored_ref": reference}) + resolver = FakeImageResolver((((reference, ImageKind.CLIENT), image),)) + + with pytest.raises(SlurmServiceError) as caught: + SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.CLIENT) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RESOLVE_IMAGE + assert resolver.calls == [(reference, ImageKind.CLIENT)] + 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 + backend = FakeBenchmarkBackend( + run_responses=((benchmark_config, benchmark_manifest),), + analysis_responses=(((benchmark_manifest.benchmark_id, True), benchmark_report),), + ) + service = SlurmBenchmarkService(backend) + + assert service.run(benchmark_config) is benchmark_manifest + assert service.analyze(benchmark_manifest.benchmark_id, refresh_state=True) is benchmark_report + assert backend.run_calls == [benchmark_config] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, True)] + backend.assert_complete() + + +def test_benchmark_service_rejects_invalid_manifest( + benchmark_config: DataDesignerSlurmBenchmarkConfig, +) -> None: + backend = FakeBenchmarkBackend(run_responses=((benchmark_config, object()),)) # type: ignore[arg-type] + + 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 == [benchmark_config] + backend.assert_complete() + + +def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkManifest) -> None: + backend = FakeBenchmarkBackend( + analysis_responses=(((benchmark_manifest.benchmark_id, False), object()),), # type: ignore[arg-type] + ) + + 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)] + 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}) + backend = FakeBenchmarkBackend(run_responses=((benchmark_config, manifest),)) + + with pytest.raises(SlurmServiceError) as caught: + SlurmBenchmarkService(backend).run(benchmark_config) + + assert caught.value.code is SlurmServiceErrorCode.INTERNAL + assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK + assert backend.run_calls == [benchmark_config] + 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"), + [("", False), ("invalid/id", False), ("benchmark-001", 1)], +) +def test_benchmark_service_validates_analysis_actions( + benchmark_id: object, + refresh_state: object, +) -> None: + service = SlurmBenchmarkService(FakeBenchmarkBackend()) + + with pytest.raises(SlurmServiceError) as caught: + service.analyze(benchmark_id, refresh_state=refresh_state) # type: ignore[arg-type] + + 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), 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)] + 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), + ("SlurmImageResolver", SlurmImageResolver), + ("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/slurm_test_fakes/__init__.py b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py index 60d45aa06..d566e22d6 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/__init__.py @@ -8,6 +8,13 @@ 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, @@ -18,12 +25,17 @@ ) __all__ = [ + "FakeBatchScriptRenderer", + "FakeBenchmarkBackend", "FakeClock", "FakeCommandResponse", "FakeDependencyInstaller", "FakeDependencyResolver", + "FakeImageResolver", "FakeInspectionEnvironment", "FakeLogicalEndpoint", + "FakeRunPlanningBackend", + "FakeScriptError", "FakeServingState", "FakeSlurmArray", "FakeSlurmJob", 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..9f05ccb4e --- /dev/null +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py @@ -0,0 +1,139 @@ +# 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 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) -> ResolvedSlurmRunPlan: + """Return the plan scripted for one exact authored config.""" + 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[DataDesignerSlurmBenchmarkConfig, BenchmarkManifest | BaseException]] = (), + analysis_responses: Iterable[tuple[tuple[Identifier, bool], BenchmarkReport | BaseException]] = (), + ) -> None: + self._run_script = _ScriptedResponses(run_responses) + self._analysis_script = _ScriptedResponses(analysis_responses) + self.run_calls = self._run_script.calls + self.analysis_calls = self._analysis_script.calls + + def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + """Return the manifest scripted for one exact benchmark config.""" + return self._run_script.next(config, operation="benchmark run") + + def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + """Return the report scripted for one benchmark and refresh action.""" + return self._analysis_script.next((benchmark_id, refresh_state), 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") From f25162ea283f9ef6b6a49016dbe0d1ceeb2b3907 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 1 Sep 2026 21:43:33 -0600 Subject: [PATCH 13/62] feat: add Slurm image lifecycle jobs (#896) --- .../data_designer/slurm/config/profiles.py | 11 + .../src/data_designer/slurm/images/errors.py | 4 + .../data_designer/slurm/images/inspection.py | 47 +- .../data_designer/slurm/images/lifecycle.py | 347 +++++++++ .../src/data_designer/slurm/images/records.py | 98 ++- .../slurm/images/resources/__init__.py | 6 + .../slurm/images/resources/enroot.rc | 6 + .../slurm/images/resources/inspect_image.py | 181 +++++ .../src/data_designer/slurm/launcher/batch.py | 57 ++ .../data_designer/slurm/launcher/client.py | 16 +- .../data_designer/slurm/launcher/renderer.py | 60 +- .../data_designer/slurm/launcher/runner.py | 16 +- .../contracts/golden/multi_node_plan.json | 9 +- .../contracts/golden/profile_catalog.json | 10 +- .../contracts/golden/single_node_plan.json | 9 +- .../tests/contracts/test_profiles.py | 3 + .../tests/images/test_inspection.py | 83 ++- .../tests/images/test_lifecycle.py | 663 ++++++++++++++++++ .../golden/finalization_chain.json | 4 +- .../tests/launcher/test_client.py | 22 + .../tests/launcher/test_runner.py | 38 + .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../tests/slurm_test_fakes/images.py | 15 +- .../tests/slurm_test_fakes/slurm.py | 10 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- 26 files changed, 1598 insertions(+), 125 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/enroot.rc create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/resources/inspect_image.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/batch.py create mode 100644 packages/data-designer-slurm/tests/images/test_lifecycle.py 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 index a5c8c91c0..f3fc57987 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -34,6 +34,17 @@ class SchedulerProfile(AuthoredConfig): 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): 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 index 7a6d426bf..db5807e21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py @@ -28,3 +28,7 @@ class ImageConflictError(SlurmImageError): 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/inspection.py b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py index d8e1fbcfc..d0bc57abf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/inspection.py @@ -7,9 +7,7 @@ import importlib.metadata import platform -import shutil import sys -from pathlib import Path from typing import Protocol from packaging.utils import canonicalize_name @@ -23,6 +21,10 @@ ) 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 = ( @@ -49,11 +51,8 @@ def get_python_abi(self) -> str: def list_distributions(self) -> tuple[InstalledDistribution, ...]: """Return the installed Python distribution inventory.""" - def get_distribution_version(self, name: str) -> str: - """Return one installed distribution version.""" - - def find_executable(self, name: str) -> str: - """Return one absolute executable path.""" + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" class SystemInspectionEnvironment: @@ -94,22 +93,14 @@ def list_distributions(self) -> tuple[InstalledDistribution, ...]: except ValidationError as error: raise ImageInspectionError("installed distribution inventory is invalid") from error - def get_distribution_version(self, name: str) -> str: - """Return one installed distribution version with a canonical missing-package error.""" + def get_distribution_console_script(self, name: str) -> tuple[str, str]: + """Return one distribution version and its owned console-script path.""" try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError as error: - raise ImageInspectionError(f"required distribution {name!r} is not installed") from error - - def find_executable(self, name: str) -> str: - """Return one resolved executable path with a canonical missing-tool error.""" - path = shutil.which(name) - if path is None: - raise ImageInspectionError(f"required executable {name!r} is not installed") - executable = Path(path) - if not executable.is_absolute(): - raise ImageInspectionError(f"required executable {name!r} did not resolve to an absolute path") - return executable.as_posix() + 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: @@ -130,14 +121,17 @@ def inspect(self, sqsh_sha256: Sha256Digest) -> ImageInspectionRecord: 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=environment.find_executable("pip"), - installer_version=versions_by_name["pip"], + installer_path=installer_path, + installer_version=installer_version, ) return ImageInspectionRecord( schema_version=1, @@ -161,11 +155,12 @@ 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=environment.get_distribution_version("vllm"), - executable_path=environment.find_executable("vllm"), + runtime_version=runtime_version, + executable_path=executable_path, ) return ImageInspectionRecord( schema_version=1, 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..cfb317f7b --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -0,0 +1,347 @@ +# 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 dataclasses import dataclass +from importlib import resources +from pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.config import ImageBuildRequest, 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 +from data_designer.slurm.images.records import ( + ImageLifecycleOperation, + ImageLifecyclePlan, + validate_enroot_mount_path, + validate_oci_source_for_lifecycle, +) +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" +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_MINIMUM_ENROOT_OCI_VERSION = (4, 0) +_MINIMUM_ENROOT_SQSH_VERSION = (3, 5) + + +@dataclass(frozen=True, slots=True) +class PreparedImageLifecycleJob: + """Persisted, checksum-bound files ready for one Slurm submission.""" + + plan: ImageLifecyclePlan + plan_file: ArtifactReference + script_file: ArtifactReference + + +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 + 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) + 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: + raise ImageLifecycleError(f"cannot prepare image lifecycle job {lifecycle_id!r}") from error + return PreparedImageLifecycleJob(plan=plan, plan_file=plan_file, script_file=script_file) + + +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 + 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" +enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" +""" + else: + minimum_major, minimum_minor = _MINIMUM_ENROOT_SQSH_VERSION + existing_sqsh_preflight = ( + f'verify_enroot_compatibility {minimum_major} {minimum_minor} "existing SQSH inspection"\n' + ) + + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +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)} + +compute_file_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$1")" + printf '%s\\n' "${{actual_sha256%% *}}" +}} + +verify_sha256() {{ + [[ "$(compute_file_sha256 "$2")" == "$1" ]] +}} + +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}}" +install -d -m 0700 \ + "${{DD_JOB_DIR}}/home" \ + "${{DD_JOB_DIR}}/enroot/cache" \ + "${{DD_JOB_DIR}}/enroot/config" \ + "${{DD_JOB_DIR}}/enroot/data" \ + "${{DD_JOB_DIR}}/enroot/tmp" \ + "${{DD_INSPECTION_DIRECTORY}}" +export HOME="${{DD_JOB_DIR}}/home" +export ENROOT_CACHE_PATH="${{DD_JOB_DIR}}/enroot/cache" +export ENROOT_CONFIG_PATH="${{DD_JOB_DIR}}/enroot/config" +export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" +export ENROOT_TEMP_PATH="${{DD_JOB_DIR}}/enroot/tmp" +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}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 +readonly DD_CONTAINER_NAME="dd-image-${{SLURM_JOB_ID}}" + +cleanup() {{ + enroot remove -f "${{DD_CONTAINER_NAME}}" >/dev/null 2>&1 || true +}} +trap cleanup EXIT + +enroot create -f --name "${{DD_CONTAINER_NAME}}" "${{DD_IMAGE_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}}" ]] +""" + + +def submit_prepared_image_lifecycle( + prepared: PreparedImageLifecycleJob, + client: SlurmCommandClient, +) -> SlurmJobSubmissionReceipt: + """Verify and submit one prepared image lifecycle script.""" + 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(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") + 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 _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: + registry_or_namespace, separator, remainder = source.partition("/") + if not separator: + return f"docker://docker.io#library/{source}" + if "." in registry_or_namespace or ":" in registry_or_namespace or registry_or_namespace == "localhost": + return f"docker://{registry_or_namespace}#{remainder}" + return f"docker://docker.io#{source}" + + +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(path: Path) -> bytes: + descriptor: int | None = None + try: + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + path_status = os.fstat(descriptor) + if not stat.S_ISREG(path_status.st_mode): + raise ImageLifecycleError("prepared image lifecycle file is not regular") + with os.fdopen(descriptor, "rb") as source: + descriptor = None + return source.read() + except OSError as error: + raise ImageLifecycleError("cannot read prepared image lifecycle file") from error + finally: + if descriptor is not None: + os.close(descriptor) 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 index 4e6589ec1..a44d24bbf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/records.py @@ -1,14 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Immutable records persisted by the Slurm image registry.""" +"""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 ImageInspectionRecord, ImageKind +from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, SelectedSlurmProfile from data_designer.slurm.contracts import ( + ArtifactReference, ContractRecord, ContractValue, Identifier, @@ -18,6 +22,96 @@ ) +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.""" 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/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 index 3c0d3e858..47d061ad1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -76,6 +76,16 @@ def submit(self, script_path: str | Path) -> SlurmJobSubmissionReceipt: output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) return parse_submission(output) + def submit_script(self, script: str) -> 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") + output = self._run( + (self._executables.sbatch, "--parsable", "--export=NIL"), + input_text=script, + ) + return parse_submission(output) + def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: """Return normalized active-queue rows for explicit managed jobs.""" requested = tuple(selectors) @@ -133,10 +143,12 @@ def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, command.append(f"--partition={partition}") return parse_gpu_counts(self._run(command)) - def _run(self, command: Sequence[str]) -> str: + def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: command_name = Path(command[0]).name try: - completed = self._runner.run(command) + completed = ( + self._runner.run(command) if input_text is None else self._runner.run(command, input_text=input_text) + ) except (OSError, subprocess.SubprocessError) as error: raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error returncode = getattr(completed, "returncode", None) 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 index e8941df39..e076f5cb1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -6,33 +6,11 @@ from __future__ import annotations import posixpath -import re -from dataclasses import dataclass +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 -_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_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: """Render a resolved generation plan as one thin deterministic entrypoint.""" @@ -41,8 +19,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") - directives = _build_generation_directives(plan) - directive_text = "\n".join(directive.render() for directive in directives) + directive_text = render_batch_directives(_build_generation_directives(plan)) attempt = f"{attempt_ordinal:04d}" return f"""#!/usr/bin/env bash @@ -50,12 +27,12 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" -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)} +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 @@ -83,7 +60,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi """ -def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDirective, ...]: +def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[tuple[str, str | None], ...]: node_indices = ( plan.client.host_node_index, *(index for deployment in plan.deployments for index in deployment.node_indices), @@ -113,21 +90,4 @@ def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[_BatchDire 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(_BatchDirective(name=name, value=value) for name, value in values if value is not None) - - -def _quote_sbatch_option_value(value: str) -> str: - _reject_control_characters(value, field_name="batch directive value") - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - return f'"{escaped}"' - - -def _quote_shell_value(value: str) -> str: - _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: - if any(ord(character) < 32 or ord(character) == 127 for character in value): - raise SlurmBatchRenderError(f"{field_name} must not contain control characters") + return tuple(values) 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 index 3679d3252..9e2e7839f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -16,7 +16,7 @@ class CommandRunner(Protocol): """Minimal command boundary implemented by production and fake runners.""" - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: """Execute one argument-vector command.""" ... @@ -51,8 +51,20 @@ def environment(self) -> Mapping[str, str]: """Return the allowlisted environment forwarded to child processes.""" return self._environment - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: """Execute an argument vector with captured text output.""" + 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=dict(self._environment), + timeout=self._timeout_seconds, + ) return subprocess.run( tuple(command), check=False, 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 index 1468aac49..4d257a8fb 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -394,7 +394,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -412,7 +412,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -422,7 +425,7 @@ "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json index 8d500c664..7fc47e93d 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -8,7 +8,10 @@ "lab-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 2, + "memory": "8G", + "partition": "cpu", + "time_limit": "02:00:00" }, "scheduler": { "account": "lab", @@ -33,7 +36,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", 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 index b2747d4c0..19b781d06 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -240,7 +240,7 @@ "schema_version": 1, "selected_profile": { "catalog_path": "/workspace/profile.json", - "catalog_sha256": "c747bdf94fe8c638c94f25522e2213d38fa904a6a448e8e1162ef1bb2eef4cc9", + "catalog_sha256": "a707e8e3b6f52170554b743d252b90273fa0deeea8141d4a87abe96e2e93215d", "cluster_name": "primary", "matched_pattern": null, "profile": { @@ -258,7 +258,10 @@ "primary-login-*" ], "image_build": { - "partition": "cpu" + "cpus_per_task": 4, + "memory": "16G", + "partition": "cpu", + "time_limit": "04:00:00" }, "scheduler": { "account": "research", @@ -268,7 +271,7 @@ "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "afaa20b6bcb7233d35b2ad4c9ca82864e7f92ce060442112f6c5e7a5c735cce4", + "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index dcf7addda..c0c67c00a 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -99,6 +99,9 @@ def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfil lambda payload: payload["clusters"]["primary"].update(extra="unknown"), lambda payload: payload["clusters"]["primary"].update(workspace_root="relative"), 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"}, diff --git a/packages/data-designer-slurm/tests/images/test_inspection.py b/packages/data-designer-slurm/tests/images/test_inspection.py index 5d12cad12..c869d6c0c 100644 --- a/packages/data-designer-slurm/tests/images/test_inspection.py +++ b/packages/data-designer-slurm/tests/images/test_inspection.py @@ -5,7 +5,7 @@ import importlib.metadata import re -import shutil +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock @@ -19,6 +19,7 @@ ServingImageInspector, SystemInspectionEnvironment, ) +from data_designer.slurm.images.resources import inspect_image as resource_inspector def _get_client_distributions() -> tuple[InstalledDistribution, ...]: @@ -151,16 +152,15 @@ def test_client_inspector_sorts_distribution_inventory() -> None: ) -def test_client_inspector_uses_installer_version_from_distribution_inventory() -> None: +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"}, ) - inspection = ClientImageInspector(environment).inspect("a" * 64).inspection - - assert inspection.installer_version == "26.1" + with pytest.raises(ImageInspectionError, match="does not match"): + ClientImageInspector(environment).inspect("a" * 64) def test_system_inspection_environment_normalizes_distribution_inventory( @@ -205,25 +205,74 @@ def test_system_inspection_environment_rejects_invalid_distribution_inventory( def test_system_inspection_environment_normalizes_missing_distribution( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr( - importlib.metadata, - "version", - Mock(side_effect=importlib.metadata.PackageNotFoundError("missing")), - ) + monkeypatch.setattr(importlib.metadata, "distributions", Mock(return_value=())) with pytest.raises(ImageInspectionError, match="not installed"): - SystemInspectionEnvironment().get_distribution_version("missing") + SystemInspectionEnvironment().get_distribution_console_script("missing") -@pytest.mark.parametrize("path", (None, "relative/pip"), ids=("missing", "relative")) -def test_system_inspection_environment_requires_absolute_executable( +def test_package_and_standalone_serving_inspectors_share_distribution_owned_console_script_semantics( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, - path: str | None, ) -> None: - monkeypatch.setattr(shutil, "which", Mock(return_value=path)) + 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,))) - with pytest.raises(ImageInspectionError): - SystemInspectionEnvironment().find_executable("pip") + 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: 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..7edd21431 --- /dev/null +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -0,0 +1,663 @@ +# 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 stat +import subprocess +from collections.abc import Sequence +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from pydantic import ValidationError +from slurm_test_fakes import FakeSlurmJob, FakeSlurmRunner + +from data_designer.slurm.config import ( + ImageBuildProfile, + ImageBuildRequest, + ImageInspectionRecord, + SchedulerProfile, + SelectedSlurmProfile, + SlurmProfile, + injected_profile, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.images.errors import ImageLifecycleError +from data_designer.slurm.images.lifecycle import ( + prepare_image_lifecycle_job, + 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.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 + + +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:release@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_JOB_DIR}/home"' in script + assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' in script + assert 'export ENROOT_MAX_PROCESSORS="${SLURM_CPUS_PER_TASK}"' in script + assert script.index('export HOME="${DD_JOB_DIR}/home"') < script.rindex("\nverify_enroot_compatibility 4 0") + assert 'version="$(enroot version)"' in script + assert 'verify_enroot_compatibility 4 0 "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", "4.0.0"), ("existing", "3.5.0")), + ids=("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' + " 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(prepared.plan.job_directory) / "home" + assert set(home_log.read_text().splitlines()) == {expected_home.as_posix()} + assert stat.S_IMODE(expected_home.stat().st_mode) == 0o700 + assert Path(prepared.plan.inspection_output_path).read_text() == "{}\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.5.0", "4.0"), ("existing", "3.4.1", "3.5")), + ids=("digest-import-before-4", "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://docker.io#vllm/vllm-openai@sha256:{'a' * 64}", + ), + ( + f"example@sha256:{'a' * 64}", + f"docker://docker.io#library/example@sha256:{'a' * 64}", + ), + ( + f"registry.example.test:5000/team/image@sha256:{'a' * 64}", + f"docker://registry.example.test:5000#team/image@sha256:{'a' * 64}", + ), + ), + ids=("docker-hub-namespace", "docker-hub-library", "explicit-registry"), +) +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()) + prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + + 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 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=NIL", + ) + ] + 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_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) -> 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-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=NIL") + assert runner.input_text == expected_script + assert Path(prepared.script_file.path).read_text() == "replaced after verification\n" + + +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) + + +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 + + def run( + self, + command: Sequence[str], + *, + input_text: str | None = None, + ) -> subprocess.CompletedProcess[str]: + self.command = tuple(command) + self.input_text = input_text + 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/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 4568b5566..c0e9887b2 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -9,7 +9,7 @@ "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" + "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" }, "run_id": "run-single", "scheduler": { @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" + "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" }, "run_id": "run-single", "schema_version": 1, diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 5155ddb02..183680d34 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -28,6 +28,18 @@ def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSl ] +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=NIL")] + assert runner.inputs == [script] + + def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") @@ -247,6 +259,16 @@ def test_client_rejects_invalid_script_path(fake_slurm_runner: FakeSlurmRunner, 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"): diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index 6bab64a83..b41811178 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -70,6 +70,44 @@ def test_subprocess_runner_default_environment_forwards_only_search_path(monkeyp assert runner.environment == {"LC_ALL": "C", "PATH": "/workspace/slurm/bin:/usr/bin"} +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", "") diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index fc36170e8..3f75a41b6 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="cd95eb995dd9705d12336b9760ca1f354cdd7614dce992cc4f92a6faa47dbfdc" +readonly DD_PLAN_SHA256="902919292da35aca426a191bee775acf140b0f8b7440b55887fb2ba60ed98a52" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index e3f8529d5..553fadd86 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="14a43d767dfab819973d1f048509ac64cb029ef8c30fcffa7aa10555b056c8d6" +readonly DD_PLAN_SHA256="cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/images.py b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py index df0d8d278..7ab4c084b 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/images.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/images.py @@ -39,16 +39,9 @@ def list_distributions(self) -> tuple[InstalledDistribution, ...]: """Return the configured distribution inventory.""" return self.distributions - def get_distribution_version(self, name: str) -> str: - """Return one configured distribution version.""" + 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] + return (self.distribution_versions[name], self.executables[name]) except KeyError: - raise ImageInspectionError(f"required distribution {name!r} is not installed") from None - - def find_executable(self, name: str) -> str: - """Return one configured executable path.""" - try: - return self.executables[name] - except KeyError: - raise ImageInspectionError(f"required executable {name!r} is not installed") from None + 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/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 22da3c51f..5a8494cfc 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -102,13 +102,21 @@ def __init__( 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) -> subprocess.CompletedProcess[str]: + def run( + self, + command: Sequence[str], + *, + check: bool = False, + input_text: 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) command_name = Path(argv[0]).name scripted = self._scripted_responses.get(command_name) if scripted: diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 9880d08e8..621d4a33e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="65032718bfc8fbd6c60700add08ebc098ba2e64858c23808decbefe5a8153d91", + expected_fixture_sha256="8238dd22393a46538169a83046e91230166aaa0f96e0a91255adf6a6074bfa20", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="58c3a2eddc9a0d377eb0a90655e70a6b2484cade0c52b2a01b9b388dcd4667e2", + expected_fixture_sha256="6153c1769d9b90315c4e7acc767bd5c469d0a5e1f5fc7843baddc1a733cb8e9b", ) From bc700fc9819407655e72f0ff36e968384608f742 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 1 Sep 2026 21:52:45 -0600 Subject: [PATCH 14/62] fix: harden Slurm image publication (#907) * fix: harden Slurm image publication Signed-off-by: Nabin Mulepati * fix: close Slurm image publication races - Hash SQSH bytes outside the global registry lock.\n- Preserve recovery state when commit synchronization fails.\n- Bind lifecycle artifacts to verified directory descriptors. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/images/filesystem.py | 18 +- .../data_designer/slurm/images/lifecycle.py | 343 ++++++- .../data_designer/slurm/images/registry.py | 267 ++++- .../src/data_designer/slurm/images/service.py | 438 +++++++- .../data_designer/slurm/serving/resolver.py | 26 +- .../tests/images/test_lifecycle.py | 966 +++++++++++++++++- .../tests/images/test_registry_store.py | 281 ++++- .../tests/serving/test_resolver.py | 21 +- 8 files changed, 2229 insertions(+), 131 deletions(-) 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 index 36441227f..b023c4ca6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py @@ -148,7 +148,7 @@ def read_regular_text(directory_descriptor: int, name: str, display_path: Path) raise OSError(f"registry path {display_path} is not a regular file") descriptor = os.open( name, - os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0), dir_fd=directory_descriptor, ) after_open = os.fstat(descriptor) @@ -157,7 +157,21 @@ def read_regular_text(directory_descriptor: int, name: str, display_path: Path) registry_file = os.fdopen(descriptor, "r", encoding="utf-8") descriptor = None with registry_file: - return registry_file.read() + 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) + + +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) 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 index cfb317f7b..ddfa231c1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -8,22 +8,30 @@ import hashlib import os import stat -from dataclasses import dataclass +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field from importlib import resources from pathlib import Path from pydantic import TypeAdapter, ValidationError -from data_designer.slurm.config import ImageBuildRequest, SelectedSlurmProfile +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 +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 @@ -36,15 +44,23 @@ _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) _MINIMUM_ENROOT_OCI_VERSION = (4, 0) _MINIMUM_ENROOT_SQSH_VERSION = (3, 5) +_MAXIMUM_INSPECTION_SIZE = 1024 * 1024 @dataclass(frozen=True, slots=True) class PreparedImageLifecycleJob: - """Persisted, checksum-bound files ready for one Slurm submission.""" + """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( @@ -73,11 +89,16 @@ def prepare_image_lifecycle_job( 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) @@ -105,8 +126,16 @@ def prepare_image_lifecycle_job( 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 - return PreparedImageLifecycleJob(plan=plan, plan_file=plan_file, script_file=script_file) + 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: @@ -143,6 +172,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: exit 73 fi verify_enroot_compatibility {minimum_major} {minimum_minor} "digest-pinned OCI imports" +DD_REMOVE_SQSH_ON_FAILURE=1 enroot import -o "${{DD_IMAGE_SQSH}}" "${{DD_OCI_SOURCE}}" """ else: @@ -150,6 +180,28 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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} @@ -165,6 +217,29 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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 + return "${{status}}" +}} +trap cleanup EXIT compute_file_sha256() {{ local actual_sha256 @@ -227,14 +302,25 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: printf '%s\\n' 'SLURM_JOB_ID must be a positive integer' >&2 exit 64 fi -readonly DD_CONTAINER_NAME="dd-image-${{SLURM_JOB_ID}}" - -cleanup() {{ - enroot remove -f "${{DD_CONTAINER_NAME}}" >/dev/null 2>&1 || true -}} -trap cleanup EXIT +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} -enroot create -f --name "${{DD_CONTAINER_NAME}}" "${{DD_IMAGE_SQSH}}" +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" \ @@ -253,6 +339,26 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: /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 +true """ @@ -261,6 +367,99 @@ def submit_prepared_image_lifecycle( 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 _verify_prepared_lifecycle( + prepared: PreparedImageLifecycleJob, + job_directory_descriptor: int, +) -> bytes: expected_artifacts = ( ( "plan", @@ -278,7 +477,11 @@ def submit_prepared_image_lifecycle( 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(expected_path) + 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 @@ -289,11 +492,28 @@ def submit_prepared_image_lifecycle( verified_script = actual_content if verified_script is None: raise ImageLifecycleError("prepared image lifecycle script was not verified") - 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) + 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: @@ -330,18 +550,95 @@ def _write_file(path: Path, content: bytes, *, mode: int) -> ArtifactReference: return ArtifactReference(path=path.as_posix(), sha256=hashlib.sha256(content).hexdigest()) -def _read_regular_file(path: Path) -> bytes: +def _read_regular_file( + directory_descriptor: int, + name: str, + display_path: Path, + *, + maximum_size: int | None = None, +) -> bytes: descriptor: int | None = None try: - descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + 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): + 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 - return source.read() + 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("cannot read prepared image lifecycle file") from 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/registry.py b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py index dd3e3ba38..be75f72fe 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py @@ -5,8 +5,10 @@ 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 @@ -37,11 +39,17 @@ _LOCK_DIRECTORY_NAME = ".locks" _REGISTRY_FILENAME = "registry.yaml" _REGISTRY_LOCK_FILENAME = "registry.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.""" + """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: @@ -91,11 +99,33 @@ def register( 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.""" + """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( @@ -105,25 +135,47 @@ def register( ), acquire_file_lock( lock_directory_descriptor, - _REGISTRY_LOCK_FILENAME, - self._lock_directory / _REGISTRY_LOCK_FILENAME, + target_lock_name, + self._lock_directory / target_lock_name, ), ): - snapshot = self._load_from_directory(image_root_descriptor) - 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") - self._validate_path_facts(snapshot, image, replacing=existing) - images = tuple(entry for entry in snapshot.images if entry.name != image.name) + (image,) - updated = ImageRegistryDocument( - schema_version=1, - images=tuple(sorted(images, key=lambda entry: entry.name)), - ) - self._save( - image_root_descriptor, - updated, - verify_before_publish=lambda: verify_before_publish(image), - ) + 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) -> RegisteredImage: @@ -144,6 +196,7 @@ def unregister(self, name: Identifier) -> RegisteredImage: 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) @@ -153,17 +206,31 @@ def unregister(self, name: Identifier) -> RegisteredImage: schema_version=1, images=tuple(image for image in snapshot.images if image.name != name), ) - self._save(image_root_descriptor, updated) + self._save(image_root_descriptor, snapshot, updated) return removed def _load(self) -> ImageRegistryDocument: try: - with open_verified_directory(self._image_root) as image_root_descriptor: - return self._load_from_directory(image_root_descriptor) + 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: @@ -180,16 +247,97 @@ def _load_from_directory(self, image_root_descriptor: int) -> ImageRegistryDocum 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=".registry.", + prefix=prefix, suffix=".tmp", ) output = os.fdopen(descriptor, "w", encoding="utf-8") @@ -203,18 +351,8 @@ def _save( ) output.flush() os.fsync(output.fileno()) - 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, - ) - os.fsync(image_root_descriptor) - except (OSError, TypeError, yaml.YAMLError) as error: - raise ImageRegistryError(f"cannot persist image registry {self._registry_path}") from error - finally: + return temporary_name + except BaseException: if descriptor is not None: os.close(descriptor) if temporary_name is not None: @@ -222,6 +360,29 @@ def _save( 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: @@ -243,6 +404,23 @@ def _open_storage(self) -> Iterator[tuple[int, int]]: 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, @@ -259,6 +437,10 @@ def _validate_path_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.""" @@ -300,3 +482,18 @@ def _validate_image_alias(name: str) -> Identifier: 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/service.py b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py index d04819708..607f7f2e0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -8,17 +8,23 @@ 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 data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, ImageRef from data_designer.slurm.contracts import Sha256Digest -from data_designer.slurm.images.errors import ImageVerificationError +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.planning import ResolvedImage _HASH_BLOCK_SIZE = 1024 * 1024 +_ARTIFACT_DIRECTORY_NAME = "artifacts" +_TEMPORARY_DIRECTORY_NAME = ".tmp" class VerifiedImageRegistry: @@ -38,15 +44,144 @@ def register_existing( 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)) + 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=path.as_posix(), + path=artifact_path.as_posix(), sqsh_sha256=inspection.sqsh_sha256, + source_oci_digest=source_oci_digest, inspection=inspection, ) - return self._registry.register(image, verify_before_publish=_verify_registered_image, replace=replace) + 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.""" @@ -79,56 +214,8 @@ def unregister(self, name: str) -> RegisteredImage: def compute_sqsh_file_sha256(path: Path) -> Sha256Digest: """Compute the SHA-256 of one regular, non-symlink SQSH file.""" - descriptor: int | None = None - try: - before_open = path.lstat() - if not stat.S_ISREG(before_open.st_mode): - raise ImageVerificationError(f"image path {path} must be a regular non-symlink file") - descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) - image_file = os.fdopen(descriptor, "rb") - descriptor = None - digest = hashlib.sha256() - with image_file: - before_read = os.fstat(image_file.fileno()) - if (before_open.st_dev, before_open.st_ino) != (before_read.st_dev, before_read.st_ino): - raise ImageVerificationError(f"image path {path} changed while it was being opened") - while block := image_file.read(_HASH_BLOCK_SIZE): - digest.update(block) - after_read = os.fstat(image_file.fileno()) - after_path = path.lstat() - before_facts = ( - before_read.st_dev, - before_read.st_ino, - before_read.st_size, - before_read.st_mtime_ns, - before_read.st_ctime_ns, - ) - after_facts = ( - after_read.st_dev, - after_read.st_ino, - after_read.st_size, - after_read.st_mtime_ns, - after_read.st_ctime_ns, - ) - if before_facts != after_facts: - raise ImageVerificationError(f"image path {path} changed while it was being verified") - path_facts = ( - after_path.st_dev, - after_path.st_ino, - after_path.st_size, - after_path.st_mtime_ns, - after_path.st_ctime_ns, - ) - if not stat.S_ISREG(after_path.st_mode) or path_facts != after_facts: - raise ImageVerificationError(f"image path {path} changed while it was being verified") - return digest.hexdigest() - except FileNotFoundError as error: - raise ImageVerificationError(f"image path {path} does not exist") from error - except OSError as error: - raise ImageVerificationError(f"cannot read image path {path}") from error - finally: - if descriptor is not None: - os.close(descriptor) + with _open_verified_sqsh(path) as snapshot: + return snapshot.sha256 def _validate_inspection( @@ -148,6 +235,247 @@ def _validate_inspection( def _verify_registered_image(image: RegisteredImage) -> None: - actual_sha256 = compute_sqsh_file_sha256(Path(image.path)) - if actual_sha256 != image.sqsh_sha256: + 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 (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/serving/resolver.py b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py index ac2bdce27..dd85211e7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py @@ -5,6 +5,8 @@ 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 @@ -25,6 +27,11 @@ ) 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.""" @@ -34,12 +41,7 @@ def resolve_vllm_server( plan: ResolvedSlurmRunPlan, deployment_id: Identifier, ) -> ResolvedVllmServerDeployment: - """Resolve one vLLM deployment and endpoint derived from the same run plan. - - Inspected runtime versions remain provenance. This resolver emits only the - baseline behavior required of every admitted vLLM image and therefore does not - maintain a package-version compatibility matrix. - """ + """Resolve one vLLM deployment admitted by the tested runtime matrix.""" try: return _resolve_vllm_server(plan, deployment_id) except ValidationError as error: @@ -72,6 +74,7 @@ def _resolve_vllm_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") @@ -208,3 +211,14 @@ def _calculate_launch_delay_seconds(server: VllmServerConfig, deployment_replica 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/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index 7edd21431..f63b0268f 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -5,11 +5,15 @@ import hashlib import json +import os +import shutil import stat import subprocess from collections.abc import 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 @@ -17,24 +21,39 @@ 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 ImageLifecycleError +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 @@ -90,6 +109,7 @@ def test_prepare_existing_sqsh_inspects_in_place_without_oci_identity(tmp_path: 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: @@ -320,12 +340,18 @@ def test_prepare_refuses_duplicate_and_invalid_lifecycle_ids(tmp_path: Path) -> workspace = tmp_path / "workspace" profile = _get_selected_profile(workspace) request = ImageBuildRequest(name="client", kind="client", source=(tmp_path / "client.sqsh").as_posix()) - prepare_image_lifecycle_job(request, profile, lifecycle_id="image-job-0004") + 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() @@ -485,6 +511,24 @@ def test_submit_rejects_modified_prepared_files(tmp_path: Path, artifact_name: s 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()), @@ -524,6 +568,858 @@ def test_submit_uses_verified_script_bytes_when_path_changes_during_submission(t 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_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, @@ -644,6 +1540,72 @@ def _get_selected_profile(workspace: Path) -> SelectedSlurmProfile: 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 = 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 diff --git a/packages/data-designer-slurm/tests/images/test_registry_store.py b/packages/data-designer-slurm/tests/images/test_registry_store.py index 55f92b6c2..8c717b352 100644 --- a/packages/data-designer-slurm/tests/images/test_registry_store.py +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -3,18 +3,23 @@ from __future__ import annotations +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.errors import ImageConflictError, ImageRegistryError +from data_designer.slurm.images.errors import 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 @@ -180,6 +185,96 @@ def test_registry_rejects_nonregular_state_file_without_blocking(tmp_path: Path) 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" @@ -243,9 +338,11 @@ def test_registry_uses_restrictive_atomic_workspace_storage(tmp_path: Path) -> N 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_remains_unpublished_until_final_verification(tmp_path: Path) -> None: +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) @@ -256,20 +353,172 @@ def verify_before_publish(_image: RegisteredImage) -> None: verification_started.set() assert finish_verification.wait(timeout=5) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit( + 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) - assert ImageRegistryStore(workspace).list_images() == () + reader = executor.submit(ImageRegistryStore(workspace).list_images) + assert not reader.done() finish_verification.set() - assert future.result().name == "client" + 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" @@ -331,6 +580,26 @@ def _write_registry_content(workspace: Path, content: bytes) -> None: 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) diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index fdd83ee99..4c7766487 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -112,13 +112,30 @@ def test_two_deployments_keep_images_and_endpoint_identities_isolated( 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"] = "vendor-vllm-build" + 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 == "vendor-vllm-build" + 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) -> None: From 4ad048b61639a6469d67d0cbb0ee64447c359c74 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 09:23:52 -0600 Subject: [PATCH 15/62] feat: add Slurm state writer (#908) * feat: add Slurm state writer Persist immutable run intent and shard manifests with restrictive descriptor-bound storage, locks, digest validation, and atomic publication. Validate monotonic attempt and readiness updates across concurrent writers, and cover crash recovery, unsafe filesystem state, tampering, and fresh-process reloads. Signed-off-by: Nabin Mulepati * fix: harden Slurm state persistence Bind loaded records back to persisted plan and location invariants, harden lock and file handling against replacement and hardlink races, and avoid duplicate immutable context reads. Add adversarial, concurrency, interruption-recovery, import, and performance regression coverage. Signed-off-by: Nabin Mulepati * fix: recover interrupted state publication Repair a committed immutable record when an interruption leaves its package-owned temporary hard link behind. Add regression coverage for the exact post-publication interruption window. Signed-off-by: Nabin Mulepati * fix: converge concurrent state recovery Treat a temporary link already removed by a recovery reader as a successful immutable publication. Cover the publisher-reader cleanup race directly. Signed-off-by: Nabin Mulepati * test: cover Slurm state reader errors Exercise invalid and missing shard, attempt, ordinal, and readiness inputs through the public state-writer API. Signed-off-by: Nabin Mulepati * fix: retry concurrent state replacement Let descriptor-bound readers retry valid atomic replacements, including replacements that race interrupted-publication cleanup. Continue rejecting unsafe record types and bound the retry loop. Signed-off-by: Nabin Mulepati * refactor: compose Slurm state persistence Separate validated reads and descriptor-bound storage from the public mutation workflows. Close concurrent attempt publication, retry durability, exact-byte reading, and interrupted recovery races with focused regression coverage. Signed-off-by: Nabin Mulepati * refactor: shorten state mutation workflows Keep error translation at public boundaries and move lock-held happy paths into small helpers. No state method now exceeds 30 lines. Signed-off-by: Nabin Mulepati * refactor: expose state collaborators within package Signed-off-by: Nabin Mulepati * refactor: harden shard-local state boundaries Keep attempt and readiness mutations scoped to their target shard, move persisted plan validation into the state layer, and share descriptor-safe filesystem primitives across state and image storage. Make the schema-v1 no-migration policy explicit and freeze compatibility fixtures. Signed-off-by: Nabin Mulepati * fix: align state boundary invariants Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/filesystem.py | 203 +++ .../data_designer/slurm/images/filesystem.py | 129 +- .../src/data_designer/slurm/integration.py | 141 +- .../src/data_designer/slurm/state/__init__.py | 38 + .../src/data_designer/slurm/state/errors.py | 22 + .../data_designer/slurm/state/filesystem.py | 301 ++++ .../slurm/state/plan_validation.py | 169 +++ .../src/data_designer/slurm/state/reader.py | 218 +++ .../src/data_designer/slurm/state/storage.py | 415 ++++++ .../src/data_designer/slurm/state/store.py | 300 ++++ .../data_designer/slurm/state/validation.py | 28 +- .../tests/images/test_registry.py | 4 +- .../tests/images/test_registry_store.py | 14 + .../tests/state/test_state_golden_records.py | 18 + .../tests/state/test_store.py | 1292 +++++++++++++++++ 15 files changed, 3060 insertions(+), 232 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/reader.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/storage.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/store.py create mode 100644 packages/data-designer-slurm/tests/state/test_store.py 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..99d8e7f11 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py @@ -0,0 +1,203 @@ +# 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, +) -> 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_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/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py index b023c4ca6..5baa2c690 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py @@ -5,86 +5,57 @@ from __future__ import annotations -import fcntl import os -import secrets import stat from collections.abc import Iterator -from contextlib import contextmanager +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 -_DIRECTORY_MODE = 0o700 -_FILE_MODE = 0o600 - @contextmanager def acquire_file_lock(directory_descriptor: int, name: str, display_path: Path) -> Iterator[None]: """Acquire one exclusive advisory lock without reopening its parent path.""" - descriptor: int | None = None - try: - # macOS can report ENOENT to the losing openat(O_CREAT) caller when two - # processes create the same lock file concurrently. Retrying then opens - # the winner's regular file with the same no-follow boundary. - for _ in range(10): - try: - descriptor = os.open( + with ExitStack() as resources: + try: + resources.enter_context( + acquire_restrictive_file_lock( + directory_descriptor, name, - os.O_CREAT | os.O_RDWR | getattr(os, "O_NOFOLLOW", 0), - _FILE_MODE, - dir_fd=directory_descriptor, + display_path, + resource_name="registry", ) - except FileNotFoundError: - continue - break - if descriptor is None: - raise OSError(f"lock path {display_path} could not be opened") - if not stat.S_ISREG(os.fstat(descriptor).st_mode): - raise OSError(f"lock path {display_path} is not a regular file") - os.fchmod(descriptor, _FILE_MODE) - fcntl.flock(descriptor, fcntl.LOCK_EX) - except OSError as error: - if descriptor is not None: - os.close(descriptor) - raise ImageRegistryError(f"cannot lock image registry target {display_path}") from error - - try: + ) + except OSError as error: + raise ImageRegistryError(f"cannot lock image registry target {display_path}") from error yield - finally: - assert descriptor is not None - try: - fcntl.flock(descriptor, fcntl.LOCK_UN) - finally: - os.close(descriptor) def ensure_private_directory(path: Path, *, parents: bool) -> None: """Create or restrict one real directory to its owning user.""" - path.mkdir(mode=_DIRECTORY_MODE, parents=parents, exist_ok=True) + path.mkdir(mode=PRIVATE_DIRECTORY_MODE, parents=parents, exist_ok=True) with open_verified_directory(path) as descriptor: - os.fchmod(descriptor, _DIRECTORY_MODE) + 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.""" - descriptor: int | None = None - try: - before_open = path.lstat() - if not stat.S_ISDIR(before_open.st_mode): - raise OSError(f"registry 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) - if (before_open.st_dev, before_open.st_ino) != (after_open.st_dev, after_open.st_ino): - raise OSError(f"registry directory {path} changed while it was being opened") - if not stat.S_ISDIR(after_open.st_mode): - raise OSError(f"registry directory {path} is not a directory") + with open_shared_directory(path, resource_name="registry") as descriptor: yield descriptor - finally: - if descriptor is not None: - os.close(descriptor) @contextmanager @@ -94,22 +65,13 @@ def open_verified_child_directory( display_path: Path, ) -> Iterator[int]: """Open one real child directory relative to an already 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"registry 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) - if (before_open.st_dev, before_open.st_ino) != (after_open.st_dev, after_open.st_ino): - raise OSError(f"registry directory {display_path} changed while it was being opened") - if not stat.S_ISDIR(after_open.st_mode): - raise OSError(f"registry directory {display_path} is not a directory") + with open_shared_child_directory( + parent_descriptor, + name, + display_path, + resource_name="registry", + ) as descriptor: yield descriptor - finally: - if descriptor is not None: - os.close(descriptor) def create_temporary_file( @@ -119,24 +81,7 @@ def create_temporary_file( suffix: str, ) -> tuple[int, str]: """Create one restrictive 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), - _FILE_MODE, - dir_fd=directory_descriptor, - ) - except FileExistsError: - continue - try: - os.fchmod(descriptor, _FILE_MODE) - except OSError: - os.close(descriptor) - raise - return descriptor, name - raise OSError("cannot allocate a unique temporary file name") + return create_restrictive_temporary_file(directory_descriptor, prefix=prefix, suffix=suffix) def read_regular_text(directory_descriptor: int, name: str, display_path: Path) -> str: @@ -163,15 +108,11 @@ def read_regular_text(directory_descriptor: int, name: str, display_path: Path) 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) + 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) - - -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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index 5f968fe73..343ea2393 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -6,139 +6,26 @@ from __future__ import annotations import posixpath -from collections.abc import Mapping -from types import MappingProxyType from data_designer.slurm.client import ClientOutcome, ClientResult -from data_designer.slurm.contracts import ArtifactReference, ShardId -from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan -from data_designer.slurm.state import ( +from data_designer.slurm.planning import PlannedShard +from data_designer.slurm.state.execution import ( AttemptLifecycleState, AttemptManifest, - AttemptReadiness, AttemptTerminalClassification, - CandidateOutputManifest, - ReadinessState, - RunManifest, - ShardManifest, - ShardWinner, - StateContractError, - validate_shard_set, ) +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner +from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError +IntegrationContractError = PlanStateContractError -class IntegrationContractError(ValueError): - """Raised when reviewed plan and state records disagree.""" - -class PlanStateValidator: +class PlanStateValidator(PersistedPlanStateValidator): """Validate records against one resolved plan with reusable derived state. This is an in-process validation service, not a persisted contract record. """ - 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} - ) - - @property - def plan(self) -> ResolvedSlurmRunPlan: - """The resolved plan used to build this validation context.""" - 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.""" - _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") - try: - validate_shard_set(run, shards) - except StateContractError as exc: - raise IntegrationContractError(str(exc)) from exc - - _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(persisted.shard_id == planned.shard_id, "state shard identity does not match planned order") - _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 shards - - def validate_initial_readiness( - self, - attempt: AttemptManifest, - readiness: AttemptReadiness, - ) -> AttemptReadiness: - """Anchor the first readiness snapshot to a fully validated planned attempt.""" - 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.revision == 1, "initial readiness must use revision 1") - _require(readiness.state is ReadinessState.PENDING, "initial readiness must be pending") - _require(readiness.updated_at >= attempt.created_at, "initial 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, "initial 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") - expected_attempt_id = f"attempt-{attempt.attempt_ordinal:04d}" - _require(attempt.attempt_id == expected_attempt_id, "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, @@ -237,22 +124,6 @@ def validate_finalization_chain( _require(winner.published_at >= attempt.updated_at, "winner publication cannot precede attempt completion") return winner - def _get_planned_shard(self, shard_id: ShardId) -> PlannedShard: - planned_shard = self._shards_by_id.get(shard_id) - if planned_shard is None: - raise IntegrationContractError("attempt shard_id does not identify a planned shard") - return planned_shard - - 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", - ) - @staticmethod def _get_expected_dataset_path( planned_shard: PlannedShard, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index 07885b69c..b4717289e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -5,6 +5,9 @@ from __future__ import annotations +import importlib +from typing import TYPE_CHECKING + from data_designer.slurm.contracts import ( ArtifactReference, AttemptId, @@ -21,6 +24,12 @@ StateRecord, StateValue, ) +from data_designer.slurm.state.errors import ( + SlurmStateError, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) from data_designer.slurm.state.execution import ( AttemptLifecycleState, AttemptManifest, @@ -61,11 +70,19 @@ 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.store import SlurmStateWriter # noqa: F401 + +_LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "SlurmStateWriter": ("data_designer.slurm.state.store", "SlurmStateWriter"), +} + __all__ = [ "ArtifactReference", "AttemptLifecycleState", @@ -98,7 +115,12 @@ "ShardManifest", "ShardId", "ShardWinner", + "SlurmStateError", + "SlurmStateWriter", + "StateConflictError", "StateContractError", + "StateCorruptionError", + "StateNotFoundError", "StateRecord", "StateValue", "reconcile_attempt_observation", @@ -108,7 +130,23 @@ "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/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py new file mode 100644 index 000000000..e43897dcc --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py @@ -0,0 +1,22 @@ +# 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.""" 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..b771da750 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py @@ -0,0 +1,301 @@ +# 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 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") + + +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 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..6b40340f0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/plan_validation.py @@ -0,0 +1,169 @@ +# 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.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.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} + ) + + @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 _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..57f0fe4a5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -0,0 +1,218 @@ +# 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.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.plan_validation import PersistedPlanStateValidator, PlanStateContractError +from data_designer.slurm.state.readiness import AttemptReadiness +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.authored_config_path.as_posix() + or run.resolved_plan.path != self._storage.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.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.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) + try: + readiness = self._storage.read_readiness(shard_id, attempt_id) + PersistedPlanStateValidator(plan).validate_readiness_snapshot(attempt, readiness) + return readiness + except FileNotFoundError as error: + raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") from error + except StateCorruptionError: + raise + except PlanStateContractError as error: + raise StateCorruptionError(f"attempt {attempt_id!r} has invalid persisted readiness") from error + except OSError as error: + raise StateCorruptionError(f"attempt {attempt_id!r} has unreadable readiness") from error + + 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/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py new file mode 100644 index 000000000..94698bf63 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -0,0 +1,415 @@ +# 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 contextmanager +from pathlib import Path +from typing import TypeVar + +from pydantic import ValidationError + +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import AttemptId, ContractRecord, 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 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.readiness import AttemptReadiness + +_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" +_LOCK_DIRECTORY_NAME = ".locks" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_ATTEMPT_NAME_PATTERN = re.compile(r"^attempt-[0-9]{4,}$") +_RecordT = TypeVar("_RecordT", bound=ContractRecord) + + +class StateStorage: + """Own descriptor-bound paths, locking, and record serialization.""" + + def __init__(self, workspace_root: Path, run_id: Identifier) -> None: + self.workspace_root = 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 + + @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 + + 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 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_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 + + 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 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 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, (AttemptManifest, AttemptReadiness)): + 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..9f106f080 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -0,0 +1,300 @@ +# 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 pathlib import Path + +from pydantic import TypeAdapter, ValidationError + +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 AttemptManifest, RunManifest, ShardManifest +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.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: Selected compute-visible workspace root. + run_id: Stable application-owned run identity. + """ + + def __init__(self, workspace_root: str | Path, run_id: Identifier) -> None: + 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 + self._storage = StateStorage(Path(normalized_root), normalized_run_id) + self._reader = StateReader(self._storage, normalized_run_id) + 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_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 _create_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) + 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 + 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) + 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.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.authored_config_path.as_posix(): + raise StateContractError("run authored config reference does not match its persisted location") + if ( + run.resolved_plan.path != self._storage.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/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py index 30fdbc066..f990b65e8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -95,10 +95,15 @@ def validate_attempt_set( 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}") - validate_attempt_manifest(run, shard, attempt) 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) @@ -115,6 +120,27 @@ def validate_attempt_set( 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, diff --git a/packages/data-designer-slurm/tests/images/test_registry.py b/packages/data-designer-slurm/tests/images/test_registry.py index 662466ebf..9fa43ea38 100644 --- a/packages/data-designer-slurm/tests/images/test_registry.py +++ b/packages/data-designer-slurm/tests/images/test_registry.py @@ -131,7 +131,7 @@ def mutate_replacement_before_registry_write(nbytes: int | None = None) -> str: with ( patch( - "data_designer.slurm.images.filesystem.secrets.token_hex", + "data_designer.slurm.filesystem.secrets.token_hex", side_effect=mutate_replacement_before_registry_write, ), pytest.raises(ImageVerificationError, match="no longer matches"), @@ -162,7 +162,7 @@ def mutate_image_before_registry_write(nbytes: int | None = None) -> str: with ( patch( - "data_designer.slurm.images.filesystem.secrets.token_hex", + "data_designer.slurm.filesystem.secrets.token_hex", side_effect=mutate_image_before_registry_write, ), pytest.raises(ImageVerificationError, match="no longer matches"), diff --git a/packages/data-designer-slurm/tests/images/test_registry_store.py b/packages/data-designer-slurm/tests/images/test_registry_store.py index 8c717b352..db1aabb31 100644 --- a/packages/data-designer-slurm/tests/images/test_registry_store.py +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -19,6 +19,7 @@ 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 ImageConflictError, ImageRegistryError, ImageVerificationError from data_designer.slurm.images.inspection import ClientImageInspector, ServingImageInspector from data_designer.slurm.images.records import RegisteredImage @@ -302,6 +303,19 @@ def test_registry_rejects_nonregular_lock_file(tmp_path: Path) -> None: 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" diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index d7b5d1adf..6599cf992 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -21,6 +21,7 @@ ) 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), @@ -39,6 +40,11 @@ ("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 = "1a51941bbb2d8aac7114c197f51105104297de1c6a367cfab0f4111b5d7eebaa" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) @@ -66,3 +72,15 @@ def test_golden_artifact_references_hash_exact_persisted_bytes() -> 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..fcd1c84ac --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -0,0 +1,1292 @@ +# 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 +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from threading import Barrier +from typing import TypeVar, cast + +import pytest + +from data_designer.slurm import filesystem as slurm_filesystem +from data_designer.slurm.config import 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.state import ( + AttemptId, + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + DeploymentReadiness, + EndpointPublicationState, + ReadinessState, + RunManifest, + SchedulerIdentity, + ShardId, + ShardManifest, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state import filesystem as state_filesystem +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 + + +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) + 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_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 _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__}") From 8fbca0c53e8016324314651f9d9c8f9df15790ca Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 09:32:29 -0600 Subject: [PATCH 16/62] feat: add one-node Slurm runtime (#909) --- .../src/data_designer/slurm/integration.py | 98 ++- .../data_designer/slurm/runtime/__init__.py | 46 ++ .../src/data_designer/slurm/runtime/bundle.py | 326 +++++++++ .../data_designer/slurm/runtime/controller.py | 608 +++++++++++++++++ .../data_designer/slurm/runtime/entrypoint.py | 122 ++++ .../src/data_designer/slurm/runtime/errors.py | 34 + .../src/data_designer/slurm/runtime/logs.py | 72 ++ .../src/data_designer/slurm/runtime/models.py | 149 ++++ .../src/data_designer/slurm/runtime/paths.py | 47 ++ .../data_designer/slurm/runtime/preflight.py | 190 ++++++ .../src/data_designer/slurm/runtime/probes.py | 38 ++ .../src/data_designer/slurm/runtime/proxy.py | 246 +++++++ .../data_designer/slurm/runtime/records.py | 71 ++ .../data_designer/slurm/runtime/signals.py | 96 +++ .../src/data_designer/slurm/runtime/steps.py | 458 +++++++++++++ .../data_designer/slurm/runtime/supervisor.py | 276 ++++++++ .../data_designer/slurm/state/readiness.py | 139 ++-- .../slurm/state/reconciliation.py | 61 +- .../test_integration_validation.py | 25 +- .../tests/runtime/conftest.py | 167 +++++ .../tests/runtime/test_bundle.py | 238 +++++++ .../tests/runtime/test_controller.py | 641 ++++++++++++++++++ .../tests/runtime/test_entrypoint.py | 15 + .../tests/runtime/test_paths.py | 47 ++ .../tests/runtime/test_preflight.py | 138 ++++ .../tests/runtime/test_probes.py | 36 + .../tests/runtime/test_proxy.py | 199 ++++++ .../tests/runtime/test_steps.py | 221 ++++++ .../tests/runtime/test_supervisor.py | 299 ++++++++ .../tests/state/test_validation.py | 29 +- 30 files changed, 5032 insertions(+), 100 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/__init__.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/logs.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/probes.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/signals.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/supervisor.py create mode 100644 packages/data-designer-slurm/tests/runtime/conftest.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_bundle.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_controller.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_entrypoint.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_paths.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_preflight.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_probes.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_proxy.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_steps.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_supervisor.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index 343ea2393..cde400e29 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -8,6 +8,7 @@ import posixpath from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import PlannedShard from data_designer.slurm.state.execution import ( AttemptLifecycleState, @@ -35,38 +36,69 @@ def validate_finalization_chain( winner: ShardWinner, ) -> ShardWinner: """Validate a complete semantic result through immutable winner publication.""" - self.validate_planned_attempt(planned_shard, attempt) + self.validate_client_candidate(planned_shard, attempt, client_result, candidate) _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", ) + self._validate_winner(planned_shard, attempt, client_result, winner) + _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") + return winner + + def validate_client_candidate( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> CandidateOutputManifest: + """Validate a complete client result and candidate before terminal publication.""" + self.validate_planned_attempt(planned_shard, attempt) _require(client_result.outcome is ClientOutcome.COMPLETE, "only complete client results may be finalized") _require(candidate.winner_eligible, "only complete candidate outputs may be finalized") + self._validate_candidate_identities(planned_shard, attempt, client_result, candidate) + self._validate_candidate_counts(planned_shard, client_result, candidate) + self._validate_candidate_location(planned_shard, attempt, client_result, candidate) + _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", + ) + return candidate - for record_name, run_id in ( - ("client result", client_result.run_id), - ("candidate", candidate.run_id), - ("winner", winner.run_id), - ): + def _validate_candidate_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), - ("winner", winner.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), - ("winner", winner.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 == winner.attempt_ordinal, - "candidate and winner attempt ordinals must match the attempt", + candidate.attempt_ordinal == attempt.attempt_ordinal, "candidate attempt ordinal does not match attempt" ) + + @staticmethod + def _validate_candidate_counts( + 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", @@ -75,11 +107,18 @@ def validate_finalization_chain( client_result.actual_records == candidate.actual_records == planned_shard.requested_records, "client and candidate actual records must complete the planned shard", ) + + def _validate_candidate_location( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> None: _require( client_result.requested_resume_mode == self.plan.invocation.authored.resume, "client requested resume mode does not match the resolved plan", ) - expected_dataset_path = self._get_expected_dataset_path(planned_shard, attempt, client_result) _require( client_result.dataset_path == expected_dataset_path, "client dataset path does not match planned intent" @@ -87,23 +126,33 @@ def validate_finalization_chain( _require( candidate.dataset_path == expected_dataset_path, "candidate dataset path does not match planned intent" ) - + candidate_reference = self._get_candidate_reference(client_result) expected_manifest_path = posixpath.join( posixpath.dirname(planned_shard.resume_workspace.path), "attempts", attempt.attempt_id, "output-manifest.json", ) - candidate_reference = client_result.candidate_output_manifest - _require(candidate_reference is not None, "complete client result has no candidate manifest reference") _require( - candidate_reference.path == expected_manifest_path, - "candidate manifest path does not match planned intent", + candidate_reference.path == expected_manifest_path, "candidate manifest path does not match planned intent" ) _require( candidate_reference.sha256 == candidate.compute_sha256(), "client candidate digest does not match the manifest", ) + + def _validate_winner( + self, + planned_shard: PlannedShard, + attempt: AttemptManifest, + client_result: ClientResult, + winner: ShardWinner, + ) -> None: + candidate_reference = self._get_candidate_reference(client_result) + _require(winner.run_id == self.plan.run_id, "winner run_id does not match the resolved plan") + _require(winner.shard_id == planned_shard.shard_id, "winner shard_id does not match the planned shard") + _require(winner.attempt_id == attempt.attempt_id, "winner attempt_id does not match the attempt") + _require(winner.attempt_ordinal == attempt.attempt_ordinal, "winner attempt ordinal does not match the attempt") _require( attempt.candidate_output == candidate_reference, "attempt candidate reference does not match client result", @@ -113,16 +162,11 @@ def validate_finalization_chain( "winner candidate reference does not match client result", ) - _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", - ) - _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") - return winner + @staticmethod + def _get_candidate_reference(client_result: ClientResult) -> ArtifactReference: + candidate_reference = client_result.candidate_output_manifest + _require(candidate_reference is not None, "complete client result has no candidate manifest reference") + return candidate_reference @staticmethod def _get_expected_dataset_path( 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/bundle.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py new file mode 100644 index 000000000..7027c496c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py @@ -0,0 +1,326 @@ +# 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$") +_ENTRYPOINT_NAME = "entrypoint.sh" +_SLURM_PACKAGE_ROOT = "data_designer/slurm" +_SOURCE_MANIFEST_NAME = f"{_SLURM_PACKAGE_ROOT}/runtime/slurm-sources.txt" +_ENTRYPOINT = b"""#!/usr/bin/env bash +set -Eeuo pipefail + +dd_slurm_run_allocation() { + if [[ $# -ne 2 ]]; then + printf '%s\\n' 'allocation runtime requires plan and attempt directory arguments' >&2 + return 64 + fi + local runtime_root + runtime_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" + PYTHONNOUSERSITE=1 PYTHONSAFEPATH=1 PYTHONPATH="${runtime_root}" \ + python3 -m data_designer.slurm.runtime.entrypoint --plan "$1" --attempt-dir "$2" +} +""" + + +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: + sources = _collect_slurm_sources(Path(__file__).parents[1]) + 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, + ): + _add_archive_file(archive, _ENTRYPOINT_NAME, _ENTRYPOINT, mode=_ENTRYPOINT_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/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py new file mode 100644 index 000000000..09a4a4cc2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -0,0 +1,608 @@ +# 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 dataclasses import dataclass +from datetime import datetime, timedelta +from pathlib import Path +from typing import Protocol + +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.preflight import AllocationPreflight +from data_designer.slurm.runtime.probes import ReadinessProber +from data_designer.slurm.runtime.records import load_complete_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.serving.resolver import resolve_vllm_server +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + AttemptTerminalClassification, + DeploymentReadiness, + EndpointPublicationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, + 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 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 only complete candidates.""" + + def __init__( + self, + context: AllocationContext, + *, + runtime_proxy_path: Path, + 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") + self._context = context + self._runtime_proxy_path = runtime_proxy_path + 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) + 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 _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 = tuple( + resolve_vllm_server(self._context.plan, deployment.deployment_id) + for deployment in self._context.plan.deployments + ) + 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, + ) + ) + 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, + ) + ) + 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, + ) + ) + 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_complete_client_candidate(self._context) + 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 - the record contract requires this for complete results + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "complete client result has no candidate reference", + ) + return candidate_reference, client_result.completed_at + + 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, + ) + ) + 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 + readiness = AttemptReadiness( + schema_version=1, + run_id=self._attempt.run_id, + shard_id=self._attempt.shard_id, + attempt_id=self._attempt.attempt_id, + revision=revision, + updated_at=self._now(), + 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 + status.last_probe = _probe_evidence( + self._now(), + ProbeOutcome.SUCCESS, + "runtime_stopped", + "allocation processes stopped", + ) + 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, + candidate_output=None, + 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/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py new file mode 100644 index 000000000..30e2fd077 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local command entrypoint loaded from the checksummed runtime bundle.""" + +from __future__ import annotations + +import argparse +import os +import sys +from collections.abc import Mapping, Sequence +from pathlib import Path + +from data_designer.slurm.planning import PlannedShard +from data_designer.slurm.runtime.controller import OneNodeAllocationController +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.preflight import SystemAllocationPreflight +from data_designer.slurm.runtime.probes import HttpReadinessProber +from data_designer.slurm.runtime.signals import TerminationSignalCoordinator +from data_designer.slurm.runtime.steps import DefaultClientStepBuilder +from data_designer.slurm.runtime.supervisor import StepSupervisor, SubprocessStepRunner, SystemRuntimeClock +from data_designer.slurm.state import SlurmStateWriter + + +def main(arguments: Sequence[str] | None = None) -> int: + """Execute one allocation and return a bounded process status.""" + parser = argparse.ArgumentParser(prog="data-designer-slurm-runtime") + parser.add_argument("--plan", required=True) + parser.add_argument("--attempt-dir", required=True) + parsed = parser.parse_args(arguments) + try: + _run(Path(parsed.plan), Path(parsed.attempt_dir), 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 _run(plan_path: Path, attempt_directory: Path, environment: Mapping[str, str]) -> None: + context, writer = _load_allocation_context(plan_path, attempt_directory, environment) + clock = SystemRuntimeClock() + signals = TerminationSignalCoordinator() + supervisor = StepSupervisor(SubprocessStepRunner(), signals=signals, clock=clock) + controller = OneNodeAllocationController( + context, + runtime_proxy_path=Path(__file__).with_name("proxy.py"), + state=writer, + supervisor=supervisor, + preflight=SystemAllocationPreflight(), + client_steps=DefaultClientStepBuilder(), + prober=HttpReadinessProber(), + clock=clock, + environment=environment, + ) + with signals.interrupt_on_termination(supervisor.cleanup): + controller.run() + + +def _load_allocation_context( + plan_path: Path, + attempt_directory: Path, + environment: Mapping[str, str], +) -> tuple[AllocationContext, SlurmStateWriter]: + 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"))) + expected_attempt_root = plan_path.parent / "shards" / shard.shard_id / "attempts" + if attempt_directory.parent != expected_attempt_root: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "attempt directory does not match the scheduler array task", + ) + attempt = writer.load_attempt(shard.shard_id, attempt_directory.name) + return AllocationContext(plan, shard, attempt, attempt_directory), writer + + +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" or plan_path.parent.parent.name != "runs": + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + workspace_root = plan_path.parent.parent.parent + run_id = plan_path.parent.name + return SlurmStateWriter(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) + + +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/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/models.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py new file mode 100644 index 000000000..50d2557fe --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -0,0 +1,149 @@ +# 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 + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) + + +class RuntimeStepRole(str, Enum): + """Lifecycle role of one allocation-local Slurm step.""" + + CLIENT_PREFLIGHT = "client_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 + + 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", + ) 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..23315a2c1 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py @@ -0,0 +1,47 @@ +# 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 + + +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) + + +__all__ = ["get_container_path"] 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..476d62a07 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -0,0 +1,190 @@ +# 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 pathlib import Path +from typing import Protocol + +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.paths import get_container_path + +_DIGEST_CHUNK_SIZE = 1024 * 1024 +_GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") + + +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) + self._verify_attempt_directory(context.attempt_directory) + get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) + self._verify_artifacts(context) + self._verify_ports(context) + 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), + } + if node_indices != {0}: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "one-node runtime received a multi-node execution plan", + ) + expected = { + "SLURM_ARRAY_JOB_ID": context.attempt.scheduler.array_job_id, + "SLURM_ARRAY_TASK_ID": context.shard.array_task_index, + "SLURM_JOB_NUM_NODES": 1, + "SLURM_NODEID": 0, + } + 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_attempt_directory(attempt_directory: Path) -> None: + 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: + 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=plan.client.image.path, sha256=plan.client.image.sha256), + *( + 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} + for reference in unique_references.values(): + _verify_artifact(reference) + + @staticmethod + def _verify_ports(context: AllocationContext) -> None: + ports = tuple(port.port for port in context.plan.client.ports) + tuple( + port.port for deployment in context.plan.deployments for port in deployment.ports + ) + 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(("127.0.0.1", port)) + except OSError as error: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.PREFLIGHT_FAILED, + "one or more resolved 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__ = ["AllocationPreflight", "SystemAllocationPreflight"] 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..64dd51fe2 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Allocation-local least-connections HTTP endpoint for one model alias.""" + +from __future__ import annotations + +import argparse +import http.client +import threading +from collections.abc import Sequence +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import SplitResult, urlsplit + +_MAXIMUM_REQUEST_BYTES = 64 * 1024 * 1024 +_MAXIMUM_RESPONSE_BYTES = 64 * 1024 * 1024 +_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 + + +class _BackendPool: + def __init__(self, backends: tuple[_Backend, ...], retry_after_seconds: int | None) -> None: + self.backends = backends + self.retry_after_seconds = retry_after_seconds + self._active = [0] * len(backends) + self._cursor = 0 + self._lock = threading.Lock() + + def acquire(self, excluded: frozenset[int]) -> int: + with self._lock: + 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 = tuple(index for index in available if self._active[index] == minimum) + index = candidates[self._cursor % len(candidates)] + self._cursor += 1 + self._active[index] += 1 + return index + + def release(self, index: int) -> None: + with self._lock: + self._active[index] -= 1 + + +class _ProxyServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, address: tuple[str, int], pool: _BackendPool) -> None: + self.pool = pool + super().__init__(address, _ProxyHandler) + + +class _ProxyHandler(BaseHTTPRequestHandler): + server: _ProxyServer + protocol_version = "HTTP/1.1" + + def do_GET(self) -> None: # noqa: N802 + if self.path == "/health": + self._write_response(200, b'{"status":"ok"}\n', {"Content-Type": "application/json"}) + return + self._forward() + + def do_POST(self) -> None: # noqa: N802 + self._forward() + + def do_DELETE(self) -> None: # noqa: N802 + self._forward() + + def do_PUT(self) -> None: # noqa: N802 + self._forward() + + def log_message(self, format: str, *args: object) -> None: + del format, args + + def _forward(self) -> None: + if not self.path.startswith("/") or self.path.startswith("//"): + self.send_error(400) + return + request_body = self._read_request_body() + if request_body is None: + return + final_response = self._try_backends(request_body) + if final_response is None: + self.send_error(503) + return + if final_response[0] == 429: + final_response = ( + final_response[0], + final_response[1], + _retry_after_headers(final_response[2], self.server.pool.retry_after_seconds), + ) + self._write_response(*final_response) + + def _try_backends(self, request_body: bytes) -> tuple[int, bytes, dict[str, str]] | None: + excluded: set[int] = set() + final_response: tuple[int, bytes, dict[str, str]] | None = None + while len(excluded) < len(self.server.pool.backends): + index = self.server.pool.acquire(frozenset(excluded)) + try: + final_response = self._request_backend(self.server.pool.backends[index], request_body) + finally: + self.server.pool.release(index) + if final_response[0] != 429: + break + excluded.add(index) + return final_response + + def _read_request_body(self) -> bytes | None: + if self.headers.get("Transfer-Encoding") is not None: + self.send_error(501) + return None + lengths = self.headers.get_all("Content-Length", failobj=[]) + if len(lengths) > 1: + self.send_error(400) + return None + raw_length = lengths[0] if lengths else "0" + try: + length = int(raw_length) + except ValueError: + self.send_error(400) + return None + if length < 0 or length > _MAXIMUM_REQUEST_BYTES: + self.send_error(413) + return None + content = self.rfile.read(length) + if len(content) != length: + self.send_error(400) + return None + return content + + def _request_backend(self, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: + request_headers = tuple(self.headers.items()) + hop_headers = _get_hop_headers(request_headers) + headers = { + name: value + for name, value in request_headers + if name.lower() not in hop_headers and name.lower() not in {"host", "content-length"} + } + headers["Content-Length"] = str(len(body)) + connection = http.client.HTTPConnection(backend.host, backend.port, timeout=300) + try: + connection.request(self.command, self.path, body=body, headers=headers) + response = connection.getresponse() + payload = response.read(_MAXIMUM_RESPONSE_BYTES + 1) + if len(payload) > _MAXIMUM_RESPONSE_BYTES: + return 502, b'{"error":"backend response too large"}\n', {"Content-Type": "application/json"} + raw_response_headers = tuple(response.getheaders()) + response_hop_headers = _get_hop_headers(raw_response_headers) + response_headers = { + name: value + for name, value in raw_response_headers + if name.lower() not in response_hop_headers and name.lower() != "content-length" + } + return response.status, payload, response_headers + except (OSError, http.client.HTTPException): + return 502, b'{"error":"backend unavailable"}\n', {"Content-Type": "application/json"} + finally: + connection.close() + + def _write_response(self, status: int, body: bytes, headers: dict[str, str]) -> None: + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("Content-Length", str(len(body))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(body) + + +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("--retry-after-seconds", type=int) + parsed = parser.parse_args(arguments) + backends = tuple(_parse_backend(value) 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") + pool = _BackendPool(backends, parsed.retry_after_seconds) + with _ProxyServer(("127.0.0.1", parsed.listen_port), pool) as server: + server.serve_forever(poll_interval=0.1) + return 0 + + +def _parse_backend(value: str) -> _Backend: + parsed: SplitResult = urlsplit(value) + try: + port = parsed.port + except ValueError as error: + raise argparse.ArgumentTypeError("backend port is invalid") from error + if ( + parsed.scheme != "http" + or parsed.hostname != "127.0.0.1" + 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 loopback HTTP origins") + return _Backend(parsed.hostname, port) + + +def _retry_after_headers(headers: dict[str, str], retry_after_seconds: int | None) -> dict[str, str]: + selected = {name: value for name, value in headers.items() if name.lower() != "retry-after"} + if retry_after_seconds is not None: + selected["Retry-After"] = str(retry_after_seconds) + return selected + + +def _get_hop_headers(headers: Sequence[tuple[str, str]]) -> frozenset[str]: + nominated = { + token.strip().lower() + for name, value in headers + if name.lower() == "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..1644fe1bf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -0,0 +1,71 @@ +# 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 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_complete_client_candidate( + context: AllocationContext, +) -> tuple[ClientResult, CandidateOutputManifest]: + """Load a complete semantic client result and its digest-bound candidate.""" + client_result = _read_record(context.attempt_directory, _CLIENT_RESULT_NAME, ClientResult) + if client_result.outcome is not ClientOutcome.COMPLETE: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.CLIENT_FAILED, + f"client generation finished with outcome {client_result.outcome.value!r}", + ) + candidate = _read_record(context.attempt_directory, _CANDIDATE_NAME, CandidateOutputManifest) + try: + PlanStateValidator(context.plan).validate_client_candidate( + context.shard, + context.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_complete_client_candidate"] 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/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py new file mode 100644 index 000000000..ed998331a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -0,0 +1,458 @@ +# 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 Protocol + +from pydantic import BaseModel + +from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.planning import PlannedShard, ResolvedSlurmRunPlan +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 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 = ( + "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], + ) -> 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], + ) -> 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], + ) -> 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, + ) + + def build_generation_step( + self, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: Mapping[str, str], + ) -> 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, + ) + + @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], + ) -> RuntimeStep: + command = _build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) + secret_names, environment = _build_client_environment(plan, source_environment) + return _build_srun_step( + step_id=step_id, + role=role, + image_path=plan.client.image.path, + command=command, + environment=environment, + container_environment=secret_names, + plan=plan, + attempt_directory=attempt_directory, + ) + + +def _build_client_command( + operation: str, + plan: ResolvedSlurmRunPlan, + shard: PlannedShard, + attempt: AttemptManifest, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], +) -> 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 ( + "python3", + "-m", + "data_designer.slurm.client.worker", + operation, + "--plan", + get_container_path(plan, plan_path(plan)), + "--shard-id", + shard.shard_id, + "--attempt-id", + attempt.attempt_id, + "--attempt-dir", + get_container_path(plan, attempt_directory.as_posix(), require_writable=True), + *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 + return secret_names, environment + + +def build_vllm_steps( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], +) -> tuple[RuntimeStep, ...]: + """Build one structured server step per resolved vLLM process.""" + return tuple( + _build_vllm_step(deployment, process, plan, attempt_directory, source_environment) + 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, +) -> tuple[tuple[RuntimeStep, RuntimeEndpoint], ...]: + """Build one package-owned logical-endpoint process per deployment.""" + steps: list[tuple[RuntimeStep, RuntimeEndpoint]] = [] + for deployment in deployments: + steps.append(_build_endpoint_step(deployment, plan, attempt_directory, source_environment, runtime_proxy_path)) + return tuple(steps) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + plan: ResolvedSlurmRunPlan, + attempt_directory: Path, + source_environment: Mapping[str, str], + runtime_proxy_path: 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) + step = _build_srun_step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + image_path=plan.client.image.path, + command=command, + environment=_base_environment(source_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, +) -> tuple[str, ...]: + backends = tuple(f"http://127.0.0.1:{backend.port}" for backend in deployment.backend_endpoints) + 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 () + return ( + "python3", + get_container_path(plan, runtime_proxy_path.as_posix()), + "--listen-port", + str(port), + *retry_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], +) -> RuntimeStep: + _validate_local_vllm_process(process) + command = _build_vllm_command(deployment, process) + environment, container_environment = _build_vllm_environment(deployment, source_environment) + 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, +) -> tuple[str, ...]: + if process.http_port is None: # pragma: no cover - validated before command construction + raise AssertionError("vLLM HTTP port is unavailable") + command: tuple[str, ...] = ( + deployment.executable_path, + "serve", + deployment.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), + ) + if deployment.launch_policy.enable_expert_parallel: + command += ("--enable-expert-parallel",) + return command + deployment.launch_policy.extra_args + + +def _build_vllm_environment( + deployment: ResolvedVllmServerDeployment, + source_environment: Mapping[str, str], +) -> 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) + return environment, tuple(container_environment) + + +def _build_srun_step( + *, + step_id: str, + role: RuntimeStepRole, + image_path: str, + 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) + _add_srun_container_options(srun_command, plan, container_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) -> list[str]: + return [ + "srun", + "--nodes=1", + "--ntasks=1", + "--exact", + "--overlap", + "--unbuffered", + "--export=ALL", + f"--cpus-per-task={plan.client.authored.cpus}", + f"--container-image={image_path}", + ] + + +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, ...], +) -> None: + mounts = _render_mounts(plan) + 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) -> 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) + return ",".join(rendered) + + +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, ...]: + names: set[str] = set() + _collect_secret_environment_names(plan.client.authored.dependencies.index_credentials, names) + _collect_secret_environment_names(plan.invocation.authored.mcp_providers, names) + return tuple(sorted(names)) + + +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/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/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py index fb88e1965..b5957da9b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -24,6 +24,7 @@ class ReadinessState(str, Enum): PENDING = "pending" + RESTARTING = "restarting" STARTING = "starting" READY = "ready" FAILED = "failed" @@ -72,32 +73,8 @@ class DeploymentReadiness(StateValue): @model_validator(mode="after") def validate_counts_and_state(self) -> DeploymentReadiness: - if self.ready_backends > self.expected_backends: - raise ValueError("ready_backends must not exceed expected_backends") - if self.endpoint_publication is EndpointPublicationState.FAILED and self.state not in { - ReadinessState.FAILED, - ReadinessState.STOPPED, - }: - raise ValueError("failed endpoint publication requires a failed or stopped deployment") - - if self.state is ReadinessState.PENDING: - if self.ready_backends != 0: - raise ValueError("pending deployments cannot have ready backends") - if self.endpoint_publication is not EndpointPublicationState.PENDING: - raise ValueError("pending deployments require pending endpoint publication") - elif self.state is ReadinessState.STARTING: - if ( - self.ready_backends == self.expected_backends - and self.endpoint_publication is EndpointPublicationState.PUBLISHED - ): - raise ValueError("fully ready published deployments must use the ready state") - elif self.state is ReadinessState.READY: - if self.ready_backends != self.expected_backends: - raise ValueError("ready deployments require every expected backend") - if self.endpoint_publication is not EndpointPublicationState.PUBLISHED: - raise ValueError("ready deployments require a published endpoint") - elif self.state is ReadinessState.STOPPED and self.ready_backends != 0: - raise ValueError("stopped deployments cannot have ready backends") + _validate_deployment_common(self) + _validate_deployment_state(self) return self @@ -116,33 +93,85 @@ class AttemptReadiness(StateRecord): @model_validator(mode="after") def validate_deployments(self) -> AttemptReadiness: - deployment_ids = [deployment.deployment_id for deployment in self.deployments] - model_aliases = [deployment.model_alias for deployment in self.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") - if any( - deployment.last_probe is not None and deployment.last_probe.observed_at > self.updated_at - for deployment in self.deployments - ): - raise ValueError("probe observations must not be later than the readiness snapshot") - - deployment_states = tuple(deployment.state for deployment in self.deployments) - if self.state is ReadinessState.PENDING: - if any(state is not ReadinessState.PENDING for state in deployment_states): - raise ValueError("pending attempts require every deployment to be pending") - elif self.state is ReadinessState.STARTING: - if any(state in {ReadinessState.FAILED, ReadinessState.STOPPED} for state in deployment_states): - raise ValueError("starting attempts cannot contain failed or stopped deployments") - if all(state is ReadinessState.READY for state in deployment_states): - raise ValueError("attempts with every deployment ready must use the ready state") - elif self.state is ReadinessState.READY: - if any(state is not ReadinessState.READY for state in deployment_states): - raise ValueError("ready attempts require every deployment to be ready") - elif self.state is ReadinessState.FAILED: - if not any(state is ReadinessState.FAILED for state in deployment_states): - raise ValueError("failed attempts require at least one failed deployment") - elif any(state is not ReadinessState.STOPPED for state in deployment_states): - raise ValueError("stopped attempts require every deployment to be stopped") + _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 index aacfacc65..0d00e7b3b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -20,12 +20,29 @@ _ALLOWED_READINESS_TRANSITIONS: dict[ReadinessState, frozenset[ReadinessState]] = { ReadinessState.PENDING: frozenset( - {ReadinessState.PENDING, ReadinessState.STARTING, ReadinessState.FAILED, ReadinessState.STOPPED} + { + ReadinessState.PENDING, + ReadinessState.RESTARTING, + ReadinessState.STARTING, + ReadinessState.FAILED, + ReadinessState.STOPPED, + } ), ReadinessState.STARTING: frozenset( - {ReadinessState.STARTING, ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED} + { + 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.READY: frozenset({ReadinessState.READY, ReadinessState.FAILED, ReadinessState.STOPPED}), ReadinessState.FAILED: frozenset({ReadinessState.FAILED, ReadinessState.STOPPED}), ReadinessState.STOPPED: frozenset({ReadinessState.STOPPED}), } @@ -74,6 +91,7 @@ def validate_readiness_transition( "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, @@ -87,20 +105,33 @@ def validate_readiness_transition( old_deployment.expected_backends == new_deployment.expected_backends, "readiness expected backend count cannot change", ) - _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", - ) + 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: + 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" diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index 17e99ea2c..2a7a7d834 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -359,6 +359,27 @@ def test_plan_state_validator_reuses_one_context_for_an_attempt_batch() -> None: ) +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"), [ @@ -446,8 +467,8 @@ def test_finalization_rejects_invalid_chains( ("winner_shard_id", "winner shard_id"), ("client_attempt_id", "client result attempt_id"), ("candidate_attempt_id", "candidate attempt_id"), - ("candidate_ordinal", "attempt ordinals"), - ("winner_ordinal", "attempt ordinals"), + ("candidate_ordinal", "attempt ordinal"), + ("winner_ordinal", "attempt ordinal"), ("client_requested_records", "requested records"), ("candidate_requested_records", "requested records"), ("client_actual_records", "actual records"), 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..f9c405e80 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import cast + +import pytest + +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, + SchedulerIdentity, + 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) + + def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: + validate_attempt_transition(self.attempt, attempt) + self.attempt = attempt + return attempt + + 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 + + +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 build_preflight_step( + self, + plan: ResolvedSlurmRunPlan, + shard: object, + attempt: object, + attempt_directory: Path, + endpoints: tuple[RuntimeEndpoint, ...], + source_environment: object, + ) -> RuntimeStep: + del plan, shard, attempt, endpoints, source_environment + 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, + ) -> RuntimeStep: + del plan, shard, attempt, endpoints, source_environment + 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_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py new file mode 100644 index 000000000..a1d67bde8 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -0,0 +1,238 @@ +# 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[0] == "entrypoint.sh" + assert names[1] == "data_designer/slurm/runtime/slurm-sources.txt" + assert names[2] == "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 + assert b'PYTHONPATH="${runtime_root}"' in entrypoint.read() + + +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("") + (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 _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_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py new file mode 100644 index 000000000..c302ce8d7 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -0,0 +1,641 @@ +# 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, replace +from datetime import timedelta +from pathlib import Path +from typing import Callable + +import pytest +from conftest import FakeClientStepBuilder, FakePreflight, FakeStateStore, RuntimeCase +from slurm_test_fakes import FakeClock + +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.contracts import ArtifactReference +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, +) + + +@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) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + 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={}, + ) + + result = controller.run() + + assert result.state is AttemptLifecycleState.SUCCEEDED + assert result.terminal_classification is AttemptTerminalClassification.SUCCEEDED + assert result.candidate_output is not None + 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 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_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={}, + ) + + 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={}, + ) + + 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={}, + ) + + 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={}, + ) + + 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={}, + ) + + 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={}, + ) + + 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={}, + ) + + with pytest.raises(SlurmRuntimeError, match="cleanup failed"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + 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={}, + ) + + 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={}, + ) + + 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={}, + ) + + with pytest.raises(SlurmRuntimeError, match="later than the allocation clock"): + controller.run() + + assert state.attempt.state is AttemptLifecycleState.FAILED + + +def test_partial_client_result_is_classified_before_candidate_loading(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={}, + ) + + with pytest.raises(SlurmRuntimeError) as raised: + controller.run() + + assert raised.value.code is SlurmRuntimeErrorCode.CLIENT_FAILED + assert not (runtime_case.context.attempt_directory / "output-manifest.json").exists() + + +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={}, + ) + + 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) -> None: + context = runtime_case.context + requested = context.shard.requested_records + dataset_path = (context.attempt_directory / "dataset").as_posix() + 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=requested, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path="part-00000.parquet", + sha256="a" * 64, + byte_size=128, + record_count=requested, + ), + ), + dataset_schema_digest="b" * 64, + provenance_digest="c" * 64, + ) + 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=requested, + outcome=ClientOutcome.COMPLETE, + 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 _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..b83f83c34 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.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 + +import pytest + +from data_designer.slurm.runtime.entrypoint import main + + +def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: + assert main(("--plan", "resolved-plan.json", "--attempt-dir", "attempt-0001")) == 64 + captured = capsys.readouterr() + assert "runtime paths must be absolute" in captured.err + assert "Traceback" not in captured.err 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..ebdd820db --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_paths.py @@ -0,0 +1,47 @@ +# 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 + + +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" + + +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_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py new file mode 100644 index 000000000..e4be62e92 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -0,0 +1,138 @@ +# 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 conftest import RuntimeCase + +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.preflight import SystemAllocationPreflight, _verify_artifact + + +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) + + +@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) + + +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) + 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) 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..0a600f1bf --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import io +from email.message import Message +from types import MethodType + +import pytest + +from data_designer.slurm.runtime import proxy as runtime_proxy +from data_designer.slurm.runtime.proxy import _Backend, _BackendPool, _parse_backend, _ProxyHandler + + +def test_proxy_retries_overload_on_another_least_active_backend() -> None: + handler = object.__new__(_ProxyHandler) + pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) + handler.server = type("Server", (), {"pool": pool})() + handler.path = "/v1/chat/completions" + handler.command = "POST" + handler.rfile = io.BytesIO(b"{}") + handler.headers = Message() + handler.headers["Content-Length"] = "2" + calls: list[int] = [] + captured: list[tuple[int, bytes, dict[str, str]]] = [] + + def request_backend(self: _ProxyHandler, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: + del self + calls.append(backend.port) + assert body == b"{}" + if backend.port == 8001: + return 429, b"overloaded", {} + return 200, b"complete", {"Content-Type": "application/json"} + + def write_response(self: _ProxyHandler, status: int, body: bytes, headers: dict[str, str]) -> None: + del self + captured.append((status, body, headers)) + + handler._request_backend = MethodType(request_backend, handler) + handler._write_response = MethodType(write_response, handler) + + handler._forward() + + assert calls == [8001, 8002] + assert captured == [(200, b"complete", {"Content-Type": "application/json"})] + + +def test_proxy_passes_final_overload_through_after_one_backend_pass() -> None: + handler = object.__new__(_ProxyHandler) + pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 3) + handler.server = type("Server", (), {"pool": pool})() + handler.path = "/v1/chat/completions" + handler.command = "POST" + handler.rfile = io.BytesIO(b"{}") + handler.headers = Message() + handler.headers["Content-Length"] = "2" + calls: list[int] = [] + captured: list[tuple[int, bytes, dict[str, str]]] = [] + + def request_backend(self: _ProxyHandler, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: + del self, body + calls.append(backend.port) + return 429, f"overloaded-{backend.port}".encode(), {"Retry-After": "99"} + + handler._request_backend = MethodType(request_backend, handler) + handler._write_response = MethodType( + lambda self, status, body, headers: captured.append((status, body, headers)), + handler, + ) + + handler._forward() + + assert calls == [8001, 8002] + assert captured == [(429, b"overloaded-8002", {"Retry-After": "3"})] + + +def test_proxy_rejects_a_truncated_fixed_length_request() -> None: + handler = object.__new__(_ProxyHandler) + handler.rfile = io.BytesIO(b"{") + handler.headers = Message() + handler.headers["Content-Length"] = "2" + errors: list[int] = [] + handler.send_error = MethodType(lambda self, status: errors.append(status), handler) + + assert handler._read_request_body() is None + assert errors == [400] + + +def test_proxy_bounds_backend_response_buffering(monkeypatch: pytest.MonkeyPatch) -> None: + class _Response: + status = 200 + + def read(self, amount: int) -> bytes: + assert amount == 5 + return b"12345" + + def getheaders(self) -> list[tuple[str, str]]: + return [] + + class _Connection: + def __init__(self, host: str, port: int, timeout: int) -> None: + del host, port, timeout + + def request(self, method: str, path: str, *, body: bytes, headers: dict[str, str]) -> None: + del method, path, body, headers + + def getresponse(self) -> _Response: + return _Response() + + def close(self) -> None: + pass + + handler = object.__new__(_ProxyHandler) + handler.command = "POST" + handler.path = "/v1/chat/completions" + handler.headers = Message() + monkeypatch.setattr(runtime_proxy, "_MAXIMUM_RESPONSE_BYTES", 4) + monkeypatch.setattr(runtime_proxy.http.client, "HTTPConnection", _Connection) + + assert handler._request_backend(_Backend("127.0.0.1", 8001), b"{}") == ( + 502, + b'{"error":"backend response too large"}\n', + {"Content-Type": "application/json"}, + ) + + +def test_proxy_strips_connection_nominated_headers_in_both_directions(monkeypatch: pytest.MonkeyPatch) -> None: + class _Response: + status = 200 + + def read(self, amount: int) -> bytes: + del amount + return b"complete" + + def getheaders(self) -> list[tuple[str, str]]: + return [ + ("Connection", "X-Backend-Hop"), + ("X-Backend-Hop", "private"), + ("Content-Type", "application/json"), + ] + + class _Connection: + forwarded_headers: dict[str, str] = {} + + def __init__(self, host: str, port: int, timeout: int) -> None: + del host, port, timeout + + def request(self, method: str, path: str, *, body: bytes, headers: dict[str, str]) -> None: + del method, path, body + type(self).forwarded_headers = headers + + def getresponse(self) -> _Response: + return _Response() + + def close(self) -> None: + pass + + handler = object.__new__(_ProxyHandler) + handler.command = "POST" + handler.path = "/v1/chat/completions" + handler.headers = Message() + handler.headers["Connection"] = "X-Client-Hop" + handler.headers["X-Client-Hop"] = "private" + handler.headers["Proxy-Connection"] = "keep-alive" + handler.headers["Authorization"] = "Bearer reviewed" + monkeypatch.setattr(runtime_proxy.http.client, "HTTPConnection", _Connection) + + status, body, headers = handler._request_backend(_Backend("127.0.0.1", 8001), b"{}") + + assert (status, body) == (200, b"complete") + assert _Connection.forwarded_headers == {"Authorization": "Bearer reviewed", "Content-Length": "2"} + assert headers == {"Content-Type": "application/json"} + + +@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_pool_selects_least_active_backend() -> None: + pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) + first = pool.acquire(frozenset()) + second = pool.acquire(frozenset()) + pool.release(second) + pool.release(first) + + assert first != second 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..b01ed41aa --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_steps.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 datetime import datetime, timezone +from pathlib import Path + +import pytest +from conftest import RuntimeCase, relocate_plan + +from data_designer.slurm.config import QueueBackpressureConfig +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.runtime.errors import SlurmRuntimeError +from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.steps import ( + DefaultClientStepBuilder, + build_endpoint_steps, + 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 = { + "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, + ) + 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) + + 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 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) + + +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, + {}, + 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, + {}, + ) + + separator = step.command.index("--") + worker = step.command[separator + 1 :] + assert worker[:5] == ( + "python3", + "-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 + + +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 + 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, + { + "PACKAGE_INDEX_TOKEN": "client-secret", + "HF_TOKEN": "server-secret", + }, + ) + + assert step.environment["PACKAGE_INDEX_TOKEN"] == "client-secret" + assert "HF_TOKEN" not in step.environment + assert "--container-env=PACKAGE_INDEX_TOKEN" in step.command + + with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): + DefaultClientStepBuilder().build_preflight_step( + plan, + shard, + attempt, + attempt_directory, + endpoints, + {"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/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index 037785899..ff0493b19 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -4,7 +4,7 @@ from __future__ import annotations import json -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import TypeVar @@ -305,6 +305,33 @@ def test_readiness_cannot_move_backward_or_change_backend_count() -> None: 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, From 3be7063f5429328ed1f26846134e7049005e75d6 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:17:26 -0300 Subject: [PATCH 17/62] feat: add allocation-local Slurm client worker (#911) * feat(slurm): add allocation-local client worker * fix(slurm): address client worker review feedback * fix(slurm): harden client worker boundaries * fix(slurm): harden client filesystem boundary --- .../data_designer/slurm/client/__init__.py | 24 +- .../data_designer/slurm/client/environment.py | 428 ++++++++++ .../src/data_designer/slurm/client/errors.py | 15 + .../data_designer/slurm/client/execution.py | 751 ++++++++++++++++++ .../data_designer/slurm/client/filesystem.py | 197 +++++ .../src/data_designer/slurm/client/plugins.py | 53 ++ .../src/data_designer/slurm/client/records.py | 168 +++- .../src/data_designer/slurm/client/worker.py | 112 +++ .../src/data_designer/slurm/config/images.py | 15 +- .../src/data_designer/slurm/contracts.py | 18 +- .../data_designer/slurm/planning/models.py | 2 +- .../tests/client/conftest.py | 151 ++++ .../tests/client/test_environment.py | 245 ++++++ .../tests/client/test_filesystem.py | 96 +++ .../tests/client/test_worker.py | 517 ++++++++++++ .../tests/contracts/test_shared_records.py | 7 + .../fake_data_designer_plugin/plugin.py | 13 + 17 files changed, 2788 insertions(+), 24 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/environment.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/errors.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/execution.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/plugins.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/worker.py create mode 100644 packages/data-designer-slurm/tests/client/conftest.py create mode 100644 packages/data-designer-slurm/tests/client/test_environment.py create mode 100644 packages/data-designer-slurm/tests/client/test_filesystem.py create mode 100644 packages/data-designer-slurm/tests/client/test_worker.py 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 index 96fe5effa..da4cf64fa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/__init__.py @@ -5,6 +5,26 @@ from __future__ import annotations -from data_designer.slurm.client.records import ClientOutcome, ClientResult +from data_designer.slurm.client.records import ( + ClientEnvironmentManifest, + ClientEnvironmentOutcome, + ClientErrorCode, + ClientInstallerOutcome, + ClientOutcome, + ClientPluginEntryPoint, + ClientProgress, + ClientProgressPhase, + ClientResult, +) -__all__ = ["ClientOutcome", "ClientResult"] +__all__ = [ + "ClientEnvironmentManifest", + "ClientEnvironmentOutcome", + "ClientErrorCode", + "ClientInstallerOutcome", + "ClientOutcome", + "ClientPluginEntryPoint", + "ClientProgress", + "ClientProgressPhase", + "ClientResult", +] 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..aa48076e8 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -0,0 +1,428 @@ +# 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 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 +from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution + +DistributionInventory = Callable[[Path | None], tuple[InstalledDistribution, ...]] +CommandRunner = Callable[[tuple[str, ...]], None] + + +@dataclass(frozen=True) +class PreparedClientEnvironment: + run_id: str + shard_id: str + attempt_id: str + attempt_dir: Path + overlay_path: Path + dependency_lock: ArtifactReference + client_image_sha256: str + python_abi: str + installer_outcome: ClientInstallerOutcome + installed_distributions: tuple[InstalledDistribution, ...] + + +@dataclass(frozen=True) +class _BootstrapInputs: + run_id: str + run_root: Path + shard_id: str + attempt_id: str + attempt_dir: Path + client_image_sha256: str + python_abi: str + installer_path: Path + inspection: dict[str, object] + dependency_lock: ArtifactReference + + +@dataclass(frozen=True) +class _VerifiedDependencies: + image_distributions: tuple[InstalledDistribution, ...] + overlay_packages: tuple[dict[str, object], ...] + + +class ClientEnvironmentBuilder: + """Prepare one immutable attempt-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, + ) -> 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, + ) + 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, + overlay_path=overlay_path, + dependency_lock=inputs.dependency_lock, + client_image_sha256=inputs.client_image_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, + ) -> _BootstrapInputs: + _validate_input_path(plan_path, "resolved plan") + _validate_input_path(attempt_dir, "attempt directory") + 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 = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) + run_id = _require_string(plan, "run_id") + run_root = plan_path.parent + if plan_path.name != "resolved-plan.json" or run_root.name != run_id or run_root.parent.name != "runs": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved plan path is not canonical") + expected_attempt = run_root / "shards" / shard_id / "attempts" / attempt_id + 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, "client") + image = _require_object(client, "image") + image_sha256 = _require_digest(image, "sha256") + inspection_record = _require_object(image, "inspection") + inspection = _require_object(inspection_record, "inspection") + if inspection.get("kind") != "client": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved client image inspection is invalid") + python_abi = _require_string(inspection, "python_abi") + if python_abi != f"{interpreter_name()}{interpreter_version()}": + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client Python ABI differs from the plan") + installer_path = Path(_require_string(inspection, "installer_path")) + if not installer_path.is_absolute() or installer_path.as_posix() == "/": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client installer path is invalid") + + lock_reference = _artifact_reference(_require_object(client, "dependency_lock")) + if Path(lock_reference.path).as_posix() != (run_root / "dependency-lock.json").as_posix(): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency lock path is not canonical") + return _BootstrapInputs( + run_id=run_id, + run_root=run_root, + shard_id=shard_id, + attempt_id=attempt_id, + attempt_dir=attempt_dir, + client_image_sha256=image_sha256, + python_abi=python_abi, + installer_path=installer_path, + inspection=inspection, + 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_image_sha256") != inputs.client_image_sha256: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency lock targets another image") + if _require_string(lock, "python_abi") != inputs.python_abi: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency lock targets another Python ABI") + + expected_image = _parse_distributions(lock.get("image_distributions")) + inspected_image = _parse_distributions(inputs.inspection.get("distributions")) + try: + actual_image = self._inventory(None) + except ClientWorkerError: + raise + except Exception as error: + raise ClientWorkerError( + ClientErrorCode.DEPENDENCY_CONFLICT, "client image inventory cannot be verified" + ) from error + if expected_image != inspected_image or actual_image != expected_image: + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client image 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.run_root / "inputs") + return _VerifiedDependencies( + image_distributions=expected_image, + 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.run_root / "dependencies", + dependencies.image_distributions, + ) + overlay_path = inputs.attempt_dir / "client-env" / "site-packages" + outcome = self._install_overlay(inputs.installer_path, wheels, expected_overlay, overlay_path) + installed = tuple(sorted((*dependencies.image_distributions, *expected_overlay), key=lambda item: item.name)) + return overlay_path, outcome, installed + + def _install_overlay( + self, + installer_path: Path, + wheels: tuple[Path, ...], + expected: tuple[InstalledDistribution, ...], + target: Path, + ) -> ClientInstallerOutcome: + ensure_private_directory(target) + existing = self._inventory(target) + if existing == expected and any(target.iterdir()): + return ClientInstallerOutcome.REUSED + if any(target.iterdir()): + raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "attempt overlay contains unexpected files") + if not wheels: + return ClientInstallerOutcome.NOT_REQUIRED + try: + self._command_runner( + ( + installer_path.as_posix(), + "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" + ) + return ClientInstallerOutcome.INSTALLED + + +def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ...]: + """Return an exact immutable distribution inventory.""" + distributions = ( + importlib.metadata.distributions() if path is None else importlib.metadata.distributions(path=[path.as_posix()]) + ) + 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, "mutable installed distributions are forbidden" + ) + names.add(name) + installed.append(InstalledDistribution(name=name, version=version)) + return tuple(sorted(installed, key=lambda item: item.name)) + + +def activate_environment(prepared: PreparedClientEnvironment) -> None: + """Activate one verified overlay and isolate Data Designer attempt state.""" + home = prepared.attempt_dir / "data-designer-home" + cache = prepared.attempt_dir / "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 _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], ...], + dependencies_root: Path, + image_distributions: tuple[InstalledDistribution, ...], +) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: + expected: list[InstalledDistribution] = [] + wheels: list[Path] = [] + image_names = {item.name for item in image_distributions} + for package in packages: + name = canonicalize_name(_require_string(package, "name")) + version = _require_string(package, "version") + if name in image_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")) + wheel = Path(artifact.path) + try: + wheel.relative_to(dependencies_root) + except ValueError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") from error + if wheel.parent != dependencies_root or wheel.suffix != ".whl": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") + 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, root: Path) -> None: + path = Path(reference.path) + try: + path.relative_to(root) + except ValueError as error: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") from error + if path.parent != root or path.suffix != ".json": + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") + 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") + ) + 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 or "vcs_info" in direct_url: + return True + 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") + return not isinstance(hashes, dict) or not isinstance(hashes.get("sha256"), str) 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..8a6d3f05c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -0,0 +1,751 @@ +# 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, compute_canonical_json_sha256 +from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest + +Clock = Callable[[], datetime] +DataDesignerFactory = Callable[..., DataDesigner] +logger = logging.getLogger(__name__) + + +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 + + 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 + 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 + + 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: + generated_records = sum(lazy.pq.read_metadata(item).num_rows for item in path.parent.glob("*.parquet")) + self.best_effort( + ClientProgressPhase.GENERATING, + completed_records=min(generated_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, + ) -> None: + self._data_designer_factory = data_designer_factory + self._clock = clock or (lambda: datetime.now(timezone.utc)) + + def preflight( + self, + plan_path: Path, + *, + prepared: PreparedClientEnvironment, + endpoints: Mapping[str, str], + plugins: tuple[ClientPluginEntryPoint, ...], + ) -> ClientEnvironmentManifest: + """Validate packages, plugins, assets, and config without generation.""" + try: + context = self._build_context(plan_path, prepared=prepared, endpoints=endpoints) + 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, ...], + ) -> 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) + 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], + ) -> _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.image_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) + providers = self._materialize_model_endpoints(plan, builder, endpoints) + self._validate_model_references(builder) + self._materialize_seed(plan, shard, builder) + mcp_providers = self._materialize_mcp_providers(plan) + self._validate_assets(plan) + requested_resume = ResumeMode(plan.invocation.authored.resume) + dataset_path = self._dataset_path(shard, prepared, requested_resume) + designer = self._data_designer_factory( + artifact_path=dataset_path.parent, + model_providers=providers, + managed_assets_path=plan.invocation.effective_input_bindings.managed_assets_path, + 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, 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: + expected_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + inspection = plan.client.image.inspection_facts + if ( + plan.run_id != prepared.run_id + or prepared.attempt_dir != expected_attempt + or plan.client.dependency_lock != prepared.dependency_lock + or plan.client.image.sha256 != prepared.client_image_sha256 + or inspection.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": seed_path}) + 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) -> None: + 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") + + @staticmethod + def _dataset_path( + 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") + if context.requested_resume is not ResumeMode.NEVER: + ensure_private_directory(context.dataset_path) + if context.requested_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, + 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, exported_path = self._place_dataset(creation, results) + candidate = self._build_candidate_manifest(context, prepared, creation, dataset_path, exported_path) + 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") + if results.requested_resume_mode is not context.requested_resume: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer resume metadata differs") + + dataset_path = Path(results.dataset_path) + effective_resume = results.effective_resume_mode + 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, Path | None]: + dataset_path = creation.dataset_path + exported_path: Path | None = None + if creation.actual_records: + exported_path = results.export(dataset_path / "part-00000.parquet", format="parquet") + if dataset_path != creation.expected_path: + ensure_private_directory(creation.expected_path.parent) + dataset_path.rename(creation.expected_path) + if exported_path is not None: + exported_path = creation.expected_path / exported_path.relative_to(dataset_path) + 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, exported_path + + def _build_candidate_manifest( + self, + context: _ExecutionContext, + prepared: PreparedClientEnvironment, + creation: _ValidatedCreation, + dataset_path: Path, + exported_path: Path | None, + ) -> CandidateOutputManifest: + files: tuple[CandidateOutputFile, ...] = () + schema_digest = context.plan.builder.content_sha256 + if exported_path is not None: + metadata = lazy.pq.read_metadata(exported_path) + if metadata.num_rows != creation.actual_records: + raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "exported record count differs") + schema_digest = hashlib.sha256(lazy.pq.read_schema(exported_path).serialize().to_pybytes()).hexdigest() + files = ( + CandidateOutputFile( + relative_path=exported_path.relative_to(dataset_path).as_posix(), + sha256=compute_file_sha256(exported_path, missing_code=ClientErrorCode.OUTPUT_INVALID), + byte_size=exported_path.stat().st_size, + record_count=creation.actual_records, + ), + ) + created_at = self._clock() + provenance_digest = compute_canonical_json_sha256( + { + "builder_sha256": context.plan.builder.content_sha256, + "client_image_sha256": prepared.client_image_sha256, + "dependency_lock_sha256": prepared.dependency_lock.sha256, + "files": [file.model_dump(mode="json") for file in files], + "attempt_id": prepared.attempt_id, + "resolved_plan_sha256": context.plan.compute_sha256(), + "run_id": context.plan.run_id, + "shard_id": context.shard.shard_id, + } + ) + 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=provenance_digest, + ) + + 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_image_sha256 != prepared.client_image_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/records.py b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py index a76290cfc..33eb2dcfa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta from enum import Enum -from typing import Annotated, Literal +from typing import TYPE_CHECKING, Annotated, Literal from pydantic import NonNegativeInt, PositiveInt, StringConstraints, field_validator, model_validator @@ -13,11 +13,170 @@ 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_image_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_image_sha256=prepared.client_image_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" @@ -66,11 +225,8 @@ def validate_message(cls, value: str | None) -> str | None: 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 != "if_possible" and self.effective_resume_mode not in { - None, - self.requested_resume_mode, - }: - raise ValueError("effective resume mode must match a fixed requested mode") + 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") 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..5139b10c0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py @@ -0,0 +1,112 @@ +# 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 signal +import sys +from collections.abc import Sequence +from datetime import datetime, timezone +from pathlib import Path + +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) + prepared = ClientEnvironmentBuilder().prepare( + arguments.plan, + shard_id=arguments.shard_id, + attempt_id=arguments.attempt_id, + attempt_dir=arguments.attempt_dir, + ) + 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() + if arguments.operation == "preflight": + worker.preflight(arguments.plan, prepared=prepared, endpoints=endpoints, plugins=plugins) + else: + worker.run(arguments.plan, prepared=prepared, endpoints=endpoints, plugins=plugins) + 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("--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/images.py b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py index 0b4e983de..4c520394f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/images.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/images.py @@ -14,20 +14,12 @@ ContractRecord, ContractValue, Identifier, + InstalledDistribution, Sha256Digest, validate_absolute_path, validate_plain_text, ) -DistributionName = Annotated[ - str, - StringConstraints( - min_length=1, - max_length=128, - pattern=r"^[a-z0-9]+(?:-[a-z0-9]+)*$", - ), -] - class ImageKind(str, Enum): CLIENT = "client" @@ -75,11 +67,6 @@ def validate_source(cls, value: str) -> str: return value -class InstalledDistribution(ContractValue): - name: DistributionName - version: Annotated[str, StringConstraints(min_length=1, max_length=128)] - - class ClientImageInspection(ContractValue): kind: Literal[ImageKind.CLIENT] python_implementation: Identifier diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index 226b9899e..42a38ecc0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -9,7 +9,7 @@ import json import posixpath from collections.abc import Mapping -from typing import TypeVar +from typing import Annotated, TypeVar from urllib.parse import urlsplit from pydantic import ( @@ -17,6 +17,7 @@ ConfigDict, NonNegativeInt, PositiveInt, + StringConstraints, field_validator, model_validator, ) @@ -107,6 +108,21 @@ 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)] + + class AuthoredConfig(ContractValue): """Base for strict authored configuration values.""" 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 index ca2a57c80..d8380cc53 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -16,7 +16,6 @@ from data_designer.slurm.config.environment import validate_no_plaintext_secrets from data_designer.slurm.config.images import ( ClientImageInspection, - DistributionName, ImageInspectionRecord, ImageKind, ImageRef, @@ -37,6 +36,7 @@ ArtifactReference, ContractRecord, ContractValue, + DistributionName, Identifier, ModelAlias, RecordRange, 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..343225a5d --- /dev/null +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -0,0 +1,151 @@ +# 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 +from dataclasses import dataclass +from pathlib import Path + +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 + +GOLDEN_DIRECTORY = Path(__file__).parents[1] / "contracts" / "golden" + + +@dataclass(frozen=True) +class ClientWorkerCase: + plan: ResolvedSlurmRunPlan + plan_path: Path + lock: ResolvedDependencyLock + attempt_dir: 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 + + 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, + **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.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) + batch_path = dataset_path / "batch.parquet" + pq.write_table(pa.table({"value": tuple(range(actual))}), 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 +def client_worker_case(tmp_path: Path) -> 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()}" + payload["client"]["image"]["inspection"]["inspection"].update( + {"python_abi": python_abi, "python_version": platform.python_version()} + ) + lock_payload = json.loads((GOLDEN_DIRECTORY / "dependency_lock_single.json").read_text()) + lock_payload["python_abi"] = python_abi + lock = ResolvedDependencyLock.model_validate_json(json.dumps(lock_payload)) + payload["client"]["dependency_lock"]["sha256"] = lock.compute_sha256() + 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" + overlay_path = attempt_dir / "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, + overlay_path=overlay_path, + dependency_lock=plan.client.dependency_lock, + client_image_sha256=plan.client.image.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.image_distributions + ), + ) + endpoints = {plan.deployments[0].authored.model_alias: f"http://127.0.0.1:{plan.client.ports[0].port}/v1"} + return ClientWorkerCase(plan, plan_path, lock, attempt_dir, prepared, endpoints) 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..9578adc7d --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -0,0 +1,245 @@ +# 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 subprocess +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest +from conftest import ClientWorkerCase + +from data_designer.slurm.client.environment import ClientEnvironmentBuilder, inspect_distributions +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.plugins import discover_plugins +from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome +from data_designer.slurm.contracts import InstalledDistribution +from data_designer.slurm.planning import ResolvedSlurmRunPlan + + +def test_environment_prepares_empty_verified_overlay(client_worker_case: ClientWorkerCase) -> None: + def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: + return client_worker_case.lock.image_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, + ) + + assert prepared.installer_outcome is ClientInstallerOutcome.NOT_REQUIRED + assert prepared.installed_distributions == client_worker_case.lock.image_distributions + + +def test_inspect_distributions_omits_path_for_active_environment(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[dict[str, object]] = [] + + def distributions(**kwargs: object) -> tuple[()]: + calls.append(kwargs) + return () + + monkeypatch.setattr("data_designer.slurm.client.environment.importlib.metadata.distributions", distributions) + + assert inspect_distributions(None) == () + assert calls == [{}] + + +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: (distribution,), + ) + + with pytest.raises(ClientWorkerError) as error: + inspect_distributions(None) + + assert error.value.code is ClientErrorCode.DEPENDENCY_CONFLICT + + +def test_environment_rejects_client_image_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, + ) + + 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, + ) + + 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.image_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, + ) + + 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.image_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, + ) + + assert prepared.installer_outcome is ClientInstallerOutcome.INSTALLED + assert "--no-index" in commands[0] + assert "--no-deps" in commands[0] + + +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]), + overlay_path=Path(sys.argv[1]), + dependency_lock=ArtifactReference(path=(Path(sys.argv[2]) / "dependency-lock.json").as_posix(), sha256="a" * 64), + client_image_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" +""" + result = subprocess.run( + [sys.executable, "-c", script, fake_plugin_overlay.as_posix(), (tmp_path / "attempt").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_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py new file mode 100644 index 000000000..7d4089df6 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -0,0 +1,517 @@ +# 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 pytest +from conftest import ClientWorkerCase, FakeDataDesigner + +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 + + +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_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(), + "--endpoint", + f"judge=v2={endpoint}", + ] + ) + + assert result == 0 + assert worker.preflight.call_args.kwargs["endpoints"] == {"judge=v2": endpoint} + + +@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 result.candidate_output_manifest.sha256 == candidate.compute_sha256() + assert progress.phase is ClientProgressPhase.COMPLETE + + +@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() + + +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 / "batch.parquet").is_file() + attempt_dir = resume_path.parent / "attempts" / "attempt-0002" + overlay_path = attempt_dir / "client-env" / "site-packages" + overlay_path.mkdir(parents=True) + prepared = replace( + client_worker_case.prepared, + attempt_id="attempt-0002", + attempt_dir=attempt_dir, + 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"}] + 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["image_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, + overlay_path=Path(sys.argv[3]), + dependency_lock=ArtifactReference(**plan["client"]["dependency_lock"]), + client_image_sha256=plan["client"]["image"]["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:17000/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") +""" + 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(), + ], + 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/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index 5b511c4c4..90009a9b4 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -156,6 +156,13 @@ def test_if_possible_uses_effective_resume_dataset_location( 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, 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 index 2dc7726b3..67ce6aae3 100644 --- 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 @@ -13,6 +13,19 @@ class FakePluginConfig(SingleColumnConfig): """Minimal custom column configuration.""" column_type: Literal["fake-slurm-column"] = "fake-slurm-column" + model_alias: str = "generator" + judge_model_alias: str = "judge" + + @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 [self.model_alias, self.judge_model_alias] class FakePluginImplementation: From a3afe2059f85dd90eb565d8806f6a073fab877cf Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 3 Sep 2026 16:14:43 -0600 Subject: [PATCH 18/62] feat: finalize Slurm shard winners (#910) * feat: finalize Slurm shard winners Lease shard dataset workspaces under reviewed locks. Validate client and candidate records plus output file digests. Atomically publish one immutable winner and seal retries. Part of #869 Signed-off-by: Nabin Mulepati * fix: unwind interrupted dataset locks Ensure an opened shard context closes when resume-lock acquisition is interrupted by any BaseException. Signed-off-by: Nabin Mulepati * fix: rebind verified candidate files Reject candidate path replacement after digesting an opened output descriptor and before winner publication. Signed-off-by: Nabin Mulepati * test: cover nested candidate paths Exercise valid nested output verification and rejection of a symlinked intermediate directory. Signed-off-by: Nabin Mulepati * fix: validate winner retry timestamps Reject non-UTC publication timestamps before the existing-winner idempotency path can ignore them. Signed-off-by: Nabin Mulepati * fix: classify invalid winners as corruption Normalize semantically invalid persisted winner chains as corruption across load, finalize, attempt, and workspace APIs. Signed-off-by: Nabin Mulepati * refactor: compose shard winner finalization Signed-off-by: Nabin Mulepati * refactor: compose finalization validation Signed-off-by: Nabin Mulepati * fix: harden candidate publication boundaries Signed-off-by: Nabin Mulepati * fix: bind finalization publication Signed-off-by: Nabin Mulepati * fix: publish allocation results before success Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/integration.py | 179 +-- .../data_designer/slurm/runtime/controller.py | 15 +- .../data_designer/slurm/runtime/records.py | 5 +- .../src/data_designer/slurm/state/__init__.py | 6 + .../data_designer/slurm/state/artifacts.py | 218 ++++ .../data_designer/slurm/state/filesystem.py | 129 ++ .../data_designer/slurm/state/finalization.py | 332 ++++++ .../slurm/state/finalization_validation.py | 244 ++++ .../src/data_designer/slurm/state/outputs.py | 6 +- .../slurm/state/plan_validation.py | 58 + .../src/data_designer/slurm/state/results.py | 86 ++ .../src/data_designer/slurm/state/storage.py | 100 +- .../src/data_designer/slurm/state/store.py | 64 +- .../golden/finalization_chain.json | 8 +- .../test_integration_validation.py | 14 +- .../tests/runtime/conftest.py | 16 + .../tests/runtime/test_controller.py | 64 +- .../tests/state/test_model_edges.py | 14 + .../tests/state/test_store.py | 1033 ++++++++++++++++- 19 files changed, 2397 insertions(+), 194 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/finalization_validation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/results.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/integration.py b/packages/data-designer-slurm/src/data_designer/slurm/integration.py index cde400e29..f19f4a04e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/integration.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/integration.py @@ -1,192 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Pure validation across Slurm execution-plan and runtime-state contracts.""" +"""Public integration validation across Slurm plan and runtime-state contracts.""" from __future__ import annotations -import posixpath - -from data_designer.slurm.client import ClientOutcome, ClientResult -from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.planning import PlannedShard -from data_designer.slurm.state.execution import ( - AttemptLifecycleState, - AttemptManifest, - AttemptTerminalClassification, -) -from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError IntegrationContractError = PlanStateContractError class PlanStateValidator(PersistedPlanStateValidator): - """Validate records against one resolved plan with reusable derived state. - - This is an in-process validation service, not a persisted contract record. - """ - - 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_client_candidate(planned_shard, attempt, client_result, candidate) - _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", - ) - self._validate_winner(planned_shard, attempt, client_result, winner) - _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") - return winner - - def validate_client_candidate( - self, - planned_shard: PlannedShard, - attempt: AttemptManifest, - client_result: ClientResult, - candidate: CandidateOutputManifest, - ) -> CandidateOutputManifest: - """Validate a complete client result and candidate before terminal publication.""" - self.validate_planned_attempt(planned_shard, attempt) - _require(client_result.outcome is ClientOutcome.COMPLETE, "only complete client results may be finalized") - _require(candidate.winner_eligible, "only complete candidate outputs may be finalized") - self._validate_candidate_identities(planned_shard, attempt, client_result, candidate) - self._validate_candidate_counts(planned_shard, client_result, candidate) - self._validate_candidate_location(planned_shard, attempt, client_result, candidate) - _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", - ) - return candidate - - def _validate_candidate_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 attempt ordinal does not match attempt" - ) - - @staticmethod - def _validate_candidate_counts( - 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 == planned_shard.requested_records, - "client and candidate actual records must complete the planned shard", - ) - - def _validate_candidate_location( - self, - planned_shard: PlannedShard, - attempt: AttemptManifest, - client_result: ClientResult, - candidate: CandidateOutputManifest, - ) -> None: - _require( - client_result.requested_resume_mode == self.plan.invocation.authored.resume, - "client requested resume mode does not match the resolved plan", - ) - expected_dataset_path = self._get_expected_dataset_path(planned_shard, attempt, client_result) - _require( - client_result.dataset_path == expected_dataset_path, "client dataset path does not match planned intent" - ) - _require( - candidate.dataset_path == expected_dataset_path, "candidate dataset path does not match planned intent" - ) - candidate_reference = self._get_candidate_reference(client_result) - expected_manifest_path = posixpath.join( - posixpath.dirname(planned_shard.resume_workspace.path), - "attempts", - attempt.attempt_id, - "output-manifest.json", - ) - _require( - candidate_reference.path == expected_manifest_path, "candidate manifest path does not match planned intent" - ) - _require( - candidate_reference.sha256 == candidate.compute_sha256(), - "client candidate digest does not match the manifest", - ) - - def _validate_winner( - self, - planned_shard: PlannedShard, - attempt: AttemptManifest, - client_result: ClientResult, - winner: ShardWinner, - ) -> None: - candidate_reference = self._get_candidate_reference(client_result) - _require(winner.run_id == self.plan.run_id, "winner run_id does not match the resolved plan") - _require(winner.shard_id == planned_shard.shard_id, "winner shard_id does not match the planned shard") - _require(winner.attempt_id == attempt.attempt_id, "winner attempt_id does not match the attempt") - _require(winner.attempt_ordinal == attempt.attempt_ordinal, "winner attempt ordinal does not match the attempt") - _require( - attempt.candidate_output == candidate_reference, - "attempt candidate reference does not match client result", - ) - _require( - winner.candidate_manifest == candidate_reference, - "winner candidate reference does not match client result", - ) - - @staticmethod - def _get_candidate_reference(client_result: ClientResult) -> ArtifactReference: - candidate_reference = client_result.candidate_output_manifest - _require(candidate_reference is not None, "complete client result has no candidate manifest reference") - return candidate_reference - - @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 IntegrationContractError(message) + """Public compatibility facade for persisted-plan validation.""" __all__ = [ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 09a4a4cc2..2341d617b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import 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 @@ -32,6 +33,7 @@ AttemptManifest, AttemptReadiness, AttemptTerminalClassification, + CandidateOutputManifest, DeploymentReadiness, EndpointPublicationState, ProbeEvidence, @@ -55,6 +57,14 @@ def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: """Persist a monotonic attempt update.""" ... + 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 write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: """Persist the exact next readiness revision.""" ... @@ -290,7 +300,7 @@ def _run_client( 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_complete_client_candidate(self._context) + client_result, candidate = load_complete_client_candidate(self._context, self._attempt) 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 - the record contract requires this for complete results @@ -298,6 +308,8 @@ def _run_client( SlurmRuntimeErrorCode.FINALIZATION_FAILED, "complete 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 def _validate_client_timestamps( @@ -540,7 +552,6 @@ def _build_terminal_attempt( self._attempt, state=AttemptLifecycleState.FAILED, terminal_classification=AttemptTerminalClassification.FAILED, - candidate_output=None, updated_at=timestamp, ) return terminal 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 index 1644fe1bf..c41454609 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -15,7 +15,7 @@ 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 CandidateOutputManifest +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" @@ -26,6 +26,7 @@ def load_complete_client_candidate( context: AllocationContext, + attempt: AttemptManifest, ) -> tuple[ClientResult, CandidateOutputManifest]: """Load a complete semantic client result and its digest-bound candidate.""" client_result = _read_record(context.attempt_directory, _CLIENT_RESULT_NAME, ClientResult) @@ -38,7 +39,7 @@ def load_complete_client_candidate( try: PlanStateValidator(context.plan).validate_client_candidate( context.shard, - context.attempt, + attempt, client_result, candidate, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index b4717289e..ce69166e3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -19,6 +19,7 @@ Sha256Digest, ShardId, ) +from data_designer.slurm.state.artifacts import compute_candidate_schema_digest from data_designer.slurm.state.base import ( SchedulerIdentity, StateRecord, @@ -38,6 +39,8 @@ ShardManifest, ) from data_designer.slurm.state.outputs import ( + CANDIDATE_OUTPUT_FORMAT, + MAXIMUM_CANDIDATE_OUTPUT_FILES, CandidateOutcome, CandidateOutputFile, CandidateOutputManifest, @@ -91,8 +94,11 @@ "AttemptReadiness", "AttemptTerminalClassification", "CandidateOutcome", + "CANDIDATE_OUTPUT_FORMAT", + "MAXIMUM_CANDIDATE_OUTPUT_FILES", "CandidateOutputFile", "CandidateOutputManifest", + "compute_candidate_schema_digest", "CollectionPlan", "CollectionShard", "ContractRecord", 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..3a3b6b466 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -0,0 +1,218 @@ +# 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 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 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() + + +class CandidateArtifactVerifier: + """Verify one manifest-bounded candidate and lease its files through publication.""" + + @contextmanager + def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandidateArtifacts]: + dataset_path = Path(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 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 + + +@contextmanager +def _open_parent_directory( + dataset_descriptor: int, + dataset_path: Path, + parts: tuple[str, ...], +) -> Iterator[tuple[int, Path]]: + parent_descriptor = os.dup(dataset_descriptor) + parent_path = dataset_path + 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 + yield parent_descriptor, parent_path + 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__ = [ + "CandidateArtifactVerifier", + "CandidateSchema", + "SerializedCandidateSchema", + "VerifiedCandidateArtifacts", + "compute_candidate_schema_digest", +] 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 index b771da750..9394889ae 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/filesystem.py @@ -5,6 +5,7 @@ from __future__ import annotations +import hashlib import os import stat from collections.abc import Iterator @@ -137,6 +138,129 @@ def read_regular_text( 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, @@ -299,3 +423,8 @@ def _repair_interrupted_publication( 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..474faaeea --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -0,0 +1,332 @@ +# 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, 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_shard_winner + + +@dataclass(frozen=True, slots=True) +class _WinnerResolution: + winner: ShardWinner + candidate: CandidateOutputManifest + plan: ResolvedSlurmRunPlan + 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() + + @contextmanager + def acquire_dataset_workspace( + self, + shard_id: ShardId, + attempt_id: AttemptId, + effective_resume_mode: Literal["never", "always"], + ) -> Iterator[Path]: + with self._storage.acquire_resume_lock(shard_id): + dataset_path = self._prepare_workspace_or_normalize(shard_id, attempt_id, effective_resume_mode) + yield dataset_path + + def finalize_winner(self, shard_id: ShardId, attempt_id: AttemptId, published_at: datetime) -> ShardWinner: + try: + with self._storage.acquire_resume_lock(shard_id): + return self._finalize_with_dataset_lease(shard_id, attempt_id, published_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, + ) -> ShardWinner: + resolution = self._resolve_under_state_locks(shard_id, attempt_id, published_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, 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, + expected: _WinnerResolution, + artifacts: VerifiedCandidateArtifacts, + ) -> ShardWinner: + with self._storage.acquire_shard_lock(shard_id): + current = self._resolve_winner(shard_id, attempt_id, published_at) + if current.already_published: + return current.winner + if current.winner != expected.winner or current.candidate != expected.candidate: + 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 + self._storage.publish_winner(expected.winner) + return expected.winner + + def _resolve_under_state_locks( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + ) -> _WinnerResolution: + with self._storage.acquire_shard_lock(shard_id): + return self._resolve_winner(shard_id, attempt_id, published_at) + + def _resolve_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + published_at: datetime, + ) -> _WinnerResolution: + 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) + existing = self.load_optional_winner(run, plan, shard, attempts) + client_result, candidate = self._load_finalization_records(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, True) + winner = self._build_winner(run, shard, attempt, client_result, published_at) + self._validate_finalization_chain(run, plan, shard, attempt, client_result, candidate, winner) + return _WinnerResolution(winner, candidate, plan, False) + + def _prepare_dataset_workspace( + self, + shard_id: ShardId, + attempt_id: AttemptId, + effective_resume_mode: Literal["never", "always"], + ) -> 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, effective_resume_mode) + dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, effective_resume_mode) + expected_path = ( + Path(shard.resume_workspace.path) + if effective_resume_mode == "always" + else self._storage.get_attempt_path(shard_id, 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, + effective_resume_mode: Literal["never", "always"], + ) -> Path: + try: + return self._prepare_dataset_workspace(shard_id, attempt_id, effective_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("complete 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, + effective_resume_mode: Literal["never", "always"], + ) -> 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 effective_resume_mode != requested: + raise StateContractError("effective 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..c5b7b6bfa --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization_validation.py @@ -0,0 +1,244 @@ +# 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 is ClientOutcome.COMPLETE, "only complete client results may be finalized") + _require(candidate.winner_eligible, "only complete candidate outputs 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 == planned_shard.requested_records, + "client and candidate actual records must complete the planned shard", + ) + 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, "complete 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, "complete 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, "complete 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/outputs.py b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py index 756e00b44..d332f7769 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -24,6 +24,10 @@ 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" @@ -54,7 +58,7 @@ class CandidateOutputManifest(StateRecord): requested_records: PositiveInt actual_records: NonNegativeInt outcome: CandidateOutcome - files: tuple[CandidateOutputFile, ...] + files: tuple[CandidateOutputFile, ...] = Field(max_length=MAXIMUM_CANDIDATE_OUTPUT_FILES) dataset_schema_digest: Sha256Digest provenance_digest: Sha256Digest 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 index 6b40340f0..1c1480641 100644 --- 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 @@ -9,9 +9,12 @@ 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 @@ -33,6 +36,7 @@ def __init__(self, plan: ResolvedSlurmRunPlan) -> None: 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: @@ -137,6 +141,60 @@ def validate_planned_attempt( ) 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: 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..1e890cb05 --- /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("complete 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/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 94698bf63..618c156ba 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -9,16 +9,17 @@ import os import re from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from pathlib import Path -from typing import TypeVar +from typing import 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 from data_designer.slurm.planning import ResolvedSlurmRunPlan -from data_designer.slurm.state.errors import StateCorruptionError, StateNotFoundError +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, @@ -31,6 +32,7 @@ replace_text, sync_directory, ) +from data_designer.slurm.state.outputs import CandidateOutputManifest, ShardWinner from data_designer.slurm.state.readiness import AttemptReadiness _RUN_FILENAME = "run.json" @@ -42,6 +44,11 @@ _SHARD_LOCK_FILENAME = "shard.lock" _ATTEMPT_FILENAME = "attempt.json" _READINESS_FILENAME = "readiness.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" _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _ATTEMPT_NAME_PATTERN = re.compile(r"^attempt-[0-9]{4,}$") @@ -75,6 +82,9 @@ def get_attempt_path(self, shard_id: str, attempt_id: str) -> Path: 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_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) @@ -109,6 +119,31 @@ def acquire_shard_lock(self, shard_id: ShardId) -> Iterator[None]: 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, @@ -239,6 +274,61 @@ def sync_attempt_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> No 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) + return dataset_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) + return dataset_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, @@ -404,7 +494,9 @@ def _get_record_path(self, name: str, record: ContractRecord) -> Path: return self.run_root / name if isinstance(record, ShardManifest): return self.get_shard_path(record.shard_id) / name - if isinstance(record, (AttemptManifest, AttemptReadiness)): + 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__}") 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 index 9f106f080..95766a8e1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -5,10 +5,15 @@ from __future__ import annotations +from collections.abc import 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 @@ -19,10 +24,13 @@ StateNotFoundError, ) from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +from data_designer.slurm.state.finalization import WinnerFinalizer +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.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.storage import StateStorage from data_designer.slurm.state.validation import ( StateContractError, @@ -57,6 +65,8 @@ def __init__(self, workspace_root: str | Path, run_id: Identifier) -> None: raise SlurmStateError("invalid persisted run location") from error self._storage = StateStorage(Path(normalized_root), normalized_run_id) self._reader = StateReader(self._storage, normalized_run_id) + self._results = AttemptResultPublisher(self._storage, self._reader) + self._finalizer = WinnerFinalizer(self._storage, self._reader) self._run_id = normalized_run_id @property @@ -154,8 +164,57 @@ def load_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptRea 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, + effective_resume_mode: Literal["never", "always"], + ) -> 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(effective_resume_mode) is not str or effective_resume_mode not in {"never", "always"}: + raise StateConflictError("effective resume mode must be 'never' or 'always'") + with self._finalizer.acquire_dataset_workspace( + normalized_shard_id, + normalized_attempt_id, + effective_resume_mode, + ) as dataset_path: + yield dataset_path + + def finalize_winner( + self, + shard_id: ShardId, + attempt_id: AttemptId, + *, + published_at: datetime, + ) -> ShardWinner: + """Validate attempt-local artifacts and publish one immutable 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") + return self._finalizer.finalize_winner(normalized_shard_id, normalized_attempt_id, published_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 _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: - with self._storage.acquire_shard_lock(attempt.shard_id): + 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) @@ -164,6 +223,7 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes 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") @@ -181,6 +241,8 @@ def _update_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes 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 diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index c0e9887b2..e9933b74c 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,7 +4,7 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + "provenance_digest": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "f2d8428ea0c36860cb835af1a2b602058c1f4b668d081c525258f5857b400398" + "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index 2a7a7d834..9f81f0226 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -391,6 +391,8 @@ def test_client_candidate_validation_supports_preterminal_runtime_state(records: ("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"), ], @@ -441,6 +443,12 @@ def test_finalization_rejects_invalid_chains( 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: @@ -602,7 +610,11 @@ def test_finalization_accepts_effective_resume_always(records: IntegrationRecord 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} + 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 diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index f9c405e80..9dc83da6d 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -11,6 +11,7 @@ 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 @@ -19,6 +20,7 @@ AttemptLifecycleState, AttemptManifest, AttemptReadiness, + CandidateOutputManifest, SchedulerIdentity, StateNotFoundError, validate_attempt_transition, @@ -43,6 +45,20 @@ def update_attempt(self, attempt: AttemptManifest) -> AttemptManifest: self.attempt = attempt return attempt + def publish_attempt_result( + self, + client_result: ClientResult, + candidate: CandidateOutputManifest, + ) -> tuple[ClientResult, CandidateOutputManifest]: + 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 + def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: if self.readiness: validate_readiness_transition(self.readiness[-1], readiness) diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index c302ce8d7..de4b363c5 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -13,6 +13,7 @@ from slurm_test_fakes import FakeClock 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.runtime.controller import OneNodeAllocationController from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode @@ -29,6 +30,9 @@ DeploymentReadiness, EndpointPublicationState, ReadinessState, + RunManifest, + ShardManifest, + SlurmStateWriter, ) @@ -178,6 +182,33 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c 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={}, + ) + + 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 + + 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) @@ -426,6 +457,7 @@ def test_cleanup_failure_prevents_false_success(runtime_case: RuntimeCase) -> No 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) @@ -586,7 +618,7 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: ), ), dataset_schema_digest="b" * 64, - provenance_digest="c" * 64, + provenance_digest=context.plan.compute_sha256(), ) candidate_path = context.attempt_directory / "output-manifest.json" _write_record(candidate_path, candidate.serialize_json()) @@ -612,6 +644,36 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: _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" diff --git a/packages/data-designer-slurm/tests/state/test_model_edges.py b/packages/data-designer-slurm/tests/state/test_model_edges.py index 0babd1f76..405102cc5 100644 --- a/packages/data-designer-slurm/tests/state/test_model_edges.py +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -11,6 +11,7 @@ from pydantic import ValidationError from data_designer.slurm.state import ( + MAXIMUM_CANDIDATE_OUTPUT_FILES, ArtifactReference, AttemptManifest, AttemptReadiness, @@ -120,6 +121,19 @@ def test_candidate_output_rejects_remaining_invalid_inventory_forms() -> None: 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") diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index fcd1c84ac..8557b113e 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -3,21 +3,26 @@ 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 +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 DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, ContractValue, compute_canonical_json_sha256, pretty_json from data_designer.slurm.planning import ResolvedSlurmRunPlan @@ -26,6 +31,10 @@ AttemptLifecycleState, AttemptManifest, AttemptReadiness, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, DeploymentReadiness, EndpointPublicationState, ReadinessState, @@ -33,13 +42,17 @@ 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) @@ -56,6 +69,49 @@ class _StateCase: 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, @@ -1102,6 +1158,870 @@ def publish(seconds: int) -> AttemptReadiness: 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_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, + ) + + +@pytest.mark.parametrize("export_format", ("jsonl", "csv")) +def test_finalization_keeps_candidate_parquet_separate_from_export_format( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + export_format: str, +) -> None: + authored = authored_run_single.model_copy( + update={"output": authored_run_single.output.model_copy(update={"format": export_format})} + ) + plan = single_node_plan.model_copy( + update={ + "authored_config": single_node_plan.authored_config.model_copy( + update={"sha256": authored.compute_sha256()} + ), + "output": single_node_plan.output.model_copy(update={"format": export_format}), + } + ) + case = _initialized_case(tmp_path, authored, 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.attempt_id == finalization.attempt.attempt_id + assert finalization.candidate.files[0].relative_path.endswith(".parquet") + + +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) + + +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_partial_client_results_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_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, @@ -1278,6 +2198,117 @@ def _readiness(case: _StateCase, attempt: AttemptManifest) -> AttemptReadiness: ) +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, + 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() + 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=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=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) From 88ff321dfa82d35e001f6374178920bafd4bd272 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:46:34 -0300 Subject: [PATCH 19/62] feat(slurm): wire M2 services and CLI (#926) * feat(slurm): wire M2 services and CLI Add package-owned execution, status, cancellation, and image service wiring. Resolve client dependencies and expose the M2 operations through the Slurm CLI. Part of #874 Signed-off-by: Andre Manoel * fix(slurm): reject swapped SQSH file types Signed-off-by: Andre Manoel * fix(slurm): clarify unavailable operations Signed-off-by: Andre Manoel * fix(slurm): resolve review findings --------- Signed-off-by: Andre Manoel --- packages/data-designer-slurm/pyproject.toml | 1 + .../src/data_designer/slurm/cli.py | 202 ++++++ .../slurm/client/dependencies.py | 351 +++++++++++ .../data_designer/slurm/config/__init__.py | 2 + .../src/data_designer/slurm/config/loading.py | 29 +- .../src/data_designer/slurm/images/service.py | 5 +- .../data_designer/slurm/services/__init__.py | 52 +- .../data_designer/slurm/services/errors.py | 7 + .../data_designer/slurm/services/images.py | 109 +++- .../data_designer/slurm/services/results.py | 113 ++++ .../src/data_designer/slurm/services/run.py | 133 +++- .../data_designer/slurm/services/wiring.py | 574 ++++++++++++++++++ .../client/test_dependency_resolution.py | 234 +++++++ .../tests/config/test_loading_builder.py | 18 + .../tests/services/test_services.py | 57 ++ .../tests/services/test_wiring.py | 457 ++++++++++++++ .../tests/slurm_test_fakes/services.py | 4 +- .../data-designer-slurm/tests/test_cli.py | 147 +++++ scripts/test_slurm_package_install.py | 2 + uv.lock | 11 + 20 files changed, 2493 insertions(+), 15 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/results.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py create mode 100644 packages/data-designer-slurm/tests/client/test_dependency_resolution.py create mode 100644 packages/data-designer-slurm/tests/services/test_wiring.py create mode 100644 packages/data-designer-slurm/tests/test_cli.py diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 9e4610eff..cc8c33046 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -38,6 +38,7 @@ bump = true dependencies = [ "data-designer=={{ version }}", "packaging>=25,<27", + "pip>=25,<27", "pydantic>=2.9.2,<3", "pyyaml>=6.0.1,<7", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index a77cc462b..7315f4d3b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -3,14 +3,41 @@ from __future__ import annotations +import re +from collections.abc import Callable +from pathlib import Path +from typing import NoReturn, TypeVar + import click import typer +from pydantic import BaseModel, ValidationError + +from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config +from data_designer.slurm.contracts import canonical_json +from data_designer.slurm.services import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, + create_slurm_image_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, +} 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) +app.add_typer(image_app, name="image") @app.callback() @@ -18,6 +45,181 @@ def slurm_callback() -> None: pass +@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"), + force: bool = typer.Option(False, "--force"), +) -> None: + """Prepare and submit one authored run.""" + operation = SlurmServiceOperation.EXECUTE_RUN + + def execute() -> BaseModel: + config = load_run_config(run_file) + service = create_slurm_run_service(profile_file=profile_file, cluster=cluster) + return service.execute(config, source_root=run_file.resolve().parent, dry_run=dry_run, force=force) + + _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: + """Show persisted M2 run status without scheduler reconciliation.""" + 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 job cancellation; status changes after reconciliation.""" + operation = SlurmServiceOperation.CANCEL_RUN + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).cancel(run_id), + ) + _emit_result(result) + + +@image_app.command("add") +def image_add_command( + source: str = typer.Argument(...), + kind: str = typer.Option(..., "--kind", help="Image role: client or serving"), + name: str | None = typer.Option(None, "--name"), + replace: bool = typer.Option(False, "--replace"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Request image import or inspection; requires IMG lifecycle support.""" + operation = SlurmServiceOperation.ADD_IMAGE + + def add() -> BaseModel: + if not source.endswith(".sqsh") and re.fullmatch(r"[^\s]+@sha256:[0-9a-f]{64}", source) is None: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "OCI image source must be digest-qualified as name@sha256:", + ) + request = ImageBuildRequest(name=name or _derive_image_name(source), kind=kind, source=source) + return create_slurm_image_service(profile_file=profile_file, cluster=cluster).add(request, replace=replace) + + _emit_result(_invoke(operation, add)) + + +@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(), + ) + _emit_json([image.model_dump(mode="json") for image in images]) + + +@image_app.command("info") +def image_info_command( + name: str = typer.Argument(...), + 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_result(result) + + +@image_app.command("rm") +def image_remove_command( + name: str = typer.Argument(...), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), +) -> None: + """Unregister one alias without deleting its SQSH artifact.""" + operation = SlurmServiceOperation.REMOVE_IMAGE + result = _invoke( + operation, + lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).remove(name), + ) + _emit_result(result) + + +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: + _emit_json( + { + "error": { + "code": error.code.value, + "message": str(error), + "operation": error.operation.value, + } + }, + err=True, + ) + raise typer.Exit(_EXIT_CODES[error.code]) + + +def _emit_result(result: BaseModel) -> None: + _emit_json(result.model_dump(mode="json")) + + +def _emit_json(value: object, *, err: bool = False) -> None: + typer.echo(canonical_json(value).decode("utf-8"), err=err) + + +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].split(":", maxsplit=1)[0] + 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]}..." + + 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/client/dependencies.py b/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py new file mode 100644 index 000000000..45b7bc39e --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py @@ -0,0 +1,351 @@ +# 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, ClientImageInspection +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution, validate_absolute_path +from data_designer.slurm.planning import LockedPackage, ResolvedDependencyLock, ResolvedImage + +_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 client-image 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_image: ResolvedImage, + *, + 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") + inspection = client_image.inspection_facts + if not isinstance(inspection, ClientImageInspection): + raise ClientDependencyResolutionError("client image lacks dependency inspection facts") + 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_image, + normalized_run_root, + normalized_source_root, + ) + else: + result = self._resolve_requirements( + dependencies, + client_image, + 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_image: ResolvedImage, + run_root: Path, + temporary_root: Path, + ) -> ResolvedClientDependencies: + requirements = dependencies.requirements + if requirements is None: + raise ClientDependencyResolutionError("client dependency requirements are unavailable") + inspection = client_image.inspection_facts + if not isinstance(inspection, ClientImageInspection): + raise ClientDependencyResolutionError("client image lacks dependency inspection facts") + image_distributions = tuple(sorted(inspection.distributions, key=lambda item: item.name)) + if not requirements: + return ResolvedClientDependencies( + lock=_build_lock(client_image, image_distributions=image_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 image_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={inspection.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, image_distributions) + try: + lock = _build_lock( + client_image, + image_distributions=image_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_image: ResolvedImage, + 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 + inspection = client_image.inspection_facts + if not isinstance(inspection, ClientImageInspection): + raise ClientDependencyResolutionError("client image lacks dependency inspection facts") + if ( + supplied.client_image_sha256 != client_image.sha256 + or supplied.python_abi != inspection.python_abi + or supplied.image_distributions != tuple(sorted(inspection.distributions, key=lambda item: item.name)) + ): + raise ClientDependencyResolutionError("client dependency lock targets another image") + + 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_image_sha256=supplied.client_image_sha256, + authored_requirements=supplied.authored_requirements, + authored_source=lock_file, + source=source_reference, + image_distributions=supplied.image_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_image: ResolvedImage, + *, + image_distributions: tuple[InstalledDistribution, ...], + authored_requirements: tuple[str, ...], + overlay_packages: tuple[LockedPackage, ...] = (), +) -> ResolvedDependencyLock: + inspection = client_image.inspection_facts + if not isinstance(inspection, ClientImageInspection): + raise ClientDependencyResolutionError("client image lacks dependency inspection facts") + return ResolvedDependencyLock( + schema_version=1, + resolver_version=_RESOLVER_VERSION, + python_abi=inspection.python_abi, + client_image_sha256=client_image.sha256, + authored_requirements=authored_requirements, + image_distributions=image_distributions, + overlay_packages=overlay_packages, + ) + + +def _load_downloaded_wheels( + wheelhouse: Path, + run_root: Path, + image_distributions: tuple[InstalledDistribution, ...], +) -> tuple[tuple[LockedPackage, ...], tuple[Path, ...]]: + image = {item.name: item.version for item in image_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 image: + if image[normalized_name] != normalized_version: + raise ClientDependencyResolutionError("client dependency conflicts with the selected image") + 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/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 6d329c23d..c08456652 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -29,6 +29,7 @@ from data_designer.slurm.config.loading import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + load_builder_payload, load_profile_catalog, load_run_config, resolve_profile, @@ -113,6 +114,7 @@ "SubmissionConfig", "VllmServerConfig", "injected_profile", + "load_builder_payload", "load_profile_catalog", "load_run_config", "resolve_profile", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 684465871..5c43b9f64 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -13,7 +13,7 @@ from typing import TypeVar, cast import yaml -from pydantic import ValidationError +from pydantic import JsonValue, ValidationError from yaml.nodes import MappingNode from data_designer.slurm._errors import format_parse_error, format_validation_error @@ -25,7 +25,7 @@ injected_profile, select_profile, ) -from data_designer.slurm.config.run import DataDesignerSlurmConfig +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" @@ -74,6 +74,31 @@ def load_profile_catalog(path: str | Path) -> SlurmProfileCatalog: 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 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, 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 index 607f7f2e0..111f85201 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -308,7 +308,10 @@ def _open_verified_sqsh( 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): + 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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 07c832814..22b603c7a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -5,24 +5,72 @@ 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 SlurmImageResolver, SlurmImageService -from data_designer.slurm.services.run import SlurmBatchScriptRenderer, SlurmRunPlanner, SlurmRunService +from data_designer.slurm.services.images import SlurmImageManager, SlurmImageResolver, SlurmImageService +from data_designer.slurm.services.results import ( + SlurmPersistedAttemptStatus, + SlurmPersistedRunStatus, + SlurmPersistedShardStatus, + 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_image_service, + create_slurm_run_service, + ) + +_LAZY_IMPORTS = { + "SlurmRunArtifactPublisher": ("data_designer.slurm.services.wiring", "SlurmRunArtifactPublisher"), + "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", + "SlurmImageManager", "SlurmImageResolver", "SlurmImageService", + "SlurmPersistedAttemptStatus", + "SlurmPersistedRunStatus", + "SlurmPersistedShardStatus", + "SlurmRunArtifactPublisher", + "SlurmRunBackend", + "SlurmRunCancellation", + "SlurmRunExecution", "SlurmRunPlanner", "SlurmRunService", "SlurmServiceError", "SlurmServiceErrorCode", "SlurmServiceOperation", + "create_slurm_image_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/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py index 1a24686ae..aa08fffed 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -29,7 +29,14 @@ class SlurmServiceOperation(str, Enum): PLAN_RUN = "plan_run" RENDER_ATTEMPT = "render_attempt" + EXECUTE_RUN = "execute_run" + STATUS_RUN = "status_run" + CANCEL_RUN = "cancel_run" 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" 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 index 848eaeced..9a2a12253 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -7,14 +7,23 @@ 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 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.""" @@ -26,18 +35,35 @@ def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedI """ +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) -> RegisteredImage: + """Unregister one alias without deleting its 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``. + resolver: Image-resolution dependency implementing ``SlurmImageResolver``. + manager: Optional image-operation dependency. Package-owned factories provide it. """ - def __init__(self, resolver: SlurmImageResolver) -> None: + 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. @@ -64,3 +90,80 @@ def resolve_image() -> ResolvedImage: 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") + 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) -> RegisteredImage: + """Unregister one alias without deleting its SQSH artifact.""" + operation = SlurmServiceOperation.REMOVE_IMAGE + normalized_name = _validate_name(name, operation) + manager = self._require_manager(operation) + + def remove_image() -> RegisteredImage: + image = manager.remove(normalized_name) + 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/results.py b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py new file mode 100644 index 000000000..51950bdcf --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -0,0 +1,113 @@ +# 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, model_validator + +from data_designer.slurm.contracts import ContractValue, Identifier, Sha256Digest +from data_designer.slurm.state import ( + AttemptManifest, + AttemptReadiness, + 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 available without 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): + """M2 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 + + +__all__ = [ + "SlurmPersistedAttemptStatus", + "SlurmPersistedRunStatus", + "SlurmPersistedShardStatus", + "SlurmRunCancellation", + "SlurmRunExecution", +] 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 index b2aaaf04e..9790b877b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -5,21 +5,34 @@ from __future__ import annotations +from pathlib import Path from typing import Protocol +from pydantic import TypeAdapter, ValidationError + from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import Identifier 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 ( + SlurmPersistedRunStatus, + SlurmRunCancellation, + SlurmRunExecution, +) + +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) class SlurmRunPlanner(Protocol): """Resolve authored run configs through a supported service dependency.""" - def plan(self, config: DataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: + 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. @@ -36,6 +49,26 @@ def __call__(self, plan: ResolvedSlurmRunPlan, *, attempt_ordinal: int) -> str: """ +class SlurmRunBackend(Protocol): + """Execute and inspect runs through a supported service dependency.""" + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + force: bool, + ) -> SlurmRunExecution: + """Render or submit one run.""" + + def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: + """Return the durable M2 status for one run.""" + + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: + """Request cancellation of active jobs without changing persisted state.""" + + class SlurmRunService: """Coordinate public run operations through package-owned boundaries. @@ -44,15 +77,26 @@ class SlurmRunService: Args: planner: Run-planning dependency implementing ``SlurmRunPlanner``. - renderer: Batch-rendering dependency implementing - ``SlurmBatchScriptRenderer``. + renderer: Batch-rendering dependency implementing ``SlurmBatchScriptRenderer``. + backend: Optional run-operation dependency. Package-owned factories provide it. """ - def __init__(self, planner: SlurmRunPlanner, renderer: SlurmBatchScriptRenderer) -> None: + 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) -> ResolvedSlurmRunPlan: + 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. @@ -63,9 +107,11 @@ def plan(self, config: DataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: 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) + 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(): @@ -100,3 +146,78 @@ def render() -> str: return script return _invoke_service_backend(operation, render) + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: str | Path = ".", + dry_run: bool = False, + force: 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 or type(force) is not bool: + raise _make_invalid_request_error(operation, "dry_run and force must be booleans") + backend = self._require_backend(operation) + + def execute_run() -> SlurmRunExecution: + result = backend.execute( + config, + source_root=Path(source_root).expanduser().resolve(), + dry_run=dry_run, + force=force, + ) + 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: + """Return persisted M2 records without scheduler reconciliation.""" + 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 _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 _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 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..ddda16ea7 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -0,0 +1,574 @@ +# 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 +from collections.abc import Callable, Iterator +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Protocol, TypeVar +from uuid import uuid4 + +from pydantic import JsonValue + +from data_designer.slurm.client.dependencies import ( + ClientDependencyResolutionError, + ClientDependencyResolver, + ResolvedClientDependencies, +) +from data_designer.slurm.config import ( + DataDesignerSlurmConfig, + ImageBuildRequest, + ImageKind, + ImageRef, + SelectedSlurmProfile, + SlurmConfigLoadError, + SlurmProfile, + SlurmProfileCatalog, + load_builder_payload, + resolve_profile, +) +from data_designer.slurm.contracts import Identifier +from data_designer.slurm.images.errors import 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 +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +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.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation +from data_designer.slurm.services.images import SlurmImageService +from data_designer.slurm.services.results import ( + SlurmPersistedAttemptStatus, + SlurmPersistedRunStatus, + SlurmPersistedShardStatus, + SlurmRunCancellation, + SlurmRunExecution, +) +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, + SlurmStateError, + SlurmStateWriter, + StateConflictError, + StateNotFoundError, +) + +RunIdFactory = Callable[[], str] +Clock = Callable[[], datetime] +_ResultT = TypeVar("_ResultT") +_MAX_VISIBLE_JOB_IDS = 16 + + +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.""" + + +@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, + ) -> 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 + + @contextmanager + def prepare( + self, + authored: DataDesignerSlurmConfig, + *, + source_root: Path, + operation: SlurmServiceOperation, + ) -> Iterator[_PreparedRun]: + stack = ExitStack() + try: + prepared = self._prepare(authored, source_root=source_root, operation=operation, stack=stack) + except BaseException: + stack.close() + raise + try: + yield prepared + finally: + stack.close() + + def _prepare( + self, + authored: DataDesignerSlurmConfig, + *, + source_root: Path, + operation: SlurmServiceOperation, + stack: ExitStack, + ) -> _PreparedRun: + try: + run_id = self._run_id_factory() + workspace_root = self._profile.profile.workspace_root + run_root = Path(workspace_root) / "runs" / run_id + client_image = self._images.resolve_for_planning(authored.client.image, expected_kind=ImageKind.CLIENT) + deployment_images = tuple( + self._images.resolve_for_planning(deployment.server.image, expected_kind=ImageKind.SERVING) + for deployment in authored.deployments + ) + builder_payload = _resolve_builder_payload(authored, source_root) + dependencies = stack.enter_context( + self._dependencies.resolve( + authored.client.dependencies, + client_image, + run_root=run_root, + source_root=source_root, + ) + ) + runtime_bundle = stage_runtime_bundle(workspace_root) + effective = resolve_slurm_config( + authored, + selected_profile=self._profile, + client_image=client_image, + deployment_images=deployment_images, + dependency_lock=dependencies.lock, + runtime_bundle=runtime_bundle, + run_id=run_id, + package_version=self._package_version, + 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 (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 | None, + clock: Clock, + ) -> None: + self._preparer = preparer + self._profile = selected_profile + self._launcher = launcher + self._publisher = publisher + self._clock = clock + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + force: 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, + ) + if self._publisher is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "run submission is not available; use --dry-run", + ) + publisher = self._publisher + self._initialize_run( + publisher, + config, + prepared, + force=force, + ) + try: + receipt = self._launcher.submit_script(prepared.batch_script) + except SlurmLauncherError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "Slurm submission is unavailable", + ) from None + try: + self._record_submission(publisher, 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 + raise + except BaseException: + try: + self._launcher.cancel(receipt.job_id) + except Exception: + pass + raise + 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 _initialize_run( + self, + publisher: SlurmRunArtifactPublisher, + config: DataDesignerSlurmConfig, + prepared: _PreparedRun, + *, + force: bool, + ) -> None: + try: + publisher.initialize( + config, + prepared.plan, + prepared.dependencies, + prepared.builder_payload, + force=force, + ) + 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, + publisher: SlurmRunArtifactPublisher, + plan: ResolvedSlurmRunPlan, + job_id: int, + ) -> None: + try: + 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 status(self, run_id: Identifier) -> SlurmPersistedRunStatus: + operation = SlurmServiceOperation.STATUS_RUN + try: + writer = SlurmStateWriter(self._profile.profile.workspace_root, run_id) + run = writer.load_run() + shards = [] + for shard in writer.load_shards(): + 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 cancel(self, run_id: Identifier) -> SlurmRunCancellation: + status = self.status(run_id) + active = {AttemptLifecycleState.SUBMITTED, AttemptLifecycleState.PENDING, AttemptLifecycleState.RUNNING} + 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 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) + + +class _RegistryImageBackend: + def __init__(self, workspace_root: str) -> None: + self._verified = VerifiedImageRegistry(workspace_root) + self._store = ImageRegistryStore(workspace_root) + + 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: + del request, replace + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.ADD_IMAGE, + "image registration is not available; use a pre-registered image", + ) + + 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) -> RegisteredImage: + return self._invoke_registry(SlurmServiceOperation.REMOVE_IMAGE, lambda: self._verified.unregister(name)) + + @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: + raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, "image registry conflict") 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, +) -> SlurmRunService: + """Create the production run service for one selected cluster profile.""" + selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + command_client = launcher or SlurmCommandClient() + 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, + clock or _utc_now, + ) + 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, +) -> SlurmImageService: + """Create the production image service for one selected cluster profile.""" + selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + backend = _RegistryImageBackend(selected.profile.workspace_root) + return SlurmImageService(backend, 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 _utc_now() -> datetime: + return datetime.now(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_image_service", + "create_slurm_run_service", +] 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..91b2f4675 --- /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.image, + 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_image_sha256 == single_node_plan.client.image.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.image, + 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.image, + 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") + image = single_node_plan.client.image.inspection_facts + supplied = ResolvedDependencyLock( + schema_version=1, + resolver_version="external-1", + python_abi=image.python_abi, + client_image_sha256=single_node_plan.client.image.sha256, + authored_requirements=("example-plugin==1.2.0",), + image_distributions=tuple(sorted(image.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.image, + 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") + image = single_node_plan.client.image.inspection_facts + supplied = ResolvedDependencyLock( + schema_version=1, + resolver_version="external-1", + python_abi=image.python_abi, + client_image_sha256=single_node_plan.client.image.sha256, + authored_requirements=("example-plugin==1.2.0",), + image_distributions=tuple(sorted(image.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.image, + 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.image, + 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.image, + run_root=tmp_path / "workspace/runs/run-002", + source_root=tmp_path, + ): + pass diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index e1cd641a2..7abca1b66 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -19,6 +19,7 @@ SlurmConfigBuilderError, SlurmConfigLoadError, SlurmProfileCatalog, + load_builder_payload, load_profile_catalog, load_run_config, resolve_profile, @@ -263,6 +264,23 @@ def test_loader_preserves_literal_interpolation_inside_builder_payload(tmp_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, diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index b937b5b06..80a5532e3 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -29,8 +29,12 @@ SlurmBatchScriptRenderer, SlurmBenchmarkBackend, SlurmBenchmarkService, + SlurmImageManager, SlurmImageResolver, SlurmImageService, + SlurmRunArtifactPublisher, + SlurmRunBackend, + SlurmRunExecution, SlurmRunPlanner, SlurmRunService, SlurmServiceError, @@ -180,6 +184,46 @@ def test_run_service_rejects_untyped_config() -> None: 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, bool]] = [] + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + force: bool, + ) -> SlurmRunExecution: + self.calls.append((config, source_root, dry_run, force)) + 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, force=True) + + assert result.run_id == "run-0001" + assert backend.calls == [(authored_run_single, tmp_path.resolve(), True, True)] + + def test_run_service_normalizes_and_redacts_unexpected_backend_errors( authored_run_single: DataDesignerSlurmConfig, ) -> None: @@ -359,6 +403,16 @@ def test_image_service_rejects_untyped_reference() -> None: 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, @@ -511,7 +565,10 @@ def test_service_errors_use_the_data_designer_error_hierarchy() -> None: [ ("SlurmBatchScriptRenderer", SlurmBatchScriptRenderer), ("SlurmBenchmarkBackend", SlurmBenchmarkBackend), + ("SlurmImageManager", SlurmImageManager), ("SlurmImageResolver", SlurmImageResolver), + ("SlurmRunArtifactPublisher", SlurmRunArtifactPublisher), + ("SlurmRunBackend", SlurmRunBackend), ("SlurmRunPlanner", SlurmRunPlanner), ], ) 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..cd8a09d2a --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +from datetime import UTC, datetime +from pathlib import Path + +import pytest + +from data_designer.slurm.client.dependencies import ResolvedClientDependencies +from data_designer.slurm.config import ( + BuilderInput, + DataDesignerSlurmConfig, + ImageBuildRequest, + SlurmProfile, + SlurmProfileCatalog, +) +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.models import SlurmJobSubmissionReceipt +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmServiceError, + SlurmServiceErrorCode, + create_slurm_image_service, + create_slurm_run_service, +) +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateWriter, + StateConflictError, +) + + +class _Launcher: + def __init__(self, gpu_counts: tuple[int, ...] = (), *, cancel_error: Exception | None = None) -> None: + self.submissions: list[str] = [] + self.cancellations: list[int] = [] + self.gpu_counts = gpu_counts + self.cancel_error = cancel_error + + def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + self.submissions.append(script) + return SlurmJobSubmissionReceipt(job_id=42) + + def cancel(self, job_id: int) -> None: + self.cancellations.append(job_id) + if self.cancel_error is not None: + raise self.cancel_error + + def query_gpu_counts(self, *, partition: str | None = None) -> tuple[int, ...]: + assert partition is not None + return self.gpu_counts + + +class _Publisher: + def __init__( + self, + *, + initialization_error: BaseException | None = None, + submission_error: BaseException | None = None, + ) -> None: + self.initializations: list[tuple[str, bool]] = [] + self.submissions: 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.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 _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 = ((authored.client.image, plan.client.image),) + 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 launcher.submissions == [] + + +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_production_wiring_submits_after_publisher_initialization( + 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() + submitted_at = datetime(2026, 9, 8, tzinfo=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, force=True) + + assert result.state == "submitted" + assert result.job_id == 42 + assert publisher.initializations == [("run-wired", True)] + assert publisher.submissions == [("run-wired", 42, submitted_at)] + assert len(launcher.submissions) == 1 + + +def test_production_wiring_stops_before_submission_without_state_publisher( + 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", + ) + + with pytest.raises(SlurmServiceError) as caught: + service.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "run submission is not available; use --dry-run" + assert launcher.submissions == [] + + +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] + + +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=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_and_lifecycle_gap( + 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 len(service.list()) == 1 + with pytest.raises(SlurmServiceError) as caught: + service.add( + ImageBuildRequest( + name="new-image", + kind="client", + source=f"registry.example/client@sha256:{'1' * 64}", + ) + ) + assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE + assert str(caught.value) == "image registration is not available; use a pre-registered image" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py index 9f05ccb4e..de32b94d1 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py @@ -7,6 +7,7 @@ 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 @@ -64,8 +65,9 @@ def __init__( self._script = _ScriptedResponses(responses) self.calls = self._script.calls - def plan(self, config: DataDesignerSlurmConfig) -> ResolvedSlurmRunPlan: + 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: 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..417657843 --- /dev/null +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -0,0 +1,147 @@ +# 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 +from click.testing import CliRunner + +import data_designer.slurm.cli as cli_module +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.services import ( + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) + + +class _RunService: + def __init__(self) -> None: + self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool, bool]] = [] + + def execute( + self, + config: DataDesignerSlurmConfig, + *, + source_root: Path, + dry_run: bool, + force: bool, + ) -> SlurmRunExecution: + self.calls.append((config, source_root, dry_run, force)) + return SlurmRunExecution( + run_id="run-0001", + state="dry_run", + plan_sha256="1" * 64, + shard_count=1, + batch_script="#!/bin/bash\n", + ) + + +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", "--force"]) + + 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, True)] + + +@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", + } + } + + +def test_cli_exposes_only_m2_run_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", "image")) + assert all(command not in result.stdout for command in ("retry", "merge", "benchmark")) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index f80247355..3ff48707c 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -191,11 +191,13 @@ def main() -> None: 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") 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.9.2,<3").specifier assert leaf_pyyaml_requirement.specifier == Requirement("pyyaml>=6.0.1,<7").specifier assert base_leaf_requirement.marker is not None diff --git a/uv.lock b/uv.lock index a995c306b..013aff91d 100644 --- a/uv.lock +++ b/uv.lock @@ -953,6 +953,7 @@ source = { editable = "packages/data-designer-slurm" } dependencies = [ { name = "data-designer" }, { name = "packaging" }, + { name = "pip" }, { name = "pydantic" }, { name = "pyyaml" }, ] @@ -961,6 +962,7 @@ dependencies = [ requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, { name = "packaging", specifier = ">=25,<27" }, + { name = "pip", specifier = ">=25,<27" }, { name = "pydantic", specifier = ">=2.9.2,<3" }, { name = "pyyaml", specifier = ">=6.0.1,<7" }, ] @@ -3406,6 +3408,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" From b5a65cc6c9038ea07faacfdc1f24f3fd9b2e7f19 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:38:24 -0300 Subject: [PATCH 20/62] fix(slurm): close one-node M2 integration (#929) * fix(slurm): close one-node M2 integration * fix(slurm): address integration review findings * fix(slurm): address smoke and review findings * docs(slurm): align status reconciliation help * fix(slurm): honor nested runtime mounts * fix(slurm): allow exclusive GPU allocations * fix(slurm): accept exclusive job GPU visibility * fix(slurm): load plans through nested mounts * fix(slurm): validate logical client paths * fix(slurm): resolve client artifacts by mount * fix(slurm): map serving model paths * fix(slurm): avoid idle GPU reaper * fix(slurm): isolate shared-node service ports * fix(slurm): validate allocation-local endpoints --- .../src/data_designer/slurm/cli.py | 4 +- .../data_designer/slurm/client/environment.py | 105 +++-- .../data_designer/slurm/client/execution.py | 68 +-- .../data_designer/slurm/config/__init__.py | 7 +- .../data_designer/slurm/config/environment.py | 24 +- .../data_designer/slurm/config/profiles.py | 11 + .../data_designer/slurm/launcher/client.py | 70 ++- .../data_designer/slurm/launcher/renderer.py | 17 +- .../data_designer/slurm/launcher/runner.py | 41 +- .../slurm/runtime/backpressure.py | 237 ++++++++++ .../data_designer/slurm/runtime/bootstrap.py | 377 ++++++++++++++++ .../src/data_designer/slurm/runtime/bundle.py | 23 +- .../data_designer/slurm/runtime/cleanup.sh | 69 +++ .../data_designer/slurm/runtime/context.py | 107 +++++ .../data_designer/slurm/runtime/controller.py | 99 ++++- .../data_designer/slurm/runtime/entrypoint.py | 411 +++++++++++++++--- .../data_designer/slurm/runtime/entrypoint.sh | 180 ++++++++ .../src/data_designer/slurm/runtime/paths.py | 40 +- .../slurm/runtime/plan_reader.sh | 197 +++++++++ .../src/data_designer/slurm/runtime/ports.py | 101 +++++ .../data_designer/slurm/runtime/preflight.py | 26 +- .../data_designer/slurm/runtime/records.py | 7 +- .../slurm/runtime/step_runner.sh | 167 +++++++ .../src/data_designer/slurm/runtime/steps.py | 90 ++-- .../data_designer/slurm/services/artifacts.py | 275 ++++++++++++ .../data_designer/slurm/services/results.py | 2 +- .../src/data_designer/slurm/services/run.py | 4 +- .../data_designer/slurm/services/wiring.py | 230 +++++++++- .../data_designer/slurm/state/artifacts.py | 7 +- .../data_designer/slurm/state/finalization.py | 153 +++++-- .../src/data_designer/slurm/state/reader.py | 8 +- .../slurm/state/reconciliation.py | 16 +- .../src/data_designer/slurm/state/storage.py | 59 ++- .../src/data_designer/slurm/state/store.py | 87 +++- .../tests/client/conftest.py | 16 +- .../tests/client/test_environment.py | 98 ++++- .../tests/client/test_worker.py | 60 ++- .../contracts/golden/multi_node_plan.json | 3 +- .../contracts/golden/profile_catalog.json | 2 + .../contracts/golden/single_node_plan.json | 3 +- .../tests/contracts/test_profiles.py | 1 + .../golden/finalization_chain.json | 12 +- .../tests/launcher/test_client.py | 51 ++- .../tests/launcher/test_renderer.py | 13 + .../tests/launcher/test_runner.py | 10 +- .../tests/runtime/conftest.py | 64 ++- .../tests/runtime/test_backpressure.py | 118 +++++ .../tests/runtime/test_bootstrap.py | 43 ++ .../tests/runtime/test_bundle.py | 34 +- .../tests/runtime/test_context.py | 117 +++++ .../tests/runtime/test_controller.py | 84 +++- .../tests/runtime/test_entrypoint.py | 194 ++++++++- .../tests/runtime/test_paths.py | 3 +- .../tests/runtime/test_ports.py | 75 ++++ .../tests/runtime/test_preflight.py | 56 ++- .../tests/runtime/test_shell_runtime.py | 312 +++++++++++++ .../tests/runtime/test_steps.py | 48 +- .../tests/services/test_wiring.py | 272 +++++++++++- .../golden/rendered/multi_node.sbatch | 11 +- .../golden/rendered/single_node.sbatch | 11 +- .../tests/slurm_test_fakes/slurm.py | 2 + .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_state_golden_records.py | 2 +- .../tests/state/test_store.py | 233 +++++++++- 64 files changed, 4878 insertions(+), 393 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/backpressure.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_backpressure.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_bootstrap.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_context.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_ports.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_shell_runtime.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 7315f4d3b..f2c9e12fd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -70,7 +70,7 @@ def status_command( profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), cluster: str | None = typer.Option(None, "--cluster"), ) -> None: - """Show persisted M2 run status without scheduler reconciliation.""" + """Reconcile scheduler observations and show persisted M2 run status.""" operation = SlurmServiceOperation.STATUS_RUN result = _invoke( operation, @@ -85,7 +85,7 @@ def cancel_command( profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), cluster: str | None = typer.Option(None, "--cluster"), ) -> None: - """Request job cancellation; status changes after reconciliation.""" + """Request cancellation of active jobs.""" operation = SlurmServiceOperation.CANCEL_RUN result = _invoke( operation, 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 index aa48076e8..6e96d15ea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -7,6 +7,7 @@ import importlib.metadata import json import os +import posixpath import re import subprocess import sys @@ -44,7 +45,7 @@ class PreparedClientEnvironment: @dataclass(frozen=True) class _BootstrapInputs: run_id: str - run_root: Path + logical_run_root: Path shard_id: str attempt_id: str attempt_dir: Path @@ -53,6 +54,7 @@ class _BootstrapInputs: installer_path: Path inspection: dict[str, object] dependency_lock: ArtifactReference + plan: dict[str, object] @dataclass(frozen=True) @@ -117,17 +119,24 @@ def _load_bootstrap_inputs( 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 = _load_json_object(plan_path, ClientErrorCode.INVALID_INPUT) - run_id = _require_string(plan, "run_id") - run_root = plan_path.parent - if plan_path.name != "resolved-plan.json" or run_root.name != run_id or run_root.parent.name != "runs": + 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 _get_container_path(plan_payload, (logical_run_root / plan_path.name).as_posix()) != plan_path.as_posix() + ): raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved plan path is not canonical") - expected_attempt = run_root / "shards" / shard_id / "attempts" / attempt_id + logical_attempt = logical_run_root / "shards" / shard_id / "attempts" / attempt_id + expected_attempt = Path(_get_container_path(plan_payload, logical_attempt.as_posix(), require_writable=True)) 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, "client") + client = _require_object(plan_payload, "client") image = _require_object(client, "image") image_sha256 = _require_digest(image, "sha256") inspection_record = _require_object(image, "inspection") @@ -142,11 +151,11 @@ def _load_bootstrap_inputs( raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client installer path is invalid") lock_reference = _artifact_reference(_require_object(client, "dependency_lock")) - if Path(lock_reference.path).as_posix() != (run_root / "dependency-lock.json").as_posix(): + 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, - run_root=run_root, + logical_run_root=logical_run_root, shard_id=shard_id, attempt_id=attempt_id, attempt_dir=attempt_dir, @@ -155,11 +164,13 @@ def _load_bootstrap_inputs( installer_path=installer_path, inspection=inspection, dependency_lock=lock_reference, + plan=plan_payload, ) def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependencies: lock_bytes = read_regular_bytes( - Path(inputs.dependency_lock.path), missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING + Path(_get_container_path(inputs.plan, 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") @@ -185,7 +196,11 @@ def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependen source = lock.get("source") if source is not None: source_reference = _artifact_reference(_as_object(source)) - _verify_input_artifact(source_reference, inputs.run_root / "inputs") + _verify_input_artifact( + source_reference, + inputs.logical_run_root / "inputs", + inputs.plan, + ) return _VerifiedDependencies( image_distributions=expected_image, overlay_packages=tuple(_as_object_list(lock.get("overlay_packages"))), @@ -198,10 +213,20 @@ def _prepare_overlay( ) -> tuple[Path, ClientInstallerOutcome, tuple[InstalledDistribution, ...]]: expected_overlay, wheels = _verify_wheels( dependencies.overlay_packages, - inputs.run_root / "dependencies", + inputs.logical_run_root / "dependencies", dependencies.image_distributions, + inputs.plan, + ) + logical_overlay = ( + inputs.logical_run_root + / "shards" + / inputs.shard_id + / "attempts" + / inputs.attempt_id + / "client-env" + / "site-packages" ) - overlay_path = inputs.attempt_dir / "client-env" / "site-packages" + overlay_path = Path(_get_container_path(inputs.plan, logical_overlay.as_posix(), require_writable=True)) outcome = self._install_overlay(inputs.installer_path, wheels, expected_overlay, overlay_path) installed = tuple(sorted((*dependencies.image_distributions, *expected_overlay), key=lambda item: item.name)) return overlay_path, outcome, installed @@ -293,8 +318,9 @@ def _run_installer(command: tuple[str, ...]) -> None: def _verify_wheels( packages: tuple[dict[str, object], ...], - dependencies_root: Path, + logical_dependencies_root: Path, image_distributions: tuple[InstalledDistribution, ...], + plan: dict[str, object], ) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: expected: list[InstalledDistribution] = [] wheels: list[Path] = [] @@ -305,13 +331,10 @@ def _verify_wheels( if name in image_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")) - wheel = Path(artifact.path) - try: - wheel.relative_to(dependencies_root) - except ValueError as error: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency wheel path is not canonical") from error - if wheel.parent != dependencies_root or wheel.suffix != ".whl": + 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 = Path(_get_container_path(plan, logical_wheel.as_posix())) if compute_file_sha256(wheel, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != artifact.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency wheel digest differs") try: @@ -332,14 +355,11 @@ def _verify_wheels( return tuple(pair[0] for pair in sorted_pairs), tuple(pair[1] for pair in sorted_pairs) -def _verify_input_artifact(reference: ArtifactReference, root: Path) -> None: - path = Path(reference.path) - try: - path.relative_to(root) - except ValueError as error: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "dependency source path is not canonical") from error - if path.parent != root or path.suffix != ".json": +def _verify_input_artifact(reference: ArtifactReference, logical_root: Path, plan: dict[str, object]) -> 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 = Path(_get_container_path(plan, logical_path.as_posix())) if compute_file_sha256(path, missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING) != reference.sha256: raise ClientWorkerError(ClientErrorCode.DEPENDENCY_DIGEST_MISMATCH, "dependency source digest differs") @@ -408,6 +428,37 @@ def _artifact_reference(value: dict[str, object]) -> ArtifactReference: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client artifact reference is invalid") from error +def _get_container_path(plan: dict[str, object], host_path: str, *, require_writable: bool = False) -> str: + if not host_path.startswith("/") or posixpath.normpath(host_path) != host_path: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is invalid") + values = plan.get("container_mounts") + if not isinstance(values, list): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") + candidates: list[tuple[str, str, bool]] = [] + for value in values: + mount = _as_object(value) + source = _require_string(mount, "source") + target = _require_string(mount, "target") + read_only = mount.get("read_only") + if ( + not source.startswith("/") + or posixpath.normpath(source) != source + or not target.startswith("/") + or posixpath.normpath(target) != target + or type(read_only) is not bool + ): + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") + if host_path == source or host_path.startswith(f"{source}/"): + candidates.append((source, target, read_only)) + if not candidates: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not mounted") + source, target, read_only = max(candidates, key=lambda item: len(item[0])) + if require_writable and read_only: + raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not writable") + relative_path = posixpath.relpath(host_path, source) + return target if relative_path == "." else posixpath.join(target, relative_path) + + def _sha256_bytes(value: bytes) -> str: return hashlib.sha256(value).hexdigest() 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 index 8a6d3f05c..5f3fa3184 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -49,8 +49,10 @@ 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, compute_canonical_json_sha256 +from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import PlannedShard, ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.paths import get_container_path, get_host_path +from data_designer.slurm.runtime.ports import resolve_allocation_plan from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest Clock = Callable[[], datetime] @@ -154,9 +156,11 @@ def __init__( *, 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, @@ -250,7 +254,8 @@ def _build_context( 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 + Path(get_container_path(plan, plan.client.dependency_lock.path)), + missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, ) ) if lock.compute_sha256() != plan.client.dependency_lock.sha256: @@ -271,17 +276,18 @@ def _build_context( raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client environment differs from the lock") builder_payload = self._load_builder(plan) builder = DataDesignerConfigBuilder.from_config(builder_payload) - providers = self._materialize_model_endpoints(plan, builder, endpoints) + 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) - self._validate_assets(plan) + managed_assets_path = self._validate_assets(plan) requested_resume = ResumeMode(plan.invocation.authored.resume) - dataset_path = self._dataset_path(shard, prepared, requested_resume) + dataset_path = self._dataset_path(plan, shard, prepared, requested_resume) designer = self._data_designer_factory( artifact_path=dataset_path.parent, model_providers=providers, - managed_assets_path=plan.invocation.effective_input_bindings.managed_assets_path, + managed_assets_path=managed_assets_path.as_posix(), mcp_providers=mcp_providers, auto_configure_logging=False, ) @@ -298,7 +304,8 @@ def _validate_prepared( shard: PlannedShard, prepared: PreparedClientEnvironment, ) -> None: - expected_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + logical_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id + expected_attempt = Path(get_container_path(plan, logical_attempt.as_posix(), require_writable=True)) inspection = plan.client.image.inspection_facts if ( plan.run_id != prepared.run_id @@ -314,7 +321,10 @@ 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) + payload = read_regular_bytes( + Path(get_container_path(plan, 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: @@ -390,15 +400,18 @@ def _materialize_seed( 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) + path = Path(get_container_path(plan, 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) + payload = read_regular_bytes( + Path(get_container_path(plan, 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": seed_path}) + source = seed.source.model_copy(update={"path": path.as_posix()}) builder.with_seed_dataset( source, sampling_strategy=seed.sampling_strategy, @@ -437,20 +450,22 @@ def _materialize_mcp_providers( return providers @staticmethod - def _validate_assets(plan: ResolvedSlurmRunPlan) -> None: + def _validate_assets(plan: ResolvedSlurmRunPlan) -> Path: value = plan.invocation.effective_input_bindings.managed_assets_path assert value is not None - path = Path(value) + path = Path(get_container_path(plan, 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) + resume_path = Path(get_container_path(plan, shard.resume_workspace.path, require_writable=True)) 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())): @@ -528,7 +543,7 @@ def _validate_creation_result( dataset_path = Path(results.dataset_path) effective_resume = results.effective_resume_mode - shared_path = Path(context.shard.resume_workspace.path) + shared_path = Path(get_container_path(context.plan, context.shard.resume_workspace.path, require_writable=True)) expected_path = shared_path if effective_resume is ResumeMode.ALWAYS else prepared.attempt_dir / "dataset" if ( not dataset_path.is_absolute() @@ -589,7 +604,9 @@ def _build_candidate_manifest( metadata = lazy.pq.read_metadata(exported_path) if metadata.num_rows != creation.actual_records: raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "exported record count differs") - schema_digest = hashlib.sha256(lazy.pq.read_schema(exported_path).serialize().to_pybytes()).hexdigest() + schema_digest = hashlib.sha256( + lazy.pq.read_schema(exported_path).remove_metadata().serialize().to_pybytes() + ).hexdigest() files = ( CandidateOutputFile( relative_path=exported_path.relative_to(dataset_path).as_posix(), @@ -599,18 +616,6 @@ def _build_candidate_manifest( ), ) created_at = self._clock() - provenance_digest = compute_canonical_json_sha256( - { - "builder_sha256": context.plan.builder.content_sha256, - "client_image_sha256": prepared.client_image_sha256, - "dependency_lock_sha256": prepared.dependency_lock.sha256, - "files": [file.model_dump(mode="json") for file in files], - "attempt_id": prepared.attempt_id, - "resolved_plan_sha256": context.plan.compute_sha256(), - "run_id": context.plan.run_id, - "shard_id": context.shard.shard_id, - } - ) return CandidateOutputManifest( schema_version=1, run_id=context.plan.run_id, @@ -618,7 +623,7 @@ def _build_candidate_manifest( attempt_id=prepared.attempt_id, attempt_ordinal=int(prepared.attempt_id.removeprefix("attempt-")), created_at=created_at, - dataset_path=dataset_path.as_posix(), + dataset_path=get_host_path(context.plan, dataset_path.as_posix(), require_writable=True), requested_records=creation.requested_records, actual_records=creation.actual_records, outcome=( @@ -630,7 +635,7 @@ def _build_candidate_manifest( ), files=files, dataset_schema_digest=schema_digest, - provenance_digest=provenance_digest, + provenance_digest=context.plan.compute_sha256(), ) def _publish_success( @@ -660,7 +665,8 @@ def _publish_success( 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() + path=get_host_path(context.plan, candidate_path.as_posix(), require_writable=True), + sha256=candidate.compute_sha256(), ), ) publish_private_text(prepared.attempt_dir / "client-result.json", result.serialize_json()) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index c08456652..00b7f9bf3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -15,7 +15,11 @@ FixedRecordPolicy, ) from data_designer.slurm.config.builder import DataDesignerSlurmConfigBuilder -from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +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, @@ -104,6 +108,7 @@ "RemoteMCPProviderConfig", "SchedulerProfile", "SecretRef", + "collect_secret_environment_names", "SelectedSlurmProfile", "ServerDeploymentConfig", "ServingImageInspection", 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 index c835c3983..006effec8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py @@ -9,7 +9,7 @@ from collections.abc import Mapping from typing import Annotated, Literal -from pydantic import Field, JsonValue, StringConstraints, field_validator +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 @@ -18,6 +18,7 @@ "EnvironmentBinding", "LiteralEnvironmentBinding", "SecretRef", + "collect_secret_environment_names", ] @@ -59,6 +60,13 @@ class SecretRef(AuthoredConfig): ] +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) @@ -112,3 +120,17 @@ 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/profiles.py b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py index f3fc57987..5567850f0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -30,6 +30,17 @@ class SchedulerProfile(AuthoredConfig): account: Identifier | None = None partition: Identifier | 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): 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 index 47d061ad1..2d822a7f1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -8,7 +8,7 @@ import re import subprocess import unicodedata -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import TypeAlias @@ -32,6 +32,7 @@ _JobSelector: TypeAlias = int | SchedulerIdentity _IDENTIFIER_PATTERN = re.compile(r"^[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 @@ -43,10 +44,11 @@ class SlurmExecutables: 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.sinfo): + 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") @@ -67,22 +69,42 @@ def __init__( 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) -> SlurmJobSubmissionReceipt: + 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 './'") - output = self._run((self._executables.sbatch, "--parsable", "--export=NIL", path)) + 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) -> SlurmJobSubmissionReceipt: + 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) output = self._run( - (self._executables.sbatch, "--parsable", "--export=NIL"), + (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), input_text=script, + environment=environment, ) return parse_submission(output) @@ -134,6 +156,10 @@ def cancel(self, selector: _JobSelector) -> 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: Identifier | None = None) -> tuple[int, ...]: """Return configured GPU counts reported for eligible node groups.""" command = [self._executables.sinfo, "--noheader", "--format=%G"] @@ -143,12 +169,23 @@ def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, command.append(f"--partition={partition}") return parse_gpu_counts(self._run(command)) - def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: + 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: - completed = ( - self._runner.run(command) if input_text is None else self._runner.run(command, input_text=input_text) - ) + 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 (OSError, subprocess.SubprocessError) as error: raise SlurmCommandError(f"{command_name} could not be executed: {_format_error_detail(error)}") from error returncode = getattr(completed, "returncode", None) @@ -162,6 +199,19 @@ def _run(self, command: Sequence[str], *, input_text: str | None = None) -> str: return stdout +def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[str, Mapping[str, str] | None]: + if not environment: + return "--export=NIL", None + names = tuple(sorted(environment)) + for name in names: + if _ENVIRONMENT_NAME_PATTERN.fullmatch(name) is None: + raise ValueError("exported environment names must be valid identifiers") + value = environment[name] + if type(value) is not str or "\0" in value: + raise ValueError("exported environment values must be strings without NUL") + return f"--export={','.join(names)}", environment + + def _format_selectors(selectors: Sequence[_JobSelector]) -> str: if not selectors: raise ValueError("at least one managed Slurm job selector is required") 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 index e076f5cb1..ad71b0af5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -11,6 +11,8 @@ from data_designer.slurm.launcher.errors import SlurmBatchRenderError from data_designer.slurm.planning import ResolvedSlurmRunPlan +_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.""" @@ -20,12 +22,16 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") directive_text = render_batch_directives(_build_generation_directives(plan)) + if plan.selected_profile.profile.gpu_request_mode == "visible": + directive_text = f"{directive_text}\n#SBATCH --exclusive" attempt = f"{attempt_ordinal:04d}" + scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path + command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" return f"""#!/usr/bin/env bash {directive_text} set -Eeuo pipefail -export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +export PATH={quote_shell_value(command_path)} readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} @@ -50,12 +56,11 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi printf -v DD_SHARD_ID 'shard-%05d' "${{DD_ARRAY_TASK_ID}}" readonly DD_SHARD_ID readonly DD_ATTEMPT_DIR="${{DD_RUN_ROOT}}/shards/${{DD_SHARD_ID}}/attempts/attempt-${{DD_ATTEMPT_ORDINAL}}" -install -d -m 0700 "${{DD_ATTEMPT_DIR}}" -DD_RUNTIME_DIR="$(mktemp -d "${{DD_ATTEMPT_DIR}}/runtime.${{DD_RUNTIME_SHA256}}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_DIR}}" +readonly DD_RUNTIME_ROOT="${{DD_ATTEMPT_DIR}}/runtime" +[[ -d ${{DD_RUNTIME_ROOT}} && ! -L ${{DD_RUNTIME_ROOT}} ]] +tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_ROOT}}" -source "${{DD_RUNTIME_DIR}}/entrypoint.sh" +source "${{DD_RUNTIME_ROOT}}/entrypoint.sh" dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" """ 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 index 9e2e7839f..de0ceb1d5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/runner.py @@ -16,7 +16,13 @@ class CommandRunner(Protocol): """Minimal command boundary implemented by production and fake runners.""" - def run(self, command: Sequence[str], *, input_text: str | None = None) -> subprocess.CompletedProcess[str]: + 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.""" ... @@ -35,9 +41,12 @@ def __init__( ) -> 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") - explicit_environment = ( - dict(environment) if environment is not None else {"PATH": os.environ.get("PATH") or os.defpath} - ) + 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") @@ -51,8 +60,18 @@ 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) -> subprocess.CompletedProcess[str]: + 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), @@ -62,7 +81,7 @@ def run(self, command: Sequence[str], *, input_text: str | None = None) -> subpr text=True, encoding="utf-8", errors="replace", - env=dict(self._environment), + env=process_environment, timeout=self._timeout_seconds, ) return subprocess.run( @@ -73,6 +92,14 @@ def run(self, command: Sequence[str], *, input_text: str | None = None) -> subpr text=True, encoding="utf-8", errors="replace", - env=dict(self._environment), + 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/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..54a053231 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed one-node step manifest produced inside the sealed client image.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Literal + +from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator + +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) +from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path +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 AllocationContext, RuntimeEndpoint, RuntimeStepRole +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.ports import resolve_allocation_deployments +from data_designer.slurm.runtime.steps import ( + build_client_command, + build_endpoint_command, + build_vllm_command, +) +from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment +from data_designer.slurm.serving.vllm import ResolvedVllmProcess +from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest + + +class RuntimeProbeSpec(ContractValue): + """One loopback readiness target monitored by the Bash controller.""" + + host: Literal["127.0.0.1"] = "127.0.0.1" + 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 + + +class RuntimeStepSpec(ContractValue): + """Container command and environment consumed by the Bash step runner.""" + + step_id: Identifier + role: RuntimeStepRole + image_path: str + command: tuple[str, ...] = Field(min_length=1) + cpus: PositiveInt + gpu_indices: tuple[NonNegativeInt, ...] = () + 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: RuntimeProbeSpec | None = None + + _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) + _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.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.role is RuntimeStepRole.SERVER and not self.gpu_indices: + raise ValueError("server runtime steps require GPUs") + if self.role is not RuntimeStepRole.SERVER and self.gpu_indices: + raise ValueError("non-server runtime steps cannot request GPUs") + return self + + +class RuntimeBootstrapManifest(ContractRecord): + """Secret-free one-node 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 + + +def build_runtime_manifest( + context: AllocationContext, + environment: Mapping[str, str], + *, + runtime_root: Path, + log_directory: Path, +) -> RuntimeBootstrapManifest: + """Build the secret-free one-node command handoff for the Bash controller.""" + plan = context.plan + runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) + 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_container_root, + log_directory, + ) + ] + for deployment in deployments: + steps.extend( + _build_server_step( + deployment, + process, + context, + runtime_root, + runtime_container_root, + log_directory, + ) + for process in deployment.processes + ) + steps.extend(_build_endpoint_step(deployment, context, runtime_root, log_directory) for deployment in deployments) + steps.append( + _build_client_step( + RuntimeStepRole.CLIENT, + "client-generation", + "client", + context, + environment, + endpoints, + runtime_container_root, + log_directory, + ) + ) + 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_container_root: str, + log_directory: Path, +) -> RuntimeStepSpec: + plan = context.plan + command = build_client_command( + "preflight" if operation == "preflight" else "run", + plan, + context.shard, + context.attempt, + context.attempt_directory, + endpoints, + ) + if operation == "client": + command = ( + "python3", + "-m", + "data_designer.slurm.runtime.entrypoint", + "client", + *command[4:6], + *command[10:], + ) + 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, + image_path=plan.client.image.path, + command=command, + cpus=plan.client.authored.cpus, + gpu_indices=(), + literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_container_root, **allocation_environment}, + secret_environment={name: name for name in secret_names}, + environment_prefixes={}, + container_environment=tuple(sorted((*secret_names, *allocation_environment, "PYTHONPATH"))), + log_directory=log_directory, + ) + + +def _build_server_step( + deployment: ResolvedVllmServerDeployment, + process: ResolvedVllmProcess, + context: AllocationContext, + runtime_root: Path, + runtime_container_root: str, + log_directory: Path, +) -> RuntimeStepSpec: + 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", + ) + 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) + ) + probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) + return _step( + step_id=process.process_id, + role=RuntimeStepRole.SERVER, + image_path=deployment.image.path, + command=build_vllm_command(deployment, process, context.plan), + cpus=context.plan.client.authored.cpus, + gpu_indices=tuple(process.gpu_indices), + literal_environment=literal_environment, + secret_environment=secret_environment, + environment_prefixes=environment_prefixes, + container_environment=tuple( + sorted( + { + *deployment.launch_policy.environment, + "PYTHONPATH", + MAX_WAITING_REQUESTS_ENVIRONMENT, + RETRY_AFTER_SECONDS_ENVIRONMENT, + } + ) + ), + log_directory=log_directory, + launch_delay_seconds=process.launch_delay_seconds, + readiness=RuntimeProbeSpec( + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ), + ) + + +def _build_endpoint_step( + deployment: ResolvedVllmServerDeployment, + context: AllocationContext, + runtime_root: Path, + log_directory: Path, +) -> RuntimeStepSpec: + proxy_path = runtime_root / "data_designer/slurm/runtime/proxy.py" + command = build_endpoint_command( + deployment, + context.plan, + proxy_path, + deployment.logical_endpoint.port, + ) + return _step( + step_id=f"{deployment.deployment_id}-endpoint", + role=RuntimeStepRole.ENDPOINT, + image_path=context.plan.client.image.path, + command=command, + cpus=context.plan.client.authored.cpus, + gpu_indices=(), + literal_environment={"LC_ALL": "C"}, + secret_environment={}, + environment_prefixes={}, + container_environment=(), + log_directory=log_directory, + readiness=RuntimeProbeSpec( + port=deployment.logical_endpoint.port, + path="/health", + deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + ), + ) + + +def _step( + *, + step_id: str, + role: RuntimeStepRole, + image_path: str, + command: tuple[str, ...], + cpus: int, + gpu_indices: tuple[int, ...], + literal_environment: dict[str, str], + secret_environment: dict[str, str], + environment_prefixes: dict[str, str], + container_environment: tuple[str, ...], + log_directory: Path, + launch_delay_seconds: int = 0, + readiness: RuntimeProbeSpec | None = None, +) -> RuntimeStepSpec: + return RuntimeStepSpec( + step_id=step_id, + role=role, + image_path=image_path, + command=command, + cpus=cpus, + gpu_indices=gpu_indices, + 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 index 7027c496c..7cf08db2f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py @@ -25,23 +25,9 @@ _SOURCE_MODE = 0o400 _MAXIMUM_SOURCE_SIZE = 16 * 1024 * 1024 _TEMPORARY_NAME_PATTERN = re.compile(r"^\.runtime\.[0-9a-f]{16}\.tmp$") -_ENTRYPOINT_NAME = "entrypoint.sh" +_SHELL_SOURCE_NAMES = ("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" -_ENTRYPOINT = b"""#!/usr/bin/env bash -set -Eeuo pipefail - -dd_slurm_run_allocation() { - if [[ $# -ne 2 ]]; then - printf '%s\\n' 'allocation runtime requires plan and attempt directory arguments' >&2 - return 64 - fi - local runtime_root - runtime_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" - PYTHONNOUSERSITE=1 PYTHONSAFEPATH=1 PYTHONPATH="${runtime_root}" \ - python3 -m data_designer.slurm.runtime.entrypoint --plan "$1" --attempt-dir "$2" -} -""" def stage_runtime_bundle(workspace_root: str | Path) -> ArtifactReference: @@ -69,13 +55,16 @@ def stage_runtime_bundle(workspace_root: str | Path) -> ArtifactReference: def _build_runtime_archive() -> bytes: - sources = _collect_slurm_sources(Path(__file__).parents[1]) + 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, ): - _add_archive_file(archive, _ENTRYPOINT_NAME, _ENTRYPOINT, mode=_ENTRYPOINT_MODE) + 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: 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..c545d4e85 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh @@ -0,0 +1,69 @@ +# 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() { + read -r -t "$1" -u "${DD_RUNTIME_TIMER[0]}" _ || true +} + +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..6817300b4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Load allocation identity from container-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.runtime.paths import get_container_path +from data_designer.slurm.state import 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], +) -> tuple[AllocationContext, SlurmStateWriter]: + """Load one scheduler-selected shard attempt through its container 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.as_posix() != get_container_path(plan, expected_plan_path.as_posix()): + 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.as_posix() != get_container_path( + plan, host_attempt_directory.as_posix(), require_writable=True + ): + 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", + ) + return AllocationContext(plan, shard, attempt, host_attempt_directory), writer + + +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 get_container_path(plan, logical_plan_path.as_posix()) != plan_path.as_posix(): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") + workspace_root = Path(get_container_path(plan, logical_workspace_root, require_writable=True)) + return SlurmStateWriter( + workspace_root, + run_id, + logical_workspace_root=logical_workspace_root, + local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), + ) + + +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 index 2341d617b..e90cbc555 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -7,16 +7,18 @@ 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 Protocol +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_complete_client_candidate @@ -27,7 +29,6 @@ ) from data_designer.slurm.runtime.supervisor import ManagedStep, RuntimeClock, StepSupervisor from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment -from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, @@ -39,6 +40,7 @@ ProbeEvidence, ProbeOutcome, ReadinessState, + ShardWinner, StateNotFoundError, ) @@ -57,6 +59,10 @@ 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, @@ -65,6 +71,26 @@ def publish_attempt_result( """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.""" ... @@ -106,6 +132,7 @@ def __init__( context: AllocationContext, *, runtime_proxy_path: Path, + runtime_root: Path | None = None, state: RuntimeStateStore, supervisor: StepSupervisor, preflight: AllocationPreflight, @@ -120,6 +147,7 @@ def __init__( raise ValueError("runtime probe intervals must be positive") 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 @@ -140,13 +168,39 @@ def run(self) -> AttemptManifest: self._record_outcome_failure(outcome) self._cleanup_runtime(outcome) self._record_stopped_readiness(outcome) - terminal = self._persist_terminal_outcome(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() @@ -214,10 +268,7 @@ def _execute(self) -> tuple[ArtifactReference, datetime]: def _prepare_runtime(self) -> _RuntimeTopology: self._validate_attempt_state() - deployments = tuple( - resolve_vllm_server(self._context.plan, deployment.deployment_id) - for deployment in self._context.plan.deployments - ) + deployments = resolve_allocation_deployments(self._context, self._environment) restarting = self._readiness is not None if restarting: self._begin_execution(deployments, ReadinessState.RESTARTING) @@ -297,20 +348,25 @@ def _run_client( self._environment, ) ) - 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_complete_client_candidate(self._context, self._attempt) - 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 - the record contract requires this for complete results - raise SlurmRuntimeError( - SlurmRuntimeErrorCode.FINALIZATION_FAILED, - "complete 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 + with self._state.acquire_dataset_workspace( + self._attempt.shard_id, + self._attempt.attempt_id, + self._context.plan.invocation.authored.resume, + ): + 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_complete_client_candidate(self._context, self._attempt) + 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 - the record contract requires this for complete results + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "complete 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 def _validate_client_timestamps( self, @@ -379,6 +435,7 @@ def _start_servers( self._context.plan, self._context.attempt_directory, self._environment, + self._runtime_root, ) ) for process, step in zip(deployment.processes, steps, strict=True): diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 30e2fd077..66ee59f21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Allocation-local command entrypoint loaded from the checksummed runtime bundle.""" +"""Typed control phases executed only inside the sealed client image.""" from __future__ import annotations @@ -9,28 +9,50 @@ import os import sys from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone from pathlib import Path -from data_designer.slurm.planning import PlannedShard -from data_designer.slurm.runtime.controller import OneNodeAllocationController +from data_designer.slurm.client.filesystem import ensure_private_directory, replace_private_text +from data_designer.slurm.client.worker import main as client_worker_main +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.paths import get_container_path +from data_designer.slurm.runtime.ports import resolve_allocation_deployments from data_designer.slurm.runtime.preflight import SystemAllocationPreflight -from data_designer.slurm.runtime.probes import HttpReadinessProber -from data_designer.slurm.runtime.signals import TerminationSignalCoordinator -from data_designer.slurm.runtime.steps import DefaultClientStepBuilder -from data_designer.slurm.runtime.supervisor import StepSupervisor, SubprocessStepRunner, SystemRuntimeClock -from data_designer.slurm.state import SlurmStateWriter +from data_designer.slurm.runtime.records import load_complete_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 allocation and return a bounded process status.""" - parser = argparse.ArgumentParser(prog="data-designer-slurm-runtime") - parser.add_argument("--plan", required=True) - parser.add_argument("--attempt-dir", required=True) - parsed = parser.parse_args(arguments) + """Execute one container-only allocation control phase.""" + parsed = _parse_arguments(arguments) try: - _run(Path(parsed.plan), Path(parsed.attempt_dir), os.environ) + 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 ( @@ -45,77 +67,326 @@ def main(arguments: Sequence[str] | None = None) -> int: return 0 -def _run(plan_path: Path, attempt_directory: Path, environment: Mapping[str, str]) -> None: - context, writer = _load_allocation_context(plan_path, attempt_directory, environment) - clock = SystemRuntimeClock() - signals = TerminationSignalCoordinator() - supervisor = StepSupervisor(SubprocessStepRunner(), signals=signals, clock=clock) - controller = OneNodeAllocationController( +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) + 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) + + +def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + _validate_attempt_is_executable(context.attempt) + SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) + SystemAllocationPreflight.verify_ports(context, environment) + readiness = _begin_attempt(context, writer, environment) + log_directory = execution_log_directory(context.attempt_directory, readiness.revision) + container_log_directory = Path(get_container_path(context.plan, log_directory.as_posix(), require_writable=True)) + ensure_private_directory(container_log_directory) + manifest = build_runtime_manifest( context, - runtime_proxy_path=Path(__file__).with_name("proxy.py"), - state=writer, - supervisor=supervisor, - preflight=SystemAllocationPreflight(), - client_steps=DefaultClientStepBuilder(), - prober=HttpReadinessProber(), - clock=clock, - environment=environment, + environment, + runtime_root=arguments.runtime_root, + log_directory=log_directory, ) - with signals.interrupt_on_termination(supervisor.cleanup): - controller.run() + expected_manifest = context.attempt_directory / "runtime-manifest.json" + if arguments.manifest.as_posix() != get_container_path( + context.plan, + expected_manifest.as_posix(), + require_writable=True, + ): + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime manifest path is invalid") + replace_private_text(arguments.manifest, manifest.serialize_json()) -def _load_allocation_context( - plan_path: Path, - attempt_directory: Path, - environment: Mapping[str, str], -) -> tuple[AllocationContext, SlurmStateWriter]: - 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", +def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, 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, + 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]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + generation_started_at = _now(context.attempt, _load_optional_readiness(context, writer)) + with writer.acquire_dataset_workspace( + context.shard.shard_id, + context.attempt.attempt_id, + context.plan.invocation.authored.resume, + ): + return_code = client_worker_main( + ( + "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(), + *(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_complete_client_candidate( + context, + context.attempt, + attempt_directory=arguments.attempt_dir, ) - shard = _select_shard(plan.shards, _scheduler_task_id(environment.get("SLURM_ARRAY_TASK_ID"))) - expected_attempt_root = plan_path.parent / "shards" / shard.shard_id / "attempts" - if attempt_directory.parent != expected_attempt_root: + 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_allocation_context(arguments.plan, arguments.attempt_dir, 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.INVALID_CONTEXT, - "attempt directory does not match the scheduler array task", + SlurmRuntimeErrorCode.FINALIZATION_FAILED, + "successful allocation has no candidate reference", ) - attempt = writer.load_attempt(shard.shard_id, attempt_directory.name) - return AllocationContext(plan, shard, attempt, attempt_directory), writer + 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 _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" or plan_path.parent.parent.name != "runs": - raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") - workspace_root = plan_path.parent.parent.parent - run_id = plan_path.parent.name - return SlurmStateWriter(workspace_root, run_id) +def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: + context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, 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 _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", +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}) ) - return selected[0] + 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 _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", +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, + 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, + 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 int(value) + ) + 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, + state=ReadinessState.STOPPED, + deployments=tuple( + deployment.model_copy( + update={ + "state": ReadinessState.STOPPED, + "ready_backends": 0, + "last_probe": _probe( + timestamp, + ProbeOutcome.SUCCESS, + "runtime_stopped", + "allocation processes stopped", + ), + } + ) + for deployment in previous.deployments + ), + ) + ) + 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 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..dc15efd34 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -0,0 +1,180 @@ +# 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}/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_slurm_run_allocation() { + if [[ $# -ne 2 ]]; then + printf '%s\n' 'allocation runtime requires plan and attempt directory arguments' >&2 + return 64 + fi + DD_PLAN_PATH=$1 + DD_ATTEMPT_PATH=$2 + DD_RUNTIME_MANIFEST=${DD_ATTEMPT_PATH}/runtime-manifest.json + dd_read_control_plan "${DD_PLAN_PATH}" + dd_read_plan_secret_names "${DD_PLAN_PATH}" + dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false + DD_PLAN_CONTAINER_PATH=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_ATTEMPT_PATH}" true + DD_ATTEMPT_CONTAINER_DIR=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_DIR}" true + DD_RUNTIME_CONTAINER_ROOT=${DD_CONTAINER_PATH} + dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_MANIFEST}" true + DD_RUNTIME_MANIFEST_CONTAINER_PATH=${DD_CONTAINER_PATH} + export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST + export DD_PLAN_CONTAINER_PATH DD_ATTEMPT_CONTAINER_DIR DD_RUNTIME_CONTAINER_ROOT + + dd_verify_host_context + dd_start_runtime_timer + trap dd_runtime_exit EXIT + trap 'exit 130' INT TERM + + DD_RUNTIME_PREPARED=1 + dd_run_control_phase prepare \ + --runtime-root "${DD_RUNTIME_DIR}" \ + --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" + dd_verify_runtime_manifest \ + "${DD_RUNTIME_MANIFEST}" \ + "${DD_PLAN_SHA256}" \ + "${DD_SHARD_ID}" \ + "attempt-${DD_ATTEMPT_ORDINAL}" + dd_require_plan_secrets + + dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client_preflight + ((${#DD_STEP_IDS[@]} == 1)) + dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + + dd_start_servers + dd_wait_for_role_readiness server + dd_start_endpoints + dd_wait_for_role_readiness endpoint + dd_require_running + dd_run_control_phase ready + + 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}") + dd_wait_for_client "${client_pid}" + dd_require_running + + dd_cleanup_steps + dd_run_control_phase succeed + DD_RUNTIME_FINALIZED=1 +} + +dd_verify_host_context() { + local tool + for tool in bash sha256sum tar jq srun scontrol getent curl; do + command -v "${tool}" >/dev/null || { + printf 'required allocation tool %q is unavailable\n' "${tool}" >&2 + return 69 + } + done + [[ -d ${DD_ATTEMPT_PATH} && ! -L ${DD_ATTEMPT_PATH} ]] + [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] + [[ ${SLURM_JOB_NUM_NODES:-} == 1 && ${SLURM_NODEID:-} == 0 ]] + dd_verify_gpu_count + dd_read_artifacts "${DD_PLAN_PATH}" "${SLURM_ARRAY_TASK_ID}" + local index path digest actual + for ((index = 0; index < ${#DD_ARTIFACT_FIELDS[@]}; index += 2)); do + path=${DD_ARTIFACT_FIELDS[index]} + digest=${DD_ARTIFACT_FIELDS[index + 1]} + actual=$(sha256sum < "${path}") + [[ ${actual%% *} == "${digest}" ]] || { + printf 'allocation artifact digest mismatch: %s\n' "${path}" >&2 + return 65 + } + done +} + +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 + 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_start_step "${DD_RUNTIME_MANIFEST}" "${step_id}" + dd_register_required_pid "${DD_LAST_PID}" + dd_require_running + 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_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 deadline + 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}" + deadline=$((SECONDS + DD_STEP_PROBE_DEADLINE)) + until curl --fail --silent --max-time 1 \ + "http://${DD_STEP_PROBE_HOST}:${DD_STEP_PROBE_PORT}${DD_STEP_PROBE_PATH}" >/dev/null 2>&1; do + dd_require_running + ((SECONDS < deadline)) || { + printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 + return 70 + } + dd_sleep 0.5 + done + done +} + +dd_runtime_exit() { + local status=$? + trap - EXIT INT TERM + set +e + dd_cleanup_steps + if ((DD_RUNTIME_PREPARED == 1 && DD_RUNTIME_FINALIZED == 0)); then + dd_run_control_phase fail >/dev/null + fi + dd_stop_runtime_timer + exit "${status}" +} 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 index 23315a2c1..1274fe054 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py @@ -44,4 +44,42 @@ def get_container_path( return validate_absolute_path(mapped) -__all__ = ["get_container_path"] +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__ = ["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..d9424183d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -0,0 +1,197 @@ +# 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 .client.host_node_index == 0 + and ([.deployments[].node_indices] | all(. == [0])) + and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) + ' "${plan}" >/dev/null + DD_CLIENT_IMAGE=$(jq -er '.client.image.path' "${plan}") + DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") + DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | 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_artifacts() { + local plan=$1 + local task_id=$2 + dd_read_null_values DD_ARTIFACT_FIELDS < <( + jq -j --argjson task_id "${task_id}" ' + [ + .runtime_bundle, + .client.dependency_lock, + {path: .client.image.path, sha256: .client.image.sha256}, + (.deployments[] | {path: .image.path, sha256: .image.sha256}), + .builder.source, + (.shards[] | select(.array_task_index == $task_id) | .input_partition) + ] + | map(select(. != null)) + | unique_by([.path, .sha256]) + | .[] + | .path, "\u0000", .sha256, "\u0000" + ' "${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) + | .image_path, "\u0000", + (.cpus | tostring), "\u0000", + .stdout_path, "\u0000", + .stderr_path, "\u0000", + (.launch_delay_seconds | tostring), "\u0000", + (.readiness.host // ""), "\u0000", + (.readiness.port // "" | tostring), "\u0000", + (.readiness.path // ""), "\u0000", + (.readiness.deadline_seconds // "" | tostring), "\u0000" + ' "${manifest}" + ) + DD_STEP_IMAGE=${DD_STEP_FIELDS[0]} + DD_STEP_CPUS=${DD_STEP_FIELDS[1]} + DD_STEP_STDOUT=${DD_STEP_FIELDS[2]} + DD_STEP_STDERR=${DD_STEP_FIELDS[3]} + DD_STEP_DELAY=${DD_STEP_FIELDS[4]} + DD_STEP_PROBE_HOST=${DD_STEP_FIELDS[5]} + DD_STEP_PROBE_PORT=${DD_STEP_FIELDS[6]} + DD_STEP_PROBE_PATH=${DD_STEP_FIELDS[7]} + DD_STEP_PROBE_DEADLINE=${DD_STEP_FIELDS[8]} + 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_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 index 476d62a07..78366c390 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -18,6 +18,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.ports import allocation_ports _DIGEST_CHUNK_SIZE = 1024 * 1024 _GPU_COUNT_PATTERN = re.compile(r"^(?:gpu(?::[^:]+)?):([0-9]+)$") @@ -38,10 +39,12 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> """Verify every launch-critical fact before model services start.""" try: self._verify_scheduler(context, environment) - self._verify_attempt_directory(context.attempt_directory) - get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) + attempt_directory = Path( + get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) + ) + self.verify_attempt_directory(attempt_directory) self._verify_artifacts(context) - self._verify_ports(context) + self.verify_ports(context, environment) except SlurmRuntimeError: raise except (OSError, ValueError) as error: @@ -74,14 +77,15 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] 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: + 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_attempt_directory(attempt_directory: Path) -> None: + 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( @@ -112,13 +116,17 @@ def _verify_artifacts(context: AllocationContext) -> None: references.extend(reference for reference in optional_references if reference is not None) unique_references = {(reference.path, reference.sha256): reference for reference in references} for reference in unique_references.values(): + if any( + reference.path == mount.source or reference.path.startswith(f"{mount.source}/") + for mount in plan.container_mounts + ): + reference = reference.model_copy(update={"path": get_container_path(plan, reference.path)}) _verify_artifact(reference) @staticmethod - def _verify_ports(context: AllocationContext) -> None: - ports = tuple(port.port for port in context.plan.client.ports) + tuple( - port.port for deployment in context.plan.deployments for port in deployment.ports - ) + def verify_ports(context: AllocationContext, environment: Mapping[str, str]) -> None: + """Verify that every planned one-node port is currently bindable.""" + ports = allocation_ports(context, environment) reservations: list[socket.socket] = [] try: for port in ports: 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 index c41454609..e61f48017 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -27,15 +27,18 @@ def load_complete_client_candidate( context: AllocationContext, attempt: AttemptManifest, + *, + attempt_directory: Path | None = None, ) -> tuple[ClientResult, CandidateOutputManifest]: """Load a complete semantic client result and its digest-bound candidate.""" - client_result = _read_record(context.attempt_directory, _CLIENT_RESULT_NAME, ClientResult) + 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 is not ClientOutcome.COMPLETE: raise SlurmRuntimeError( SlurmRuntimeErrorCode.CLIENT_FAILED, f"client generation finished with outcome {client_result.outcome.value!r}", ) - candidate = _read_record(context.attempt_directory, _CANDIDATE_NAME, CandidateOutputManifest) + candidate = _read_record(record_directory, _CANDIDATE_NAME, CandidateOutputManifest) try: PlanStateValidator(context.plan).validate_client_candidate( context.shard, 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..7dc5aeae3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh @@ -0,0 +1,167 @@ +# 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_CONTAINER_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= + DD_STEP_VISIBLE_GPUS= + DD_SRUN_COMMAND=( + srun + --nodes=1 + --ntasks=1 + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_STEP_CPUS}" + "--container-image=${DD_STEP_IMAGE}" + ) + if ((${#DD_STEP_GPU_INDICES[@]})); then + if [[ ${DD_GPU_REQUEST_MODE} == gres ]]; 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+=( + "--gpus-per-task=${#DD_STEP_GPU_INDICES[@]}" + "--gpu-bind=mask_gpu:${gpu_mask}" + ) + 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 + [[ -z ${DD_CONTAINER_MOUNTS} ]] || DD_SRUN_COMMAND+=("--container-mounts=${DD_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 +} + +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[@]}"}") + return "${status}" +} + +dd_run_control_phase() { + local operation=$1 + shift + local -a command=( + srun + --nodes=1 + --ntasks=1 + --exact + --overlap + --unbuffered + --export=ALL + "--cpus-per-task=${DD_CLIENT_CPUS}" + "--container-image=${DD_CLIENT_IMAGE}" + --gres=none + ) + [[ -z ${DD_CONTAINER_MOUNTS} ]] || command+=("--container-mounts=${DD_CONTAINER_MOUNTS}") + command+=( + --container-env=PYTHONPATH,SLURM_JOB_GPUS + -- + python3 + -m + data_designer.slurm.runtime.entrypoint + "${operation}" + --plan + "${DD_PLAN_CONTAINER_PATH}" + --attempt-dir + "${DD_ATTEMPT_CONTAINER_DIR}" + "$@" + ) + ( + 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 index ed998331a..8aeb96eb2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -10,10 +10,16 @@ from pathlib import Path from typing import Protocol -from pydantic import BaseModel - -from data_designer.slurm.config.environment import LiteralEnvironmentBinding, SecretRef +from data_designer.slurm.config.environment import ( + LiteralEnvironmentBinding, + SecretRef, + collect_secret_environment_names, +) 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 get_container_path @@ -125,21 +131,22 @@ def _build_step( endpoints: tuple[RuntimeEndpoint, ...], source_environment: Mapping[str, str], ) -> RuntimeStep: - command = _build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) + command = build_client_command(operation, plan, shard, attempt, attempt_directory, endpoints) secret_names, environment = _build_client_environment(plan, source_environment) + allocation_environment = ("SLURM_JOB_GPUS",) if plan.selected_profile.profile.gpu_request_mode == "gres" else () return _build_srun_step( step_id=step_id, role=role, image_path=plan.client.image.path, command=command, environment=environment, - container_environment=secret_names, + container_environment=(*secret_names, *allocation_environment), plan=plan, attempt_directory=attempt_directory, ) -def _build_client_command( +def build_client_command( operation: str, plan: ResolvedSlurmRunPlan, shard: PlannedShard, @@ -191,10 +198,11 @@ def build_vllm_steps( plan: ResolvedSlurmRunPlan, attempt_directory: Path, source_environment: Mapping[str, str], + runtime_root: Path, ) -> tuple[RuntimeStep, ...]: """Build one structured server step per resolved vLLM process.""" return tuple( - _build_vllm_step(deployment, process, plan, attempt_directory, source_environment) + _build_vllm_step(deployment, process, plan, attempt_directory, source_environment, runtime_root) for process in deployment.processes ) @@ -226,7 +234,7 @@ def _build_endpoint_step( host="127.0.0.1", port=deployment.logical_endpoint.port, ) - command = _build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) + command = build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) step = _build_srun_step( step_id=f"{deployment.deployment_id}-endpoint", role=RuntimeStepRole.ENDPOINT, @@ -240,7 +248,7 @@ def _build_endpoint_step( return step, endpoint -def _build_endpoint_command( +def build_endpoint_command( deployment: ResolvedVllmServerDeployment, plan: ResolvedSlurmRunPlan, runtime_proxy_path: Path, @@ -270,10 +278,16 @@ def _build_vllm_step( 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) - environment, container_environment = _build_vllm_environment(deployment, source_environment) + 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, @@ -295,16 +309,18 @@ def _validate_local_vllm_process(process: ResolvedVllmProcess) -> None: ) -def _build_vllm_command( +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", - deployment.model, + model, "--served-model-name", deployment.served_model_name, "--host", @@ -313,6 +329,8 @@ def _build_vllm_command( 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",) @@ -322,6 +340,8 @@ def _build_vllm_command( 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] = [] @@ -339,6 +359,17 @@ def _build_vllm_environment( 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) @@ -438,21 +469,18 @@ def _base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: def _collect_client_secret_environment_names(plan: ResolvedSlurmRunPlan) -> tuple[str, ...]: - names: set[str] = set() - _collect_secret_environment_names(plan.client.authored.dependencies.index_credentials, names) - _collect_secret_environment_names(plan.invocation.authored.mcp_providers, names) - return tuple(sorted(names)) - - -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) + 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/services/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py new file mode 100644 index 000000000..5ae28389c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -0,0 +1,275 @@ +# 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, +) + +_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") + created_at = self._clock() + writer = SlurmStateWriter(self._workspace_root, plan.run_id) + 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, *, failed_at: datetime) -> None: + """Mark every initial attempt failed after its held job is cancelled.""" + 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 + 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/results.py b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py index 51950bdcf..2417b000d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -40,7 +40,7 @@ def validate_state(self) -> SlurmRunExecution: class SlurmPersistedAttemptStatus(ContractValue): - """Persisted attempt state available without scheduler reconciliation.""" + """Persisted attempt state returned after scheduler reconciliation.""" attempt: AttemptManifest readiness: AttemptReadiness | None = None 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 index 9790b877b..68eeaa90a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -66,7 +66,7 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: """Return the durable M2 status for one run.""" def cancel(self, run_id: Identifier) -> SlurmRunCancellation: - """Request cancellation of active jobs without changing persisted state.""" + """Request cancellation of active jobs.""" class SlurmRunService: @@ -179,7 +179,7 @@ def execute_run() -> SlurmRunExecution: return _invoke_service_backend(operation, execute_run) def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: - """Return persisted M2 records without scheduler reconciliation.""" + """Reconcile scheduler observations and return persisted M2 records.""" operation = SlurmServiceOperation.STATUS_RUN normalized_run_id = _validate_run_id(run_id, operation) backend = self._require_backend(operation) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index ddda16ea7..b4d7dbd69 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -6,10 +6,11 @@ from __future__ import annotations import importlib.metadata -from collections.abc import Callable, Iterator +import os +from collections.abc import Callable, Iterator, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime, timezone from pathlib import Path from typing import Protocol, TypeVar from uuid import uuid4 @@ -21,6 +22,8 @@ ClientDependencyResolver, ResolvedClientDependencies, ) +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.filesystem import ensure_private_directory from data_designer.slurm.config import ( DataDesignerSlurmConfig, ImageBuildRequest, @@ -30,6 +33,7 @@ SlurmConfigLoadError, SlurmProfile, SlurmProfileCatalog, + collect_secret_environment_names, load_builder_payload, resolve_profile, ) @@ -46,6 +50,7 @@ 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.services.artifacts import StateRunArtifactPublisher from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.images import SlurmImageService from data_designer.slurm.services.results import ( @@ -59,16 +64,34 @@ from data_designer.slurm.serving.resolver import resolve_vllm_server from data_designer.slurm.state import ( AttemptLifecycleState, + AttemptManifest, + AttemptTerminalClassification, + EffectiveAttemptState, + SchedulerObservation, + SchedulerState, SlurmStateError, SlurmStateWriter, StateConflictError, StateNotFoundError, + reconcile_attempt_observation, ) 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): @@ -88,6 +111,9 @@ def initialize( def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: """Persist the submitted scheduler identity for every initial attempt.""" + def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: + """Mark initial attempts failed after a held submission is cancelled.""" + @dataclass(frozen=True, slots=True) class _PreparedRun: @@ -237,14 +263,16 @@ def __init__( preparer: _RunPreparer, selected_profile: SelectedSlurmProfile, launcher: SlurmCommandClient, - publisher: SlurmRunArtifactPublisher | None, + publisher: SlurmRunArtifactPublisher, clock: Clock, + source_environment: Mapping[str, str], ) -> None: self._preparer = preparer self._profile = selected_profile self._launcher = launcher self._publisher = publisher self._clock = clock + self._source_environment = source_environment def execute( self, @@ -268,12 +296,8 @@ def execute( shard_count=len(plan.shards), batch_script=prepared.batch_script, ) - if self._publisher is None: - raise SlurmServiceError( - SlurmServiceErrorCode.UNAVAILABLE, - SlurmServiceOperation.EXECUTE_RUN, - "run submission is not available; use --dry-run", - ) + self._materialize_default_managed_assets(config, plan) + export_environment = self._build_export_environment(config) publisher = self._publisher self._initialize_run( publisher, @@ -282,7 +306,11 @@ def execute( force=force, ) try: - receipt = self._launcher.submit_script(prepared.batch_script) + receipt = self._launcher.submit_script( + prepared.batch_script, + hold=True, + export_environment=export_environment, + ) except SlurmLauncherError: raise SlurmServiceError( SlurmServiceErrorCode.UNAVAILABLE, @@ -300,13 +328,50 @@ def execute( 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(publisher, plan) + 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(publisher, plan) + 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(publisher, plan) + 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 return SlurmRunExecution( run_id=plan.run_id, state="submitted", @@ -315,6 +380,53 @@ def execute( job_id=receipt.job_id, ) + @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"} + 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, publisher: SlurmRunArtifactPublisher, @@ -377,13 +489,29 @@ def _record_submission( "submission state cannot be recorded", ) from None + def _record_submission_failure(self, publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan) -> None: + try: + publisher.record_submission_failure(plan, 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 writer.load_shards(): + for shard in persisted_shards: + writer.resume_incomplete_finalization(shard.shard_id, published_at=self._clock()) attempts = tuple( SlurmPersistedAttemptStatus( attempt=attempt, @@ -410,16 +538,83 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: 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 + identities = tuple(attempt.scheduler for attempt in active if attempt.scheduler is not None) + try: + queue = {entry.job_identity: entry.state for entry in self._launcher.query_queue(identities)} + missing = tuple(identity for identity in identities if identity not in queue) + accounting = ( + {entry.job_identity: entry.state for entry in self._launcher.query_accounting(missing)} + if missing + else {} + ) + except SlurmLauncherError: + return + for attempt in active: + scheduler = attempt.scheduler + assert scheduler is not None + scheduler_state = queue.get(scheduler, accounting.get(scheduler)) + if scheduler_state is None: + continue + readiness = _load_optional(lambda: writer.load_readiness(attempt.shard_id, attempt.attempt_id)) + observed_at = max( + self._clock(), + attempt.updated_at, + readiness.updated_at if readiness is not None else attempt.updated_at, + ) + observation = SchedulerObservation( + schema_version=1, + scheduler=scheduler, + observed_at=observed_at, + state=scheduler_state, + ) + effective = reconcile_attempt_observation( + attempt, + readiness, + observation, + current_time=observed_at, + ) + update: dict[str, object] = {"updated_at": observed_at} + if effective is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: + update["state"] = AttemptLifecycleState.PENDING + elif effective is EffectiveAttemptState.RUNNING and attempt.state in { + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + }: + update["state"] = AttemptLifecycleState.RUNNING + elif effective is EffectiveAttemptState.FAILED: + update.update( + state=AttemptLifecycleState.FAILED, + terminal_classification=_FAILURE_CLASSIFICATIONS.get( + scheduler_state, + AttemptTerminalClassification.UNKNOWN, + ), + ) + else: + continue + try: + writer.update_attempt(attempt.model_copy(update=update)) + except StateConflictError: + continue + def cancel(self, run_id: Identifier) -> SlurmRunCancellation: status = self.status(run_id) - active = {AttemptLifecycleState.SUBMITTED, AttemptLifecycleState.PENDING, AttemptLifecycleState.RUNNING} 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 and attempt.attempt.scheduler is not None + if attempt.attempt.state in _ACTIVE_ATTEMPT_STATES and attempt.attempt.scheduler is not None } ) ) @@ -502,10 +697,12 @@ def create_slurm_run_service( run_id_factory: RunIdFactory | None = None, clock: Clock | None = None, package_version: str | None = None, + source_environment: Mapping[str, str] | 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) command_client = launcher or SlurmCommandClient() + selected_clock = clock or _utc_now preparer = _RunPreparer( selected, VerifiedImageRegistry(selected.profile.workspace_root), @@ -518,8 +715,9 @@ def create_slurm_run_service( preparer, selected, command_client, - artifact_publisher, - clock or _utc_now, + artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock), + selected_clock, + dict(os.environ if source_environment is None else source_environment), ) return SlurmRunService(_SystemRunPlanner(preparer), render_generation_attempt_script, backend) @@ -558,7 +756,7 @@ def _new_run_id() -> str: def _utc_now() -> datetime: - return datetime.now(UTC) + return datetime.now(timezone.utc) def _format_job_ids(job_ids: list[int]) -> str: 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 index 3a3b6b466..c397e50f6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -12,7 +12,7 @@ from contextlib import ExitStack, contextmanager from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Protocol +from typing import Callable, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.state.filesystem import ( @@ -104,9 +104,12 @@ def rebind(self) -> None: 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 = Path(candidate.dataset_path) + 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) 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 index 474faaeea..38fb75138 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -24,12 +24,22 @@ StateCorruptionError, StateNotFoundError, ) -from data_designer.slurm.state.execution import AttemptLifecycleState, AttemptManifest, RunManifest, ShardManifest +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_shard_winner +from data_designer.slurm.state.validation import ( + StateContractError, + validate_attempt_transition, + validate_shard_winner, +) @dataclass(frozen=True, slots=True) @@ -37,6 +47,8 @@ class _WinnerResolution: winner: ShardWinner candidate: CandidateOutputManifest plan: ResolvedSlurmRunPlan + persisted_attempt: AttemptManifest + successful_attempt: AttemptManifest already_published: bool @@ -46,23 +58,29 @@ class WinnerFinalizer: def __init__(self, storage: StateStorage, reader: StateReader) -> None: self._storage = storage self._reader = reader - self._artifacts = CandidateArtifactVerifier() + self._artifacts = CandidateArtifactVerifier(storage.get_local_path) @contextmanager def acquire_dataset_workspace( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + 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, effective_resume_mode) + 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) -> ShardWinner: + 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) + 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: @@ -124,14 +142,22 @@ def _finalize_with_dataset_lease( 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) + 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, resolution, artifacts) + return self._publish_verified_resolution( + shard_id, + attempt_id, + published_at, + completed_at, + resolution, + artifacts, + ) def _open_candidate_artifacts( self, @@ -148,20 +174,35 @@ def _publish_verified_resolution( 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) + 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: + 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 - self._storage.publish_winner(expected.winner) + 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( @@ -169,46 +210,98 @@ def _resolve_under_state_locks( 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) + 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) - attempt = self._reader.get_attempt(attempts, attempt_id) + 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(attempt) + 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, True) - winner = self._build_winner(run, shard, attempt, client_result, published_at) - self._validate_finalization_chain(run, plan, shard, attempt, client_result, candidate, winner) - return _WinnerResolution(winner, candidate, plan, False) + 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, - effective_resume_mode: Literal["never", "always"], + 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, effective_resume_mode) - dataset_path = self._storage.ensure_dataset_directory(shard_id, attempt_id, effective_resume_mode) + 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 = ( - Path(shard.resume_workspace.path) - if effective_resume_mode == "always" - else self._storage.get_attempt_path(shard_id, attempt_id) / "dataset" + 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") @@ -218,10 +311,10 @@ def _prepare_workspace_or_normalize( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> Path: try: - return self._prepare_dataset_workspace(shard_id, attempt_id, effective_resume_mode) + return self._prepare_dataset_workspace(shard_id, attempt_id, resume_mode) except (StateConflictError, StateCorruptionError, StateNotFoundError, SlurmStateError): raise except (PlanStateContractError, StateContractError) as error: @@ -309,7 +402,7 @@ def _validate_artifact_metadata( def _validate_workspace_mode( plan: ResolvedSlurmRunPlan, attempt: AttemptManifest, - effective_resume_mode: Literal["never", "always"], + resume_mode: Literal["never", "always", "if_possible"], ) -> None: if attempt.state not in { AttemptLifecycleState.SUBMITTED, @@ -318,8 +411,8 @@ def _validate_workspace_mode( }: raise StateContractError("dataset workspace requires an active submitted attempt") requested = plan.invocation.authored.resume - if requested != "if_possible" and effective_resume_mode != requested: - raise StateContractError("effective resume mode does not match the resolved plan") + if requested != "if_possible" and resume_mode != requested: + raise StateContractError("resume mode does not match the resolved plan") def _get_winner_attempt( self, 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 index 57f0fe4a5..67e23c2a9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -33,8 +33,8 @@ def load_run(self) -> RunManifest: run = self._storage.read_run() if ( run.run_id != self._run_id - or run.authored_config.path != self._storage.authored_config_path.as_posix() - or run.resolved_plan.path != self._storage.resolved_plan_path.as_posix() + 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 @@ -52,7 +52,7 @@ def load_authored_config(self, run: RunManifest | None = None) -> DataDesignerSl 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.authored_config_path.as_posix() + 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") @@ -65,7 +65,7 @@ def load_resolved_plan(self, run: RunManifest | None = None) -> ResolvedSlurmRun 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.resolved_plan_path.as_posix() + 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") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 0d00e7b3b..63c0d2a1a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -145,7 +145,7 @@ def validate_readiness_transition( def reconcile_attempt_observation( attempt: AttemptManifest, - readiness: AttemptReadiness, + readiness: AttemptReadiness | None, scheduler: SchedulerObservation, *, current_time: datetime, @@ -154,14 +154,16 @@ def reconcile_attempt_observation( _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") - _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") + 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") - _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") + if readiness is not None: + _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") if scheduler.state in _SCHEDULER_FAILURE_STATES: return EffectiveAttemptState.FAILED @@ -180,7 +182,7 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if scheduler.state is SchedulerState.UNKNOWN: return EffectiveAttemptState.UNKNOWN - if readiness.state is ReadinessState.FAILED: + if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED if scheduler.state is SchedulerState.PENDING: return EffectiveAttemptState.PENDING diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 618c156ba..f940c3238 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -11,13 +11,13 @@ from collections.abc import Iterator from contextlib import ExitStack, contextmanager from pathlib import Path -from typing import Literal, TypeVar +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 +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 @@ -48,22 +48,34 @@ _CANDIDATE_OUTPUT_FILENAME = "output-manifest.json" _WINNER_FILENAME = "winner.json" _DATASET_DIRECTORY_NAME = "dataset" +_RUNTIME_DIRECTORY_NAME = "runtime" _RESUME_LOCK_FILENAME = "resume.lock" _LOCK_DIRECTORY_NAME = ".locks" _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) -> None: + 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: @@ -73,6 +85,27 @@ def authored_config_path(self) -> Path: 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 @@ -249,6 +282,13 @@ def publish_attempt(self, attempt: AttemptManifest) -> None: ) as attempt_descriptor: self._publish_immutable_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) + def ensure_runtime_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> None: + attempt_root = self.get_attempt_path(shard_id, attempt_id) + runtime_root = attempt_root / _RUNTIME_DIRECTORY_NAME + with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: + ensure_private_child_directory(attempt_descriptor, _RUNTIME_DIRECTORY_NAME, runtime_root) + sync_directory(attempt_descriptor) + 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) @@ -284,11 +324,20 @@ def ensure_dataset_directory( 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) - return 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) - return 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, 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 index 95766a8e1..9076e1955 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from datetime import datetime, timedelta from pathlib import Path @@ -23,7 +23,7 @@ StateCorruptionError, StateNotFoundError, ) -from data_designer.slurm.state.execution import AttemptManifest, RunManifest, ShardManifest +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, ShardWinner from data_designer.slurm.state.plan_validation import PersistedPlanStateValidator, PlanStateContractError @@ -53,17 +53,34 @@ class SlurmStateWriter: audit snapshots. Args: - workspace_root: Selected compute-visible workspace root. + 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) -> None: + 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) + 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) @@ -177,17 +194,17 @@ def acquire_dataset_workspace( self, shard_id: ShardId, attempt_id: AttemptId, - effective_resume_mode: Literal["never", "always"], + 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(effective_resume_mode) is not str or effective_resume_mode not in {"never", "always"}: - raise StateConflictError("effective resume mode must be 'never' or 'always'") + 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, - effective_resume_mode, + resume_mode, ) as dataset_path: yield dataset_path @@ -197,8 +214,9 @@ def finalize_winner( attempt_id: AttemptId, *, published_at: datetime, + completed_at: datetime | None = None, ) -> ShardWinner: - """Validate attempt-local artifacts and publish one immutable winner.""" + """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 ( @@ -207,12 +225,51 @@ def finalize_winner( or published_at.utcoffset() != timedelta(0) ): raise StateConflictError("winner publication timestamp must be timezone-aware UTC") - return self._finalizer.finalize_winner(normalized_shard_id, normalized_attempt_id, published_at) + 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) @@ -221,6 +278,7 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes if existing is not None: if existing != attempt: raise StateConflictError(f"attempt {attempt.attempt_id!r} already contains different state") + self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) return existing self._finalizer.require_no_winner(run, plan, shard, shard_attempts) @@ -230,6 +288,7 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes self._reader.validate_attempt_against_plan(run, plan, shard, attempt) validate_shard_attempt_set(run, shard, shard_attempts + (attempt,)) self._storage.publish_attempt(attempt) + self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) return attempt def _update_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: @@ -325,16 +384,16 @@ def _validate_initial_bindings( ) -> 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.workspace_root.as_posix(): + 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.authored_config_path.as_posix(): + 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.resolved_plan_path.as_posix() + 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") diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index 343225a5d..95220553e 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -18,6 +18,7 @@ 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" @@ -104,6 +105,11 @@ def create( 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) -> ClientWorkerCase: workspace = tmp_path / "workspace" @@ -118,6 +124,9 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: lock_payload["python_abi"] = python_abi 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"] ) @@ -147,5 +156,10 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: for distribution in lock.image_distributions ), ) - endpoints = {plan.deployments[0].authored.model_alias: f"http://127.0.0.1:{plan.client.ports[0].port}/v1"} + 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, prepared, endpoints) diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 9578adc7d..c42e284fd 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -9,16 +9,19 @@ import subprocess import sys from pathlib import Path +from typing import cast from unittest.mock import Mock import pytest -from conftest import ClientWorkerCase +from conftest import ClientWorkerCase, FakeDataDesigner from data_designer.slurm.client.environment import ClientEnvironmentBuilder, 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.contracts import InstalledDistribution +from data_designer.slurm.config import SlurmProfile +from data_designer.slurm.contracts import InstalledDistribution, compute_canonical_json_sha256 from data_designer.slurm.planning import ResolvedSlurmRunPlan @@ -37,6 +40,97 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: assert prepared.installed_distributions == client_worker_case.lock.image_distributions +def test_client_runs_through_non_identity_workspace_mount(client_worker_case: ClientWorkerCase) -> None: + physical_workspace = client_worker_case.plan_path.parents[2] + logical_workspace = "/host/workspace" + payload = cast( + dict[str, object], + json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), + ) + selected = cast(dict[str, object], payload["selected_profile"]) + profile_payload = cast(dict[str, object], selected["profile"]) + mount = {"source": logical_workspace, "target": physical_workspace.as_posix(), "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.image_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, + ) + 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(logical_workspace) + assert result.candidate_output_manifest is not None + assert result.candidate_output_manifest.path.startswith(logical_workspace) + + +def test_environment_maps_each_artifact_through_nested_mounts( + 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.image_distributions if path is None else () + + prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( + physical_plan, + shard_id=shard_id, + attempt_id="attempt-0001", + attempt_dir=attempt_dir, + ) + + assert prepared.attempt_dir == attempt_dir + + def test_inspect_distributions_omits_path_for_active_environment(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[dict[str, object]] = [] diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 7d4089df6..c2175577b 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -12,8 +12,10 @@ from types import SimpleNamespace from unittest.mock import Mock +import pyarrow as pa +import pyarrow.parquet as pq import pytest -from conftest import ClientWorkerCase, FakeDataDesigner +from conftest import ClientWorkerCase, FakeCreationResults, FakeDataDesigner import data_designer.slurm.client.worker as worker_module from data_designer.config import ResumeMode @@ -30,7 +32,7 @@ ) 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 +from data_designer.slurm.state import CandidateOutputManifest, compute_candidate_schema_digest def test_preflight_materializes_endpoint_and_ready_environment(client_worker_case: ClientWorkerCase) -> None: @@ -58,6 +60,18 @@ def factory(**kwargs: object) -> FakeDataDesigner: ) +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() @@ -167,6 +181,46 @@ def test_run_persists_semantic_result_and_candidate( assert progress.phase is ClientProgressPhase.COMPLETE +def test_run_normalizes_parquet_schema_metadata( + client_worker_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def export_with_metadata(results: FakeCreationResults, path: Path, *, format: str) -> Path: + assert format == "parquet" + schema = pa.schema( + [("value", pa.int64())], + metadata={b"producer": b"data-designer"}, + ) + table = pa.Table.from_arrays( + [pa.array(range(results.actual_num_records))], + schema=schema, + ) + pq.write_table(table, path) + return path + + monkeypatch.setattr(FakeCreationResults, "export", export_with_metadata) + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + 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/part-00000.parquet") + assert candidate.dataset_schema_digest == compute_candidate_schema_digest(schema) + + @pytest.mark.parametrize( ("failure", "expected_code"), ( @@ -438,7 +492,7 @@ def test_preflight_rejects_plugin_secondary_model_alias( ClientWorker().preflight( plan_path, prepared=prepared, - endpoints={"generator": "http://127.0.0.1:17000/v1"}, + endpoints={"generator": "http://127.0.0.1:10000/v1"}, plugins=plugins, ) except ClientWorkerError as error: 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 index 4d257a8fb..a16835e3f 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -419,13 +419,14 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json index 7fc47e93d..baf71e3b7 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json +++ b/packages/data-designer-slurm/tests/contracts/golden/profile_catalog.json @@ -15,6 +15,7 @@ }, "scheduler": { "account": "lab", + "bin_path": null, "mem_per_gpu": null, "partition": "gpu" }, @@ -43,6 +44,7 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, 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 index 19b781d06..2ac3d4def 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -265,13 +265,14 @@ }, "scheduler": { "account": "research", + "bin_path": null, "mem_per_gpu": null, "partition": "batch" }, "schema_version": 1, "workspace_root": "/workspace/primary" }, - "profile_sha256": "80e50266bc5748469c0429533731e1973856f6f58259d89ab6cde6d4dc0e5bc0", + "profile_sha256": "657ae8b87f9f6d271b059e0987b571a83f06622fab1633447dae1326f28f232d", "schema_version": 1, "selection_source": "explicit" }, diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index c0c67c00a..29a7f2dda 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -98,6 +98,7 @@ def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfil 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"), diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index e9933b74c..6798f7507 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" + "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604", + "provenance_digest": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" + "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e9a71197aa03233d233c9d208f8f13d076fa1e8fff9762365d61e7ad476a1e3f" + "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 183680d34..55b8258ff 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -4,7 +4,7 @@ from __future__ import annotations import subprocess -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import pytest from slurm_test_fakes import FakeCommandResponse, FakeSlurmJob, FakeSlurmRunner @@ -40,6 +40,38 @@ def test_client_submits_verified_script_text_through_standard_input() -> None: 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=NIL"), + ("scontrol", "release", "5101"), + ] + + +def test_client_exports_only_explicit_environment_names_without_values_in_argv() -> None: + runner = _EnvironmentRunner() + client = SlurmCommandClient(runner) + + client.submit_script( + "#!/usr/bin/env bash\n", + export_environment={"HF_TOKEN": "secret-value", "SLURM_EXPORT_ENV": "ALL"}, + ) + + assert runner.command == ( + "sbatch", + "--parsable", + "--export=HF_TOKEN,SLURM_EXPORT_ENV", + ) + assert runner.environment == {"HF_TOKEN": "secret-value", "SLURM_EXPORT_ENV": "ALL"} + assert "secret-value" not in " ".join(runner.command) + + def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: FakeSlurmRunner) -> None: client = SlurmCommandClient(fake_slurm_runner) client.submit("run.sbatch") @@ -281,6 +313,23 @@ def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: 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]) -> subprocess.CompletedProcess[str]: raise subprocess.TimeoutExpired(command, 30.0) diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 08d93d747..f4ed66d29 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -73,6 +73,7 @@ def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fiel 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 @@ -85,6 +86,17 @@ def test_renderer_emits_mem_per_gpu_for_gres_mode(single_node_plan: ResolvedSlur 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 + + @pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( single_node_plan: ResolvedSlurmRunPlan, @@ -192,6 +204,7 @@ def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_run ) -> None: run_root = tmp_path / "run" run_root.mkdir() + (run_root / "shards/shard-00000/attempts/attempt-0001/runtime").mkdir(parents=True) captured_plan_path = tmp_path / "captured-plan.json" entrypoint_path = tmp_path / "entrypoint.sh" entrypoint_path.write_text( diff --git a/packages/data-designer-slurm/tests/launcher/test_runner.py b/packages/data-designer-slurm/tests/launcher/test_runner.py index b41811178..98c74b636 100644 --- a/packages/data-designer-slurm/tests/launcher/test_runner.py +++ b/packages/data-designer-slurm/tests/launcher/test_runner.py @@ -61,13 +61,18 @@ def fake_run( } -def test_subprocess_runner_default_environment_forwards_only_search_path(monkeypatch: pytest.MonkeyPatch) -> None: +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"} + 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: @@ -110,6 +115,7 @@ def fake_run( 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() diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index 9dc83da6d..475a56ff5 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -4,10 +4,12 @@ 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 cast +from typing import Literal, cast import pytest @@ -20,8 +22,10 @@ AttemptLifecycleState, AttemptManifest, AttemptReadiness, + AttemptTerminalClassification, CandidateOutputManifest, SchedulerIdentity, + ShardWinner, StateNotFoundError, validate_attempt_transition, validate_readiness_transition, @@ -39,6 +43,9 @@ class RuntimeCase: 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) @@ -50,6 +57,7 @@ def publish_attempt_result( 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 @@ -59,6 +67,55 @@ def publish_attempt_result( 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) @@ -75,6 +132,11 @@ def load_readiness(self, shard_id: str, attempt_id: str) -> AttemptReadiness: 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: 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..4478a3690 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -0,0 +1,118 @@ +# 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: + 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..51a56aa2c --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from conftest import RuntimeCase + +from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest +from data_designer.slurm.runtime.models import RuntimeStepRole + + +def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: + context = runtime_case.context + runtime_root = context.attempt_directory / "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, + ) + 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] == ( + "python3", + "-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 diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index a1d67bde8..40174aba1 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -36,8 +36,9 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as archive: names = archive.getnames() assert names[0] == "entrypoint.sh" - assert names[1] == "data_designer/slurm/runtime/slurm-sources.txt" - assert names[2] == "data_designer/slurm/__init__.py" + assert names[1:4] == ["plan_reader.sh", "step_runner.sh", "cleanup.sh"] + assert names[4] == "data_designer/slurm/runtime/slurm-sources.txt" + assert names[5] == "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 @@ -46,7 +47,7 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p assert all(archive.getmember(name).uid == 0 for name in names) entrypoint = archive.extractfile("entrypoint.sh") assert entrypoint is not None - assert b'PYTHONPATH="${runtime_root}"' in entrypoint.read() + assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint.read() def test_runtime_bundle_recursively_collects_and_imports_nested_packages( @@ -62,6 +63,8 @@ def test_runtime_bundle_recursively_collects_and_imports_nested_packages( (source_root / "__init__.py").write_text("") (runtime_root / "__init__.py").write_text("") (runtime_root / "bundle.py").write_text("") + for name in ("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()) @@ -233,6 +236,31 @@ def test_extracted_bundle_runtime_takes_precedence_over_installed_sources(tmp_pa 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..a0d0e7d19 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -0,0 +1,117 @@ +# 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 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.state import ( + AttemptLifecycleState, + AttemptManifest, + RunManifest, + SchedulerIdentity, + ShardManifest, + SlurmStateWriter, +) + + +def test_allocation_context_reads_and_updates_state_through_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) + + context, runtime_writer = load_allocation_context( + plan_path, + attempt_directory, + {"SLURM_ARRAY_TASK_ID": "0", "SLURM_ARRAY_JOB_ID": "4101"}, + ) + runtime_writer.update_attempt(context.attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING})) + with runtime_writer.acquire_dataset_workspace(shard.shard_id, attempt.attempt_id, "never") as dataset_path: + assert dataset_path == attempt_directory / "dataset" + + assert context.attempt_directory.as_posix().startswith(logical_workspace) + assert host_writer.load_attempt(shard.shard_id, attempt.attempt_id).state is AttemptLifecycleState.RUNNING diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index de4b363c5..45ff94b19 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib from dataclasses import dataclass, replace from datetime import timedelta from pathlib import Path @@ -12,6 +13,7 @@ from conftest import FakeClientStepBuilder, FakePreflight, FakeStateStore, RuntimeCase 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 @@ -34,6 +36,9 @@ ShardManifest, SlurmStateWriter, ) +from data_designer.slurm.state.artifacts import compute_candidate_schema_digest + +_ALLOCATION_ENVIRONMENT = {"SLURM_JOB_GPUS": "0"} @dataclass(slots=True) @@ -149,7 +154,12 @@ def _supervisor( 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) - runner = _FakeRunner(generation_hook=lambda: _write_complete_result(runtime_case, clock)) + + 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, @@ -160,7 +170,7 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -168,6 +178,9 @@ def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_c 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, @@ -198,7 +211,7 @@ def test_controller_publishes_result_before_success_with_real_state_writer( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -207,6 +220,36 @@ def test_controller_publishes_result_before_success_with_real_state_writer( 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_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: RuntimeCase) -> None: @@ -223,7 +266,7 @@ def test_preflight_failure_starts_no_process_and_fails_attempt(runtime_case: Run client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="injected preflight failure") as raised: @@ -288,7 +331,7 @@ def test_requeued_running_attempt_publishes_restart_epoch_and_uses_fresh_logs( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) result = controller.run() @@ -322,7 +365,7 @@ def test_required_server_exit_fails_and_cleans_partial_start(runtime_case: Runti client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="required runtime step"): @@ -353,7 +396,7 @@ def write_readiness(self, readiness: AttemptReadiness) -> AttemptReadiness: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="required runtime step") as raised: @@ -377,7 +420,7 @@ def test_readiness_timeout_fails_and_terminates_server(runtime_case: RuntimeCase client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=False, clock=clock, advance_on_failure=10_000), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="readiness timed out"): @@ -413,7 +456,7 @@ def test_managed_step_failure_fails_attempt_and_cleans_started_services( client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="status 23"): @@ -450,7 +493,7 @@ def test_cleanup_failure_prevents_false_success(runtime_case: RuntimeCase) -> No client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="cleanup failed"): @@ -479,7 +522,7 @@ def test_cleanup_failure_is_retained_as_a_note_on_the_primary_failure(runtime_ca client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="status 23") as raised: @@ -510,7 +553,7 @@ def test_incomplete_cleanup_does_not_publish_stopped_readiness(runtime_case: Run client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="cleanup failed"): @@ -539,7 +582,7 @@ def write_future_result() -> None: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="later than the allocation clock"): @@ -561,7 +604,7 @@ def test_partial_client_result_is_classified_before_candidate_loading(runtime_ca client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError) as raised: @@ -585,7 +628,7 @@ def test_stale_candidate_from_an_earlier_generation_cannot_succeed(runtime_case: client_steps=FakeClientStepBuilder(), prober=_FakeProber(ready=True, clock=clock), clock=clock, - environment={}, + environment=_ALLOCATION_ENVIRONMENT, ) with pytest.raises(SlurmRuntimeError, match="predates the current generation"): @@ -598,6 +641,11 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> 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) + table = lazy.pa.table({"record_id": range(requested)}) + lazy.pq.write_table(table, output_path) + output_bytes = output_path.read_bytes() candidate = CandidateOutputManifest( schema_version=1, run_id=context.plan.run_id, @@ -612,12 +660,12 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: files=( CandidateOutputFile( relative_path="part-00000.parquet", - sha256="a" * 64, - byte_size=128, + sha256=hashlib.sha256(output_bytes).hexdigest(), + byte_size=len(output_bytes), record_count=requested, ), ), - dataset_schema_digest="b" * 64, + dataset_schema_digest=compute_candidate_schema_digest(table.schema), provenance_digest=context.plan.compute_sha256(), ) candidate_path = context.attempt_directory / "output-manifest.json" diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index b83f83c34..86439bcc8 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -3,13 +3,203 @@ from __future__ import annotations +from pathlib import Path +from types import SimpleNamespace + import pytest +from conftest import FakeStateStore, RuntimeCase + +import data_designer.slurm.runtime.entrypoint as entrypoint +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.state import AttemptLifecycleState, ReadinessState -from data_designer.slurm.runtime.entrypoint import main + +class _InjectedFailure(Exception): + pass def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: - assert main(("--plan", "resolved-plan.json", "--attempt-dir", "attempt-0001")) == 64 + assert ( + entrypoint.main( + ( + "prepare", + "--plan", + "resolved-plan.json", + "--attempt-dir", + "attempt-0001", + "--runtime-root", + "runtime", + "--manifest", + "runtime-manifest.json", + ) + ) + == 64 + ) captured = capsys.readouterr() assert "runtime paths must be absolute" in captured.err assert "Traceback" not in captured.err + + +def test_container_phases_use_the_container_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, {}) + + monkeypatch.setattr(entrypoint, "client_worker_main", lambda arguments: 0) + + def load_candidate(*args: object, attempt_directory: Path | None = None) -> None: + assert attempt_directory == container_attempt_directory + raise _InjectedFailure + + monkeypatch.setattr(entrypoint, "load_complete_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, {}) + + +def test_control_phases_record_running_ready_and_failed( + 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 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 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] + + +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_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: (runtime_case.context, state)) + monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) + monkeypatch.setenv("SLURM_JOB_GPUS", "0") + + +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()) + return arguments diff --git a/packages/data-designer-slurm/tests/runtime/test_paths.py b/packages/data-designer-slurm/tests/runtime/test_paths.py index ebdd820db..7e55a3806 100644 --- a/packages/data-designer-slurm/tests/runtime/test_paths.py +++ b/packages/data-designer-slurm/tests/runtime/test_paths.py @@ -8,7 +8,7 @@ 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 +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: @@ -26,6 +26,7 @@ def test_container_path_uses_most_specific_mount_and_preserves_relative_path(run 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: 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 index e4be62e92..4cc0890e6 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -10,9 +10,11 @@ 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 SystemAllocationPreflight, _verify_artifact @@ -28,6 +30,18 @@ def test_scheduler_preflight_accepts_exact_one_node_gpu_shape(runtime_case: Runt 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"), ( @@ -127,7 +141,7 @@ def close(self) -> None: monkeypatch.setattr("data_designer.slurm.runtime.preflight.socket.socket", _UnavailableSocket) with pytest.raises(SlurmRuntimeError, match="ports are unavailable"): - SystemAllocationPreflight._verify_ports(runtime_case.context) + SystemAllocationPreflight.verify_ports(runtime_case.context, {"SLURM_JOB_GPUS": "0"}) assert _UnavailableSocket.closed @@ -135,4 +149,42 @@ def test_attempt_directory_must_be_restrictive(runtime_case: RuntimeCase) -> Non runtime_case.context.attempt_directory.chmod(0o755) with pytest.raises(SlurmRuntimeError, match="restrictive directory"): - SystemAllocationPreflight._verify_attempt_directory(runtime_case.context.attempt_directory) + SystemAllocationPreflight.verify_attempt_directory(runtime_case.context.attempt_directory) + + +def test_preflight_translates_host_paths_into_container( + runtime_case: RuntimeCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + logical_attempt = runtime_case.context.attempt_directory + container_workspace = tmp_path / "container-workspace" + container_attempt = container_workspace / logical_attempt.relative_to(runtime_case.workspace) + logical_attempt.rmdir() + container_attempt.mkdir(parents=True, mode=0o700) + 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_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(container_workspace.as_posix()) for path in verified_paths) + assert not any(path.startswith(runtime_case.workspace.as_posix()) for path in verified_paths) 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..ef877ff6a --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -0,0 +1,312 @@ +# 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 subprocess +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" + 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_root, artifacts, gpu_request_mode) + 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\nexit 0\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 +readonly DD_PLAN={shlex.quote(plan_path.as_posix())} +readonly DD_ATTEMPT_DIR={shlex.quote(attempt_directory.as_posix())} +source {shlex.quote((runtime_root / "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", + "SLURM_ARRAY_JOB_ID": "4101", + "SLURM_ARRAY_TASK_ID": "0", + "SLURM_JOB_NUM_NODES": "1", + "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 marker_path.exists() + assert "supersecret" not in completed.stdout + completed.stderr + manifest_path.read_text() + + +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 ("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 + + +@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 _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, +) -> 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]}, + "image": {"path": client[0], "sha256": client[1]}, + }, + "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]: + _, _, client, 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", client[0], "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", + client[0], + "fake-endpoint", + readiness={"host": "127.0.0.1", "port": 17000, "path": "/health", "deadline_seconds": 2}, + ), + _step( + attempt_directory, + "client-generation", + "client", + client[0], + "fake-client", + secret_environment={"SOURCE_TOKEN": "SOURCE_TOKEN"}, + container_environment=["SOURCE_TOKEN"], + ), + ], + } + + +def _step( + attempt_directory: Path, + step_id: str, + role: str, + image_path: str, + 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, + "image_path": image_path, + "command": [command], + "cpus": 1, + "gpu_indices": gpu_indices or [], + "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, + } + + +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-env=PYTHONPATH,SLURM_JOB_GPUS* ]] + 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 index b01ed41aa..bda517d92 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -3,20 +3,27 @@ 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 -from data_designer.slurm.contracts import ArtifactReference +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.errors import SlurmRuntimeError from data_designer.slurm.runtime.models import RuntimeEndpoint, RuntimeStepRole 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 @@ -37,6 +44,7 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti context.plan, context.attempt_directory, source_environment, + context.attempt_directory / "runtime", ) endpoint_steps = build_endpoint_steps( (deployment,), @@ -82,11 +90,44 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti 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) +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" + + 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) @@ -203,12 +244,13 @@ def test_client_receives_client_secrets_without_server_only_secrets( { "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 "--container-env=PACKAGE_INDEX_TOKEN" in step.command + assert "--container-env=PACKAGE_INDEX_TOKEN,SLURM_JOB_GPUS" in step.command with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): DefaultClientStepBuilder().build_preflight_step( diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index cd8a09d2a..8e5c1976a 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -4,7 +4,8 @@ from __future__ import annotations import hashlib -from datetime import UTC, datetime +from collections.abc import Mapping +from datetime import datetime, timezone from pathlib import Path import pytest @@ -14,13 +15,20 @@ BuilderInput, DataDesignerSlurmConfig, ImageBuildRequest, + SecretRef, SlurmProfile, SlurmProfileCatalog, ) 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.models import SlurmJobSubmissionReceipt +from data_designer.slurm.launcher.errors import SlurmLauncherError +from data_designer.slurm.launcher.models import ( + SlurmAccountingEntry, + SlurmJobSubmissionReceipt, + SlurmProcessExitCode, + SlurmQueueEntry, +) from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services import ( SlurmServiceError, @@ -31,23 +39,46 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptTerminalClassification, RunManifest, SchedulerIdentity, + SchedulerState, ShardManifest, SlurmStateWriter, StateConflictError, + StateNotFoundError, ) class _Launcher: - def __init__(self, gpu_counts: tuple[int, ...] = (), *, cancel_error: Exception | None = None) -> None: + def __init__( + self, + gpu_counts: tuple[int, ...] = (), + *, + cancel_error: Exception | None = None, + release_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.gpu_counts = gpu_counts self.cancel_error = cancel_error + self.release_error = release_error - def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + 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 {})) return SlurmJobSubmissionReceipt(job_id=42) def cancel(self, job_id: int) -> None: @@ -55,10 +86,23 @@ def cancel(self, job_id: int) -> None: 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 + class _Publisher: def __init__( @@ -69,6 +113,7 @@ def __init__( ) -> None: self.initializations: list[tuple[str, bool]] = [] self.submissions: list[tuple[str, int, datetime]] = [] + self.submission_failures: list[tuple[str, datetime]] = [] self.initialization_error = initialization_error self.submission_error = submission_error @@ -93,6 +138,9 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte raise self.submission_error self.submissions.append((plan.run_id, job_id, submitted_at)) + def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: + self.submission_failures.append((plan.run_id, failed_at)) + def _profile(tmp_path: Path, profile_catalog: SlurmProfileCatalog) -> SlurmProfile: return profile_catalog.clusters["primary"].model_copy(update={"workspace_root": tmp_path.as_posix()}) @@ -182,7 +230,7 @@ def test_production_wiring_submits_after_publisher_initialization( _register_images(tmp_path, authored_run_single, single_node_plan) launcher = _Launcher() publisher = _Publisher() - submitted_at = datetime(2026, 9, 8, tzinfo=UTC) + 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] @@ -199,9 +247,13 @@ def test_production_wiring_submits_after_publisher_initialization( assert publisher.initializations == [("run-wired", True)] 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"}] + assert (tmp_path / "managed-assets").is_dir() -def test_production_wiring_stops_before_submission_without_state_publisher( +def test_production_publisher_rejects_force_before_submission( tmp_path: Path, profile_catalog: SlurmProfileCatalog, authored_run_single: DataDesignerSlurmConfig, @@ -216,14 +268,137 @@ def test_production_wiring_stops_before_submission_without_state_publisher( package_version="0.9.2", ) - with pytest.raises(SlurmServiceError) as caught: - service.execute(authored_run_single, source_root=tmp_path) + with pytest.raises(SlurmServiceError, match="different inputs") as caught: + service.execute(authored_run_single, source_root=tmp_path, force=True) - assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE - assert str(caught.value) == "run submission is not available; use --dry-run" + assert caught.value.code is SlurmServiceErrorCode.CONFLICT assert launcher.submissions == [] +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", "SOURCE_TOKEN": "secret-value"}] + + +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_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_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -291,6 +466,81 @@ def test_recording_conflict_cancels_the_accepted_job( assert caught.value.code is SlurmServiceErrorCode.CONFLICT assert launcher.cancellations == [42] + assert len(publisher.submission_failures) == 1 + + +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( @@ -352,7 +602,7 @@ def test_production_status_and_cancel_use_only_persisted_m2_records( package_version="0.9.2", ) plan = service.plan(authored_run_single) - now = datetime(2026, 9, 8, tzinfo=UTC) + 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(), diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 3f75a41b6..c38572c6f 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="902919292da35aca426a191bee775acf140b0f8b7440b55887fb2ba60ed98a52" +readonly DD_PLAN_SHA256="69ff0f0e29331a513d732f9ad9f5fc979f5ed833b75061d8fdf769feff96ab11" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" @@ -33,10 +33,9 @@ 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}" -install -d -m 0700 "${DD_ATTEMPT_DIR}" -DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" +readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" +[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" -source "${DD_RUNTIME_DIR}/entrypoint.sh" +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 index 553fadd86..72a565bd4 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -13,7 +13,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="cc5ce7314c57fa4e784e149d6d47aafce721118e9f249d6d91518dc20c8fb604" +readonly DD_PLAN_SHA256="b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" @@ -33,10 +33,9 @@ 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}" -install -d -m 0700 "${DD_ATTEMPT_DIR}" -DD_RUNTIME_DIR="$(mktemp -d "${DD_ATTEMPT_DIR}/runtime.${DD_RUNTIME_SHA256}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_DIR}" +readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" +[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] +tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" -source "${DD_RUNTIME_DIR}/entrypoint.sh" +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/slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py index 5a8494cfc..2350b5ab2 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/slurm.py @@ -110,6 +110,7 @@ def run( *, 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: @@ -117,6 +118,7 @@ def run( 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: diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 621d4a33e..6c5578a70 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="8238dd22393a46538169a83046e91230166aaa0f96e0a91255adf6a6074bfa20", + expected_fixture_sha256="3523fe78e80edd4b4acd1e090d9888b3a9879430f1460fecd3576755a4ca923e", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="6153c1769d9b90315c4e7acc767bd5c469d0a5e1f5fc7843baddc1a733cb8e9b", + expected_fixture_sha256="e69a4303113764670742d8d0ce12b1ebe4cc80751734bf68341b535b578108bf", ) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 6599cf992..057ba5bf9 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "1a51941bbb2d8aac7114c197f51105104297de1c6a367cfab0f4111b5d7eebaa" +_COMPATIBILITY_FIXTURE_DIGEST = "16460710f383bb9f4650f17e3bac95e4e4a495ca3d80873d7282aa3953533f97" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 8557b113e..9a59c86cd 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -23,9 +23,10 @@ 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 DataDesignerSlurmConfig, SlurmProfile +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, @@ -187,6 +188,9 @@ def track_read( attempt = _submitted_attempt(case) case.writer.create_attempt(attempt) + runtime_directory = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/runtime" + assert runtime_directory.is_dir() + assert runtime_directory.stat().st_mode & 0o777 == 0o700 readiness = _readiness(case, attempt) case.writer.write_readiness(readiness) record_names.clear() @@ -1236,6 +1240,208 @@ def test_result_publication_binds_candidate_before_success_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, @@ -1963,6 +2169,25 @@ def second_writer() -> Path: 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, @@ -2217,6 +2442,7 @@ def _persist_complete_result( 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( @@ -2236,6 +2462,7 @@ def _persist_complete_result( 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, @@ -2243,7 +2470,7 @@ def _persist_complete_result( attempt_id=attempt.attempt_id, attempt_ordinal=attempt.attempt_ordinal, created_at=case.created_at + timedelta(minutes=3), - dataset_path=dataset_path.as_posix(), + dataset_path=persisted_dataset_path.as_posix(), requested_records=requested_records, actual_records=requested_records, outcome=CandidateOutcome.COMPLETE, @@ -2272,7 +2499,7 @@ def _persist_complete_result( requested_records=requested_records, actual_records=requested_records, outcome=ClientOutcome.COMPLETE, - dataset_path=dataset_path.as_posix(), + dataset_path=persisted_dataset_path.as_posix(), early_shutdown=False, requested_resume_mode=case.plan.invocation.authored.resume, effective_resume_mode="never", From de38564325cacceb9c6f7a430b92a1deddb9f219 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 09:58:38 -0600 Subject: [PATCH 21/62] feat: reconcile persisted Slurm state (#913) * feat: reconcile persisted Slurm state Persist normalized scheduler observations and compose fresh-process run, shard, attempt, readiness, generation, and winner status. Preserve bounded accounting lag and immutable terminal evidence for status and benchmark refresh consumers. Part of #869 Signed-off-by: Nabin Mulepati * fix(slurm): harden state reconciliation Signed-off-by: Nabin Mulepati * fix(slurm): bound preemption reconciliation Signed-off-by: Nabin Mulepati * fix(slurm): bound public preemption status Signed-off-by: Nabin Mulepati * feat: add Slurm retry and deterministic collection (#915) * feat(slurm): add retry and deterministic collection Signed-off-by: Nabin Mulepati * fix(slurm): recover ambiguous submissions Signed-off-by: Nabin Mulepati * fix(slurm): preserve collection snapshots Signed-off-by: Nabin Mulepati * fix(slurm): bind retries and collection ownership Signed-off-by: Nabin Mulepati * fix(slurm): preserve scheduler path on retry Render retry allocations with the same profile-configured Slurm command path as initial generation attempts so cluster-local srun and scontrol binaries remain discoverable. Signed-off-by: Nabin Mulepati * fix(slurm): isolate attempt observation clocks Apply persisted timestamp floors per scheduler identity so a clock-ahead sibling cannot expire another attempt's preemption-requeue window. Persist each active attempt with its own reconciled observation time. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/client/execution.py | 43 +- .../src/data_designer/slurm/client/worker.py | 19 +- .../data_designer/slurm/launcher/client.py | 164 +- .../slurm/launcher/collection.py | 103 + .../data_designer/slurm/launcher/errors.py | 12 + .../data_designer/slurm/launcher/models.py | 27 +- .../data_designer/slurm/launcher/parsing.py | 34 +- .../data_designer/slurm/launcher/renderer.py | 130 +- .../data_designer/slurm/runtime/bootstrap.py | 23 +- .../data_designer/slurm/runtime/context.py | 41 +- .../data_designer/slurm/runtime/controller.py | 17 +- .../data_designer/slurm/runtime/entrypoint.py | 43 +- .../data_designer/slurm/runtime/entrypoint.sh | 30 +- .../src/data_designer/slurm/runtime/models.py | 25 +- .../src/data_designer/slurm/runtime/steps.py | 25 +- .../data_designer/slurm/services/wiring.py | 102 +- .../src/data_designer/slurm/state/__init__.py | 50 + .../data_designer/slurm/state/artifacts.py | 128 +- .../slurm/state/attempt_identity.py | 91 + .../src/data_designer/slurm/state/base.py | 5 + .../data_designer/slurm/state/collection.py | 385 ++++ .../slurm/state/collection_filesystem.py | 357 ++++ .../slurm/state/collection_inputs.py | 81 + .../slurm/state/collection_merge.py | 338 ++++ .../slurm/state/collection_records.py | 112 ++ .../slurm/state/collection_storage.py | 339 ++++ .../slurm/state/collection_validation.py | 113 ++ .../slurm/state/collection_worker.py | 292 +++ .../data_designer/slurm/state/destinations.py | 92 + .../data_designer/slurm/state/observation.py | 263 +++ .../src/data_designer/slurm/state/observer.py | 308 +++ .../src/data_designer/slurm/state/outputs.py | 44 + .../src/data_designer/slurm/state/reader.py | 64 +- .../slurm/state/reconciliation.py | 22 +- .../src/data_designer/slurm/state/retry.py | 481 +++++ .../slurm/state/retry_records.py | 76 + .../slurm/state/retry_storage.py | 210 ++ .../data_designer/slurm/state/scheduler.py | 28 +- .../src/data_designer/slurm/state/status.py | 209 ++ .../src/data_designer/slurm/state/storage.py | 54 + .../src/data_designer/slurm/state/store.py | 20 +- .../slurm/state/submission_recovery.py | 85 + .../data_designer/slurm/state/validation.py | 14 +- .../tests/client/test_worker.py | 74 + .../tests/launcher/test_client.py | 107 +- .../tests/launcher/test_collection.py | 205 ++ .../tests/runtime/conftest.py | 9 + .../tests/runtime/test_bootstrap.py | 39 + .../tests/runtime/test_bundle.py | 5 +- .../tests/runtime/test_context.py | 60 + .../tests/runtime/test_controller.py | 44 + .../tests/runtime/test_entrypoint.py | 2 +- .../tests/services/test_wiring.py | 87 +- .../tests/state/test_observer.py | 709 +++++++ .../tests/state/test_retry_collection.py | 1780 +++++++++++++++++ .../tests/state/test_store.py | 16 + .../tests/state/test_submission_recovery.py | 144 ++ .../tests/state/test_validation.py | 2 + scripts/test_slurm_package_install.py | 18 +- 59 files changed, 8222 insertions(+), 178 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/attempt_identity.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/observation.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/observer.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry_records.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/status.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/submission_recovery.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_collection.py create mode 100644 packages/data-designer-slurm/tests/state/test_observer.py create mode 100644 packages/data-designer-slurm/tests/state/test_retry_collection.py create mode 100644 packages/data-designer-slurm/tests/state/test_submission_recovery.py 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 index 5f3fa3184..31011a4dd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -79,6 +79,7 @@ class _ExecutionContext: builder: DataDesignerConfigBuilder designer: DataDesigner requested_resume: ResumeMode + retry_resume: ResumeMode | None dataset_path: Path @@ -169,10 +170,16 @@ def preflight( 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) + 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) @@ -200,12 +207,18 @@ def run( 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) + 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) @@ -245,6 +258,7 @@ def _build_context( *, prepared: PreparedClientEnvironment, endpoints: Mapping[str, str], + retry_resume: ResumeMode | None = None, ) -> _ExecutionContext: try: plan = ResolvedSlurmRunPlan.model_validate_json( @@ -283,7 +297,14 @@ def _build_context( mcp_providers = self._materialize_mcp_providers(plan) managed_assets_path = self._validate_assets(plan) requested_resume = ResumeMode(plan.invocation.authored.resume) - dataset_path = self._dataset_path(plan, shard, prepared, requested_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, @@ -292,7 +313,7 @@ def _build_context( auto_configure_logging=False, ) designer.set_run_config(RunConfig.model_validate(plan.invocation.effective_run_config)) - return _ExecutionContext(plan, shard, builder, designer, requested_resume, dataset_path) + return _ExecutionContext(plan, shard, builder, designer, requested_resume, retry_resume, dataset_path) except ClientWorkerError: raise except Exception as error: @@ -485,9 +506,10 @@ def _prepare_dataset_workspace( 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") - if context.requested_resume is not ResumeMode.NEVER: + 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 context.requested_resume is ResumeMode.NEVER and context.dataset_path.exists(): + 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") @@ -503,7 +525,7 @@ def _generate_dataset( context.builder, num_records=context.shard.requested_records, dataset_name=context.dataset_path.name, - resume=context.requested_resume, + 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, ), @@ -538,11 +560,16 @@ def _validate_creation_result( 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") - if results.requested_resume_mode is not context.requested_resume: + 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(get_container_path(context.plan, context.shard.resume_workspace.path, require_writable=True)) expected_path = shared_path if effective_resume is ResumeMode.ALWAYS else prepared.attempt_dir / "dataset" if ( 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 index 5139b10c0..d34fb5dcf 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py @@ -11,6 +11,7 @@ from datetime import datetime, timezone from pathlib import Path +from data_designer.config import ResumeMode from data_designer.slurm.client.environment import ( ClientEnvironmentBuilder, PreparedClientEnvironment, @@ -47,10 +48,23 @@ def main(argv: Sequence[str] | None = None) -> int: 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) + 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) + 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: @@ -78,6 +92,7 @@ def _parse_arguments(argv: Sequence[str] | None) -> argparse.Namespace: 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) 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 index 2d822a7f1..58e203105 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -5,35 +5,43 @@ from __future__ import annotations +import os import re import subprocess import unicodedata from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timedelta from pathlib import Path -from typing import TypeAlias from data_designer.slurm.contracts import Identifier -from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, - SlurmObservedJobIdentity, + SlurmNamedJobEntry, SlurmQueueEntry, + SlurmSubmissionMatch, ) from data_designer.slurm.launcher.parsing import ( parse_accounting, parse_gpu_counts, + parse_named_jobs, parse_queue, parse_submission, ) from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner -from data_designer.slurm.state import SchedulerIdentity +from data_designer.slurm.state import SchedulerIdentity, SchedulerJobIdentity -_JobSelector: TypeAlias = int | SchedulerIdentity _IDENTIFIER_PATTERN = re.compile(r"^[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) @@ -101,26 +109,40 @@ def submit_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) - output = self._run( - (self._executables.sbatch, "--parsable", *hold_arguments, export_argument), - input_text=script, - environment=environment, - ) - return parse_submission(output) + 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[_JobSelector]) -> tuple[SlurmQueueEntry, ...]: + 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) - output = self._run( - ( - self._executables.squeue, - "--noheader", - "--array", - "--format=%i|%T", - f"--jobs={jobs}", + 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), @@ -129,7 +151,7 @@ def query_queue(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmQueueEntr ) return tuple(entry for entry in entries if entry.job_identity not in ignored) - def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAccountingEntry, ...]: + def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[SlurmAccountingEntry, ...]: """Return normalized accounting rows for explicit managed jobs.""" requested = tuple(selectors) jobs = _format_selectors(requested) @@ -152,7 +174,48 @@ def query_accounting(self, selectors: Sequence[_JobSelector]) -> tuple[SlurmAcco ) return tuple(entry for entry in entries if entry.job_identity not in ignored) - def cancel(self, selector: _JobSelector) -> None: + 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))) @@ -186,13 +249,26 @@ def _run( ) else: completed = self._runner.run(command, input_text=input_text, environment=environment) - except (OSError, subprocess.SubprocessError) as error: + 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") + 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}") @@ -212,13 +288,13 @@ def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[s return f"--export={','.join(names)}", environment -def _format_selectors(selectors: Sequence[_JobSelector]) -> str: +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: _JobSelector) -> str: +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: @@ -234,16 +310,16 @@ def _format_job_id(value: object) -> str: def _validate_observed_job_identities( - job_identities: Sequence[SlurmObservedJobIdentity], - selectors: Sequence[_JobSelector], + job_identities: Sequence[SchedulerJobIdentity], + selectors: Sequence[SchedulerJobIdentity], *, command: str, -) -> frozenset[SlurmObservedJobIdentity]: +) -> 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[SlurmObservedJobIdentity] = set() + ignored: set[SchedulerJobIdentity] = set() for job_identity in job_identities: if type(job_identity) is int and job_identity in selected_job_ids: continue @@ -275,3 +351,31 @@ 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..5072d14d9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -0,0 +1,103 @@ +# 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 pathlib import PurePosixPath + +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") + if collection_plan.container_destination != destination.container_path: + raise SlurmBatchRenderError("collection container 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"), + ) + ) + workspace_root = resolved_plan.selected_profile.profile.workspace_root + state_mount = f"{workspace_root}:{workspace_root}" + output_mount = f"{destination.mount.source}:{destination.mount.target}" + mount_arguments = _render_mount_arguments( + ("DD_STATE_MOUNT", workspace_root, workspace_root), + ("DD_OUTPUT_MOUNT", destination.mount.source, destination.mount.target), + ) + return f"""#!/usr/bin/env bash +{directives} +set -Eeuo pipefail +export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" + +readonly DD_CLIENT_IMAGE={quote_shell_value(resolved_plan.client.image.path)} +readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(resolved_plan.client.image.sha256)} +readonly DD_COLLECTION_PLAN={quote_shell_value(collection_plan_path)} +readonly DD_COLLECTION_PLAN_SHA256={quote_shell_value(collection_plan.compute_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)} +readonly DD_STATE_MOUNT={quote_shell_value(state_mount)} +readonly DD_OUTPUT_MOUNT={quote_shell_value(output_mount)} + +verify_sha256() {{ + local actual_sha256 + actual_sha256="$(sha256sum < "$2")" + [[ "${{actual_sha256%% *}}" == "$1" ]] +}} + +verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" +verify_sha256 "${{DD_COLLECTION_PLAN_SHA256}}" "${{DD_COLLECTION_PLAN}}" +DD_ENROOT_MOUNTS=({mount_arguments}) +readonly DD_ENROOT_MOUNTS +exec enroot start --root "${{DD_ENROOT_MOUNTS[@]}}" "${{DD_CLIENT_IMAGE}}" \ + python -m data_designer.slurm.state.collection_worker \ + --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" --collection-id "${{DD_COLLECTION_ID}}" +""" + + +def _render_mount_arguments(*mounts: tuple[str, str, str]) -> str: + unique: dict[str, tuple[str, str, str]] = {} + targets: dict[str, str] = {} + for variable, source, target in mounts: + mount = f"{source}:{target}" + existing_source = targets.get(target) + if existing_source is not None and existing_source != source: + raise SlurmBatchRenderError("collection state and output mounts cannot share a target") + targets[target] = source + unique.setdefault(mount, (variable, source, target)) + ordered = sorted(unique.values(), key=lambda item: len(PurePosixPath(item[2]).parts)) + return " ".join(f'--mount "${{{variable}}}"' for variable, _, _ in ordered) + + +__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 index 4a611152b..6e483d86e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/errors.py @@ -13,6 +13,18 @@ class SlurmLauncherError(RuntimeError): 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.""" 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 index 7d1168255..64111099e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/models.py @@ -6,11 +6,8 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TypeAlias -from data_designer.slurm.state import SchedulerIdentity, SchedulerState - -SlurmObservedJobIdentity: TypeAlias = int | SchedulerIdentity +from data_designer.slurm.state import SchedulerJobIdentity, SchedulerState @dataclass(frozen=True) @@ -32,7 +29,7 @@ class SlurmProcessExitCode: class SlurmQueueEntry: """One transient normalized active-queue entry.""" - job_identity: SlurmObservedJobIdentity + job_identity: SchedulerJobIdentity state: SchedulerState @@ -40,6 +37,24 @@ class SlurmQueueEntry: class SlurmAccountingEntry: """One transient normalized accounting entry.""" - job_identity: SlurmObservedJobIdentity + 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 index adf0e13c2..32766d782 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -6,16 +6,17 @@ 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, - SlurmObservedJobIdentity, + SlurmNamedJobEntry, SlurmProcessExitCode, SlurmQueueEntry, ) -from data_designer.slurm.state import SchedulerIdentity, SchedulerState +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]*$") @@ -71,7 +72,7 @@ def parse_submission(output: str) -> SlurmJobSubmissionReceipt: def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: """Parse ``squeue --format=%i|%T`` rows.""" entries: list[SlurmQueueEntry] = [] - identities: set[SlurmObservedJobIdentity] = set() + identities: set[SchedulerJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 2: @@ -85,7 +86,7 @@ def parse_queue(output: str) -> tuple[SlurmQueueEntry, ...]: def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: """Parse job and array-task rows from ``sacct --format=JobID,State,ExitCode``.""" entries: list[SlurmAccountingEntry] = [] - identities: set[SlurmObservedJobIdentity] = set() + identities: set[SchedulerJobIdentity] = set() for line_number, line in _collect_nonempty_lines(output): fields = line.split("|") if len(fields) != 3: @@ -102,6 +103,25 @@ def parse_accounting(output: str) -> tuple[SlurmAccountingEntry, ...]: 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] = [] @@ -184,7 +204,7 @@ def _parse_array_identity(value: str, *, command: str, line_number: int) -> Sche ) -def _parse_job_identity(value: str, *, command: str, line_number: int) -> SlurmObservedJobIdentity: +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) @@ -218,8 +238,8 @@ def _parse_decimal(value: str, *, message: str) -> int: def _reject_duplicate( - job_identity: SlurmObservedJobIdentity, - identities: set[SlurmObservedJobIdentity], + job_identity: SchedulerJobIdentity, + identities: set[SchedulerJobIdentity], *, command: str, line_number: int, 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 index ad71b0af5..efc16549a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -7,9 +7,11 @@ 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" @@ -25,13 +27,10 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi if plan.selected_profile.profile.gpu_request_mode == "visible": directive_text = f"{directive_text}\n#SBATCH --exclusive" attempt = f"{attempt_ordinal:04d}" - scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path - command_path = _SYSTEM_PATH if scheduler_bin_path is None else f"{scheduler_bin_path}:{_SYSTEM_PATH}" - return f"""#!/usr/bin/env bash {directive_text} set -Eeuo pipefail -export PATH={quote_shell_value(command_path)} +export PATH={quote_shell_value(_get_command_path(plan))} readonly DD_RUNTIME_ARCHIVE={quote_shell_value(plan.runtime_bundle.path)} readonly DD_RUNTIME_SHA256={quote_shell_value(plan.runtime_bundle.sha256)} @@ -65,26 +64,135 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi """ -def _build_generation_directives(plan: ResolvedSlurmRunPlan) -> tuple[tuple[str, str | None], ...]: +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}" + directives = render_batch_directives( + _build_generation_directives(plan, array=array_tasks, job_name=retry.submission_job_name) + ) + 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))} + +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)} +readonly DD_CLIENT_IMAGE={quote_shell_value(plan.client.image.path)} +readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(plan.client.image.sha256)} +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}}" +verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" +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}}" +enroot start --root --mount "${{DD_WORKSPACE_ROOT}}:${{DD_WORKSPACE_ROOT}}" "${{DD_CLIENT_IMAGE}}" \\ + 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}}" +DD_RUNTIME_DIR="$(mktemp -d "${{DD_ATTEMPT_DIR}}/runtime.${{DD_RUNTIME_SHA256}}.XXXXXX")" +readonly DD_RUNTIME_DIR +tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_DIR}}" + +source "${{DD_RUNTIME_DIR}}/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.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, +) -> 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 - array = "0" - if plan.array_tasks.count > 1: - array = f"0-{plan.array_tasks.count - 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: - array = f"{array}%{plan.array_tasks.max_concurrent}" + resolved_array = f"{resolved_array}%{plan.array_tasks.max_concurrent}" values: list[tuple[str, str | None]] = [ - ("job-name", plan.submission.job_name), + ("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)), ("cpus-per-task", str(plan.client.authored.cpus)), ("time", plan.submission.time_limit), - ("array", array), + ("array", resolved_array), ] profile = plan.selected_profile.profile if profile.gpu_request_mode == "gres": 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 index 54a053231..264bfdd4a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -203,6 +203,7 @@ def _build_client_step( log_directory: Path, ) -> 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, @@ -210,15 +211,35 @@ def _build_client_step( 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 = ( "python3", "-m", "data_designer.slurm.runtime.entrypoint", "client", *command[4:6], - *command[10:], + *command[10:12], + *retry_binding, + *endpoint_arguments, ) secret_names = collect_secret_environment_names( (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) 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 index 6817300b4..cbb32b181 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -14,7 +14,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path -from data_designer.slurm.state import SlurmStateWriter +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 @@ -24,6 +24,10 @@ 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 its container paths.""" writer = _load_state_writer(plan_path, attempt_directory) @@ -50,7 +54,40 @@ def load_allocation_context( SlurmRuntimeErrorCode.INVALID_CONTEXT, "scheduler array job does not match the persisted attempt", ) - return AllocationContext(plan, shard, attempt, host_attempt_directory), writer + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index e90cbc555..c8abeda28 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -295,6 +295,7 @@ def _prepare_runtime(self) -> _RuntimeTopology: self._context.attempt_directory, endpoints, self._environment, + retry_resume_mode=self._retry_resume_mode, ) ) self._supervisor.wait(self._supervisor.start(client_preflight)) @@ -346,17 +347,26 @@ def _run_client( 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, - self._context.plan.invocation.authored.resume, + 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_complete_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 - the record contract requires this for complete results @@ -368,6 +378,11 @@ def _run_client( 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, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 66ee59f21..db0462981 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -85,10 +85,13 @@ def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: 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_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) _validate_attempt_is_executable(context.attempt) SystemAllocationPreflight.verify_attempt_directory(arguments.attempt_dir) SystemAllocationPreflight.verify_ports(context, environment) @@ -113,7 +116,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + 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) @@ -143,12 +146,17 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + context, writer = _load_context(arguments, environment) 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, - context.plan.invocation.authored.resume, + resume_mode, ): return_code = client_worker_main( ( @@ -161,6 +169,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No 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)), ) ) @@ -171,6 +180,14 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No 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( @@ -186,7 +203,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + 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: @@ -204,7 +221,7 @@ def _succeed(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: - context, writer = load_allocation_context(arguments.plan, arguments.attempt_dir, environment) + 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 @@ -220,6 +237,20 @@ def _fail(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None ) +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, 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 index dc15efd34..cbba85438 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -10,15 +10,24 @@ source "${DD_RUNTIME_DIR}/cleanup.sh" DD_RUNTIME_PREPARED=0 DD_RUNTIME_FINALIZED=0 +DD_RUNTIME_RETRY_ARGUMENTS=() dd_slurm_run_allocation() { - if [[ $# -ne 2 ]]; then - printf '%s\n' 'allocation runtime requires plan and attempt directory arguments' >&2 + 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=() + if [[ $# -eq 5 ]]; then + DD_RUNTIME_RETRY_ARGUMENTS=( + --retry-id "$3" + --retry-plan-sha256 "$4" + --effective-resume-mode "$5" + ) + fi dd_read_control_plan "${DD_PLAN_PATH}" dd_read_plan_secret_names "${DD_PLAN_PATH}" dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false @@ -38,7 +47,7 @@ dd_slurm_run_allocation() { trap 'exit 130' INT TERM DD_RUNTIME_PREPARED=1 - dd_run_control_phase prepare \ + dd_run_bound_control_phase prepare \ --runtime-root "${DD_RUNTIME_DIR}" \ --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" dd_verify_runtime_manifest \ @@ -57,7 +66,7 @@ dd_slurm_run_allocation() { dd_start_endpoints dd_wait_for_role_readiness endpoint dd_require_running - dd_run_control_phase ready + dd_run_bound_control_phase ready dd_read_step_ids "${DD_RUNTIME_MANIFEST}" client ((${#DD_STEP_IDS[@]} == 1)) @@ -68,10 +77,19 @@ dd_slurm_run_allocation() { dd_require_running dd_cleanup_steps - dd_run_control_phase succeed + dd_run_bound_control_phase succeed DD_RUNTIME_FINALIZED=1 } +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; do @@ -173,7 +191,7 @@ dd_runtime_exit() { set +e dd_cleanup_steps if ((DD_RUNTIME_PREPARED == 1 && DD_RUNTIME_FINALIZED == 0)); then - dd_run_control_phase fail >/dev/null + dd_run_bound_control_phase fail >/dev/null fi dd_stop_runtime_timer exit "${status}" 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 index 50d2557fe..2aa93ab16 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -16,7 +16,7 @@ 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 +from data_designer.slurm.state import AttemptManifest, RetryPlan _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) @@ -117,6 +117,7 @@ class AllocationContext: shard: PlannedShard attempt: AttemptManifest attempt_directory: Path + retry_plan: RetryPlan | None = None def __post_init__(self) -> None: if not isinstance(self.plan, ResolvedSlurmRunPlan): @@ -147,3 +148,25 @@ def __post_init__(self) -> None: 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/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index 8aeb96eb2..74ca803ea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -8,7 +8,7 @@ import os from collections.abc import Mapping from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from data_designer.slurm.config.environment import ( LiteralEnvironmentBinding, @@ -55,6 +55,8 @@ def build_preflight_step( 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.""" ... @@ -67,6 +69,8 @@ def build_generation_step( 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.""" ... @@ -83,6 +87,8 @@ def build_preflight_step( 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( @@ -95,6 +101,7 @@ def build_preflight_step( attempt_directory, endpoints, source_environment, + retry_resume_mode, ) def build_generation_step( @@ -105,6 +112,8 @@ def build_generation_step( 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( @@ -117,6 +126,7 @@ def build_generation_step( attempt_directory, endpoints, source_environment, + retry_resume_mode, ) @staticmethod @@ -130,8 +140,17 @@ def _build_step( 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) + command = build_client_command( + operation, + plan, + shard, + attempt, + attempt_directory, + endpoints, + retry_resume_mode, + ) secret_names, environment = _build_client_environment(plan, source_environment) allocation_environment = ("SLURM_JOB_GPUS",) if plan.selected_profile.profile.gpu_request_mode == "gres" else () return _build_srun_step( @@ -153,6 +172,7 @@ def build_client_command( attempt: AttemptManifest, attempt_directory: Path, endpoints: tuple[RuntimeEndpoint, ...], + retry_resume_mode: Literal["never", "always"] | None, ) -> tuple[str, ...]: endpoint_arguments = tuple( argument @@ -172,6 +192,7 @@ def build_client_command( attempt.attempt_id, "--attempt-dir", get_container_path(plan, attempt_directory.as_posix(), require_writable=True), + *(() if retry_resume_mode is None else ("--resume-mode", retry_resume_mode)), *endpoint_arguments, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index b4d7dbd69..b13834f56 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -65,15 +65,15 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptStatus, AttemptTerminalClassification, EffectiveAttemptState, - SchedulerObservation, SchedulerState, SlurmStateError, + SlurmStateReconciler, SlurmStateWriter, StateConflictError, StateNotFoundError, - reconcile_attempt_observation, ) RunIdFactory = Callable[[], str] @@ -548,63 +548,53 @@ def _reconcile_attempts( ) if not active: return - identities = tuple(attempt.scheduler for attempt in active if attempt.scheduler is not None) try: - queue = {entry.job_identity: entry.state for entry in self._launcher.query_queue(identities)} - missing = tuple(identity for identity in identities if identity not in queue) - accounting = ( - {entry.job_identity: entry.state for entry in self._launcher.query_accounting(missing)} - if missing - else {} + 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, + ), ) - except SlurmLauncherError: + else: + return + try: + writer.update_attempt(attempt.model_copy(update=update)) + except StateConflictError: return - for attempt in active: - scheduler = attempt.scheduler - assert scheduler is not None - scheduler_state = queue.get(scheduler, accounting.get(scheduler)) - if scheduler_state is None: - continue - readiness = _load_optional(lambda: writer.load_readiness(attempt.shard_id, attempt.attempt_id)) - observed_at = max( - self._clock(), - attempt.updated_at, - readiness.updated_at if readiness is not None else attempt.updated_at, - ) - observation = SchedulerObservation( - schema_version=1, - scheduler=scheduler, - observed_at=observed_at, - state=scheduler_state, - ) - effective = reconcile_attempt_observation( - attempt, - readiness, - observation, - current_time=observed_at, - ) - update: dict[str, object] = {"updated_at": observed_at} - if effective is EffectiveAttemptState.PENDING and attempt.state is AttemptLifecycleState.SUBMITTED: - update["state"] = AttemptLifecycleState.PENDING - elif effective is EffectiveAttemptState.RUNNING and attempt.state in { - AttemptLifecycleState.SUBMITTED, - AttemptLifecycleState.PENDING, - }: - update["state"] = AttemptLifecycleState.RUNNING - elif effective is EffectiveAttemptState.FAILED: - update.update( - state=AttemptLifecycleState.FAILED, - terminal_classification=_FAILURE_CLASSIFICATIONS.get( - scheduler_state, - AttemptTerminalClassification.UNKNOWN, - ), - ) - else: - continue - try: - writer.update_attempt(attempt.model_copy(update=update)) - except StateConflictError: - continue def cancel(self, run_id: Identifier) -> SlurmRunCancellation: status = self.status(run_id) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index ce69166e3..1fc65422c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -22,9 +22,16 @@ 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 ( SlurmStateError, StateConflictError, @@ -38,6 +45,12 @@ 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, @@ -46,6 +59,8 @@ CandidateOutputManifest, CollectionPlan, CollectionShard, + RetryPlan, + RetryShard, ShardWinner, ) from data_designer.slurm.state.readiness import ( @@ -61,11 +76,19 @@ 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, @@ -80,9 +103,15 @@ ) 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"), } @@ -92,36 +121,57 @@ "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", "Sha256Digest", "ShardManifest", + "ShardStatus", "ShardId", "ShardWinner", "SlurmStateError", + "SlurmCollectionCoordinator", + "SlurmRetryCoordinator", + "SlurmStateReconciler", "SlurmStateWriter", "StateConflictError", "StateContractError", 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 index c397e50f6..c68bb8738 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/artifacts.py @@ -12,7 +12,7 @@ from contextlib import ExitStack, contextmanager from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Callable, Protocol +from typing import BinaryIO, Callable, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.state.filesystem import ( @@ -72,6 +72,7 @@ def rebind(self, dataset_descriptor: int, dataset_path: Path) -> None: 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: @@ -83,6 +84,13 @@ def validate_lease(self) -> None: 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.""" @@ -101,6 +109,16 @@ def rebind(self) -> None: 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.""" @@ -127,6 +145,57 @@ def verify(self, candidate: CandidateOutputManifest) -> Iterator[VerifiedCandida _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.""" @@ -141,7 +210,7 @@ def _open_output_file( output_file: CandidateOutputFile, ) -> tuple[int, str, _FileBinding]: parts = PurePosixPath(output_file.relative_path).parts - parent_descriptor, parent_path = resources.enter_context( + parent_descriptor, parent_path, _ = resources.enter_context( _open_parent_directory(dataset_descriptor, dataset_path, parts[:-1]) ) name = parts[-1] @@ -165,14 +234,63 @@ def _open_output_file( 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]]: +) -> 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 @@ -186,7 +304,8 @@ def _open_parent_directory( os.close(parent_descriptor) parent_descriptor = next_descriptor parent_path = child_path - yield parent_descriptor, parent_path + directory_identities.append(_identity(os.fstat(parent_descriptor))) + yield parent_descriptor, parent_path, tuple(directory_identities) finally: os.close(parent_descriptor) @@ -213,6 +332,7 @@ def _is_safe_file(status: os.stat_result) -> bool: __all__ = [ + "CandidateArtifactSnapshot", "CandidateArtifactVerifier", "CandidateSchema", "SerializedCandidateSchema", 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 index e9825bd95..dc3ffcf19 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/base.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/base.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import datetime, timedelta +from typing import TypeAlias from pydantic import NonNegativeInt, PositiveInt @@ -44,10 +45,14 @@ class SchedulerIdentity(StateValue): array_task_id: NonNegativeInt +SchedulerJobIdentity: TypeAlias = PositiveInt | SchedulerIdentity + + __all__ = [ "ArtifactReference", "Identifier", "SchedulerIdentity", + "SchedulerJobIdentity", "Sha256Digest", "StateRecord", "StateValue", 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..396ccdcf5 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -0,0 +1,385 @@ +# 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 + +from datetime import datetime, timezone +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.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) -> 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, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + 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 + + def submit( + self, + *, + destination: str | Path | 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) + 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.mount.source), + ) + run, resolved_plan, _ = self._reader.load_context() + resolved_destination = self._destinations.resolve(resolved_plan, destination) + 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=self._inputs.get_winner_shards(), + host_destination=resolved_destination.host_path, + container_destination=resolved_destination.container_path, + num_partitions=resolved_plan.output.partitions, + ) + self._inputs.resolve(collection_plan) + prepare_collection_destination( + Path(collection_plan.host_destination), + Path(resolved_destination.mount.source), + ) + 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, 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.mount.source), + ) + 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.mount.source), + ) + 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.mount.source), + ) + 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) + 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) + requested = self._destinations.resolve(resolved_plan, requested_destination) + if requested != resolved: + raise StateCorruptionError("persisted collection destination does not match the requested destination") + if plan.host_destination != resolved.host_path or plan.container_destination != resolved.container_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: + self._inputs.resolve(plan) + result = self._collections.read_result(plan) + resolved_plan = self._reader.load_resolved_plan() + validated = validate_collection_result( + plan, + result, + expected_records=resolved_plan.invocation.authored.num_records, + output_format=resolved_plan.output.format, + ) + 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, + Path(plan.host_destination), + 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 _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..72330c1c0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_filesystem.py @@ -0,0 +1,357 @@ +# 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 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 + + +@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) + _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) + _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: + 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 _identity(status: os.stat_result) -> tuple[int, int]: + return status.st_dev, status.st_ino + + +__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..764cf1ac3 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_inputs.py @@ -0,0 +1,81 @@ +# 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.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) + + 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 + 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..d7ec2803a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_merge.py @@ -0,0 +1,338 @@ +# 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 AbstractContextManager, contextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import BinaryIO, Protocol, TextIO + +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.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 + output_format: str + schema: object + record_counts: tuple[int, ...] + + +class CollectionMerger: + """Merge ordered winner files while retaining only bounded data and descriptors.""" + + def __init__( + self, + output_format: str, + *, + verifier: CandidateArtifactVerifier | None = None, + completed_at: datetime | None = None, + ) -> None: + self._output_format = output_format + 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) + 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, + output_format=self._output_format, + schema=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): + suffix = "jsonl" if layout.output_format == "jsonl" else layout.output_format + relative_path = f"part-{partition_index:05d}.{suffix}" + output_path = layout.stage_path / relative_path + with _open_partition_writer(output_path, layout.output_format, 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 | TextIO) -> 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) -> None: + self._writer = writer + self._output = output + self._closed = False + + def write(self, batch: object) -> None: + self._writer.write_batch(batch.replace_schema_metadata(None)) + + 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() + + +class _TextWriter: + def __init__(self, output: TextIO, output_format: str, bound_output: _BoundOutput) -> None: + self._output = output + self._bound_output = bound_output + self._format = output_format + self._is_first = True + + def write(self, batch: object) -> None: + frame = batch.to_pandas() + if self._format == "csv": + frame.to_csv(self._output, header=self._is_first, index=False) + else: + content = frame.to_json(orient="records", lines=True, force_ascii=False, date_format="iso") + if content: + self._output.write(content) + self._is_first = False + + def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: + return self._bound_output.describe(relative_path, record_count) + + +def _open_partition_writer( + output_path: Path, + output_format: str, + schema: object, +) -> AbstractContextManager[CollectionPartitionWriter]: + if output_format == "parquet": + return _open_parquet_writer(output_path, schema) + if output_format in {"csv", "jsonl"}: + return _open_text_writer(output_path, output_format) + raise ValueError(f"unsupported collection output format {output_format!r}") + + +@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)) + try: + yield partition_writer + finally: + partition_writer.close() + _sync_output(output) + + +@contextmanager +def _open_text_writer(output_path: Path, output_format: str) -> Generator[CollectionPartitionWriter, None, None]: + with output_path.open("x+", encoding="utf-8", newline="") as output: + os.fchmod(output.fileno(), 0o600) + try: + yield _TextWriter(output, output_format, _BoundOutput(output_path, output)) + finally: + _sync_output(output) + + +def _sync_output(output: BinaryIO | TextIO) -> 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_storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py new file mode 100644 index 000000000..7da068a42 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_storage.py @@ -0,0 +1,339 @@ +# 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: + return self.read_result_from(plan, Path(plan.host_destination)) + + def read_result_from(self, plan: CollectionPlan, destination: Path) -> CollectionResult: + """Read a result through an explicitly selected host or container view.""" + if destination.as_posix() not in {plan.host_destination, plan.container_destination}: + raise OSError("collection result destination does not match its immutable plan") + 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, + destination: Path, + *, + verify_digests: bool = True, + ) -> None: + """Verify exact inventory with bounded metadata or full output digests.""" + if destination.as_posix() not in {plan.host_destination, plan.container_destination}: + raise OSError("collection result destination does not match its immutable plan") + 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..369b4cf6d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py @@ -0,0 +1,113 @@ +# 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, complete 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 not complete") + + schema_digests = {candidate.dataset_schema_digest for candidate in candidates} + _require(len(schema_digests) == 1, "collection candidates have incompatible schemas") + _require( + all(candidate.provenance_digest == resolved_plan.compute_sha256() for candidate in candidates), + "collection candidate provenance does not match the resolved plan", + ) + + expected_records = resolved_plan.invocation.authored.num_records + actual_records = sum(candidate.actual_records for candidate in candidates) + _require(actual_records == expected_records, "collection candidate rows do not match requested run rows") + 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, + output_format: str, +) -> 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") + suffix = "jsonl" if output_format == "jsonl" else output_format + expected_paths = tuple(f"part-{index:05d}.{suffix}" 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..8f464ee6c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_worker.py @@ -0,0 +1,292 @@ +# 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.container_destination), + derive_collection_staging_directory(plan), + Path(destination.mount.target), + ) as staged: + merger = CollectionMerger(resolved_plan.output.format, 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_from(plan, Path(plan.container_destination)) + resolved_plan = self._reader.load_resolved_plan() + validated = validate_collection_result( + plan, + result, + expected_records=resolved_plan.invocation.authored.num_records, + output_format=resolved_plan.output.format, + ) + 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, Path(plan.container_destination)) + 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..78c8a2bc4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Profile-authorized host and container collection destinations.""" + +from __future__ import annotations + +import posixpath +from dataclasses import dataclass +from pathlib import Path + +from data_designer.slurm.config import ContainerMount +from data_designer.slurm.contracts import is_path_below, validate_absolute_path +from data_designer.slurm.images.records import validate_enroot_mount_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 output directory and its authorized writable mount.""" + + host_path: str + container_path: str + mount: ContainerMount + + +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 output path in host and container namespaces.""" + raw_destination = plan.output.root 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 + workspace_root = plan.selected_profile.profile.workspace_root + workspace_mount = ContainerMount(source=workspace_root, target=workspace_root, read_only=False) + authorized = { + (mount.source, mount.target): mount + for mount in (workspace_mount, *plan.container_mounts) + if not mount.read_only + } + writable = tuple( + mount + for mount in authorized.values() + if host_path == mount.source or is_path_below(host_path, mount.source) + ) + if not writable: + raise StateConflictError("collection destination is not covered by a profile-authorized writable mount") + longest = max(len(mount.source) for mount in writable) + matches = tuple(mount for mount in writable if len(mount.source) == longest) + if len(matches) != 1: + raise StateConflictError("collection destination has an ambiguous writable mount mapping") + mount = matches[0] + if host_path == mount.source: + raise StateConflictError("collection destination must be below its writable mount source") + try: + validate_enroot_mount_path(workspace_root) + validate_enroot_mount_path(mount.source) + validate_enroot_mount_path(mount.target) + except ValueError as error: + raise StateConflictError("collection paths cannot be represented as safe Enroot mounts") from error + relative = posixpath.relpath(host_path, mount.source) + container_path = mount.target if relative == "." else posixpath.join(mount.target, relative) + return CollectionDestination(host_path, validate_absolute_path(container_path), mount) + + 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.host_destination != destination.host_path: + raise StateConflictError("collection host destination no longer matches resolved intent") + if collection_plan.container_destination != destination.container_path: + raise StateConflictError("collection container destination no longer matches resolved intent") + if collection_plan.num_partitions != resolved_plan.output.partitions: + raise StateConflictError("collection partition count no longer matches resolved intent") + return destination + + +__all__ = ["CollectionDestination", "CollectionDestinationResolver"] 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..4b018d726 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -0,0 +1,308 @@ +# 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 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 ( + 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. + """ + 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, + ) + for shard in shards + ) + return self._compose_run_status(run, timestamp, 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, + ) -> ShardStatus: + try: + with self._storage.acquire_shard_lock(expected.shard.shard_id): + 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, + ) + 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, + ) -> 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] + self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) + effective_state = reconcile_attempt_observation( + attempt, + readiness, + scheduler, + current_time=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 StateCorruptionError("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 index d332f7769..13aa11887 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -115,6 +115,45 @@ class CollectionShard(StateValue): 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.""" @@ -131,6 +170,11 @@ class CollectionPlan(StateRecord): _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) _destinations_are_safe = field_validator("host_destination", "container_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: shard_ids = [shard.shard_id for shard in self.planned_shards] 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 index 67e23c2a9..a20828838 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -5,13 +5,16 @@ 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, @@ -139,18 +142,65 @@ def load_readiness(self, shard_id: ShardId, attempt_id: AttemptId) -> AttemptRea 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(shard_id, attempt_id) + readiness = self._storage.read_readiness(attempt.shard_id, attempt.attempt_id) PersistedPlanStateValidator(plan).validate_readiness_snapshot(attempt, readiness) return readiness - except FileNotFoundError as error: - raise StateNotFoundError(f"attempt {attempt_id!r} has no readiness snapshot") from error - except StateCorruptionError: - raise + except FileNotFoundError: + return None except PlanStateContractError as error: - raise StateCorruptionError(f"attempt {attempt_id!r} has invalid persisted readiness") from 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_id!r} has unreadable readiness") from 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 and validate a complete producer result pair when present.""" + 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, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 63c0d2a1a..473150f43 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -15,6 +15,7 @@ EffectiveAttemptState, SchedulerObservation, SchedulerState, + is_scheduler_failure_state, ) from data_designer.slurm.state.validation import StateContractError @@ -59,18 +60,6 @@ EndpointPublicationState.FAILED: frozenset({EndpointPublicationState.FAILED}), } -_SCHEDULER_FAILURE_STATES = frozenset( - { - SchedulerState.FAILED, - SchedulerState.CANCELLED, - SchedulerState.TIMED_OUT, - SchedulerState.NODE_FAILED, - SchedulerState.PREEMPTED, - SchedulerState.REQUEUED, - SchedulerState.OUT_OF_MEMORY, - } -) - def validate_readiness_transition( previous: AttemptReadiness, @@ -165,7 +154,7 @@ def reconcile_attempt_observation( if readiness is not None: _require(readiness.updated_at >= attempt.created_at, "readiness update cannot precede attempt creation") - if scheduler.state in _SCHEDULER_FAILURE_STATES: + if is_scheduler_failure_state(scheduler.state): return EffectiveAttemptState.FAILED if attempt.state is AttemptLifecycleState.FAILED: return EffectiveAttemptState.FAILED @@ -184,7 +173,12 @@ def reconcile_attempt_observation( return EffectiveAttemptState.UNKNOWN if readiness is not None and readiness.state is ReadinessState.FAILED: return EffectiveAttemptState.FAILED - if scheduler.state is SchedulerState.PENDING: + 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 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..607ab3fc9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -0,0 +1,481 @@ +# 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, 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) -> 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, + ) -> None: + root, normalized_run_id = _validate_location(workspace_root, run_id) + self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() + 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 _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) + 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..489af26b0 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry_storage.py @@ -0,0 +1,210 @@ +# 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.""" + return f"retry-{len(self.list_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.""" + return ArtifactReference(path=self.get_plan_path(plan.retry_id).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 + + @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 index 3245c75fa..6961a1295 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/scheduler.py @@ -9,7 +9,7 @@ from pydantic import field_validator, model_validator from data_designer.slurm.state.base import ( - SchedulerIdentity, + SchedulerJobIdentity, StateRecord, validate_optional_utc_timestamp, validate_utc_timestamp, @@ -34,7 +34,7 @@ class SchedulerState(str, Enum): class SchedulerObservation(StateRecord): """Normalized scheduler observation used for deterministic reconciliation.""" - scheduler: SchedulerIdentity + scheduler: SchedulerJobIdentity observed_at: datetime state: SchedulerState reconciliation_deadline: datetime | None = None @@ -47,10 +47,10 @@ 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") - if self.reconciliation_deadline < self.observed_at: - raise ValueError("reconciliation deadline must not precede the observation") - elif self.reconciliation_deadline is not None: - raise ValueError("only accounting lag may have 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 @@ -61,3 +61,19 @@ class EffectiveAttemptState(str, Enum): 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 index f940c3238..4f6f60694 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -34,6 +34,7 @@ ) 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" @@ -44,6 +45,7 @@ _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" @@ -115,6 +117,9 @@ def get_attempt_path(self, shard_id: str, attempt_id: str) -> Path: 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 @@ -310,6 +315,55 @@ 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) 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 index 9076e1955..6ab827ffe 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -25,12 +25,13 @@ ) 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, ShardWinner +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, @@ -84,6 +85,7 @@ def __init__( 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 @@ -124,6 +126,22 @@ 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() 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 index f990b65e8..1e20fc7c4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/validation.py @@ -15,7 +15,11 @@ CollectionPlan, ShardWinner, ) -from data_designer.slurm.state.scheduler import SchedulerObservation, SchedulerState +from data_designer.slurm.state.scheduler import ( + SchedulerObservation, + SchedulerState, + is_scheduler_terminal_state, +) _ATTEMPT_STATE_ORDER = { AttemptLifecycleState.CREATED: 0, @@ -183,6 +187,8 @@ def validate_scheduler_observation_transition( """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 @@ -197,6 +203,12 @@ def validate_scheduler_observation_transition( 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 diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index c2175577b..7da326e0c 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -135,6 +135,8 @@ def test_main_preserves_equals_in_model_alias( client_worker_case.prepared.attempt_id, "--attempt-dir", client_worker_case.attempt_dir.as_posix(), + "--resume-mode", + "always", "--endpoint", f"judge=v2={endpoint}", ] @@ -142,6 +144,7 @@ def test_main_preserves_equals_in_model_alias( 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( @@ -348,6 +351,77 @@ def test_if_possible_without_workspace_uses_attempt_dataset(client_worker_case: 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" diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 55b8258ff..d17e5eb25 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -5,12 +5,13 @@ 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 +from data_designer.slurm.launcher.errors import SlurmCommandError, SlurmCommandOutputError, SlurmSubmissionError from data_designer.slurm.state import SchedulerIdentity, SchedulerState @@ -72,6 +73,23 @@ def test_client_exports_only_explicit_environment_names_without_values_in_argv() assert "secret-value" not in " ".join(runner.command) +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") @@ -97,6 +115,51 @@ def test_client_queries_accounting_and_cancels_one_array_task(fake_slurm_runner: ] +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") @@ -208,6 +271,39 @@ def test_client_normalizes_command_failures(fake_slurm_runner: FakeSlurmRunner) 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", @@ -331,7 +427,14 @@ def run( class _TimeoutRunner: - def run(self, command: Sequence[str]) -> subprocess.CompletedProcess[str]: + 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) 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..a79a8187a --- /dev/null +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -0,0 +1,205 @@ +# 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 +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +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_authorized_mounts_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, + ), + ), + ), + host_destination=destination.host_path, + container_destination=destination.container_path, + num_partitions=plan.output.partitions, + ) + + 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_STATE_MOUNT="/workspace/primary:/workspace/primary"' in script + assert 'readonly DD_OUTPUT_MOUNT="/workspace/primary/runs/run-001:/exports"' in script + assert ( + 'readonly DD_COLLECTION_PLAN="/workspace/primary/runs/run-001/collections/collection-0001/plan.json"' in script + ) + assert "data_designer.slurm.state.collection_worker" in script + assert "#SBATCH --gres=" not in script + assert "#SBATCH --gpus=" not in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( + multi_node_plan: ResolvedSlurmRunPlan, +) -> 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 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_DIR") + assert "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("data_designer.slurm.state.attempt_identity") < script.index("DD_RUNTIME_DIR") + assert '"${DD_RETRY_PLAN_SHA256}" "${DD_EFFECTIVE_RESUME_MODE}"' in script + assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 + + +def test_retry_renderer_uses_profile_slurm_bin_path(multi_node_plan: ResolvedSlurmRunPlan) -> 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}) + 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 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, + ), + ), + ), + host_destination=requested, + container_destination=requested, + num_partitions=plan.output.partitions, + ) + + assert destination.mount == ContainerMount(source=workspace_root, target=workspace_root, read_only=False) + assert resolver.validate_persisted(plan, collection) == destination + + +def test_destination_requires_one_unique_most_specific_mapping( + 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}) + + with pytest.raises(StateConflictError, match="ambiguous writable mount"): + CollectionDestinationResolver().resolve(plan, requested) diff --git a/packages/data-designer-slurm/tests/runtime/conftest.py b/packages/data-designer-slurm/tests/runtime/conftest.py index 475a56ff5..253f7a771 100644 --- a/packages/data-designer-slurm/tests/runtime/conftest.py +++ b/packages/data-designer-slurm/tests/runtime/conftest.py @@ -151,6 +151,9 @@ def verify(self, context: AllocationContext, environment: object) -> None: class FakeClientStepBuilder: + def __init__(self) -> None: + self.retry_resume_modes: list[Literal["never", "always"] | None] = [] + def build_preflight_step( self, plan: ResolvedSlurmRunPlan, @@ -159,8 +162,11 @@ def build_preflight_step( 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( @@ -171,8 +177,11 @@ def build_generation_step( 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) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 51a56aa2c..10fb6c277 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -3,10 +3,13 @@ from __future__ import annotations +from dataclasses import replace + from conftest import RuntimeCase from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.state import RetryPlan, RetryShard def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: @@ -41,3 +44,39 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert "--attempt-id" not in manifest.steps[-1].command assert "--plan" in manifest.steps[-1].command assert "--attempt-dir" in manifest.steps[-1].command + + +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=context.attempt_directory / "runtime", + log_directory=context.attempt_directory / "logs/execution-00000002", + ) + + 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 diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index 40174aba1..5b73ad071 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -47,7 +47,10 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p assert all(archive.getmember(name).uid == 0 for name in names) entrypoint = archive.extractfile("entrypoint.sh") assert entrypoint is not None - assert b"python3 -m data_designer.slurm.runtime.entrypoint" not in entrypoint.read() + entrypoint_content = entrypoint.read() + 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( diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py index a0d0e7d19..ee16de9f2 100644 --- a/packages/data-designer-slurm/tests/runtime/test_context.py +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -7,14 +7,21 @@ 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, @@ -22,6 +29,59 @@ ) +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_reads_and_updates_state_through_remapped_workspace( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index 45ff94b19..df3439944 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -32,6 +32,8 @@ DeploymentReadiness, EndpointPublicationState, ReadinessState, + RetryPlan, + RetryShard, RunManifest, ShardManifest, SlurmStateWriter, @@ -252,6 +254,48 @@ def finalize_winner(self, *args: object, **kwargs: object) -> 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) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 86439bcc8..17c12005d 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -179,7 +179,7 @@ def _patch_runtime_context( runtime_case: RuntimeCase, state: FakeStateStore, ) -> None: - monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args: (runtime_case.context, state)) + monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) monkeypatch.setenv("SLURM_JOB_GPUS", "0") diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 8e5c1976a..bbe2b39d1 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -5,7 +5,7 @@ import hashlib from collections.abc import Mapping -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path import pytest @@ -399,6 +399,91 @@ def test_status_reconciles_cancelled_scheduler_attempt( 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})} + ) + _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 + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, 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..0a49de6be --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -0,0 +1,709 @@ +# 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, FakeSlurmRunner + +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, + AttemptTerminalClassification, + CandidateOutcome, + CandidateOutputFile, + CandidateOutputManifest, + EffectiveAttemptState, + EffectiveRunState, + GenerationState, + 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_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_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)) + + +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_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: + 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=1, + ) + planned_shard = relocated_plan.shards[0] + shard = ShardManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=planned_shard.shard_id, + shard_index=planned_shard.shard_index, + record_range=planned_shard.record_range, + input_partition=planned_shard.input_partition, + resume_workspace=planned_shard.resume_workspace, + created_at=created_at, + ) + writer = SlurmStateWriter(workspace, relocated_plan.run_id) + writer.initialize_run(authored_config, relocated_plan, run, (shard,)) + attempt = AttemptManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, + scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0) if submitted else None, + created_at=created_at + timedelta(minutes=1), + updated_at=created_at + timedelta(minutes=2), + ) + writer.create_attempt(attempt) + return _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, 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 = 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) -> tuple[AttemptManifest, ShardWinner]: + running = _copy_attempt( + case.attempt, + state=AttemptLifecycleState.RUNNING, + updated_at=case.created_at + timedelta(minutes=3), + ) + case.writer.update_attempt(running) + candidate_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/output-manifest.json" + dataset_path = candidate_path.parent / "dataset" + requested = case.plan.shards[0].requested_records + candidate = CandidateOutputManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=running.shard_id, + attempt_id=running.attempt_id, + attempt_ordinal=running.attempt_ordinal, + created_at=case.created_at + timedelta(minutes=4), + 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=case.created_at + timedelta(minutes=5), + 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=case.created_at + timedelta(minutes=6), + ) + 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=case.created_at + timedelta(minutes=7), + ) + winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + winner_path.write_text(winner.serialize_json()) + winner_path.chmod(0o600) + return completed, winner + + +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_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py new file mode 100644 index 000000000..441744846 --- /dev/null +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -0,0 +1,1780 @@ +# 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 json +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 + +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_filesystem import derive_collection_staging_directory +from data_designer.slurm.state.collection_inputs import CollectionInputResolver +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.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 + 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_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),)) + coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) + + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + + assert submitted.state is CollectionState.SUBMITTED + assert submitted.scheduler == 5101 + script = cast(str, runner.inputs[-1]) + assert "data_designer.slurm.state.collection_worker" in script + assert "--gpus" not in script + assert "--gres" not in script + stale_stage = Path(case.plan.output.root).parent / submitted.staging_directory + stale_stage.mkdir(mode=0o700) + (stale_stage / "partial").write_text("incomplete") + unrelated_stage = Path(case.plan.output.root).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 = Path(case.plan.output.root) + assert result.actual_records == case.plan.invocation.authored.num_records + assert len(result.files) == case.plan.output.partitions + 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_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: + workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) + output = multi_node_plan.output.model_copy( + update={"root": (workspace_root / "new" / "nested" / "output").as_posix()} + ) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, plan) + destination = Path(case.plan.output.root) + 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(submitted_at=case.created_at + timedelta(minutes=10)) + + assert destination.parent.is_dir() + assert not destination.exists() + 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 Path(case.plan.output.root).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)) + + 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.FAILED + submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=11)) + assert submitted.collection_id == "collection-0002" + assert submitted.state is CollectionState.SUBMITTED + + +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 = Path(case.plan.output.root) + 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 = Path(case.plan.output.root) + 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 = Path(case.plan.output.root).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 Path(case.plan.output.root).exists() + + +def test_collection_worker_reauthorizes_persisted_partition_intent( + 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)) + storage = CollectionStorage(StateStorage(case.workspace, case.plan.run_id)) + plan = storage.read_plan(submitted.collection_id) + tampered = plan.model_copy(update={"num_partitions": plan.num_partitions + 1}) + storage.get_plan_path(plan.collection_id).write_text(tampered.serialize_json()) + storage.replace_status( + submitted.model_copy( + update={ + "collection_plan": storage.get_plan_reference(tampered), + "staging_directory": derive_collection_staging_directory(tampered), + } + ) + ) + + with pytest.raises(StateConflictError, match="partition count"): + 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 Path(case.plan.output.root).exists() + + +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 = Path(case.plan.output.root) / 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 = Path(case.plan.output.root) / 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 = Path(case.plan.output.root) + 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 = Path(case.plan.output.root) + 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_detects_destination_parent_replacement_before_success( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) + output = multi_node_plan.output.model_copy(update={"root": (workspace_root / "exports" / "output").as_posix()}) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, 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 = Path(case.plan.output.root) + 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 Path(case.plan.output.root).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") + + +@pytest.mark.parametrize("output_format", ["csv", "jsonl"]) +def test_collection_exports_non_parquet_formats_in_deterministic_partitions( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + output_format: str, +) -> None: + output = multi_node_plan.output.model_copy(update={"format": output_format, "partitions": 3}) + plan = ResolvedSlurmRunPlan.model_validate_json( + json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) + ) + case = _initialize_run(tmp_path, authored_run, 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)) + + 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 len(result.files) == 3 + requested_records = case.plan.invocation.authored.num_records + floor_count = requested_records // 3 + assert tuple(file.record_count for file in result.files) == ( + floor_count, + floor_count, + requested_records - 2 * floor_count, + ) + values: list[int] = [] + for output_file in result.files: + path = Path(case.plan.output.root) / output_file.relative_path + if output_format == "csv": + with path.open(newline="") as source: + values.extend(int(row["record_id"]) for row in csv.DictReader(source)) + else: + values.extend(int(json.loads(line)["record_id"]) for line in path.read_text().splitlines()) + assert values == list(range(requested_records)) + + +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) -> 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) + 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) -> None: + output_path = dataset_path / "part-00000.parquet" + values = range(shard.record_range.start_index, shard.record_range.end_index_exclusive) + table = lazy.pa.table({"record_id": values}) + 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=shard.record_range.record_count, + outcome=CandidateOutcome.COMPLETE, + files=( + CandidateOutputFile( + relative_path=output_path.name, + sha256=hashlib.sha256(content).hexdigest(), + byte_size=len(content), + record_count=shard.record_range.record_count, + ), + ), + 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=shard.record_range.record_count, + 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=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_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 9a59c86cd..a7865afc4 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1633,6 +1633,22 @@ def track_open_file( 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())]) 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_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index ff0493b19..8817b4339 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -569,6 +569,8 @@ def test_reconciliation_covers_nonterminal_and_fallback_states() -> None: 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), diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 3ff48707c..c2c696633 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -136,8 +136,24 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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 RunManifest +from data_designer.slurm.state import ( + CollectionResult, + RetryPlan, + RunManifest, + RunStatus, + SchedulerObservationCollector, + SlurmCollectionCoordinator, + SlurmRetryCoordinator, + SlurmStateReconciler, +) +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 From 10601e43977639b2c48372182710fe7f3eb58eea Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 12:29:23 -0600 Subject: [PATCH 22/62] fix(slurm): restore plugin activation boundary (#930) * fix(slurm): isolate plugin-aware client generation Signed-off-by: Nabin Mulepati * refactor(slurm): expose client worker boundary Signed-off-by: Nabin Mulepati * refactor(slurm): compose client process boundary Signed-off-by: Nabin Mulepati * fix(slurm): refresh client overlay inventory Signed-off-by: Nabin Mulepati * test(slurm): exercise plugin worker handoff Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/client/environment.py | 2 + .../src/data_designer/slurm/client/process.py | 36 +++ .../data_designer/slurm/runtime/entrypoint.py | 12 +- .../tests/client/test_environment.py | 15 + .../tests/client/test_process.py | 26 ++ .../tests/client/test_worker.py | 9 +- .../implementation.py | 17 ++ .../fake_data_designer_plugin/plugin.py | 16 +- .../tests/runtime/test_entrypoint.py | 274 +++++++++++++++++- 9 files changed, 386 insertions(+), 21 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/process.py create mode 100644 packages/data-designer-slurm/tests/client/test_process.py create mode 100644 packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py 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 index 6e96d15ea..2d7780948 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -272,6 +272,8 @@ def _install_overlay( def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ...]: """Return an exact immutable distribution inventory.""" + if path is not None: + importlib.invalidate_caches() distributions = ( importlib.metadata.distributions() if path is None else importlib.metadata.distributions(path=[path.as_posix()]) ) 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/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index db0462981..dd43a06f2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -13,7 +13,7 @@ from pathlib import Path from data_designer.slurm.client.filesystem import ensure_private_directory, replace_private_text -from data_designer.slurm.client.worker import main as client_worker_main +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 @@ -145,8 +145,14 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non ) -def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: +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 @@ -158,7 +164,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No context.attempt.attempt_id, resume_mode, ): - return_code = client_worker_main( + return_code = worker_process.run( ( "run", "--plan", diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index c42e284fd..51c920589 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -6,6 +6,7 @@ import hashlib import importlib import json +import os import subprocess import sys from pathlib import Path @@ -144,6 +145,20 @@ def distributions(**kwargs: object) -> tuple[()]: assert calls == [{}] +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"}, 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_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 7da326e0c..bb5805573 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -500,7 +500,14 @@ def test_preflight_rejects_plugin_secondary_model_alias( ) -> 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"}] + 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 = ( 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 index 67ce6aae3..dea69fc26 100644 --- 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 @@ -13,8 +13,8 @@ class FakePluginConfig(SingleColumnConfig): """Minimal custom column configuration.""" column_type: Literal["fake-slurm-column"] = "fake-slurm-column" - model_alias: str = "generator" - judge_model_alias: str = "judge" + model_alias: str | None = None + judge_model_alias: str | None = None @property def required_columns(self) -> list[str]: @@ -25,19 +25,11 @@ def side_effect_columns(self) -> list[str]: return [] def get_model_aliases(self) -> list[str]: - return [self.model_alias, self.judge_model_alias] - - -class FakePluginImplementation: - """Minimal loadable implementation for entry-point verification.""" - - def generate(self, data: dict[str, object]) -> dict[str, object]: - """Return the provided record unchanged.""" - return data + 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.plugin.FakePluginImplementation", + impl_qualified_name="fake_data_designer_plugin.implementation.FakePluginImplementation", plugin_type=PluginType.COLUMN_GENERATOR, ) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 17c12005d..561588810 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -3,21 +3,50 @@ from __future__ import annotations +import csv +import hashlib +import io +import json +import os +import sys +import venv +import zipfile +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace +from typing import cast +import pyarrow.parquet as pq import pytest from conftest import FakeStateStore, RuntimeCase +from packaging.tags import interpreter_name, interpreter_version import data_designer.slurm.runtime.entrypoint as entrypoint -from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.state import AttemptLifecycleState, ReadinessState +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( @@ -65,8 +94,6 @@ def verify_attempt_directory(path: Path) -> None: with pytest.raises(_InjectedFailure): entrypoint._prepare(prepare, {}) - monkeypatch.setattr(entrypoint, "client_worker_main", lambda arguments: 0) - def load_candidate(*args: object, attempt_directory: Path | None = None) -> None: assert attempt_directory == container_attempt_directory raise _InjectedFailure @@ -76,7 +103,58 @@ def load_candidate(*args: object, attempt_directory: Path | None = None) -> None _phase_arguments("client", runtime_case, attempt_directory=container_attempt_directory) ) with pytest.raises(_InjectedFailure): - entrypoint._client(client, {}) + 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) + 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 def test_control_phases_record_running_ready_and_failed( @@ -184,6 +262,192 @@ def _patch_runtime_context( monkeypatch.setenv("SLURM_JOB_GPUS", "0") +def _prepare_plugin_runtime_case( + runtime_case: RuntimeCase, + fake_plugin_overlay: Path, + tmp_path: Path, +) -> _PluginRuntimeCase: + image_distributions = runtime_case.context.plan.client.image.inspection_facts.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_inspection(client, image_distributions, _write_test_installer(tmp_path)) + lock = _plugin_dependency_lock(client, image_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(image_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_inspection( + client: dict[str, object], + image_distributions: tuple[InstalledDistribution, ...], + installer: Path, +) -> None: + inspection = cast( + dict[str, object], + cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], + ) + inspection["python_abi"] = f"{interpreter_name()}{interpreter_version()}" + inspection["python_version"] = sys.version.split()[0] + inspection["installer_path"] = installer.as_posix() + inspection["distributions"] = [item.model_dump(mode="json") for item in image_distributions] + 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], + image_distributions: tuple[InstalledDistribution, ...], + wheel_path: Path, +) -> ResolvedDependencyLock: + inspection = cast( + dict[str, object], + cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], + ) + return ResolvedDependencyLock.model_validate( + { + "schema_version": 1, + "resolver_version": "resolver-1", + "python_abi": inspection["python_abi"], + "client_image_sha256": cast(dict[str, object], client["image"])["sha256"], + "authored_requirements": ("fake-data-designer-plugin==1.0.0",), + "authored_source": None, + "source": None, + "image_distributions": image_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_test_installer(tmp_path: Path) -> Path: + """Create an offline pip executable without changing the test checkout.""" + environment = tmp_path / "installer-environment" + venv.EnvBuilder(with_pip=True, symlinks=True).create(environment) + return environment / "bin/pip" + + +def _write_inventory_bootstrap( + distributions: tuple[InstalledDistribution, ...], + tmp_path: Path, +) -> Path: + """Present plan-recorded image 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): + if kwargs.get("path") is not None: + 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 image 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, From 9dd6fc0aa520baf8bc11ed796cfe44c9d31cfe1b Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 12:29:51 -0600 Subject: [PATCH 23/62] fix: harden Slurm release artifacts (#916) * harden Slurm release artifacts Add sanitized credential and infrastructure scans for source, fixtures, runtime archives, logs, and wheels. Include Apache license text in the Slurm wheel and redact sensitive scheduler diagnostics. Signed-off-by: Nabin Mulepati * fix Slurm security audit gaps Signed-off-by: Nabin Mulepati * harden adjacent secret redaction Signed-off-by: Nabin Mulepati * make Slurm redaction fail closed Signed-off-by: Nabin Mulepati * redact underscore secret options Signed-off-by: Nabin Mulepati * harden Slurm security audit boundaries Signed-off-by: Nabin Mulepati * fix control-obscured Slurm secrets Signed-off-by: Nabin Mulepati * fix control-split secret redaction Signed-off-by: Nabin Mulepati * fix whitespace-control secret suffixes Signed-off-by: Nabin Mulepati * harden archive content classification Detect explicitly supplied ZIP and tar inputs from their content when filenames use opaque suffixes. Reject unsupported tar special members while continuing to allow ordinary directories. Signed-off-by: Nabin Mulepati * close Slurm security boundary gaps Handle serialized authorization headers and Unicode whitespace without exposing credential suffixes. Keep public-artifact scanning fail-closed for non-regular files and unreadable ZIP members, and narrow test sentinels to complete values. Signed-off-by: Nabin Mulepati * redact complete Authorization diagnostics Treat an Authorization field as credential-bearing through its line boundary even when its value is quoted. This keeps control-obscured suffixes from surviving the protected redaction pass. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- packages/data-designer-slurm/LICENSE | 201 ++++++++ packages/data-designer-slurm/pyproject.toml | 1 + .../data_designer/slurm/launcher/client.py | 6 +- .../src/data_designer/slurm/security.py | 145 ++++++ .../tests/launcher/test_client.py | 356 ++++++++++++++ .../tests/runtime/test_bundle.py | 1 + .../data-designer-slurm/tests/test_package.py | 6 + .../tests/test_public_artifacts.py | 374 +++++++++++++++ plans/870/slurm-early-security-review.md | 60 +++ scripts/audit_slurm_public_artifacts.py | 435 ++++++++++++++++++ scripts/test_slurm_package_install.py | 11 + 11 files changed, 1593 insertions(+), 3 deletions(-) create mode 100644 packages/data-designer-slurm/LICENSE create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/security.py create mode 100644 packages/data-designer-slurm/tests/test_public_artifacts.py create mode 100644 plans/870/slurm-early-security-review.md create mode 100644 scripts/audit_slurm_public_artifacts.py 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/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index cc8c33046..f691c26dd 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -5,6 +5,7 @@ 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", 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 index 58e203105..cc40fbeff 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -8,7 +8,6 @@ import os import re import subprocess -import unicodedata from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta @@ -31,6 +30,7 @@ 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}$") @@ -342,8 +342,8 @@ def _validate_argument(value: str, *, field_name: str) -> None: def _normalize_bounded_text(value: str, *, limit: int = 512) -> str: - sanitized = "".join(" " if unicodedata.category(character).startswith("C") else character for character in value) - normalized = " ".join(sanitized.split()) + redacted = redact_sensitive_diagnostic(value) + normalized = " ".join(redacted.split()) return normalized if len(normalized) <= limit else f"{normalized[: limit - 3]}..." 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/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index d17e5eb25..e4be4ea4a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -317,6 +317,362 @@ def test_client_removes_terminal_controls_from_command_failures(fake_slurm_runne 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) diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index 5b73ad071..abfc4a885 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -48,6 +48,7 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p 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 diff --git a/packages/data-designer-slurm/tests/test_package.py b/packages/data-designer-slurm/tests/test_package.py index 2544dd973..24a4b4955 100644 --- a/packages/data-designer-slurm/tests/test_package.py +++ b/packages/data-designer-slurm/tests/test_package.py @@ -17,6 +17,12 @@ def test_slurm_uses_shared_namespace() -> 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() 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/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/scripts/audit_slurm_public_artifacts.py b/scripts/audit_slurm_public_artifacts.py new file mode 100644 index 000000000..236f8608b --- /dev/null +++ b/scripts/audit_slurm_public_artifacts.py @@ -0,0 +1,435 @@ +# 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, 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/tests", + "plans/850/data-designer-contract.md", + "plans/870/slurm-early-security-review.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/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index c2c696633..412a1b00b 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -46,6 +46,16 @@ def wheel_metadata(path: Path) -> Message: 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 build_wheels(uv: str, wheel_directory: Path) -> dict[str, Path]: for package_path in PACKAGE_PATHS: run( @@ -199,6 +209,7 @@ def main() -> None: base_wheel = wheels["data-designer"] leaf_wheel = wheels["data-designer-slurm"] + audit_public_artifacts(leaf_wheel) base_metadata = wheel_metadata(base_wheel) leaf_metadata = wheel_metadata(leaf_wheel) version = base_metadata["Version"] From a8726a95b7a878e92ebd65d7bbabb0a3fdbc2d39 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 10 Sep 2026 13:42:40 -0600 Subject: [PATCH 24/62] feat(slurm): harden distributed allocation runtime (#914) * feat(slurm): compose distributed runtime launch Layer multi-node topology and follower-failure handling onto the production allocation bootstrap introduced by #929. Preserve the one-node path while adding host-scoped steps, coordinated node workers, remote readiness probes, and distributed preflight coverage. Signed-off-by: Nabin Mulepati * map distributed model paths Apply the resolved container-mount mapping to absolute model paths in both coordinated multi-node workers and remote-node serving commands. Cover differing host and container roots in the bootstrap manifest regression. Signed-off-by: Nabin Mulepati * adapt retry bootstrap coverage Pass the one-node allocation layout required by the distributed runtime manifest after composing the merged retry coverage onto the new shared base. Signed-off-by: Nabin Mulepati * fix distributed launch preflight Select the multi-process vLLM executor whenever pipeline or tensor parallelism spans more than one worker. Carry mapped absolute model paths through the node-worker contract so every assigned node verifies accessibility before launch. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/runtime/bootstrap.py | 219 +++----------- .../slurm/runtime/distributed.py | 123 ++++++++ .../data_designer/slurm/runtime/entrypoint.py | 10 +- .../data_designer/slurm/runtime/entrypoint.sh | 68 ++++- .../data_designer/slurm/runtime/manifest.py | 129 +++++++++ .../src/data_designer/slurm/runtime/models.py | 1 + .../data_designer/slurm/runtime/network.py | 27 ++ .../data_designer/slurm/runtime/node_spec.py | 221 ++++++++++++++ .../slurm/runtime/node_worker.py | 272 ++++++++++++++++++ .../slurm/runtime/plan_reader.sh | 31 +- .../data_designer/slurm/runtime/preflight.py | 84 +++++- .../src/data_designer/slurm/runtime/proxy.py | 20 +- .../slurm/runtime/server_manifest.py | 244 ++++++++++++++++ .../slurm/runtime/step_runner.sh | 15 +- .../src/data_designer/slurm/runtime/steps.py | 19 +- .../tests/runtime/test_bootstrap.py | 120 +++++++- .../tests/runtime/test_entrypoint.py | 12 +- .../tests/runtime/test_node_worker.py | 214 ++++++++++++++ .../tests/runtime/test_preflight.py | 37 ++- .../tests/runtime/test_proxy.py | 6 + .../tests/runtime/test_shell_runtime.py | 63 +++- 21 files changed, 1704 insertions(+), 231 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/network.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_node_worker.py 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 index 264bfdd4a..2f425b3f7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -1,120 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed one-node step manifest produced inside the sealed client image.""" +"""Typed allocation step manifest produced inside the sealed client image.""" from __future__ import annotations -import os from collections.abc import Mapping from pathlib import Path -from typing import Literal - -from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator from data_designer.slurm.config.environment import ( - LiteralEnvironmentBinding, SecretRef, collect_secret_environment_names, ) -from data_designer.slurm.contracts import ContractRecord, ContractValue, validate_absolute_path -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.manifest import RuntimeBootstrapManifest, RuntimeProbeSpec, RuntimeStepSpec from data_designer.slurm.runtime.models import AllocationContext, RuntimeEndpoint, RuntimeStepRole from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_deployments -from data_designer.slurm.runtime.steps import ( - build_client_command, - build_endpoint_command, - build_vllm_command, -) +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 -from data_designer.slurm.serving.vllm import ResolvedVllmProcess -from data_designer.slurm.types import EnvironmentName, Identifier, NetworkPort, Sha256Digest - - -class RuntimeProbeSpec(ContractValue): - """One loopback readiness target monitored by the Bash controller.""" - - host: Literal["127.0.0.1"] = "127.0.0.1" - 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 - - -class RuntimeStepSpec(ContractValue): - """Container command and environment consumed by the Bash step runner.""" - - step_id: Identifier - role: RuntimeStepRole - image_path: str - command: tuple[str, ...] = Field(min_length=1) - cpus: PositiveInt - gpu_indices: tuple[NonNegativeInt, ...] = () - 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: RuntimeProbeSpec | None = None - - _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) - _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.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.role is RuntimeStepRole.SERVER and not self.gpu_indices: - raise ValueError("server runtime steps require GPUs") - if self.role is not RuntimeStepRole.SERVER and self.gpu_indices: - raise ValueError("non-server runtime steps cannot request GPUs") - return self - - -class RuntimeBootstrapManifest(ContractRecord): - """Secret-free one-node 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 def build_runtime_manifest( @@ -123,9 +28,11 @@ def build_runtime_manifest( *, runtime_root: Path, log_directory: Path, + layout: AllocationLayout, ) -> RuntimeBootstrapManifest: - """Build the secret-free one-node command handoff for the Bash controller.""" + """Build the secret-free command handoff for the Bash controller.""" plan = context.plan + validate_allocation_layout(plan, layout) runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) deployments = resolve_allocation_deployments(context, environment) endpoints = tuple( @@ -147,21 +54,13 @@ def build_runtime_manifest( endpoints, runtime_container_root, log_directory, + layout, ) ] - for deployment in deployments: - steps.extend( - _build_server_step( - deployment, - process, - context, - runtime_root, - runtime_container_root, - log_directory, - ) - for process in deployment.processes - ) - steps.extend(_build_endpoint_step(deployment, context, runtime_root, log_directory) for deployment in deployments) + steps.extend(build_server_steps(deployments, context, runtime_container_root, 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, @@ -172,6 +71,7 @@ def build_runtime_manifest( endpoints, runtime_container_root, log_directory, + layout, ) ) secret_names = set(collect_secret_environment_names(plan)) @@ -201,6 +101,7 @@ def _build_client_step( endpoints: tuple[RuntimeEndpoint, ...], runtime_container_root: str, 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 @@ -261,70 +162,7 @@ def _build_client_step( environment_prefixes={}, container_environment=tuple(sorted((*secret_names, *allocation_environment, "PYTHONPATH"))), log_directory=log_directory, - ) - - -def _build_server_step( - deployment: ResolvedVllmServerDeployment, - process: ResolvedVllmProcess, - context: AllocationContext, - runtime_root: Path, - runtime_container_root: str, - log_directory: Path, -) -> RuntimeStepSpec: - 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", - ) - 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) - ) - probe = next(item for item in deployment.readiness_probes if item.port == process.http_port) - return _step( - step_id=process.process_id, - role=RuntimeStepRole.SERVER, - image_path=deployment.image.path, - command=build_vllm_command(deployment, process, context.plan), - cpus=context.plan.client.authored.cpus, - gpu_indices=tuple(process.gpu_indices), - literal_environment=literal_environment, - secret_environment=secret_environment, - environment_prefixes=environment_prefixes, - container_environment=tuple( - sorted( - { - *deployment.launch_policy.environment, - "PYTHONPATH", - MAX_WAITING_REQUESTS_ENVIRONMENT, - RETRY_AFTER_SECONDS_ENVIRONMENT, - } - ) - ), - log_directory=log_directory, - launch_delay_seconds=process.launch_delay_seconds, - readiness=RuntimeProbeSpec( - port=probe.port, - path=probe.path, - deadline_seconds=probe.deadline_seconds, - ), + node_hosts=(layout.get_host(plan.client.host_node_index),), ) @@ -333,13 +171,20 @@ def _build_endpoint_step( 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", @@ -353,10 +198,14 @@ def _build_endpoint_step( environment_prefixes={}, container_environment=(), log_directory=log_directory, - readiness=RuntimeProbeSpec( - port=deployment.logical_endpoint.port, - path="/health", - deadline_seconds=deployment.launch_policy.startup_timeout_seconds, + 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, + ), ), ) @@ -369,13 +218,15 @@ def _step( 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: RuntimeProbeSpec | None = None, + readiness: tuple[RuntimeProbeSpec, ...] = (), ) -> RuntimeStepSpec: return RuntimeStepSpec( step_id=step_id, @@ -384,6 +235,8 @@ def _step( 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, 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..88236233f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py @@ -0,0 +1,123 @@ +# 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, NodeSpec, NodeWorkerSpec +from data_designer.slurm.runtime.paths import get_container_path +from data_designer.slurm.runtime.preflight import AllocationLayout +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, +) -> NodeWorkerSpec: + """Build the validated work assigned to each node in one deployment.""" + model = _resolve_model(deployment, plan) + nodes = tuple( + NodeSpec( + node_index=node_index, + host=layout.get_host(node_index), + ports=_get_node_ports(deployment, node_index), + processes=tuple( + NodeProcessSpec( + process_id=process.process_id, + command=build_vllm_process_command(deployment, process, plan, layout), + gpu_indices=tuple(process.gpu_indices), + launch_delay_seconds=process.launch_delay_seconds, + ) + for process in deployment.processes + if process.node_index == node_index + ), + ) + for node_index in deployment.node_indices + ) + return NodeWorkerSpec( + schema_version=1, + resolved_gpus_per_node=deployment.gpus_per_node, + required_model_path=model if deployment.model.startswith("/") else None, + nodes=nodes, + ) + + +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",) + 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 index dd43a06f2..28a0ea265 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -21,7 +21,11 @@ from data_designer.slurm.runtime.models import AllocationContext from data_designer.slurm.runtime.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_deployments -from data_designer.slurm.runtime.preflight import SystemAllocationPreflight +from data_designer.slurm.runtime.preflight import ( + AllocationLayout, + SystemAllocationPreflight, + validate_allocation_layout, +) from data_designer.slurm.runtime.records import load_complete_client_candidate from data_designer.slurm.serving.deployment import ResolvedVllmServerDeployment from data_designer.slurm.state import ( @@ -76,6 +80,7 @@ def _parse_arguments(arguments: Sequence[str] | None) -> argparse.Namespace: 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=[]) @@ -93,6 +98,8 @@ def _add_context_arguments(parser: argparse.ArgumentParser) -> None: 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_ports(context, environment) readiness = _begin_attempt(context, writer, environment) @@ -104,6 +111,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N environment, runtime_root=arguments.runtime_root, log_directory=log_directory, + layout=layout, ) expected_manifest = context.attempt_directory / "runtime-manifest.json" if arguments.manifest.as_posix() != get_container_path( 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 index cbba85438..4015b4f34 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -47,9 +47,15 @@ dd_slurm_run_allocation() { trap 'exit 130' INT TERM DD_RUNTIME_PREPARED=1 + local host + local -a host_arguments=() + for host in "${DD_ALLOCATION_HOSTS[@]}"; do + host_arguments+=(--node-host "${host}") + done dd_run_bound_control_phase prepare \ --runtime-root "${DD_RUNTIME_DIR}" \ - --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" + --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" \ + "${host_arguments[@]}" dd_verify_runtime_manifest \ "${DD_RUNTIME_MANIFEST}" \ "${DD_PLAN_SHA256}" \ @@ -61,6 +67,7 @@ dd_slurm_run_allocation() { ((${#DD_STEP_IDS[@]} == 1)) dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" + dd_run_role_steps server_preflight dd_start_servers dd_wait_for_role_readiness server dd_start_endpoints @@ -100,7 +107,11 @@ dd_verify_host_context() { done [[ -d ${DD_ATTEMPT_PATH} && ! -L ${DD_ATTEMPT_PATH} ]] [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] - [[ ${SLURM_JOB_NUM_NODES:-} == 1 && ${SLURM_NODEID:-} == 0 ]] + [[ ${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_read_artifacts "${DD_PLAN_PATH}" "${SLURM_ARRAY_TASK_ID}" local index path digest actual @@ -115,6 +126,24 @@ dd_verify_host_context() { done } +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 @@ -156,6 +185,15 @@ dd_start_servers() { 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 @@ -168,19 +206,25 @@ dd_start_endpoints() { dd_wait_for_role_readiness() { local role=$1 - local step_id deadline + local step_id index host port path deadline_seconds deadline 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}" - deadline=$((SECONDS + DD_STEP_PROBE_DEADLINE)) - until curl --fail --silent --max-time 1 \ - "http://${DD_STEP_PROBE_HOST}:${DD_STEP_PROBE_PORT}${DD_STEP_PROBE_PATH}" >/dev/null 2>&1; do - dd_require_running - ((SECONDS < deadline)) || { - printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 - return 70 - } - dd_sleep 0.5 + ((${#DD_STEP_PROBE_FIELDS[@]} % 4 == 0 && ${#DD_STEP_PROBE_FIELDS[@]} > 0)) || return 65 + for ((index = 0; index < ${#DD_STEP_PROBE_FIELDS[@]}; index += 4)); do + host=${DD_STEP_PROBE_FIELDS[index]} + port=${DD_STEP_PROBE_FIELDS[index + 1]} + path=${DD_STEP_PROBE_FIELDS[index + 2]} + deadline_seconds=${DD_STEP_PROBE_FIELDS[index + 3]} + deadline=$((SECONDS + deadline_seconds)) + until curl --fail --silent --max-time 1 "http://${host}:${port}${path}" >/dev/null 2>&1; do + dd_require_running + ((SECONDS < deadline)) || { + printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 + return 70 + } + dd_sleep 0.5 + done done done } 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..728c7d8b9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -0,0 +1,129 @@ +# 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 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): + """Container command and placement consumed by the Bash step runner.""" + + step_id: Identifier + role: RuntimeStepRole + image_path: str + command: tuple[str, ...] = Field(min_length=1) + cpus: PositiveInt + gpu_indices: tuple[NonNegativeInt, ...] = () + 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, ...] = () + + _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) + _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") + self._validate_placement() + self._validate_readiness() + return self + + def _validate_placement(self) -> None: + server_roles = {RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER} + 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 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 index 2aa93ab16..c91370b38 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/models.py @@ -25,6 +25,7 @@ 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" 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..b225328dd --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_spec.py @@ -0,0 +1,221 @@ +# 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 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 NodeProcessSpec: + """One resolved lane process owned by a node worker.""" + + process_id: str + command: tuple[str, ...] + gpu_indices: tuple[int, ...] + launch_delay_seconds: int + + 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") + + +@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") + + +@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, ...] + + def __post_init__(self) -> None: + 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.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): + raise ValueError("node processes must not share GPUs") + + +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, + } + for process in node.processes + ], + } + for node in spec.nodes + ], + } + 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: + root = _require_mapping(payload, {"schema_version", "resolved_gpus_per_node", "required_model_path", "nodes"}) + 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, + ) + + +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: + value = _require_mapping(payload, {"process_id", "command", "gpu_indices", "launch_delay_seconds"}) + 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"]), + ) + + +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 + + +__all__ = [ + "NodeProcessSpec", + "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..d8e8a0eff --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/node_worker.py @@ -0,0 +1,272 @@ +# 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 data_designer.slurm.runtime.node_spec import ( + NodeProcessSpec as NodeProcessSpec, +) +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, +) + +_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._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 process in node.processes: + self._wait_for_launch(process.launch_delay_seconds, started_at) + if self._stop_requested: + return 1 + self._children.append(self._start_process(process, visible_gpus)) + 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 + self._cleanup_complete = all(child.poll() is not None for child in self._children) + if not self._cleanup_complete: + raise RuntimeError("node process cleanup is incomplete") + + 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) + 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) -> None: + while not self._stop_requested: + remaining = delay_seconds - (time.monotonic() - started_at) + if remaining <= 0: + return + time.sleep(min(remaining, _POLL_INTERVAL_SECONDS)) + + 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) + _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/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh index d9424183d..a16c54392 100644 --- 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 @@ -18,13 +18,15 @@ dd_read_control_plan() { local plan=$1 jq -e ' .schema_version == 1 - and .client.host_node_index == 0 - and ([.deployments[].node_indices] | all(. == [0])) + and ([.deployments[].node_indices[]] | length > 0) + and ([.deployments[].node_indices[]] | all(type == "number" and . >= 0)) and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) ' "${plan}" >/dev/null DD_CLIENT_IMAGE=$(jq -er '.client.image.path' "${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)] @@ -140,10 +142,7 @@ dd_read_step() { .stdout_path, "\u0000", .stderr_path, "\u0000", (.launch_delay_seconds | tostring), "\u0000", - (.readiness.host // ""), "\u0000", - (.readiness.port // "" | tostring), "\u0000", - (.readiness.path // ""), "\u0000", - (.readiness.deadline_seconds // "" | tostring), "\u0000" + (.kill_on_bad_exit | tostring), "\u0000" ' "${manifest}" ) DD_STEP_IMAGE=${DD_STEP_FIELDS[0]} @@ -151,10 +150,7 @@ dd_read_step() { DD_STEP_STDOUT=${DD_STEP_FIELDS[2]} DD_STEP_STDERR=${DD_STEP_FIELDS[3]} DD_STEP_DELAY=${DD_STEP_FIELDS[4]} - DD_STEP_PROBE_HOST=${DD_STEP_FIELDS[5]} - DD_STEP_PROBE_PORT=${DD_STEP_FIELDS[6]} - DD_STEP_PROBE_PATH=${DD_STEP_FIELDS[7]} - DD_STEP_PROBE_DEADLINE=${DD_STEP_FIELDS[8]} + DD_STEP_KILL_ON_BAD_EXIT=${DD_STEP_FIELDS[5]} dd_read_null_values DD_STEP_COMMAND < <( jq -j --arg step_id "${step_id}" '.steps[] | select(.step_id == $step_id) | .command[] | ., "\u0000"' \ "${manifest}" @@ -163,6 +159,21 @@ dd_read_step() { 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}" 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 index 78366c390..1b3204d41 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -11,19 +11,70 @@ 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.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.paths import get_container_path -from data_designer.slurm.runtime.ports import allocation_ports +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.""" @@ -59,16 +110,17 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] context.plan.client.host_node_index, *(index for deployment in context.plan.deployments for index in deployment.node_indices), } - if node_indices != {0}: + node_count = max(node_indices) + 1 + if node_indices != set(range(node_count)): raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one-node runtime received a multi-node execution plan", + "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": 1, - "SLURM_NODEID": 0, + "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: @@ -125,19 +177,26 @@ def _verify_artifacts(context: AllocationContext) -> None: @staticmethod def verify_ports(context: AllocationContext, environment: Mapping[str, str]) -> None: - """Verify that every planned one-node port is currently bindable.""" - ports = allocation_ports(context, environment) + """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(("127.0.0.1", port)) + reservation.bind(("0.0.0.0", port)) except OSError as error: raise SlurmRuntimeError( SlurmRuntimeErrorCode.PREFLIGHT_FAILED, - "one or more resolved allocation ports are unavailable", + "one or more local allocation ports are unavailable", ) from error finally: for reservation in reservations: @@ -195,4 +254,9 @@ def _parse_gpu_count(value: str | None) -> int: return len(values) -__all__ = ["AllocationPreflight", "SystemAllocationPreflight"] +__all__ = [ + "AllocationLayout", + "AllocationPreflight", + "SystemAllocationPreflight", + "validate_allocation_layout", +] 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 index 64dd51fe2..402677767 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -13,6 +13,8 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import SplitResult, urlsplit +from data_designer.slurm.runtime.network import validate_host_name + _MAXIMUM_REQUEST_BYTES = 64 * 1024 * 1024 _MAXIMUM_RESPONSE_BYTES = 64 * 1024 * 1024 _HOP_HEADERS = frozenset( @@ -191,9 +193,14 @@ def main(arguments: Sequence[str] | None = None) -> int: 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) parsed = parser.parse_args(arguments) - backends = tuple(_parse_backend(value) for value in parsed.backend) + 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: @@ -204,15 +211,18 @@ def main(arguments: Sequence[str] | None = None) -> int: return 0 -def _parse_backend(value: str) -> _Backend: +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 parsed.hostname != "127.0.0.1" + 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 {"", "/"} @@ -220,8 +230,8 @@ def _parse_backend(value: str) -> _Backend: or parsed.fragment or port is None ): - raise argparse.ArgumentTypeError("backends must be loopback HTTP origins") - return _Backend(parsed.hostname, port) + raise argparse.ArgumentTypeError("backends must be allowed HTTP origins") + return _Backend(host, port) def _retry_after_headers(headers: dict[str, str], retry_after_seconds: int | None) -> dict[str, str]: 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..4cd6b4401 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/server_manifest.py @@ -0,0 +1,244 @@ +# 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.""" + preflight_steps = tuple( + build_distributed_server_step( + deployment, + context, + runtime_container_root, + log_directory, + layout, + operation="preflight", + ) + for deployment in deployments + if _requires_server_preflight(deployment, context) + ) + serving_steps = tuple( + step + for deployment in deployments + for step in _build_deployment_serving_steps( + deployment, + context, + runtime_container_root, + log_directory, + layout, + ) + ) + return 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"], +) -> 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) + readiness = ( + tuple( + RuntimeProbeSpec( + host=layout.get_host(probe.node_index), + port=probe.port, + path=probe.path, + deadline_seconds=probe.deadline_seconds, + ) + for probe in deployment.readiness_probes + ) + if operation == "serve" + else () + ) + step_id = f"{deployment.deployment_id}-{operation}" + return RuntimeStepSpec( + step_id=step_id, + role=RuntimeStepRole.SERVER if operation == "serve" else RuntimeStepRole.SERVER_PREFLIGHT, + 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(deployment.gpus_per_node)), + 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: + 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, + 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/step_runner.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh index 7dc5aeae3..34e5c6aeb 100644 --- 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 @@ -66,12 +66,17 @@ dd_materialize_step_environment() { } dd_build_srun_command() { - local index gpu_mask=0 container_names= + local index gpu_mask=0 container_names= node_list= + ((${#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=1 - --ntasks=1 + "--nodes=${#DD_STEP_NODE_HOSTS[@]}" + "--ntasks=${#DD_STEP_NODE_HOSTS[@]}" + --ntasks-per-node=1 + "--nodelist=${node_list}" --exact --overlap --unbuffered @@ -79,6 +84,9 @@ dd_build_srun_command() { "--cpus-per-task=${DD_STEP_CPUS}" "--container-image=${DD_STEP_IMAGE}" ) + 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 for index in "${DD_STEP_GPU_INDICES[@]+"${DD_STEP_GPU_INDICES[@]}"}"; do @@ -138,6 +146,7 @@ dd_run_control_phase() { srun --nodes=1 --ntasks=1 + "--nodelist=${DD_CLIENT_HOST}" --exact --overlap --unbuffered 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 index 74ca803ea..555c0a5db 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -274,8 +274,24 @@ def build_endpoint_command( plan: ResolvedSlurmRunPlan, runtime_proxy_path: Path, port: int, + *, + backend_hosts: tuple[str, ...] | None = None, ) -> tuple[str, ...]: - backends = tuple(f"http://127.0.0.1:{backend.port}" for backend in deployment.backend_endpoints) + 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 () return ( @@ -284,6 +300,7 @@ def build_endpoint_command( "--listen-port", str(port), *retry_arguments, + *allowed_host_arguments, *(argument for backend in backends for argument in ("--backend", backend)), ) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 10fb6c277..c9cedca44 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -4,11 +4,19 @@ from __future__ import annotations from dataclasses import replace +from pathlib import Path -from conftest import RuntimeCase +from conftest import RuntimeCase, relocate_plan +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest -from data_designer.slurm.runtime.models import RuntimeStepRole +from data_designer.slurm.runtime.distributed import build_vllm_process_command +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 @@ -22,6 +30,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru {"SLURM_JOB_GPUS": "0"}, runtime_root=runtime_root, log_directory=log_directory, + layout=AllocationLayout(("compute-001",)), ) reloaded = RuntimeBootstrapManifest.model_validate_json(manifest.serialize_json()) @@ -44,6 +53,8 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru 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) def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runtime_case: RuntimeCase) -> None: @@ -71,6 +82,7 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti {"SLURM_JOB_GPUS": "0"}, runtime_root=context.attempt_directory / "runtime", log_directory=context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), ) preflight = manifest.steps[0].command @@ -80,3 +92,107 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti 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, + } + ), + "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=context.attempt_directory / "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-replica-00000-rank-00000") + 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 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.command[2] == "/workspace/primary/models/model-1" + executor_index = remote_server.command.index("--distributed-executor-backend") + assert remote_server.command[executor_index + 1] == "uni" + assert tuple(probe.host for probe in remote_server.readiness) == ("compute-003",) + + +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_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 561588810..3c27e87c9 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -60,6 +60,8 @@ def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.Capt "runtime", "--manifest", "runtime-manifest.json", + "--node-host", + "compute-001", ) ) == 64 @@ -465,5 +467,13 @@ def _phase_arguments( ) 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()) + 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..a9511f1d8 --- /dev/null +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.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 sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from data_designer.slurm.runtime import node_worker as runtime_node_worker +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 + + +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] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + nonlocal calls + del arguments, keywords + 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 + + +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_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 4cc0890e6..604e3e737 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -6,6 +6,8 @@ import hashlib import os from pathlib import Path +from types import SimpleNamespace +from typing import cast import pytest from conftest import RuntimeCase @@ -15,7 +17,40 @@ 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 SystemAllocationPreflight, _verify_artifact +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: diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py index 0a600f1bf..5387aa4ec 100644 --- a/packages/data-designer-slurm/tests/runtime/test_proxy.py +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -189,6 +189,12 @@ def test_proxy_rejects_non_loopback_or_malformed_backend(value: str) -> None: _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() -> None: pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) first = pool.acquire(frozenset()) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index ef877ff6a..555b003d3 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -35,7 +35,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( _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\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 @@ -57,6 +57,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "SLURM_ARRAY_JOB_ID": "4101", "SLURM_ARRAY_TASK_ID": "0", "SLURM_JOB_NUM_NODES": "1", + "SLURM_JOB_NODELIST": "compute-001", "SLURM_JOB_GPUS": "0", "SLURM_NODEID": "0", } @@ -134,6 +135,62 @@ def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: assert completed.returncode == 0, 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_CPUS=4 +DD_STEP_IMAGE=/images/server.sqsh +DD_STEP_KILL_ON_BAD_EXIT=true +DD_GPU_REQUEST_MODE=gres +DD_CONTAINER_MOUNTS= +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* ]] +""" + + 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 _artifact(root: Path, name: str) -> tuple[str, str]: path = root / f"{name}.artifact" content = name.encode() @@ -245,6 +302,8 @@ def _step( "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": {}, @@ -252,7 +311,7 @@ def _step( "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, + "readiness": [readiness] if readiness is not None else [], } From 6636e2c61cedf35c85643ffed7767f57194e2b7c Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:10:29 -0300 Subject: [PATCH 25/62] feat(slurm): add profile setup and image lifecycle wiring (#933) * feat(slurm): add profile setup and image lifecycle wiring * fix(slurm): address profile and lifecycle review * fix(slurm): require terminal cancellation evidence --- .../src/data_designer/slurm/cli.py | 52 +++ .../data_designer/slurm/launcher/client.py | 14 +- .../data_designer/slurm/launcher/parsing.py | 21 + .../data_designer/slurm/services/__init__.py | 12 + .../data_designer/slurm/services/errors.py | 2 + .../slurm/services/image_lifecycle.py | 221 ++++++++++ .../data_designer/slurm/services/images.py | 9 +- .../data_designer/slurm/services/profiles.py | 393 +++++++++++++++++ .../data_designer/slurm/services/wiring.py | 25 +- .../data-designer-slurm/tests/conftest.py | 2 +- .../tests/launcher/test_client.py | 4 +- .../tests/launcher/test_parsing.py | 17 + .../tests/services/test_image_wiring.py | 403 ++++++++++++++++++ .../tests/services/test_profile_service.py | 339 +++++++++++++++ .../tests/services/test_wiring.py | 13 +- .../golden/slurm/sinfo_gres.txt | 2 +- .../tests/slurm_test_fakes/test_slurm.py | 2 +- .../data-designer-slurm/tests/test_cli.py | 89 +++- scripts/test_slurm_package_install.py | 25 ++ 19 files changed, 1612 insertions(+), 33 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py create mode 100644 packages/data-designer-slurm/tests/services/test_image_wiring.py create mode 100644 packages/data-designer-slurm/tests/services/test_profile_service.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index f2c9e12fd..ebe7de55e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -14,11 +14,13 @@ from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json +from data_designer.slurm.images.records import validate_oci_source_for_lifecycle from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, create_slurm_image_service, + create_slurm_profile_service, create_slurm_run_service, ) @@ -37,7 +39,9 @@ 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() @@ -94,6 +98,46 @@ def cancel_command( _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"), + 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(...), @@ -113,6 +157,14 @@ def add() -> BaseModel: operation, "OCI image source must be digest-qualified as name@sha256:", ) + try: + validate_oci_source_for_lifecycle(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(source), kind=kind, source=source) return create_slurm_image_service(profile_file=profile_file, cluster=cluster).add(request, replace=replace) 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 index cc40fbeff..e9c5fbf06 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -24,6 +24,7 @@ ) from data_designer.slurm.launcher.parsing import ( parse_accounting, + parse_default_partition_gpu_counts, parse_gpu_counts, parse_named_jobs, parse_queue, @@ -224,12 +225,13 @@ def release(self, job_id: int) -> None: self._run((self._executables.scontrol, "release", _format_job_id(job_id))) def query_gpu_counts(self, *, partition: Identifier | None = None) -> tuple[int, ...]: - """Return configured GPU counts reported for eligible node groups.""" - command = [self._executables.sinfo, "--noheader", "--format=%G"] - if partition is not None: - if type(partition) is not str or _IDENTIFIER_PATTERN.fullmatch(partition) is None: - raise ValueError("Slurm partition must be a valid identifier") - command.append(f"--partition={partition}") + """Return configured GPU counts for the requested or default partition.""" + 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 _IDENTIFIER_PATTERN.fullmatch(partition) is None: + raise ValueError("Slurm partition must be a valid identifier") + command = (self._executables.sinfo, "--noheader", "--format=%G", f"--partition={partition}") return parse_gpu_counts(self._run(command)) def _run( 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 index 32766d782..56e4f97a6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -146,6 +146,27 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: 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("+") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 22b603c7a..5c0ae9c7e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -15,6 +15,13 @@ 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 ( SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, @@ -47,6 +54,10 @@ "SlurmPersistedAttemptStatus", "SlurmPersistedRunStatus", "SlurmPersistedShardStatus", + "SlurmProfileInitialization", + "SlurmProfileMatch", + "SlurmProfileService", + "SlurmProfileValidation", "SlurmRunArtifactPublisher", "SlurmRunBackend", "SlurmRunCancellation", @@ -57,6 +68,7 @@ "SlurmServiceErrorCode", "SlurmServiceOperation", "create_slurm_image_service", + "create_slurm_profile_service", "create_slurm_run_service", ] 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 index aa08fffed..7e7e19814 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -32,6 +32,8 @@ class SlurmServiceOperation(str, Enum): EXECUTE_RUN = "execute_run" STATUS_RUN = "status_run" CANCEL_RUN = "cancel_run" + INIT_PROFILE = "init_profile" + VALIDATE_PROFILE = "validate_profile" RESOLVE_IMAGE = "resolve_image" ADD_IMAGE = "add_image" LIST_IMAGES = "list_images" 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..a335e248f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/image_lifecycle.py @@ -0,0 +1,221 @@ +# 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 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, + 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.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] + +_POLL_INTERVAL_SECONDS = 300.0 +_ACCOUNTING_EXIT_LAG = timedelta(minutes=5) + + +class _TerminalLifecycleError(RuntimeError): + pass + + +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, + ) -> 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 + + 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: + 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 + + try: + self._wait_for_success(receipt.job_id) + return publish_completed_image_lifecycle(prepared, replace=replace) + 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) + 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) + raise + + def _wait_for_success(self, job_id: int) -> None: + client = _RecordingObservationClient(self._launcher) + observations = SchedulerObservationCollector(client) + previous: SchedulerObservation | None = None + completed_without_accounting_deadline: datetime | None = None + while True: + observed_at = self._clock() + observation = observations.collect( + (job_id,), + observed_at=observed_at, + previous={job_id: previous}, + )[0] + 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: + 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 + ): + raise _TerminalLifecycleError("completed image lifecycle job has no successful exit evidence") + return + if is_scheduler_failure_state(observation.state): + 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 _cancel_and_cleanup(self, job_id: int, prepared: PreparedImageLifecycleJob) -> None: + try: + self._launcher.cancel(job_id) + if not self._wait_for_termination(job_id): + return + except (SlurmLauncherError, SlurmStateError, OSError, ValueError): + return + _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] + 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) + + +__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 index 9a2a12253..531c468c5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -12,7 +12,7 @@ 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 +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, @@ -98,6 +98,13 @@ def add(self, request: ImageBuildRequest, *, replace: bool = False) -> Registere 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: 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..550af5562 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -0,0 +1,393 @@ +# 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 +from collections.abc import Callable, Mapping +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 create_restrictive_temporary_file, 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, + ) + selected_patterns = host_patterns or _normalize_hostnames(self._hostname_resolver()) + 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=selected_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_verified_directory(path.parent, resource_name="profile") 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 + + +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/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index b13834f56..f48757534 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -52,6 +52,7 @@ from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.services.artifacts import StateRunArtifactPublisher 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 ( SlurmPersistedAttemptStatus, @@ -624,9 +625,10 @@ def cancel(self, run_id: Identifier) -> SlurmRunCancellation: class _RegistryImageBackend: - def __init__(self, workspace_root: str) -> None: + 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: @@ -645,12 +647,7 @@ def resolve(self, reference: ImageRef, *, expected_kind: ImageKind) -> ResolvedI ) from None def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: - del request, replace - raise SlurmServiceError( - SlurmServiceErrorCode.UNAVAILABLE, - SlurmServiceOperation.ADD_IMAGE, - "image registration is not available; use a pre-registered image", - ) + return self._lifecycle.add(request, replace=replace) def list(self) -> tuple[RegisteredImage, ...]: return self._invoke_registry(SlurmServiceOperation.LIST_IMAGES, self._store.list_images) @@ -718,10 +715,22 @@ def create_slurm_image_service( 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, ) -> SlurmImageService: """Create the production image service for one selected cluster profile.""" selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) - backend = _RegistryImageBackend(selected.profile.workspace_root) + command_client = launcher or SlurmCommandClient() + lifecycle = SlurmImageLifecycleManager( + selected, + command_client, + lifecycle_id_factory=lifecycle_id_factory, + clock=clock, + sleep=sleep, + ) + backend = _RegistryImageBackend(selected.profile.workspace_root, lifecycle) return SlurmImageService(backend, backend) diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index a46fcb160..ed14a4820 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -197,7 +197,7 @@ def fake_slurm_runner() -> FakeSlurmRunner: ), ), sinfo_responses={ - ("sinfo", "--noheader", "--format=%G"): FakeCommandResponse( + ("sinfo", "--noheader", "--format=%P|%G"): FakeCommandResponse( stdout=(SLURM_GOLDEN_DIRECTORY / "sinfo_gres.txt").read_text() ) }, diff --git a/packages/data-designer-slurm/tests/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index e4be4ea4a..a9c247eff 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -236,11 +236,11 @@ def test_client_rejects_unbounded_or_invalid_job_selectors(fake_slurm_runner: Fa assert fake_slurm_runner.calls == [] -def test_client_queries_bounded_gpu_inventory(fake_slurm_runner: FakeSlurmRunner) -> None: +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=%G")] + assert fake_slurm_runner.calls == [("sinfo", "--noheader", "--format=%P|%G")] def test_client_queries_partition_scoped_gpu_inventory() -> None: diff --git a/packages/data-designer-slurm/tests/launcher/test_parsing.py b/packages/data-designer-slurm/tests/launcher/test_parsing.py index 7509f93c2..c6d1f955e 100644 --- a/packages/data-designer-slurm/tests/launcher/test_parsing.py +++ b/packages/data-designer-slurm/tests/launcher/test_parsing.py @@ -12,6 +12,7 @@ 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, @@ -174,6 +175,22 @@ def test_parse_gpu_counts_normalizes_configured_gres(output: str, expected: tupl 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", ( 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..e7b858353 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -0,0 +1,403 @@ +# 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 _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 == [300.0, 300.0] + assert not _job_directory(workspace).exists() + + +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 == [300.0, 300.0, 300.0] + 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 == [300.0, 300.0, 300.0] + 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 == [300.0, 300.0, 300.0] + 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="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..1440ca609 --- /dev/null +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -0,0 +1,339 @@ +# 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 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_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, + hostname_resolver=lambda: ("Node.EXAMPLE.test", "node.example.test", "node-short"), + ).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 == [ + "node.example.test", + "node-short", + ] + + +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.yml")), + ids=("unsupported-suffix", "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_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index bbe2b39d1..475bfd105 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -14,7 +14,6 @@ from data_designer.slurm.config import ( BuilderInput, DataDesignerSlurmConfig, - ImageBuildRequest, SecretRef, SlurmProfile, SlurmProfileCatalog, @@ -765,7 +764,7 @@ def test_production_status_maps_unknown_run_to_not_found( assert caught.value.code is SlurmServiceErrorCode.NOT_FOUND -def test_production_image_registry_operations_and_lifecycle_gap( +def test_production_image_registry_operations( tmp_path: Path, profile_catalog: SlurmProfileCatalog, authored_run_single: DataDesignerSlurmConfig, @@ -780,13 +779,3 @@ def test_production_image_registry_operations_and_lifecycle_gap( assert selected == removed == images[0] assert len(service.list()) == 1 - with pytest.raises(SlurmServiceError) as caught: - service.add( - ImageBuildRequest( - name="new-image", - kind="client", - source=f"registry.example/client@sha256:{'1' * 64}", - ) - ) - assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE - assert str(caught.value) == "image registration is not available; use a pre-registered image" 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 index 25175f090..3e5fe27b7 100644 --- 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 @@ -1 +1 @@ -gpu:2 +batch*|gpu:2 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py index c47b21432..3051ca5ee 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_slurm.py @@ -245,7 +245,7 @@ def test_fake_slurm_runner_exposes_retry_terminal_spellings() -> None: def test_fake_slurm_runner_bounds_sinfo_queries(fake_slurm_runner: FakeSlurmRunner) -> None: - response = fake_slurm_runner.run(("/usr/bin/sinfo", "--noheader", "--format=%G"), check=True) + 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"): diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 417657843..392e3e976 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -139,9 +139,96 @@ def test_image_add_rejects_mutable_oci_source(source: str) -> None: } +def test_profile_init_creates_starter_and_emits_validation_command(tmp_path: Path) -> None: + profile_file = tmp_path / "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 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_only_m2_run_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", "image")) + assert all(command in result.stdout for command in ("execute", "status", "cancel", "image", "profile")) assert all(command not in result.stdout for command in ("retry", "merge", "benchmark")) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 412a1b00b..6be54bcdd 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -113,6 +113,8 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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 @@ -171,6 +173,29 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange assert StateResumeWorkspace is ContractResumeWorkspace +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) From 7999f73db77a4b4649837eece173fefb080e9304 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:21:56 -0300 Subject: [PATCH 26/62] feat(slurm): wire retry and merge commands (#931) * feat(slurm): wire retry and merge commands * fix(slurm): harden retry and collection actions * fix(slurm): preserve retry confirmation intent * fix(slurm): recover active retry previews * fix(slurm): address retry and merge review feedback --- .../src/data_designer/slurm/cli.py | 92 +++++ .../data_designer/slurm/services/__init__.py | 4 + .../data_designer/slurm/services/errors.py | 2 + .../data_designer/slurm/services/results.py | 54 ++- .../slurm/services/retry_collection.py | 311 ++++++++++++++++ .../src/data_designer/slurm/services/run.py | 147 +++++++- .../data_designer/slurm/services/wiring.py | 40 ++- .../data_designer/slurm/state/collection.py | 4 +- .../data_designer/slurm/state/destinations.py | 20 +- .../src/data_designer/slurm/state/observer.py | 20 +- .../src/data_designer/slurm/state/retry.py | 46 +++ .../slurm/state/retry_storage.py | 12 +- .../tests/launcher/test_collection.py | 30 ++ .../tests/services/test_services.py | 116 ++++++ .../tests/services/test_wiring.py | 333 +++++++++++++++++- .../tests/state/test_observer.py | 23 ++ .../tests/state/test_retry_collection.py | 30 ++ .../data-designer-slurm/tests/test_cli.py | 165 ++++++++- scripts/test_slurm_package_install.py | 2 + 19 files changed, 1434 insertions(+), 17 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index ebe7de55e..730ce513c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -5,6 +5,8 @@ import re from collections.abc import Callable +from enum import Enum +from functools import partial from pathlib import Path from typing import NoReturn, TypeVar @@ -33,6 +35,13 @@ SlurmServiceErrorCode.INTERNAL: 1, } + +class _RetryResumeMode(str, Enum): + NEVER = "never" + ALWAYS = "always" + IF_POSSIBLE = "if_possible" + + app = typer.Typer( name="slurm", help="Run Data Designer workloads on Slurm", @@ -98,6 +107,89 @@ def cancel_command( _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_json({"operation": operation.value, "state": "declined"}) + return + shard_ids = planned.shard_ids + resume = _RetryResumeMode(planned.effective_resume_mode) + 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( + input_path: Path = typer.Option(..., "--input-path", file_okay=False), + output_path: Path = typer.Option(..., "--output-path", file_okay=False), + num_partitions: int | None = typer.Option(None, "--num-partitions", min=1), + 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 + result = _invoke( + operation, + lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).collect( + 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), diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 5c0ae9c7e..79722df62 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -23,9 +23,11 @@ create_slurm_profile_service, ) from data_designer.slurm.services.results import ( + SlurmCollectionExecution, SlurmPersistedAttemptStatus, SlurmPersistedRunStatus, SlurmPersistedShardStatus, + SlurmRetryExecution, SlurmRunCancellation, SlurmRunExecution, ) @@ -48,6 +50,7 @@ "SlurmBatchScriptRenderer", "SlurmBenchmarkBackend", "SlurmBenchmarkService", + "SlurmCollectionExecution", "SlurmImageManager", "SlurmImageResolver", "SlurmImageService", @@ -58,6 +61,7 @@ "SlurmProfileMatch", "SlurmProfileService", "SlurmProfileValidation", + "SlurmRetryExecution", "SlurmRunArtifactPublisher", "SlurmRunBackend", "SlurmRunCancellation", 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 index 7e7e19814..ecf336f49 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/errors.py @@ -32,6 +32,8 @@ class SlurmServiceOperation(str, Enum): 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" 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 index 2417b000d..5d25a51f9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -7,12 +7,20 @@ from typing import Literal -from pydantic import PositiveInt, model_validator - -from data_designer.slurm.contracts import ContractValue, Identifier, Sha256Digest +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, @@ -104,10 +112,50 @@ def validate_jobs(self) -> SlurmRunCancellation: 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..46de9f076 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -0,0 +1,311 @@ +# 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 +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.destinations import CollectionDestinationResolver +from data_designer.slurm.state.outputs import RetryPlan + +_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]) -> None: + self._workspace_root = workspace_root + self._launcher = launcher + self._clock = clock + + 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) + 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: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + operation = SlurmServiceOperation.COLLECT_RUN + run_id = self._resolve_run_input_path(input_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.output.partitions: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "num_partitions must match the persisted run output partitions", + ) + effective_partitions = plan.output.partitions + status = SlurmCollectionCoordinator(self._workspace_root, run_id, self._launcher).submit( + destination=destination, + submitted_at=self._clock(), + ) + if status.scheduler is None: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + operation, + "Slurm collection submission was not accepted", + ) + return SlurmCollectionExecution( + run_id=run_id, + collection_id=status.collection_id, + state=status.state, + job_id=status.scheduler, + output_path=destination.as_posix(), + num_partitions=effective_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 index 68eeaa90a..f40ea2614 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -5,13 +5,15 @@ from __future__ import annotations +from collections.abc import Sequence +from functools import partial from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from pydantic import TypeAdapter, ValidationError from data_designer.slurm.config import DataDesignerSlurmConfig -from data_designer.slurm.contracts import Identifier +from data_designer.slurm.contracts import Identifier, ShardId from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services.errors import ( SlurmServiceError, @@ -21,12 +23,15 @@ _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): @@ -68,6 +73,25 @@ def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: 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: Path, + *, + destination: Path, + num_partitions: int | None, + ) -> SlurmCollectionExecution: + """Submit or recover one winner-driven collection.""" + class SlurmRunService: """Coordinate public run operations through package-owned boundaries. @@ -206,6 +230,64 @@ def cancel_run() -> SlurmRunCancellation: 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, + num_partitions: int | None = None, + ) -> SlurmCollectionExecution: + """Submit or recover collection for one managed run directory.""" + operation = SlurmServiceOperation.COLLECT_RUN + if not isinstance(input_path, str | Path) or not isinstance(destination, str | Path): + raise _make_invalid_request_error(operation, "input_path and destination must be paths") + 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 = Path(input_path).expanduser().resolve() + normalized_destination = 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( @@ -216,8 +298,69 @@ def _require_backend(self, operation: SlurmServiceOperation) -> SlurmRunBackend: 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: Path, + *, + destination: Path, + 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 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 index f48757534..691849bc4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Protocol, TypeVar +from typing import Literal, Protocol, TypeVar from uuid import uuid4 from pydantic import JsonValue @@ -37,7 +37,7 @@ load_builder_payload, resolve_profile, ) -from data_designer.slurm.contracts import Identifier +from data_designer.slurm.contracts import Identifier, ShardId from data_designer.slurm.images.errors import ImageConflictError, ImageNotFoundError, SlurmImageError from data_designer.slurm.images.records import RegisteredImage from data_designer.slurm.images.registry import ImageRegistryStore @@ -55,12 +55,15 @@ 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 ( @@ -274,6 +277,11 @@ def __init__( self._publisher = publisher self._clock = clock self._source_environment = source_environment + self._retry_collection = RunRetryCollectionBackend( + selected_profile.profile.workspace_root, + launcher, + clock, + ) def execute( self, @@ -623,6 +631,34 @@ def cancel(self, run_id: Identifier) -> SlurmRunCancellation: ) 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: Path, + *, + destination: Path, + 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: 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 index 396ccdcf5..78b1b0fab 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -293,9 +293,11 @@ def _validate_existing_destination( 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 not None and 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 StateCorruptionError("persisted collection destination does not match the requested destination") + raise StateConflictError("persisted collection destination does not match the requested destination") if plan.host_destination != resolved.host_path or plan.container_destination != resolved.container_path: raise StateCorruptionError("persisted collection destination does not match the resolved plan") if status.state is CollectionState.SUCCEEDED: 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 index 78c8a2bc4..1c738a2d6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -10,7 +10,7 @@ from pathlib import Path from data_designer.slurm.config import ContainerMount -from data_designer.slurm.contracts import is_path_below, validate_absolute_path +from data_designer.slurm.contracts import is_path_below, paths_overlap, validate_absolute_path from data_designer.slurm.images.records import validate_enroot_mount_path from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.errors import StateConflictError @@ -41,6 +41,7 @@ def resolve( except ValueError as error: raise StateConflictError("collection destination must be a normalized absolute path") from error workspace_root = plan.selected_profile.profile.workspace_root + _validate_destination_overlap(plan, host_path) workspace_mount = ContainerMount(source=workspace_root, target=workspace_root, read_only=False) authorized = { (mount.source, mount.target): mount @@ -89,4 +90,21 @@ def validate_persisted( 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, "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/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index 4b018d726..1c49d2952 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -5,6 +5,7 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -92,6 +93,13 @@ def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: 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() @@ -114,6 +122,7 @@ def refresh(self, *, observed_at: datetime | None = None) -> RunStatus: self._refresh_shard( _ShardSnapshot(run, plan, shard, attempts_by_shard[shard.shard_id]), batch, + persist=persist, ) for shard in shards ) @@ -172,9 +181,12 @@ def _refresh_shard( self, expected: _ShardSnapshot, batch: _ObservationBatch, + *, + persist: bool, ) -> ShardStatus: try: - with self._storage.acquire_shard_lock(expected.shard.shard_id): + 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( @@ -193,6 +205,7 @@ def _refresh_shard( batch, attempt, winner, + persist=persist, ) for attempt in attempts ) @@ -214,13 +227,16 @@ def _build_attempt_status( 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] - self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) + if persist: + self._persist_observation(attempt, batch.previous[attempt.scheduler], scheduler) effective_state = reconcile_attempt_observation( attempt, readiness, 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 index 607ab3fc9..a8b9188cd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -114,6 +114,52 @@ def retry( 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: 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 index 489af26b0..cbb10803d 100644 --- 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 @@ -49,7 +49,10 @@ def acquire_lock(self) -> Iterator[None]: def get_next_retry_id(self) -> Identifier: """Return the next monotonic retry identity.""" - return f"retry-{len(self.list_retry_ids()) + 1:04d}" + 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.""" @@ -156,6 +159,13 @@ def read_status(self, retry_id: Identifier) -> RetryStatus: 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: diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index a79a8187a..478e202d9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -191,6 +191,36 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( 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"), + ], +) +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_output_subtree(multi_node_plan: ResolvedSlurmRunPlan) -> None: + requested = (Path(multi_node_plan.output.root) / "collected").as_posix() + + assert CollectionDestinationResolver().resolve(multi_node_plan, requested).host_path == requested + + def test_destination_requires_one_unique_most_specific_mapping( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index 80a5532e3..b3df00eb7 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -5,6 +5,7 @@ import pickle from pathlib import Path +from unittest.mock import Mock import pytest from slurm_test_fakes import ( @@ -29,9 +30,11 @@ SlurmBatchScriptRenderer, SlurmBenchmarkBackend, SlurmBenchmarkService, + SlurmCollectionExecution, SlurmImageManager, SlurmImageResolver, SlurmImageService, + SlurmRetryExecution, SlurmRunArtifactPublisher, SlurmRunBackend, SlurmRunExecution, @@ -41,6 +44,7 @@ SlurmServiceErrorCode, SlurmServiceOperation, ) +from data_designer.slurm.state import CollectionState GOLDEN_DIRECTORY = Path(__file__).parents[1] / "slurm_test_fakes" / "golden" / "rendered" @@ -224,6 +228,118 @@ def cancel(self, run_id): assert backend.calls == [(authored_run_single, tmp_path.resolve(), True, 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_uses_persisted_collection_partitions_when_omitted(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=2, + ) + 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, + ) + + +@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: diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 475bfd105..cd453f598 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -7,9 +7,11 @@ from collections.abc import Mapping 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 from data_designer.slurm.client.dependencies import ResolvedClientDependencies from data_designer.slurm.config import ( BuilderInput, @@ -27,11 +29,13 @@ SlurmJobSubmissionReceipt, SlurmProcessExitCode, SlurmQueueEntry, + SlurmSubmissionMatch, ) from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, + SlurmServiceOperation, create_slurm_image_service, create_slurm_run_service, ) @@ -39,6 +43,7 @@ AttemptLifecycleState, AttemptManifest, AttemptTerminalClassification, + CollectionState, RunManifest, SchedulerIdentity, SchedulerState, @@ -54,6 +59,7 @@ def __init__( self, gpu_counts: tuple[int, ...] = (), *, + submission_job_ids: tuple[int, ...] = (42,), cancel_error: Exception | None = None, release_error: Exception | None = None, ) -> None: @@ -64,6 +70,8 @@ def __init__( 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 = submission_job_ids self.gpu_counts = gpu_counts self.cancel_error = cancel_error self.release_error = release_error @@ -78,7 +86,8 @@ def submit_script( self.submissions.append(script) self.held_submissions.append(hold) self.exported_environments.append(dict(export_environment or {})) - return SlurmJobSubmissionReceipt(job_id=42) + 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) @@ -102,6 +111,15 @@ 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__( @@ -442,7 +460,10 @@ def test_status_does_not_expire_requeue_window_from_another_attempt_clock( single_node_plan: ResolvedSlurmRunPlan, ) -> None: authored = authored_run_single.model_copy( - update={"array_tasks": authored_run_single.array_tasks.model_copy(update={"count": 2})} + 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() @@ -483,6 +504,314 @@ def test_status_does_not_expire_requeue_window_from_another_attempt_clock( assert requeued.shards[0].attempts[0].attempt.state is AttemptLifecycleState.PENDING +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, +) -> 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(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, 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 + + +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=2) + + 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, submitted_at: datetime) -> SimpleNamespace: + del destination, submitted_at + return next(statuses) + + monkeypatch.setattr(retry_collection_module, "SlurmCollectionCoordinator", lambda *_: Coordinator()) + output_path = tmp_path / "collected" + + active = service.collect(tmp_path / "runs/run-wired", destination=output_path) + completed = service.collect(tmp_path / "runs/run-wired", destination=output_path) + + assert active.state is CollectionState.RUNNING + assert completed.state is CollectionState.SUCCEEDED + assert active.job_id == completed.job_id == 43 + + def test_auto_gpu_resolution_rejects_mixed_node_shapes( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 0a49de6be..41bc6dc73 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -442,6 +442,29 @@ def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[S SlurmStateReconciler(case.workspace, case.plan.run_id, MutatingClient()).refresh(observed_at=observed_at) +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) diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 441744846..87ee5f001 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -129,6 +129,34 @@ def test_retry_refreshes_failed_shard_and_publishes_exact_next_attempt( 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, @@ -710,6 +738,8 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( assert submitted.state is CollectionState.SUBMITTED assert submitted.scheduler == 5101 + with pytest.raises(StateConflictError, match="requested destination"): + coordinator.submit(destination=Path(case.plan.output.root).parent / "other") script = cast(str, runner.inputs[-1]) assert "data_designer.slurm.state.collection_worker" in script assert "--gpus" not in script diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 392e3e976..aad4367ee 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -12,16 +12,21 @@ import data_designer.slurm.cli as cli_module from data_designer.slurm.config import DataDesignerSlurmConfig from data_designer.slurm.services import ( + SlurmCollectionExecution, + SlurmRetryExecution, SlurmRunExecution, SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, ) +from data_designer.slurm.state import CollectionState class _RunService: def __init__(self) -> None: self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool, bool]] = [] + self.retry_calls: list[tuple[str, tuple[str, ...] | None, str, bool]] = [] + self.collection_calls: list[tuple[Path, Path, int | None]] = [] def execute( self, @@ -40,6 +45,43 @@ def execute( batch_script="#!/bin/bash\n", ) + 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: Path, + *, + destination: Path, + 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.resolve().as_posix(), + num_partitions=1 if num_partitions is None else num_partitions, + ) + def test_execute_emits_deterministic_json_and_forwards_actions( tmp_path: Path, @@ -65,6 +107,121 @@ def test_execute_emits_deterministic_json_and_forwards_actions( assert service.calls == [(authored_run_single, tmp_path, True, True)] +def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: + service = _RunService() + monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) + + 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) + + result = CliRunner().invoke(cli_module.create_cli(), ["retry", "42", "--force"]) + + assert result.exit_code == 0 + assert service.retry_calls == [("42", None, "if_possible", False)] + + +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) + 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)] + + @pytest.mark.parametrize( ("code", "exit_code"), [ @@ -226,9 +383,11 @@ def test_image_add_rejects_credential_bearing_oci_source() -> None: assert "secret" not in result.stderr -def test_cli_exposes_only_m2_run_commands() -> None: +def test_cli_exposes_m3c_run_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", "image", "profile")) - assert all(command not in result.stdout for command in ("retry", "merge", "benchmark")) + assert all( + command in result.stdout for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile") + ) + assert "benchmark" not in result.stdout diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 6be54bcdd..b5f54d241 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -135,6 +135,8 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non statement += f""" slurm_help_result = CliRunner().invoke(app, ["slurm", "--help"]) assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) +assert all(command in slurm_help_result.output for command in ("execute", "status", "cancel", "retry", "merge", "image")) +assert "benchmark" not in slurm_help_result.output assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference From cf3f51efce8897a581b9ef2e60f3b542725471b8 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:36:15 -0300 Subject: [PATCH 27/62] feat(slurm): add deterministic benchmark workflows (#932) * feat(slurm): add deterministic benchmark workflows * fix(slurm): correct benchmark concurrency and accounting * refactor(slurm): flatten benchmark CLI callbacks --- fern/versions/latest.yml | 4 + .../latest/pages/slurm/benchmarks.mdx | 94 +++ packages/data-designer-slurm/pyproject.toml | 2 +- .../data_designer/slurm/benchmark/__init__.py | 10 + .../data_designer/slurm/benchmark/analysis.py | 346 +++++++++++ .../data_designer/slurm/benchmark/compiler.py | 177 ++++++ .../slurm/benchmark/execution.py | 292 +++++++++ .../data_designer/slurm/benchmark/observer.py | 222 +++++++ .../data_designer/slurm/benchmark/records.py | 23 +- .../data_designer/slurm/benchmark/store.py | 208 +++++++ .../src/data_designer/slurm/cli.py | 4 + .../src/data_designer/slurm/cli_benchmark.py | 103 ++++ .../data_designer/slurm/config/__init__.py | 2 + .../src/data_designer/slurm/config/loading.py | 29 +- .../data_designer/slurm/runtime/controller.py | 10 +- .../data_designer/slurm/runtime/entrypoint.py | 10 +- .../data_designer/slurm/services/__init__.py | 6 + .../data_designer/slurm/services/artifacts.py | 23 +- .../data_designer/slurm/services/benchmark.py | 54 +- .../data_designer/slurm/services/wiring.py | 100 +++- .../src/data_designer/slurm/state/__init__.py | 2 + .../src/data_designer/slurm/state/errors.py | 4 + .../src/data_designer/slurm/state/observer.py | 3 +- .../data_designer/slurm/state/readiness.py | 8 + .../slurm/state/reconciliation.py | 1 + .../src/data_designer/slurm/state/storage.py | 15 + .../benchmark/test_benchmark_compiler.py | 77 +++ .../tests/benchmark/test_workflow.py | 558 ++++++++++++++++++ .../tests/config/test_loading_builder.py | 12 + .../tests/contracts/test_shared_records.py | 16 + .../tests/runtime/test_controller.py | 4 + .../tests/services/test_services.py | 45 +- .../tests/services/test_wiring.py | 171 +++++- .../tests/slurm_test_fakes/services.py | 29 +- .../tests/state/test_observer.py | 481 +++++++++++++-- .../data-designer-slurm/tests/test_cli.py | 55 +- scripts/test_slurm_package_install.py | 12 +- uv.lock | 2 +- 38 files changed, 3101 insertions(+), 113 deletions(-) create mode 100644 fern/versions/latest/pages/slurm/benchmarks.mdx create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py create mode 100644 packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py create mode 100644 packages/data-designer-slurm/tests/benchmark/test_workflow.py diff --git a/fern/versions/latest.yml b/fern/versions/latest.yml index 0d44b79b0..2f7cfed32 100644 --- a/fern/versions/latest.yml +++ b/fern/versions/latest.yml @@ -79,6 +79,10 @@ navigation: path: ./latest/pages/notebooks/5-generating-images.mdx - page: Image-to-Image Editing path: ./latest/pages/notebooks/6-editing-images-with-image-context.mdx + - section: Slurm + contents: + - page: Benchmarks + path: ./latest/pages/slurm/benchmarks.mdx - section: Recipes contents: - page: Recipe Cards diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx new file mode 100644 index 000000000..f1161462d --- /dev/null +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -0,0 +1,94 @@ +# Slurm benchmarks + +Use a benchmark to compare concurrency and deployment topology while keeping each case as an ordinary Data Designer Slurm run. No benchmark controller stays resident after submission. + +Install the Slurm extension: + +```bash +pip install "data-designer[slurm]" +``` + +## Define cases + +Create `benchmark.yaml` next to an existing `run.yaml`: + +```yaml +schema_version: 1 +name: generator-scaling +base_run: run.yaml +model_aliases: + - generator +concurrency_values: + - 32 + - 64 +deployment_cases: + - name: two-independent-replicas + deployments: + generator: + nodes: 2 + nodes_per_replica: 1 + - name: one-two-node-replica + deployments: + generator: + nodes: 2 + nodes_per_replica: 2 +record_policy: + type: adaptive + base_records: 1000 + max_records: 5000 + records_per_concurrency: 1.0 +analysis: + target_total_records: 1000000 + target_runtime: 4h +``` + +The compiler expands the authored order deterministically. For this example, the first cases are `two-independent-replicas-c32` and `one-two-node-replica-c32`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. + +## Run + +```bash +data-designer slurm benchmark run benchmark.yaml \ + --profile-file ~/.data-designer-slurm-profile.yml \ + --cluster primary +``` + +The command writes `benchmarks//config.json`, the resolved `base-run.json`, and `benchmark.json` before submitting any child. The manifest is an immutable ordered mapping from case IDs to normal run IDs. Each child keeps its authored config, state, attempts, scheduler evidence, and results under `runs/`. + +If submission stops partway through, rerun the same command. Children with scheduler submission evidence are not submitted again, and missing children are attempted in the original order. If a child has initialized inputs but no submission evidence, rerun with `--force` to resume it. Force never replaces benchmark or child metadata. Partial scheduler evidence across a child's shards is preserved as a conflict for operator investigation and cannot be overwritten with force. + +## Observe and analyze + +Inspect any child with the normal run command: + +```bash +data-designer slurm status +``` + +Write a point-in-time benchmark report from persisted child state: + +```bash +data-designer slurm benchmark analyze --refresh-state +``` + +`--refresh-state` performs one fresh-process scheduler reconciliation per child before analysis. Add `--fail-if-incomplete` when automation should return a conflict after the report is persisted if any case is not successful. + +Reports retain every manifest case in order. Outcomes include `pending`, `accounting_lag`, `succeeded`, `failed`, `incomplete`, `missing`, `stale`, and `scheduler_inconsistent`. Successful cases can be infeasible when boot time consumes the runtime budget. Boot and wall timing begin when execution starts inside each allocation, so queue wait does not affect topology comparisons. For array runs, `rows_per_second` sums the independently measured shard rates, while target jobs and GPU hours count individual task allocations. `generation_seconds` is the target runtime remaining after the slowest shard boot. Only successful feasible cases participate in Pareto, minimum-job, and minimum-GPU-hour recommendations. + +## Python API + +```python +from pathlib import Path + +from data_designer.slurm.config import load_benchmark_config +from data_designer.slurm.services import create_slurm_benchmark_service + +benchmark_file = Path("benchmark.yaml").resolve() +config = load_benchmark_config(benchmark_file) +service = create_slurm_benchmark_service( + profile_file="~/.data-designer-slurm-profile.yml", + cluster="primary", +) + +manifest = service.run(config, source_root=benchmark_file.parent) +report = service.analyze(manifest.benchmark_id, refresh_state=True) +``` diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index f691c26dd..326e052ab 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "data-designer=={{ version }}", "packaging>=25,<27", "pip>=25,<27", - "pydantic>=2.9.2,<3", + "pydantic>=2.12,<3", "pyyaml>=6.0.1,<7", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py index 40acb8f5b..52a7ed22f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/__init__.py @@ -5,6 +5,12 @@ from __future__ import annotations +from data_designer.slurm.benchmark.compiler import ( + BenchmarkCompiler, + CompiledBenchmark, + CompiledBenchmarkCase, + resolve_requested_records, +) from data_designer.slurm.benchmark.records import ( BenchmarkCaseResult, BenchmarkChildRun, @@ -16,6 +22,7 @@ ) __all__ = [ + "BenchmarkCompiler", "BenchmarkCaseResult", "BenchmarkChildRun", "BenchmarkManifest", @@ -23,4 +30,7 @@ "BenchmarkRecommendation", "BenchmarkRecommendationKind", "BenchmarkReport", + "CompiledBenchmark", + "CompiledBenchmarkCase", + "resolve_requested_records", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py new file mode 100644 index 000000000..28eaec966 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Point-in-time analysis of ordinary benchmark child runs.""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol + +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.records import ( + BenchmarkCaseResult, + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkRecommendation, + BenchmarkRecommendationKind, + BenchmarkReport, +) +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig +from data_designer.slurm.config.utils import convert_duration_to_seconds +from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan + + +@dataclass(frozen=True, slots=True) +class BenchmarkShardMeasurements: + """Complete timing and output facts for one successful shard.""" + + actual_records: int + boot_seconds: float + generation_seconds: float + wall_seconds: float + + +@dataclass(frozen=True, slots=True) +class BenchmarkMeasurements: + """Complete timing and output facts for one successful ordinary run.""" + + shards: tuple[BenchmarkShardMeasurements, ...] + + @property + def actual_records(self) -> int: + return sum(shard.actual_records for shard in self.shards) + + @property + def boot_seconds(self) -> float: + return max(shard.boot_seconds for shard in self.shards) + + @property + def generation_seconds(self) -> float: + rate = sum(shard.actual_records / shard.generation_seconds for shard in self.shards) + return self.actual_records / rate + + @property + def wall_seconds(self) -> float: + return max(shard.wall_seconds for shard in self.shards) + + +@dataclass(frozen=True, slots=True) +class BenchmarkRunObservation: + """Verified ordinary run inputs and one status snapshot.""" + + authored_config: DataDesignerSlurmConfig + resolved_plan: ResolvedSlurmRunPlan + outcome: BenchmarkOutcome + measurements: BenchmarkMeasurements | None = None + + +class BenchmarkRunObserver(Protocol): + """Observe one ordinary child run from durable state.""" + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + """Return one verified child snapshot or raise an observation failure.""" + + +class BenchmarkObservationFailure(Exception): + """Preserve a non-success child outcome without aborting the report.""" + + def __init__(self, outcome: BenchmarkOutcome) -> None: + if outcome not in { + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + }: + raise ValueError("invalid benchmark observation failure outcome") + self.outcome = outcome + super().__init__(outcome.value) + + +class BenchmarkManifestMismatchError(ValueError): + """Raised when persisted benchmark children do not match compilation.""" + + +class BenchmarkAnalyzer: + """Analyze one immutable benchmark manifest in manifest order.""" + + def __init__(self, observer: BenchmarkRunObserver, clock: Callable[[], datetime]) -> None: + self._observer = observer + self._clock = clock + + def analyze( + self, + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, + manifest_reference: ArtifactReference, + *, + refresh_state: bool, + ) -> BenchmarkReport: + requested_records = _validate_and_get_requested_records(config, base_run, manifest) + results = tuple( + self._analyze_case( + config, + child, + requested_records[index], + refresh_state=refresh_state, + ) + for index, child in enumerate(manifest.children) + ) + recommendations = _recommend(results) + created_at = self._clock() + analysis_digest = compute_canonical_json_sha256( + { + "benchmark_id": manifest.benchmark_id, + "created_at": created_at.isoformat(), + "cases": [case.model_dump(mode="json") for case in results], + "recommendations": [item.model_dump(mode="json") for item in recommendations], + } + ) + return BenchmarkReport( + schema_version=1, + benchmark_id=manifest.benchmark_id, + analysis_id=f"analysis-{analysis_digest[:32]}", + benchmark_manifest=manifest_reference, + created_at=created_at, + cases=results, + recommendations=recommendations, + ) + + def _analyze_case( + self, + config: DataDesignerSlurmBenchmarkConfig, + child: BenchmarkChildRun, + requested_records: int, + *, + refresh_state: bool, + ) -> BenchmarkCaseResult: + fallback = _fallback_facts(child, requested_records) + try: + observed = self._observer.observe(child.child_run_id, refresh_state=refresh_state) + except BenchmarkObservationFailure as error: + return _update_result(fallback, outcome=error.outcome) + if ( + observed.authored_config.compute_sha256() != child.child_authored_config.sha256 + or observed.resolved_plan.run_id != child.child_run_id + ): + return _update_result(fallback, outcome=BenchmarkOutcome.STALE) + + result = _resolved_facts(child, observed.authored_config, observed.resolved_plan) + if observed.outcome is not BenchmarkOutcome.SUCCEEDED: + return _update_result(result, outcome=observed.outcome) + if observed.measurements is None: + return _update_result(result, outcome=BenchmarkOutcome.INCOMPLETE) + measurements = observed.measurements + actual_records = measurements.actual_records + if ( + not measurements.shards + or len(measurements.shards) != len(observed.resolved_plan.shards) + or any( + shard.actual_records <= 0 + or shard.boot_seconds < 0 + or shard.generation_seconds <= 0 + or shard.wall_seconds <= 0 + for shard in measurements.shards + ) + ): + return _update_result( + result, + outcome=BenchmarkOutcome.INCOMPLETE, + actual_records=actual_records, + ) + boot_seconds = measurements.boot_seconds + generation_seconds = measurements.generation_seconds + wall_seconds = measurements.wall_seconds + requested_records = observed.authored_config.invocation.num_records + if actual_records != requested_records: + return _update_result( + result, + outcome=BenchmarkOutcome.INCOMPLETE, + actual_records=actual_records, + ) + gpus_per_job = result.gpus_per_job + if gpus_per_job is None: + raise ValueError("resolved benchmark case has no GPU count") + rows_per_second = actual_records / generation_seconds + budget_seconds = convert_duration_to_seconds(config.analysis.target_runtime) + effective_generation_seconds = max(0.0, budget_seconds - boot_seconds) + feasible = boot_seconds < budget_seconds + gpu_hours_per_job = gpus_per_job * budget_seconds / 3600 + if not feasible: + return _update_result( + result, + outcome=BenchmarkOutcome.SUCCEEDED, + actual_records=actual_records, + boot_seconds=boot_seconds, + generation_seconds=effective_generation_seconds, + wall_seconds=wall_seconds, + rows_per_second=rows_per_second, + gpu_hours_per_job=gpu_hours_per_job, + feasible=False, + ) + records_per_job = sum( + shard.actual_records / shard.generation_seconds * (budget_seconds - shard.boot_seconds) + for shard in measurements.shards + ) / len(measurements.shards) + target_jobs = math.ceil(config.analysis.target_total_records / records_per_job) + return _update_result( + result, + outcome=BenchmarkOutcome.SUCCEEDED, + actual_records=actual_records, + boot_seconds=boot_seconds, + generation_seconds=effective_generation_seconds, + wall_seconds=wall_seconds, + rows_per_second=rows_per_second, + gpu_hours_per_job=gpu_hours_per_job, + total_gpu_hours=gpu_hours_per_job * target_jobs, + target_jobs=target_jobs, + feasible=True, + ) + + +def _fallback_facts( + child: BenchmarkChildRun, + requested_records: int, +) -> BenchmarkCaseResult: + return BenchmarkCaseResult( + case_id=child.case_id, + child_run_id=child.child_run_id, + outcome=BenchmarkOutcome.MISSING, + requested_records=requested_records, + ) + + +def _resolved_facts( + child: BenchmarkChildRun, + authored: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, +) -> BenchmarkCaseResult: + nodes = sum(len(deployment.node_indices) for deployment in plan.deployments) + gpus = sum(len(deployment.node_indices) * deployment.gpus_per_node for deployment in plan.deployments) + topology_digest = compute_canonical_json_sha256( + [ + { + "deployment_id": deployment.deployment_id, + "model_alias": deployment.authored.model_alias, + "node_indices": deployment.node_indices, + "gpus_per_node": deployment.gpus_per_node, + "topology": deployment.topology.model_dump(mode="json"), + } + for deployment in plan.deployments + ] + ) + return BenchmarkCaseResult( + case_id=child.case_id, + child_run_id=child.child_run_id, + outcome=BenchmarkOutcome.INCOMPLETE, + topology_digest=topology_digest, + requested_records=authored.invocation.num_records, + gpus_per_job=gpus, + nodes_per_job=nodes, + ) + + +def _validate_and_get_requested_records( + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, +) -> tuple[int, ...]: + try: + compiled = BenchmarkCompiler.compile(config, base_run) + except ValueError as error: + raise BenchmarkManifestMismatchError("benchmark configuration cannot reproduce its children") from error + expected = tuple( + (case.case_id, case.child_run_id, case.child_run_config.compute_sha256()) for case in compiled.cases + ) + actual = tuple( + (child.case_id, child.child_run_id, child.child_authored_config.sha256) for child in manifest.children + ) + if manifest.benchmark_id != compiled.benchmark_id or actual != expected: + raise BenchmarkManifestMismatchError("benchmark manifest does not match deterministic child expansion") + return tuple(case.child_run_config.invocation.num_records for case in compiled.cases) + + +def _update_result(result: BenchmarkCaseResult, **updates: object) -> BenchmarkCaseResult: + return BenchmarkCaseResult.model_validate(result.model_dump(mode="python") | updates) + + +def _recommend(cases: tuple[BenchmarkCaseResult, ...]) -> tuple[BenchmarkRecommendation, ...]: + eligible = tuple( + case + for case in cases + if case.outcome is BenchmarkOutcome.SUCCEEDED + and case.feasible is True + and case.target_jobs is not None + and case.total_gpu_hours is not None + ) + if not eligible: + return () + pareto = tuple( + case + for case in eligible + if not any( + other.case_id != case.case_id + and other.target_jobs <= case.target_jobs + and other.total_gpu_hours <= case.total_gpu_hours + and (other.target_jobs < case.target_jobs or other.total_gpu_hours < case.total_gpu_hours) + for other in eligible + ) + ) + minimum_jobs = min(eligible, key=lambda case: (case.target_jobs, case.total_gpu_hours, case.case_id)) + minimum_gpu_hours = min(eligible, key=lambda case: (case.total_gpu_hours, case.target_jobs, case.case_id)) + return tuple( + BenchmarkRecommendation(kind=BenchmarkRecommendationKind.PARETO, case_id=case.case_id) for case in pareto + ) + ( + BenchmarkRecommendation(kind=BenchmarkRecommendationKind.MINIMUM_JOBS, case_id=minimum_jobs.case_id), + BenchmarkRecommendation( + kind=BenchmarkRecommendationKind.MINIMUM_GPU_HOURS, + case_id=minimum_gpu_hours.case_id, + ), + ) + + +__all__ = [ + "BenchmarkAnalyzer", + "BenchmarkMeasurements", + "BenchmarkShardMeasurements", + "BenchmarkManifestMismatchError", + "BenchmarkObservationFailure", + "BenchmarkRunObservation", + "BenchmarkRunObserver", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py new file mode 100644 index 000000000..568533af9 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic expansion of benchmark intent into ordinary run configs.""" + +from __future__ import annotations + +import math +import posixpath + +from pydantic import Field + +from data_designer.slurm.config import ( + AdaptiveRecordPolicy, + BenchmarkDeploymentCase, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, +) +from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 + + +class CompiledBenchmarkCase(ContractValue): + """One stable benchmark case and its ordinary authored run config.""" + + case_id: Identifier + child_run_id: Identifier + child_run_config: DataDesignerSlurmConfig + + +class CompiledBenchmark(ContractValue): + """Ordered immutable result of benchmark expansion.""" + + benchmark_id: Identifier + cases: tuple[CompiledBenchmarkCase, ...] = Field(min_length=1) + + +class BenchmarkCompiler: + """Compile a benchmark declaration without filesystem or scheduler access.""" + + @staticmethod + def compile( + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + ) -> CompiledBenchmark: + if not isinstance(config, DataDesignerSlurmBenchmarkConfig): + raise TypeError("config must be a DataDesignerSlurmBenchmarkConfig") + if not isinstance(base_run, DataDesignerSlurmConfig): + raise TypeError("base_run must be a DataDesignerSlurmConfig") + + aliases = tuple(deployment.model_alias for deployment in base_run.deployments) + selected_aliases = aliases if config.model_aliases == "all" else tuple(config.model_aliases) + unknown_aliases = set(selected_aliases).difference(aliases) + if unknown_aliases: + raise ValueError(f"benchmark references unknown model aliases: {', '.join(sorted(unknown_aliases))}") + for deployment_case in config.deployment_cases: + unknown_deployments = set(deployment_case.deployments).difference(aliases) + if unknown_deployments: + raise ValueError( + f"benchmark deployment case references unknown aliases: {', '.join(sorted(unknown_deployments))}" + ) + + identity_digest = compute_canonical_json_sha256( + { + "benchmark": config.model_dump(mode="json"), + "base_run": base_run.model_dump(mode="json"), + } + ) + benchmark_id = f"benchmark-{identity_digest[:32]}" + expanded: list[tuple[str, DataDesignerSlurmConfig]] = [] + for concurrency in config.concurrency_values: + for deployment_case in config.deployment_cases: + case_id = derive_benchmark_case_id(deployment_case.name, concurrency) + expanded.append( + ( + case_id, + _compile_child( + base_run, + benchmark_id=benchmark_id, + case_id=case_id, + concurrency=concurrency, + selected_aliases=selected_aliases, + deployment_case=deployment_case, + requested_records=resolve_requested_records(config, concurrency), + ), + ) + ) + + cases = tuple( + CompiledBenchmarkCase( + case_id=case_id, + child_run_id=f"run-{identity_digest[:16]}-{index:04d}-{child.compute_sha256()[:12]}", + child_run_config=child, + ) + for index, (case_id, child) in enumerate(expanded) + ) + return CompiledBenchmark(benchmark_id=benchmark_id, cases=cases) + + +def _compile_child( + base_run: DataDesignerSlurmConfig, + *, + benchmark_id: str, + case_id: str, + concurrency: int, + selected_aliases: tuple[str, ...], + deployment_case: BenchmarkDeploymentCase, + requested_records: int, +) -> DataDesignerSlurmConfig: + model_concurrency = dict(base_run.invocation.model_concurrency) + model_concurrency.update(dict.fromkeys(selected_aliases, concurrency)) + invocation = base_run.invocation.model_copy( + update={ + "num_records": requested_records, + "model_concurrency": model_concurrency, + } + ) + deployments = [] + for deployment in base_run.deployments: + override = deployment_case.deployments.get(deployment.model_alias) + if override is None: + deployments.append(deployment) + continue + deployments.append( + deployment.model_copy( + update={ + "resources": deployment.resources.model_copy(update={"nodes": override.nodes}), + "topology": deployment.topology.model_copy( + update={"nodes_per_replica": override.nodes_per_replica} + ), + } + ) + ) + + output_root = base_run.output.root + output = base_run.output + if output_root is not None: + output = output.model_copy(update={"root": posixpath.join(output_root, benchmark_id, case_id)}) + return DataDesignerSlurmConfig.model_validate( + base_run.model_copy( + update={ + "name": _bounded_identifier(f"{base_run.name}-{case_id}"), + "invocation": invocation, + "deployments": deployments, + "output": output, + } + ).model_dump(mode="python") + ) + + +def resolve_requested_records(config: DataDesignerSlurmBenchmarkConfig, concurrency: int) -> int: + """Resolve the deterministic record count for one concurrency value.""" + policy = config.record_policy + if not isinstance(policy, AdaptiveRecordPolicy): + return policy.records + scaled = math.ceil(concurrency * policy.records_per_concurrency) + return min(policy.max_records, max(policy.base_records, scaled)) + + +def derive_benchmark_case_id(case_name: str, concurrency: int) -> str: + """Return the stable bounded identity for one cross-product case.""" + return _bounded_identifier(f"{case_name}-c{concurrency}") + + +def _bounded_identifier(value: str) -> str: + if len(value) <= 128: + return value + digest = compute_canonical_json_sha256(value)[:12] + return f"{value[:115]}-{digest}" + + +__all__ = [ + "BenchmarkCompiler", + "CompiledBenchmark", + "CompiledBenchmarkCase", + "derive_benchmark_case_id", + "resolve_requested_records", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py new file mode 100644 index 000000000..6841a453c --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark orchestration through ordinary public Slurm run services.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkAnalyzer, + BenchmarkManifestMismatchError, + BenchmarkRunObserver, +) +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.records import ( + BenchmarkChildRun, + BenchmarkManifest, + BenchmarkOutcome, + BenchmarkReport, +) +from data_designer.slurm.benchmark.store import ( + BenchmarkConflictError, + BenchmarkNotFoundError, + BenchmarkStore, + BenchmarkStoreError, +) +from data_designer.slurm.config import ( + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + SlurmConfigLoadError, + load_run_config, +) +from data_designer.slurm.contracts import ArtifactReference, Identifier +from data_designer.slurm.services.errors import ( + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) +from data_designer.slurm.services.run import SlurmRunService +from data_designer.slurm.state import SlurmStateError, SlurmStateWriter, StateNotFoundError +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +ChildRunServiceFactory = Callable[[Identifier], SlurmRunService] +ChildConfigLoader = Callable[[Identifier], DataDesignerSlurmConfig | None] +ChildSubmissionLoader = Callable[[Identifier], bool] + + +class SystemBenchmarkBackend: + """Persist benchmark intent and submit each case as an ordinary run.""" + + def __init__( + self, + workspace_root: str | Path, + run_service_factory: ChildRunServiceFactory, + observer: BenchmarkRunObserver, + clock: Callable[[], datetime], + child_config_loader: ChildConfigLoader | None = None, + child_submission_loader: ChildSubmissionLoader | None = None, + ) -> None: + self._workspace_root = Path(workspace_root) + self._run_service_factory = run_service_factory + self._observer = observer + self._clock = clock + self._child_config_loader = child_config_loader or self._load_child_config + self._child_submission_loader = child_submission_loader or self._load_child_submission + + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: + try: + base_run, child_source_root = _resolve_base_run(config, source_root) + compiled = BenchmarkCompiler.compile(config, base_run) + except (SlurmConfigLoadError, ValueError): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark configuration cannot be compiled", + ) from None + store = BenchmarkStore(self._workspace_root, compiled.benchmark_id) + children = tuple( + BenchmarkChildRun( + case_id=case.case_id, + child_run_id=case.child_run_id, + child_authored_config=ArtifactReference( + path=(self._workspace_root / "runs" / case.child_run_id / "authored-config.json").as_posix(), + sha256=case.child_run_config.compute_sha256(), + ), + ) + for case in compiled.cases + ) + manifest = store.build_manifest(config, children) + try: + store.publish(config, base_run, manifest) + except BenchmarkConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark metadata conflicts with persisted state", + ) from None + except BenchmarkStoreError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark metadata cannot be persisted", + ) from None + + failures: list[SlurmServiceErrorCode] = [] + for case in compiled.cases: + try: + with StateStorage(self._workspace_root, case.child_run_id).acquire_submission_lock(): + if self._child_exists(case.child_run_id, case.child_run_config, force=force): + continue + execution = self._run_service_factory(case.child_run_id).execute( + case.child_run_config, + source_root=child_source_root, + dry_run=False, + force=False, + ) + if execution.run_id != case.child_run_id or execution.state != "submitted": + raise SlurmServiceError( + SlurmServiceErrorCode.INTERNAL, + SlurmServiceOperation.RUN_BENCHMARK, + "ordinary run service returned an invalid benchmark child", + ) + except SlurmServiceError as error: + failures.append(error.code) + except Exception: + failures.append(SlurmServiceErrorCode.INTERNAL) + if failures: + code = _partial_failure_code(failures) + raise SlurmServiceError( + code, + SlurmServiceOperation.RUN_BENCHMARK, + f"{len(failures)} of {len(compiled.cases)} benchmark child runs could not be submitted", + ) + return manifest + + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: + store = BenchmarkStore(self._workspace_root, benchmark_id) + try: + config, base_run, manifest = store.load() + report = BenchmarkAnalyzer(self._observer, self._clock).analyze( + config, + base_run, + manifest, + store.manifest_reference(manifest), + refresh_state=refresh_state, + ) + store.publish_report(report) + except BenchmarkNotFoundError: + raise SlurmServiceError( + SlurmServiceErrorCode.NOT_FOUND, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata was not found", + ) from None + except BenchmarkConflictError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata conflicts with persisted state", + ) from None + except BenchmarkStoreError: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark metadata cannot be read or persisted", + ) from None + except BenchmarkManifestMismatchError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + "benchmark manifest does not match its configuration", + ) from None + if fail_if_incomplete: + incomplete = sum(case.outcome is not BenchmarkOutcome.SUCCEEDED for case in report.cases) + if incomplete: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.ANALYZE_BENCHMARK, + f"benchmark analysis contains {incomplete} incomplete cases", + ) + return report + + def _child_exists( + self, + run_id: Identifier, + expected: DataDesignerSlurmConfig, + *, + force: bool, + ) -> bool: + persisted = self._child_config_loader(run_id) + if persisted is None: + return False + if persisted != expected: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state contains different inputs", + ) + if self._child_submission_loader(run_id): + return True + if force: + return False + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child inputs exist without submission evidence; rerun with force", + ) + + def _load_child_config(self, run_id: Identifier) -> DataDesignerSlurmConfig | None: + try: + return SlurmStateWriter(self._workspace_root, run_id).load_authored_config() + except StateNotFoundError: + return None + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state is invalid", + ) from None + + def _load_child_submission(self, run_id: Identifier) -> bool: + try: + storage = StateStorage(self._workspace_root, run_id) + reader = StateReader(storage, run_id) + run, plan, shards = reader.load_context() + attempts = reader.load_validated_attempts(run, plan, shards) + submitted = tuple( + bool(attempts[shard.shard_id] and attempts[shard.shard_id][0].scheduler is not None) for shard in shards + ) + if all(submitted): + return True + if any(submitted): + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child contains partial submission evidence", + ) + return False + except SlurmServiceError: + raise + except SlurmStateError: + raise SlurmServiceError( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceOperation.RUN_BENCHMARK, + "benchmark child run state is invalid", + ) from None + + +def _resolve_base_run( + config: DataDesignerSlurmBenchmarkConfig, + source_root: Path, +) -> tuple[DataDesignerSlurmConfig, Path]: + if config.base_run.inline is not None: + return config.base_run.inline, source_root + assert config.base_run.source is not None + source_path = source_root / config.base_run.source + return load_run_config(source_path), source_path.resolve().parent + + +def _partial_failure_code(codes: list[SlurmServiceErrorCode]) -> SlurmServiceErrorCode: + for code in ( + SlurmServiceErrorCode.CONFLICT, + SlurmServiceErrorCode.INVALID_REQUEST, + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceErrorCode.INTERNAL, + ): + if code in codes: + return code + return SlurmServiceErrorCode.INTERNAL + + +__all__ = [ + "ChildConfigLoader", + "ChildRunServiceFactory", + "ChildSubmissionLoader", + "SystemBenchmarkBackend", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py new file mode 100644 index 000000000..3ae540698 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fresh-process observation of ordinary benchmark child state.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import datetime +from pathlib import Path + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkMeasurements, + BenchmarkObservationFailure, + BenchmarkRunObservation, + BenchmarkShardMeasurements, +) +from data_designer.slurm.benchmark.records import BenchmarkOutcome +from data_designer.slurm.client import ClientOutcome, ClientResult +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.state import ( + AttemptLifecycleState, + AttemptManifest, + AttemptReadiness, + CandidateOutputManifest, + EffectiveRunState, + ProbeOutcome, + ReadinessState, + RunManifest, + RunStatus, + SchedulerObservationClient, + SchedulerState, + SchedulerStateConflictError, + ShardManifest, + SlurmStateError, + SlurmStateReconciler, + StateConflictError, + StateCorruptionError, + StateNotFoundError, +) +from data_designer.slurm.state.finalization import WinnerFinalizer +from data_designer.slurm.state.reader import StateReader +from data_designer.slurm.state.storage import StateStorage + +_MeasurementRecord = tuple[ + AttemptManifest, + AttemptReadiness | None, + tuple[ClientResult, CandidateOutputManifest] | None, +] + + +class PersistedBenchmarkRunObserver: + """Reconstruct benchmark child facts from the ordinary run state tree.""" + + def __init__( + self, + workspace_root: str | Path, + scheduler: SchedulerObservationClient, + clock: Callable[[], datetime], + ) -> None: + self._workspace_root = Path(workspace_root) + self._scheduler = scheduler + self._clock = clock + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + try: + storage = StateStorage(self._workspace_root, run_id) + reader = StateReader(storage, run_id) + run, plan, shards = reader.load_context() + authored = reader.load_authored_config(run) + if refresh_state: + status = SlurmStateReconciler( + self._workspace_root, + run_id, + self._scheduler, + ).refresh(observed_at=self._clock()) + outcome = _status_outcome(status.effective_state) + measurements = _measure_status(status) if outcome is BenchmarkOutcome.SUCCEEDED else None + else: + outcome, measurements = _load_persisted_outcome(storage, reader, run, plan, shards) + return BenchmarkRunObservation( + authored_config=authored, + resolved_plan=plan, + outcome=outcome, + measurements=measurements, + ) + except StateNotFoundError: + raise BenchmarkObservationFailure(BenchmarkOutcome.MISSING) from None + except SchedulerStateConflictError: + raise BenchmarkObservationFailure(BenchmarkOutcome.SCHEDULER_INCONSISTENT) from None + except (StateConflictError, StateCorruptionError): + raise BenchmarkObservationFailure(BenchmarkOutcome.STALE) from None + except (OSError, SlurmStateError): + raise BenchmarkObservationFailure(BenchmarkOutcome.STALE) from None + + +def _load_persisted_outcome( + storage: StateStorage, + reader: StateReader, + run: RunManifest, + plan: ResolvedSlurmRunPlan, + shards: tuple[ShardManifest, ...], +) -> tuple[BenchmarkOutcome, BenchmarkMeasurements | None]: + attempts_by_shard = reader.load_validated_attempts(run, plan, shards) + finalizer = WinnerFinalizer(storage, reader) + winning_records: list[_MeasurementRecord] = [] + latest_states: list[AttemptLifecycleState] = [] + scheduler_states: list[SchedulerState] = [] + for shard in shards: + attempts = attempts_by_shard[shard.shard_id] + if not attempts: + return BenchmarkOutcome.INCOMPLETE, None + winner = finalizer.load_optional_winner(run, plan, shard, attempts) + if winner is None: + latest = attempts[-1] + latest_states.append(latest.state) + observation = reader.load_optional_scheduler_observation(latest) + if observation is not None: + scheduler_states.append(observation.state) + continue + attempt = next(item for item in attempts if item.attempt_id == winner.attempt_id) + observation = reader.load_optional_scheduler_observation(attempt) + if observation is not None and observation.state in { + SchedulerState.FAILED, + SchedulerState.CANCELLED, + SchedulerState.TIMED_OUT, + SchedulerState.NODE_FAILED, + SchedulerState.OUT_OF_MEMORY, + }: + raise SchedulerStateConflictError("persisted winner conflicts with scheduler evidence") + result = reader.load_optional_attempt_result(plan, shard, attempt) + readiness = reader.load_optional_readiness(plan, attempt) + winning_records.append((attempt, readiness, result)) + + if len(winning_records) == len(shards): + return BenchmarkOutcome.SUCCEEDED, _measure_records(tuple(winning_records)) + if SchedulerState.ACCOUNTING_LAG in scheduler_states: + return BenchmarkOutcome.ACCOUNTING_LAG, None + if any( + state in {SchedulerState.UNKNOWN, SchedulerState.PREEMPTED, SchedulerState.REQUEUED} + for state in scheduler_states + ): + return BenchmarkOutcome.STALE, None + if any( + state + in { + AttemptLifecycleState.CREATED, + AttemptLifecycleState.SUBMITTED, + AttemptLifecycleState.PENDING, + AttemptLifecycleState.RUNNING, + } + for state in latest_states + ): + return BenchmarkOutcome.PENDING, None + if all(state is AttemptLifecycleState.FAILED for state in latest_states): + return BenchmarkOutcome.FAILED, None + return BenchmarkOutcome.INCOMPLETE, None + + +def _status_outcome(state: EffectiveRunState) -> BenchmarkOutcome: + if state in {EffectiveRunState.PENDING, EffectiveRunState.RUNNING}: + return BenchmarkOutcome.PENDING + if state is EffectiveRunState.ACCOUNTING_LAG: + return BenchmarkOutcome.ACCOUNTING_LAG + if state is EffectiveRunState.SUCCEEDED: + return BenchmarkOutcome.SUCCEEDED + if state is EffectiveRunState.FAILED: + return BenchmarkOutcome.FAILED + return BenchmarkOutcome.STALE + + +def _measure_status(status: RunStatus) -> BenchmarkMeasurements | None: + records: list[_MeasurementRecord] = [] + for shard in status.shards: + winner = next((attempt for attempt in shard.attempts if attempt.is_winner), None) + if winner is None: + return None + result = ( + None + if winner.client_result is None or winner.candidate_output is None + else (winner.client_result, winner.candidate_output) + ) + records.append((winner.attempt, winner.readiness, result)) + return _measure_records(tuple(records)) + + +def _measure_records(records: tuple[_MeasurementRecord, ...]) -> BenchmarkMeasurements | None: + measurements: list[BenchmarkShardMeasurements] = [] + for attempt, readiness, result in records: + if result is None: + return None + client_result, candidate = result + if ( + client_result.outcome is not ClientOutcome.COMPLETE + or candidate.actual_records != client_result.actual_records + ): + return None + if readiness is None or readiness.state not in {ReadinessState.READY, ReadinessState.STOPPED}: + return None + if readiness.started_at is None: + return None + probes = tuple(deployment.last_probe for deployment in readiness.deployments) + if any(probe is None or probe.outcome is not ProbeOutcome.SUCCESS for probe in probes): + return None + started_at = readiness.started_at + ready_at = max(probe.observed_at for probe in probes if probe is not None) + completed_at = client_result.completed_at + stopped_at = attempt.updated_at + if ready_at < started_at or completed_at <= ready_at or stopped_at < completed_at or stopped_at <= started_at: + return None + measurements.append( + BenchmarkShardMeasurements( + actual_records=client_result.actual_records or 0, + boot_seconds=(ready_at - started_at).total_seconds(), + generation_seconds=(completed_at - ready_at).total_seconds(), + wall_seconds=(stopped_at - started_at).total_seconds(), + ) + ) + return BenchmarkMeasurements(shards=tuple(measurements)) + + +__all__ = ["PersistedBenchmarkRunObserver"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py index 11c797527..40d98b232 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -57,13 +57,16 @@ class BenchmarkOutcome(str, Enum): SUCCEEDED = "succeeded" FAILED = "failed" INCOMPLETE = "incomplete" + MISSING = "missing" + STALE = "stale" + SCHEDULER_INCONSISTENT = "scheduler_inconsistent" class BenchmarkCaseResult(ContractValue): case_id: Identifier child_run_id: Identifier outcome: BenchmarkOutcome - topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] + topology_digest: Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] | None = None requested_records: PositiveInt actual_records: NonNegativeInt | None = None boot_seconds: NonNegativeFloat | None = None @@ -72,8 +75,8 @@ class BenchmarkCaseResult(ContractValue): rows_per_second: NonNegativeFloat | None = None request_count: NonNegativeInt | None = None token_count: NonNegativeInt | None = None - gpus_per_job: PositiveInt - nodes_per_job: PositiveInt + gpus_per_job: PositiveInt | None = None + nodes_per_job: PositiveInt | None = None gpu_hours_per_job: NonNegativeFloat | None = None total_gpu_hours: NonNegativeFloat | None = None target_jobs: PositiveInt | None = None @@ -89,9 +92,9 @@ def validate_metrics(self) -> BenchmarkCaseResult: self.generation_seconds, self.wall_seconds, self.rows_per_second, + self.gpus_per_job, + self.nodes_per_job, self.gpu_hours_per_job, - self.total_gpu_hours, - self.target_jobs, self.feasible, ) if self.outcome is BenchmarkOutcome.SUCCEEDED and any(value is None for value in required): @@ -99,8 +102,14 @@ def validate_metrics(self) -> BenchmarkCaseResult: if self.outcome is BenchmarkOutcome.SUCCEEDED: if self.actual_records != self.requested_records: raise ValueError("successful benchmark cases require the requested record count") - if self.generation_seconds == 0 or self.wall_seconds == 0 or self.rows_per_second == 0: - raise ValueError("successful benchmark generation, wall time, and throughput must be positive") + if self.wall_seconds == 0 or self.rows_per_second == 0: + raise ValueError("successful benchmark wall time and throughput must be positive") + if self.feasible and self.generation_seconds == 0: + raise ValueError("feasible benchmark generation time must be positive") + if self.feasible and (self.total_gpu_hours is None or self.target_jobs is None): + raise ValueError("feasible benchmark cases require target jobs and total GPU hours") + if not self.feasible and (self.total_gpu_hours is not None or self.target_jobs is not None): + raise ValueError("infeasible benchmark cases cannot have target jobs or total GPU hours") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py new file mode 100644 index 000000000..fa7027d2a --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/store.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Immutable benchmark metadata beneath the selected Slurm workspace.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TypeVar + +from pydantic import TypeAdapter, ValidationError + +from data_designer.slurm.benchmark.records import BenchmarkChildRun, BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig +from data_designer.slurm.contracts import ArtifactReference, ContractRecord, Identifier, validate_absolute_path +from data_designer.slurm.images.filesystem import ensure_private_directory +from data_designer.slurm.state.filesystem import open_verified_directory, publish_immutable_text, read_regular_text + +_CONFIG_FILENAME = "config.json" +_BASE_RUN_FILENAME = "base-run.json" +_MANIFEST_FILENAME = "benchmark.json" +_REPORTS_DIRECTORY = "reports" +_MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 +_IDENTIFIER_ADAPTER = TypeAdapter(Identifier) +_RecordT = TypeVar("_RecordT", bound=ContractRecord) + + +class BenchmarkStoreError(Exception): + """Base error for benchmark metadata persistence.""" + + +class BenchmarkNotFoundError(BenchmarkStoreError): + """A benchmark metadata record is missing.""" + + +class BenchmarkConflictError(BenchmarkStoreError): + """Benchmark metadata is unsafe, corrupt, or conflicts with immutable state.""" + + +class BenchmarkStore: + """Persist one benchmark config, manifest, and point-in-time reports.""" + + def __init__(self, workspace_root: str | Path, benchmark_id: Identifier) -> None: + try: + root = Path(validate_absolute_path(Path(workspace_root).as_posix())) + normalized_id = _IDENTIFIER_ADAPTER.validate_python(benchmark_id, strict=True) + except (ValidationError, ValueError) as error: + raise BenchmarkStoreError("invalid benchmark location") from error + self.workspace_root = root + self.benchmark_id = normalized_id + self.benchmarks_root = root / "benchmarks" + self.benchmark_root = self.benchmarks_root / normalized_id + self.config_path = self.benchmark_root / _CONFIG_FILENAME + self.base_run_path = self.benchmark_root / _BASE_RUN_FILENAME + self.manifest_path = self.benchmark_root / _MANIFEST_FILENAME + self.reports_root = self.benchmark_root / _REPORTS_DIRECTORY + + def build_manifest( + self, + config: DataDesignerSlurmBenchmarkConfig, + children: tuple[BenchmarkChildRun, ...], + ) -> BenchmarkManifest: + """Build a manifest bound to this store's immutable locations.""" + return BenchmarkManifest( + schema_version=1, + benchmark_id=self.benchmark_id, + benchmark_config=ArtifactReference(path=self.config_path.as_posix(), sha256=config.compute_sha256()), + children=children, + ) + + def publish( + self, + config: DataDesignerSlurmBenchmarkConfig, + base_run: DataDesignerSlurmConfig, + manifest: BenchmarkManifest, + ) -> BenchmarkManifest: + """Convergently publish resolved inputs then the manifest commit record.""" + self._validate_manifest(config, manifest) + try: + ensure_private_directory(self.benchmark_root, parents=True) + with open_verified_directory(self.benchmark_root, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + _BASE_RUN_FILENAME, + base_run.serialize_json(), + self.base_run_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + publish_immutable_text( + descriptor, + _CONFIG_FILENAME, + config.serialize_json(), + self.config_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + publish_immutable_text( + descriptor, + _MANIFEST_FILENAME, + manifest.serialize_json(), + self.manifest_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + return manifest + except FileExistsError as error: + raise BenchmarkConflictError("benchmark already contains different immutable metadata") from error + except OSError as error: + raise BenchmarkStoreError("benchmark metadata cannot be persisted") from error + + def load(self) -> tuple[DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, BenchmarkManifest]: + """Load and verify the committed benchmark metadata.""" + try: + with open_verified_directory(self.benchmark_root, require_private=True) as descriptor: + config = self._read_record( + descriptor, + _CONFIG_FILENAME, + self.config_path, + DataDesignerSlurmBenchmarkConfig, + ) + base_run = self._read_record( + descriptor, + _BASE_RUN_FILENAME, + self.base_run_path, + DataDesignerSlurmConfig, + ) + manifest = self._read_record( + descriptor, + _MANIFEST_FILENAME, + self.manifest_path, + BenchmarkManifest, + ) + self._validate_manifest(config, manifest) + return config, base_run, manifest + except FileNotFoundError as error: + raise BenchmarkNotFoundError(f"benchmark {self.benchmark_id!r} is not initialized") from error + except BenchmarkConflictError: + raise + except (OSError, ValidationError, ValueError) as error: + raise BenchmarkConflictError(f"benchmark {self.benchmark_id!r} contains invalid metadata") from error + + def publish_report(self, report: BenchmarkReport) -> Path: + """Persist one immutable point-in-time report.""" + if report.benchmark_id != self.benchmark_id: + raise BenchmarkConflictError("benchmark report identity does not match its store") + _, _, manifest = self.load() + expected_reference = ArtifactReference(path=self.manifest_path.as_posix(), sha256=manifest.compute_sha256()) + if report.benchmark_manifest != expected_reference: + raise BenchmarkConflictError("benchmark report does not bind the persisted manifest") + report_root = self.reports_root / report.analysis_id + report_path = report_root / "report.json" + try: + ensure_private_directory(self.reports_root, parents=False) + ensure_private_directory(report_root, parents=False) + with open_verified_directory(report_root, require_private=True) as descriptor: + publish_immutable_text( + descriptor, + "report.json", + report.serialize_json(), + report_path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + return report_path + except FileExistsError as error: + raise BenchmarkConflictError("benchmark report identity already contains different bytes") from error + except OSError as error: + raise BenchmarkStoreError("benchmark report cannot be persisted") from error + + def manifest_reference(self, manifest: BenchmarkManifest) -> ArtifactReference: + """Return the immutable reference for a verified manifest.""" + return ArtifactReference(path=self.manifest_path.as_posix(), sha256=manifest.compute_sha256()) + + def _validate_manifest( + self, + config: DataDesignerSlurmBenchmarkConfig, + manifest: BenchmarkManifest, + ) -> None: + expected = ArtifactReference(path=self.config_path.as_posix(), sha256=config.compute_sha256()) + if manifest.benchmark_id != self.benchmark_id or manifest.benchmark_config != expected: + raise BenchmarkConflictError("benchmark manifest does not match its immutable config") + for child in manifest.children: + expected_path = self.workspace_root / "runs" / child.child_run_id / "authored-config.json" + if child.child_authored_config.path != expected_path.as_posix(): + raise BenchmarkConflictError("benchmark child reference is outside its ordinary run state") + + @staticmethod + def _read_record( + directory_descriptor: int, + name: str, + path: Path, + record_type: type[_RecordT], + ) -> _RecordT: + content = read_regular_text( + directory_descriptor, + name, + path, + maximum_size=_MAXIMUM_RECORD_SIZE, + ) + record = record_type.model_validate_json(content) + if record.serialize_json() != content: + raise BenchmarkConflictError("benchmark metadata is not canonical") + return record + + +__all__ = [ + "BenchmarkConflictError", + "BenchmarkNotFoundError", + "BenchmarkStore", + "BenchmarkStoreError", +] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 730ce513c..18d8a4c2f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -14,6 +14,7 @@ import typer from pydantic import BaseModel, ValidationError +from data_designer.slurm.cli_benchmark import create_benchmark_app from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json from data_designer.slurm.images.records import validate_oci_source_for_lifecycle @@ -364,6 +365,9 @@ def _bounded_error_message(message: str, *, fallback: str) -> str: return sanitized if len(sanitized) <= 512 else f"{sanitized[:509]}..." +app.add_typer(create_benchmark_app(_invoke, _emit_result), name="benchmark") + + def create_cli() -> click.Command: """Create the Slurm CLI group.""" return typer.main.get_command(app) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py new file mode 100644 index 000000000..83c3f56b4 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Thin benchmark commands for the Slurm CLI.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import partial +from pathlib import Path + +import typer +from pydantic import BaseModel + +from data_designer.slurm.config import load_benchmark_config +from data_designer.slurm.services import create_slurm_benchmark_service +from data_designer.slurm.services.errors import SlurmServiceOperation + +_Invoke = Callable[[SlurmServiceOperation, Callable[[], BaseModel]], BaseModel] +_Emit = Callable[[BaseModel], None] + + +def _run_benchmark( + benchmark_file: Path, + *, + profile_file: Path | None, + cluster: str | None, + force: bool, +) -> BaseModel: + config = load_benchmark_config(benchmark_file) + service = create_slurm_benchmark_service(profile_file=profile_file, cluster=cluster) + return service.run(config, source_root=benchmark_file.resolve().parent, force=force) + + +def _analyze_benchmark( + benchmark: str, + *, + profile_file: Path | None, + cluster: str | None, + refresh: bool, + fail_if_incomplete: bool, +) -> BaseModel: + benchmark_id = Path(benchmark.rstrip("/")).name + service = create_slurm_benchmark_service(profile_file=profile_file, cluster=cluster) + return service.analyze( + benchmark_id, + refresh_state=refresh, + fail_if_incomplete=fail_if_incomplete, + ) + + +def create_benchmark_app(invoke: _Invoke, emit: _Emit) -> typer.Typer: + """Create benchmark commands using the root CLI error and output policy.""" + benchmark_app = typer.Typer(help="Run and analyze Slurm benchmarks", no_args_is_help=True) + + @benchmark_app.command("run") + def run_command( + benchmark_file: Path = typer.Argument(..., exists=True, dir_okay=False, readable=True), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), + force: bool = typer.Option(False, "--force"), + ) -> None: + """Expand and submit ordinary child runs.""" + emit( + invoke( + SlurmServiceOperation.RUN_BENCHMARK, + partial( + _run_benchmark, + benchmark_file, + profile_file=profile_file, + cluster=cluster, + force=force, + ), + ) + ) + + @benchmark_app.command("analyze") + def analyze_command( + benchmark: str = typer.Argument(..., help="Managed benchmark ID or benchmark directory"), + profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), + cluster: str | None = typer.Option(None, "--cluster"), + refresh: bool = typer.Option(False, "--refresh-state", "--refresh"), + fail_if_incomplete: bool = typer.Option(False, "--fail-if-incomplete"), + ) -> None: + """Write one point-in-time report from ordinary child state.""" + emit( + invoke( + SlurmServiceOperation.ANALYZE_BENCHMARK, + partial( + _analyze_benchmark, + benchmark, + profile_file=profile_file, + cluster=cluster, + refresh=refresh, + fail_if_incomplete=fail_if_incomplete, + ), + ) + ) + + return benchmark_app + + +__all__ = ["create_benchmark_app"] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 00b7f9bf3..3fca0ea0f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -33,6 +33,7 @@ from data_designer.slurm.config.loading import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + load_benchmark_config, load_builder_payload, load_profile_catalog, load_run_config, @@ -120,6 +121,7 @@ "VllmServerConfig", "injected_profile", "load_builder_payload", + "load_benchmark_config", "load_profile_catalog", "load_run_config", "resolve_profile", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index 5c43b9f64..b4cad4be5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -17,6 +17,7 @@ from yaml.nodes import MappingNode from data_designer.slurm._errors import format_parse_error, format_validation_error +from data_designer.slurm.config.benchmark import DataDesignerSlurmBenchmarkConfig from data_designer.slurm.config.errors import SlurmConfigLoadError from data_designer.slurm.config.profiles import ( SelectedSlurmProfile, @@ -30,7 +31,12 @@ PROFILE_FILE_ENVIRONMENT = "DATA_DESIGNER_SLURM_PROFILE_FILE" DEFAULT_PROFILE_FILE_NAME = ".data-designer-slurm-profile.yml" -_ConfigT = TypeVar("_ConfigT", DataDesignerSlurmConfig, SlurmProfileCatalog) +_ConfigT = TypeVar( + "_ConfigT", + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, + SlurmProfileCatalog, +) _HostnameResolver = Callable[[], tuple[str, ...]] @@ -69,6 +75,11 @@ def load_run_config(path: str | Path) -> DataDesignerSlurmConfig: return _load_config(path, DataDesignerSlurmConfig) +def load_benchmark_config(path: str | Path) -> DataDesignerSlurmBenchmarkConfig: + """Load one strict local YAML or JSON benchmark declaration.""" + return _load_config(path, DataDesignerSlurmBenchmarkConfig) + + def load_profile_catalog(path: str | Path) -> SlurmProfileCatalog: """Load one strict local YAML or JSON cluster-profile catalog.""" return _load_config(path, SlurmProfileCatalog) @@ -161,6 +172,8 @@ def _load_config(path: str | Path, config_type: type[_ConfigT]) -> _ConfigT: payload = _parse_mapping(contents, suffix=resolved_path.suffix) if config_type is DataDesignerSlurmConfig: _reject_run_environment_interpolation(payload) + elif config_type is DataDesignerSlurmBenchmarkConfig: + _reject_benchmark_environment_interpolation(payload) else: _reject_environment_interpolation(payload) return config_type.model_validate(payload) @@ -223,6 +236,20 @@ def _reject_run_environment_interpolation(payload: Mapping[str, object]) -> None _reject_environment_interpolation(builder_value) +def _reject_benchmark_environment_interpolation(payload: Mapping[str, object]) -> None: + for key, value in payload.items(): + _reject_environment_interpolation(key) + if key != "base_run" or not isinstance(value, Mapping): + _reject_environment_interpolation(value) + continue + for base_key, base_value in value.items(): + _reject_environment_interpolation(base_key) + if base_key == "inline" and isinstance(base_value, Mapping): + _reject_run_environment_interpolation(base_value) + else: + _reject_environment_interpolation(base_value) + + def _resolve_profile_path( explicit_path: str | Path | None, *, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index c8abeda28..9f83c0c92 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -531,13 +531,15 @@ def _publish_readiness(self, state: ReadinessState) -> None: if self._readiness is not None and not _can_advance_readiness(self._readiness.state, state): return revision = 1 if self._readiness is None else self._readiness.revision + 1 + timestamp = self._now() readiness = AttemptReadiness( schema_version=1, run_id=self._attempt.run_id, shard_id=self._attempt.shard_id, attempt_id=self._attempt.attempt_id, revision=revision, - updated_at=self._now(), + updated_at=timestamp, + started_at=self._attempt.updated_at if self._readiness is None else self._readiness.started_at, state=state, deployments=tuple( DeploymentReadiness( @@ -578,12 +580,6 @@ def _publish_stopped_readiness(self) -> None: for status in self._statuses: status.state = ReadinessState.STOPPED status.ready_backends = 0 - status.last_probe = _probe_evidence( - self._now(), - ProbeOutcome.SUCCESS, - "runtime_stopped", - "allocation processes stopped", - ) self._publish_readiness(ReadinessState.STOPPED) def _publish_terminal_attempt( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 28a0ea265..d7f125a36 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -136,6 +136,7 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non attempt_id=context.attempt.attempt_id, revision=previous.revision + 1, updated_at=timestamp, + started_at=previous.started_at, state=ReadinessState.READY, deployments=tuple( DeploymentReadiness( @@ -299,6 +300,7 @@ def _readiness( attempt_id=context.attempt.attempt_id, revision=1 if previous is None else previous.revision + 1, updated_at=timestamp, + started_at=timestamp if previous is None else previous.started_at, state=state, deployments=tuple( DeploymentReadiness( @@ -327,6 +329,7 @@ def _write_failed_and_stopped_readiness(context: AllocationContext, writer: Slur attempt_id=previous.attempt_id, revision=previous.revision + 1, updated_at=failed_at, + started_at=previous.started_at, state=ReadinessState.FAILED, deployments=tuple( deployment.model_copy( @@ -369,18 +372,13 @@ def _write_stopped_readiness( attempt_id=previous.attempt_id, revision=previous.revision + 1, updated_at=timestamp, + started_at=previous.started_at, state=ReadinessState.STOPPED, deployments=tuple( deployment.model_copy( update={ "state": ReadinessState.STOPPED, "ready_backends": 0, - "last_probe": _probe( - timestamp, - ProbeOutcome.SUCCESS, - "runtime_stopped", - "allocation processes stopped", - ), } ) for deployment in previous.deployments diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py index 79722df62..e6a3107f1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/__init__.py @@ -36,12 +36,17 @@ if TYPE_CHECKING: from data_designer.slurm.services.wiring import ( # noqa: F401 SlurmRunArtifactPublisher, + create_slurm_benchmark_service, create_slurm_image_service, create_slurm_run_service, ) _LAZY_IMPORTS = { "SlurmRunArtifactPublisher": ("data_designer.slurm.services.wiring", "SlurmRunArtifactPublisher"), + "create_slurm_benchmark_service": ( + "data_designer.slurm.services.wiring", + "create_slurm_benchmark_service", + ), "create_slurm_image_service": ("data_designer.slurm.services.wiring", "create_slurm_image_service"), "create_slurm_run_service": ("data_designer.slurm.services.wiring", "create_slurm_run_service"), } @@ -71,6 +76,7 @@ "SlurmServiceError", "SlurmServiceErrorCode", "SlurmServiceOperation", + "create_slurm_benchmark_service", "create_slurm_image_service", "create_slurm_profile_service", "create_slurm_run_service", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py index 5ae28389c..6e01c2d46 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/artifacts.py @@ -38,6 +38,7 @@ publish_immutable_text, sync_directory, ) +from data_designer.slurm.state.storage import StateStorage _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _TEMPORARY_PREFIX = ".artifact." @@ -62,8 +63,16 @@ def initialize( ) -> None: if force: raise StateConflictError("force cannot replace durable run state") - created_at = self._clock() writer = SlurmStateWriter(self._workspace_root, plan.run_id) + try: + created_at = writer.load_run().created_at + except StateNotFoundError: + try: + created_at = ( + StateStorage(self._workspace_root, plan.run_id).read_shard(plan.shards[0].shard_id).created_at + ) + except FileNotFoundError: + created_at = self._clock() plan_reference = ArtifactReference( path=(writer.run_root / "resolved-plan.json").as_posix(), sha256=plan.compute_sha256(), @@ -119,14 +128,22 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte ) ) - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - """Mark every initial attempt failed after its held job is cancelled.""" + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + """Fail initial attempts owned by the cancelled held job.""" writer = SlurmStateWriter(self._workspace_root, plan.run_id) for shard in plan.shards: try: attempt = writer.load_attempt(shard.shard_id, "attempt-0001") except StateNotFoundError: continue + if attempt.scheduler is None or attempt.scheduler.array_job_id != job_id: + continue writer.update_attempt( attempt.model_copy( update={ diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py index 4c4fd2192..ceb04ec03 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/benchmark.py @@ -5,6 +5,7 @@ from __future__ import annotations +from pathlib import Path from typing import Protocol from pydantic import TypeAdapter, ValidationError @@ -24,13 +25,25 @@ class SlurmBenchmarkBackend(Protocol): """Process benchmarks through a supported service dependency.""" - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: """Persist and start the ordinary child runs for one benchmark. Any non-``INTERNAL`` service error must contain a caller-safe message. """ - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: """Return one point-in-time benchmark report. Any non-``INTERNAL`` service error must contain a caller-safe message. @@ -49,7 +62,13 @@ class SlurmBenchmarkService: def __init__(self, backend: SlurmBenchmarkBackend) -> None: self._backend = backend - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: str | Path = ".", + force: bool = False, + ) -> BenchmarkManifest: """Start all ordinary child runs and return their immutable mapping. The returned manifest must reference the exact serialized benchmark config. @@ -60,9 +79,17 @@ def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: operation = SlurmServiceOperation.RUN_BENCHMARK if not isinstance(config, DataDesignerSlurmBenchmarkConfig): raise _make_invalid_request_error(operation, "config must be a DataDesignerSlurmBenchmarkConfig") + if not isinstance(source_root, str | Path): + raise _make_invalid_request_error(operation, "source_root must be a path") + if type(force) is not bool: + raise _make_invalid_request_error(operation, "force must be a boolean") def run_benchmark() -> BenchmarkManifest: - manifest = self._backend.run(config) + manifest = self._backend.run( + config, + source_root=Path(source_root).expanduser().resolve(), + force=force, + ) if not isinstance(manifest, BenchmarkManifest): raise TypeError("benchmark backend returned an invalid manifest") if manifest.benchmark_config.sha256 != config.compute_sha256(): @@ -71,12 +98,19 @@ def run_benchmark() -> BenchmarkManifest: return _invoke_service_backend(operation, run_benchmark) - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool = False, + fail_if_incomplete: bool = False, + ) -> BenchmarkReport: """Analyze one persisted benchmark without a resident monitor. Args: benchmark_id: Persisted benchmark identity to analyze. refresh_state: Request one point-in-time state refresh before analysis. + fail_if_incomplete: Fail after persisting a report containing non-success cases. Raises: SlurmServiceError: If the request is invalid or analysis fails. @@ -86,11 +120,15 @@ def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> B validated_id = _IDENTIFIER_ADAPTER.validate_python(benchmark_id, strict=True) except ValidationError: raise _make_invalid_request_error(operation, "benchmark_id must be a valid identifier") from None - if type(refresh_state) is not bool: - raise _make_invalid_request_error(operation, "refresh_state must be a boolean") + if type(refresh_state) is not bool or type(fail_if_incomplete) is not bool: + raise _make_invalid_request_error(operation, "analysis actions must be booleans") def analyze_benchmark() -> BenchmarkReport: - report = self._backend.analyze(validated_id, refresh_state=refresh_state) + report = self._backend.analyze( + validated_id, + refresh_state=refresh_state, + fail_if_incomplete=fail_if_incomplete, + ) if not isinstance(report, BenchmarkReport): raise TypeError("benchmark backend returned an invalid report") if report.benchmark_id != validated_id: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 691849bc4..6055e6604 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -17,6 +17,8 @@ from pydantic import JsonValue +from data_designer.slurm.benchmark.execution import SystemBenchmarkBackend +from data_designer.slurm.benchmark.observer import PersistedBenchmarkRunObserver from data_designer.slurm.client.dependencies import ( ClientDependencyResolutionError, ClientDependencyResolver, @@ -51,6 +53,7 @@ from data_designer.slurm.runtime.bundle import stage_runtime_bundle from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.services.artifacts import StateRunArtifactPublisher +from data_designer.slurm.services.benchmark import SlurmBenchmarkService from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.image_lifecycle import SlurmImageLifecycleManager from data_designer.slurm.services.images import SlurmImageService @@ -115,8 +118,14 @@ def initialize( def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: """Persist the submitted scheduler identity for every initial attempt.""" - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - """Mark initial attempts failed after a held submission is cancelled.""" + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + """Fail initial attempts owned by the cancelled held job.""" @dataclass(frozen=True, slots=True) @@ -338,7 +347,7 @@ def execute( f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", ) from error try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -353,7 +362,7 @@ def execute( pass else: try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except BaseException: pass raise @@ -369,7 +378,7 @@ def execute( f"held Slurm job {receipt.job_id} could not be released or cancelled", ) from error try: - self._record_submission_failure(publisher, plan) + self._record_submission_failure(publisher, plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -498,9 +507,14 @@ def _record_submission( "submission state cannot be recorded", ) from None - def _record_submission_failure(self, publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan) -> None: + def _record_submission_failure( + self, + publisher: SlurmRunArtifactPublisher, + plan: ResolvedSlurmRunPlan, + job_id: int, + ) -> None: try: - publisher.record_submission_failure(plan, failed_at=self._clock()) + publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) except (StateNotFoundError, StateConflictError, SlurmStateError): raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -724,6 +738,29 @@ def create_slurm_run_service( ) -> SlurmRunService: """Create the production run service for one selected cluster profile.""" selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + return _create_slurm_run_service( + selected, + artifact_publisher=artifact_publisher, + dependency_resolver=dependency_resolver, + launcher=launcher, + run_id_factory=run_id_factory, + clock=clock, + package_version=package_version, + source_environment=source_environment, + ) + + +def _create_slurm_run_service( + selected: SelectedSlurmProfile, + *, + artifact_publisher: SlurmRunArtifactPublisher | None = None, + dependency_resolver: ClientDependencyResolver | None = None, + launcher: SlurmCommandClient | None = None, + run_id_factory: RunIdFactory | None = None, + clock: Clock | None = None, + package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, +) -> SlurmRunService: command_client = launcher or SlurmCommandClient() selected_clock = clock or _utc_now preparer = _RunPreparer( @@ -770,6 +807,54 @@ def create_slurm_image_service( return SlurmImageService(backend, backend) +def create_slurm_benchmark_service( + *, + profile: SlurmProfile | None = None, + catalog: SlurmProfileCatalog | None = None, + profile_file: str | Path | None = None, + cluster: str | None = None, + artifact_publisher: SlurmRunArtifactPublisher | None = None, + dependency_resolver: ClientDependencyResolver | None = None, + launcher: SlurmCommandClient | None = None, + clock: Clock | None = None, + package_version: str | None = None, + source_environment: Mapping[str, str] | None = None, +) -> SlurmBenchmarkService: + """Create the production benchmark service for one selected cluster profile.""" + selected = resolve_profile(profile=profile, catalog=catalog, profile_file=profile_file, cluster=cluster) + command_client = launcher or SlurmCommandClient() + selected_clock = clock or _utc_now + publisher = artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock) + dependencies = dependency_resolver or ClientDependencyResolver() + environment = dict(os.environ if source_environment is None else source_environment) + version = package_version or importlib.metadata.version("data-designer-slurm") + + def create_child_service(run_id: Identifier) -> SlurmRunService: + return _create_slurm_run_service( + selected, + artifact_publisher=publisher, + dependency_resolver=dependencies, + launcher=command_client, + run_id_factory=lambda: run_id, + clock=selected_clock, + package_version=version, + source_environment=environment, + ) + + observer = PersistedBenchmarkRunObserver( + selected.profile.workspace_root, + command_client, + selected_clock, + ) + backend = SystemBenchmarkBackend( + selected.profile.workspace_root, + create_child_service, + observer, + selected_clock, + ) + return SlurmBenchmarkService(backend) + + def _resolve_builder_payload( authored: DataDesignerSlurmConfig, source_root: Path, @@ -802,6 +887,7 @@ def _format_job_ids(job_ids: list[int]) -> str: __all__ = [ "SlurmRunArtifactPublisher", + "create_slurm_benchmark_service", "create_slurm_image_service", "create_slurm_run_service", ] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py index 1fc65422c..15985b2a3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/__init__.py @@ -33,6 +33,7 @@ CollectionStatus, ) from data_designer.slurm.state.errors import ( + SchedulerStateConflictError, SlurmStateError, StateConflictError, StateCorruptionError, @@ -163,6 +164,7 @@ "SchedulerQueueRecord", "SchedulerObservation", "SchedulerState", + "SchedulerStateConflictError", "Sha256Digest", "ShardManifest", "ShardStatus", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py index e43897dcc..c3d5bb528 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/errors.py @@ -20,3 +20,7 @@ class StateConflictError(SlurmStateError): class StateCorruptionError(SlurmStateError): """Raised when persisted state cannot be safely read or validated.""" + + +class SchedulerStateConflictError(StateCorruptionError): + """Raised when persisted winners conflict with scheduler evidence.""" diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index 1c49d2952..aa459a6de 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -16,6 +16,7 @@ from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.base import SchedulerIdentity, SchedulerJobIdentity from data_designer.slurm.state.errors import ( + SchedulerStateConflictError, SlurmStateError, StateConflictError, StateCorruptionError, @@ -302,7 +303,7 @@ def _validate_winner_scheduler_consistency( return winning = next(status for status in statuses if status.attempt.attempt_id == winner.attempt_id) if winning.effective_state is not EffectiveAttemptState.SUCCEEDED: - raise StateCorruptionError("persisted winner conflicts with terminal scheduler evidence") + raise SchedulerStateConflictError("persisted winner conflicts with terminal scheduler evidence") def _validate_location(workspace_root: str | Path, run_id: Identifier) -> tuple[Path, Identifier]: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py index b5957da9b..27fc325a2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -86,13 +86,21 @@ class AttemptReadiness(StateRecord): attempt_id: AttemptId revision: PositiveInt updated_at: datetime + started_at: datetime | None = Field(default=None, exclude_if=lambda value: value is None) state: ReadinessState deployments: tuple[DeploymentReadiness, ...] = Field(min_length=1) _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) + @field_validator("started_at") + @classmethod + def validate_started_at(cls, value: datetime | None) -> datetime | None: + return None if value is None else validate_utc_timestamp(value) + @model_validator(mode="after") def validate_deployments(self) -> AttemptReadiness: + if self.started_at is not None and self.started_at > self.updated_at: + raise ValueError("readiness started_at cannot follow updated_at") _validate_deployment_uniqueness(self.deployments) _validate_probe_chronology(self.deployments, self.updated_at) _validate_attempt_state(self.state, self.deployments) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py index 473150f43..3dd9ab644 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reconciliation.py @@ -70,6 +70,7 @@ def validate_readiness_transition( _require(previous.shard_id == current.shard_id, "readiness shard_id cannot change") _require(previous.attempt_id == current.attempt_id, "readiness attempt_id cannot change") _require(current.revision == previous.revision + 1, "readiness revision must increase by exactly one") + _require(current.started_at == previous.started_at, "readiness started_at cannot change") _require(current.updated_at >= previous.updated_at, "readiness updated_at cannot move backward") _require( current.state in _ALLOWED_READINESS_TRANSITIONS[previous.state], diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 4f6f60694..39e2e6f71 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -53,6 +53,7 @@ _RUNTIME_DIRECTORY_NAME = "runtime" _RESUME_LOCK_FILENAME = "resume.lock" _LOCK_DIRECTORY_NAME = ".locks" +_SUBMISSION_LOCK_PREFIX = "submission" _MAXIMUM_RECORD_SIZE = 16 * 1024 * 1024 _ATTEMPT_NAME_PATTERN = re.compile(r"^attempt-[0-9]{4,}$") _RecordT = TypeVar("_RecordT", bound=ContractRecord) @@ -144,6 +145,20 @@ def acquire_run_lock(self) -> Iterator[None]: except FileNotFoundError as error: raise StateNotFoundError(f"run {self.run_id!r} is not initialized") from error + @contextmanager + def acquire_submission_lock(self) -> Iterator[None]: + """Serialize the initial submission decision for one run identity.""" + self.ensure_storage() + with open_verified_directory(self.runs_root, require_private=True) as runs_descriptor: + with open_verified_child_directory( + runs_descriptor, + _LOCK_DIRECTORY_NAME, + self.locks_root, + ) as locks_descriptor: + lock_name = f"{_SUBMISSION_LOCK_PREFIX}-{self.run_id}.lock" + with acquire_file_lock(locks_descriptor, lock_name, self.locks_root / lock_name): + yield + @contextmanager def acquire_shard_lock(self, shard_id: ShardId) -> Iterator[None]: try: diff --git a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py new file mode 100644 index 000000000..907acd5ea --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from data_designer.slurm.benchmark import BenchmarkCompiler +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig + + +def test_compiler_expands_stable_ordered_ordinary_runs( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + compiled = BenchmarkCompiler.compile(benchmark_config, authored_run) + equivalent = DataDesignerSlurmBenchmarkConfig.model_validate(benchmark_config.model_dump(mode="json")) + + assert compiled == BenchmarkCompiler.compile(equivalent, authored_run) + assert tuple(case.case_id for case in compiled.cases) == ( + "two-independent-replicas-c32", + "one-two-node-replica-c32", + "two-independent-replicas-c64", + "one-two-node-replica-c64", + "two-independent-replicas-c128", + "one-two-node-replica-c128", + ) + assert len({case.child_run_id for case in compiled.cases}) == 6 + assert all(case.child_run_config.invocation.num_records == 1000 for case in compiled.cases) + assert [deployment.resources.nodes for deployment in compiled.cases[0].child_run_config.deployments] == [2, 1] + assert [deployment.topology.nodes_per_replica for deployment in compiled.cases[1].child_run_config.deployments] == [ + 2, + 1, + ] + assert compiled.cases[0].child_run_config.invocation.model_concurrency == {"generator": 32, "judge": 32} + + +def test_compiled_child_maps_are_immutable( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + compiled = BenchmarkCompiler.compile(benchmark_config, authored_run) + + with pytest.raises(TypeError, match="frozen dictionary"): + compiled.cases[0].child_run_config.invocation.model_concurrency["generator"] = 1 + + +def test_compiler_rejects_unknown_alias_at_the_boundary( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + invalid = benchmark_config.model_copy(update={"model_aliases": ["missing"]}) + + with pytest.raises(ValueError, match="unknown model aliases"): + BenchmarkCompiler.compile(invalid, authored_run) + + +@pytest.mark.parametrize( + ("concurrency", "expected_records"), + [(1, 1000), (2000, 2000), (6000, 5000)], +) +def test_adaptive_record_policy_is_bounded( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + concurrency: int, + expected_records: int, +) -> None: + config = benchmark_config.model_copy( + update={ + "concurrency_values": [concurrency], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + + compiled = BenchmarkCompiler.compile(config, authored_run) + + assert compiled.cases[0].child_run_config.invocation.num_records == expected_records diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py new file mode 100644 index 000000000..99e23af5f --- /dev/null +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from threading import Event, Lock + +import pytest + +from data_designer.slurm.benchmark.analysis import ( + BenchmarkAnalyzer, + BenchmarkMeasurements, + BenchmarkObservationFailure, + BenchmarkRunObservation, + BenchmarkShardMeasurements, +) +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler +from data_designer.slurm.benchmark.execution import SystemBenchmarkBackend +from data_designer.slurm.benchmark.records import ( + BenchmarkChildRun, + BenchmarkOutcome, + BenchmarkRecommendationKind, +) +from data_designer.slurm.benchmark.store import BenchmarkConflictError, BenchmarkStore +from data_designer.slurm.config import ( + BenchmarkBaseRun, + DataDesignerSlurmBenchmarkConfig, + DataDesignerSlurmConfig, +) +from data_designer.slurm.contracts import ArtifactReference +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services import ( + SlurmRunExecution, + SlurmServiceError, + SlurmServiceErrorCode, + SlurmServiceOperation, +) + +NOW = datetime(2026, 8, 19, 12, tzinfo=timezone.utc) + + +class _Observer: + def __init__(self, responses=None) -> None: + self.responses = {} if responses is None else responses + self.calls = [] + + def observe(self, run_id: str, *, refresh_state: bool) -> BenchmarkRunObservation: + self.calls.append((run_id, refresh_state)) + response = self.responses[run_id] + if isinstance(response, Exception): + raise response + return response + + +class _RunService: + def __init__(self, run_id, configs, submissions, calls, benchmark_root, failures) -> None: + self.run_id = run_id + self.configs = configs + self.submissions = submissions + self.calls = calls + self.benchmark_root = benchmark_root + self.failures = failures + + def execute(self, config, *, source_root, dry_run, force): + assert (self.benchmark_root / "benchmark.json").is_file() + self.calls.append((self.run_id, config, source_root, dry_run, force)) + self.configs[self.run_id] = config + if self.run_id in self.failures: + raise SlurmServiceError( + SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceOperation.EXECUTE_RUN, + "submission unavailable", + ) + self.submissions.add(self.run_id) + return SlurmRunExecution( + run_id=self.run_id, + state="submitted", + plan_sha256="1" * 64, + shard_count=1, + job_id=4101, + ) + + +class _ConcurrentRunService(_RunService): + def __init__(self, *args, second_check: Event) -> None: + super().__init__(*args) + self.second_check = second_check + + def execute(self, config, *, source_root, dry_run, force): + self.second_check.wait(timeout=1) + return super().execute(config, source_root=source_root, dry_run=dry_run, force=force) + + +class _ConcurrentSubmissionLoader: + def __init__(self, submissions: set[str], second_check: Event) -> None: + self._submissions = submissions + self._second_check = second_check + self._lock = Lock() + self._checks = 0 + + def __call__(self, run_id: str) -> bool: + with self._lock: + self._checks += 1 + if self._checks == 2: + self._second_check.set() + return run_id in self._submissions + + +def _inline_config( + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> DataDesignerSlurmBenchmarkConfig: + return benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run), + "concurrency_values": [32], + } + ) + + +def _manifest(workspace: Path, config, compiled): + store = BenchmarkStore(workspace, compiled.benchmark_id) + children = tuple( + BenchmarkChildRun( + case_id=case.case_id, + child_run_id=case.child_run_id, + child_authored_config=ArtifactReference( + path=(workspace / "runs" / case.child_run_id / "authored-config.json").as_posix(), + sha256=case.child_run_config.compute_sha256(), + ), + ) + for case in compiled.cases + ) + return store, store.build_manifest(config, children) + + +def _measurements( + actual_records: int, + boot_seconds: float, + generation_seconds: float, + wall_seconds: float, + shard_count: int, +) -> BenchmarkMeasurements: + records_per_shard = actual_records // shard_count + return BenchmarkMeasurements( + shards=tuple( + BenchmarkShardMeasurements( + actual_records=( + actual_records - records_per_shard * (shard_count - 1) + if index == shard_count - 1 + else records_per_shard + ), + boot_seconds=boot_seconds, + generation_seconds=generation_seconds, + wall_seconds=wall_seconds, + ) + for index in range(shard_count) + ) + ) + + +def test_store_is_convergent_and_rejects_tampered_metadata( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + + assert store.publish(config, authored_run, manifest) == manifest + assert store.publish(config, authored_run, manifest) == manifest + assert store.load() == (config, authored_run, manifest) + + store.config_path.write_text("{}\n") + with pytest.raises(BenchmarkConflictError): + store.load() + + +def test_run_persists_manifest_before_submission_and_is_idempotent( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + configs = {} + submissions = set() + calls = [] + failures = set() + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + first = backend.run(config, source_root=tmp_path, force=False) + second = backend.run(config, source_root=tmp_path, force=False) + + assert first == second + assert len(calls) == len(compiled.cases) + assert tuple(configs) == tuple(case.child_run_id for case in compiled.cases) + assert all(force is False for *_, force in calls) + + +def test_concurrent_runs_submit_each_deterministic_child_once( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + configs = {} + submissions = set() + calls = [] + failures = set() + second_check = Event() + submission_loader = _ConcurrentSubmissionLoader(submissions, second_check) + store, manifest = _manifest(tmp_path, config, compiled) + store.publish(config, authored_run, manifest) + + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _ConcurrentRunService( + run_id, + configs, + submissions, + calls, + store.benchmark_root, + failures, + second_check=second_check, + ), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submission_loader, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = tuple(executor.submit(backend.run, config, source_root=tmp_path, force=False) for _ in range(2)) + assert tuple(future.result() for future in futures) == (manifest, manifest) + + assert len(calls) == len(compiled.cases) + + +def test_initialized_child_without_submission_requires_force_to_resume( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + missing_id = compiled.cases[0].child_run_id + configs = {} + submissions = set() + calls = [] + failures = {missing_id} + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + with pytest.raises(SlurmServiceError, match="1 of 2"): + backend.run(config, source_root=tmp_path, force=False) + failures.clear() + with pytest.raises(SlurmServiceError, match="1 of 2") as conflict: + backend.run(config, source_root=tmp_path, force=False) + assert conflict.value.code is SlurmServiceErrorCode.CONFLICT + + backend.run(config, source_root=tmp_path, force=True) + + assert [run_id for run_id, *_ in calls].count(missing_id) == 2 + + +def test_partial_submission_attempts_all_children_and_resumes_only_missing( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + missing_id = compiled.cases[0].child_run_id + configs = {} + submissions = set() + calls = [] + failures = {missing_id} + benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id + backend = SystemBenchmarkBackend( + tmp_path, + lambda run_id: _RunService(run_id, configs, submissions, calls, benchmark_root, failures), + _Observer(), + lambda: NOW, + child_config_loader=configs.get, + child_submission_loader=submissions.__contains__, + ) + + with pytest.raises(SlurmServiceError, match="1 of 2"): + backend.run(config, source_root=tmp_path, force=True) + assert len(calls) == 2 + + failures.clear() + backend.run(config, source_root=tmp_path, force=True) + + assert [run_id for run_id, *_ in calls].count(missing_id) == 2 + assert len(calls) == 3 + assert all(force is False for *_, force in calls) + + +def test_analysis_preserves_mixed_outcomes_and_stable_order( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + successful = compiled.cases[0] + plan = multi_node_plan.model_copy(update={"run_id": successful.child_run_id}) + observer = _Observer( + { + successful.child_run_id: BenchmarkRunObservation( + authored_config=successful.child_run_config, + resolved_plan=plan, + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements(1000, 60, 120, 180, len(plan.shards)), + ), + compiled.cases[1].child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING), + } + ) + + report = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=True, + ) + repeated = BenchmarkAnalyzer(observer, lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=True, + ) + + assert tuple(case.outcome for case in report.cases) == ( + BenchmarkOutcome.SUCCEEDED, + BenchmarkOutcome.MISSING, + ) + assert report.cases[0].rows_per_second == 1000 / 120 + assert report.cases[0].generation_seconds == 14340 + assert report.cases[0].target_jobs == 17 + assert report.cases[0].total_gpu_hours == report.cases[0].gpu_hours_per_job * 17 + assert report.serialize_json() == repeated.serialize_json() + assert observer.calls == [(case.child_run_id, True) for case in compiled.cases] * 2 + + +def test_analysis_keeps_measured_infeasible_case_successful( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run).model_copy( + update={"deployment_cases": [benchmark_config.deployment_cases[0]]} + ) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + case = compiled.cases[0] + observation = BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements( + case.child_run_config.invocation.num_records, + 15000, + 120, + 15120, + len(multi_node_plan.shards), + ), + ) + + report = BenchmarkAnalyzer(_Observer({case.child_run_id: observation}), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert report.cases[0].outcome is BenchmarkOutcome.SUCCEEDED + assert report.cases[0].feasible is False + assert report.cases[0].generation_seconds == 0 + assert report.cases[0].target_jobs is None + assert report.cases[0].total_gpu_hours is None + + +def test_analysis_rejects_manifest_rewritten_away_from_compiled_children( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + first = manifest.children[0] + rewritten_id = f"{first.child_run_id}-other" + rewritten = BenchmarkChildRun( + case_id=first.case_id, + child_run_id=rewritten_id, + child_authored_config=ArtifactReference( + path=(tmp_path / "runs" / rewritten_id / "authored-config.json").as_posix(), + sha256="f" * 64, + ), + ) + tampered = manifest.model_copy(update={"children": (rewritten, *manifest.children[1:])}) + + with pytest.raises(ValueError, match="deterministic child expansion"): + BenchmarkAnalyzer(_Observer(), lambda: NOW).analyze( + config, + authored_run, + tampered, + store.manifest_reference(tampered), + refresh_state=False, + ) + + +def test_analysis_recommendations_are_stable_and_select_dominating_case( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + responses = {} + for index, case in enumerate(compiled.cases, start=1): + responses[case.child_run_id] = BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=BenchmarkOutcome.SUCCEEDED, + measurements=_measurements( + case.child_run_config.invocation.num_records, + 60, + 120 * index, + 180 * index, + len(multi_node_plan.shards), + ), + ) + + report = BenchmarkAnalyzer(_Observer(responses), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert tuple(item.kind for item in report.recommendations) == ( + BenchmarkRecommendationKind.PARETO, + BenchmarkRecommendationKind.MINIMUM_JOBS, + BenchmarkRecommendationKind.MINIMUM_GPU_HOURS, + ) + assert {item.case_id for item in report.recommendations} == {compiled.cases[0].case_id} + + +@pytest.mark.parametrize( + "outcome", + [ + BenchmarkOutcome.PENDING, + BenchmarkOutcome.ACCOUNTING_LAG, + BenchmarkOutcome.FAILED, + BenchmarkOutcome.INCOMPLETE, + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + ], +) +def test_analysis_keeps_each_non_success_outcome_explicit( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, + outcome: BenchmarkOutcome, +) -> None: + config = _inline_config(benchmark_config, authored_run).model_copy( + update={"deployment_cases": [benchmark_config.deployment_cases[0]]} + ) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + case = compiled.cases[0] + response = ( + BenchmarkObservationFailure(outcome) + if outcome + in { + BenchmarkOutcome.MISSING, + BenchmarkOutcome.STALE, + BenchmarkOutcome.SCHEDULER_INCONSISTENT, + } + else BenchmarkRunObservation( + authored_config=case.child_run_config, + resolved_plan=multi_node_plan.model_copy(update={"run_id": case.child_run_id}), + outcome=outcome, + ) + ) + + report = BenchmarkAnalyzer(_Observer({case.child_run_id: response}), lambda: NOW).analyze( + config, + authored_run, + manifest, + store.manifest_reference(manifest), + refresh_state=False, + ) + + assert report.cases[0].outcome is outcome + + +def test_fail_if_incomplete_persists_the_point_in_time_report_first( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run: DataDesignerSlurmConfig, +) -> None: + config = _inline_config(benchmark_config, authored_run) + compiled = BenchmarkCompiler.compile(config, authored_run) + store, manifest = _manifest(tmp_path, config, compiled) + store.publish(config, authored_run, manifest) + observer = _Observer( + {case.child_run_id: BenchmarkObservationFailure(BenchmarkOutcome.MISSING) for case in compiled.cases} + ) + backend = SystemBenchmarkBackend( + tmp_path, + lambda _: pytest.fail("analysis must not create a run service"), + observer, + lambda: NOW, + ) + + with pytest.raises(SlurmServiceError, match="2 incomplete cases"): + backend.analyze( + manifest.benchmark_id, + refresh_state=False, + fail_if_incomplete=True, + ) + + assert len(tuple(store.reports_root.iterdir())) == 1 diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 7abca1b66..c2ff3087f 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -13,12 +13,14 @@ from data_designer.slurm.config import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, DataDesignerSlurmConfigBuilder, ProfileSelectionSource, SlurmConfigBuilderError, SlurmConfigLoadError, SlurmProfileCatalog, + load_benchmark_config, load_builder_payload, load_profile_catalog, load_run_config, @@ -79,6 +81,16 @@ def test_builder_write_config_round_trips_supported_formats(tmp_path: Path, suff assert load_run_config(path) == builder.build() +def test_benchmark_loader_round_trips_strict_config( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, +) -> None: + path = tmp_path / "benchmark.json" + path.write_text(benchmark_config.serialize_json()) + + assert load_benchmark_config(path) == benchmark_config + + def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: with pytest.raises(SlurmConfigBuilderError, match="must end"): _config_builder().write_config(tmp_path / "run.txt") diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index 90009a9b4..e35ba4f98 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -262,6 +262,22 @@ def test_successful_benchmark_case_requires_complete_positive_output(field: str, BenchmarkReport.model_validate_json(json.dumps(payload)) +def test_successful_infeasible_benchmark_case_allows_zero_budgeted_generation() -> None: + payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) + case = payload["cases"][0] + case.update( + feasible=False, + generation_seconds=0, + target_jobs=None, + total_gpu_hours=None, + ) + payload["recommendations"] = [] + + report = BenchmarkReport.model_validate_json(json.dumps(payload)) + + assert report.cases[0].feasible is False + + def test_benchmark_report_allows_pareto_frontier_but_singleton_minima() -> None: payload = json.loads((GOLDEN_DIR / "benchmark_report.json").read_text()) second_case = deepcopy(payload["cases"][0]) diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index df3439944..b8bf75f5f 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -193,6 +193,10 @@ def write_result_under_dataset_lease() -> None: assert ReadinessState.READY in [item.state for item in state.readiness] assert state.readiness[-1].state is ReadinessState.STOPPED assert state.readiness[-1].deployments[0].endpoint_publication is EndpointPublicationState.PUBLISHED + assert state.readiness[-1].deployments[0].last_probe is not None + assert state.readiness[-1].deployments[0].last_probe.reason_code == "endpoint_ready" + assert state.readiness[0].started_at is not None + assert {item.started_at for item in state.readiness} == {state.readiness[0].started_at} assert all(process.poll() is not None for process in runner.processes) assert {step.stdout_path.parent.name for step in runner.steps} == {"execution-00000001"} diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index b3df00eb7..df0b0e47e 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -558,23 +558,28 @@ def test_benchmark_service_delegates_run_and_analysis( benchmark_report: BenchmarkReport, ) -> None: benchmark_manifest = manifest_with_matching_config_digest + source_root = Path.cwd().resolve() backend = FakeBenchmarkBackend( - run_responses=((benchmark_config, benchmark_manifest),), - analysis_responses=(((benchmark_manifest.benchmark_id, True), benchmark_report),), + run_responses=(((benchmark_config, source_root, True), benchmark_manifest),), + analysis_responses=(((benchmark_manifest.benchmark_id, True, True), benchmark_report),), ) service = SlurmBenchmarkService(backend) - assert service.run(benchmark_config) is benchmark_manifest - assert service.analyze(benchmark_manifest.benchmark_id, refresh_state=True) is benchmark_report - assert backend.run_calls == [benchmark_config] - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, True)] + assert service.run(benchmark_config, force=True) is benchmark_manifest + assert ( + service.analyze(benchmark_manifest.benchmark_id, refresh_state=True, fail_if_incomplete=True) + is benchmark_report + ) + assert backend.run_calls == [(benchmark_config, source_root, True)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, True, True)] backend.assert_complete() def test_benchmark_service_rejects_invalid_manifest( benchmark_config: DataDesignerSlurmBenchmarkConfig, ) -> None: - backend = FakeBenchmarkBackend(run_responses=((benchmark_config, object()),)) # type: ignore[arg-type] + request = (benchmark_config, Path.cwd().resolve(), False) + backend = FakeBenchmarkBackend(run_responses=((request, object()),)) # type: ignore[arg-type] with pytest.raises(SlurmServiceError) as caught: SlurmBenchmarkService(backend).run(benchmark_config) @@ -582,13 +587,13 @@ def test_benchmark_service_rejects_invalid_manifest( assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK assert str(caught.value) == "run benchmark failed" - assert backend.run_calls == [benchmark_config] + assert backend.run_calls == [request] backend.assert_complete() def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkManifest) -> None: backend = FakeBenchmarkBackend( - analysis_responses=(((benchmark_manifest.benchmark_id, False), object()),), # type: ignore[arg-type] + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), object()),), # type: ignore[arg-type] ) with pytest.raises(SlurmServiceError) as caught: @@ -597,7 +602,7 @@ def test_benchmark_service_rejects_invalid_report(benchmark_manifest: BenchmarkM assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK assert str(caught.value) == "analyze benchmark failed" - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False, False)] backend.assert_complete() @@ -607,14 +612,15 @@ def test_benchmark_service_rejects_manifest_for_another_config( ) -> None: reference = manifest_with_matching_config_digest.benchmark_config.model_copy(update={"sha256": "0" * 64}) manifest = manifest_with_matching_config_digest.model_copy(update={"benchmark_config": reference}) - backend = FakeBenchmarkBackend(run_responses=((benchmark_config, manifest),)) + request = (benchmark_config, Path.cwd().resolve(), False) + backend = FakeBenchmarkBackend(run_responses=((request, manifest),)) with pytest.raises(SlurmServiceError) as caught: SlurmBenchmarkService(backend).run(benchmark_config) assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.RUN_BENCHMARK - assert backend.run_calls == [benchmark_config] + assert backend.run_calls == [request] backend.assert_complete() @@ -629,17 +635,22 @@ def test_benchmark_service_rejects_untyped_config() -> None: @pytest.mark.parametrize( - ("benchmark_id", "refresh_state"), - [("", False), ("invalid/id", False), ("benchmark-001", 1)], + ("benchmark_id", "refresh_state", "fail_if_incomplete"), + [("", False, False), ("invalid/id", False, False), ("benchmark-001", 1, False), ("benchmark-001", False, 1)], ) def test_benchmark_service_validates_analysis_actions( benchmark_id: object, refresh_state: object, + fail_if_incomplete: object, ) -> None: service = SlurmBenchmarkService(FakeBenchmarkBackend()) with pytest.raises(SlurmServiceError) as caught: - service.analyze(benchmark_id, refresh_state=refresh_state) # type: ignore[arg-type] + service.analyze( # type: ignore[arg-type] + benchmark_id, + refresh_state=refresh_state, + fail_if_incomplete=fail_if_incomplete, + ) assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK @@ -651,7 +662,7 @@ def test_benchmark_service_rejects_uncorrelated_report( ) -> None: report = benchmark_report.model_copy(update={"benchmark_id": "other-benchmark"}) backend = FakeBenchmarkBackend( - analysis_responses=(((benchmark_manifest.benchmark_id, False), report),), + analysis_responses=(((benchmark_manifest.benchmark_id, False, False), report),), ) with pytest.raises(SlurmServiceError) as caught: @@ -659,7 +670,7 @@ def test_benchmark_service_rejects_uncorrelated_report( assert caught.value.code is SlurmServiceErrorCode.INTERNAL assert caught.value.operation is SlurmServiceOperation.ANALYZE_BENCHMARK - assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False)] + assert backend.analysis_calls == [(benchmark_manifest.benchmark_id, False, False)] backend.assert_complete() diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index cd453f598..02f072ebe 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -12,13 +12,17 @@ import pytest import data_designer.slurm.services.retry_collection as retry_collection_module +from data_designer.slurm.benchmark.compiler import BenchmarkCompiler from data_designer.slurm.client.dependencies import ResolvedClientDependencies from data_designer.slurm.config import ( + BenchmarkBaseRun, BuilderInput, + DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, SecretRef, SlurmProfile, SlurmProfileCatalog, + select_profile, ) from data_designer.slurm.contracts import ArtifactReference, canonical_json from data_designer.slurm.images.records import RegisteredImage @@ -36,6 +40,7 @@ SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, + create_slurm_benchmark_service, create_slurm_image_service, create_slurm_run_service, ) @@ -59,9 +64,11 @@ def __init__( self, gpu_counts: tuple[int, ...] = (), *, - submission_job_ids: tuple[int, ...] = (42,), + submission_job_ids: tuple[int, ...] | None = None, cancel_error: Exception | None = None, + job_id: int = 42, release_error: Exception | None = None, + submission_error: Exception | None = None, ) -> None: self.submissions: list[str] = [] self.cancellations: list[int] = [] @@ -71,10 +78,11 @@ def __init__( self.queue_entries: tuple[SlurmQueueEntry, ...] = () self.accounting_entries: tuple[SlurmAccountingEntry, ...] = () self.submission_matches: tuple[SlurmSubmissionMatch, ...] = () - self.submission_job_ids = submission_job_ids + self.submission_job_ids = (job_id,) if submission_job_ids is None else submission_job_ids self.gpu_counts = gpu_counts self.cancel_error = cancel_error self.release_error = release_error + self.submission_error = submission_error def submit_script( self, @@ -86,6 +94,8 @@ def submit_script( self.submissions.append(script) self.held_submissions.append(hold) self.exported_environments.append(dict(export_environment or {})) + if self.submission_error is not None: + raise self.submission_error index = min(len(self.submissions) - 1, len(self.submission_job_ids) - 1) return SlurmJobSubmissionReceipt(job_id=self.submission_job_ids[index]) @@ -129,8 +139,9 @@ def __init__( submission_error: BaseException | None = None, ) -> None: self.initializations: list[tuple[str, bool]] = [] + self.plans: list[ResolvedSlurmRunPlan] = [] self.submissions: list[tuple[str, int, datetime]] = [] - self.submission_failures: list[tuple[str, datetime]] = [] + self.submission_failures: list[tuple[str, int, datetime]] = [] self.initialization_error = initialization_error self.submission_error = submission_error @@ -148,6 +159,7 @@ def initialize( assert plan.authored_config.sha256 == authored.compute_sha256() assert builder_payload is None assert all(path.is_file() for path in dependencies.wheel_sources) + self.plans.append(plan) self.initializations.append((plan.run_id, force)) def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitted_at: datetime) -> None: @@ -155,8 +167,14 @@ def record_submission(self, plan: ResolvedSlurmRunPlan, job_id: int, *, submitte raise self.submission_error self.submissions.append((plan.run_id, job_id, submitted_at)) - def record_submission_failure(self, plan: ResolvedSlurmRunPlan, *, failed_at: datetime) -> None: - self.submission_failures.append((plan.run_id, failed_at)) + def record_submission_failure( + self, + plan: ResolvedSlurmRunPlan, + job_id: int, + *, + failed_at: datetime, + ) -> None: + self.submission_failures.append((plan.run_id, job_id, failed_at)) def _profile(tmp_path: Path, profile_catalog: SlurmProfileCatalog) -> SlurmProfile: @@ -213,6 +231,116 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( assert launcher.submissions == [] +def test_production_benchmark_wiring_submits_each_case_as_an_ordinary_run( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher() + publisher = _Publisher() + service = create_slurm_benchmark_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + artifact_publisher=publisher, + package_version="0.9.2", + ) + + manifest = service.run(config, source_root=tmp_path, force=True) + + assert len(manifest.children) == 1 + assert publisher.initializations == [(manifest.children[0].child_run_id, False)] + assert len(launcher.submissions) == 1 + assert (tmp_path / "benchmarks" / manifest.benchmark_id / "benchmark.json").is_file() + + +def test_production_benchmark_children_preserve_catalog_profile_provenance( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + catalog = profile_catalog.model_copy( + update={ + "clusters": { + name: profile.model_copy(update={"workspace_root": tmp_path.as_posix()}) + for name, profile in profile_catalog.clusters.items() + } + } + ) + _register_images(tmp_path, authored_run_single, single_node_plan) + publisher = _Publisher() + service = create_slurm_benchmark_service( + catalog=catalog, + cluster="primary", + launcher=_Launcher(), # type: ignore[arg-type] + artifact_publisher=publisher, + package_version="0.9.2", + ) + + service.run(config, source_root=tmp_path) + + assert publisher.plans[0].selected_profile == select_profile(catalog, cluster="primary") + + +@pytest.mark.parametrize("drop_run_manifest", [False, True]) +def test_production_benchmark_force_resumes_initialized_child_without_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + drop_run_manifest: bool, +) -> None: + config = benchmark_config.model_copy( + update={ + "base_run": BenchmarkBaseRun(inline=authored_run_single), + "concurrency_values": [32], + "deployment_cases": [benchmark_config.deployment_cases[0]], + } + ) + child_run_id = BenchmarkCompiler.compile(config, authored_run_single).cases[0].child_run_id + _register_images(tmp_path, authored_run_single, single_node_plan) + launcher = _Launcher(submission_error=SlurmLauncherError("unavailable")) + service = create_slurm_benchmark_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, # type: ignore[arg-type] + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError, match="1 of 1") as unavailable: + service.run(config, source_root=tmp_path) + assert unavailable.value.code is SlurmServiceErrorCode.UNAVAILABLE + run_root = tmp_path / "runs" / child_run_id + if drop_run_manifest: + (run_root / "run.json").unlink() + launcher.submission_error = None + + manifest = service.run(config, source_root=tmp_path, force=True) + + writer = SlurmStateWriter(tmp_path, manifest.children[0].child_run_id) + shard = writer.load_shards()[0] + assert len(launcher.submissions) == 2 + assert writer.load_attempts(shard.shard_id)[0].scheduler is not None + + def test_public_plan_resolves_builder_relative_to_explicit_source_root( tmp_path: Path, profile_catalog: SlurmProfileCatalog, @@ -882,6 +1010,39 @@ def test_recording_conflict_cancels_the_accepted_job( assert len(publisher.submission_failures) == 1 +def test_recording_conflict_does_not_fail_another_submission_attempt( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + _register_images(tmp_path, authored_run_single, single_node_plan) + profile = _profile(tmp_path, profile_catalog) + first = create_slurm_run_service( + profile=profile, + launcher=_Launcher(job_id=41), # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + first.execute(authored_run_single, source_root=tmp_path) + launcher = _Launcher(job_id=42) + second = create_slurm_run_service( + profile=profile, + launcher=launcher, # type: ignore[arg-type] + run_id_factory=lambda: "run-wired", + package_version="0.9.2", + ) + + with pytest.raises(SlurmServiceError) as caught: + second.execute(authored_run_single, source_root=tmp_path) + + assert caught.value.code is SlurmServiceErrorCode.CONFLICT + assert launcher.cancellations == [42] + attempt = SlurmStateWriter(tmp_path, "run-wired").load_attempt("shard-00000", "attempt-0001") + assert attempt.state is AttemptLifecycleState.SUBMITTED + assert attempt.scheduler == SchedulerIdentity(array_job_id=41, array_task_id=0) + + def test_partial_submission_recording_failure_cancels_job_and_fails_created_attempts( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py index de32b94d1..9db2d7af0 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/services.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/services.py @@ -119,21 +119,38 @@ class FakeBenchmarkBackend(SlurmBenchmarkBackend): def __init__( self, *, - run_responses: Iterable[tuple[DataDesignerSlurmBenchmarkConfig, BenchmarkManifest | BaseException]] = (), - analysis_responses: Iterable[tuple[tuple[Identifier, bool], BenchmarkReport | BaseException]] = (), + run_responses: Iterable[ + tuple[tuple[DataDesignerSlurmBenchmarkConfig, Path, bool], BenchmarkManifest | BaseException] + ] = (), + analysis_responses: Iterable[tuple[tuple[Identifier, bool, bool], BenchmarkReport | BaseException]] = (), ) -> None: self._run_script = _ScriptedResponses(run_responses) self._analysis_script = _ScriptedResponses(analysis_responses) self.run_calls = self._run_script.calls self.analysis_calls = self._analysis_script.calls - def run(self, config: DataDesignerSlurmBenchmarkConfig) -> BenchmarkManifest: + def run( + self, + config: DataDesignerSlurmBenchmarkConfig, + *, + source_root: Path, + force: bool, + ) -> BenchmarkManifest: """Return the manifest scripted for one exact benchmark config.""" - return self._run_script.next(config, operation="benchmark run") + return self._run_script.next((config, source_root, force), operation="benchmark run") - def analyze(self, benchmark_id: Identifier, *, refresh_state: bool = False) -> BenchmarkReport: + def analyze( + self, + benchmark_id: Identifier, + *, + refresh_state: bool, + fail_if_incomplete: bool, + ) -> BenchmarkReport: """Return the report scripted for one benchmark and refresh action.""" - return self._analysis_script.next((benchmark_id, refresh_state), operation="benchmark analysis") + return self._analysis_script.next( + (benchmark_id, refresh_state, fail_if_incomplete), + operation="benchmark analysis", + ) def assert_complete(self) -> None: """Assert that every scripted benchmark response was consumed.""" diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 41bc6dc73..4781fdccf 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -12,8 +12,11 @@ from typing import cast import pytest -from slurm_test_fakes import FakeCommandResponse, FakeSlurmRunner +from slurm_test_fakes import FakeCommandResponse, FakeSlurmArray, FakeSlurmRunner, FakeSlurmTask +from data_designer.slurm.benchmark.analysis import BenchmarkObservationFailure +from data_designer.slurm.benchmark.observer import PersistedBenchmarkRunObserver +from data_designer.slurm.benchmark.records import BenchmarkOutcome from data_designer.slurm.client import ClientOutcome, ClientResult from data_designer.slurm.config import DataDesignerSlurmConfig, SlurmProfile from data_designer.slurm.contracts import ArtifactReference, compute_canonical_json_sha256 @@ -23,13 +26,19 @@ from data_designer.slurm.state import ( AttemptLifecycleState, AttemptManifest, + AttemptReadiness, AttemptTerminalClassification, CandidateOutcome, CandidateOutputFile, CandidateOutputManifest, + DeploymentReadiness, EffectiveAttemptState, EffectiveRunState, + EndpointPublicationState, GenerationState, + ProbeEvidence, + ProbeOutcome, + ReadinessState, RunManifest, SchedulerIdentity, SchedulerJobIdentity, @@ -193,6 +202,290 @@ def test_fresh_process_refresh_persists_one_fixed_accounting_lag_deadline( assert stat.S_IMODE(scheduler_path.stat().st_mode) == 0o600 +def test_benchmark_observer_refreshes_from_a_fresh_process( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state=None) + observed_at = case.created_at + timedelta(minutes=3) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: observed_at, + ).observe(case.plan.run_id, refresh_state=True) + + assert observation.outcome is BenchmarkOutcome.ACCOUNTING_LAG + assert observation.authored_config == authored_run_single + assert observation.resolved_plan == case.plan + + +def test_benchmark_observer_preserves_missing_and_tampered_child_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan, submitted=False) + observer = PersistedBenchmarkRunObserver(case.workspace, _StaticSchedulerClient((), ()), lambda: case.created_at) + + with pytest.raises(BenchmarkObservationFailure) as missing: + observer.observe("run-missing", refresh_state=False) + assert missing.value.outcome is BenchmarkOutcome.MISSING + + authored_path = case.writer.run_root / "authored-config.json" + authored_path.write_text("{}\n") + with pytest.raises(BenchmarkObservationFailure) as stale: + observer.observe(case.plan.run_id, refresh_state=False) + assert stale.value.outcome is BenchmarkOutcome.STALE + + +@pytest.mark.parametrize( + ("attempt_state", "terminal_classification", "expected"), + [ + (AttemptLifecycleState.SUBMITTED, None, BenchmarkOutcome.PENDING), + ( + AttemptLifecycleState.FAILED, + AttemptTerminalClassification.FAILED, + BenchmarkOutcome.FAILED, + ), + ], +) +def test_benchmark_observer_reconstructs_persisted_non_success_outcomes( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + attempt_state: AttemptLifecycleState, + terminal_classification: AttemptTerminalClassification | None, + expected: BenchmarkOutcome, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + if attempt_state is not AttemptLifecycleState.SUBMITTED: + case.writer.update_attempt( + _copy_attempt( + case.attempt, + state=attempt_state, + terminal_classification=terminal_classification, + ) + ) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + _StaticSchedulerClient((), ()), + lambda: case.created_at, + ).observe(case.plan.run_id, refresh_state=False) + + assert observation.outcome is expected + assert observation.measurements is None + + +def test_benchmark_observer_requires_measurements_for_persisted_winner( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + _publish_winner_state(case) + + observation = PersistedBenchmarkRunObserver( + case.workspace, + _StaticSchedulerClient((), ()), + lambda: case.created_at, + ).observe(case.plan.run_id, refresh_state=False) + + assert observation.outcome is BenchmarkOutcome.SUCCEEDED + assert observation.measurements is None + + +def test_benchmark_observer_reads_complete_measurements_from_ordinary_state( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + fake_slurm_runner: FakeSlurmRunner, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + ready_at = case.created_at + timedelta(minutes=2) + deployment = case.plan.deployments[0] + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=1, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.PENDING, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.PENDING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ), + ), + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=2, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.STARTING, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.STARTING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ), + ), + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id=case.attempt.attempt_id, + revision=3, + updated_at=ready_at, + started_at=case.attempt.created_at, + state=ReadinessState.READY, + deployments=( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.READY, + expected_backends=deployment.topology.replica_count, + ready_backends=deployment.topology.replica_count, + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=ProbeEvidence( + observed_at=ready_at, + outcome=ProbeOutcome.SUCCESS, + reason_code="backend_ready", + redacted_message="Backend is ready", + ), + ), + ), + ) + ) + _publish_winner_state(case) + SlurmCommandClient(fake_slurm_runner).submit("run.sbatch") + scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + fake_slurm_runner.set_task_state(scheduler, queue_state=None, accounting_state="COMPLETED") + + observation = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: case.created_at + timedelta(minutes=8), + ).observe(case.plan.run_id, refresh_state=True) + + assert observation.outcome is BenchmarkOutcome.SUCCEEDED + assert observation.measurements is not None + assert observation.measurements.actual_records == authored_run_single.invocation.num_records + assert observation.measurements.boot_seconds == 60 + assert observation.measurements.generation_seconds == 180 + assert observation.measurements.wall_seconds == 300 + + +def test_benchmark_observer_excludes_queue_stagger_from_cross_shard_timings( + tmp_path: Path, + authored_run: DataDesignerSlurmConfig, + multi_node_plan: ResolvedSlurmRunPlan, +) -> None: + cases = _initialized_cases(tmp_path, authored_run, multi_node_plan) + for index, case in enumerate(cases): + started_at = case.created_at + timedelta(minutes=2, hours=index) + _publish_ready_state(case, case.attempt, started_at=started_at) + _publish_winner_state( + case, + completed_at=started_at + timedelta(minutes=5), + stopped_at=started_at + timedelta(minutes=6), + ) + + observation = PersistedBenchmarkRunObserver( + cases[0].workspace, + _StaticSchedulerClient((), ()), + lambda: cases[-1].created_at + timedelta(hours=2), + ).observe(cases[0].plan.run_id, refresh_state=False) + + assert observation.measurements is not None + assert observation.measurements.actual_records == authored_run.invocation.num_records + assert observation.measurements.boot_seconds == 60 + assert observation.measurements.generation_seconds == 240 + assert observation.measurements.wall_seconds == 360 + + +def test_benchmark_observer_ignores_superseded_attempt_evidence_after_retry_wins( + tmp_path: Path, + authored_run_single: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, +) -> None: + case = _initialized_case(tmp_path, authored_run_single, single_node_plan) + case.writer.update_attempt( + _copy_attempt( + case.attempt, + state=AttemptLifecycleState.FAILED, + terminal_classification=AttemptTerminalClassification.PREEMPTED, + updated_at=case.attempt.created_at + timedelta(minutes=2), + ) + ) + retry = AttemptManifest( + schema_version=1, + run_id=case.plan.run_id, + shard_id=case.shard.shard_id, + attempt_id="attempt-0002", + attempt_ordinal=2, + resolved_plan=case.run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED, + scheduler=SchedulerIdentity(array_job_id=4102, array_task_id=0), + created_at=case.attempt.created_at + timedelta(minutes=3), + updated_at=case.attempt.created_at + timedelta(minutes=3), + ) + case.writer.create_attempt(retry) + _publish_ready_state(case, retry) + _, winner = _publish_winner_state(case, retry) + runner = FakeSlurmRunner( + arrays=( + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4101, array_task_id=0)),)), + FakeSlurmArray(tasks=(FakeSlurmTask(SchedulerIdentity(array_job_id=4102, array_task_id=0)),)), + ) + ) + observer = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(runner), + lambda: winner.published_at + timedelta(minutes=1), + ) + + persisted = observer.observe(case.plan.run_id, refresh_state=False) + first_scheduler = cast(SchedulerIdentity, case.attempt.scheduler) + retry_scheduler = cast(SchedulerIdentity, retry.scheduler) + SlurmCommandClient(runner).submit("first.sbatch") + SlurmCommandClient(runner).submit("retry.sbatch") + runner.set_task_state(first_scheduler, queue_state=None, accounting_state="PREEMPTED") + runner.set_task_state(retry_scheduler, queue_state=None, accounting_state="COMPLETED") + refreshed = observer.observe(case.plan.run_id, refresh_state=True) + + assert persisted.outcome is BenchmarkOutcome.SUCCEEDED + assert persisted.measurements is not None + assert persisted.measurements.boot_seconds == 60 + assert refreshed.outcome is BenchmarkOutcome.SUCCEEDED + assert refreshed.measurements == persisted.measurements + + def test_refresh_uses_terminal_accounting_over_stale_active_queue_state( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -352,6 +645,15 @@ def test_refresh_rejects_winner_that_conflicts_with_terminal_scheduler_evidence( SlurmCommandClient(fake_slurm_runner), ).refresh(observed_at=winner.published_at + timedelta(minutes=1)) + observer = PersistedBenchmarkRunObserver( + case.workspace, + SlurmCommandClient(fake_slurm_runner), + lambda: winner.published_at + timedelta(minutes=1), + ) + with pytest.raises(BenchmarkObservationFailure) as conflict: + observer.observe(case.plan.run_id, refresh_state=True) + assert conflict.value.outcome is BenchmarkOutcome.SCHEDULER_INCONSISTENT + def test_refresh_reports_a_validated_winner_as_succeeded( tmp_path: Path, @@ -589,6 +891,18 @@ def _initialized_case( *, submitted: bool = True, ) -> _ReconciliationCase: + cases = _initialized_cases(tmp_path, authored_config, plan, submitted=submitted) + assert len(cases) == 1 + return cases[0] + + +def _initialized_cases( + tmp_path: Path, + authored_config: DataDesignerSlurmConfig, + plan: ResolvedSlurmRunPlan, + *, + submitted: bool = True, +) -> tuple[_ReconciliationCase, ...]: workspace = tmp_path / "workspace" workspace.mkdir() relocated_plan = _relocate_plan(plan, workspace) @@ -603,35 +917,44 @@ def _initialized_case( path=(run_root / "resolved-plan.json").as_posix(), sha256=relocated_plan.compute_sha256(), ), - shard_count=1, + shard_count=len(relocated_plan.shards), ) - planned_shard = relocated_plan.shards[0] - shard = ShardManifest( - schema_version=1, - run_id=relocated_plan.run_id, - shard_id=planned_shard.shard_id, - shard_index=planned_shard.shard_index, - record_range=planned_shard.record_range, - input_partition=planned_shard.input_partition, - resume_workspace=planned_shard.resume_workspace, - created_at=created_at, + shards = tuple( + ShardManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=planned_shard.shard_id, + shard_index=planned_shard.shard_index, + record_range=planned_shard.record_range, + input_partition=planned_shard.input_partition, + resume_workspace=planned_shard.resume_workspace, + created_at=created_at, + ) + for planned_shard in relocated_plan.shards ) writer = SlurmStateWriter(workspace, relocated_plan.run_id) - writer.initialize_run(authored_config, relocated_plan, run, (shard,)) - attempt = AttemptManifest( - schema_version=1, - run_id=relocated_plan.run_id, - shard_id=shard.shard_id, - attempt_id="attempt-0001", - attempt_ordinal=1, - resolved_plan=run.resolved_plan, - state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, - scheduler=SchedulerIdentity(array_job_id=4101, array_task_id=0) if submitted else None, - created_at=created_at + timedelta(minutes=1), - updated_at=created_at + timedelta(minutes=2), + writer.initialize_run(authored_config, relocated_plan, run, shards) + attempts = tuple( + AttemptManifest( + schema_version=1, + run_id=relocated_plan.run_id, + shard_id=shard.shard_id, + attempt_id="attempt-0001", + attempt_ordinal=1, + resolved_plan=run.resolved_plan, + state=AttemptLifecycleState.SUBMITTED if submitted else AttemptLifecycleState.CREATED, + scheduler=(SchedulerIdentity(array_job_id=4101, array_task_id=shard.shard_index) if submitted else None), + created_at=created_at + timedelta(minutes=1), + updated_at=created_at + timedelta(minutes=2), + ) + for shard in shards + ) + for attempt in attempts: + writer.create_attempt(attempt) + return tuple( + _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, writer, created_at) + for shard, attempt in zip(shards, attempts, strict=True) ) - writer.create_attempt(attempt) - return _ReconciliationCase(workspace, relocated_plan, run, shard, attempt, writer, created_at) def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurmRunPlan: @@ -646,23 +969,36 @@ def _relocate_plan(plan: ResolvedSlurmRunPlan, workspace: Path) -> ResolvedSlurm return ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) -def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, ShardWinner]: +def _publish_winner_state( + case: _ReconciliationCase, + attempt: AttemptManifest | None = None, + *, + completed_at: datetime | None = None, + stopped_at: datetime | None = None, +) -> tuple[AttemptManifest, ShardWinner]: + attempt = case.attempt if attempt is None else attempt + result_completed_at = completed_at or attempt.created_at + timedelta(minutes=4) + attempt_stopped_at = stopped_at or attempt.created_at + timedelta(minutes=5) running = _copy_attempt( - case.attempt, + attempt, state=AttemptLifecycleState.RUNNING, - updated_at=case.created_at + timedelta(minutes=3), + updated_at=attempt.created_at + timedelta(minutes=2), ) case.writer.update_attempt(running) - candidate_path = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/output-manifest.json" + candidate_path = ( + case.writer.run_root / f"shards/{attempt.shard_id}/attempts/{attempt.attempt_id}/output-manifest.json" + ) dataset_path = candidate_path.parent / "dataset" - requested = case.plan.shards[0].requested_records + requested = next(shard.requested_records for shard in case.plan.shards if shard.shard_id == case.shard.shard_id) candidate = CandidateOutputManifest( schema_version=1, run_id=case.plan.run_id, shard_id=running.shard_id, attempt_id=running.attempt_id, attempt_ordinal=running.attempt_ordinal, - created_at=case.created_at + timedelta(minutes=4), + created_at=( + result_completed_at - timedelta(minutes=1) if completed_at else attempt.created_at + timedelta(minutes=3) + ), dataset_path=dataset_path.as_posix(), requested_records=requested, actual_records=requested, @@ -684,7 +1020,7 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S run_id=case.plan.run_id, shard_id=running.shard_id, attempt_id=running.attempt_id, - completed_at=case.created_at + timedelta(minutes=5), + completed_at=result_completed_at, requested_records=requested, actual_records=requested, outcome=ClientOutcome.COMPLETE, @@ -700,7 +1036,7 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S state=AttemptLifecycleState.SUCCEEDED, terminal_classification=AttemptTerminalClassification.SUCCEEDED, candidate_output=candidate_reference, - updated_at=case.created_at + timedelta(minutes=6), + updated_at=attempt_stopped_at, ) case.writer.update_attempt(completed) winner = ShardWinner( @@ -710,14 +1046,87 @@ def _publish_winner_state(case: _ReconciliationCase) -> tuple[AttemptManifest, S attempt_id=completed.attempt_id, attempt_ordinal=completed.attempt_ordinal, candidate_manifest=candidate_reference, - published_at=case.created_at + timedelta(minutes=7), + published_at=attempt_stopped_at + timedelta(minutes=1), ) - winner_path = case.writer.run_root / "shards/shard-00000/winner.json" + winner_path = case.writer.run_root / f"shards/{attempt.shard_id}/winner.json" winner_path.write_text(winner.serialize_json()) winner_path.chmod(0o600) return completed, winner +def _publish_ready_state( + case: _ReconciliationCase, + attempt: AttemptManifest, + *, + started_at: datetime | None = None, +) -> None: + started_at = started_at or attempt.created_at + timedelta(minutes=2) + ready_at = started_at + timedelta(minutes=1) + pending = AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=1, + updated_at=started_at, + started_at=started_at, + state=ReadinessState.PENDING, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.PENDING, + expected_backends=deployment.topology.replica_count, + ready_backends=0, + endpoint_publication=EndpointPublicationState.PENDING, + ) + for deployment in case.plan.deployments + ), + ) + case.writer.write_readiness(pending) + case.writer.write_readiness( + pending.model_copy( + update={ + "revision": 2, + "state": ReadinessState.STARTING, + "deployments": tuple( + deployment.model_copy(update={"state": ReadinessState.STARTING}) + for deployment in pending.deployments + ), + } + ) + ) + case.writer.write_readiness( + AttemptReadiness( + schema_version=1, + run_id=case.plan.run_id, + shard_id=attempt.shard_id, + attempt_id=attempt.attempt_id, + revision=3, + updated_at=ready_at, + started_at=started_at, + state=ReadinessState.READY, + deployments=tuple( + DeploymentReadiness( + deployment_id=deployment.deployment_id, + model_alias=deployment.authored.model_alias, + state=ReadinessState.READY, + expected_backends=deployment.topology.replica_count, + ready_backends=deployment.topology.replica_count, + endpoint_publication=EndpointPublicationState.PUBLISHED, + last_probe=ProbeEvidence( + observed_at=ready_at, + outcome=ProbeOutcome.SUCCESS, + reason_code="backend_ready", + redacted_message="Backend is ready", + ), + ) + for deployment in case.plan.deployments + ), + ) + ) + + def _copy_attempt(attempt: AttemptManifest, **updates: object) -> AttemptManifest: payload = attempt.model_dump(mode="json") payload.update(updates) diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index aad4367ee..dfa23f957 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -10,7 +10,9 @@ from click.testing import CliRunner import data_designer.slurm.cli as cli_module -from data_designer.slurm.config import DataDesignerSlurmConfig +import data_designer.slurm.cli_benchmark as benchmark_cli_module +from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig from data_designer.slurm.services import ( SlurmCollectionExecution, SlurmRetryExecution, @@ -83,6 +85,22 @@ def collect( ) +class _BenchmarkService: + def __init__(self, manifest: BenchmarkManifest, report: BenchmarkReport) -> None: + self.manifest = manifest + self.report = report + self.run_calls = [] + self.analysis_calls = [] + + def run(self, config, *, source_root, force): + self.run_calls.append((config, source_root, force)) + return self.manifest + + def analyze(self, benchmark_id, *, refresh_state, fail_if_incomplete): + self.analysis_calls.append((benchmark_id, refresh_state, fail_if_incomplete)) + return self.report + + def test_execute_emits_deterministic_json_and_forwards_actions( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, @@ -107,6 +125,35 @@ def test_execute_emits_deterministic_json_and_forwards_actions( assert service.calls == [(authored_run_single, tmp_path, True, True)] +def test_benchmark_cli_forwards_run_and_analysis_actions( + tmp_path: Path, + benchmark_config: DataDesignerSlurmBenchmarkConfig, + benchmark_manifest: BenchmarkManifest, + benchmark_report: BenchmarkReport, + monkeypatch, +) -> None: + benchmark_file = tmp_path / "benchmark.json" + benchmark_file.write_text(benchmark_config.serialize_json()) + service = _BenchmarkService(benchmark_manifest, benchmark_report) + monkeypatch.setattr(benchmark_cli_module, "create_slurm_benchmark_service", lambda **_: service) + + run_result = CliRunner().invoke( + cli_module.create_cli(), + ["benchmark", "run", str(benchmark_file), "--force"], + ) + analyze_result = CliRunner().invoke( + cli_module.create_cli(), + ["benchmark", "analyze", "/workspace/benchmarks/benchmark-001", "--refresh", "--fail-if-incomplete"], + ) + + assert run_result.exit_code == 0 + assert json.loads(run_result.stdout)["benchmark_id"] == benchmark_manifest.benchmark_id + assert analyze_result.exit_code == 0 + assert json.loads(analyze_result.stdout)["analysis_id"] == benchmark_report.analysis_id + assert service.run_calls == [(benchmark_config, tmp_path, True)] + assert service.analysis_calls == [("benchmark-001", True, True)] + + def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: service = _RunService() monkeypatch.setattr(cli_module, "create_slurm_run_service", lambda **_: service) @@ -383,11 +430,11 @@ def test_image_add_rejects_credential_bearing_oci_source() -> None: assert "secret" not in result.stderr -def test_cli_exposes_m3c_run_commands() -> None: +def test_cli_exposes_m3_public_commands() -> None: result = CliRunner().invoke(cli_module.create_cli(), ["--help"]) assert result.exit_code == 0 assert all( - command in result.stdout for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile") + command in result.stdout + for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile", "benchmark") ) - assert "benchmark" not in result.stdout diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index b5f54d241..82a7cc4b5 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -135,10 +135,13 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non statement += f""" slurm_help_result = CliRunner().invoke(app, ["slurm", "--help"]) assert slurm_help_result.exit_code == 0, (slurm_help_result.output, repr(slurm_help_result.exception)) -assert all(command in slurm_help_result.output for command in ("execute", "status", "cancel", "retry", "merge", "image")) -assert "benchmark" not in slurm_help_result.output +assert all( + command in slurm_help_result.output + for command in ("execute", "status", "cancel", "retry", "merge", "image", "profile", "benchmark") +) assert "data_designer.slurm.cli" in sys.modules assert version("data-designer-slurm") == {version!r} +from data_designer.slurm.benchmark import BenchmarkCompiler from data_designer.slurm.contracts import ArtifactReference as ContractArtifactReference from data_designer.slurm.contracts import RecordRange as ContractRecordRange from data_designer.slurm.contracts import ResumeWorkspace as ContractResumeWorkspace @@ -160,6 +163,9 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non SlurmRetryCoordinator, SlurmStateReconciler, ) +from data_designer.slurm.services import create_slurm_benchmark_service +assert BenchmarkCompiler.__name__ == "BenchmarkCompiler" +assert callable(create_slurm_benchmark_service) assert CollectionResult.__name__ == "CollectionResult" assert RetryPlan.__name__ == "RetryPlan" assert RunManifest.__name__ == "RunManifest" @@ -252,7 +258,7 @@ def main() -> None: assert str(leaf_base_requirement.specifier) == f"=={version}" assert leaf_packaging_requirement.specifier == Requirement("packaging>=25,<27").specifier assert leaf_pip_requirement.specifier == Requirement("pip>=25,<27").specifier - assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.9.2,<3").specifier + assert leaf_pydantic_requirement.specifier == Requirement("pydantic>=2.12,<3").specifier assert leaf_pyyaml_requirement.specifier == Requirement("pyyaml>=6.0.1,<7").specifier assert base_leaf_requirement.marker is not None assert base_leaf_requirement.marker.evaluate({"extra": "slurm"}) diff --git a/uv.lock b/uv.lock index 013aff91d..84e85921e 100644 --- a/uv.lock +++ b/uv.lock @@ -963,7 +963,7 @@ requires-dist = [ { name = "data-designer", editable = "packages/data-designer" }, { name = "packaging", specifier = ">=25,<27" }, { name = "pip", specifier = ">=25,<27" }, - { name = "pydantic", specifier = ">=2.9.2,<3" }, + { name = "pydantic", specifier = ">=2.12,<3" }, { name = "pyyaml", specifier = ">=6.0.1,<7" }, ] From 7b55030c0ed4fa1cf25d918d3c9efebdd51ff910 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 15 Sep 2026 10:50:49 -0300 Subject: [PATCH 28/62] fix(slurm): use allocation-local scratch for runtime (#938) * fix: use allocation-local scratch for Slurm runtime * fix: propagate staged runtime to Slurm endpoints * fix: launch retry identity through Slurm * fix: expose staged runtime to retries * fix: isolate retry verification state * fix: preserve logical retry plan references --- .../data_designer/slurm/client/environment.py | 44 +++- .../src/data_designer/slurm/client/worker.py | 5 + .../data_designer/slurm/config/profiles.py | 9 + .../data_designer/slurm/images/lifecycle.py | 23 +- .../slurm/launcher/collection.py | 10 +- .../data_designer/slurm/launcher/renderer.py | 69 ++++-- .../data_designer/slurm/runtime/bootstrap.py | 32 ++- .../src/data_designer/slurm/runtime/bundle.py | 2 +- .../data_designer/slurm/runtime/cleanup.sh | 6 +- .../data_designer/slurm/runtime/controller.py | 1 + .../data_designer/slurm/runtime/entrypoint.sh | 16 +- .../src/data_designer/slurm/runtime/paths.py | 4 +- .../slurm/runtime/plan_reader.sh | 7 +- .../data_designer/slurm/runtime/scratch.sh | 166 +++++++++++++ .../slurm/runtime/step_runner.sh | 8 +- .../src/data_designer/slurm/runtime/steps.py | 69 +++++- .../slurm/state/retry_storage.py | 3 +- .../src/data_designer/slurm/state/storage.py | 8 - .../src/data_designer/slurm/state/store.py | 2 - .../tests/client/conftest.py | 11 +- .../tests/client/test_environment.py | 46 +++- .../tests/client/test_worker.py | 5 +- .../tests/contracts/test_profiles.py | 9 + .../tests/images/test_lifecycle.py | 14 +- .../tests/launcher/test_collection.py | 20 +- .../tests/launcher/test_renderer.py | 30 ++- .../tests/runtime/test_bootstrap.py | 33 ++- .../tests/runtime/test_bundle.py | 10 +- .../tests/runtime/test_entrypoint.py | 3 + .../tests/runtime/test_shell_runtime.py | 232 +++++++++++++++++- .../tests/runtime/test_steps.py | 41 +++- .../golden/rendered/multi_node.sbatch | 13 +- .../golden/rendered/single_node.sbatch | 13 +- .../slurm_test_fakes/test_rendered_scripts.py | 7 +- .../tests/state/test_retry_collection.py | 11 +- .../tests/state/test_store.py | 3 +- 36 files changed, 860 insertions(+), 125 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/runtime/scratch.sh 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 index 2d7780948..7db43dba8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -9,6 +9,7 @@ import os import posixpath import re +import stat import subprocess import sys from collections.abc import Callable @@ -34,6 +35,7 @@ class PreparedClientEnvironment: shard_id: str attempt_id: str attempt_dir: Path + scratch_root: Path overlay_path: Path dependency_lock: ArtifactReference client_image_sha256: str @@ -49,6 +51,7 @@ class _BootstrapInputs: shard_id: str attempt_id: str attempt_dir: Path + scratch_root: Path client_image_sha256: str python_abi: str installer_path: Path @@ -64,7 +67,7 @@ class _VerifiedDependencies: class ClientEnvironmentBuilder: - """Prepare one immutable attempt-local package environment.""" + """Prepare one immutable allocation-local package environment.""" def __init__( self, @@ -82,6 +85,7 @@ def prepare( 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( @@ -89,6 +93,7 @@ def prepare( 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) @@ -97,6 +102,7 @@ def prepare( 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_image_sha256=inputs.client_image_sha256, @@ -112,9 +118,11 @@ def _load_bootstrap_inputs( 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: @@ -159,6 +167,7 @@ def _load_bootstrap_inputs( shard_id=shard_id, attempt_id=attempt_id, attempt_dir=attempt_dir, + scratch_root=scratch_root, client_image_sha256=image_sha256, python_abi=python_abi, installer_path=installer_path, @@ -217,16 +226,7 @@ def _prepare_overlay( dependencies.image_distributions, inputs.plan, ) - logical_overlay = ( - inputs.logical_run_root - / "shards" - / inputs.shard_id - / "attempts" - / inputs.attempt_id - / "client-env" - / "site-packages" - ) - overlay_path = Path(_get_container_path(inputs.plan, logical_overlay.as_posix(), require_writable=True)) + overlay_path = inputs.scratch_root / "client-env" / "site-packages" outcome = self._install_overlay(inputs.installer_path, wheels, expected_overlay, overlay_path) installed = tuple(sorted((*dependencies.image_distributions, *expected_overlay), key=lambda item: item.name)) return overlay_path, outcome, installed @@ -296,8 +296,8 @@ def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ... def activate_environment(prepared: PreparedClientEnvironment) -> None: """Activate one verified overlay and isolate Data Designer attempt state.""" - home = prepared.attempt_dir / "data-designer-home" - cache = prepared.attempt_dir / "cache" + 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() @@ -314,6 +314,24 @@ def _validate_input_path(path: Path, label: str) -> None: 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) 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 index d34fb5dcf..87a8f9ca7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/worker.py @@ -5,6 +5,7 @@ import argparse import importlib +import os import signal import sys from collections.abc import Sequence @@ -35,11 +36,15 @@ def main(argv: Sequence[str] | None = None) -> int: 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") 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 index 5567850f0..e1530eec8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -20,6 +20,8 @@ validate_plain_text, ) +_ALLOCATION_SCRATCH_CONTAINER_ROOT = "/run/data-designer-slurm" + class GpuRequestMode(str, Enum): GRES = "gres" @@ -98,6 +100,13 @@ 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 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 index ddfa231c1..a14eab697 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -41,6 +41,7 @@ _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 = (4, 0) _MINIMUM_ENROOT_SQSH_VERSION = (3, 5) @@ -147,6 +148,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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}"), @@ -208,6 +210,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: set -Eeuo pipefail 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)} @@ -237,9 +240,14 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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 @@ -271,18 +279,9 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: verify_sha256 "${{DD_INSPECTOR_SHA256}}" "${{DD_INSPECTOR}}" verify_sha256 "${{DD_ENROOT_RC_SHA256}}" "${{DD_ENROOT_RC}}" -install -d -m 0700 \ - "${{DD_JOB_DIR}}/home" \ - "${{DD_JOB_DIR}}/enroot/cache" \ - "${{DD_JOB_DIR}}/enroot/config" \ - "${{DD_JOB_DIR}}/enroot/data" \ - "${{DD_JOB_DIR}}/enroot/tmp" \ - "${{DD_INSPECTION_DIRECTORY}}" -export HOME="${{DD_JOB_DIR}}/home" -export ENROOT_CACHE_PATH="${{DD_JOB_DIR}}/enroot/cache" -export ENROOT_CONFIG_PATH="${{DD_JOB_DIR}}/enroot/config" -export ENROOT_DATA_PATH="${{DD_JOB_DIR}}/enroot/data" -export ENROOT_TEMP_PATH="${{DD_JOB_DIR}}/enroot/tmp" +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 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 index 5072d14d9..abe663892 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -6,6 +6,7 @@ from __future__ import annotations import posixpath +from importlib import resources from pathlib import PurePosixPath from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives @@ -55,11 +56,13 @@ def render_collection_script( ("DD_STATE_MOUNT", workspace_root, workspace_root), ("DD_OUTPUT_MOUNT", destination.mount.source, destination.mount.target), ) + scratch_source = resources.files("data_designer.slurm.runtime").joinpath("scratch.sh").read_text() return f"""#!/usr/bin/env bash {directives} set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" +{scratch_source} readonly DD_CLIENT_IMAGE={quote_shell_value(resolved_plan.client.image.path)} readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(resolved_plan.client.image.sha256)} readonly DD_COLLECTION_PLAN={quote_shell_value(collection_plan_path)} @@ -78,9 +81,14 @@ def render_collection_script( verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" verify_sha256 "${{DD_COLLECTION_PLAN_SHA256}}" "${{DD_COLLECTION_PLAN}}" +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 +export HOME="${{DD_SCRATCH_ROOT}}/home" DD_ENROOT_MOUNTS=({mount_arguments}) readonly DD_ENROOT_MOUNTS -exec enroot start --root "${{DD_ENROOT_MOUNTS[@]}}" "${{DD_CLIENT_IMAGE}}" \ +enroot start --root "${{DD_ENROOT_MOUNTS[@]}}" "${{DD_CLIENT_IMAGE}}" \ python -m data_designer.slurm.state.collection_worker \ --workspace-root "${{DD_WORKSPACE_ROOT}}" --run-id "${{DD_RUN_ID}}" --collection-id "${{DD_COLLECTION_ID}}" """ 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 index efc16549a..b1a7543a0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -23,10 +23,16 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi run_root = posixpath.dirname(plan.authored_config.path) plan_path = posixpath.join(run_root, "resolved-plan.json") - directive_text = render_batch_directives(_build_generation_directives(plan)) + 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 plan.selected_profile.profile.gpu_request_mode == "visible": directive_text = f"{directive_text}\n#SBATCH --exclusive" - attempt = f"{attempt_ordinal:04d}" return f"""#!/usr/bin/env bash {directive_text} set -Eeuo pipefail @@ -55,9 +61,13 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi 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_RUNTIME_ROOT="${{DD_ATTEMPT_DIR}}/runtime" -[[ -d ${{DD_RUNTIME_ROOT}} && ! -L ${{DD_RUNTIME_ROOT}} ]] -tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_ROOT}}" +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}}" @@ -85,8 +95,15 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) 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) + _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"), + ) ) attempt_cases = "\n".join( f" {shard.array_task_index}) DD_ATTEMPT_ORDINAL={quote_shell_value(f'{shard.attempt_ordinal:04d}')} ;;" @@ -104,8 +121,8 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) 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)} -readonly DD_CLIENT_IMAGE={quote_shell_value(plan.client.image.path)} -readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(plan.client.image.sha256)} +readonly DD_VERIFIED_CLIENT_IMAGE={quote_shell_value(plan.client.image.path)} +readonly DD_VERIFIED_CLIENT_IMAGE_SHA256={quote_shell_value(plan.client.image.sha256)} 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)} @@ -118,7 +135,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" verify_sha256 "${{DD_PLAN_SHA256}}" "${{DD_PLAN}}" -verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" +verify_sha256 "${{DD_VERIFIED_CLIENT_IMAGE_SHA256}}" "${{DD_VERIFIED_CLIENT_IMAGE}}" if [[ ! ${{SLURM_ARRAY_TASK_ID:-}} =~ ^[0-9]+$ ]]; then printf '%s\\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 exit 64 @@ -147,16 +164,33 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) fi readonly DD_ARRAY_JOB_ID="${{SLURM_ARRAY_JOB_ID}}" readonly DD_ATTEMPT_ID="attempt-${{DD_ATTEMPT_ORDINAL}}" -enroot start --root --mount "${{DD_WORKSPACE_ROOT}}:${{DD_WORKSPACE_ROOT}}" "${{DD_CLIENT_IMAGE}}" \\ - python -m data_designer.slurm.state.attempt_identity \\ +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_SCRATCH_CONTAINER_ROOT}}/runtime" +srun \\ + --nodes=1 \\ + --ntasks=1 \\ + --cpus-per-task=1 \\ + --exact \\ + --overlap \\ + --unbuffered \\ + --gres=none \\ + --export=ALL \\ + --container-image="${{DD_VERIFIED_CLIENT_IMAGE}}" \\ + --container-mounts="${{DD_WORKSPACE_ROOT}}:${{DD_WORKSPACE_ROOT}},${{DD_SCRATCH_ROOT}}:${{DD_SCRATCH_CONTAINER_ROOT}}" \\ + --container-env=PYTHONPATH \\ + -- \\ + python3 -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}}" -DD_RUNTIME_DIR="$(mktemp -d "${{DD_ATTEMPT_DIR}}/runtime.${{DD_RUNTIME_SHA256}}.XXXXXX")" -readonly DD_RUNTIME_DIR -tar -xzf "${{DD_RUNTIME_ARCHIVE}}" -C "${{DD_RUNTIME_DIR}}" -source "${{DD_RUNTIME_DIR}}/entrypoint.sh" +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}}" @@ -173,6 +207,8 @@ def _build_generation_directives( *, 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, @@ -193,6 +229,9 @@ def _build_generation_directives( ("cpus-per-task", str(plan.client.authored.cpus)), ("time", plan.submission.time_limit), ("array", resolved_array), + ("chdir", posixpath.dirname(plan.authored_config.path)), + ("output", output_path), + ("error", error_path), ] profile = plan.selected_profile.profile if profile.gpu_request_mode == "gres": 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 index 2f425b3f7..10449cba6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -12,9 +12,10 @@ 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 get_container_path +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 @@ -33,7 +34,12 @@ def build_runtime_manifest( """Build the secret-free command handoff for the Bash controller.""" plan = context.plan validate_allocation_layout(plan, layout) - runtime_container_root = get_container_path(plan, runtime_root.as_posix(), require_writable=True) + runtime_container_root = Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime" + if runtime_root != runtime_container_root: + raise SlurmRuntimeError( + SlurmRuntimeErrorCode.INVALID_CONTEXT, + "runtime root is outside allocation-local scratch", + ) deployments = resolve_allocation_deployments(context, environment) endpoints = tuple( RuntimeEndpoint( @@ -52,12 +58,12 @@ def build_runtime_manifest( context, environment, endpoints, - runtime_container_root, + runtime_container_root.as_posix(), log_directory, layout, ) ] - steps.extend(build_server_steps(deployments, context, runtime_container_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 ) @@ -69,7 +75,7 @@ def build_runtime_manifest( context, environment, endpoints, - runtime_container_root, + runtime_container_root.as_posix(), log_directory, layout, ) @@ -157,10 +163,17 @@ def _build_client_step( command=command, cpus=plan.client.authored.cpus, gpu_indices=(), - literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_container_root, **allocation_environment}, + literal_environment={ + "DATA_DESIGNER_SLURM_SCRATCH_ROOT": ALLOCATION_SCRATCH_CONTAINER_ROOT, + "LC_ALL": "C", + "PYTHONPATH": runtime_container_root, + **allocation_environment, + }, secret_environment={name: name for name in secret_names}, environment_prefixes={}, - container_environment=tuple(sorted((*secret_names, *allocation_environment, "PYTHONPATH"))), + container_environment=tuple( + sorted((*secret_names, *allocation_environment, "DATA_DESIGNER_SLURM_SCRATCH_ROOT", "PYTHONPATH")) + ), log_directory=log_directory, node_hosts=(layout.get_host(plan.client.host_node_index),), ) @@ -185,6 +198,7 @@ def _build_endpoint_step( proxy_path, deployment.logical_endpoint.port, backend_hosts=backend_hosts, + runtime_proxy_is_container_path=True, ) return _step( step_id=f"{deployment.deployment_id}-endpoint", @@ -193,10 +207,10 @@ def _build_endpoint_step( command=command, cpus=context.plan.client.authored.cpus, gpu_indices=(), - literal_environment={"LC_ALL": "C"}, + literal_environment={"LC_ALL": "C", "PYTHONPATH": runtime_root.as_posix()}, secret_environment={}, environment_prefixes={}, - container_environment=(), + container_environment=("PYTHONPATH",), log_directory=log_directory, node_hosts=(layout.get_host(context.plan.client.host_node_index),), readiness=( 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 index 7cf08db2f..8e70ffc14 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bundle.py @@ -25,7 +25,7 @@ _SOURCE_MODE = 0o400 _MAXIMUM_SOURCE_SIZE = 16 * 1024 * 1024 _TEMPORARY_NAME_PATTERN = re.compile(r"^\.runtime\.[0-9a-f]{16}\.tmp$") -_SHELL_SOURCE_NAMES = ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh") +_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" 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 index c545d4e85..c5e43336a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/cleanup.sh @@ -10,7 +10,11 @@ dd_start_runtime_timer() { } dd_sleep() { - read -r -t "$1" -u "${DD_RUNTIME_TIMER[0]}" _ || true + 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() { diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 9f83c0c92..625050f2e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -284,6 +284,7 @@ def _prepare_runtime(self) -> _RuntimeTopology: self._context.attempt_directory, self._environment, self._runtime_proxy_path, + self._runtime_root, ) ) endpoints = tuple(endpoint for _, endpoint in endpoint_steps) 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 index 4015b4f34..327214f77 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -4,6 +4,7 @@ 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" @@ -34,17 +35,16 @@ dd_slurm_run_allocation() { DD_PLAN_CONTAINER_PATH=${DD_CONTAINER_PATH} dd_read_container_path "${DD_PLAN_PATH}" "${DD_ATTEMPT_PATH}" true DD_ATTEMPT_CONTAINER_DIR=${DD_CONTAINER_PATH} - dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_DIR}" true - DD_RUNTIME_CONTAINER_ROOT=${DD_CONTAINER_PATH} + DD_RUNTIME_CONTAINER_ROOT=${DD_SCRATCH_CONTAINER_ROOT}/runtime dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_MANIFEST}" true DD_RUNTIME_MANIFEST_CONTAINER_PATH=${DD_CONTAINER_PATH} export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST export DD_PLAN_CONTAINER_PATH DD_ATTEMPT_CONTAINER_DIR DD_RUNTIME_CONTAINER_ROOT - dd_verify_host_context - dd_start_runtime_timer trap dd_runtime_exit EXIT trap 'exit 130' INT TERM + dd_verify_host_context + dd_start_runtime_timer DD_RUNTIME_PREPARED=1 local host @@ -53,7 +53,7 @@ dd_slurm_run_allocation() { host_arguments+=(--node-host "${host}") done dd_run_bound_control_phase prepare \ - --runtime-root "${DD_RUNTIME_DIR}" \ + --runtime-root "${DD_RUNTIME_CONTAINER_ROOT}" \ --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" \ "${host_arguments[@]}" dd_verify_runtime_manifest \ @@ -99,12 +99,15 @@ dd_run_bound_control_phase() { dd_verify_host_context() { local tool - for tool in bash sha256sum tar jq srun scontrol getent curl; do + 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} ]] [[ ${SLURM_ARRAY_TASK_ID:-} =~ ^[0-9]+$ ]] [[ ${DD_EXPECTED_NODES} =~ ^[1-9][0-9]*$ ]] @@ -237,6 +240,7 @@ dd_runtime_exit() { 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/paths.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py index 1274fe054..edf1593f3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/paths.py @@ -11,6 +11,8 @@ 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, @@ -82,4 +84,4 @@ def get_host_path( return host_path -__all__ = ["get_container_path", "get_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 index a16c54392..0564e6a00 100644 --- 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 @@ -21,7 +21,12 @@ dd_read_control_plan() { and ([.deployments[].node_indices[]] | length > 0) and ([.deployments[].node_indices[]] | all(type == "number" and . >= 0)) and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) - ' "${plan}" >/dev/null + 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_IMAGE=$(jq -er '.client.image.path' "${plan}") DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | tostring' "${plan}") 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/step_runner.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/step_runner.sh index 34e5c6aeb..22c19afde 100644 --- 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 @@ -67,6 +67,7 @@ dd_materialize_step_environment() { 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%,} @@ -106,7 +107,8 @@ dd_build_srun_command() { else DD_SRUN_COMMAND+=(--gres=none) fi - [[ -z ${DD_CONTAINER_MOUNTS} ]] || DD_SRUN_COMMAND+=("--container-mounts=${DD_CONTAINER_MOUNTS}") + [[ -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%,}") @@ -142,6 +144,7 @@ dd_run_step() { dd_run_control_phase() { local operation=$1 shift + local container_mounts="${DD_SCRATCH_ROOT}:${DD_SCRATCH_CONTAINER_ROOT}" local -a command=( srun --nodes=1 @@ -155,7 +158,8 @@ dd_run_control_phase() { "--container-image=${DD_CLIENT_IMAGE}" --gres=none ) - [[ -z ${DD_CONTAINER_MOUNTS} ]] || command+=("--container-mounts=${DD_CONTAINER_MOUNTS}") + [[ -z ${DD_CONTAINER_MOUNTS} ]] || container_mounts="${DD_CONTAINER_MOUNTS},${container_mounts}" + command+=("--container-mounts=${container_mounts}") command+=( --container-env=PYTHONPATH,SLURM_JOB_GPUS -- 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 index 555c0a5db..7492ad685 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -15,6 +15,7 @@ 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, @@ -22,12 +23,13 @@ ) 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 get_container_path +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", @@ -159,7 +161,11 @@ def _build_step( image_path=plan.client.image.path, command=command, environment=environment, - container_environment=(*secret_names, *allocation_environment), + container_environment=( + *secret_names, + *allocation_environment, + "DATA_DESIGNER_SLURM_SCRATCH_ROOT", + ), plan=plan, attempt_directory=attempt_directory, ) @@ -211,6 +217,15 @@ def _build_client_environment( 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"] = ALLOCATION_SCRATCH_CONTAINER_ROOT return secret_names, environment @@ -234,11 +249,22 @@ def build_endpoint_steps( 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)) + steps.append( + _build_endpoint_step( + deployment, + plan, + attempt_directory, + source_environment, + runtime_proxy_path, + resolved_runtime_root, + ) + ) return tuple(steps) @@ -248,6 +274,7 @@ def _build_endpoint_step( 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, @@ -256,13 +283,15 @@ def _build_endpoint_step( port=deployment.logical_endpoint.port, ) command = build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) + environment = _base_environment(source_environment) + environment["PYTHONPATH"] = get_container_path(plan, runtime_root.as_posix()) step = _build_srun_step( step_id=f"{deployment.deployment_id}-endpoint", role=RuntimeStepRole.ENDPOINT, image_path=plan.client.image.path, command=command, - environment=_base_environment(source_environment), - container_environment=(), + environment=environment, + container_environment=("PYTHONPATH",), plan=plan, attempt_directory=attempt_directory, ) @@ -276,6 +305,7 @@ def build_endpoint_command( port: int, *, backend_hosts: tuple[str, ...] | None = None, + runtime_proxy_is_container_path: bool = False, ) -> tuple[str, ...]: selected_hosts = backend_hosts or ("127.0.0.1",) * len(deployment.backend_endpoints) if len(selected_hosts) != len(deployment.backend_endpoints): @@ -294,9 +324,17 @@ def build_endpoint_command( ) 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: + container_proxy_path = ( + validate_absolute_path(runtime_proxy_path.as_posix()) + if runtime_proxy_is_container_path + else get_container_path(plan, runtime_proxy_path.as_posix()) + ) + except ValueError as error: + raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime proxy path is invalid") from error return ( "python3", - get_container_path(plan, runtime_proxy_path.as_posix()), + container_proxy_path, "--listen-port", str(port), *retry_arguments, @@ -425,7 +463,7 @@ def _build_srun_step( ) -> RuntimeStep: srun_command = _build_srun_prefix(plan, image_path) _add_srun_resources(srun_command, gpu_indices) - _add_srun_container_options(srun_command, plan, container_environment) + _add_srun_container_options(srun_command, plan, container_environment, environment) srun_command.extend(("--", *command)) log_root = attempt_directory / "logs" return RuntimeStep( @@ -468,8 +506,9 @@ def _add_srun_container_options( srun_command: list[str], plan: ResolvedSlurmRunPlan, container_environment: tuple[str, ...], + environment: Mapping[str, str], ) -> None: - mounts = _render_mounts(plan) + mounts = _render_mounts(plan, environment.get("DD_SCRATCH_ROOT")) if mounts: srun_command.append(f"--container-mounts={mounts}") if container_environment: @@ -480,7 +519,7 @@ def _get_gpu_mask(gpu_indices: tuple[int, ...]) -> int: return sum(1 << index for index in gpu_indices) -def _render_mounts(plan: ResolvedSlurmRunPlan) -> str: +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 (",", ":")): @@ -492,9 +531,21 @@ def _render_mounts(plan: ResolvedSlurmRunPlan) -> str: 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, 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 index cbb10803d..8847516e2 100644 --- 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 @@ -100,7 +100,8 @@ def get_plan_path(self, retry_id: Identifier) -> Path: def get_plan_reference(self, plan: RetryPlan) -> ArtifactReference: """Return the exact path and digest bound by retry status.""" - return ArtifactReference(path=self.get_plan_path(plan.retry_id).as_posix(), sha256=plan.compute_sha256()) + 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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py index 39e2e6f71..606aab299 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/storage.py @@ -50,7 +50,6 @@ _CANDIDATE_OUTPUT_FILENAME = "output-manifest.json" _WINNER_FILENAME = "winner.json" _DATASET_DIRECTORY_NAME = "dataset" -_RUNTIME_DIRECTORY_NAME = "runtime" _RESUME_LOCK_FILENAME = "resume.lock" _LOCK_DIRECTORY_NAME = ".locks" _SUBMISSION_LOCK_PREFIX = "submission" @@ -302,13 +301,6 @@ def publish_attempt(self, attempt: AttemptManifest) -> None: ) as attempt_descriptor: self._publish_immutable_record(attempt_descriptor, _ATTEMPT_FILENAME, attempt) - def ensure_runtime_directory(self, shard_id: ShardId, attempt_id: AttemptId) -> None: - attempt_root = self.get_attempt_path(shard_id, attempt_id) - runtime_root = attempt_root / _RUNTIME_DIRECTORY_NAME - with self.open_attempt_directory(shard_id, attempt_id) as attempt_descriptor: - ensure_private_child_directory(attempt_descriptor, _RUNTIME_DIRECTORY_NAME, runtime_root) - sync_directory(attempt_descriptor) - 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) 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 index 6ab827ffe..270264e99 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/store.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/store.py @@ -296,7 +296,6 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes if existing is not None: if existing != attempt: raise StateConflictError(f"attempt {attempt.attempt_id!r} already contains different state") - self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) self._storage.sync_attempt_directory(attempt.shard_id, attempt.attempt_id) return existing self._finalizer.require_no_winner(run, plan, shard, shard_attempts) @@ -306,7 +305,6 @@ def _create_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifes self._reader.validate_attempt_against_plan(run, plan, shard, attempt) validate_shard_attempt_set(run, shard, shard_attempts + (attempt,)) self._storage.publish_attempt(attempt) - self._storage.ensure_runtime_directory(attempt.shard_id, attempt.attempt_id) return attempt def _update_attempt_with_locks(self, attempt: AttemptManifest) -> AttemptManifest: diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index 95220553e..67bd3c517 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -29,6 +29,7 @@ class ClientWorkerCase: plan_path: Path lock: ResolvedDependencyLock attempt_dir: Path + scratch_root: Path prepared: PreparedClientEnvironment endpoints: dict[str, str] @@ -111,7 +112,7 @@ def allocation_environment(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.fixture -def client_worker_case(tmp_path: Path) -> ClientWorkerCase: +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()) @@ -139,13 +140,17 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: 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" - overlay_path = attempt_dir / "client-env" / "site-packages" + 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_image_sha256=plan.client.image.sha256, @@ -162,4 +167,4 @@ def client_worker_case(tmp_path: Path) -> ClientWorkerCase: 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, prepared, endpoints) + return ClientWorkerCase(plan, plan_path, lock, attempt_dir, scratch_root, prepared, endpoints) diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 51c920589..8013f5f35 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -16,7 +16,7 @@ import pytest from conftest import ClientWorkerCase, FakeDataDesigner -from data_designer.slurm.client.environment import ClientEnvironmentBuilder, inspect_distributions +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 @@ -35,10 +35,47 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: 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.image_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_client_runs_through_non_identity_workspace_mount(client_worker_case: ClientWorkerCase) -> None: @@ -66,6 +103,7 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: 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( @@ -127,6 +165,7 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: shard_id=shard_id, attempt_id="attempt-0001", attempt_dir=attempt_dir, + scratch_root=client_worker_case.scratch_root, ) assert prepared.attempt_dir == attempt_dir @@ -185,6 +224,7 @@ def test_environment_rejects_client_image_inventory_conflict(client_worker_case: 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 @@ -211,6 +251,7 @@ def test_environment_rejects_invalid_identity_before_creating_attempt( 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 @@ -244,6 +285,7 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: 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 @@ -283,6 +325,7 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: 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 @@ -323,6 +366,7 @@ def test_worker_bootstrap_loads_overlay_plugin_before_config_import( shard_id="shard-00000", attempt_id="attempt-0001", attempt_dir=Path(sys.argv[2]), + scratch_root=Path(sys.argv[1]).parent, overlay_path=Path(sys.argv[1]), dependency_lock=ArtifactReference(path=(Path(sys.argv[2]) / "dependency-lock.json").as_posix(), sha256="a" * 64), client_image_sha256="b" * 64, diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index bb5805573..a1544d121 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -448,12 +448,14 @@ def test_if_possible_interruption_preserves_workspace_for_retry(client_worker_ca resume_path = Path(plan.shards[0].resume_workspace.path) assert (resume_path / "batch.parquet").is_file() attempt_dir = resume_path.parent / "attempts" / "attempt-0002" - overlay_path = attempt_dir / "client-env" / "site-packages" + 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)) @@ -554,6 +556,7 @@ def test_preflight_rejects_plugin_secondary_model_alias( shard_id=plan["shards"][0]["shard_id"], attempt_id="attempt-0001", attempt_dir=attempt_dir, + scratch_root=Path(sys.argv[3]).parent, overlay_path=Path(sys.argv[3]), dependency_lock=ArtifactReference(**plan["client"]["dependency_lock"]), client_image_sha256=plan["client"]["image"]["sha256"], diff --git a/packages/data-designer-slurm/tests/contracts/test_profiles.py b/packages/data-designer-slurm/tests/contracts/test_profiles.py index 29a7f2dda..93e7aecd2 100644 --- a/packages/data-designer-slurm/tests/contracts/test_profiles.py +++ b/packages/data-designer-slurm/tests/contracts/test_profiles.py @@ -109,6 +109,15 @@ def test_explicit_selection_rejects_unknown_cluster(profile_catalog: SlurmProfil {"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( diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index f63b0268f..f16d18c09 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -135,10 +135,11 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ 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_JOB_DIR}/home"' in script - assert 'export ENROOT_CONFIG_PATH="${DD_JOB_DIR}/enroot/config"' 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_JOB_DIR}/home"') < script.rindex("\nverify_enroot_compatibility 4 0") + assert script.index('export HOME="${DD_SCRATCH_ROOT}/home"') < script.rindex("\nverify_enroot_compatibility 4 0") assert 'version="$(enroot version)"' in script assert 'verify_enroot_compatibility 4 0 "digest-pinned OCI imports"' in script assert f'--mount "{prepared.plan.job_directory}:' not in script @@ -212,9 +213,10 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( ) assert completed.returncode == 0, completed.stderr - expected_home = Path(prepared.plan.job_directory) / "home" - assert set(home_log.read_text().splitlines()) == {expected_home.as_posix()} - assert stat.S_IMODE(expected_home.stat().st_mode) == 0o700 + 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" if source_kind == "oci": assert enroot_log.read_text().splitlines() == ["version", "import", "create", "start", "remove"] diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 478e202d9..29ef1070c 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -82,6 +82,7 @@ def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( assert "data_designer.slurm.state.collection_worker" 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 @@ -118,10 +119,25 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( 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_DIR") + 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="${DD_VERIFIED_CLIENT_IMAGE}"' in script + assert "readonly DD_CLIENT_IMAGE=" not in script + assert ( + '--container-mounts="${DD_WORKSPACE_ROOT}:${DD_WORKSPACE_ROOT},${DD_SCRATCH_ROOT}:${DD_SCRATCH_CONTAINER_ROOT}"' + in script + ) + assert 'export PYTHONPATH="${DD_SCRATCH_CONTAINER_ROOT}/runtime"' in script + assert "--container-env=PYTHONPATH" in script + assert "--gres=none" in script + assert "--export=ALL" in script + assert "python3 -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("data_designer.slurm.state.attempt_identity") < script.index("DD_RUNTIME_DIR") + 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 diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index f4ed66d29..50dc506a1 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -32,8 +32,10 @@ def test_renderer_matches_contract_bound_goldens( request: pytest.FixtureRequest, ) -> None: plan = request.getfixturevalue(plan_fixture) + script = render_generation_attempt_script(plan, attempt_ordinal=1) - assert render_generation_attempt_script(plan, attempt_ordinal=1) == (GOLDEN_DIRECTORY / fixture_name).read_text() + 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: @@ -194,7 +196,7 @@ def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) - assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script - assert len(script.splitlines()) <= 42 + assert len(script.splitlines()) <= 48 assert script.endswith("\n") @@ -204,16 +206,29 @@ def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_run ) -> None: run_root = tmp_path / "run" run_root.mkdir() - (run_root / "shards/shard-00000/attempts/attempt-0001/runtime").mkdir(parents=True) + (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( @@ -224,6 +239,7 @@ def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_run 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" @@ -231,7 +247,13 @@ def test_rendered_script_verifies_exact_persisted_plan_bytes_before_sourcing_run 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"} + 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",), diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index c9cedca44..6e6d06153 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -6,14 +6,17 @@ 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 ResolvedSlurmRunPlan from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, 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.paths import ALLOCATION_SCRATCH_CONTAINER_ROOT 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 @@ -22,7 +25,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: context = runtime_case.context - runtime_root = context.attempt_directory / "runtime" + runtime_root = Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime" log_directory = context.attempt_directory / "logs/execution-00000002" manifest = build_runtime_manifest( @@ -55,6 +58,28 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru 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) + 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 == ("PYTHONPATH",) + 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"] == ALLOCATION_SCRATCH_CONTAINER_ROOT + for step in client_steps + ) + assert "SLURM_TMPDIR" not in manifest.serialize_json() + + +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: @@ -80,7 +105,7 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti manifest = build_runtime_manifest( retry_context, {"SLURM_JOB_GPUS": "0"}, - runtime_root=context.attempt_directory / "runtime", + runtime_root=Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime", log_directory=context.attempt_directory / "logs/execution-00000002", layout=AllocationLayout(("compute-001",)), ) @@ -118,7 +143,7 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( manifest = build_runtime_manifest( context, {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, - runtime_root=context.attempt_directory / "runtime", + runtime_root=Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime", log_directory=context.attempt_directory / "logs/execution-00000002", layout=layout, ) @@ -142,6 +167,8 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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"] == f"{ALLOCATION_SCRATCH_CONTAINER_ROOT}/runtime" + assert endpoint.container_environment == ("PYTHONPATH",) 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" diff --git a/packages/data-designer-slurm/tests/runtime/test_bundle.py b/packages/data-designer-slurm/tests/runtime/test_bundle.py index abfc4a885..ad8361a77 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bundle.py +++ b/packages/data-designer-slurm/tests/runtime/test_bundle.py @@ -35,10 +35,10 @@ def test_runtime_bundle_is_deterministic_content_addressed_and_restrictive(tmp_p 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[0] == "entrypoint.sh" - assert names[1:4] == ["plan_reader.sh", "step_runner.sh", "cleanup.sh"] - assert names[4] == "data_designer/slurm/runtime/slurm-sources.txt" - assert names[5] == "data_designer/slurm/__init__.py" + 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 @@ -67,7 +67,7 @@ def test_runtime_bundle_recursively_collects_and_imports_nested_packages( (source_root / "__init__.py").write_text("") (runtime_root / "__init__.py").write_text("") (runtime_root / "bundle.py").write_text("") - for name in ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh"): + 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") diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 3c27e87c9..5e01d6040 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -115,6 +115,9 @@ def test_client_phase_installs_and_runs_plugin_in_default_worker_process( 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) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 555b003d3..649af45dd 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -7,7 +7,9 @@ import json import os import shlex +import shutil import subprocess +import tarfile from pathlib import Path import pytest @@ -19,6 +21,9 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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" @@ -27,7 +32,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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_root, artifacts, gpu_request_mode) + plan = _plan(tmp_path, runtime_copy, artifacts, gpu_request_mode) 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))) @@ -42,9 +47,11 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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())} -source {shlex.quote((runtime_root / "entrypoint.sh").as_posix())} +source {shlex.quote((runtime_copy / "entrypoint.sh").as_posix())} dd_slurm_run_allocation "${{DD_PLAN}}" "${{DD_ATTEMPT_DIR}}" """ environment = { @@ -54,8 +61,10 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "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", @@ -72,6 +81,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( ) 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() @@ -79,13 +89,204 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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 ("entrypoint.sh", "plan_reader.sh", "step_runner.sh", "cleanup.sh")) + 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)), @@ -148,6 +349,8 @@ def test_step_runner_builds_one_coordinated_srun_across_selected_nodes() -> None 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* ]] @@ -156,6 +359,7 @@ def test_step_runner_builds_one_coordinated_srun_across_selected_nodes() -> None [[ $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) @@ -191,6 +395,28 @@ def test_shell_resolves_scheduler_hosts_in_planner_order(tmp_path: Path) -> None 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() diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index bda517d92..2e5f025e9 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -34,6 +34,7 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti 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", @@ -82,6 +83,14 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti 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 all_steps + ) server = server_steps[0] assert server.role is RuntimeStepRole.SERVER @@ -135,7 +144,7 @@ def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(run (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) @@ -146,7 +155,7 @@ def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(run context.attempt, context.attempt_directory, endpoints, - {}, + {"DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0"}, ) separator = step.command.index("--") @@ -163,6 +172,24 @@ def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(run 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) @@ -183,6 +210,8 @@ def test_endpoint_step_uses_resolved_retry_policy_and_backends(runtime_case: Run 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" in step.command for backend in deployment.backend_endpoints: assert f"http://127.0.0.1:{backend.port}" in step.command @@ -242,6 +271,7 @@ def test_client_receives_client_secrets_without_server_only_secrets( 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", @@ -250,7 +280,7 @@ def test_client_receives_client_secrets_without_server_only_secrets( assert step.environment["PACKAGE_INDEX_TOKEN"] == "client-secret" assert "HF_TOKEN" not in step.environment - assert "--container-env=PACKAGE_INDEX_TOKEN,SLURM_JOB_GPUS" in step.command + assert "--container-env=DATA_DESIGNER_SLURM_SCRATCH_ROOT,PACKAGE_INDEX_TOKEN,SLURM_JOB_GPUS" in step.command with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): DefaultClientStepBuilder().build_preflight_step( @@ -259,5 +289,8 @@ def test_client_receives_client_secrets_without_server_only_secrets( attempt, attempt_directory, endpoints, - {"HF_TOKEN": "server-secret"}, + { + "DD_SCRATCH_ROOT": "/tmp/data-designer-slurm-4101-0", + "HF_TOKEN": "server-secret", + }, ) diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index c38572c6f..296eb0511 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -6,6 +6,9 @@ #SBATCH --cpus-per-task=32 #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 set -Eeuo pipefail export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" @@ -33,9 +36,13 @@ 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}" -readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" -[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" +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 index 72a565bd4..95d4f45c6 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -6,6 +6,9 @@ #SBATCH --cpus-per-task=32 #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" @@ -33,9 +36,13 @@ 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}" -readonly DD_RUNTIME_ROOT="${DD_ATTEMPT_DIR}/runtime" -[[ -d ${DD_RUNTIME_ROOT} && ! -L ${DD_RUNTIME_ROOT} ]] -tar -xzf "${DD_RUNTIME_ARCHIVE}" -C "${DD_RUNTIME_ROOT}" +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/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 6c5578a70..6b3df11a6 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="3523fe78e80edd4b4acd1e090d9888b3a9879430f1460fecd3576755a4ca923e", + expected_fixture_sha256="dd74e01e7e2471a5c91343fe6fe771af4dd19976670fe3513e65e0d6568d8993", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="e69a4303113764670742d8d0ce12b1ebe4cc80751734bf68341b535b578108bf", + expected_fixture_sha256="5a76d1d73bc96112484689064e1cb9aad71a8e86970cc2f728c74ba47d415f9f", ) @@ -67,6 +67,9 @@ def _assert_script_matches_plan( assert f"#SBATCH --cpus-per-task={plan.client.authored.cpus}\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 diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 87ee5f001..db0ceb27e 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -6,6 +6,7 @@ import csv import hashlib import json +import shutil from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -101,6 +102,14 @@ def test_retry_refreshes_failed_shard_and_publishes_exact_next_attempt( 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, @@ -743,7 +752,7 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( script = cast(str, runner.inputs[-1]) assert "data_designer.slurm.state.collection_worker" in script assert "--gpus" not in script - assert "--gres" not in script + assert "#SBATCH --gres" not in script stale_stage = Path(case.plan.output.root).parent / submitted.staging_directory stale_stage.mkdir(mode=0o700) (stale_stage / "partial").write_text("incomplete") diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index a7865afc4..470aa1383 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -189,8 +189,7 @@ def track_read( attempt = _submitted_attempt(case) case.writer.create_attempt(attempt) runtime_directory = case.writer.run_root / "shards/shard-00000/attempts/attempt-0001/runtime" - assert runtime_directory.is_dir() - assert runtime_directory.stat().st_mode & 0o777 == 0o700 + assert not runtime_directory.exists() readiness = _readiness(case, attempt) case.writer.write_readiness(readiness) record_names.clear() From 9c721df1228e77974dcd7cd0dae7829587819c36 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:58:54 -0300 Subject: [PATCH 29/62] docs(slurm): add public workflow and release checks (#939) * docs(slurm): add public workflow and release checks Signed-off-by: Andre Manoel * fix(slurm): address release review findings Signed-off-by: Andre Manoel * fix(slurm): harden release validation Signed-off-by: Andre Manoel * docs(slurm): sanitize coverage source Signed-off-by: Andre Manoel * docs(slurm): fix walkthrough findings Signed-off-by: Andre Manoel * docs(slurm): harden onboarding guidance Signed-off-by: Andre Manoel * fix(slurm): preserve allocation failure diagnostics * docs(slurm): clarify verified onboarding behavior Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- fern/versions/latest.yml | 20 ++ .../latest/pages/slurm/benchmarks.mdx | 12 +- .../latest/pages/slurm/cli-reference.mdx | 42 ++++ .../pages/slurm/dependencies-and-tools.mdx | 115 ++++++++++ .../latest/pages/slurm/getting-started.mdx | 205 ++++++++++++++++++ fern/versions/latest/pages/slurm/images.mdx | 80 +++++++ .../latest/pages/slurm/operations.mdx | 71 ++++++ fern/versions/latest/pages/slurm/overview.mdx | 39 ++++ fern/versions/latest/pages/slurm/profiles.mdx | 105 +++++++++ .../pages/slurm/retry-and-collection.mdx | 56 +++++ .../latest/pages/slurm/run-configuration.mdx | 186 ++++++++++++++++ .../latest/pages/slurm/troubleshooting.mdx | 84 +++++++ packages/data-designer-slurm/README.md | 2 + .../examples/benchmark.yaml | 27 +++ .../data-designer-slurm/examples/builder.yaml | 10 + .../examples/profile-catalog.yaml | 34 +++ .../data-designer-slurm/examples/run.yaml | 29 +++ .../slurm/benchmark/execution.py | 1 - .../src/data_designer/slurm/cli.py | 7 +- .../src/data_designer/slurm/config/run.py | 2 +- .../data_designer/slurm/runtime/entrypoint.sh | 29 ++- .../slurm/runtime/step_runner.sh | 4 + .../data_designer/slurm/services/results.py | 2 +- .../src/data_designer/slurm/services/run.py | 11 +- .../data_designer/slurm/services/wiring.py | 3 +- .../tests/benchmark/test_workflow.py | 9 +- .../tests/contracts/test_config_records.py | 4 +- .../tests/runtime/test_shell_runtime.py | 69 ++++++ .../tests/services/test_services.py | 9 +- .../tests/services/test_wiring.py | 26 +-- .../data-designer-slurm/tests/test_cli.py | 23 +- plans/874/slurm-documentation-coverage.md | 36 +++ scripts/audit_slurm_public_artifacts.py | 6 +- scripts/test_slurm_package_install.py | 148 ++++++++++++- 34 files changed, 1439 insertions(+), 67 deletions(-) create mode 100644 fern/versions/latest/pages/slurm/cli-reference.mdx create mode 100644 fern/versions/latest/pages/slurm/dependencies-and-tools.mdx create mode 100644 fern/versions/latest/pages/slurm/getting-started.mdx create mode 100644 fern/versions/latest/pages/slurm/images.mdx create mode 100644 fern/versions/latest/pages/slurm/operations.mdx create mode 100644 fern/versions/latest/pages/slurm/overview.mdx create mode 100644 fern/versions/latest/pages/slurm/profiles.mdx create mode 100644 fern/versions/latest/pages/slurm/retry-and-collection.mdx create mode 100644 fern/versions/latest/pages/slurm/run-configuration.mdx create mode 100644 fern/versions/latest/pages/slurm/troubleshooting.mdx create mode 100644 packages/data-designer-slurm/examples/benchmark.yaml create mode 100644 packages/data-designer-slurm/examples/builder.yaml create mode 100644 packages/data-designer-slurm/examples/profile-catalog.yaml create mode 100644 packages/data-designer-slurm/examples/run.yaml create mode 100644 plans/874/slurm-documentation-coverage.md diff --git a/fern/versions/latest.yml b/fern/versions/latest.yml index 2f7cfed32..b903bad0c 100644 --- a/fern/versions/latest.yml +++ b/fern/versions/latest.yml @@ -81,8 +81,28 @@ navigation: 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/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx index f1161462d..059f46e6a 100644 --- a/fern/versions/latest/pages/slurm/benchmarks.mdx +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -2,6 +2,8 @@ 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. + Install the Slurm extension: ```bash @@ -22,12 +24,12 @@ concurrency_values: - 32 - 64 deployment_cases: - - name: two-independent-replicas + - name: single-node-groups deployments: generator: nodes: 2 nodes_per_replica: 1 - - name: one-two-node-replica + - name: two-node-group deployments: generator: nodes: 2 @@ -42,7 +44,9 @@ analysis: target_runtime: 4h ``` -The compiler expands the authored order deterministically. For this example, the first cases are `two-independent-replicas-c32` and `one-two-node-replica-c32`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. +The compiler expands the authored order deterministically. For this example, the first cases are `single-node-groups-c32` and `two-node-group-c32`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. + +`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)`, so it also depends on the selected profile and the base run's deployment topology. ## Run @@ -92,3 +96,5 @@ 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..8566d34fa --- /dev/null +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -0,0 +1,42 @@ +# 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` | Import or inspect an image through a CPU Slurm job and publish an 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 without deleting its SQSH 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` | Submit a zero-GPU winner collection job for a managed run directory. | +| `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 | +| --- | --- | --- | +| `--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. | +| `--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 one canonical JSON value on stdout. Failure is one object on stderr: + +```json +{"error":{"code":"invalid_request","message":"...","operation":"execute"}} +``` + +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..25f2affcd --- /dev/null +++ b/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx @@ -0,0 +1,115 @@ +# Dependencies, plugins, MCP, and secrets + +The client image provides Data Designer and its base environment. A run may add immutable pure Python wheels as an allocation-local overlay. The overlay is reconstructed for each allocation and is not written back into the image. + +## Add package requirements + +```yaml +client: + image: + name: dd-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. + +Do not use the overlay to replace `data-designer`, `data-designer-config`, `data-designer-engine`, or `data-designer-slurm`. Build a new client image when these packages need to change. + +## Reuse a dependency lock + +Use exactly one of `requirements` or `lock_file`: + +```yaml +client: + image: + name: dd-client + dependencies: + requirements: null + lock_file: inputs/dependency-lock.json +``` + +The lock must target the registered client image 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, bake the package into the client image, 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 client image 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 client image 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..f5cd08449 --- /dev/null +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -0,0 +1,205 @@ +# 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 digest-qualified client OCI image or absolute SQSH path containing the same Data Designer package versions as the submit environment +- A digest-qualified serving OCI image or absolute SQSH path containing a supported vLLM installation +- A model available to the serving image and compute nodes + +The image references and `example/generator` model below are syntax examples. They are intentionally not runnable values. + +## 1. Create and activate an environment + +Use a shared path instead of `$HOME` when required by site policy. 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. + +## 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 \ + --host-pattern 'login*.example.test' +``` + +The command writes the profile with restrictive permissions and never overwrites an existing file. Before submitting work, review every generated site value. In particular, confirm the account, partitions, hostname pattern, 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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. If Slurm commands are not on the batch-job `PATH`, also set `scheduler.bin_path` to their absolute directory. 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 --cluster primary +``` + +Validation checks strict profile loading, cluster selection, workspace access, and scheduler facts. Keep `--cluster primary` explicit when running from a compute host whose hostname does not match the login-host pattern. See [Profiles and clusters](/slurm/profiles-and-clusters) for non-default files and multiple clusters. + +## 6. Register a client image + +Use a digest-qualified OCI reference: + +```bash +data-designer slurm image add \ + registry.example.test/data-designer/client@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ + --kind client \ + --name dd-client \ + --cluster primary +``` + +Replace the entire source with the digest-qualified OCI reference or absolute SQSH path supplied for your site. See [Prepare compatible images](/slurm/images#prepare-compatible-images) before registering a client image. + +## 7. Register a serving image + +```bash +data-designer slurm image add \ + registry.example.test/inference/vllm@sha256:1111111111111111111111111111111111111111111111111111111111111111 \ + --kind serving \ + --name vllm \ + --cluster primary +``` + +Each `image add` command submits a CPU Slurm job, waits for it to finish, inspects the result, and only then publishes the alias. Keep the terminal open. You can also register an existing absolute `.sqsh` path. See [Images](/slurm/images). + +## 8. Create a Data Designer builder config + +Create `builder.yaml`: + +```yaml +data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + columns: + - name: greeting + column_type: llm-text + prompt: Write one short greeting. + model_alias: generator +``` + +The model alias in this file must match a deployment in the Slurm run config. + +## 9. 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 +client: + image: + name: dd-client +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 +output: + format: parquet + partitions: 1 +``` + +The repository also contains these 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. + +## 10. Dry-run the submission + +```bash +data-designer slurm execute run.yaml --dry-run --cluster primary +``` + +Dry-run resolves the selected profile, verifies registered image contents, loads the builder envelope, resolves client dependencies, and renders the batch script without calling `sbatch`. Full builder and plugin validation runs inside the prepared client 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`. + +## 11. Submit the run + +```bash +data-designer slurm execute run.yaml --cluster primary +``` + +The command prints one JSON object. Save its `run_id`; all lifecycle commands use the managed run ID rather than a raw Slurm job ID. + +## 12. Inspect or cancel the run + +```bash +data-designer slurm status --cluster primary +data-designer slurm cancel --cluster primary +``` + +Status reconciles active queue and accounting evidence with durable run state. Cancellation targets only jobs recorded for the managed run. + +## 13. 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 an explicit destination, submit a managed collection job: + +```bash +data-designer slurm merge \ + --input-path /workspace/alice/data-designer-slurm/runs/ \ + --output-path /workspace/alice/exports/greetings \ + --cluster primary +``` + +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..0c6f921a0 --- /dev/null +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -0,0 +1,80 @@ +# Slurm images + +Data Designer Slurm uses two image roles: + +- `client` images run Data Designer and installed plugins. +- `serving` images run vLLM model servers. + +Use separate aliases even if both roles originate from related images. Inspection verifies role-specific contents before an alias becomes available to planning. + +## Prepare compatible images + +Image construction is an operator prerequisite. Data Designer Slurm imports or inspects images, but does not build a client or serving environment from a Dockerfile. + +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. A client image must install the same release used by the submit environment: + +```bash +python -m pip install "data-designer[slurm]==" +``` + +For an unpublished build, install the aligned `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` wheels into the client image. Do not mix versions. A serving 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 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/data-designer/client@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ + --kind client \ + --name dd-client +``` + +OCI sources must use `repository@sha256:<64 lowercase hex characters>`. Schemes, embedded credentials, tags without a digest, 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. + +## Register an existing SQSH + +```bash +data-designer slurm image add \ + /workspace/alice/images/client.sqsh \ + --kind client \ + --name dd-client +``` + +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 dd-client +``` + +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/client-v2.sqsh \ + --kind client \ + --name dd-client \ + --replace +``` + +Removing an alias does not delete its SQSH artifact: + +```bash +data-designer slurm image rm dd-client +``` + +Published OCI imports are content-addressed. Identical imports can share an artifact, while conflicting bytes or inspection facts fail closed. + +## Compatibility rules + +The client image inventory must contain the same versions of `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` used to submit the run. Dependency overlays may add pure Python wheels, but cannot replace these protected distributions. 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..283a22faa --- /dev/null +++ b/fern/versions/latest/pages/slurm/operations.mdx @@ -0,0 +1,71 @@ +# Slurm run operations + +The CLI prints canonical JSON to stdout. Errors are canonical JSON on stderr, so automation can consume the same commands as an operator. + +## 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. + +Use `--profile-file` and `--cluster` on any command when the defaults are not appropriate. + +## 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..4bd91d2f7 --- /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 Enroot 3.5 or later for existing SQSH images. Digest-pinned OCI imports need Enroot 4.0 or later. The client image must contain the same Data Designer packages as the submit environment. 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 images, start vLLM and client containers, 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, collection state, and default output. + +Runtime bundle extraction, the client dependency overlay, Data Designer home and cache directories, and Enroot state use allocation-local scratch. The runtime mounts that scratch at `/run/data-designer-slurm` inside containers and 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, client and 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..b7fc903b0 --- /dev/null +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -0,0 +1,105 @@ +# 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 +mkdir -p /workspace/alice/config +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 \ + --host-pattern 'login*.example.test' +``` + +The starter sets `gpus_per_node: auto`, `gpu_request_mode: gres`, and CPU image-build resources of 2 CPUs, 8 GiB, and four hours. The parent directory of a custom `--profile-file` must already exist. Review the generated account, partitions, host patterns, workspace, GPU settings, and image-build resources against site policy before validation or image registration. + +## 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 + bin_path: /opt/slurm/bin + 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`. + +`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. + +Set `scheduler.bin_path` when Slurm executables are installed outside the generated batch job's default `PATH`. It must be one absolute directory and is prepended to `PATH`; 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..7816bc75a --- /dev/null +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -0,0 +1,56 @@ +# 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 + +Default successful output is already managed beneath the run directory. Use `merge` to submit a zero-GPU CPU job that collects shard winners to an explicit compute-visible destination: + +```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. The output path must be absolute, writable, compute-visible, and authorized by the profile mounts. + +If supplied, `--num-partitions` must equal the partition count persisted in the run: + +```bash +data-designer slurm merge \ + --input-path /workspace/alice/data-designer-slurm/runs/ \ + --output-path /workspace/alice/exports/final-dataset \ + --num-partitions 20 +``` + +Repeating collection returns or reuses managed state when safe. Conflicting destinations or partition counts are rejected. 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..225d35aad --- /dev/null +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -0,0 +1,186 @@ +# 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 +client: + cpus: 32 + image: + name: dd-client +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 +output: + format: parquet + partitions: 1 + require_exact_record_count: false +``` + +`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, while `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. Output formats are `parquet`, `jsonl`, and `csv`. + +## 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. + +`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 + +### Arrays + +Split records into independent shards and limit how many allocations run at once: + +```yaml +invocation: + num_records: 100000 + dataset_name: large-run +array_tasks: + count: 20 + max_concurrent: 4 +output: + partitions: 20 +``` + +`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. + +Runs with more than one array task require Parquet output. They 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 keep one vLLM replica per node: + +```yaml +resources: + nodes: 4 +topology: + tensor_parallel: 8 + 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 +``` + +`tensor_parallel` must divide the GPUs available on each node. With the eight GPUs per node in the profile example, `tensor_parallel: 8` and `nodes_per_replica: 2` use 16 GPUs per replica. `nodes_per_replica` must divide `resources.nodes`. 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 +``` + +## 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(image={"name": "dd-client"}) + .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..1ccc3e061 --- /dev/null +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -0,0 +1,84 @@ +# Slurm troubleshooting + +Start with the JSON error on stderr, then inspect the managed run directory and Slurm state. 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 + +Set `scheduler.bin_path` in the selected profile to the absolute directory containing `srun`, `scontrol`, and the other Slurm executables. The generated batch script prepends this 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. + +## Image registration fails + +- Confirm the image-build partition accepts CPU jobs. +- Confirm Enroot is available on the allocated node. +- Use Enroot 4.0 or later for OCI import and 3.5 or later for existing SQSH inspection. +- Use a digest-qualified, credential-free OCI reference. +- 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. Retry missing shards, then use `merge` with the original run directory. When `require_exact_record_count` is true, a short result remains a failure rather than being silently accepted. + +## 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-slurm/README.md b/packages/data-designer-slurm/README.md index 04d29ccbd..1c000bd69 100644 --- a/packages/data-designer-slurm/README.md +++ b/packages/data-designer-slurm/README.md @@ -7,3 +7,5 @@ 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..120120654 --- /dev/null +++ b/packages/data-designer-slurm/examples/benchmark.yaml @@ -0,0 +1,27 @@ +schema_version: 1 +name: generator-scaling +base_run: run.yaml +model_aliases: + - generator +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 +record_policy: + type: adaptive + base_records: 1000 + max_records: 5000 + records_per_concurrency: 1.0 +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..9e2ca35d8 --- /dev/null +++ b/packages/data-designer-slurm/examples/builder.yaml @@ -0,0 +1,10 @@ +data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + columns: + - name: greeting + column_type: llm-text + prompt: Write one short greeting. + model_alias: generator 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..77ce34607 --- /dev/null +++ b/packages/data-designer-slurm/examples/profile-catalog.yaml @@ -0,0 +1,34 @@ +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 + bin_path: /opt/slurm/bin + 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..6879a4a96 --- /dev/null +++ b/packages/data-designer-slurm/examples/run.yaml @@ -0,0 +1,29 @@ +schema_version: 1 +name: greeting-run +builder: + source: builder.yaml +invocation: + num_records: 8 + dataset_name: greetings + model_concurrency: + generator: 8 +client: + image: + name: dd-client +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 +output: + format: parquet + partitions: 1 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 index 6841a453c..871d88e2d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -122,7 +122,6 @@ def run( case.child_run_config, source_root=child_source_root, dry_run=False, - force=False, ) if execution.run_id != case.child_run_id or execution.state != "submitted": raise SlurmServiceError( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 18d8a4c2f..8a0ee3b27 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -65,7 +65,6 @@ def execute_command( 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"), - force: bool = typer.Option(False, "--force"), ) -> None: """Prepare and submit one authored run.""" operation = SlurmServiceOperation.EXECUTE_RUN @@ -73,7 +72,7 @@ def execute_command( def execute() -> BaseModel: config = load_run_config(run_file) service = create_slurm_run_service(profile_file=profile_file, cluster=cluster) - return service.execute(config, source_root=run_file.resolve().parent, dry_run=dry_run, force=force) + return service.execute(config, source_root=run_file.resolve().parent, dry_run=dry_run) _emit_result(_invoke(operation, execute)) @@ -84,7 +83,7 @@ def status_command( 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 M2 run status.""" + """Reconcile scheduler observations and show persisted run status.""" operation = SlurmServiceOperation.STATUS_RUN result = _invoke( operation, @@ -240,7 +239,7 @@ def image_add_command( profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), cluster: str | None = typer.Option(None, "--cluster"), ) -> None: - """Request image import or inspection; requires IMG lifecycle support.""" + """Import or inspect an image and register its alias.""" operation = SlurmServiceOperation.ADD_IMAGE def add() -> BaseModel: 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 index 0d9380a92..f561a397d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -85,7 +85,7 @@ def validate_input(self) -> BuilderInput: 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 Big Iron fields: {', '.join(sorted(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") 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 index 327214f77..3a88d6515 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -11,6 +11,7 @@ source "${DD_RUNTIME_DIR}/cleanup.sh" DD_RUNTIME_PREPARED=0 DD_RUNTIME_FINALIZED=0 +DD_RUNTIME_PHASE=initialization DD_RUNTIME_RETRY_ARGUMENTS=() dd_slurm_run_allocation() { @@ -29,21 +30,28 @@ dd_slurm_run_allocation() { --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_PHASE=translate_plan_path dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false DD_PLAN_CONTAINER_PATH=${DD_CONTAINER_PATH} + DD_RUNTIME_PHASE=translate_attempt_path dd_read_container_path "${DD_PLAN_PATH}" "${DD_ATTEMPT_PATH}" true DD_ATTEMPT_CONTAINER_DIR=${DD_CONTAINER_PATH} DD_RUNTIME_CONTAINER_ROOT=${DD_SCRATCH_CONTAINER_ROOT}/runtime + DD_RUNTIME_PHASE=translate_manifest_path dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_MANIFEST}" true DD_RUNTIME_MANIFEST_CONTAINER_PATH=${DD_CONTAINER_PATH} export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST export DD_PLAN_CONTAINER_PATH DD_ATTEMPT_CONTAINER_DIR DD_RUNTIME_CONTAINER_ROOT - trap dd_runtime_exit EXIT - trap 'exit 130' INT TERM + DD_RUNTIME_PHASE=verify_host_context dd_verify_host_context + DD_RUNTIME_PHASE=start_runtime_timer dd_start_runtime_timer DD_RUNTIME_PREPARED=1 @@ -52,29 +60,40 @@ dd_slurm_run_allocation() { 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_CONTAINER_ROOT}" \ --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" \ "${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_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_run_role_steps server_preflight + DD_RUNTIME_PHASE=start_servers dd_start_servers + DD_RUNTIME_PHASE=server_readiness dd_wait_for_role_readiness server + 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_read_step_ids "${DD_RUNTIME_MANIFEST}" client ((${#DD_STEP_IDS[@]} == 1)) dd_start_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" @@ -83,9 +102,12 @@ dd_slurm_run_allocation() { dd_wait_for_client "${client_pid}" 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_run_bound_control_phase() { @@ -236,6 +258,9 @@ 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 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 index 22c19afde..f1c0fb2b5 100644 --- 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 @@ -138,6 +138,10 @@ dd_run_step() { 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}" } 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 index 5d25a51f9..cecbbbf14 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/results.py @@ -83,7 +83,7 @@ def validate_shard(self) -> SlurmPersistedShardStatus: class SlurmPersistedRunStatus(ContractValue): - """M2 status assembled only from durable run records.""" + """Status assembled only from durable run records.""" run: RunManifest shards: tuple[SlurmPersistedShardStatus, ...] 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 index f40ea2614..e241b3f4f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -63,12 +63,11 @@ def execute( *, source_root: Path, dry_run: bool, - force: bool, ) -> SlurmRunExecution: """Render or submit one run.""" def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: - """Return the durable M2 status for one run.""" + """Return the durable status for one run.""" def cancel(self, run_id: Identifier) -> SlurmRunCancellation: """Request cancellation of active jobs.""" @@ -177,7 +176,6 @@ def execute( *, source_root: str | Path = ".", dry_run: bool = False, - force: bool = False, ) -> SlurmRunExecution: """Render or submit one run through package-owned production wiring.""" operation = SlurmServiceOperation.EXECUTE_RUN @@ -185,8 +183,8 @@ def execute( 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 or type(force) is not bool: - raise _make_invalid_request_error(operation, "dry_run and force must be booleans") + 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: @@ -194,7 +192,6 @@ def execute_run() -> SlurmRunExecution: config, source_root=Path(source_root).expanduser().resolve(), dry_run=dry_run, - force=force, ) if not isinstance(result, SlurmRunExecution): raise TypeError("run backend returned an invalid execution result") @@ -203,7 +200,7 @@ def execute_run() -> SlurmRunExecution: return _invoke_service_backend(operation, execute_run) def status(self, run_id: Identifier) -> SlurmPersistedRunStatus: - """Reconcile scheduler observations and return persisted M2 records.""" + """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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 6055e6604..b4bfced83 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -298,7 +298,6 @@ def execute( *, source_root: Path, dry_run: bool, - force: bool, ) -> SlurmRunExecution: with self._preparer.prepare( config, @@ -321,7 +320,7 @@ def execute( publisher, config, prepared, - force=force, + force=False, ) try: receipt = self._launcher.submit_script( diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index 99e23af5f..735ba900f 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -64,9 +64,9 @@ def __init__(self, run_id, configs, submissions, calls, benchmark_root, failures self.benchmark_root = benchmark_root self.failures = failures - def execute(self, config, *, source_root, dry_run, force): + 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, force)) + 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( @@ -89,9 +89,9 @@ def __init__(self, *args, second_check: Event) -> None: super().__init__(*args) self.second_check = second_check - def execute(self, config, *, source_root, dry_run, force): + 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, force=force) + return super().execute(config, source_root=source_root, dry_run=dry_run) class _ConcurrentSubmissionLoader: @@ -207,7 +207,6 @@ def test_run_persists_manifest_before_submission_and_is_idempotent( assert first == second assert len(calls) == len(compiled.cases) assert tuple(configs) == tuple(case.child_run_id for case in compiled.cases) - assert all(force is False for *_, force in calls) def test_concurrent_runs_submit_each_deterministic_child_once( diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index 58651662a..f74dc6a75 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -602,9 +602,11 @@ def test_run_rejects_retired_builder_fields(authored_run: DataDesignerSlurmConfi payload = authored_run.model_dump(mode="json") payload["builder"]["inline"]["server_configs"] = [] - with pytest.raises(ValidationError, match="retired"): + 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", diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 649af45dd..b4cafd17d 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -336,6 +336,75 @@ def test_shell_helpers_handle_empty_and_sparse_arrays() -> None: 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""" diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index df0b0e47e..b9b6855c9 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -194,7 +194,7 @@ def test_run_service_delegates_execute_actions( ) -> None: class RunBackend: def __init__(self) -> None: - self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool, bool]] = [] + self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool]] = [] def execute( self, @@ -202,9 +202,8 @@ def execute( *, source_root: Path, dry_run: bool, - force: bool, ) -> SlurmRunExecution: - self.calls.append((config, source_root, dry_run, force)) + self.calls.append((config, source_root, dry_run)) return SlurmRunExecution( run_id="run-0001", state="dry_run", @@ -222,10 +221,10 @@ def cancel(self, run_id): backend = RunBackend() service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) - result = service.execute(authored_run_single, source_root=tmp_path, dry_run=True, force=True) + 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, True)] + assert backend.calls == [(authored_run_single, tmp_path.resolve(), True)] def test_run_service_delegates_retry_with_stable_shard_order() -> None: diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 02f072ebe..8af4bd7ae 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -385,11 +385,11 @@ def test_production_wiring_submits_after_publisher_initialization( package_version="0.9.2", ) - result = service.execute(authored_run_single, source_root=tmp_path, force=True) + result = service.execute(authored_run_single, source_root=tmp_path) assert result.state == "submitted" assert result.job_id == 42 - assert publisher.initializations == [("run-wired", True)] + assert publisher.initializations == [("run-wired", False)] assert publisher.submissions == [("run-wired", 42, submitted_at)] assert len(launcher.submissions) == 1 assert launcher.held_submissions == [True] @@ -398,28 +398,6 @@ def test_production_wiring_submits_after_publisher_initialization( assert (tmp_path / "managed-assets").is_dir() -def test_production_publisher_rejects_force_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) - 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="different inputs") as caught: - service.execute(authored_run_single, source_root=tmp_path, force=True) - - assert caught.value.code is SlurmServiceErrorCode.CONFLICT - assert launcher.submissions == [] - - def test_production_wiring_exports_referenced_secrets_to_the_allocation( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index dfa23f957..0cf0b83e1 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -26,7 +26,7 @@ class _RunService: def __init__(self) -> None: - self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool, bool]] = [] + self.calls: list[tuple[DataDesignerSlurmConfig, Path, bool]] = [] self.retry_calls: list[tuple[str, tuple[str, ...] | None, str, bool]] = [] self.collection_calls: list[tuple[Path, Path, int | None]] = [] @@ -36,9 +36,8 @@ def execute( *, source_root: Path, dry_run: bool, - force: bool, ) -> SlurmRunExecution: - self.calls.append((config, source_root, dry_run, force)) + self.calls.append((config, source_root, dry_run)) return SlurmRunExecution( run_id="run-0001", state="dry_run", @@ -101,6 +100,20 @@ def analyze(self, benchmark_id, *, refresh_state, fail_if_incomplete): return self.report +@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, @@ -111,7 +124,7 @@ def test_execute_emits_deterministic_json_and_forwards_actions( 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", "--force"]) + result = CliRunner().invoke(cli_module.create_cli(), ["execute", str(run_file), "--dry-run"]) assert result.exit_code == 0 assert json.loads(result.stdout) == { @@ -122,7 +135,7 @@ def test_execute_emits_deterministic_json_and_forwards_actions( "shard_count": 1, "state": "dry_run", } - assert service.calls == [(authored_run_single, tmp_path, True, True)] + assert service.calls == [(authored_run_single, tmp_path, True)] def test_benchmark_cli_forwards_run_and_analysis_actions( 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/scripts/audit_slurm_public_artifacts.py b/scripts/audit_slurm_public_artifacts.py index 236f8608b..41bdcbc85 100644 --- a/scripts/audit_slurm_public_artifacts.py +++ b/scripts/audit_slurm_public_artifacts.py @@ -4,7 +4,7 @@ """Scan public Slurm artifacts without rendering matched sensitive content. The default scope covers deployable source, test source and fixtures, package -metadata, and release scripts. A small path-scoped allowlist masks only exact +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. """ @@ -32,9 +32,13 @@ "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", ) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 82a7cc4b5..7da8135ef 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import os import shutil import statistics @@ -10,9 +11,14 @@ 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 @@ -25,15 +31,34 @@ "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 -def run(command: list[str], *, cwd: Path, check: bool = True) -> subprocess.CompletedProcess[str]: +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) @@ -56,6 +81,19 @@ def audit_public_artifacts(*paths: Path) -> None: 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( @@ -95,6 +133,7 @@ def install(uv: str, python: Path, wheel_directory: Path, package: str, *, cwd: "--python", str(python), "--prerelease=allow", + "--no-index", "--find-links", str(wheel_directory), package, @@ -103,6 +142,79 @@ def install(uv: str, python: Path, wheel_directory: Path, package: str, *, 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 @@ -142,6 +254,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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 @@ -181,6 +294,14 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert StateArtifactReference is ContractArtifactReference assert StateRecordRange is ContractRecordRange assert StateResumeWorkspace is ContractResumeWorkspace +examples = Path("examples") +assert load_builder_payload(examples / "builder.yaml")["data_designer"]["columns"][0]["name"] == "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: @@ -233,27 +354,42 @@ 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) - base_wheel = wheels["data-designer"] + 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 = wheel_metadata(base_wheel) - leaf_metadata = wheel_metadata(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") + 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 @@ -272,7 +408,9 @@ def main() -> None: verify_install(base_python, version, slurm=False, cwd=root) extra_python = create_environment(uv, root / "extra", cwd=root) - install(uv, extra_python, wheel_directory, f"data-designer[slurm]=={version}", 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 From ca4dd9d684c661b8a8490a6f189c2a3ca507ced5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 28 Sep 2026 16:02:38 -0600 Subject: [PATCH 30/62] fix(slurm): run collection with native pinned runtime (#965) * fix(slurm): run collection natively Bind collection jobs to the submitting Python environment and the digest-pinned runtime bundle so CPU-only collection does not depend on the client container image. Preserve no-overwrite publication on filesystems that do not support atomic no-replace rename by reserving the destination before promotion. Fixes #964 Signed-off-by: Nabin Mulepati * fix(slurm): recover collection plans and outputs Read existing v1 collection plans without changing their persisted bytes and write new native plans as v2. Journal publication reservations so interrupted collection jobs can safely retry their output destination. Signed-off-by: Nabin Mulepati * fix(slurm): serialize plan versions precisely Omit only the execution field absent from each collection plan version so future nullable fields retain their serialized meaning. Signed-off-by: Nabin Mulepati * fix(slurm): verify reservation identity Write the recovery marker only after the destination exists and bind it to that directory inode. Preserve unrelated or ambiguous destinations when retrying collection. Drop the extra collection plan version because Slurm has not yet been released from this feature branch. Signed-off-by: Nabin Mulepati * fix(slurm): bind reservations to change time Require the original directory change time as well as its device and inode before recovery removes a private empty reservation. Signed-off-by: Nabin Mulepati * fix(slurm): preserve changed reservations Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../slurm/launcher/collection.py | 46 +--- .../data_designer/slurm/state/collection.py | 31 ++- .../slurm/state/collection_filesystem.py | 170 +++++++++++++ .../slurm/state/collection_storage.py | 12 +- .../slurm/state/collection_worker.py | 8 +- .../data_designer/slurm/state/destinations.py | 44 +--- .../src/data_designer/slurm/state/outputs.py | 4 +- .../tests/launcher/test_collection.py | 24 +- .../tests/state/golden/collection_plan.json | 2 +- .../tests/state/test_model_edges.py | 8 + .../tests/state/test_retry_collection.py | 235 ++++++++++++++++++ .../tests/state/test_state_golden_records.py | 2 +- 12 files changed, 478 insertions(+), 108 deletions(-) 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 index abe663892..3ca48ade0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -7,7 +7,6 @@ import posixpath from importlib import resources -from pathlib import PurePosixPath from data_designer.slurm.launcher.batch import quote_shell_value, render_batch_directives from data_designer.slurm.launcher.errors import SlurmBatchRenderError @@ -26,8 +25,6 @@ def render_collection_script( 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") - if collection_plan.container_destination != destination.container_path: - raise SlurmBatchRenderError("collection container destination does not match its resolved mount") collection_root = posixpath.join( posixpath.dirname(resolved_plan.authored_config.path), @@ -49,29 +46,23 @@ def render_collection_script( ("error", f"{collection_root}/slurm-%j.err"), ) ) - workspace_root = resolved_plan.selected_profile.profile.workspace_root - state_mount = f"{workspace_root}:{workspace_root}" - output_mount = f"{destination.mount.source}:{destination.mount.target}" - mount_arguments = _render_mount_arguments( - ("DD_STATE_MOUNT", workspace_root, workspace_root), - ("DD_OUTPUT_MOUNT", destination.mount.source, destination.mount.target), - ) scratch_source = resources.files("data_designer.slurm.runtime").joinpath("scratch.sh").read_text() 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_CLIENT_IMAGE={quote_shell_value(resolved_plan.client.image.path)} -readonly DD_CLIENT_IMAGE_SHA256={quote_shell_value(resolved_plan.client.image.sha256)} 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(resolved_plan.package_version)} +readonly DD_RUNTIME_ARCHIVE={quote_shell_value(resolved_plan.runtime_bundle.path)} +readonly DD_RUNTIME_SHA256={quote_shell_value(resolved_plan.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)} -readonly DD_STATE_MOUNT={quote_shell_value(state_mount)} -readonly DD_OUTPUT_MOUNT={quote_shell_value(output_mount)} verify_sha256() {{ local actual_sha256 @@ -79,33 +70,20 @@ def render_collection_script( [[ "${{actual_sha256%% *}}" == "$1" ]] }} -verify_sha256 "${{DD_CLIENT_IMAGE_SHA256}}" "${{DD_CLIENT_IMAGE}}" 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 -export HOME="${{DD_SCRATCH_ROOT}}/home" -DD_ENROOT_MOUNTS=({mount_arguments}) -readonly DD_ENROOT_MOUNTS -enroot start --root "${{DD_ENROOT_MOUNTS[@]}}" "${{DD_CLIENT_IMAGE}}" \ - python -m data_designer.slurm.state.collection_worker \ +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}}" """ -def _render_mount_arguments(*mounts: tuple[str, str, str]) -> str: - unique: dict[str, tuple[str, str, str]] = {} - targets: dict[str, str] = {} - for variable, source, target in mounts: - mount = f"{source}:{target}" - existing_source = targets.get(target) - if existing_source is not None and existing_source != source: - raise SlurmBatchRenderError("collection state and output mounts cannot share a target") - targets[target] = source - unique.setdefault(mount, (variable, source, target)) - ordered = sorted(unique.values(), key=lambda item: len(PurePosixPath(item[2]).parts)) - return " ".join(f'--mount "${{{variable}}}"' for variable, _, _ in ordered) - - __all__ = ["render_collection_script"] 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 index 78b1b0fab..9995b61c1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -5,6 +5,7 @@ from __future__ import annotations +import sys from datetime import datetime, timezone from pathlib import Path from typing import Protocol @@ -80,6 +81,7 @@ def __init__( self._destinations = CollectionDestinationResolver() self._collector = SchedulerObservationCollector(self._scheduler) self._run_id = normalized_run_id + self._python_executable = _resolve_python_executable() def submit( self, @@ -107,7 +109,7 @@ def submit( remove_collection_stage( Path(current_plan.host_destination), current.staging_directory, - Path(resolved_destination.mount.source), + Path(resolved_destination.authorized_root), ) run, resolved_plan, _ = self._reader.load_context() resolved_destination = self._destinations.resolve(resolved_plan, destination) @@ -118,14 +120,14 @@ def submit( created_at=timestamp, resolved_plan=run.resolved_plan, planned_shards=self._inputs.get_winner_shards(), + python_executable=self._python_executable, host_destination=resolved_destination.host_path, - container_destination=resolved_destination.container_path, num_partitions=resolved_plan.output.partitions, ) self._inputs.resolve(collection_plan) prepare_collection_destination( Path(collection_plan.host_destination), - Path(resolved_destination.mount.source), + Path(resolved_destination.authorized_root), ) self._collections.ensure_collection(collection_plan.collection_id) self._collections.publish_plan(collection_plan) @@ -173,7 +175,7 @@ def refresh( remove_collection_stage( Path(plan.host_destination), previous.staging_directory, - Path(destination.mount.source), + Path(destination.authorized_root), ) return previous if previous.state is CollectionState.PREPARED: @@ -182,7 +184,7 @@ def refresh( remove_collection_stage( Path(plan.host_destination), previous.staging_directory, - Path(destination.mount.source), + Path(destination.authorized_root), ) return previous assert previous.scheduler is not None @@ -208,7 +210,7 @@ def refresh( remove_collection_stage( Path(plan.host_destination), current.staging_directory, - Path(destination.mount.source), + Path(destination.authorized_root), ) return current except (StateConflictError, StateCorruptionError, SlurmStateError): @@ -298,7 +300,7 @@ def _validate_existing_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 or plan.container_destination != resolved.container_path: + 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) @@ -334,12 +336,7 @@ def _load_valid_result( ) 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, - Path(plan.host_destination), - verify_digests=False, - ) + self._collections.verify_result_files(plan, validated, verify_digests=False) return validated def _publish_succeeded( @@ -369,6 +366,14 @@ def _load_bound_plan(self, status: CollectionStatus) -> CollectionPlan: 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()) 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 index 72330c1c0..fdc239f79 100644 --- 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 @@ -7,6 +7,7 @@ import ctypes import errno +import hashlib import os import stat import sys @@ -23,6 +24,14 @@ _RENAME_NOREPLACE = 1 _RENAME_EXCL = 0x00000004 +_UNSUPPORTED_NO_REPLACE_ERRORS = frozenset( + { + errno.EINVAL, + errno.ENOSYS, + errno.ENOTSUP, + errno.EOPNOTSUPP, + } +) @dataclass(frozen=True, slots=True) @@ -118,6 +127,7 @@ def remove_collection_stage(destination: Path, staging_directory: str, authorize 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 @@ -141,6 +151,7 @@ def stage_collection( 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) @@ -313,6 +324,30 @@ def _rename_without_overwrite( 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) @@ -343,10 +378,145 @@ def _rename_without_overwrite( 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", 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 index 7da068a42..3d53de2e0 100644 --- 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 @@ -198,12 +198,8 @@ def read_status(self, collection_id: Identifier) -> CollectionStatus: return status def read_result(self, plan: CollectionPlan) -> CollectionResult: - return self.read_result_from(plan, Path(plan.host_destination)) - - def read_result_from(self, plan: CollectionPlan, destination: Path) -> CollectionResult: - """Read a result through an explicitly selected host or container view.""" - if destination.as_posix() not in {plan.host_destination, plan.container_destination}: - raise OSError("collection result destination does not match its immutable plan") + """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, @@ -216,13 +212,11 @@ def verify_result_files( self, plan: CollectionPlan, result: CollectionResult, - destination: Path, *, verify_digests: bool = True, ) -> None: """Verify exact inventory with bounded metadata or full output digests.""" - if destination.as_posix() not in {plan.host_destination, plan.container_destination}: - raise OSError("collection result destination does not match its immutable plan") + 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") 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 index 8f464ee6c..fa2019f2a 100644 --- 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 @@ -97,9 +97,9 @@ def _execute_collection( try: destination = self._destinations.validate_persisted(resolved_plan, plan) with stage_collection( - Path(plan.container_destination), + Path(plan.host_destination), derive_collection_staging_directory(plan), - Path(destination.mount.target), + Path(destination.authorized_root), ) as staged: merger = CollectionMerger(resolved_plan.output.format, completed_at=completed_at) result = merger.merge( @@ -138,7 +138,7 @@ def _load_valid_result( plan: CollectionPlan, status: CollectionStatus | None = None, ) -> CollectionResult: - result = self._collections.read_result_from(plan, Path(plan.container_destination)) + result = self._collections.read_result(plan) resolved_plan = self._reader.load_resolved_plan() validated = validate_collection_result( plan, @@ -148,7 +148,7 @@ def _load_valid_result( ) 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, Path(plan.container_destination)) + self._collections.verify_result_files(plan, validated) return validated def _publish_success_status( 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 index 1c738a2d6..746483d0e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Profile-authorized host and container collection destinations.""" +"""Profile-authorized host collection destinations.""" from __future__ import annotations @@ -9,9 +9,7 @@ from dataclasses import dataclass from pathlib import Path -from data_designer.slurm.config import ContainerMount from data_designer.slurm.contracts import is_path_below, paths_overlap, validate_absolute_path -from data_designer.slurm.images.records import validate_enroot_mount_path from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.state.errors import StateConflictError from data_designer.slurm.state.outputs import CollectionPlan @@ -19,11 +17,10 @@ @dataclass(frozen=True, slots=True) class CollectionDestination: - """One normalized output directory and its authorized writable mount.""" + """One normalized host output directory and its authorized writable root.""" host_path: str - container_path: str - mount: ContainerMount + authorized_root: str class CollectionDestinationResolver: @@ -34,7 +31,7 @@ def resolve( plan: ResolvedSlurmRunPlan, requested_destination: str | Path | None = None, ) -> CollectionDestination: - """Return one authorized output path in host and container namespaces.""" + """Return one authorized host output path.""" raw_destination = plan.output.root if requested_destination is None else Path(requested_destination).as_posix() try: host_path = validate_absolute_path(raw_destination) @@ -42,35 +39,14 @@ def resolve( raise StateConflictError("collection destination must be a normalized absolute path") from error workspace_root = plan.selected_profile.profile.workspace_root _validate_destination_overlap(plan, host_path) - workspace_mount = ContainerMount(source=workspace_root, target=workspace_root, read_only=False) - authorized = { - (mount.source, mount.target): mount - for mount in (workspace_mount, *plan.container_mounts) - if not mount.read_only - } - writable = tuple( - mount - for mount in authorized.values() - if host_path == mount.source or is_path_below(host_path, mount.source) - ) + 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") - longest = max(len(mount.source) for mount in writable) - matches = tuple(mount for mount in writable if len(mount.source) == longest) - if len(matches) != 1: - raise StateConflictError("collection destination has an ambiguous writable mount mapping") - mount = matches[0] - if host_path == mount.source: + authorized_root = max(writable, key=len) + if host_path == authorized_root: raise StateConflictError("collection destination must be below its writable mount source") - try: - validate_enroot_mount_path(workspace_root) - validate_enroot_mount_path(mount.source) - validate_enroot_mount_path(mount.target) - except ValueError as error: - raise StateConflictError("collection paths cannot be represented as safe Enroot mounts") from error - relative = posixpath.relpath(host_path, mount.source) - container_path = mount.target if relative == "." else posixpath.join(mount.target, relative) - return CollectionDestination(host_path, validate_absolute_path(container_path), mount) + return CollectionDestination(host_path, authorized_root) def validate_persisted( self, @@ -83,8 +59,6 @@ def validate_persisted( raise StateConflictError("collection run identity does not match the resolved plan") if collection_plan.host_destination != destination.host_path: raise StateConflictError("collection host destination no longer matches resolved intent") - if collection_plan.container_destination != destination.container_path: - raise StateConflictError("collection container destination no longer matches resolved intent") if collection_plan.num_partitions != resolved_plan.output.partitions: raise StateConflictError("collection partition count no longer matches resolved intent") return destination 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 index 13aa11887..c83ec6089 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -162,13 +162,13 @@ class CollectionPlan(StateRecord): created_at: datetime resolved_plan: ArtifactReference planned_shards: tuple[CollectionShard, ...] = Field(min_length=1) + python_executable: str host_destination: str - container_destination: str num_partitions: PositiveInt overwrite: Literal[False] = False _created_at_is_utc = field_validator("created_at")(validate_utc_timestamp) - _destinations_are_safe = field_validator("host_destination", "container_destination")(validate_absolute_path) + _paths_are_safe = field_validator("python_executable", "host_destination")(validate_absolute_path) @property def submission_job_name(self) -> Identifier: diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 29ef1070c..2a56d9d0a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -20,7 +20,7 @@ from data_designer.slurm.state.errors import StateConflictError -def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( +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) @@ -63,8 +63,8 @@ def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( ), ), ), + python_executable="/shared/data-designer/bin/python", host_destination=destination.host_path, - container_destination=destination.container_path, num_partitions=plan.output.partitions, ) @@ -74,12 +74,19 @@ def test_collection_renderer_uses_authorized_mounts_and_no_gpu_directives( 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_STATE_MOUNT="/workspace/primary:/workspace/primary"' in script - assert 'readonly DD_OUTPUT_MOUNT="/workspace/primary/runs/run-001:/exports"' 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.2"' in script + assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"' in script + assert f'readonly DD_RUNTIME_SHA256="{plan.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") @@ -198,12 +205,12 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( ), ), ), + python_executable="/shared/data-designer/bin/python", host_destination=requested, - container_destination=requested, num_partitions=plan.output.partitions, ) - assert destination.mount == ContainerMount(source=workspace_root, target=workspace_root, read_only=False) + assert destination.authorized_root == workspace_root assert resolver.validate_persisted(plan, collection) == destination @@ -237,7 +244,7 @@ def test_destination_allows_current_run_output_subtree(multi_node_plan: Resolved assert CollectionDestinationResolver().resolve(multi_node_plan, requested).host_path == requested -def test_destination_requires_one_unique_most_specific_mapping( +def test_destination_ignores_duplicate_container_targets_for_one_host_root( multi_node_plan: ResolvedSlurmRunPlan, ) -> None: requested = "/host/exports/run-001" @@ -247,5 +254,4 @@ def test_destination_requires_one_unique_most_specific_mapping( ) plan = multi_node_plan.model_copy(update={"container_mounts": mounts}) - with pytest.raises(StateConflictError, match="ambiguous writable mount"): - CollectionDestinationResolver().resolve(plan, requested) + assert CollectionDestinationResolver().resolve(plan, requested).authorized_root == "/host/exports" diff --git a/packages/data-designer-slurm/tests/state/golden/collection_plan.json b/packages/data-designer-slurm/tests/state/golden/collection_plan.json index 1d528f0d1..1c02ee94a 100644 --- a/packages/data-designer-slurm/tests/state/golden/collection_plan.json +++ b/packages/data-designer-slurm/tests/state/golden/collection_plan.json @@ -1,6 +1,5 @@ { "collection_id": "collection-0001", - "container_destination": "/mnt/data/run-0001.parquet", "created_at": "2026-08-18T12:06:00Z", "host_destination": "/workspace/runs/run-0001/collection/run-0001.parquet", "num_partitions": 1, @@ -14,6 +13,7 @@ } } ], + "python_executable": "/shared/data-designer/bin/python", "resolved_plan": { "path": "/workspace/runs/run-0001/resolved-plan.json", "sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" diff --git a/packages/data-designer-slurm/tests/state/test_model_edges.py b/packages/data-designer-slurm/tests/state/test_model_edges.py index 405102cc5..148c34f25 100644 --- a/packages/data-designer-slurm/tests/state/test_model_edges.py +++ b/packages/data-designer-slurm/tests/state/test_model_edges.py @@ -151,6 +151,14 @@ def test_collection_plan_rejects_duplicate_shards_and_winner_paths() -> None: 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"), ( diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index db0ceb27e..49f3bdf44 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -4,9 +4,12 @@ from __future__ import annotations import csv +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 @@ -750,7 +753,15 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( with pytest.raises(StateConflictError, match="requested destination"): coordinator.submit(destination=Path(case.plan.output.root).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 "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 = Path(case.plan.output.root).parent / submitted.staging_directory @@ -1445,6 +1456,230 @@ def collide( 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 = Path(case.plan.output.root) + 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 = Path(case.plan.output.root) + 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, diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 057ba5bf9..540c20082 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "16460710f383bb9f4650f17e3bac95e4e4a495ca3d80873d7282aa3953533f97" +_COMPATIBILITY_FIXTURE_DIGEST = "922a14db7134472ac63f5929e74afdc00b611b522c4ac15d087a12139fb5e8f1" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) From ec00abaa10b0e26718c25b6ede80080e7c1ff6ad Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 29 Sep 2026 15:03:10 -0600 Subject: [PATCH 31/62] feat(slurm): harden inference proxy Stream allocation-local inference responses over pooled connections with bounded replica failover and aggregate health. Preserve backend overload signals, validate client-image aiohttp compatibility during planning, and document the image requirement. Closes #966 Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/images.mdx | 2 +- packages/data-designer-slurm/pyproject.toml | 1 + .../src/data_designer/slurm/images/service.py | 28 +- .../src/data_designer/slurm/runtime/proxy.py | 511 ++++++++++++------ .../src/data_designer/slurm/runtime/steps.py | 2 + .../golden/client_image_inspection.json | 4 + .../contracts/golden/dependency_lock.json | 4 + .../golden/dependency_lock_single.json | 4 + .../contracts/golden/multi_node_plan.json | 6 +- .../contracts/golden/single_node_plan.json | 6 +- .../tests/images/test_inspection.py | 1 + .../tests/images/test_lifecycle.py | 20 +- .../tests/images/test_registry.py | 22 +- .../golden/finalization_chain.json | 12 +- .../tests/runtime/test_backpressure.py | 1 + .../tests/runtime/test_proxy.py | 489 +++++++++++------ .../tests/runtime/test_steps.py | 2 + .../tests/services/test_image_wiring.py | 5 +- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_state_golden_records.py | 2 +- scripts/test_slurm_package_install.py | 4 + uv.lock | 2 + 24 files changed, 805 insertions(+), 331 deletions(-) diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index 0c6f921a0..becb947a9 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -77,4 +77,4 @@ Published OCI imports are content-addressed. Identical imports can share an arti ## Compatibility rules -The client image inventory must contain the same versions of `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` used to submit the run. Dependency overlays may add pure Python wheels, but cannot replace these protected distributions. Serving images must expose a supported vLLM executable and version. +The client image inventory must contain the same versions of `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` used to submit the run. It must also contain `aiohttp>=3.14.3,<4` for the inference proxy; planning rejects an image without a compatible version. Rebuild and replace an existing client image if it lacks that dependency. Dependency overlays may add pure Python wheels, but cannot replace these protected distributions. Serving images must expose a supported vLLM executable and version. diff --git a/packages/data-designer-slurm/pyproject.toml b/packages/data-designer-slurm/pyproject.toml index 326e052ab..5d4a08779 100644 --- a/packages/data-designer-slurm/pyproject.toml +++ b/packages/data-designer-slurm/pyproject.toml @@ -37,6 +37,7 @@ bump = true [tool.hatch.metadata.hooks.uv-dynamic-versioning] dependencies = [ + "aiohttp>=3.14.3,<4", "data-designer=={{ version }}", "packaging>=25,<27", "pip>=25,<27", 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 index 111f85201..64acaa924 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -13,7 +13,16 @@ from dataclasses import dataclass from pathlib import Path -from data_designer.slurm.config import ImageBuildRequest, ImageInspectionRecord, ImageKind, ImageRef +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion + +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildRequest, + ImageInspectionRecord, + ImageKind, + ImageRef, +) 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 @@ -25,6 +34,7 @@ _HASH_BLOCK_SIZE = 1024 * 1024 _ARTIFACT_DIRECTORY_NAME = "artifacts" _TEMPORARY_DIRECTORY_NAME = ".tmp" +_PROXY_HTTP_REQUIREMENT = SpecifierSet(">=3.14.3,<4") class VerifiedImageRegistry: @@ -200,6 +210,8 @@ def resolve_for_planning(self, reference: ImageRef, *, expected_kind: ImageKind) 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, @@ -234,6 +246,20 @@ def _validate_inspection( ) +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) 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 index 402677767..82de94e16 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -1,22 +1,34 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Allocation-local least-connections HTTP endpoint for one model alias.""" +"""Allocation-local streaming HTTP endpoint for one model alias.""" from __future__ import annotations import argparse -import http.client -import threading -from collections.abc import Sequence +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 http.server import BaseHTTPRequestHandler, ThreadingHTTPServer 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 -_MAXIMUM_REQUEST_BYTES = 64 * 1024 * 1024 -_MAXIMUM_RESPONSE_BYTES = 64 * 1024 * 1024 +_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", @@ -37,155 +49,347 @@ 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 - self._lock = threading.Lock() def acquire(self, excluded: frozenset[int]) -> int: - with self._lock: - 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 = tuple(index for index in available if self._active[index] == minimum) - index = candidates[self._cursor % len(candidates)] - self._cursor += 1 - self._active[index] += 1 - return index + 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: - with self._lock: - self._active[index] -= 1 - - -class _ProxyServer(ThreadingHTTPServer): - daemon_threads = True - - def __init__(self, address: tuple[str, int], pool: _BackendPool) -> None: - self.pool = pool - super().__init__(address, _ProxyHandler) - - -class _ProxyHandler(BaseHTTPRequestHandler): - server: _ProxyServer - protocol_version = "HTTP/1.1" - - def do_GET(self) -> None: # noqa: N802 - if self.path == "/health": - self._write_response(200, b'{"status":"ok"}\n', {"Content-Type": "application/json"}) - return - self._forward() - - def do_POST(self) -> None: # noqa: N802 - self._forward() - - def do_DELETE(self) -> None: # noqa: N802 - self._forward() - - def do_PUT(self) -> None: # noqa: N802 - self._forward() - - def log_message(self, format: str, *args: object) -> None: - del format, args - - def _forward(self) -> None: - if not self.path.startswith("/") or self.path.startswith("//"): - self.send_error(400) - return - request_body = self._read_request_body() - if request_body is None: - return - final_response = self._try_backends(request_body) - if final_response is None: - self.send_error(503) - return - if final_response[0] == 429: - final_response = ( - final_response[0], - final_response[1], - _retry_after_headers(final_response[2], self.server.pool.retry_after_seconds), - ) - self._write_response(*final_response) - - def _try_backends(self, request_body: bytes) -> tuple[int, bytes, dict[str, str]] | 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, 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() - final_response: tuple[int, bytes, dict[str, str]] | None = None - while len(excluded) < len(self.server.pool.backends): - index = self.server.pool.acquire(frozenset(excluded)) + 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: - final_response = self._request_backend(self.server.pool.backends[index], request_body) - finally: - self.server.pool.release(index) - if final_response[0] != 429: - break - excluded.add(index) - return final_response - - def _read_request_body(self) -> bytes | None: - if self.headers.get("Transfer-Encoding") is not None: - self.send_error(501) - return None - lengths = self.headers.get_all("Content-Length", failobj=[]) - if len(lengths) > 1: - self.send_error(400) - return None - raw_length = lengths[0] if lengths else "0" - try: - length = int(raw_length) - except ValueError: - self.send_error(400) - return None - if length < 0 or length > _MAXIMUM_REQUEST_BYTES: - self.send_error(413) - return None - content = self.rfile.read(length) - if len(content) != length: - self.send_error(400) - return None - return content - - def _request_backend(self, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: - request_headers = tuple(self.headers.items()) - hop_headers = _get_hop_headers(request_headers) - headers = { - name: value - for name, value in request_headers - if name.lower() not in hop_headers and name.lower() not in {"host", "content-length"} - } - headers["Content-Length"] = str(len(body)) - connection = http.client.HTTPConnection(backend.host, backend.port, timeout=300) + response = await self._request_backend(request, backend, body, index) + except (ClientError, TimeoutError) as error: + self.metrics.connection_failed += 1 + reason = "timeout" if isinstance(error, 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: - connection.request(self.command, self.path, body=body, headers=headers) - response = connection.getresponse() - payload = response.read(_MAXIMUM_RESPONSE_BYTES + 1) - if len(payload) > _MAXIMUM_RESPONSE_BYTES: - return 502, b'{"error":"backend response too large"}\n', {"Content-Type": "application/json"} - raw_response_headers = tuple(response.getheaders()) - response_hop_headers = _get_hop_headers(raw_response_headers) - response_headers = { - name: value - for name, value in raw_response_headers - if name.lower() not in response_hop_headers and name.lower() != "content-length" - } - return response.status, payload, response_headers - except (OSError, http.client.HTTPException): - return 502, b'{"error":"backend unavailable"}\n', {"Content-Type": "application/json"} + result = await self._stream_response(request, final_response) + self._record_final(final_response.status, outcomes, started_at) + return result finally: - connection.close() + 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 _write_response(self, status: int, body: bytes, headers: dict[str, str]) -> None: - self.send_response(status) - for name, value in headers.items(): - self.send_header(name, value) - self.send_header("Content-Length", str(len(body))) - self.send_header("Connection", "close") - self.end_headers() - self.wfile.write(body) +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: @@ -195,6 +399,7 @@ def main(arguments: Sequence[str] | None = None) -> 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",))) @@ -205,9 +410,18 @@ def main(arguments: Sequence[str] | None = None) -> int: 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") - pool = _BackendPool(backends, parsed.retry_after_seconds) - with _ProxyServer(("127.0.0.1", parsed.listen_port), pool) as server: - server.serve_forever(poll_interval=0.1) + 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, + ) return 0 @@ -234,18 +448,11 @@ def _parse_backend(value: str, allowed_hosts: frozenset[str] = frozenset({"127.0 return _Backend(host, port) -def _retry_after_headers(headers: dict[str, str], retry_after_seconds: int | None) -> dict[str, str]: - selected = {name: value for name, value in headers.items() if name.lower() != "retry-after"} - if retry_after_seconds is not None: - selected["Retry-After"] = str(retry_after_seconds) - return selected - - def _get_hop_headers(headers: Sequence[tuple[str, str]]) -> frozenset[str]: nominated = { - token.strip().lower() + token.strip().casefold() for name, value in headers - if name.lower() == "connection" + if name.casefold() == "connection" for token in value.split(",") if token.strip() } 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 index 7492ad685..1e706828c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -337,6 +337,8 @@ def build_endpoint_command( container_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)), 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 index 5e83b994b..17f97038c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json +++ b/packages/data-designer-slurm/tests/contracts/golden/client_image_inspection.json @@ -1,6 +1,10 @@ { "inspection": { "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, { "name": "data-designer", "version": "0.9.2" diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json index d776429b0..c8cb7d7d9 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -5,6 +5,10 @@ "authored_source": null, "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "image_distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, { "name": "data-designer", "version": "0.9.2" 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 index d7d0836c3..2219e03cf 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -3,6 +3,10 @@ "authored_source": null, "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", "image_distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, { "name": "data-designer", "version": "0.9.2" 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 index a16835e3f..f6cddb883 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -55,7 +55,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "fb536259ea0c206eca90ee8a4ad0a95345f97eb869e349732d1d19db857ecbc3" + "sha256": "f2eea2f07a52e23a60b7d5de317fbc6c6762456a8467401d620d430fe8fe6344" }, "gpu_count": 0, "host_node_index": 0, @@ -67,6 +67,10 @@ "inspection": { "inspection": { "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, { "name": "data-designer", "version": "0.9.2" 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 index 2ac3d4def..ad8a6773c 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -42,7 +42,7 @@ }, "dependency_lock": { "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "ce671fcd812110bc09bd39a21c2c294754c8255d42030c1b90e250a1b513a311" + "sha256": "abd08d979eaf61f088680847887c8c0def48f6210b532cb1c08b445120f57742" }, "gpu_count": 0, "host_node_index": 0, @@ -54,6 +54,10 @@ "inspection": { "inspection": { "distributions": [ + { + "name": "aiohttp", + "version": "3.14.3" + }, { "name": "data-designer", "version": "0.9.2" diff --git a/packages/data-designer-slurm/tests/images/test_inspection.py b/packages/data-designer-slurm/tests/images/test_inspection.py index c869d6c0c..0b682c6e7 100644 --- a/packages/data-designer-slurm/tests/images/test_inspection.py +++ b/packages/data-designer-slurm/tests/images/test_inspection.py @@ -24,6 +24,7 @@ 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"), diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index f16d18c09..c881ab1a6 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -1569,15 +1569,19 @@ def _write_completed_inspection( def _get_client_inspection(content: bytes) -> ImageInspectionRecord: - distributions = tuple( - InstalledDistribution(name=name, version="0.9.2") - for name in ( - "data-designer", - "data-designer-config", - "data-designer-engine", - "data-designer-slurm", + 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"),) + + (InstalledDistribution(name="pip", version="26.1"),) + ) return ImageInspectionRecord( schema_version=1, inspector_version="inspector-1", diff --git a/packages/data-designer-slurm/tests/images/test_registry.py b/packages/data-designer-slurm/tests/images/test_registry.py index 9fa43ea38..93c794985 100644 --- a/packages/data-designer-slurm/tests/images/test_registry.py +++ b/packages/data-designer-slurm/tests/images/test_registry.py @@ -66,6 +66,21 @@ def test_register_and_resolve_existing_sqsh_by_alias_and_path(tmp_path: Path) -> 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"): @@ -304,9 +319,10 @@ def _write_sqsh(path: Path, content: bytes) -> Path: return path -def _get_client_environment() -> FakeInspectionEnvironment: +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"), @@ -318,8 +334,8 @@ def _get_client_environment() -> FakeInspectionEnvironment: ) -def _inspect_client(path: Path) -> ImageInspectionRecord: - return ClientImageInspector(_get_client_environment()).inspect(compute_sqsh_file_sha256(path)) +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: diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 6798f7507..5e116d9a5 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" + "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" + "sha256": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3", + "provenance_digest": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" + "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -90,7 +90,7 @@ "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" + "sha256": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "ff5ff7cc468c22e2d91302b5f400b2275cf045b4d1366ff23e237e363f55d7b0" + "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/runtime/test_backpressure.py b/packages/data-designer-slurm/tests/runtime/test_backpressure.py index 4478a3690..e0235d97f 100644 --- a/packages/data-designer-slurm/tests/runtime/test_backpressure.py +++ b/packages/data-designer-slurm/tests/runtime/test_backpressure.py @@ -16,6 +16,7 @@ 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", diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py index 5387aa4ec..f3364d78f 100644 --- a/packages/data-designer-slurm/tests/runtime/test_proxy.py +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -4,174 +4,353 @@ from __future__ import annotations import argparse -import io -from email.message import Message -from types import MethodType +import asyncio +import socket +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager import pytest +from aiohttp import ClientSession, 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, _ProxyHandler - - -def test_proxy_retries_overload_on_another_least_active_backend() -> None: - handler = object.__new__(_ProxyHandler) - pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) - handler.server = type("Server", (), {"pool": pool})() - handler.path = "/v1/chat/completions" - handler.command = "POST" - handler.rfile = io.BytesIO(b"{}") - handler.headers = Message() - handler.headers["Content-Length"] = "2" - calls: list[int] = [] - captured: list[tuple[int, bytes, dict[str, str]]] = [] - - def request_backend(self: _ProxyHandler, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: - del self - calls.append(backend.port) - assert body == b"{}" - if backend.port == 8001: - return 429, b"overloaded", {} - return 200, b"complete", {"Content-Type": "application/json"} - - def write_response(self: _ProxyHandler, status: int, body: bytes, headers: dict[str, str]) -> None: - del self - captured.append((status, body, headers)) - - handler._request_backend = MethodType(request_backend, handler) - handler._write_response = MethodType(write_response, handler) - - handler._forward() - - assert calls == [8001, 8002] - assert captured == [(200, b"complete", {"Content-Type": "application/json"})] - - -def test_proxy_passes_final_overload_through_after_one_backend_pass() -> None: - handler = object.__new__(_ProxyHandler) - pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 3) - handler.server = type("Server", (), {"pool": pool})() - handler.path = "/v1/chat/completions" - handler.command = "POST" - handler.rfile = io.BytesIO(b"{}") - handler.headers = Message() - handler.headers["Content-Length"] = "2" - calls: list[int] = [] - captured: list[tuple[int, bytes, dict[str, str]]] = [] - - def request_backend(self: _ProxyHandler, backend: _Backend, body: bytes) -> tuple[int, bytes, dict[str, str]]: - del self, body - calls.append(backend.port) - return 429, f"overloaded-{backend.port}".encode(), {"Retry-After": "99"} - - handler._request_backend = MethodType(request_backend, handler) - handler._write_response = MethodType( - lambda self, status, body, headers: captured.append((status, body, headers)), - handler, - ) - - handler._forward() +from data_designer.slurm.runtime.proxy import _Backend, _BackendPool, _parse_backend, _ProxyApplication - assert calls == [8001, 8002] - assert captured == [(429, b"overloaded-8002", {"Retry-After": "3"})] +_Handler = Callable[[web.Request], Awaitable[web.StreamResponse]] -def test_proxy_rejects_a_truncated_fixed_length_request() -> None: - handler = object.__new__(_ProxyHandler) - handler.rfile = io.BytesIO(b"{") - handler.headers = Message() - handler.headers["Content-Length"] = "2" - errors: list[int] = [] - handler.send_error = MethodType(lambda self, status: errors.append(status), handler) +@asynccontextmanager +async def _serve(application: web.Application) -> AsyncIterator[TestServer]: + server = TestServer(application) + await server.start_server() + try: + yield server + finally: + await server.close() - assert handler._read_request_body() is None - assert errors == [400] +def _application(handler: _Handler) -> web.Application: + application = web.Application() + application.router.add_route("*", "/{path:.*}", handler) + return application -def test_proxy_bounds_backend_response_buffering(monkeypatch: pytest.MonkeyPatch) -> None: - class _Response: - status = 200 - def read(self, amount: int) -> bytes: - assert amount == 5 - return b"12345" +def _backend(server: TestServer) -> _Backend: + assert server.host is not None + assert server.port is not None + return _Backend(server.host, server.port) - def getheaders(self) -> list[tuple[str, str]]: - return [] - class _Connection: - def __init__(self, host: str, port: int, timeout: int) -> None: - del host, port, timeout +@pytest.mark.asyncio +async def test_proxy_retries_overload_on_another_least_active_backend() -> None: + calls: list[str] = [] - def request(self, method: str, path: str, *, body: bytes, headers: dict[str, str]) -> None: - del method, path, body, headers - - def getresponse(self) -> _Response: - return _Response() - - def close(self) -> None: - pass - - handler = object.__new__(_ProxyHandler) - handler.command = "POST" - handler.path = "/v1/chat/completions" - handler.headers = Message() - monkeypatch.setattr(runtime_proxy, "_MAXIMUM_RESPONSE_BYTES", 4) - monkeypatch.setattr(runtime_proxy.http.client, "HTTPConnection", _Connection) - - assert handler._request_backend(_Backend("127.0.0.1", 8001), b"{}") == ( - 502, - b'{"error":"backend response too large"}\n', - {"Content-Type": "application/json"}, - ) + 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"}) -def test_proxy_strips_connection_nominated_headers_in_both_directions(monkeypatch: pytest.MonkeyPatch) -> None: - class _Response: - status = 200 + 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"} - def read(self, amount: int) -> bytes: - del amount - return b"complete" + assert calls == ["overloaded", "available"] - def getheaders(self) -> list[tuple[str, str]]: - return [ - ("Connection", "X-Backend-Hop"), - ("X-Backend-Hop", "private"), - ("Content-Type", "application/json"), - ] - class _Connection: - forwarded_headers: dict[str, str] = {} - - def __init__(self, host: str, port: int, timeout: int) -> None: - del host, port, timeout - - def request(self, method: str, path: str, *, body: bytes, headers: dict[str, str]) -> None: - del method, path, body - type(self).forwarded_headers = headers - - def getresponse(self) -> _Response: - return _Response() - - def close(self) -> None: - pass - - handler = object.__new__(_ProxyHandler) - handler.command = "POST" - handler.path = "/v1/chat/completions" - handler.headers = Message() - handler.headers["Connection"] = "X-Client-Hop" - handler.headers["X-Client-Hop"] = "private" - handler.headers["Proxy-Connection"] = "keep-alive" - handler.headers["Authorization"] = "Bearer reviewed" - monkeypatch.setattr(runtime_proxy.http.client, "HTTPConnection", _Connection) - - status, body, headers = handler._request_backend(_Backend("127.0.0.1", 8001), b"{}") +@pytest.mark.asyncio +async def test_proxy_limits_failover_to_three_distinct_backends_and_preserves_final_retry_after() -> None: + calls: list[int] = [] - assert (status, body) == (200, b"complete") - assert _Connection.forwarded_headers == {"Authorization": "Bearer reviewed", "Content-Length": "2"} - assert headers == {"Content-Type": "application/json"} + 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( @@ -195,11 +374,17 @@ def test_proxy_matches_allowed_backend_hosts_case_insensitively() -> None: assert backend == _Backend("compute-001", 8000) -def test_pool_selects_least_active_backend() -> None: - pool = _BackendPool((_Backend("127.0.0.1", 8001), _Backend("127.0.0.1", 8002)), 1) +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(second) pool.release(first) + third = pool.acquire(frozenset()) + pool.release(second) + pool.release(third) - assert first != second + assert (first, second, third) == (0, 1, 2) + assert pool.active_counts() == (0, 0, 0) diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index 2e5f025e9..e5e6a3fb8 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -115,6 +115,8 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti 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: diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index e7b858353..89d95b7a0 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -390,7 +390,10 @@ def _write_inspection(workspace: Path, content: bytes) -> None: "data-designer-slurm", ) ) - + (InstalledDistribution(name="pip", version="26.1"),), + + ( + InstalledDistribution(name="aiohttp", version="3.14.3"), + InstalledDistribution(name="pip", version="26.1"), + ), installer_path="/usr/bin/pip", installer_version="26.1", ), diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 296eb0511..3afccceab 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -16,7 +16,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="69ff0f0e29331a513d732f9ad9f5fc979f5ed833b75061d8fdf769feff96ab11" +readonly DD_PLAN_SHA256="302dfcdcdd61b12f7e227dce10a307e2205e0165ec66969ab04d1e7a028790c1" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 95d4f45c6..7a471e039 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -16,7 +16,7 @@ export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" 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="b3f09f41e45348794055bd3c006c0dbc2c1c741f433b957676f99be24d38b7e3" +readonly DD_PLAN_SHA256="7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 6b3df11a6..29ef2837e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="dd74e01e7e2471a5c91343fe6fe771af4dd19976670fe3513e65e0d6568d8993", + expected_fixture_sha256="202f1070fe8351f8afd3f71ca0353d431b8b4b66c1643e9767a2650c20efbe20", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="5a76d1d73bc96112484689064e1cb9aad71a8e86970cc2f728c74ba47d415f9f", + expected_fixture_sha256="ecffd8264aa022dce66c2eeee447fd71babb507287cfe059f35a9fd1ae009953", ) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 540c20082..1f0ba524f 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "922a14db7134472ac63f5929e74afdc00b611b522c4ac15d087a12139fb5e8f1" +_COMPATIBILITY_FIXTURE_DIGEST = "4283c4f3b6d5d5181c91651ad10377a43737f1c9d86949e76e9eb4d37626ed0e" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index 7da8135ef..a6809e178 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -263,6 +263,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non 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 @@ -278,6 +279,7 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non ) 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" @@ -387,6 +389,7 @@ def main() -> None: 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}" @@ -396,6 +399,7 @@ def main() -> None: 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": ""}) diff --git a/uv.lock b/uv.lock index 84e85921e..3e6d228f9 100644 --- a/uv.lock +++ b/uv.lock @@ -951,6 +951,7 @@ requires-dist = [ name = "data-designer-slurm" source = { editable = "packages/data-designer-slurm" } dependencies = [ + { name = "aiohttp" }, { name = "data-designer" }, { name = "packaging" }, { name = "pip" }, @@ -960,6 +961,7 @@ dependencies = [ [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" }, From 981d0a5f12a0a9ee1a62a4616057c9d4fa6a2be8 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 5 Oct 2026 10:25:04 -0600 Subject: [PATCH 32/62] feat(slurm): run client natively with serving image (#970) * feat(slurm): run client natively with serving image Remove the generation client-image requirement and bind native client and proxy steps to the installed Python runtime. Resolve official versioned vLLM sources to immutable digests before import and verify the inspected version before alias publication. Update Slurm examples, documentation, and regression coverage. Closes #969 Signed-off-by: Nabin Mulepati * preserve actionable runtime inspection errors * ignore submit-only pythonpath in runtime inventory --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/cli-reference.mdx | 2 +- .../pages/slurm/dependencies-and-tools.mdx | 16 +- .../latest/pages/slurm/getting-started.mdx | 46 ++-- fern/versions/latest/pages/slurm/images.mdx | 49 ++--- fern/versions/latest/pages/slurm/overview.mdx | 8 +- .../latest/pages/slurm/run-configuration.mdx | 6 +- .../latest/pages/slurm/troubleshooting.mdx | 3 +- .../data-designer-slurm/examples/run.yaml | 3 - .../src/data_designer/slurm/cli.py | 35 ++- .../slurm/client/dependencies.py | 74 +++---- .../data_designer/slurm/client/environment.py | 200 ++++++++++-------- .../data_designer/slurm/client/execution.py | 30 ++- .../src/data_designer/slurm/client/records.py | 4 +- .../src/data_designer/slurm/client/runtime.py | 99 +++++++++ .../src/data_designer/slurm/config/builder.py | 8 +- .../src/data_designer/slurm/config/run.py | 4 +- .../src/data_designer/slurm/contracts.py | 2 + .../src/data_designer/slurm/images/service.py | 15 +- .../data_designer/slurm/images/vllm_source.py | 112 ++++++++++ .../data_designer/slurm/launcher/renderer.py | 13 +- .../data_designer/slurm/planning/__init__.py | 2 + .../data_designer/slurm/planning/compiler.py | 2 +- .../data_designer/slurm/planning/models.py | 62 +++--- .../slurm/planning/resolution.py | 36 ++-- .../slurm/planning/validation.py | 14 +- .../data_designer/slurm/runtime/bootstrap.py | 41 ++-- .../data_designer/slurm/runtime/context.py | 21 +- .../data_designer/slurm/runtime/entrypoint.py | 15 +- .../data_designer/slurm/runtime/entrypoint.sh | 16 +- .../data_designer/slurm/runtime/manifest.py | 19 +- .../slurm/runtime/plan_reader.sh | 19 +- .../data_designer/slurm/runtime/preflight.py | 29 ++- .../slurm/runtime/server_manifest.py | 2 + .../slurm/runtime/step_runner.sh | 32 ++- .../src/data_designer/slurm/runtime/steps.py | 49 ++--- .../data_designer/slurm/services/wiring.py | 11 +- .../tests/client/conftest.py | 19 +- .../client/test_dependency_resolution.py | 32 +-- .../tests/client/test_environment.py | 147 ++++++++++--- .../tests/client/test_runtime.py | 107 ++++++++++ .../tests/client/test_worker.py | 9 +- .../tests/config/test_loading_builder.py | 7 +- .../tests/contracts/golden/authored_run.json | 4 - .../contracts/golden/authored_run_single.json | 4 - .../contracts/golden/dependency_lock.json | 4 +- .../golden/dependency_lock_single.json | 4 +- .../contracts/golden/multi_node_plan.json | 90 ++++---- .../contracts/golden/single_node_plan.json | 90 ++++---- .../tests/contracts/test_planning_records.py | 10 +- .../tests/images/test_lifecycle.py | 44 ++++ .../tests/images/test_vllm_source.py | 99 +++++++++ .../golden/finalization_chain.json | 14 +- .../tests/launcher/test_collection.py | 15 +- .../tests/launcher/test_renderer.py | 2 +- .../tests/planning/test_compiler.py | 20 +- .../tests/runtime/test_bootstrap.py | 42 +++- .../tests/runtime/test_context.py | 19 +- .../tests/runtime/test_entrypoint.py | 61 +++--- .../tests/runtime/test_preflight.py | 34 ++- .../tests/runtime/test_shell_runtime.py | 26 ++- .../tests/runtime/test_steps.py | 11 +- .../tests/services/test_services.py | 39 ++-- .../tests/services/test_wiring.py | 37 +++- .../tests/serving/test_resolver.py | 10 +- .../tests/serving/test_serving_records.py | 6 +- .../golden/rendered/multi_node.sbatch | 3 +- .../golden/rendered/single_node.sbatch | 3 +- .../slurm_test_fakes/test_dependencies.py | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_state_golden_records.py | 2 +- .../data-designer-slurm/tests/test_cli.py | 34 ++- 71 files changed, 1398 insertions(+), 755 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/images/vllm_source.py create mode 100644 packages/data-designer-slurm/tests/client/test_runtime.py create mode 100644 packages/data-designer-slurm/tests/images/test_vllm_source.py diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx index 8566d34fa..07a48ab9c 100644 --- a/fern/versions/latest/pages/slurm/cli-reference.mdx +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -8,7 +8,7 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | --- | --- | | `profile init` | Create a restrictive starter profile catalog without overwriting. | | `profile validate` | Validate catalog selection, shared workspace access, and scheduler facts. | -| `image add` | Import or inspect an image through a CPU Slurm job and publish an alias. | +| `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 without deleting its SQSH artifact. | diff --git a/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx b/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx index 25f2affcd..157d9a196 100644 --- a/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx +++ b/fern/versions/latest/pages/slurm/dependencies-and-tools.mdx @@ -1,13 +1,11 @@ # Dependencies, plugins, MCP, and secrets -The client image provides Data Designer and its base environment. A run may add immutable pure Python wheels as an allocation-local overlay. The overlay is reconstructed for each allocation and is not written back into the image. +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: - image: - name: dd-client dependencies: requirements: - example-plugin==1.2.3 @@ -15,7 +13,7 @@ client: 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. -Do not use the overlay to replace `data-designer`, `data-designer-config`, `data-designer-engine`, or `data-designer-slurm`. Build a new client image when these packages need to change. +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 @@ -23,18 +21,16 @@ Use exactly one of `requirements` or `lock_file`: ```yaml client: - image: - name: dd-client dependencies: requirements: null lock_file: inputs/dependency-lock.json ``` -The lock must target the registered client image inventory and reference available wheel artifacts. +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, bake the package into the client image, use an immutable HTTPS wheel URL that does not embed credentials, or supply a previously resolved compatible lock and wheel artifacts. +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 @@ -70,7 +66,7 @@ Export referenced variables in the trusted submit environment before running `ex ## Use plugins -Install plugin packages into the client image 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. +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 @@ -110,6 +106,6 @@ invocation: environment: TOOL_API_KEY ``` -The command is one executable token and arguments are passed without a shell. The client image or dependency overlay must provide the executable. See [Tool Use and MCP](/concepts/tool-use-mcp/overview) for builder-side tool configuration. +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 index f5cd08449..d97c82a2f 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -8,15 +8,14 @@ Obtain these site-specific values from your Slurm or platform administrator befo - A Python 3.10 or later executable - An account, GPU partition, CPU image-build partition, and shared workspace path -- A digest-qualified client OCI image or absolute SQSH path containing the same Data Designer package versions as the submit environment -- A digest-qualified serving OCI image or absolute SQSH path containing a supported vLLM installation +- 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 image references and `example/generator` model below are syntax examples. They are intentionally not runnable values. +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 -Use a shared path instead of `$HOME` when required by site policy. 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. +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 @@ -41,6 +40,8 @@ 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 @@ -75,33 +76,19 @@ data-designer slurm profile validate --cluster primary Validation checks strict profile loading, cluster selection, workspace access, and scheduler facts. Keep `--cluster primary` explicit when running from a compute host whose hostname does not match the login-host pattern. See [Profiles and clusters](/slurm/profiles-and-clusters) for non-default files and multiple clusters. -## 6. Register a client image - -Use a digest-qualified OCI reference: - -```bash -data-designer slurm image add \ - registry.example.test/data-designer/client@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ - --kind client \ - --name dd-client \ - --cluster primary -``` - -Replace the entire source with the digest-qualified OCI reference or absolute SQSH path supplied for your site. See [Prepare compatible images](/slurm/images#prepare-compatible-images) before registering a client image. - -## 7. Register a serving image +## 6. Register a serving image ```bash data-designer slurm image add \ - registry.example.test/inference/vllm@sha256:1111111111111111111111111111111111111111111111111111111111111111 \ + vllm/vllm-openai:v0.22.0 \ --kind serving \ --name vllm \ --cluster primary ``` -Each `image add` command submits a CPU Slurm job, waits for it to finish, inspects the result, and only then publishes the alias. Keep the terminal open. You can also register an existing absolute `.sqsh` path. See [Images](/slurm/images). +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). -## 8. Create a Data Designer builder config +## 7. Create a Data Designer builder config Create `builder.yaml`: @@ -120,7 +107,7 @@ data_designer: The model alias in this file must match a deployment in the Slurm run config. -## 9. Create a Slurm run config +## 8. Create a Slurm run config Create `run.yaml` next to `builder.yaml`: @@ -134,9 +121,6 @@ invocation: dataset_name: greetings model_concurrency: generator: 8 -client: - image: - name: dd-client deployments: - model_alias: generator served_model_name: example/generator @@ -158,15 +142,15 @@ output: The repository also contains these 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. -## 10. Dry-run the submission +## 9. Dry-run the submission ```bash data-designer slurm execute run.yaml --dry-run --cluster primary ``` -Dry-run resolves the selected profile, verifies registered image contents, loads the builder envelope, resolves client dependencies, and renders the batch script without calling `sbatch`. Full builder and plugin validation runs inside the prepared client 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`. +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`. -## 11. Submit the run +## 10. Submit the run ```bash data-designer slurm execute run.yaml --cluster primary @@ -174,7 +158,7 @@ data-designer slurm execute run.yaml --cluster primary The command prints one JSON object. Save its `run_id`; all lifecycle commands use the managed run ID rather than a raw Slurm job ID. -## 12. Inspect or cancel the run +## 11. Inspect or cancel the run ```bash data-designer slurm status --cluster primary @@ -183,7 +167,7 @@ data-designer slurm cancel --cluster primary Status reconciles active queue and accounting evidence with durable run state. Cancellation targets only jobs recorded for the managed run. -## 13. Inspect results +## 12. Inspect results Each successful shard publishes a canonical winner at: diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index becb947a9..cc37d8d37 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -1,38 +1,35 @@ # Slurm images -Data Designer Slurm uses two image roles: +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. -- `client` images run Data Designer and installed plugins. -- `serving` images run vLLM model servers. +## Choose a serving image -Use separate aliases even if both roles originate from related images. Inspection verifies role-specific contents before an alias becomes available to planning. - -## Prepare compatible images - -Image construction is an operator prerequisite. Data Designer Slurm imports or inspects images, but does not build a client or serving environment from a Dockerfile. - -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. A client image must install the same release used by the submit environment: +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 -python -m pip install "data-designer[slurm]==" +data-designer slurm image add vllm/vllm-openai:v0.22.0 \ + --kind serving \ + --name vllm ``` -For an unpublished build, install the aligned `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` wheels into the client image. Do not mix versions. A serving image must provide a supported vLLM executable and all runtime dependencies for the selected model. +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. 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 digest-pinned OCI image +## 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/data-designer/client@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ - --kind client \ - --name dd-client + registry.example.test/inference/vllm@sha256:0000000000000000000000000000000000000000000000000000000000000000 \ + --kind serving \ + --name vllm ``` -OCI sources must use `repository@sha256:<64 lowercase hex characters>`. Schemes, embedded credentials, tags without a digest, query strings, and fragments are rejected. Registry credentials belong in the compute environment or Enroot configuration, not in the command or profile. +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. @@ -40,9 +37,9 @@ The command stages package-owned inspection code, submits a CPU job using `image ```bash data-designer slurm image add \ - /workspace/alice/images/client.sqsh \ - --kind client \ - --name dd-client + /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. @@ -51,7 +48,7 @@ The absolute SQSH path must be visible on the image-build node and all later com ```bash data-designer slurm image ls -data-designer slurm image info dd-client +data-designer slurm image info vllm ``` 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`. @@ -61,20 +58,20 @@ Aliases are cluster-profile scoped because each profile has its own workspace. R An alias is immutable unless replacement is explicit: ```bash -data-designer slurm image add /workspace/alice/images/client-v2.sqsh \ - --kind client \ - --name dd-client \ +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 dd-client +data-designer slurm image rm vllm ``` Published OCI imports are content-addressed. Identical imports can share an artifact, while conflicting bytes or inspection facts fail closed. ## Compatibility rules -The client image inventory must contain the same versions of `data-designer`, `data-designer-config`, `data-designer-engine`, and `data-designer-slurm` used to submit the run. It must also contain `aiohttp>=3.14.3,<4` for the inference proxy; planning rejects an image without a compatible version. Rebuild and replace an existing client image if it lacks that dependency. Dependency overlays may add pure Python wheels, but cannot replace these protected distributions. Serving images must expose a supported vLLM executable and version. +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/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index 4bd91d2f7..d884db31c 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -12,7 +12,7 @@ The submit environment needs: - `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 Enroot 3.5 or later for existing SQSH images. Digest-pinned OCI imports need Enroot 4.0 or later. The client image must contain the same Data Designer packages as the submit environment. Serving images must contain vLLM. +Compute nodes need access to the same shared Python environment at the exact path used by the submit command. They also need Enroot 3.5 or later for existing SQSH images; digest-pinned OCI imports need Enroot 4.0 or later. 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. @@ -22,7 +22,7 @@ Your Slurm administrator can confirm the account, GPU partition, CPU image-build | --- | --- | | 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 images, start vLLM and client containers, generate records, and write managed state and results. | +| 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. @@ -30,10 +30,10 @@ Run Slurm commands from a site-approved submit environment with the cluster CLI The configured `workspace_root` is durable and shared. It contains image metadata, immutable run inputs, Slurm logs, attempt state, winning results, collection state, and default output. -Runtime bundle extraction, the client dependency overlay, Data Designer home and cache directories, and Enroot state use allocation-local scratch. The runtime mounts that scratch at `/run/data-designer-slurm` inside containers and 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. +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, client and 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. +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/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 225d35aad..d7201563e 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -16,8 +16,6 @@ invocation: generator: 64 client: cpus: 32 - image: - name: dd-client deployments: - model_alias: generator served_model_name: example/generator @@ -46,6 +44,8 @@ output: 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. Output formats are `parquet`, `jsonl`, and `csv`. +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 32 and controls the client's CPU allocation, not its execution image. + ## Configure vLLM Each server supports readiness and launch controls: @@ -171,7 +171,7 @@ slurm = ( dataset_name="generated", model_concurrency={"generator": 16}, ) - .with_client(image={"name": "dd-client"}) + .with_client(cpus=32) .with_deployment( { "model_alias": "generator", diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index 1ccc3e061..f404fece0 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -47,7 +47,8 @@ Set `scheduler.bin_path` in the selected profile to the absolute directory conta - Confirm the image-build partition accepts CPU jobs. - Confirm Enroot is available on the allocated node. - Use Enroot 4.0 or later for OCI import and 3.5 or later for existing SQSH inspection. -- Use a digest-qualified, credential-free OCI reference. +- 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. diff --git a/packages/data-designer-slurm/examples/run.yaml b/packages/data-designer-slurm/examples/run.yaml index 6879a4a96..32158ce8b 100644 --- a/packages/data-designer-slurm/examples/run.yaml +++ b/packages/data-designer-slurm/examples/run.yaml @@ -7,9 +7,6 @@ invocation: dataset_name: greetings model_concurrency: generator: 8 -client: - image: - name: dd-client deployments: - model_alias: generator served_model_name: example/generator diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 8a0ee3b27..eea437382 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -18,6 +18,12 @@ from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config from data_designer.slurm.contracts import canonical_json from data_designer.slurm.images.records import validate_oci_source_for_lifecycle +from data_designer.slurm.images.vllm_source import ( + VllmSourceResolutionError, + VllmTagNotFoundError, + is_versioned_vllm_tag, + resolve_versioned_vllm_source, +) from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, @@ -232,8 +238,8 @@ def profile_validate_command( @image_app.command("add") def image_add_command( - source: str = typer.Argument(...), - kind: str = typer.Option(..., "--kind", help="Image role: client or serving"), + source: str = typer.Argument(..., help="Digest-pinned OCI, versioned official vLLM image, or SQSH"), + kind: str = typer.Option("serving", "--kind", help="Image role (defaults to serving)"), name: str | None = typer.Option(None, "--name"), replace: bool = typer.Option(False, "--replace"), profile_file: Path | None = typer.Option(None, "--profile-file", dir_okay=False), @@ -243,21 +249,38 @@ def image_add_command( operation = SlurmServiceOperation.ADD_IMAGE def add() -> BaseModel: - if not source.endswith(".sqsh") and re.fullmatch(r"[^\s]+@sha256:[0-9a-f]{64}", source) is None: + resolved_source = source + if is_versioned_vllm_tag(source): + if kind != "serving": + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "versioned vLLM images require --kind serving", + ) + 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(source) + 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(source), kind=kind, source=source) + request = ImageBuildRequest(name=name or _derive_image_name(resolved_source), kind=kind, source=resolved_source) return create_slurm_image_service(profile_file=profile_file, cluster=cluster).add(request, replace=replace) _emit_result(_invoke(operation, add)) @@ -353,7 +376,7 @@ 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].split(":", maxsplit=1)[0] + candidate = source.rpartition("@sha256:")[0].rsplit("/", maxsplit=1)[-1].replace(":", "-") return re.sub(r"[^A-Za-z0-9._-]+", "-", candidate).strip("-._") or "image" 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 index 45b7bc39e..aee3843d9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/dependencies.py @@ -20,9 +20,9 @@ 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, ClientImageInspection +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, ResolvedDependencyLock, ResolvedImage +from data_designer.slurm.planning import LockedPackage, ResolvedClientRuntime, ResolvedDependencyLock _RESOLVER_VERSION = "pip-pure-wheel-1" _DEFAULT_INDEX_URL = "https://pypi.org/simple" @@ -46,7 +46,7 @@ class ResolvedClientDependencies: class ClientDependencyResolver: - """Resolve pure wheels against the immutable client-image inventory.""" + """Resolve pure wheels against the immutable native client inventory.""" def __init__( self, @@ -63,7 +63,7 @@ def __init__( def resolve( self, dependencies: ClientDependencies, - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, *, run_root: str | Path, source_root: str | Path, @@ -71,9 +71,6 @@ def resolve( """Yield one verified resolution while temporary wheel sources remain live.""" if not isinstance(dependencies, ClientDependencies): raise ClientDependencyResolutionError("client dependencies are invalid") - inspection = client_image.inspection_facts - if not isinstance(inspection, ClientImageInspection): - raise ClientDependencyResolutionError("client image lacks dependency inspection facts") 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())) @@ -86,14 +83,14 @@ def resolve( if dependencies.requirements is None: result = self._load_lock( dependencies, - client_image, + client_runtime, normalized_run_root, normalized_source_root, ) else: result = self._resolve_requirements( dependencies, - client_image, + client_runtime, normalized_run_root, temporary_root, ) @@ -106,20 +103,17 @@ def resolve( def _resolve_requirements( self, dependencies: ClientDependencies, - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, run_root: Path, temporary_root: Path, ) -> ResolvedClientDependencies: requirements = dependencies.requirements if requirements is None: raise ClientDependencyResolutionError("client dependency requirements are unavailable") - inspection = client_image.inspection_facts - if not isinstance(inspection, ClientImageInspection): - raise ClientDependencyResolutionError("client image lacks dependency inspection facts") - image_distributions = tuple(sorted(inspection.distributions, key=lambda item: item.name)) + base_distributions = tuple(sorted(client_runtime.distributions, key=lambda item: item.name)) if not requirements: return ResolvedClientDependencies( - lock=_build_lock(client_image, image_distributions=image_distributions, authored_requirements=()), + lock=_build_lock(client_runtime, base_distributions=base_distributions, authored_requirements=()), wheel_sources=(), ) @@ -129,7 +123,7 @@ def _resolve_requirements( 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 image_distributions), + "".join(f"{item.name}=={item.version}\n" for item in base_distributions), encoding="utf-8", ) command = ( @@ -144,7 +138,7 @@ def _resolve_requirements( "--platform=any", "--implementation=py", "--abi=none", - f"--python-version={inspection.python_version.rsplit('.', maxsplit=1)[0]}", + 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()}", @@ -157,11 +151,11 @@ def _resolve_requirements( except Exception as error: raise ClientDependencyResolutionError("client dependency resolution failed") from error - packages, sources = _load_downloaded_wheels(wheelhouse, run_root, image_distributions) + packages, sources = _load_downloaded_wheels(wheelhouse, run_root, base_distributions) try: lock = _build_lock( - client_image, - image_distributions=image_distributions, + client_runtime, + base_distributions=base_distributions, authored_requirements=tuple(requirements), overlay_packages=packages, ) @@ -172,7 +166,7 @@ def _resolve_requirements( def _load_lock( self, dependencies: ClientDependencies, - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, run_root: Path, source_root: Path, ) -> ResolvedClientDependencies: @@ -185,15 +179,12 @@ def _load_lock( supplied = ResolvedDependencyLock.model_validate_json(content, strict=True) except Exception as error: raise ClientDependencyResolutionError("client dependency lock is invalid") from error - inspection = client_image.inspection_facts - if not isinstance(inspection, ClientImageInspection): - raise ClientDependencyResolutionError("client image lacks dependency inspection facts") if ( - supplied.client_image_sha256 != client_image.sha256 - or supplied.python_abi != inspection.python_abi - or supplied.image_distributions != tuple(sorted(inspection.distributions, key=lambda item: item.name)) + 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 image") + raise ClientDependencyResolutionError("client dependency lock targets another runtime") packages: list[LockedPackage] = [] sources: list[Path] = [] @@ -220,11 +211,11 @@ def _load_lock( schema_version=1, resolver_version=supplied.resolver_version, python_abi=supplied.python_abi, - client_image_sha256=supplied.client_image_sha256, + client_runtime_sha256=supplied.client_runtime_sha256, authored_requirements=supplied.authored_requirements, authored_source=lock_file, source=source_reference, - image_distributions=supplied.image_distributions, + base_distributions=supplied.base_distributions, overlay_packages=tuple(packages), ) except ValueError as error: @@ -248,22 +239,19 @@ def _resolution_environment(self, dependencies: ClientDependencies) -> dict[str, def _build_lock( - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, *, - image_distributions: tuple[InstalledDistribution, ...], + base_distributions: tuple[InstalledDistribution, ...], authored_requirements: tuple[str, ...], overlay_packages: tuple[LockedPackage, ...] = (), ) -> ResolvedDependencyLock: - inspection = client_image.inspection_facts - if not isinstance(inspection, ClientImageInspection): - raise ClientDependencyResolutionError("client image lacks dependency inspection facts") return ResolvedDependencyLock( schema_version=1, resolver_version=_RESOLVER_VERSION, - python_abi=inspection.python_abi, - client_image_sha256=client_image.sha256, + python_abi=client_runtime.python_abi, + client_runtime_sha256=client_runtime.runtime_sha256, authored_requirements=authored_requirements, - image_distributions=image_distributions, + base_distributions=base_distributions, overlay_packages=overlay_packages, ) @@ -271,9 +259,9 @@ def _build_lock( def _load_downloaded_wheels( wheelhouse: Path, run_root: Path, - image_distributions: tuple[InstalledDistribution, ...], + base_distributions: tuple[InstalledDistribution, ...], ) -> tuple[tuple[LockedPackage, ...], tuple[Path, ...]]: - image = {item.name: item.version for item in image_distributions} + 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): @@ -288,9 +276,9 @@ def _load_downloaded_wheels( 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 image: - if image[normalized_name] != normalized_version: - raise ClientDependencyResolutionError("client dependency conflicts with the selected image") + 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") 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 index 7db43dba8..5144ac409 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -7,8 +7,8 @@ import importlib.metadata import json import os -import posixpath import re +import site import stat import subprocess import sys @@ -23,7 +23,7 @@ from data_designer.slurm.client.errors import ClientWorkerError from data_designer.slurm.client.filesystem import compute_file_sha256, ensure_private_directory, read_regular_bytes from data_designer.slurm.client.records import ClientErrorCode, ClientInstallerOutcome -from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution +from data_designer.slurm.contracts import ArtifactReference, InstalledDistribution, canonical_json DistributionInventory = Callable[[Path | None], tuple[InstalledDistribution, ...]] CommandRunner = Callable[[tuple[str, ...]], None] @@ -38,7 +38,7 @@ class PreparedClientEnvironment: scratch_root: Path overlay_path: Path dependency_lock: ArtifactReference - client_image_sha256: str + client_runtime_sha256: str python_abi: str installer_outcome: ClientInstallerOutcome installed_distributions: tuple[InstalledDistribution, ...] @@ -52,17 +52,16 @@ class _BootstrapInputs: attempt_id: str attempt_dir: Path scratch_root: Path - client_image_sha256: str + client_runtime_sha256: str python_abi: str - installer_path: Path - inspection: dict[str, object] + python_executable: Path + runtime: dict[str, object] dependency_lock: ArtifactReference - plan: dict[str, object] @dataclass(frozen=True) class _VerifiedDependencies: - image_distributions: tuple[InstalledDistribution, ...] + base_distributions: tuple[InstalledDistribution, ...] overlay_packages: tuple[dict[str, object], ...] @@ -105,7 +104,7 @@ def prepare( scratch_root=inputs.scratch_root, overlay_path=overlay_path, dependency_lock=inputs.dependency_lock, - client_image_sha256=inputs.client_image_sha256, + client_runtime_sha256=inputs.client_runtime_sha256, python_abi=inputs.python_abi, installer_outcome=installer_outcome, installed_distributions=installed, @@ -135,28 +134,30 @@ def _load_bootstrap_inputs( plan_path.name != "resolved-plan.json" or logical_run_root.name != run_id or logical_run_root.parent.name != "runs" - or _get_container_path(plan_payload, (logical_run_root / plan_path.name).as_posix()) != plan_path.as_posix() + 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 = Path(_get_container_path(plan_payload, logical_attempt.as_posix(), require_writable=True)) + 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") - image = _require_object(client, "image") - image_sha256 = _require_digest(image, "sha256") - inspection_record = _require_object(image, "inspection") - inspection = _require_object(inspection_record, "inspection") - if inspection.get("kind") != "client": - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "resolved client image inspection is invalid") - python_abi = _require_string(inspection, "python_abi") + 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") - installer_path = Path(_require_string(inspection, "installer_path")) - if not installer_path.is_absolute() or installer_path.as_posix() == "/": - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client installer path is invalid") + 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(): @@ -168,39 +169,40 @@ def _load_bootstrap_inputs( attempt_id=attempt_id, attempt_dir=attempt_dir, scratch_root=scratch_root, - client_image_sha256=image_sha256, + client_runtime_sha256=runtime_sha256, python_abi=python_abi, - installer_path=installer_path, - inspection=inspection, + python_executable=python_executable, + runtime=runtime, dependency_lock=lock_reference, - plan=plan_payload, ) def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependencies: lock_bytes = read_regular_bytes( - Path(_get_container_path(inputs.plan, inputs.dependency_lock.path)), + 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_image_sha256") != inputs.client_image_sha256: - raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "dependency lock targets another image") + 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_image = _parse_distributions(lock.get("image_distributions")) - inspected_image = _parse_distributions(inputs.inspection.get("distributions")) + expected_base = _parse_distributions(lock.get("base_distributions")) + inspected_base = _parse_distributions(inputs.runtime.get("distributions")) try: - actual_image = self._inventory(None) + actual_base = self._inventory(None) except ClientWorkerError: raise except Exception as error: raise ClientWorkerError( - ClientErrorCode.DEPENDENCY_CONFLICT, "client image inventory cannot be verified" + ClientErrorCode.DEPENDENCY_CONFLICT, "client runtime inventory cannot be verified" ) from error - if expected_image != inspected_image or actual_image != expected_image: - raise ClientWorkerError(ClientErrorCode.DEPENDENCY_CONFLICT, "client image inventory differs from the lock") + 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: @@ -208,10 +210,9 @@ def _verify_dependency_lock(self, inputs: _BootstrapInputs) -> _VerifiedDependen _verify_input_artifact( source_reference, inputs.logical_run_root / "inputs", - inputs.plan, ) return _VerifiedDependencies( - image_distributions=expected_image, + base_distributions=expected_base, overlay_packages=tuple(_as_object_list(lock.get("overlay_packages"))), ) @@ -223,17 +224,16 @@ def _prepare_overlay( expected_overlay, wheels = _verify_wheels( dependencies.overlay_packages, inputs.logical_run_root / "dependencies", - dependencies.image_distributions, - inputs.plan, + dependencies.base_distributions, ) overlay_path = inputs.scratch_root / "client-env" / "site-packages" - outcome = self._install_overlay(inputs.installer_path, wheels, expected_overlay, overlay_path) - installed = tuple(sorted((*dependencies.image_distributions, *expected_overlay), key=lambda item: item.name)) + 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, - installer_path: Path, + python_executable: Path, wheels: tuple[Path, ...], expected: tuple[InstalledDistribution, ...], target: Path, @@ -249,7 +249,9 @@ def _install_overlay( try: self._command_runner( ( - installer_path.as_posix(), + python_executable.as_posix(), + "-m", + "pip", "install", "--disable-pip-version-check", "--no-deps", @@ -274,9 +276,8 @@ def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ... """Return an exact immutable distribution inventory.""" if path is not None: importlib.invalidate_caches() - distributions = ( - importlib.metadata.distributions() if path is None else importlib.metadata.distributions(path=[path.as_posix()]) - ) + 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: @@ -287,13 +288,49 @@ def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ... 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, "mutable installed distributions are forbidden" + ClientErrorCode.DEPENDENCY_CONFLICT, + f"mutable installed distribution {name!r} is forbidden", ) names.add(name) - installed.append(InstalledDistribution(name=name, version=version)) + 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 submit-only ``PYTHONPATH`` entries from the client runtime inventory.""" + configured_paths = { + _normalize_import_path(entry) for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep) if entry + } + if not configured_paths: + return sys.path.copy() + + site_paths = site.getsitepackages() + if site.ENABLE_USER_SITE: + user_site_path = site.getusersitepackages() + if isinstance(user_site_path, str): + site_paths.append(user_site_path) + else: + site_paths.extend(user_site_path) + standard_site_paths = {_normalize_import_path(entry) for entry in site_paths} + search_paths: list[str] = [] + for entry in sys.path: + normalized_entry = _normalize_import_path(entry) + if normalized_entry in configured_paths and normalized_entry not in standard_site_paths: + continue + search_paths.append(entry) + return search_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" @@ -339,22 +376,21 @@ def _run_installer(command: tuple[str, ...]) -> None: def _verify_wheels( packages: tuple[dict[str, object], ...], logical_dependencies_root: Path, - image_distributions: tuple[InstalledDistribution, ...], - plan: dict[str, object], + base_distributions: tuple[InstalledDistribution, ...], ) -> tuple[tuple[InstalledDistribution, ...], tuple[Path, ...]]: expected: list[InstalledDistribution] = [] wheels: list[Path] = [] - image_names = {item.name for item in image_distributions} + 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 image_names or name in {item.name for item in expected}: + 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 = Path(_get_container_path(plan, logical_wheel.as_posix())) + 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: @@ -375,11 +411,11 @@ def _verify_wheels( 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, plan: dict[str, object]) -> None: +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 = Path(_get_container_path(plan, logical_path.as_posix())) + 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") @@ -387,7 +423,11 @@ def _verify_input_artifact(reference: ArtifactReference, logical_root: Path, pla def _parse_distributions(value: object) -> tuple[InstalledDistribution, ...]: parsed = tuple( InstalledDistribution( - name=canonicalize_name(_require_string(item, "name")), version=_require_string(item, "version") + 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) ) @@ -448,37 +488,6 @@ def _artifact_reference(value: dict[str, object]) -> ArtifactReference: raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client artifact reference is invalid") from error -def _get_container_path(plan: dict[str, object], host_path: str, *, require_writable: bool = False) -> str: - if not host_path.startswith("/") or posixpath.normpath(host_path) != host_path: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is invalid") - values = plan.get("container_mounts") - if not isinstance(values, list): - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") - candidates: list[tuple[str, str, bool]] = [] - for value in values: - mount = _as_object(value) - source = _require_string(mount, "source") - target = _require_string(mount, "target") - read_only = mount.get("read_only") - if ( - not source.startswith("/") - or posixpath.normpath(source) != source - or not target.startswith("/") - or posixpath.normpath(target) != target - or type(read_only) is not bool - ): - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan mounts are invalid") - if host_path == source or host_path.startswith(f"{source}/"): - candidates.append((source, target, read_only)) - if not candidates: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not mounted") - source, target, read_only = max(candidates, key=lambda item: len(item[0])) - if require_writable and read_only: - raise ClientWorkerError(ClientErrorCode.INVALID_INPUT, "client plan path is not writable") - relative_path = posixpath.relpath(host_path, source) - return target if relative_path == "." else posixpath.join(target, relative_path) - - def _sha256_bytes(value: bytes) -> str: return hashlib.sha256(value).hexdigest() @@ -488,12 +497,27 @@ def _is_mutable_direct_url(value: str) -> bool: direct_url = _as_object(json.loads(value)) except (json.JSONDecodeError, TypeError): return True - if "dir_info" in direct_url or "vcs_info" in direct_url: + 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") - return not isinstance(hashes, dict) or not isinstance(hashes.get("sha256"), str) + 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/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 31011a4dd..0cc43a577 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -51,7 +51,6 @@ 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.paths import get_container_path, get_host_path from data_designer.slurm.runtime.ports import resolve_allocation_plan from data_designer.slurm.state import CandidateOutcome, CandidateOutputFile, CandidateOutputManifest @@ -268,7 +267,7 @@ def _build_context( self._validate_prepared(plan, shard, prepared) lock = ResolvedDependencyLock.model_validate_json( read_regular_bytes( - Path(get_container_path(plan, plan.client.dependency_lock.path)), + Path(plan.client.dependency_lock.path), missing_code=ClientErrorCode.DEPENDENCY_ARTIFACT_MISSING, ) ) @@ -277,7 +276,7 @@ def _build_context( expected_distributions = tuple( sorted( ( - *lock.image_distributions, + *lock.base_distributions, *( InstalledDistribution(name=package.name, version=package.version) for package in lock.overlay_packages @@ -326,14 +325,13 @@ def _validate_prepared( prepared: PreparedClientEnvironment, ) -> None: logical_attempt = Path(shard.resume_workspace.path).parent / "attempts" / prepared.attempt_id - expected_attempt = Path(get_container_path(plan, logical_attempt.as_posix(), require_writable=True)) - inspection = plan.client.image.inspection_facts + 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.image.sha256 != prepared.client_image_sha256 - or inspection.python_abi != prepared.python_abi + 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") @@ -343,7 +341,7 @@ def _load_builder(plan: ResolvedSlurmRunPlan) -> dict[str, object]: return cast(dict[str, object], plan.builder.inline) assert plan.builder.source is not None payload = read_regular_bytes( - Path(get_container_path(plan, plan.builder.source.path)), + Path(plan.builder.source.path), missing_code=ClientErrorCode.INVALID_INPUT, ) if hashlib.sha256(payload).hexdigest() != plan.builder.source.sha256: @@ -421,13 +419,13 @@ def _materialize_seed( 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(get_container_path(plan, seed_path)) + 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(get_container_path(plan, shard.input_partition.path)), + Path(shard.input_partition.path), missing_code=ClientErrorCode.INVALID_INPUT, ) if hashlib.sha256(payload).hexdigest() != shard.input_partition.sha256: @@ -474,7 +472,7 @@ def _materialize_mcp_providers( def _validate_assets(plan: ResolvedSlurmRunPlan) -> Path: value = plan.invocation.effective_input_bindings.managed_assets_path assert value is not None - path = Path(get_container_path(plan, value)) + 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 @@ -486,7 +484,7 @@ def _dataset_path( prepared: PreparedClientEnvironment, resume: ResumeMode, ) -> Path: - resume_path = Path(get_container_path(plan, shard.resume_workspace.path, require_writable=True)) + 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())): @@ -570,7 +568,7 @@ def _validate_creation_result( raise ClientWorkerError( ClientErrorCode.OUTPUT_INVALID, "Data Designer effective resume mode differs from retry intent" ) - shared_path = Path(get_container_path(context.plan, context.shard.resume_workspace.path, require_writable=True)) + 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() @@ -650,7 +648,7 @@ def _build_candidate_manifest( attempt_id=prepared.attempt_id, attempt_ordinal=int(prepared.attempt_id.removeprefix("attempt-")), created_at=created_at, - dataset_path=get_host_path(context.plan, dataset_path.as_posix(), require_writable=True), + dataset_path=dataset_path.as_posix(), requested_records=creation.requested_records, actual_records=creation.actual_records, outcome=( @@ -692,7 +690,7 @@ def _publish_success( requested_resume_mode=context.requested_resume.value, effective_resume_mode=creation.effective_resume.value, candidate_output_manifest=ArtifactReference( - path=get_host_path(context.plan, candidate_path.as_posix(), require_writable=True), + path=candidate_path.as_posix(), sha256=candidate.compute_sha256(), ), ) @@ -720,7 +718,7 @@ def _validate_environment_manifest( 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_image_sha256 != prepared.client_image_sha256 + 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 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 index 33eb2dcfa..ce8cf1cd7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/records.py @@ -81,7 +81,7 @@ class ClientEnvironmentManifest(ContractRecord): created_at: datetime outcome: ClientEnvironmentOutcome dependency_lock: ArtifactReference - client_image_sha256: Sha256Digest + client_runtime_sha256: Sha256Digest python_abi: Identifier overlay_path: str installer_outcome: ClientInstallerOutcome @@ -111,7 +111,7 @@ def from_prepared( created_at=created_at, outcome=outcome, dependency_lock=prepared.dependency_lock, - client_image_sha256=prepared.client_image_sha256, + client_runtime_sha256=prepared.client_runtime_sha256, python_abi=prepared.python_abi, overlay_path=prepared.overlay_path.as_posix(), installer_outcome=prepared.installer_outcome, 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..bfcd6774d --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py @@ -0,0 +1,99 @@ +# 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: {', '.join(sorted(missing))}" + ) + 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/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 679d818b3..ccdecd384 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -15,7 +15,6 @@ 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.images import ImageRef from data_designer.slurm.config.run import ( ArrayTasksConfig, BuilderInput, @@ -98,7 +97,6 @@ def with_invocation( def with_client( self, *, - image: ImageRef | Mapping[str, object], cpus: int = 32, dependencies: ClientDependencies | Mapping[str, object] | None = None, ) -> DataDesignerSlurmConfigBuilder: @@ -107,7 +105,6 @@ def with_client( ClientConfig, { "cpus": cpus, - "image": image, "dependencies": {} if dependencies is None else dependencies, }, ) @@ -149,14 +146,11 @@ def build(self) -> DataDesignerSlurmConfig: missing = [] if self._invocation is None: missing.append("invocation") - if self._client is None: - missing.append("client") if not self._deployments: missing.append("deployment") if missing: raise SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") assert self._invocation is not None - assert self._client is not None return _validate_model( DataDesignerSlurmConfig, { @@ -164,7 +158,7 @@ def build(self) -> DataDesignerSlurmConfig: "name": self._name, "builder": self._builder, "invocation": self._invocation, - "client": self._client, + "client": self._client or ClientConfig(), "deployments": self._deployments, "array_tasks": self._array_tasks, "submission": self._submission, 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 index f561a397d..9fb2cc4f8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -30,7 +30,6 @@ validate_environment_bindings, validate_no_plaintext_secrets, ) -from data_designer.slurm.config.images import ImageRef from data_designer.slurm.config.vllm import QueueBackpressureConfig, VllmServerConfig from data_designer.slurm.contracts import ( AuthoredConfig, @@ -259,7 +258,6 @@ def validate_source(self) -> ClientDependencies: class ClientConfig(AuthoredConfig): cpus: PositiveInt = 32 - image: ImageRef dependencies: ClientDependencies = Field(default_factory=ClientDependencies) @@ -356,7 +354,7 @@ class DataDesignerSlurmConfig(AuthoredConfig): name: Identifier builder: BuilderInput invocation: InvocationConfig - client: ClientConfig + 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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index 42a38ecc0..e613c067d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -15,6 +15,7 @@ from pydantic import ( BaseModel, ConfigDict, + Field, NonNegativeInt, PositiveInt, StringConstraints, @@ -121,6 +122,7 @@ def freeze_collections(cls, value: object) -> object: 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): 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 index 64acaa924..65d091b3d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -14,7 +14,7 @@ from pathlib import Path from packaging.specifiers import SpecifierSet -from packaging.version import InvalidVersion +from packaging.version import InvalidVersion, Version from data_designer.slurm.config import ( ClientImageInspection, @@ -22,6 +22,7 @@ ImageInspectionRecord, ImageKind, ImageRef, + ServingImageInspection, ) from data_designer.slurm.contracts import Sha256Digest from data_designer.slurm.images.errors import ImageConflictError, ImageVerificationError @@ -29,6 +30,7 @@ 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 @@ -123,6 +125,17 @@ def publish_imported( "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( 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/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index b1a7543a0..17f5128ad 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -37,6 +37,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi {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)} @@ -113,6 +114,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) {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)} @@ -121,8 +123,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) 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)} -readonly DD_VERIFIED_CLIENT_IMAGE={quote_shell_value(plan.client.image.path)} -readonly DD_VERIFIED_CLIENT_IMAGE_SHA256={quote_shell_value(plan.client.image.sha256)} +readonly 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)} @@ -135,7 +136,6 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) verify_sha256 "${{DD_RUNTIME_SHA256}}" "${{DD_RUNTIME_ARCHIVE}}" verify_sha256 "${{DD_PLAN_SHA256}}" "${{DD_PLAN}}" -verify_sha256 "${{DD_VERIFIED_CLIENT_IMAGE_SHA256}}" "${{DD_VERIFIED_CLIENT_IMAGE}}" if [[ ! ${{SLURM_ARRAY_TASK_ID:-}} =~ ^[0-9]+$ ]]; then printf '%s\\n' 'SLURM_ARRAY_TASK_ID must be a non-negative integer' >&2 exit 64 @@ -171,7 +171,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) dd_initialize_local_scratch dd_stage_allocation_runtime readonly DD_RUNTIME_ROOT="${{DD_SCRATCH_ROOT}}/runtime" -export PYTHONPATH="${{DD_SCRATCH_CONTAINER_ROOT}}/runtime" +export PYTHONPATH="${{DD_RUNTIME_ROOT}}" srun \\ --nodes=1 \\ --ntasks=1 \\ @@ -181,11 +181,8 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) --unbuffered \\ --gres=none \\ --export=ALL \\ - --container-image="${{DD_VERIFIED_CLIENT_IMAGE}}" \\ - --container-mounts="${{DD_WORKSPACE_ROOT}}:${{DD_WORKSPACE_ROOT}},${{DD_SCRATCH_ROOT}}:${{DD_SCRATCH_CONTAINER_ROOT}}" \\ - --container-env=PYTHONPATH \\ -- \\ - python3 -m data_designer.slurm.state.attempt_identity \\ + "${{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}}" 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 index 46dfcbcea..c4f3d0067 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -13,6 +13,7 @@ RecordRange, ResolvedBuilderInput, ResolvedClient, + ResolvedClientRuntime, ResolvedDependencyLock, ResolvedDeployment, ResolvedImage, @@ -32,6 +33,7 @@ "RecordRange", "ResolvedBuilderInput", "ResolvedClient", + "ResolvedClientRuntime", "ResolvedDependencyLock", "ResolvedDeployment", "ResolvedImage", 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 index 964a282c3..c3fdb03a0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -187,7 +187,7 @@ def _compile_client( ) return ResolvedClient( authored=effective.authored.client, - image=effective.client_image, + runtime=effective.client_runtime, dependency_lock=ArtifactReference( path=posixpath.join(_run_root(effective), "dependency-lock.json"), sha256=effective.dependency_lock.compute_sha256(), 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 index d8380cc53..9c602b3fb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -14,14 +14,7 @@ from data_designer.config import RunConfig from data_designer.slurm.config.environment import validate_no_plaintext_secrets -from data_designer.slurm.config.images import ( - ClientImageInspection, - ImageInspectionRecord, - ImageKind, - ImageRef, - InstalledDistribution, - ServingImageInspection, -) +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, @@ -38,11 +31,13 @@ ContractValue, DistributionName, Identifier, + InstalledDistribution, ModelAlias, RecordRange, ResumeWorkspace, Sha256Digest, ShardId, + compute_canonical_json_sha256, compute_serialized_json_sha256, derive_managed_assets_path, is_path_below, @@ -84,11 +79,34 @@ def kind(self) -> ImageKind: return self.inspection.inspection.kind @property - def inspection_facts(self) -> ClientImageInspection | ServingImageInspection: + 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)] @@ -102,15 +120,15 @@ def validate_artifact(self) -> LockedPackage: class ResolvedDependencyLock(ContractRecord): - """Immutable client dependency resolution against one fixed image inventory.""" + """Immutable client dependency resolution against one fixed Python runtime.""" resolver_version: Identifier python_abi: Identifier - client_image_sha256: Sha256Digest + client_runtime_sha256: Sha256Digest authored_requirements: tuple[str, ...] authored_source: str | None = None source: ArtifactReference | None = None - image_distributions: tuple[InstalledDistribution, ...] + base_distributions: tuple[InstalledDistribution, ...] overlay_packages: tuple[LockedPackage, ...] @field_validator("authored_source") @@ -128,16 +146,16 @@ 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)) - image_names = tuple(distribution.name for distribution in self.image_distributions) + base_names = tuple(distribution.name for distribution in self.base_distributions) overlay_names = tuple(package.name for package in self.overlay_packages) - if image_names != tuple(sorted(image_names)) or overlay_names != tuple(sorted(overlay_names)): + 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(image_names) != len(set(image_names)) or len(overlay_names) != len(set(overlay_names)): + 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(image_names).intersection(overlay_names) + overlap = set(base_names).intersection(overlay_names) if overlap: - raise ValueError(f"overlay packages overlap image-owned distributions: {', '.join(sorted(overlap))}") - image_packages = {distribution.name: distribution for distribution in self.image_distributions} + 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) @@ -148,7 +166,7 @@ def validate_packages(self) -> ResolvedDependencyLock: 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 image_packages.get(name) + 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: @@ -322,7 +340,7 @@ def validate_deployment(self) -> ResolvedDeployment: class ResolvedClient(ContractValue): authored: ClientConfig - image: ResolvedImage + runtime: ResolvedClientRuntime dependency_lock: ArtifactReference host_node_index: NonNegativeInt gpu_count: Literal[0] @@ -330,10 +348,6 @@ class ResolvedClient(ContractValue): @model_validator(mode="after") def validate_client(self) -> ResolvedClient: - if self.image.kind is not ImageKind.CLIENT: - raise ValueError("Data Designer client requires a client image") - if self.image.authored_ref != self.authored.image: - raise ValueError("resolved client image does not match the authored image reference") 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): 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 index 3fc47b05f..e3a7ca5cb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -12,7 +12,7 @@ from data_designer.config import RunConfig from data_designer.slurm._errors import format_validation_error -from data_designer.slurm.config.images import ClientImageInspection, ImageKind +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 ( @@ -31,6 +31,7 @@ from data_designer.slurm.planning.errors import SlurmConfigResolutionError from data_designer.slurm.planning.models import ( ResolvedBuilderInput, + ResolvedClientRuntime, ResolvedDependencyLock, ResolvedImage, ResolvedInvocation, @@ -88,7 +89,7 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): builder: ResolvedBuilderInput builder_payload: dict[str, JsonValue] | None = None invocation: ResolvedInvocation - client_image: ResolvedImage + client_runtime: ResolvedClientRuntime deployment_images: tuple[ResolvedImage, ...] dependency_lock: ResolvedDependencyLock submission: ResolvedSubmission @@ -100,7 +101,7 @@ def resolve_slurm_config( authored: DataDesignerSlurmConfig, *, selected_profile: SelectedSlurmProfile, - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, deployment_images: tuple[ResolvedImage, ...], dependency_lock: ResolvedDependencyLock, runtime_bundle: ArtifactReference, @@ -130,7 +131,7 @@ def resolve_slurm_config( builder=builder, builder_payload=resolved_builder_payload, invocation=_materialize_invocation(authored, workspace_root), - client_image=client_image, + client_runtime=client_runtime, deployment_images=deployment_images, dependency_lock=dependency_lock, submission=_materialize_submission(authored, selected_profile), @@ -181,8 +182,8 @@ def validate_effective_slurm_config( if effective.output != expected_output: raise SlurmConfigResolutionError("resolved output does not match the authored output") - _validate_dependency_resolution(authored, effective.client_image, effective.dependency_lock) - _validate_resolved_images(authored, effective.client_image, effective.deployment_images) + _validate_dependency_resolution(authored, effective.client_runtime, effective.dependency_lock) + _validate_resolved_images(authored, effective.deployment_images) _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 @@ -349,18 +350,15 @@ def _validate_sharding_constraints( def _validate_dependency_resolution( authored: DataDesignerSlurmConfig, - client_image: ResolvedImage, + client_runtime: ResolvedClientRuntime, dependency_lock: ResolvedDependencyLock, ) -> None: - inspection = client_image.inspection.inspection - if not isinstance(inspection, ClientImageInspection): - raise SlurmConfigResolutionError("resolved client image lacks dependency inspection facts") - if dependency_lock.client_image_sha256 != client_image.sha256: - raise SlurmConfigResolutionError("dependency lock does not match the resolved client image") - if dependency_lock.python_abi != inspection.python_abi: - raise SlurmConfigResolutionError("dependency lock Python ABI does not match the client image") - if dependency_lock.image_distributions != inspection.distributions: - raise SlurmConfigResolutionError("dependency lock inventory does not match the client image") + 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: @@ -373,14 +371,8 @@ def _validate_dependency_resolution( def _validate_resolved_images( authored: DataDesignerSlurmConfig, - client_image: ResolvedImage, deployment_images: tuple[ResolvedImage, ...], ) -> None: - _validate_resolved_image_identity(client_image) - if client_image.kind is not ImageKind.CLIENT: - raise SlurmConfigResolutionError("resolved client image must contain client inspection facts") - if client_image.authored_ref != authored.client.image: - raise SlurmConfigResolutionError("resolved client image does not match the authored reference") 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): 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 index 5cf7ee9ba..dc0fa3245 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -8,7 +8,6 @@ from pydantic import JsonValue from data_designer.config import RunConfig -from data_designer.slurm.config.images import ClientImageInspection 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 @@ -107,17 +106,16 @@ def validate_resolved_plan( "dependency lock digest does not match the resolved plan", ) _require( - dependency_lock.client_image_sha256 == plan.client.image.sha256, - "dependency lock client image digest does not match the resolved client image", + dependency_lock.client_runtime_sha256 == plan.client.runtime.runtime_sha256, + "dependency lock runtime fingerprint does not match the resolved client runtime", ) - inspection = plan.client.image.inspection.inspection - _require(isinstance(inspection, ClientImageInspection), "resolved client image lacks client inspection facts") _require( - dependency_lock.python_abi == inspection.python_abi, "dependency lock Python ABI does not match client image" + dependency_lock.python_abi == plan.client.runtime.python_abi, + "dependency lock Python ABI does not match client runtime", ) _require( - dependency_lock.image_distributions == inspection.distributions, - "dependency lock image inventory does not match client image inspection", + 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: 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 index 10449cba6..214711b7f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -1,12 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed allocation step manifest produced inside the sealed client image.""" +"""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, @@ -34,12 +36,16 @@ def build_runtime_manifest( """Build the secret-free command handoff for the Bash controller.""" plan = context.plan validate_allocation_layout(plan, layout) - runtime_container_root = Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime" - if runtime_root != runtime_container_root: + 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( @@ -58,7 +64,7 @@ def build_runtime_manifest( context, environment, endpoints, - runtime_container_root.as_posix(), + runtime_root, log_directory, layout, ) @@ -75,7 +81,7 @@ def build_runtime_manifest( context, environment, endpoints, - runtime_container_root.as_posix(), + runtime_root, log_directory, layout, ) @@ -105,7 +111,7 @@ def _build_client_step( context: AllocationContext, environment: Mapping[str, str], endpoints: tuple[RuntimeEndpoint, ...], - runtime_container_root: str, + runtime_host_root: Path, log_directory: Path, layout: AllocationLayout, ) -> RuntimeStepSpec: @@ -139,7 +145,7 @@ def _build_client_step( ) ) command = ( - "python3", + plan.client.runtime.python_executable, "-m", "data_designer.slurm.runtime.entrypoint", "client", @@ -159,21 +165,20 @@ def _build_client_step( return _step( step_id=step_id, role=role, - image_path=plan.client.image.path, + execution="native", + image_path=None, command=command, cpus=plan.client.authored.cpus, gpu_indices=(), literal_environment={ - "DATA_DESIGNER_SLURM_SCRATCH_ROOT": ALLOCATION_SCRATCH_CONTAINER_ROOT, + "DATA_DESIGNER_SLURM_SCRATCH_ROOT": runtime_host_root.parent.as_posix(), "LC_ALL": "C", - "PYTHONPATH": runtime_container_root, + "PYTHONPATH": runtime_host_root.as_posix(), **allocation_environment, }, secret_environment={name: name for name in secret_names}, environment_prefixes={}, - container_environment=tuple( - sorted((*secret_names, *allocation_environment, "DATA_DESIGNER_SLURM_SCRATCH_ROOT", "PYTHONPATH")) - ), + container_environment=(), log_directory=log_directory, node_hosts=(layout.get_host(plan.client.host_node_index),), ) @@ -198,19 +203,19 @@ def _build_endpoint_step( proxy_path, deployment.logical_endpoint.port, backend_hosts=backend_hosts, - runtime_proxy_is_container_path=True, ) return _step( step_id=f"{deployment.deployment_id}-endpoint", role=RuntimeStepRole.ENDPOINT, - image_path=context.plan.client.image.path, + 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=("PYTHONPATH",), + container_environment=(), log_directory=log_directory, node_hosts=(layout.get_host(context.plan.client.host_node_index),), readiness=( @@ -228,7 +233,8 @@ def _step( *, step_id: str, role: RuntimeStepRole, - image_path: str, + execution: Literal["native", "container"], + image_path: str | None, command: tuple[str, ...], cpus: int, gpu_indices: tuple[int, ...], @@ -245,6 +251,7 @@ def _step( return RuntimeStepSpec( step_id=step_id, role=role, + execution=execution, image_path=image_path, command=command, cpus=cpus, 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 index cbb32b181..0d54156fa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/context.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Load allocation identity from container-visible persisted state.""" +"""Load allocation identity from native compute-visible persisted state.""" from __future__ import annotations @@ -13,7 +13,6 @@ 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.runtime.paths import get_container_path from data_designer.slurm.state import RetryPlan, SlurmStateError, SlurmStateWriter from data_designer.slurm.state.filesystem import open_verified_directory, read_regular_text @@ -29,20 +28,18 @@ def load_allocation_context( retry_plan_sha256: str | None = None, effective_resume_mode: str | None = None, ) -> tuple[AllocationContext, SlurmStateWriter]: - """Load one scheduler-selected shard attempt through its container paths.""" + """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.as_posix() != get_container_path(plan, expected_plan_path.as_posix()): + 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.as_posix() != get_container_path( - plan, host_attempt_directory.as_posix(), require_writable=True - ): + if attempt_directory != host_attempt_directory: raise SlurmRuntimeError( SlurmRuntimeErrorCode.INVALID_CONTEXT, "attempt directory does not match the scheduler array task", @@ -111,15 +108,9 @@ def _load_state_writer(plan_path: Path, attempt_directory: Path) -> SlurmStateWr 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 get_container_path(plan, logical_plan_path.as_posix()) != plan_path.as_posix(): + if logical_plan_path != plan_path: raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "resolved plan path is invalid") - workspace_root = Path(get_container_path(plan, logical_workspace_root, require_writable=True)) - return SlurmStateWriter( - workspace_root, - run_id, - logical_workspace_root=logical_workspace_root, - local_path_resolver=lambda path: get_container_path(plan, path, require_writable=True), - ) + return SlurmStateWriter(logical_workspace_root, run_id) def _select_shard(shards: tuple[PlannedShard, ...], task_id: int) -> PlannedShard: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index d7f125a36..67d74fd15 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed control phases executed only inside the sealed client image.""" +"""Typed control phases executed with the resolved native client runtime.""" from __future__ import annotations @@ -19,7 +19,6 @@ 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.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_deployments from data_designer.slurm.runtime.preflight import ( AllocationLayout, @@ -44,7 +43,7 @@ def main(arguments: Sequence[str] | None = None) -> int: - """Execute one container-only allocation control phase.""" + """Execute one native allocation control phase.""" parsed = _parse_arguments(arguments) try: if parsed.operation == "prepare": @@ -101,11 +100,11 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N 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_ports(context, environment) readiness = _begin_attempt(context, writer, environment) log_directory = execution_log_directory(context.attempt_directory, readiness.revision) - container_log_directory = Path(get_container_path(context.plan, log_directory.as_posix(), require_writable=True)) - ensure_private_directory(container_log_directory) + ensure_private_directory(log_directory) manifest = build_runtime_manifest( context, environment, @@ -114,11 +113,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N layout=layout, ) expected_manifest = context.attempt_directory / "runtime-manifest.json" - if arguments.manifest.as_posix() != get_container_path( - context.plan, - expected_manifest.as_posix(), - require_writable=True, - ): + if arguments.manifest != expected_manifest: raise SlurmRuntimeError(SlurmRuntimeErrorCode.INVALID_CONTEXT, "runtime manifest path is invalid") replace_private_text(arguments.manifest, manifest.serialize_json()) 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 index 3a88d6515..8974914aa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -36,18 +36,9 @@ dd_slurm_run_allocation() { dd_read_control_plan "${DD_PLAN_PATH}" DD_RUNTIME_PHASE=read_plan_secrets dd_read_plan_secret_names "${DD_PLAN_PATH}" - DD_RUNTIME_PHASE=translate_plan_path - dd_read_container_path "${DD_PLAN_PATH}" "${DD_PLAN_PATH}" false - DD_PLAN_CONTAINER_PATH=${DD_CONTAINER_PATH} - DD_RUNTIME_PHASE=translate_attempt_path - dd_read_container_path "${DD_PLAN_PATH}" "${DD_ATTEMPT_PATH}" true - DD_ATTEMPT_CONTAINER_DIR=${DD_CONTAINER_PATH} DD_RUNTIME_CONTAINER_ROOT=${DD_SCRATCH_CONTAINER_ROOT}/runtime - DD_RUNTIME_PHASE=translate_manifest_path - dd_read_container_path "${DD_PLAN_PATH}" "${DD_RUNTIME_MANIFEST}" true - DD_RUNTIME_MANIFEST_CONTAINER_PATH=${DD_CONTAINER_PATH} export DD_PLAN_PATH DD_ATTEMPT_PATH DD_RUNTIME_MANIFEST - export DD_PLAN_CONTAINER_PATH DD_ATTEMPT_CONTAINER_DIR DD_RUNTIME_CONTAINER_ROOT + export DD_RUNTIME_ROOT DD_RUNTIME_CONTAINER_ROOT DD_RUNTIME_PHASE=verify_host_context dd_verify_host_context @@ -62,8 +53,8 @@ dd_slurm_run_allocation() { done DD_RUNTIME_PHASE=prepare dd_run_bound_control_phase prepare \ - --runtime-root "${DD_RUNTIME_CONTAINER_ROOT}" \ - --manifest "${DD_RUNTIME_MANIFEST_CONTAINER_PATH}" \ + --runtime-root "${DD_RUNTIME_ROOT}" \ + --manifest "${DD_RUNTIME_MANIFEST}" \ "${host_arguments[@]}" DD_RUNTIME_PHASE=verify_runtime_manifest dd_verify_runtime_manifest \ @@ -131,6 +122,7 @@ dd_verify_host_context() { [[ -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} ]] 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 index 728c7d8b9..5569f38c2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -6,6 +6,7 @@ from __future__ import annotations from pathlib import Path +from typing import Literal from pydantic import Field, NonNegativeInt, PositiveInt, field_validator, model_validator @@ -37,11 +38,12 @@ def validate_host(cls, value: str) -> str: class RuntimeStepSpec(ContractValue): - """Container command and placement consumed by the Bash step runner.""" + """Native or container command and placement consumed by the Bash step runner.""" step_id: Identifier role: RuntimeStepRole - image_path: str + execution: Literal["native", "container"] + image_path: str | None = None command: tuple[str, ...] = Field(min_length=1) cpus: PositiveInt gpu_indices: tuple[NonNegativeInt, ...] = () @@ -56,7 +58,11 @@ class RuntimeStepSpec(ContractValue): launch_delay_seconds: NonNegativeInt = 0 readiness: tuple[RuntimeProbeSpec, ...] = () - _image_path_is_absolute = field_validator("image_path")(validate_absolute_path) + @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) @@ -77,12 +83,19 @@ def validate_step(self) -> RuntimeStepSpec: 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: 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 index 0564e6a00..bf39dcbf6 100644 --- 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 @@ -27,7 +27,7 @@ dd_read_control_plan() { and (. as $target | ("/run/data-designer-slurm" | startswith($target + "/") | not)) )) ' "${plan}" >/dev/null || return 65 - DD_CLIENT_IMAGE=$(jq -er '.client.image.path' "${plan}") + 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}") @@ -72,7 +72,6 @@ dd_read_artifacts() { [ .runtime_bundle, .client.dependency_lock, - {path: .client.image.path, sha256: .client.image.sha256}, (.deployments[] | {path: .image.path, sha256: .image.sha256}), .builder.source, (.shards[] | select(.array_task_index == $task_id) | .input_partition) @@ -142,7 +141,8 @@ dd_read_step() { jq -j --arg step_id "${step_id}" ' .steps[] | select(.step_id == $step_id) - | .image_path, "\u0000", + | .execution, "\u0000", + (.image_path // ""), "\u0000", (.cpus | tostring), "\u0000", .stdout_path, "\u0000", .stderr_path, "\u0000", @@ -150,12 +150,13 @@ dd_read_step() { (.kill_on_bad_exit | tostring), "\u0000" ' "${manifest}" ) - DD_STEP_IMAGE=${DD_STEP_FIELDS[0]} - DD_STEP_CPUS=${DD_STEP_FIELDS[1]} - DD_STEP_STDOUT=${DD_STEP_FIELDS[2]} - DD_STEP_STDERR=${DD_STEP_FIELDS[3]} - DD_STEP_DELAY=${DD_STEP_FIELDS[4]} - DD_STEP_KILL_ON_BAD_EXIT=${DD_STEP_FIELDS[5]} + 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_read_null_values DD_STEP_COMMAND < <( jq -j --arg step_id "${step_id}" '.steps[] | select(.step_id == $step_id) | .command[] | ., "\u0000"' \ "${manifest}" 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 index 1b3204d41..480759039 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -15,12 +15,12 @@ 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.paths import get_container_path from data_designer.slurm.runtime.ports import resolve_allocation_plan _DIGEST_CHUNK_SIZE = 1024 * 1024 @@ -90,10 +90,9 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> """Verify every launch-critical fact before model services start.""" try: self._verify_scheduler(context, environment) - attempt_directory = Path( - get_container_path(context.plan, context.attempt_directory.as_posix(), require_writable=True) - ) + 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: @@ -135,6 +134,22 @@ def _verify_scheduler(context: AllocationContext, environment: Mapping[str, str] "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.""" @@ -155,7 +170,6 @@ def _verify_artifacts(context: AllocationContext) -> None: ), plan.runtime_bundle, plan.client.dependency_lock, - ArtifactReference(path=plan.client.image.path, sha256=plan.client.image.sha256), *( ArtifactReference(path=deployment.image.path, sha256=deployment.image.sha256) for deployment in plan.deployments @@ -168,11 +182,6 @@ def _verify_artifacts(context: AllocationContext) -> None: references.extend(reference for reference in optional_references if reference is not None) unique_references = {(reference.path, reference.sha256): reference for reference in references} for reference in unique_references.values(): - if any( - reference.path == mount.source or reference.path.startswith(f"{mount.source}/") - for mount in plan.container_mounts - ): - reference = reference.model_copy(update={"path": get_container_path(plan, reference.path)}) _verify_artifact(reference) @staticmethod 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 index 4cd6b4401..3968a5dbb 100644 --- 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 @@ -91,6 +91,7 @@ def build_distributed_server_step( 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", @@ -171,6 +172,7 @@ def _build_local_server_step( 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, 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 index f1c0fb2b5..e4de7b973 100644 --- 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 @@ -23,7 +23,7 @@ dd_scope_control_environment() { done unset CUDA_VISIBLE_DEVICES export LC_ALL=C - export PYTHONPATH=${DD_RUNTIME_CONTAINER_ROOT} + export PYTHONPATH=${DD_RUNTIME_ROOT} } dd_materialize_step_environment() { @@ -83,7 +83,6 @@ dd_build_srun_command() { --unbuffered --export=ALL "--cpus-per-task=${DD_STEP_CPUS}" - "--container-image=${DD_STEP_IMAGE}" ) if [[ ${DD_STEP_KILL_ON_BAD_EXIT} == true ]]; then DD_SRUN_COMMAND+=(--kill-on-bad-exit=1) @@ -107,11 +106,17 @@ dd_build_srun_command() { else DD_SRUN_COMMAND+=(--gres=none) fi - [[ -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%,}") + 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 } @@ -148,7 +153,6 @@ dd_run_step() { dd_run_control_phase() { local operation=$1 shift - local container_mounts="${DD_SCRATCH_ROOT}:${DD_SCRATCH_CONTAINER_ROOT}" local -a command=( srun --nodes=1 @@ -159,22 +163,16 @@ dd_run_control_phase() { --unbuffered --export=ALL "--cpus-per-task=${DD_CLIENT_CPUS}" - "--container-image=${DD_CLIENT_IMAGE}" --gres=none - ) - [[ -z ${DD_CONTAINER_MOUNTS} ]] || container_mounts="${DD_CONTAINER_MOUNTS},${container_mounts}" - command+=("--container-mounts=${container_mounts}") - command+=( - --container-env=PYTHONPATH,SLURM_JOB_GPUS -- - python3 + "${DD_CLIENT_PYTHON}" -m data_designer.slurm.runtime.entrypoint "${operation}" --plan - "${DD_PLAN_CONTAINER_PATH}" + "${DD_PLAN_PATH}" --attempt-dir - "${DD_ATTEMPT_CONTAINER_DIR}" + "${DD_ATTEMPT_PATH}" "$@" ) ( 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 index 1e706828c..c6f2be674 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -153,19 +153,14 @@ def _build_step( endpoints, retry_resume_mode, ) - secret_names, environment = _build_client_environment(plan, source_environment) - allocation_environment = ("SLURM_JOB_GPUS",) if plan.selected_profile.profile.gpu_request_mode == "gres" else () + _, environment = _build_client_environment(plan, source_environment) return _build_srun_step( step_id=step_id, role=role, - image_path=plan.client.image.path, + image_path=None, command=command, environment=environment, - container_environment=( - *secret_names, - *allocation_environment, - "DATA_DESIGNER_SLURM_SCRATCH_ROOT", - ), + container_environment=(), plan=plan, attempt_directory=attempt_directory, ) @@ -186,18 +181,18 @@ def build_client_command( for argument in ("--endpoint", f"{endpoint.model_alias}=http://{endpoint.host}:{endpoint.port}/v1") ) return ( - "python3", + plan.client.runtime.python_executable, "-m", "data_designer.slurm.client.worker", operation, "--plan", - get_container_path(plan, plan_path(plan)), + plan_path(plan), "--shard-id", shard.shard_id, "--attempt-id", attempt.attempt_id, "--attempt-dir", - get_container_path(plan, attempt_directory.as_posix(), require_writable=True), + attempt_directory.as_posix(), *(() if retry_resume_mode is None else ("--resume-mode", retry_resume_mode)), *endpoint_arguments, ) @@ -225,7 +220,7 @@ def _build_client_environment( "allocation scratch is unavailable", ) from None environment["DD_SCRATCH_ROOT"] = scratch_root - environment["DATA_DESIGNER_SLURM_SCRATCH_ROOT"] = ALLOCATION_SCRATCH_CONTAINER_ROOT + environment["DATA_DESIGNER_SLURM_SCRATCH_ROOT"] = scratch_root return secret_names, environment @@ -284,14 +279,14 @@ def _build_endpoint_step( ) command = build_endpoint_command(deployment, plan, runtime_proxy_path, endpoint.port) environment = _base_environment(source_environment) - environment["PYTHONPATH"] = get_container_path(plan, runtime_root.as_posix()) + environment["PYTHONPATH"] = runtime_root.as_posix() step = _build_srun_step( step_id=f"{deployment.deployment_id}-endpoint", role=RuntimeStepRole.ENDPOINT, - image_path=plan.client.image.path, + image_path=None, command=command, environment=environment, - container_environment=("PYTHONPATH",), + container_environment=(), plan=plan, attempt_directory=attempt_directory, ) @@ -305,7 +300,6 @@ def build_endpoint_command( port: int, *, backend_hosts: tuple[str, ...] | None = None, - runtime_proxy_is_container_path: bool = False, ) -> tuple[str, ...]: selected_hosts = backend_hosts or ("127.0.0.1",) * len(deployment.backend_endpoints) if len(selected_hosts) != len(deployment.backend_endpoints): @@ -325,16 +319,12 @@ def build_endpoint_command( 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: - container_proxy_path = ( - validate_absolute_path(runtime_proxy_path.as_posix()) - if runtime_proxy_is_container_path - else get_container_path(plan, runtime_proxy_path.as_posix()) - ) + 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 ( - "python3", - container_proxy_path, + plan.client.runtime.python_executable, + host_proxy_path, "--listen-port", str(port), "--health-path", @@ -455,7 +445,7 @@ def _build_srun_step( *, step_id: str, role: RuntimeStepRole, - image_path: str, + image_path: str | None, command: tuple[str, ...], environment: Mapping[str, str], container_environment: tuple[str, ...], @@ -465,7 +455,8 @@ def _build_srun_step( ) -> RuntimeStep: srun_command = _build_srun_prefix(plan, image_path) _add_srun_resources(srun_command, gpu_indices) - _add_srun_container_options(srun_command, plan, container_environment, environment) + 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( @@ -478,8 +469,8 @@ def _build_srun_step( ) -def _build_srun_prefix(plan: ResolvedSlurmRunPlan, image_path: str) -> list[str]: - return [ +def _build_srun_prefix(plan: ResolvedSlurmRunPlan, image_path: str | None) -> list[str]: + command = [ "srun", "--nodes=1", "--ntasks=1", @@ -488,8 +479,10 @@ def _build_srun_prefix(plan: ResolvedSlurmRunPlan, image_path: str) -> list[str] "--unbuffered", "--export=ALL", f"--cpus-per-task={plan.client.authored.cpus}", - f"--container-image={image_path}", ] + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index b4bfced83..5a5abb6f8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -26,6 +26,7 @@ ) 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, @@ -145,6 +146,7 @@ def __init__( launcher: SlurmCommandClient, run_id_factory: RunIdFactory, package_version: str, + runtime_inspector: ClientRuntimeInspector | None = None, ) -> None: self._profile = selected_profile self._images = image_registry @@ -152,6 +154,7 @@ def __init__( 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( @@ -184,7 +187,7 @@ def _prepare( run_id = self._run_id_factory() workspace_root = self._profile.profile.workspace_root run_root = Path(workspace_root) / "runs" / run_id - client_image = self._images.resolve_for_planning(authored.client.image, expected_kind=ImageKind.CLIENT) + client_runtime = self._runtime_inspector.inspect() deployment_images = tuple( self._images.resolve_for_planning(deployment.server.image, expected_kind=ImageKind.SERVING) for deployment in authored.deployments @@ -193,7 +196,7 @@ def _prepare( dependencies = stack.enter_context( self._dependencies.resolve( authored.client.dependencies, - client_image, + client_runtime, run_root=run_root, source_root=source_root, ) @@ -202,7 +205,7 @@ def _prepare( effective = resolve_slurm_config( authored, selected_profile=self._profile, - client_image=client_image, + client_runtime=client_runtime, deployment_images=deployment_images, dependency_lock=dependencies.lock, runtime_bundle=runtime_bundle, @@ -230,6 +233,8 @@ def _prepare( 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" diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index 67bd3c517..4f49bb6e8 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -5,6 +5,7 @@ import json import platform +import sys from dataclasses import dataclass from pathlib import Path @@ -118,11 +119,21 @@ def client_worker_case(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Clien (GOLDEN_DIRECTORY / "single_node_plan.json").read_text().replace("/workspace/primary", workspace.as_posix()) ) python_abi = f"{interpreter_name()}{interpreter_version()}" - payload["client"]["image"]["inspection"]["inspection"].update( - {"python_abi": python_abi, "python_version": platform.python_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} @@ -153,12 +164,12 @@ def client_worker_case(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Clien scratch_root=scratch_root, overlay_path=overlay_path, dependency_lock=plan.client.dependency_lock, - client_image_sha256=plan.client.image.sha256, + 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.image_distributions + for distribution in lock.base_distributions ), ) allocation_plan = resolve_allocation_plan(plan, {"SLURM_JOB_GPUS": "0"}) diff --git a/packages/data-designer-slurm/tests/client/test_dependency_resolution.py b/packages/data-designer-slurm/tests/client/test_dependency_resolution.py index 91b2f4675..bc5dec95f 100644 --- a/packages/data-designer-slurm/tests/client/test_dependency_resolution.py +++ b/packages/data-designer-slurm/tests/client/test_dependency_resolution.py @@ -30,13 +30,13 @@ def unexpected_runner(*_args: object) -> None: with resolver.resolve( dependencies, - single_node_plan.client.image, + 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_image_sha256 == single_node_plan.client.image.sha256 + assert resolved.lock.client_runtime_sha256 == single_node_plan.client.runtime.runtime_sha256 assert resolved.wheel_sources == () assert resolved.lock_source is None @@ -75,7 +75,7 @@ def download(command: tuple[str, ...], _cwd: Path, environment: dict[str, str]) with resolver.resolve( dependencies, - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=run_root, source_root=tmp_path, ) as resolved: @@ -105,7 +105,7 @@ def download(command: tuple[str, ...], _cwd: Path, _environment: dict[str, str]) pytest.raises(ClientDependencyResolutionError, match="platform-independent"), resolver.resolve( ClientDependencies(requirements=["example-plugin==1.2.0"]), - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=tmp_path / "workspace/runs/run-001", source_root=tmp_path, ), @@ -119,14 +119,14 @@ def test_supplied_lock_is_verified_and_rebound_to_run_inputs( ) -> None: wheel = tmp_path / "example_plugin-1.2.0-py3-none-any.whl" wheel.write_bytes(b"plugin") - image = single_node_plan.client.image.inspection_facts + runtime = single_node_plan.client.runtime supplied = ResolvedDependencyLock( schema_version=1, resolver_version="external-1", - python_abi=image.python_abi, - client_image_sha256=single_node_plan.client.image.sha256, + python_abi=runtime.python_abi, + client_runtime_sha256=runtime.runtime_sha256, authored_requirements=("example-plugin==1.2.0",), - image_distributions=tuple(sorted(image.distributions, key=lambda item: item.name)), + base_distributions=tuple(sorted(runtime.distributions, key=lambda item: item.name)), overlay_packages=( LockedPackage( name="example-plugin", @@ -144,7 +144,7 @@ def test_supplied_lock_is_verified_and_rebound_to_run_inputs( with ClientDependencyResolver().resolve( ClientDependencies(requirements=None, lock_file="lock.json"), - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=run_root, source_root=tmp_path, ) as resolved: @@ -167,14 +167,14 @@ def test_supplied_lock_rejects_platform_specific_wheels( ) -> None: wheel = tmp_path / "example_plugin-1.2.0-cp312-cp312-manylinux_2_28_x86_64.whl" wheel.write_bytes(b"plugin") - image = single_node_plan.client.image.inspection_facts + runtime = single_node_plan.client.runtime supplied = ResolvedDependencyLock( schema_version=1, resolver_version="external-1", - python_abi=image.python_abi, - client_image_sha256=single_node_plan.client.image.sha256, + python_abi=runtime.python_abi, + client_runtime_sha256=runtime.runtime_sha256, authored_requirements=("example-plugin==1.2.0",), - image_distributions=tuple(sorted(image.distributions, key=lambda item: item.name)), + base_distributions=tuple(sorted(runtime.distributions, key=lambda item: item.name)), overlay_packages=( LockedPackage( name="example-plugin", @@ -193,7 +193,7 @@ def test_supplied_lock_rejects_platform_specific_wheels( pytest.raises(ClientDependencyResolutionError, match="artifact differs"), ClientDependencyResolver().resolve( ClientDependencies(requirements=None, lock_file="lock.json"), - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=tmp_path / "workspace/runs/run-001", source_root=tmp_path, ), @@ -213,7 +213,7 @@ def fail(*_args: object) -> None: with pytest.raises(ClientDependencyResolutionError, match="resolution failed") as captured: with resolver.resolve( dependencies, - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=tmp_path / "workspace/runs/run-001", source_root=tmp_path, ): @@ -227,7 +227,7 @@ def fail(*_args: object) -> None: with pytest.raises(ClientDependencyResolutionError, match="credential is unavailable"): with resolver.resolve( credentialed, - single_node_plan.client.image, + single_node_plan.client.runtime, run_root=tmp_path / "workspace/runs/run-002", source_root=tmp_path, ): diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 8013f5f35..9386229e0 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -16,19 +16,28 @@ 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, compute_canonical_json_sha256 +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.image_distributions if path is None else () + return client_worker_case.lock.base_distributions if path is None else () prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( client_worker_case.plan_path, @@ -39,7 +48,7 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: ) assert prepared.installer_outcome is ClientInstallerOutcome.NOT_REQUIRED - assert prepared.installed_distributions == client_worker_case.lock.image_distributions + 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() @@ -78,16 +87,12 @@ def test_environment_rejects_symlink_scratch_root( assert error.value.code is ClientErrorCode.INVALID_INPUT -def test_client_runs_through_non_identity_workspace_mount(client_worker_case: ClientWorkerCase) -> None: +def test_native_client_ignores_container_mount_mapping(client_worker_case: ClientWorkerCase) -> None: physical_workspace = client_worker_case.plan_path.parents[2] - logical_workspace = "/host/workspace" - payload = cast( - dict[str, object], - json.loads(client_worker_case.plan.serialize_json().replace(physical_workspace.as_posix(), logical_workspace)), - ) + 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": logical_workspace, "target": physical_workspace.as_posix(), "read_only": False} + 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) @@ -96,7 +101,7 @@ def test_client_runs_through_non_identity_workspace_mount(client_worker_case: Cl client_worker_case.plan_path.write_text(plan.serialize_json()) def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: - return client_worker_case.lock.image_distributions if path is None else () + return client_worker_case.lock.base_distributions if path is None else () prepared = ClientEnvironmentBuilder(inventory=inventory).prepare( client_worker_case.plan_path, @@ -120,12 +125,12 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: plugins=(), ) - assert result.dataset_path.startswith(logical_workspace) + assert result.dataset_path.startswith(physical_workspace.as_posix()) assert result.candidate_output_manifest is not None - assert result.candidate_output_manifest.path.startswith(logical_workspace) + assert result.candidate_output_manifest.path.startswith(physical_workspace.as_posix()) -def test_environment_maps_each_artifact_through_nested_mounts( +def test_native_environment_rejects_container_only_plan_path( client_worker_case: ClientWorkerCase, tmp_path: Path, ) -> None: @@ -158,30 +163,45 @@ def test_environment_maps_each_artifact_through_nested_mounts( 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.image_distributions if path is None else () + return client_worker_case.lock.base_distributions if path is None else () - prepared = 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, - ) + 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 prepared.attempt_dir == attempt_dir + assert error.value.code is ClientErrorCode.INVALID_INPUT -def test_inspect_distributions_omits_path_for_active_environment(monkeypatch: pytest.MonkeyPatch) -> None: +@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 = "/users/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.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) == () - assert calls == [{}] + expected_paths = [interpreter_path, *([user_site_path] if user_site_enabled else []), pth_path] + assert calls == [{"path": expected_paths}] def test_inspect_distributions_refreshes_overlay_with_unchanged_directory_mtime(tmp_path: Path) -> None: @@ -206,7 +226,61 @@ def test_inspect_distributions_rejects_unhashed_direct_url(monkeypatch: pytest.M distribution.read_text.return_value = '{"url":"https://example.test/example.whl"}' monkeypatch.setattr( "data_designer.slurm.client.environment.importlib.metadata.distributions", - lambda: (distribution,), + 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: @@ -215,7 +289,7 @@ def test_inspect_distributions_rejects_unhashed_direct_url(monkeypatch: pytest.M assert error.value.code is ClientErrorCode.DEPENDENCY_CONFLICT -def test_environment_rejects_client_image_inventory_conflict(client_worker_case: ClientWorkerCase) -> None: +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: @@ -277,7 +351,7 @@ def test_environment_rejects_missing_locked_wheel(client_worker_case: ClientWork client_worker_case.plan_path.write_text(plan.serialize_json()) def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: - return lock.image_distributions if path is None else () + return lock.base_distributions if path is None else () with pytest.raises(ClientWorkerError) as error: ClientEnvironmentBuilder(inventory=inventory).prepare( @@ -317,7 +391,7 @@ def test_environment_installs_verified_wheel_overlay(client_worker_case: ClientW def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: if path is None: - return lock.image_distributions + 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( @@ -366,10 +440,10 @@ def test_worker_bootstrap_loads_overlay_plugin_before_config_import( shard_id="shard-00000", attempt_id="attempt-0001", attempt_dir=Path(sys.argv[2]), - scratch_root=Path(sys.argv[1]).parent, + 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_image_sha256="b" * 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"),), @@ -388,8 +462,17 @@ def test_worker_bootstrap_loads_overlay_plugin_before_config_import( ) 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()], + [ + 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, 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..64827e4cc --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_runtime.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +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 + + +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 index a1544d121..dd25a3c92 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -545,7 +545,7 @@ def test_preflight_rejects_plugin_secondary_model_alias( installed = tuple( sorted( ( - *(InstalledDistribution(**item) for item in lock["image_distributions"]), + *(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, @@ -556,10 +556,10 @@ def test_preflight_rejects_plugin_secondary_model_alias( shard_id=plan["shards"][0]["shard_id"], attempt_id="attempt-0001", attempt_dir=attempt_dir, - scratch_root=Path(sys.argv[3]).parent, + scratch_root=Path(sys.argv[4]), overlay_path=Path(sys.argv[3]), dependency_lock=ArtifactReference(**plan["client"]["dependency_lock"]), - client_image_sha256=plan["client"]["image"]["sha256"], + client_runtime_sha256=plan["client"]["runtime"]["runtime_sha256"], python_abi=lock["python_abi"], installer_outcome=ClientInstallerOutcome.REUSED, installed_distributions=installed, @@ -585,6 +585,8 @@ def test_preflight_rejects_plugin_secondary_model_alias( 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, @@ -593,6 +595,7 @@ def test_preflight_rejects_plugin_secondary_model_alias( 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, diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index c2ff3087f..4238b402e 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -41,7 +41,7 @@ def _config_builder(*, prompt: str | None = None) -> DataDesignerSlurmConfigBuil dataset_name="generated", model_concurrency={"generator": 8}, ) - .with_client(image={"name": "dd-client"}) + .with_client() .with_deployment( { "model_alias": "generator", @@ -67,7 +67,7 @@ def test_builder_builds_without_file_discovery_or_serialization(tmp_path: Path) def test_builder_requires_complete_authored_intent() -> None: builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") - with pytest.raises(SlurmConfigBuilderError, match="invocation, client, deployment"): + with pytest.raises(SlurmConfigBuilderError, match="invocation, deployment"): builder.build() @@ -122,7 +122,7 @@ def test_builder_validation_errors_hide_secret_inputs() -> None: secret = "super-secret-token" with pytest.raises(SlurmConfigBuilderError) as error: - _config_builder().with_client(image={"name": "dd-client", "api_key": secret}) + _config_builder().with_client(cpus=secret) # type: ignore[arg-type] assert secret not in str(error.value) assert error.value.__cause__ is None @@ -167,7 +167,6 @@ def test_builder_custom_validation_errors_hide_secret_values() -> None: with pytest.raises(SlurmConfigBuilderError) as error: _config_builder().with_client( - image={"name": "dd-client"}, dependencies={"requirements": [f"pkg=={secret}"]}, ) diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index c1bb0d43f..a0e1a6e70 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -36,10 +36,6 @@ "requirements": [ "data-designer-speech==0.2.0" ] - }, - "image": { - "name": "dd-client-0.9", - "path": null } }, "deployments": [ 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 index 7a852dbcd..f1e81db8d 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -24,10 +24,6 @@ "index_credentials": {}, "lock_file": null, "requirements": [] - }, - "image": { - "name": "dd-client-0.9", - "path": null } }, "deployments": [ diff --git a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json index c8cb7d7d9..1f44f2042 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock.json @@ -3,8 +3,7 @@ "data-designer-speech==0.2.0" ], "authored_source": null, - "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", - "image_distributions": [ + "base_distributions": [ { "name": "aiohttp", "version": "3.14.3" @@ -30,6 +29,7 @@ "version": "26.1" } ], + "client_runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1", "overlay_packages": [ { "artifact": { 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 index 2219e03cf..b75c6b608 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/dependency_lock_single.json @@ -1,8 +1,7 @@ { "authored_requirements": [], "authored_source": null, - "client_image_sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc", - "image_distributions": [ + "base_distributions": [ { "name": "aiohttp", "version": "3.14.3" @@ -28,6 +27,7 @@ "version": "26.1" } ], + "client_runtime_sha256": "d18ec68ea21778533f769192b6bd0d4a735b3ad782dd6b53900a12ef7fd05be1", "overlay_packages": [], "python_abi": "cp312", "resolver_version": "resolver-1", 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 index f6cddb883..5595a228e 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-001/authored-config.json", - "sha256": "74f30c766e96181f72642dff1da72ae10786d5e404b4f65b927bdf27fa0c6a61" + "sha256": "fac5dec4c06558f13e761426bc52699a06017d5ba34514af422eda040fe43137" }, "builder": { "authored_source": null, @@ -47,65 +47,14 @@ "requirements": [ "data-designer-speech==0.2.0" ] - }, - "image": { - "name": "dd-client-0.9", - "path": null } }, "dependency_lock": { "path": "/workspace/primary/runs/run-001/dependency-lock.json", - "sha256": "f2eea2f07a52e23a60b7d5de317fbc6c6762456a8467401d620d430fe8fe6344" + "sha256": "d318c2a934694f9c5186236c98a4e66c58b5572bcb4d1aec1309344fea0766be" }, "gpu_count": 0, "host_node_index": 0, - "image": { - "authored_ref": { - "name": "dd-client-0.9", - "path": null - }, - "inspection": { - "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" - }, - "path": "/images/dd-client-0.9.sqsh", - "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" - }, "ports": [ { "name": "deployment-00000-logical-endpoint", @@ -119,7 +68,40 @@ "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": [ { 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 index ad8a6773c..03893958b 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" + "sha256": "7eb39d9f28b7c073bc520c6596bdf7f6e2afc5e7fa7881945c5bcb34b19f9920" }, "builder": { "authored_source": null, @@ -34,65 +34,14 @@ "index_credentials": {}, "lock_file": null, "requirements": [] - }, - "image": { - "name": "dd-client-0.9", - "path": null } }, "dependency_lock": { "path": "/workspace/primary/runs/run-single/dependency-lock.json", - "sha256": "abd08d979eaf61f088680847887c8c0def48f6210b532cb1c08b445120f57742" + "sha256": "df1cf487016d89f1ab4d0d2df5551d0d899fd695804cab0cd3d99a7a936e2f3f" }, "gpu_count": 0, "host_node_index": 0, - "image": { - "authored_ref": { - "name": "dd-client-0.9", - "path": null - }, - "inspection": { - "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" - }, - "path": "/images/dd-client-0.9.sqsh", - "sha256": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" - }, "ports": [ { "name": "deployment-00000-logical-endpoint", @@ -100,7 +49,40 @@ "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": [ { diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index fa45b8291..4d04b9488 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -350,7 +350,7 @@ def test_resolved_submission_preserves_authored_validation(update: dict[str, obj def test_resolved_image_rejects_digest_mismatch(multi_node_plan: ResolvedSlurmRunPlan) -> None: payload = multi_node_plan.model_dump(mode="json") - payload["client"]["image"]["inspection"]["sqsh_sha256"] = "a" * 64 + payload["deployments"][0]["image"]["inspection"]["sqsh_sha256"] = "a" * 64 with pytest.raises(ValidationError, match="digest"): ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) @@ -366,7 +366,7 @@ def test_resolved_image_rejects_digest_mismatch(multi_node_plan: ResolvedSlurmRu "artifact": {"path": "/wheels/data_designer.whl", "sha256": "a" * 64}, } ), - lambda payload: payload.update(image_distributions=list(reversed(payload["image_distributions"]))), + 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( @@ -526,12 +526,12 @@ def test_cross_record_validation_rejects_python_abi( validate_resolved_plan(authored_run, invalid_lock, invalid_plan) -def test_cross_record_validation_rejects_image_inventory( +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={"image_distributions": ()}) + 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( @@ -541,7 +541,7 @@ def test_cross_record_validation_rejects_image_inventory( ) invalid_plan = multi_node_plan.model_copy(update={"client": client}) - with pytest.raises(SlurmPlanContractError, match="image inventory"): + with pytest.raises(SlurmPlanContractError, match="runtime inspection"): validate_resolved_plan(authored_run, invalid_lock, invalid_plan) diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index c881ab1a6..bc402eaee 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -610,6 +610,50 @@ def test_publish_completed_lifecycle_registers_digest_bound_image_in_fresh_proce 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( 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 index 5e116d9a5..f34ed1d52 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" + "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" + "sha256": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac", + "provenance_digest": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" + "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "5b7614af448e9c647a2244257da37923784c3897e221a15f0a050b3e5f0c0324" + "sha256": "7eb39d9f28b7c073bc520c6596bdf7f6e2afc5e7fa7881945c5bcb34b19f9920" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" + "sha256": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "e1eb191572ad51ddfc1e14306d5efaed56bd56e7ad61cd77bb8351c85d593b0b" + "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 2a56d9d0a..fc5fe7041 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -129,17 +129,14 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( 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="${DD_VERIFIED_CLIENT_IMAGE}"' in script - assert "readonly DD_CLIENT_IMAGE=" not in script - assert ( - '--container-mounts="${DD_WORKSPACE_ROOT}:${DD_WORKSPACE_ROOT},${DD_SCRATCH_ROOT}:${DD_SCRATCH_CONTAINER_ROOT}"' - in script - ) - assert 'export PYTHONPATH="${DD_SCRATCH_CONTAINER_ROOT}/runtime"' in script - assert "--container-env=PYTHONPATH" in script + assert "--container-image=" not in script + assert "--container-mounts=" not in script + assert "--container-env=" not in script + assert f'readonly DD_CLIENT_PYTHON="{multi_node_plan.client.runtime.python_executable}"' in script + assert 'export PYTHONPATH="${DD_RUNTIME_ROOT}"' in script assert "--gres=none" in script assert "--export=ALL" in script - assert "python3 -m data_designer.slurm.state.attempt_identity" 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") diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 50dc506a1..465480439 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -196,7 +196,7 @@ def test_renderer_is_a_thin_entrypoint(single_node_plan: ResolvedSlurmRunPlan) - assert script.count("dd_slurm_run_allocation") == 1 assert 'readonly DD_ATTEMPT_ORDINAL="0012"' in script - assert len(script.splitlines()) <= 48 + assert len(script.splitlines()) <= 49 assert script.endswith("\n") diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index f9a2389fa..3a26bba22 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -47,7 +47,7 @@ def _resolve_fixture( ) -> EffectiveDataDesignerSlurmConfig: values = { "selected_profile": expected.selected_profile, - "client_image": expected.client.image, + "client_runtime": expected.client.runtime, "deployment_images": tuple(deployment.image for deployment in expected.deployments), "dependency_lock": dependency_lock, "runtime_bundle": expected.runtime_bundle, @@ -462,14 +462,6 @@ def test_compiler_rejects_resolved_image_identity_drift( single_node_plan: ResolvedSlurmRunPlan, ) -> None: effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) - client_image = effective.client_image.model_copy(update={"sha256": "a" * 64}) - dependency_lock = effective.dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) - - with pytest.raises(SlurmPlanCompilationError, match="inspection record"): - SlurmRunCompiler.compile( - effective.model_copy(update={"client_image": client_image, "dependency_lock": dependency_lock}) - ) - deployment_images = ( effective.deployment_images[0].model_copy(update={"sha256": "a" * 64}), *effective.deployment_images[1:], @@ -478,7 +470,7 @@ def test_compiler_rejects_resolved_image_identity_drift( SlurmRunCompiler.compile(effective.model_copy(update={"deployment_images": deployment_images})) -@pytest.mark.parametrize("invalid_field", ["client_image", "runtime_bundle"]) +@pytest.mark.parametrize("invalid_field", ["client_runtime", "runtime_bundle"]) def test_compiler_revalidates_nested_effective_contracts( authored_run_single: DataDesignerSlurmConfig, dependency_lock_single: ResolvedDependencyLock, @@ -486,8 +478,8 @@ def test_compiler_revalidates_nested_effective_contracts( invalid_field: str, ) -> None: effective = _resolve_fixture(authored_run_single, dependency_lock_single, single_node_plan) - if invalid_field == "client_image": - invalid_value = effective.client_image.model_copy(update={"path": "/workspace/images/client.txt"}) + 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"} @@ -860,9 +852,9 @@ def test_resolution_rejects_dependency_identity_mismatch( dependency_lock_single: ResolvedDependencyLock, single_node_plan: ResolvedSlurmRunPlan, ) -> None: - wrong_lock = dependency_lock_single.model_copy(update={"client_image_sha256": "a" * 64}) + wrong_lock = dependency_lock_single.model_copy(update={"client_runtime_sha256": "a" * 64}) - with pytest.raises(SlurmConfigResolutionError, match="client image"): + with pytest.raises(SlurmConfigResolutionError, match="client runtime"): _resolve_fixture(authored_run_single, wrong_lock, single_node_plan) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 6e6d06153..593e1a75f 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -11,12 +11,11 @@ from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import ResolvedSlurmRunPlan -from data_designer.slurm.runtime.bootstrap import RuntimeBootstrapManifest, build_runtime_manifest +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.paths import ALLOCATION_SCRATCH_CONTAINER_ROOT 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 @@ -25,7 +24,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: context = runtime_case.context - runtime_root = Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime" + runtime_root = Path("/tmp/data-designer-slurm-4101-0/runtime") log_directory = context.attempt_directory / "logs/execution-00000002" manifest = build_runtime_manifest( @@ -47,7 +46,7 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru 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] == ( - "python3", + context.plan.client.runtime.python_executable, "-m", "data_designer.slurm.runtime.entrypoint", "client", @@ -60,17 +59,40 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert all(step.role is not RuntimeStepRole.SERVER_PREFLIGHT for step in manifest.steps) 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 == ("PYTHONPATH",) + 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"] == ALLOCATION_SCRATCH_CONTAINER_ROOT + 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("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( @@ -105,7 +127,7 @@ def test_bootstrap_manifest_binds_retry_plan_to_control_and_client_workers(runti manifest = build_runtime_manifest( retry_context, {"SLURM_JOB_GPUS": "0"}, - runtime_root=Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime", + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), log_directory=context.attempt_directory / "logs/execution-00000002", layout=AllocationLayout(("compute-001",)), ) @@ -143,7 +165,7 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( manifest = build_runtime_manifest( context, {"SLURM_JOB_GPUS": "0,1,2,3,4,5,6,7"}, - runtime_root=Path(ALLOCATION_SCRATCH_CONTAINER_ROOT) / "runtime", + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), log_directory=context.attempt_directory / "logs/execution-00000002", layout=layout, ) @@ -167,8 +189,8 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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"] == f"{ALLOCATION_SCRATCH_CONTAINER_ROOT}/runtime" - assert endpoint.container_environment == ("PYTHONPATH",) + 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" diff --git a/packages/data-designer-slurm/tests/runtime/test_context.py b/packages/data-designer-slurm/tests/runtime/test_context.py index ee16de9f2..fe9553100 100644 --- a/packages/data-designer-slurm/tests/runtime/test_context.py +++ b/packages/data-designer-slurm/tests/runtime/test_context.py @@ -82,7 +82,7 @@ def test_runtime_retry_binding_rejects_tampered_arguments( ) -def test_allocation_context_reads_and_updates_state_through_remapped_workspace( +def test_allocation_context_rejects_container_remapped_workspace( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, single_node_plan: ResolvedSlurmRunPlan, @@ -164,14 +164,9 @@ def test_allocation_context_reads_and_updates_state_through_remapped_workspace( attempt_directory = fast_attempts / attempt.attempt_id attempt_directory.mkdir(parents=True, mode=0o700) - context, runtime_writer = load_allocation_context( - plan_path, - attempt_directory, - {"SLURM_ARRAY_TASK_ID": "0", "SLURM_ARRAY_JOB_ID": "4101"}, - ) - runtime_writer.update_attempt(context.attempt.model_copy(update={"state": AttemptLifecycleState.RUNNING})) - with runtime_writer.acquire_dataset_workspace(shard.shard_id, attempt.attempt_id, "never") as dataset_path: - assert dataset_path == attempt_directory / "dataset" - - assert context.attempt_directory.as_posix().startswith(logical_workspace) - assert host_writer.load_attempt(shard.shard_id, attempt.attempt_id).state is AttemptLifecycleState.RUNNING + 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_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 5e01d6040..de9ed5760 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -9,7 +9,6 @@ import json import os import sys -import venv import zipfile from dataclasses import dataclass from pathlib import Path @@ -71,7 +70,7 @@ def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.Capt assert "Traceback" not in captured.err -def test_container_phases_use_the_container_attempt_directory( +def test_native_phases_use_the_supplied_host_attempt_directory( monkeypatch: pytest.MonkeyPatch, runtime_case: RuntimeCase, ) -> None: @@ -263,7 +262,7 @@ def _patch_runtime_context( state: FakeStateStore, ) -> None: monkeypatch.setattr(entrypoint, "load_allocation_context", lambda *args, **kwargs: (runtime_case.context, state)) - monkeypatch.setattr(entrypoint, "get_container_path", lambda plan, path, **kwargs: path) + monkeypatch.setattr(entrypoint.SystemAllocationPreflight, "verify_client_runtime", lambda context: None) monkeypatch.setenv("SLURM_JOB_GPUS", "0") @@ -272,7 +271,7 @@ def _prepare_plugin_runtime_case( fake_plugin_overlay: Path, tmp_path: Path, ) -> _PluginRuntimeCase: - image_distributions = runtime_case.context.plan.client.image.inspection_facts.distributions + 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" ) @@ -281,8 +280,8 @@ def _prepare_plugin_runtime_case( _configure_identity_mount(payload, runtime_case.workspace) _configure_plugin_builder(payload) client = cast(dict[str, object], payload["client"]) - _configure_client_inspection(client, image_distributions, _write_test_installer(tmp_path)) - lock = _plugin_dependency_lock(client, image_distributions, wheel_path) + _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)) @@ -296,7 +295,7 @@ def _prepare_plugin_runtime_case( } ) context = AllocationContext(plan, plan.shards[0], attempt, runtime_case.context.attempt_directory) - image_metadata_directory = _write_inventory_bootstrap(image_distributions, tmp_path) + image_metadata_directory = _write_inventory_bootstrap(runtime_distributions, tmp_path) return _PluginRuntimeCase(context, plan_path, image_metadata_directory) @@ -319,42 +318,43 @@ def _configure_plugin_builder(payload: dict[str, object]) -> None: resolved_builder["content_sha256"] = compute_serialized_json_sha256(builder) -def _configure_client_inspection( +def _configure_client_runtime( client: dict[str, object], - image_distributions: tuple[InstalledDistribution, ...], - installer: Path, + runtime_distributions: tuple[InstalledDistribution, ...], ) -> None: - inspection = cast( - dict[str, object], - cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], + 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"} ) - inspection["python_abi"] = f"{interpreter_name()}{interpreter_version()}" - inspection["python_version"] = sys.version.split()[0] - inspection["installer_path"] = installer.as_posix() - inspection["distributions"] = [item.model_dump(mode="json") for item in image_distributions] 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], - image_distributions: tuple[InstalledDistribution, ...], + runtime_distributions: tuple[InstalledDistribution, ...], wheel_path: Path, ) -> ResolvedDependencyLock: - inspection = cast( - dict[str, object], - cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], - ) + runtime = cast(dict[str, object], client["runtime"]) return ResolvedDependencyLock.model_validate( { "schema_version": 1, "resolver_version": "resolver-1", - "python_abi": inspection["python_abi"], - "client_image_sha256": cast(dict[str, object], client["image"])["sha256"], + "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, - "image_distributions": image_distributions, + "base_distributions": runtime_distributions, "overlay_packages": ( { "name": "fake-data-designer-plugin", @@ -399,18 +399,11 @@ def _build_plugin_wheel(source: Path, destination: Path) -> None: archive.writestr(name, content) -def _write_test_installer(tmp_path: Path) -> Path: - """Create an offline pip executable without changing the test checkout.""" - environment = tmp_path / "installer-environment" - venv.EnvBuilder(with_pip=True, symlinks=True).create(environment) - return environment / "bin/pip" - - def _write_inventory_bootstrap( distributions: tuple[InstalledDistribution, ...], tmp_path: Path, ) -> Path: - """Present plan-recorded image metadata while imports use the editable checkout.""" + """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) @@ -445,7 +438,7 @@ def _controlled_distributions(**kwargs): def _use_controlled_image_inventory(monkeypatch: pytest.MonkeyPatch, image_metadata_directory: Path) -> None: - """Limit child-process inventory discovery to the plan image and installed overlay.""" + """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) diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 604e3e737..4c95f24c5 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -187,16 +187,37 @@ def test_attempt_directory_must_be_restrictive(runtime_case: RuntimeCase) -> Non SystemAllocationPreflight.verify_attempt_directory(runtime_case.context.attempt_directory) -def test_preflight_translates_host_paths_into_container( +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" - container_attempt = container_workspace / logical_attempt.relative_to(runtime_case.workspace) - logical_attempt.rmdir() - container_attempt.mkdir(parents=True, mode=0o700) plan = runtime_case.context.plan.model_copy( update={ "container_mounts": ( @@ -215,11 +236,12 @@ def test_preflight_translates_host_paths_into_container( ) 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(container_workspace.as_posix()) for path in verified_paths) - assert not any(path.startswith(runtime_case.workspace.as_posix()) for path in 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_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index b4cafd17d..62b6cd3ec 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -32,7 +32,9 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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) + 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))) @@ -51,6 +53,7 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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}}" """ @@ -413,6 +416,7 @@ def test_step_runner_builds_one_coordinated_srun_across_selected_nodes() -> None 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 @@ -498,6 +502,8 @@ def _plan( 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 = [ @@ -511,7 +517,7 @@ def _plan( "host_node_index": 0, "authored": {"cpus": 1}, "dependency_lock": {"path": lock[0], "sha256": lock[1]}, - "image": {"path": client[0], "sha256": client[1]}, + "runtime": {"python_executable": python_executable}, }, "resolved_gpus_per_node": 1, "selected_profile": {"profile": {"gpu_request_mode": gpu_request_mode}}, @@ -535,7 +541,7 @@ def _manifest( artifacts: tuple[tuple[str, str], ...], plan_sha256: str, ) -> dict[str, object]: - _, _, client, server = artifacts + _, _, _, server = artifacts return { "schema_version": 1, "run_id": "run-shell", @@ -544,7 +550,7 @@ def _manifest( "plan_sha256": plan_sha256, "all_secret_environment_names": ["SERVER_TOKEN", "SOURCE_TOKEN"], "steps": [ - _step(attempt_directory, "client-preflight", "client_preflight", client[0], "fake-preflight"), + _step(attempt_directory, "client-preflight", "client_preflight", None, "fake-preflight"), _step( attempt_directory, "deployment-00000-replica-00000-rank-00000", @@ -560,7 +566,7 @@ def _manifest( attempt_directory, "deployment-00000-endpoint", "endpoint", - client[0], + None, "fake-endpoint", readiness={"host": "127.0.0.1", "port": 17000, "path": "/health", "deadline_seconds": 2}, ), @@ -568,10 +574,9 @@ def _manifest( attempt_directory, "client-generation", "client", - client[0], + None, "fake-client", secret_environment={"SOURCE_TOKEN": "SOURCE_TOKEN"}, - container_environment=["SOURCE_TOKEN"], ), ], } @@ -581,7 +586,7 @@ def _step( attempt_directory: Path, step_id: str, role: str, - image_path: str, + image_path: str | None, command: str, *, gpu_indices: list[int] | None = None, @@ -593,6 +598,7 @@ def _step( 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, @@ -623,10 +629,10 @@ def _fake_srun() -> str: [[ ${arguments[index]} == -- ]] && break done command=${arguments[index + 1]} -if [[ ${command} == python3 ]]; then +if [[ ${command} == */python3 ]]; then operation=${arguments[index + 4]} [[ ! ${SOURCE_TOKEN+x} && ! ${SERVER_TOKEN+x} && ! ${CUDA_VISIBLE_DEVICES+x} ]] - [[ ${arguments[*]} == *--container-env=PYTHONPATH,SLURM_JOB_GPUS* ]] + [[ ${arguments[*]} != *--container-* ]] if [[ ${operation} == prepare ]]; then for ((position = index + 5; position < ${#arguments[@]}; position++)); do if [[ ${arguments[position]} == --manifest ]]; then diff --git a/packages/data-designer-slurm/tests/runtime/test_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index e5e6a3fb8..91a8f2657 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -89,8 +89,10 @@ def test_all_processes_use_structured_srun_steps_and_sanitized_environment(runti for argument in step.command if argument.startswith("--container-mounts=") ) - for step in all_steps + 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 @@ -163,7 +165,7 @@ def test_client_worker_receives_only_persisted_identity_and_logical_endpoint(run separator = step.command.index("--") worker = step.command[separator + 1 :] assert worker[:5] == ( - "python3", + context.plan.client.runtime.python_executable, "-m", "data_designer.slurm.client.worker", "run", @@ -213,7 +215,8 @@ def test_endpoint_step_uses_resolved_retry_policy_and_backends(runtime_case: Run ] assert "--max-waiting-requests" not in step.command assert step.environment["PYTHONPATH"].endswith("/runtime") - assert "--container-env=PYTHONPATH" in step.command + 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 @@ -282,7 +285,7 @@ def test_client_receives_client_secrets_without_server_only_secrets( assert step.environment["PACKAGE_INDEX_TOKEN"] == "client-secret" assert "HF_TOKEN" not in step.environment - assert "--container-env=DATA_DESIGNER_SLURM_SCRATCH_ROOT,PACKAGE_INDEX_TOKEN,SLURM_JOB_GPUS" in step.command + assert not any(argument.startswith("--container-") for argument in step.command) with pytest.raises(SlurmRuntimeError, match="PACKAGE_INDEX_TOKEN"): DefaultClientStepBuilder().build_preflight_step( diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index b9b6855c9..659af84b1 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -22,6 +22,7 @@ from data_designer.slurm.config import ( DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, + ImageInspectionRecord, ImageKind, ImageRef, ) @@ -473,28 +474,28 @@ def test_service_boundary_does_not_swallow_cancellation_signals( def test_image_service_returns_correlated_resolution( single_node_plan: ResolvedSlurmRunPlan, ) -> None: - reference = single_node_plan.client.image.authored_ref - image = single_node_plan.client.image - resolver = FakeImageResolver((((reference, ImageKind.CLIENT), image),)) + 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.CLIENT) + result = SlurmImageService(resolver).resolve(reference, expected_kind=ImageKind.SERVING) assert result is image - assert resolver.calls == [(reference, ImageKind.CLIENT)] + 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.client.image.authored_ref - resolver = FakeImageResolver((((reference, ImageKind.CLIENT), object()),)) # type: ignore[arg-type] + 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.CLIENT) + 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.CLIENT)] + assert resolver.calls == [(reference, ImageKind.SERVING)] resolver.assert_complete() @@ -502,7 +503,7 @@ def test_image_service_rejects_untyped_image_kind(single_node_plan: ResolvedSlur service = SlurmImageService(FakeImageResolver(())) with pytest.raises(SlurmServiceError) as caught: - service.resolve(single_node_plan.client.image.authored_ref, expected_kind="client") # type: ignore[arg-type] + 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 @@ -512,7 +513,7 @@ def test_image_service_rejects_untyped_reference() -> None: service = SlurmImageService(FakeImageResolver(())) with pytest.raises(SlurmServiceError) as caught: - service.resolve(object(), expected_kind=ImageKind.CLIENT) # type: ignore[arg-type] + 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 @@ -531,23 +532,23 @@ def test_image_service_requires_package_owned_manager() -> None: @pytest.mark.parametrize("mismatch", ["reference", "kind"]) def test_image_service_rejects_uncorrelated_results( single_node_plan: ResolvedSlurmRunPlan, + client_image_inspection: ImageInspectionRecord, mismatch: str, ) -> None: - reference = single_node_plan.client.image.authored_ref + source = single_node_plan.deployments[0].image + reference = source.authored_ref if mismatch == "reference": - source = single_node_plan.client.image - image = source.model_copy(update={"authored_ref": ImageRef(name="other-client")}) + image = source.model_copy(update={"authored_ref": ImageRef(name="other-serving")}) else: - source = single_node_plan.deployments[0].image - image = source.model_copy(update={"authored_ref": reference}) - resolver = FakeImageResolver((((reference, ImageKind.CLIENT), image),)) + 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.CLIENT) + 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.CLIENT)] + assert resolver.calls == [(reference, ImageKind.SERVING)] resolver.assert_complete() diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 8af4bd7ae..f1a8b52c7 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -14,6 +14,7 @@ import data_designer.slurm.services.retry_collection as retry_collection_module from data_designer.slurm.benchmark.compiler import BenchmarkCompiler from data_designer.slurm.client.dependencies import ResolvedClientDependencies +from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector from data_designer.slurm.config import ( BenchmarkBaseRun, BuilderInput, @@ -59,6 +60,14 @@ ) +@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) + + class _Launcher: def __init__( self, @@ -187,7 +196,7 @@ def _register_images( plan: ResolvedSlurmRunPlan, ) -> None: store = ImageRegistryStore(tmp_path) - pairs = ((authored.client.image, plan.client.image),) + tuple( + pairs = tuple( (deployment.server.image, resolved.image) for deployment, resolved in zip(authored.deployments, plan.deployments, strict=True) ) @@ -231,6 +240,30 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( 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, @@ -1246,4 +1279,4 @@ def test_production_image_registry_operations( removed = service.remove(images[0].name) assert selected == removed == images[0] - assert len(service.list()) == 1 + assert service.list() == () diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index 4c7766487..28155e9fc 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -10,7 +10,7 @@ import pytest import data_designer.slurm.serving.resolver as resolver_module -from data_designer.slurm.config import QueueBackpressureConfig +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 @@ -138,8 +138,12 @@ def test_resolution_rejects_unsupported_runtime_versions( _resolve(plan, placement.deployment_id) -def test_resolution_rejects_image_inspection_mismatch(single_node_plan: ResolvedSlurmRunPlan) -> None: - placement = single_node_plan.deployments[0].model_copy(update={"image": single_node_plan.client.image}) +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"): diff --git a/packages/data-designer-slurm/tests/serving/test_serving_records.py b/packages/data-designer-slurm/tests/serving/test_serving_records.py index 3e8a59cc3..a79086ca9 100644 --- a/packages/data-designer-slurm/tests/serving/test_serving_records.py +++ b/packages/data-designer-slurm/tests/serving/test_serving_records.py @@ -8,6 +8,7 @@ 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 @@ -96,12 +97,15 @@ def test_process_specs_reject_one_field_topology_drift( 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"] = multi_node_plan.client.image.model_dump(mode="json") + 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": diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 3afccceab..e08096373 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -12,11 +12,12 @@ #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-001/resolved-plan.json" -readonly DD_PLAN_SHA256="302dfcdcdd61b12f7e227dce10a307e2205e0165ec66969ab04d1e7a028790c1" +readonly DD_PLAN_SHA256="3984ca2a246f948b98aec22323e2d6724383739a90db0666918c8aaf3b41a89c" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 7a471e039..544d56739 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -12,11 +12,12 @@ #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="7db47928d64daf85b6fd28b510cfd3b1bebcc1672a1ea1163ecb2d86051dc3ac" +readonly DD_PLAN_SHA256="fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py index 324a8948d..594ca9bbb 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_dependencies.py @@ -47,7 +47,7 @@ def test_fake_dependency_installer_scripts_success_and_digest_mismatch( ) -> None: first_target = tmp_path / "compatible" second_target = tmp_path / "digest-mismatch" - mismatched_lock = dependency_lock.model_copy(update={"client_image_sha256": "a" * 64}) + 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), diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 29ef2837e..024dcbf36 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="202f1070fe8351f8afd3f71ca0353d431b8b4b66c1643e9767a2650c20efbe20", + expected_fixture_sha256="f130dd88efe97ebb060eeef743a62f1f1d27e8f16f0f71d6ec31791b4ccee6b3", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="ecffd8264aa022dce66c2eeee447fd71babb507287cfe059f35a9fd1ae009953", + expected_fixture_sha256="eefbd439a72c59dc09a97f5ebc73f9fdb5ee13a2f41f1447ab560e8e2d589cca", ) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 1f0ba524f..f591869bc 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "4283c4f3b6d5d5181c91651ad10377a43737f1c9d86949e76e9eb4d37626ed0e" +_COMPATIBILITY_FIXTURE_DIGEST = "182770715c75496f62ed196cb04f088638a01c2d7ea7504547371cfce74c2b04" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 0cf0b83e1..380ec0ebe 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -12,7 +12,7 @@ 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 +from data_designer.slurm.config import DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, ImageBuildRequest from data_designer.slurm.services import ( SlurmCollectionExecution, SlurmRetryExecution, @@ -356,6 +356,38 @@ def test_image_add_rejects_mutable_oci_source(source: str) -> None: } +def test_image_add_resolves_versioned_vllm_before_import(monkeypatch: pytest.MonkeyPatch) -> 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, "create_slurm_image_service", lambda **_kwargs: ImageService()) + + 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" + + +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 / "profile.yml" workspace = tmp_path / "workspace" From 2ad7e7148cf7a7de38e52cf4e8c168228c388216 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 5 Oct 2026 16:20:54 -0600 Subject: [PATCH 33/62] test(slurm): repair runtime test fixtures (#980) Signed-off-by: Nabin Mulepati --- packages/data-designer-slurm/tests/client/test_environment.py | 2 +- packages/data-designer-slurm/tests/runtime/test_entrypoint.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 9386229e0..167050459 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -190,7 +190,7 @@ def distributions(**kwargs: object) -> tuple[()]: submit_path = "/submit/venv/site-packages" interpreter_path = "/control/venv/site-packages" - user_site_path = "/users/alice/.local/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]) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index de9ed5760..b8963513a 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -424,7 +424,8 @@ def _write_inventory_bootstrap( def _controlled_distributions(**kwargs): - if kwargs.get("path") is not None: + 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")) From d270f6d7d09995b4cde0a285b6f5cc973c8126e2 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 6 Oct 2026 09:05:03 -0600 Subject: [PATCH 34/62] fix(slurm): show CLI submission progress (#979) * fix(slurm): show CLI submission progress Signed-off-by: Nabin Mulepati * fix(slurm): keep progress failures from stopping submission Signed-off-by: Nabin Mulepati * feat(slurm): add opt-in image job log following Signed-off-by: Nabin Mulepati * fix(slurm): preserve and bound followed log output Signed-off-by: Nabin Mulepati * fix(slurm): flush final and carriage-return logs Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/images.mdx | 9 + .../slurm/benchmark/execution.py | 26 +- .../src/data_designer/slurm/cli.py | 40 ++- .../src/data_designer/slurm/cli_benchmark.py | 8 +- .../data_designer/slurm/images/lifecycle.py | 36 ++ .../slurm/services/image_lifecycle.py | 167 +++++++++- .../data_designer/slurm/services/wiring.py | 6 + .../tests/benchmark/test_workflow.py | 48 +++ .../tests/services/test_image_wiring.py | 313 +++++++++++++++++- .../data-designer-slurm/tests/test_cli.py | 126 ++++++- 10 files changed, 754 insertions(+), 25 deletions(-) diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index cc37d8d37..c46dcfd90 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -33,6 +33,15 @@ Custom OCI sources must use `repository@sha256:<64 lowercase hex characters>`. T 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 scheduler state every 30 seconds. To also see the image job's standard output and error while it runs, add `--follow-logs`: + +```bash +data-designer slurm image add vllm/vllm-openai:v0.22.0 \ + --kind serving --name vllm --follow-logs +``` + +The log stream is opt-in, bounded to 1 MiB, and written to standard error; the final machine-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. If you need durable logs, capture standard error in your own private file. + ## Register an existing SQSH ```bash 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 index 871d88e2d..5bbe370f2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -47,6 +47,7 @@ ChildRunServiceFactory = Callable[[Identifier], SlurmRunService] ChildConfigLoader = Callable[[Identifier], DataDesignerSlurmConfig | None] ChildSubmissionLoader = Callable[[Identifier], bool] +ProgressReporter = Callable[[str], None] class SystemBenchmarkBackend: @@ -60,6 +61,7 @@ def __init__( 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 @@ -67,6 +69,7 @@ def __init__( 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, @@ -113,10 +116,16 @@ def run( ) from None failures: list[SlurmServiceErrorCode] = [] - for case in compiled.cases: + 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, @@ -129,10 +138,17 @@ def run( 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) + self._report_progress(f"Benchmark case {index}/{total}: submission failed; continuing.") except Exception: failures.append(SlurmServiceErrorCode.INTERNAL) + self._report_progress(f"Benchmark case {index}/{total}: submission failed; continuing.") + self._report_progress( + f"Benchmark launch: {submitted} submitted, {reused} already submitted, {len(failures)} failed." + ) if failures: code = _partial_failure_code(failures) raise SlurmServiceError( @@ -142,6 +158,14 @@ def run( ) 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, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index eea437382..c480702ad 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -4,6 +4,7 @@ from __future__ import annotations import re +import sys from collections.abc import Callable from enum import Enum from functools import partial @@ -76,6 +77,7 @@ def execute_command( 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) return service.execute(config, source_root=run_file.resolve().parent, dry_run=dry_run) @@ -162,6 +164,8 @@ def retry_command( 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( @@ -185,6 +189,7 @@ def merge_command( ) -> None: """Submit winner-driven collection as a zero-GPU Slurm job.""" operation = SlurmServiceOperation.COLLECT_RUN + _emit_progress("Preparing and submitting collection job...") result = _invoke( operation, lambda: create_slurm_run_service(profile_file=profile_file, cluster=cluster).collect( @@ -242,10 +247,11 @@ def image_add_command( kind: str = typer.Option("serving", "--kind", help="Image role (defaults to serving)"), name: str | None = typer.Option(None, "--name"), replace: bool = typer.Option(False, "--replace"), + follow_logs: bool = typer.Option(False, "--follow-logs", help="Show bounded image-job logs on stderr"), 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.""" + """Import or inspect an image and register its alias, reporting progress in terminals.""" operation = SlurmServiceOperation.ADD_IMAGE def add() -> BaseModel: @@ -257,6 +263,7 @@ def add() -> BaseModel: 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: @@ -281,7 +288,13 @@ def add() -> BaseModel: "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) - return create_slurm_image_service(profile_file=profile_file, cluster=cluster).add(request, replace=replace) + _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 else None, + ).add(request, replace=replace) _emit_result(_invoke(operation, add)) @@ -372,6 +385,27 @@ 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(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(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 _derive_image_name(source: str) -> str: if source.endswith(".sqsh"): candidate = Path(source).stem @@ -387,7 +421,7 @@ def _bounded_error_message(message: str, *, fallback: str) -> str: return sanitized if len(sanitized) <= 512 else f"{sanitized[:509]}..." -app.add_typer(create_benchmark_app(_invoke, _emit_result), name="benchmark") +app.add_typer(create_benchmark_app(_invoke, _emit_result, _emit_progress), name="benchmark") def create_cli() -> click.Command: 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 index 83c3f56b4..b01cc7667 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py @@ -18,6 +18,7 @@ _Invoke = Callable[[SlurmServiceOperation, Callable[[], BaseModel]], BaseModel] _Emit = Callable[[BaseModel], None] +_Progress = Callable[[str], None] def _run_benchmark( @@ -26,9 +27,11 @@ def _run_benchmark( 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) + 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) @@ -49,7 +52,7 @@ def _analyze_benchmark( ) -def create_benchmark_app(invoke: _Invoke, emit: _Emit) -> typer.Typer: +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) @@ -70,6 +73,7 @@ def run_command( profile_file=profile_file, cluster=cluster, force=force, + progress=progress, ), ) ) 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 index a14eab697..b8c6f48be 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -13,6 +13,7 @@ from dataclasses import dataclass, field from importlib import resources from pathlib import Path +from typing import Literal from pydantic import TypeAdapter, ValidationError @@ -455,6 +456,41 @@ def cleanup_prepared_image_lifecycle(prepared: PreparedImageLifecycleJob) -> Non 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 _verify_prepared_lifecycle( prepared: PreparedImageLifecycleJob, job_directory_descriptor: int, 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 index a335e248f..1e84db3b5 100644 --- 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 @@ -5,6 +5,8 @@ from __future__ import annotations +import codecs +import re import time from collections.abc import Callable, Sequence from datetime import datetime, timedelta, timezone @@ -23,6 +25,7 @@ cleanup_prepared_image_lifecycle, prepare_image_lifecycle_job, publish_completed_image_lifecycle, + read_image_lifecycle_log, submit_prepared_image_lifecycle, ) from data_designer.slurm.images.records import RegisteredImage @@ -42,15 +45,123 @@ LifecycleIdFactory = Callable[[], str] Clock = Callable[[], datetime] Sleeper = Callable[[float], None] +ProgressReporter = Callable[[str], None] +LogReporter = Callable[[str], None] -_POLL_INTERVAL_SECONDS = 300.0 +_POLL_INTERVAL_SECONDS = 30.0 _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]") 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) -> None: + """Print newly appended bytes; after termination, drain the bounded remainder.""" + if self._limit_reported: + return + 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._offsets[stream] = offset + self._source_bytes += len(content) + if content: + received = True + self._append(stream, content) + if self._limit_reported: + return + 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() + + 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: + self._emit(f"[image {stream}] {_UNSAFE_TERMINAL_CHARACTERS.sub('?', line)}") + + 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 @@ -75,12 +186,16 @@ def __init__( 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.""" @@ -95,6 +210,7 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: raise _unavailable("image lifecycle job cannot be prepared") from None try: + self._report_progress("Submitting image inspection job...") receipt = submit_prepared_image_lifecycle(prepared, self._launcher) except SlurmSubmissionError as error: if not error.may_have_succeeded: @@ -107,8 +223,11 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: _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: - self._wait_for_success(receipt.job_id) + self._report_progress("Image inspection job submitted; checking every 30 seconds.") + self._wait_for_success(receipt.job_id, follower) + self._report_progress("Inspection complete; verifying and registering image...") return publish_completed_image_lifecycle(prepared, replace=replace) except (ImageConflictError, ImageVerificationError): raise SlurmServiceError( @@ -128,13 +247,13 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: _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) + 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) + self._cancel_and_cleanup(receipt.job_id, prepared, follower) raise - def _wait_for_success(self, job_id: int) -> None: + def _wait_for_success(self, job_id: int, follower: _JobLogFollower | None) -> None: client = _RecordingObservationClient(self._launcher) observations = SchedulerObservationCollector(client) previous: SchedulerObservation | None = None @@ -146,9 +265,13 @@ def _wait_for_success(self, job_id: int) -> None: observed_at=observed_at, previous={job_id: previous}, )[0] + self._report_progress(f"Image inspection job: {observation.state.value.replace('_', ' ')}.") + if follower is not None: + follower.drain() 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: + self._report_progress("Waiting for successful job exit evidence...") if completed_without_accounting_deadline is None: completed_without_accounting_deadline = observed_at + _ACCOUNTING_EXIT_LAG elif observed_at > completed_without_accounting_deadline: @@ -160,23 +283,46 @@ def _wait_for_success(self, job_id: int) -> None: 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 _cancel_and_cleanup(self, job_id: int, prepared: PreparedImageLifecycleJob) -> None: + 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) - if not self._wait_for_termination(job_id): - return + terminated = self._wait_for_termination(job_id) except (SlurmLauncherError, SlurmStateError, OSError, ValueError): - return - _cleanup_failed_lifecycle(prepared) + 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) @@ -190,6 +336,7 @@ def _wait_for_termination(self, job_id: int) -> bool: 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 5a5abb6f8..2b731bed8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -796,6 +796,8 @@ def create_slurm_image_service( 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) @@ -806,6 +808,8 @@ def create_slurm_image_service( 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) @@ -823,6 +827,7 @@ def create_slurm_benchmark_service( 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) @@ -855,6 +860,7 @@ def create_child_service(run_id: Identifier) -> SlurmRunService: create_child_service, observer, selected_clock, + progress=progress, ) return SlurmBenchmarkService(backend) diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index 735ba900f..10f31d5ac 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -191,6 +191,7 @@ def test_run_persists_manifest_before_submission_and_is_idempotent( submissions = set() calls = [] failures = set() + progress: list[str] = [] benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id backend = SystemBenchmarkBackend( tmp_path, @@ -199,6 +200,7 @@ def test_run_persists_manifest_before_submission_and_is_idempotent( 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) @@ -207,6 +209,48 @@ def test_run_persists_manifest_before_submission_and_is_idempotent( 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( @@ -295,6 +339,7 @@ def test_partial_submission_attempts_all_children_and_resumes_only_missing( submissions = set() calls = [] failures = {missing_id} + progress: list[str] = [] benchmark_root = tmp_path / "benchmarks" / compiled.benchmark_id backend = SystemBenchmarkBackend( tmp_path, @@ -303,11 +348,14 @@ def test_partial_submission_attempts_all_children_and_resumes_only_missing( lambda: NOW, child_config_loader=configs.get, child_submission_loader=submissions.__contains__, + progress=progress.append, ) with pytest.raises(SlurmServiceError, match="1 of 2"): backend.run(config, source_root=tmp_path, force=True) assert len(calls) == 2 + assert "Benchmark case 1/2: submission failed; 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) diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index 89d95b7a0..776002597 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -120,6 +120,21 @@ def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ... return super().query_accounting(selectors) +class _TransitionLauncher(_Launcher): + def __init__(self, *, on_submit: Callable[[], None]) -> None: + super().__init__(queue_state=SchedulerState.PENDING, accounting_state=None, on_submit=on_submit) + self.queue_queries = 0 + + 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 == 3: + 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) @@ -187,10 +202,274 @@ def test_default_image_add_waits_for_completed_job_accounting_evidence(tmp_path: assert registered.path == source.as_posix() assert launcher.queue_queries == 3 assert launcher.accounting_queries == 3 - assert clock.sleep_calls == [300.0, 300.0] + assert clock.sleep_calls == [30.0, 30.0] + assert not _job_directory(workspace).exists() + + +def test_image_add_reports_progress_on_each_thirty_second_check(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 inspection job...", + "Image inspection job submitted; checking every 30 seconds.", + "Image inspection job: completed.", + "Waiting for successful job exit evidence...", + "Image inspection job: completed.", + "Waiting for successful job exit evidence...", + "Image inspection job: completed.", + "Inspection complete; verifying and registering image...", + ] + + +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)) + 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 "Image inspection job: pending." in progress + assert "Image inspection job: running." in progress + assert "Image inspection job: completed." in progress + + +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\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", + ] 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" @@ -287,7 +566,7 @@ def test_default_image_add_waits_for_cancelled_job_before_cleanup(tmp_path: Path assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE assert launcher.cancellations == [_JOB_ID] - assert clock.sleep_calls == [300.0, 300.0, 300.0] + assert clock.sleep_calls == [30.0] * 12 assert not _job_directory(workspace).exists() @@ -304,7 +583,7 @@ def test_default_image_add_retains_state_without_terminal_cancellation_evidence( ) assert launcher.cancellations == [_JOB_ID] - assert clock.sleep_calls == [300.0, 300.0, 300.0] + assert clock.sleep_calls == [30.0] * 21 assert _job_directory(workspace).is_dir() @@ -322,7 +601,33 @@ def test_default_image_add_retains_completed_job_without_accounting_exit_evidenc assert caught.value.code is SlurmServiceErrorCode.UNAVAILABLE assert launcher.cancellations == [_JOB_ID] - assert clock.sleep_calls == [300.0, 300.0, 300.0] + 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() diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 380ec0ebe..24b5ede33 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -40,10 +40,11 @@ def execute( self.calls.append((config, source_root, dry_run)) return SlurmRunExecution( run_id="run-0001", - state="dry_run", + state="dry_run" if dry_run else "submitted", plan_sha256="1" * 64, shard_count=1, - batch_script="#!/bin/bash\n", + batch_script="#!/bin/bash\n" if dry_run else None, + job_id=None if dry_run else 43, ) def retry( @@ -136,6 +137,61 @@ def test_execute_emits_deterministic_json_and_forwards_actions( "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 result.stderr == "Preparing and submitting run...\n" + + +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 message.startswith("Preparing and submitting"): + 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( @@ -148,7 +204,16 @@ def test_benchmark_cli_forwards_run_and_analysis_actions( benchmark_file = tmp_path / "benchmark.json" benchmark_file.write_text(benchmark_config.serialize_json()) service = _BenchmarkService(benchmark_manifest, benchmark_report) - monkeypatch.setattr(benchmark_cli_module, "create_slurm_benchmark_service", lambda **_: service) + + 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(cli_module, "_progress_enabled", lambda: True) run_result = CliRunner().invoke( cli_module.create_cli(), @@ -165,6 +230,7 @@ def test_benchmark_cli_forwards_run_and_analysis_actions( 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 run_result.stderr == "Preparing benchmark launch...\nBenchmark case 1/1: submitted.\n" def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: @@ -216,11 +282,13 @@ def test_retry_auto_selects_tasks_and_confirms_submission(monkeypatch) -> None: 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 result.stderr == "Preparing and submitting retry...\n" def test_retry_decline_emits_stable_result(monkeypatch) -> None: @@ -254,6 +322,7 @@ def test_retry_noninteractive_confirmation_is_invalid_request(monkeypatch) -> No 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" @@ -280,6 +349,7 @@ def test_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypat "state": "submitted", } assert service.collection_calls == [(input_path, output_path, 2)] + assert result.stderr == "Preparing and submitting collection job...\n" @pytest.mark.parametrize( @@ -356,7 +426,11 @@ def test_image_add_rejects_mutable_oci_source(source: str) -> None: } -def test_image_add_resolves_versioned_vllm_before_import(monkeypatch: pytest.MonkeyPatch) -> None: +@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] = [] @@ -368,7 +442,15 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest return request monkeypatch.setattr(cli_module, "resolve_versioned_vllm_source", lambda _source: pinned) - monkeypatch.setattr(cli_module, "create_slurm_image_service", lambda **_kwargs: ImageService()) + monkeypatch.setattr(cli_module, "_progress_enabled", lambda: interactive) + + def create_service(**kwargs: object) -> ImageService: + progress = kwargs["progress"] + assert callable(progress) + progress("Image inspection job: pending.") + return ImageService() + + monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) result = CliRunner().invoke(cli_module.create_cli(), ["image", "add", source]) @@ -377,6 +459,40 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest 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 result.stderr.splitlines() == ( + [ + "Resolving vLLM image tag...", + "Preparing image registration...", + "Image inspection job: pending.", + ] + 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 result.stderr == "[image stderr] inspection started\n" def test_image_add_rejects_versioned_vllm_as_client() -> None: From 6e8ce60753fa1aa42efa3f7a0840cd29ca5d7167 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 6 Oct 2026 15:26:57 -0600 Subject: [PATCH 35/62] fix(slurm): support Enroot 3.5 OCI imports (#988) Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/images.mdx | 2 +- fern/versions/latest/pages/slurm/overview.mdx | 2 +- .../latest/pages/slurm/troubleshooting.mdx | 2 +- .../data_designer/slurm/images/lifecycle.py | 18 +++++--- .../tests/images/test_lifecycle.py | 43 ++++++++++++++----- 5 files changed, 48 insertions(+), 19 deletions(-) diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index c46dcfd90..da2e5a102 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -12,7 +12,7 @@ data-designer slurm image add vllm/vllm-openai:v0.22.0 \ --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. 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. +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. diff --git a/fern/versions/latest/pages/slurm/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index d884db31c..fba10bde7 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -12,7 +12,7 @@ The submit environment needs: - `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. They also need Enroot 3.5 or later for existing SQSH images; digest-pinned OCI imports need Enroot 4.0 or later. Serving images must contain vLLM. +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. diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index f404fece0..ecaf89215 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -46,7 +46,7 @@ Set `scheduler.bin_path` in the selected profile to the absolute directory conta - Confirm the image-build partition accepts CPU jobs. - Confirm Enroot is available on the allocated node. -- Use Enroot 4.0 or later for OCI import and 3.5 or later for existing SQSH inspection. +- 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. 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 index b8c6f48be..f64319136 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -44,7 +44,7 @@ _RESOURCE_PACKAGE = "data_designer.slurm.images.resources" _SCRATCH_RESOURCE_PACKAGE = "data_designer.slurm.runtime" _IDENTIFIER_ADAPTER = TypeAdapter(Identifier) -_MINIMUM_ENROOT_OCI_VERSION = (4, 0) +_MINIMUM_ENROOT_OCI_VERSION = (3, 5) _MINIMUM_ENROOT_SQSH_VERSION = (3, 5) _MAXIMUM_INSPECTION_SIZE = 1024 * 1024 @@ -557,12 +557,20 @@ def _stage_resource(job_directory: Path, filename: str) -> ArtifactReference: def _format_enroot_oci_uri(source: str) -> str: - registry_or_namespace, separator, remainder = source.partition("/") + 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://docker.io#library/{source}" + 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": - return f"docker://{registry_or_namespace}#{remainder}" - return f"docker://docker.io#{source}" + 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: diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index bc402eaee..ea27d6d70 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -130,7 +130,7 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ 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:release@sha256:{'b' * 64}" + 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 @@ -139,9 +139,9 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ 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 4 0") + 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 4 0 "digest-pinned OCI imports"' 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 @@ -149,8 +149,8 @@ def test_image_lifecycle_renderer_uses_explicit_cpu_profile_and_safe_mounts(tmp_ @pytest.mark.parametrize( ("source_kind", "enroot_version"), - (("oci", "4.0.0"), ("existing", "3.5.0")), - ids=("oci-import-enroot-4", "existing-sqsh-enroot-3.5"), + (("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, @@ -180,6 +180,8 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( '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' @@ -227,8 +229,8 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( @pytest.mark.parametrize( ("source_kind", "enroot_version", "required_version"), - (("oci", "3.5.0", "4.0"), ("existing", "3.4.1", "3.5")), - ids=("digest-import-before-4", "existing-sqsh-before-3.5"), + (("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, @@ -311,18 +313,37 @@ def test_prepare_rejects_oci_sources_that_could_persist_credentials(tmp_path: Pa ( ( f"vllm/vllm-openai@sha256:{'a' * 64}", - f"docker://docker.io#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://docker.io#library/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"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", ), - ids=("docker-hub-namespace", "docker-hub-library", "explicit-registry"), ) def test_renderer_normalizes_digest_qualified_sources_for_enroot( tmp_path: Path, From b7892667333263fd0eed01d28db48e731dae0120 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 6 Oct 2026 19:16:30 -0600 Subject: [PATCH 36/62] feat(slurm): improve CLI output and image logs (#992) * feat(slurm): improve CLI output and image logs Make operator-facing Slurm results readable while preserving opt-in JSON for automation. Surface bounded, redacted sbatch rejection details and print status/log guidance after successful generation submission. Follow image lifecycle logs by default in interactive terminals, with explicit flags to enable or disable streaming. Update the Slurm docs and cover both output modes and submission diagnostics in tests. Signed-off-by: Nabin Mulepati * fix(slurm): harden CLI guidance and output Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/cli-reference.mdx | 6 +- .../latest/pages/slurm/getting-started.mdx | 2 +- fern/versions/latest/pages/slurm/images.mdx | 6 +- .../latest/pages/slurm/operations.mdx | 11 +- .../latest/pages/slurm/troubleshooting.mdx | 15 +- .../src/data_designer/slurm/cli.py | 115 +++++++++++++--- .../slurm/services/image_lifecycle.py | 4 +- .../data_designer/slurm/services/wiring.py | 57 +++++++- .../tests/services/test_image_wiring.py | 5 +- .../tests/services/test_wiring.py | 130 +++++++++++++++++- .../data-designer-slurm/tests/test_cli.py | 119 ++++++++++++++++ 11 files changed, 438 insertions(+), 32 deletions(-) diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx index 07a48ab9c..3eeff803f 100644 --- a/fern/versions/latest/pages/slurm/cli-reference.mdx +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -24,6 +24,8 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | 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. | @@ -33,10 +35,10 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel ## Output and errors -Success is one canonical JSON value on stdout. Failure is one object on stderr: +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"}} +{"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/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index d97c82a2f..9397457f7 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -156,7 +156,7 @@ Dry-run resolves the selected profile, records the active shared Python runtime, data-designer slurm execute run.yaml --cluster primary ``` -The command prints one JSON object. Save its `run_id`; all lifecycle commands use the managed run ID rather than a raw Slurm job ID. +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 --cluster primary`. ## 11. Inspect or cancel the run diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index da2e5a102..481fadb19 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -33,14 +33,14 @@ Custom OCI sources must use `repository@sha256:<64 lowercase hex characters>`. T 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 scheduler state every 30 seconds. To also see the image job's standard output and error while it runs, add `--follow-logs`: +In an interactive terminal, `image add` reports the scheduler state every 30 seconds and 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 --follow-logs + --kind serving --name vllm ``` -The log stream is opt-in, bounded to 1 MiB, and written to standard error; the final machine-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. If you need durable logs, capture standard error in your own private file. +The log stream is bounded to 1 MiB and written to standard error; 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 diff --git a/fern/versions/latest/pages/slurm/operations.mdx b/fern/versions/latest/pages/slurm/operations.mdx index 283a22faa..2dae74437 100644 --- a/fern/versions/latest/pages/slurm/operations.mdx +++ b/fern/versions/latest/pages/slurm/operations.mdx @@ -1,6 +1,13 @@ # Slurm run operations -The CLI prints canonical JSON to stdout. Errors are canonical JSON on stderr, so automation can consume the same commands as an operator. +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; do not parse them as the command result. ## Resolve without submitting @@ -18,6 +25,8 @@ 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. ## Inspect status diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index ecaf89215..41530960b 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -1,6 +1,6 @@ # Slurm troubleshooting -Start with the JSON error on stderr, then inspect the managed run directory and Slurm state. The CLI uses these exit codes: +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 | | --- | --- | --- | @@ -42,6 +42,19 @@ Set `scheduler.bin_path` in the selected profile to the absolute directory conta `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. diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index c480702ad..138fad4ef 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -5,7 +5,7 @@ import re import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from enum import Enum from functools import partial from pathlib import Path @@ -25,6 +25,7 @@ is_versioned_vllm_tag, resolve_versioned_vllm_source, ) +from data_designer.slurm.security import redact_sensitive_diagnostic from data_designer.slurm.services import ( SlurmServiceError, SlurmServiceErrorCode, @@ -50,6 +51,11 @@ class _RetryResumeMode(str, Enum): IF_POSSIBLE = "if_possible" +class _OutputFormat(str, Enum): + HUMAN = "human" + JSON = "json" + + app = typer.Typer( name="slurm", help="Run Data Designer workloads on Slurm", @@ -62,8 +68,17 @@ class _RetryResumeMode(str, Enum): @app.callback() -def slurm_callback() -> None: - pass +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 @app.command("execute") @@ -79,7 +94,7 @@ def execute_command( 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) + 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)) @@ -160,7 +175,7 @@ def retry_command( ) ) if not confirmed: - _emit_json({"operation": operation.value, "state": "declined"}) + _emit_value({"operation": operation.value, "state": "declined"}) return shard_ids = planned.shard_ids resume = _RetryResumeMode(planned.effective_resume_mode) @@ -247,7 +262,11 @@ def image_add_command( kind: str = typer.Option("serving", "--kind", help="Image role (defaults to serving)"), name: str | None = typer.Option(None, "--name"), replace: bool = typer.Option(False, "--replace"), - follow_logs: bool = typer.Option(False, "--follow-logs", help="Show bounded image-job logs on stderr"), + 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: @@ -293,7 +312,7 @@ def add() -> BaseModel: profile_file=profile_file, cluster=cluster, progress=_emit_progress, - logs=_emit_image_log if follow_logs else None, + logs=_emit_image_log if (follow_logs if follow_logs is not None else _progress_enabled()) else None, ).add(request, replace=replace) _emit_result(_invoke(operation, add)) @@ -310,7 +329,7 @@ def image_list_command( operation, lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).list(), ) - _emit_json([image.model_dump(mode="json") for image in images]) + _emit_value([image.model_dump(mode="json") for image in images]) @image_app.command("info") @@ -364,21 +383,79 @@ def _invoke(operation: SlurmServiceOperation, call: Callable[[], _ResultT]) -> _ def _fail(error: SlurmServiceError) -> NoReturn: - _emit_json( - { - "error": { - "code": error.code.value, - "message": str(error), - "operation": error.operation.value, - } - }, - err=True, - ) + 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: - _emit_json(result.model_dump(mode="json")) + _emit_value(result.model_dump(mode="json")) + + +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 _OutputFormat.HUMAN + + +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: 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 index 1e84db3b5..dd3b3f6eb 100644 --- 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 @@ -32,6 +32,7 @@ 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, @@ -134,7 +135,8 @@ def _append(self, stream: str, content: bytes, *, final: bool = False) -> None: 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: - self._emit(f"[image {stream}] {_UNSAFE_TERMINAL_CHARACTERS.sub('?', line)}") + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 2b731bed8..835cc0b11 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -7,6 +7,7 @@ import importlib.metadata import os +import shlex from collections.abc import Callable, Iterator, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass @@ -46,13 +47,14 @@ 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 +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError from data_designer.slurm.launcher.renderer import render_generation_attempt_script 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 @@ -284,6 +286,7 @@ def __init__( publisher: SlurmRunArtifactPublisher, clock: Clock, source_environment: Mapping[str, str], + progress: Callable[[str], None] | None = None, ) -> None: self._preparer = preparer self._profile = selected_profile @@ -291,6 +294,7 @@ def __init__( self._publisher = publisher self._clock = clock self._source_environment = source_environment + self._progress = progress self._retry_collection = RunRetryCollectionBackend( selected_profile.profile.workspace_root, launcher, @@ -327,12 +331,19 @@ def execute( prepared, force=False, ) + 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, @@ -394,6 +405,27 @@ def execute( 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 logs: " + + shlex.join( + ( + "tail", + "-F", + (run_root / f"{log_stem}.out").as_posix(), + (run_root / f"{log_stem}.err").as_posix(), + ) + ) + ) + if len(plan.shards) > 1: + self._report_progress("For another shard, replace _0 in the log filenames with its task index.") return SlurmRunExecution( run_id=plan.run_id, state="submitted", @@ -402,6 +434,14 @@ def execute( 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, @@ -739,6 +779,7 @@ def create_slurm_run_service( 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) @@ -751,6 +792,7 @@ def create_slurm_run_service( clock=clock, package_version=package_version, source_environment=source_environment, + progress=progress, ) @@ -764,6 +806,7 @@ def _create_slurm_run_service( 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 @@ -782,6 +825,7 @@ def _create_slurm_run_service( 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) @@ -885,6 +929,17 @@ def _new_run_id() -> str: return f"run-{uuid4().hex}" +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) diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index 776002597..0bb3965d9 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -272,7 +272,7 @@ def test_image_add_follows_bounded_job_logs_before_success_cleanup(tmp_path: Pat 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\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( @@ -287,7 +287,8 @@ def on_submit() -> None: assert logs == [ "[image stdout] importing", "[image stdout] ready", - "[image stderr] warning?[31m", + "[image stderr] warning [31m", + "[image stderr] token=", ] assert not _job_directory(workspace).exists() diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index f1a8b52c7..1df7fcd20 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -28,7 +28,7 @@ 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 +from data_designer.slurm.launcher.errors import SlurmLauncherError, SlurmSubmissionError from data_designer.slurm.launcher.models import ( SlurmAccountingEntry, SlurmJobSubmissionReceipt, @@ -522,6 +522,134 @@ def test_production_wiring_publishes_initial_state_before_releasing_submission( 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] + + +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, diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 24b5ede33..43053aba7 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -101,6 +101,77 @@ def analyze(self, 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_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 == "(none)\n" + + +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"), [ @@ -495,6 +566,54 @@ def create_service(**kwargs: object) -> ImageService: assert result.stderr == "[image stderr] inspection started\n" +@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"] From e8b65555e6c3c34a50c0c5ff8fa4a0b3bb00011e Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 7 Oct 2026 11:18:08 -0600 Subject: [PATCH 37/62] fix(slurm): stabilize job and GPU replica startup (#993) * fix(slurm): discover scheduler tools for jobs Capture the submit-host Slurm executable directory in the immutable run plan so controlled batch PATHs can invoke srun and scontrol without a profile override. Fail before submission with an actionable diagnostic when those tools are unavailable, and document the optional override. * fix(slurm): preserve scheduler command symlinks Signed-off-by: Nabin Mulepati * fix(slurm): isolate single-node serving replicas Run multiple vLLM replicas inside one GPU-owning Slurm step so overlapping steps cannot be assigned the same device. Preserve the existing local vLLM command for each process and cover TP=1 and TP=4 assignments. Signed-off-by: Nabin Mulepati * fix(slurm): bind shared serving worker to node address Signed-off-by: Nabin Mulepati * fix(slurm): reserve CPUs across GPU sockets Signed-off-by: Nabin Mulepati * fix(slurm): keep local replicas on loopback Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/getting-started.mdx | 2 +- fern/versions/latest/pages/slurm/profiles.mdx | 3 +- .../latest/pages/slurm/troubleshooting.mdx | 2 +- .../data_designer/slurm/launcher/renderer.py | 19 +++-- .../slurm/launcher/scheduler_path.py | 44 +++++++++++ .../data_designer/slurm/planning/compiler.py | 1 + .../data_designer/slurm/planning/models.py | 13 ++++ .../slurm/planning/resolution.py | 3 + .../slurm/runtime/distributed.py | 7 +- .../slurm/runtime/server_manifest.py | 4 +- .../data_designer/slurm/services/wiring.py | 6 ++ .../tests/launcher/test_collection.py | 2 + .../tests/launcher/test_renderer.py | 33 +++++++- .../tests/launcher/test_scheduler_path.py | 77 +++++++++++++++++++ .../tests/runtime/test_bootstrap.py | 68 ++++++++++++++-- .../tests/runtime/test_node_worker.py | 30 ++++++++ .../tests/services/test_wiring.py | 33 ++++++++ .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 7 +- 20 files changed, 331 insertions(+), 27 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/launcher/scheduler_path.py create mode 100644 packages/data-designer-slurm/tests/launcher/test_scheduler_path.py diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index 9397457f7..b732f25ab 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -66,7 +66,7 @@ data-designer slurm profile init \ --host-pattern 'login*.example.test' ``` -The command writes the profile with restrictive permissions and never overwrites an existing file. Before submitting work, review every generated site value. In particular, confirm the account, partitions, hostname pattern, 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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. If Slurm commands are not on the batch-job `PATH`, also set `scheduler.bin_path` to their absolute directory. The starter image-build time limit is four hours and may exceed site policy. Edit it before `image add` when necessary. +The command writes the profile with restrictive permissions and never overwrites an existing file. Before submitting work, review every generated site value. In particular, confirm the account, partitions, hostname pattern, 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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. 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 diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index b7fc903b0..c8a48a533 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -42,7 +42,6 @@ clusters: account: research partition: gpu mem_per_gpu: 80G - bin_path: /opt/slurm/bin gpus_per_node: 8 workspace_root: /workspace/alice/data-designer-slurm image_build: @@ -79,7 +78,7 @@ data-designer slurm execute run.yaml --cluster lab `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. -Set `scheduler.bin_path` when Slurm executables are installed outside the generated batch job's default `PATH`. It must be one absolute directory and is prepended to `PATH`; do not provide a colon-separated list. +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 diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index 41530960b..78ccaefca 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -36,7 +36,7 @@ Check that the profile file and workspace are accessible from the submit host, h ## Slurm commands are unavailable in a batch job -Set `scheduler.bin_path` in the selected profile to the absolute directory containing `srun`, `scontrol`, and the other Slurm executables. The generated batch script prepends this directory to its controlled `PATH`. +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 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 index 17f5128ad..d933dcea8 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -195,7 +195,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) def _get_command_path(plan: ResolvedSlurmRunPlan) -> str: - scheduler_bin_path = plan.selected_profile.profile.scheduler.bin_path + 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}" @@ -218,23 +218,32 @@ def _build_generation_directives( 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)), - ("cpus-per-task", str(plan.client.authored.cpus)), + cpu_directive, ("time", plan.submission.time_limit), ("array", resolved_array), ("chdir", posixpath.dirname(plan.authored_config.path)), ("output", output_path), ("error", error_path), ] - profile = plan.selected_profile.profile if profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) - elif profile.scheduler.mem_per_gpu is not None: - raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") 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: 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/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index c3fdb03a0..2b296c39a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -72,6 +72,7 @@ def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan 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, 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 index 9c602b3fb..617a52d4f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -400,6 +400,7 @@ class ResolvedSlurmRunPlan(ContractRecord): 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 @@ -412,9 +413,21 @@ class ResolvedSlurmRunPlan(ContractRecord): 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: 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 index e3a7ca5cb..3ecf5db90 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -85,6 +85,7 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): 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 @@ -107,6 +108,7 @@ def resolve_slurm_config( 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: @@ -127,6 +129,7 @@ def resolve_slurm_config( 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, 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 index 88236233f..1aabb4b6b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py @@ -9,6 +9,7 @@ from data_designer.slurm.runtime.node_spec import NodeProcessSpec, 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 @@ -28,7 +29,11 @@ def build_node_worker_spec( processes=tuple( NodeProcessSpec( process_id=process.process_id, - command=build_vllm_process_command(deployment, process, plan, layout), + command=( + build_vllm_command(deployment, process, plan) + if len(layout.node_hosts) == 1 + else build_vllm_process_command(deployment, process, plan, layout) + ), gpu_indices=tuple(process.gpu_indices), launch_delay_seconds=process.launch_delay_seconds, ) 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 index 3968a5dbb..43cc0c3d2 100644 --- 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 @@ -77,7 +77,7 @@ def build_distributed_server_step( readiness = ( tuple( RuntimeProbeSpec( - host=layout.get_host(probe.node_index), + 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, @@ -121,7 +121,7 @@ def _build_deployment_serving_steps( log_directory: Path, layout: AllocationLayout, ) -> tuple[RuntimeStepSpec, ...]: - if len(deployment.node_indices) > 1: + if len(deployment.node_indices) > 1 or len(deployment.processes) > 1: return ( build_distributed_server_step( deployment, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 835cc0b11..e34e70060 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -49,6 +49,7 @@ 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 @@ -187,6 +188,10 @@ def _prepare( ) -> _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() @@ -213,6 +218,7 @@ def _prepare( 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, ) diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index fc5fe7041..9f034ee54 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -119,6 +119,8 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( 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 diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 465480439..10111c99c 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -99,18 +99,43 @@ def test_renderer_uses_profile_slurm_bin_path(single_node_plan: ResolvedSlurmRun assert "#SBATCH --exclusive" not in script -@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) -def test_renderer_reserves_client_cpus_for_each_gpu_request_mode( +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": 17})} + 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}) - assert "#SBATCH --cpus-per-task=17\n" in render_generation_attempt_script(plan, attempt_ordinal=1) + 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( 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/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 593e1a75f..98187c753 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -10,7 +10,7 @@ from conftest import RuntimeCase, relocate_plan from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.planning import ResolvedSlurmRunPlan +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 @@ -57,6 +57,8 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru 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 == () @@ -70,6 +72,54 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert "SLURM_TMPDIR" not in manifest.serialize_json() +@pytest.mark.parametrize(("tensor_parallel", "replicas"), ((1, 8), (2, 4), (4, 2))) +def test_one_node_multi_replica_manifest_uses_one_gpu_owned_worker( + runtime_case: RuntimeCase, tensor_parallel: int, replicas: 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})} + ) + topology = ResolvedTopology.derive( + node_count=1, 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=0, port=18000 + index) + for index in range(replicas) + ) + placement = placement.model_copy(update={"authored": authored, "topology": topology, "ports": ports}) + plan = ResolvedSlurmRunPlan.model_validate(plan.model_copy(update={"deployments": (placement,)}).model_dump()) + context = replace(runtime_case.context, plan=plan) + + 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=AllocationLayout(("compute-001",)), + ) + + 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"} + worker = decode_node_worker_spec(servers[0].command[-1]) + processes = worker.nodes[0].processes + assert len(processes) == replicas + assert [process.gpu_indices for process in processes] == [ + tuple(range(index * tensor_parallel, (index + 1) * tensor_parallel)) for index in range(replicas) + ] + assert all( + ("--host", "127.0.0.1") == 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 "compute-001" not in endpoint.command + assert " ".join(endpoint.command).count("http://127.0.0.1:") == replicas + + @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 @@ -174,7 +224,7 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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-replica-00000-rank-00000") + 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") @@ -195,10 +245,16 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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.command[2] == "/workspace/primary/models/model-1" - executor_index = remote_server.command.index("--distributed-executor-backend") - assert remote_server.command[executor_index + 1] == "uni" - assert tuple(probe.host for probe in remote_server.readiness) == ("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 def test_pipeline_parallel_process_uses_multi_process_executor( diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py index a9511f1d8..3809ef195 100644 --- a/packages/data-designer-slurm/tests/runtime/test_node_worker.py +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -101,6 +101,36 @@ def terminate(child: _FakeProcess, selected: int) -> None: assert supervisor.cleanup_complete +def test_node_worker_assigns_distinct_visible_gpus_to_two_tp4_replicas(monkeypatch: pytest.MonkeyPatch) -> None: + visible: list[str] = [] + + def start(*arguments: object, **keywords: object) -> _FakeProcess: + del arguments + environment = keywords["env"] + assert isinstance(environment, dict) + visible.append(environment["CUDA_VISIBLE_DEVICES"]) + 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"] + + def test_follower_failure_terminates_sibling_without_orphan(tmp_path: Path) -> None: ready = tmp_path / "ready" stopped = tmp_path / "stopped" diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 1df7fcd20..01213c32e 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -12,6 +12,7 @@ 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 @@ -36,6 +37,7 @@ 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, @@ -66,6 +68,7 @@ def _use_persisted_client_runtime( 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: @@ -237,6 +240,36 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( 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 == [] diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index e08096373..33a4cef2b 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -3,7 +3,7 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=3 -#SBATCH --cpus-per-task=32 +#SBATCH --cpus-per-gpu=4 #SBATCH --time=03:55:00 #SBATCH --array=0-1%2 #SBATCH --chdir="/workspace/primary/runs/run-001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 544d56739..12b7f1b89 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -3,7 +3,7 @@ #SBATCH --account=research #SBATCH --partition=batch #SBATCH --nodes=1 -#SBATCH --cpus-per-task=32 +#SBATCH --cpus-per-gpu=4 #SBATCH --time=03:55:00 #SBATCH --array=0 #SBATCH --chdir="/workspace/primary/runs/run-single" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 024dcbf36..d5d52ec38 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="f130dd88efe97ebb060eeef743a62f1f1d27e8f16f0f71d6ec31791b4ccee6b3", + expected_fixture_sha256="e50d3a25a6f1aa05bd4b14818c02809007347a49073c0af8fd453dc1f9b8183a", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="eefbd439a72c59dc09a97f5ebc73f9fdb5ee13a2f41f1447ab560e8e2d589cca", + expected_fixture_sha256="5d06a7a7bd3bb8be9a62fdb251b8956368d27c9d2d9c4ab7d2af418c9bff889a", ) @@ -64,7 +64,8 @@ def _assert_script_matches_plan( assert f"#SBATCH --account={plan.submission.account}\n" in script assert f"#SBATCH --partition={plan.submission.partition}\n" in script assert f"#SBATCH --nodes={node_count}\n" in script - assert f"#SBATCH --cpus-per-task={plan.client.authored.cpus}\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 From d967bf73d522edb553496a1318048974ee39d401 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 7 Oct 2026 12:47:50 -0600 Subject: [PATCH 38/62] feat(slurm): support multiple GPU partitions (#998) * feat(slurm): support multiple GPU partitions Signed-off-by: Nabin Mulepati * fix(slurm): validate every selected GPU partition Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/profiles.mdx | 13 +++++++ .../latest/pages/slurm/run-configuration.mdx | 9 +++++ .../src/data_designer/slurm/cli.py | 2 +- .../data_designer/slurm/config/profiles.py | 3 +- .../src/data_designer/slurm/config/run.py | 4 +-- .../src/data_designer/slurm/contracts.py | 2 ++ .../data_designer/slurm/launcher/client.py | 25 +++++++++---- .../data_designer/slurm/launcher/parsing.py | 23 ++++++++++++ .../data_designer/slurm/planning/models.py | 4 +-- .../src/data_designer/slurm/types.py | 24 ++++++++++++- .../tests/config/test_partition_selection.py | 36 +++++++++++++++++++ .../tests/launcher/test_client.py | 31 ++++++++++++++-- .../tests/launcher/test_renderer.py | 12 +++++++ .../tests/services/test_profile_service.py | 24 +++++++++++++ 14 files changed, 195 insertions(+), 17 deletions(-) create mode 100644 packages/data-designer-slurm/tests/config/test_partition_selection.py diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index c8a48a533..ea124c09b 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -76,8 +76,21 @@ 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. +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 diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index d7201563e..6e938f12f 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -46,6 +46,15 @@ Every `model_alias` in `model_concurrency` and the builder must have one deploym 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 32 and controls the client's CPU allocation, not its execution image. +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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 138fad4ef..94cc15b0e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -223,7 +223,7 @@ def profile_init_command( 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"), + 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.""" 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 index e1530eec8..4c16b4610 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/profiles.py @@ -19,6 +19,7 @@ validate_absolute_path, validate_plain_text, ) +from data_designer.slurm.types import PartitionSelection _ALLOCATION_SCRATCH_CONTAINER_ROOT = "/run/data-designer-slurm" @@ -30,7 +31,7 @@ class GpuRequestMode(str, Enum): class SchedulerProfile(AuthoredConfig): account: Identifier | None = None - partition: 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 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 index 9fb2cc4f8..a9d10855d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -38,7 +38,7 @@ validate_plain_text, validate_url, ) -from data_designer.slurm.types import EnvironmentName, Identifier, SchemaVersion +from data_designer.slurm.types import EnvironmentName, Identifier, PartitionSelection, SchemaVersion __all__ = [ "ArrayTasksConfig", @@ -315,7 +315,7 @@ def validate_concurrency(self) -> ArrayTasksConfig: class SubmissionConfig(AuthoredConfig): account: Identifier | None = None - partition: 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py index e613c067d..4e511d6ca 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/contracts.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/contracts.py @@ -30,6 +30,7 @@ Identifier, NetworkPort, NonNegativeDuration, + PartitionSelection, SchemaVersion, Sha256Digest, ShardId, @@ -325,6 +326,7 @@ class ResumeWorkspace(ContractValue): "ModelAlias", "NetworkPort", "NonNegativeDuration", + "PartitionSelection", "RecordRange", "ResumeWorkspace", "SchemaVersion", 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 index e9c5fbf06..e5f3fa6ad 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -13,7 +13,7 @@ from datetime import datetime, timedelta from pathlib import Path -from data_designer.slurm.contracts import Identifier +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, @@ -28,6 +28,7 @@ parse_gpu_counts, parse_named_jobs, parse_queue, + parse_selected_partition_gpu_counts, parse_submission, ) from data_designer.slurm.launcher.runner import CommandRunner, SubprocessRunner @@ -35,6 +36,7 @@ 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( @@ -224,15 +226,24 @@ 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: Identifier | None = None) -> tuple[int, ...]: - """Return configured GPU counts for the requested or default partition.""" + 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 _IDENTIFIER_PATTERN.fullmatch(partition) is None: - raise ValueError("Slurm partition must be a valid identifier") - command = (self._executables.sinfo, "--noheader", "--format=%G", f"--partition={partition}") - return parse_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, 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 index 56e4f97a6..41a9f865b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/parsing.py @@ -146,6 +146,29 @@ def parse_gpu_counts(output: str) -> tuple[int, ...]: 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 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 index 617a52d4f..bc26bb085 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -47,7 +47,7 @@ validate_plain_text, ) from data_designer.slurm.planning.builder_identity import get_declared_model_aliases -from data_designer.slurm.types import NetworkPort +from data_designer.slurm.types import NetworkPort, PartitionSelection class ResolvedImage(ContractValue): @@ -373,7 +373,7 @@ def requested_records(self) -> int: class ResolvedSubmission(ContractValue): account: Identifier | None = None - partition: Identifier | None = None + partition: PartitionSelection | None = None job_name: Identifier time_limit: str comment: str | None = None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/types.py b/packages/data-designer-slurm/src/data_designer/slurm/types.py index aba456062..2eafa605f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/types.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/types.py @@ -7,7 +7,7 @@ from typing import Annotated, Literal -from pydantic import Field, StringConstraints +from pydantic import BeforeValidator, Field, StringConstraints Identifier = Annotated[ str, @@ -17,6 +17,27 @@ 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] @@ -33,6 +54,7 @@ "Identifier", "NetworkPort", "NonNegativeDuration", + "PartitionSelection", "SchemaVersion", "Sha256Digest", "ShardId", 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/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index a9c247eff..7beb3225e 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -251,11 +251,36 @@ def test_client_queries_partition_scoped_gpu_inventory() -> None: assert runner.calls == [command] -def test_client_rejects_invalid_gpu_partition_without_running_command(fake_slurm_runner: FakeSlurmRunner) -> None: +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="valid identifier"): - client.query_gpu_counts(partition="batch,other") + with pytest.raises(ValueError, match="comma-separated list"): + client.query_gpu_counts(partition=selection) assert fake_slurm_runner.calls == [] diff --git a/packages/data-designer-slurm/tests/launcher/test_renderer.py b/packages/data-designer-slurm/tests/launcher/test_renderer.py index 10111c99c..91757b78a 100644 --- a/packages/data-designer-slurm/tests/launcher/test_renderer.py +++ b/packages/data-designer-slurm/tests/launcher/test_renderer.py @@ -79,6 +79,18 @@ def test_renderer_omits_gres_for_visible_mode_and_emits_optional_submission_fiel 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")} diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py index 1440ca609..b543f6645 100644 --- a/packages/data-designer-slurm/tests/services/test_profile_service.py +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -91,6 +91,30 @@ def test_profile_init_creates_deterministic_private_starter_without_side_effects assert "container_mounts" not in profile_file.read_text() +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" + + result = service.validate() + + assert result.gpus_per_node == 8 + 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" From 8f78c4c04e1a11d87d9912c35aa37d7caa0f5327 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 7 Oct 2026 15:51:05 -0600 Subject: [PATCH 39/62] fix(slurm): create profile parent directories (#999) * fix(slurm): create profile parent directories Signed-off-by: Nabin Mulepati * fix(slurm): verify new profile directories Keep profile creation bound to verified directory descriptors and restore private permissions when a restrictive umask is active. Update the profile guide to reflect automatic parent creation. Signed-off-by: Nabin Mulepati * fix(slurm): harden profile directory setup Normalize unsupported permission operations and sync newly created directory metadata before publishing the profile catalog. Cover replacement and error-boundary behavior. Signed-off-by: Nabin Mulepati * fix(slurm): support Linux directory chmod Use an opened directory for permission repair when Linux rejects nofollow fchmodat. Handle a restrictive umask through a Linux path descriptor while keeping directory traversal bound to verified descriptors. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/profiles.mdx | 3 +- .../data_designer/slurm/services/profiles.py | 81 ++++++++++++++++- .../tests/services/test_profile_service.py | 90 ++++++++++++++++++- .../data-designer-slurm/tests/test_cli.py | 3 +- 4 files changed, 169 insertions(+), 8 deletions(-) diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index ea124c09b..1aabf8dbd 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -15,7 +15,6 @@ Pass `--cluster NAME` to select an entry explicitly. Without it, hostname patter ## Initialize a catalog ```bash -mkdir -p /workspace/alice/config data-designer slurm profile init \ --profile-file /workspace/alice/config/slurm-profiles.yml \ --workspace-root /workspace/alice/data-designer-slurm \ @@ -26,7 +25,7 @@ data-designer slurm profile init \ --host-pattern 'login*.example.test' ``` -The starter sets `gpus_per_node: auto`, `gpu_request_mode: gres`, and CPU image-build resources of 2 CPUs, 8 GiB, and four hours. The parent directory of a custom `--profile-file` must already exist. Review the generated account, partitions, host patterns, workspace, GPU settings, and image-build resources against site policy before validation or image registration. +The starter sets `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, host patterns, workspace, GPU settings, and image-build resources against site policy before validation or image registration. ## Configure multiple clusters 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 index 550af5562..652330d5c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -9,7 +9,9 @@ import os import shlex import socket -from collections.abc import Callable, Mapping +import stat +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack, contextmanager from fnmatch import fnmatchcase from pathlib import Path @@ -27,7 +29,12 @@ resolve_profile, ) from data_designer.slurm.contracts import ContractValue, Identifier, compute_canonical_json_sha256 -from data_designer.slurm.filesystem import create_restrictive_temporary_file, open_verified_directory +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 @@ -283,7 +290,7 @@ def _serialize_catalog(payload: dict[str, object], *, suffix: str) -> bytes: def _create_profile_file(path: Path, content: bytes) -> None: - with open_verified_directory(path.parent, resource_name="profile") as parent_descriptor: + with _open_or_create_profile_parent(path.parent) as parent_descriptor: descriptor, temporary_name = create_restrictive_temporary_file( parent_descriptor, prefix=f".{path.name}.", @@ -313,6 +320,74 @@ def _create_profile_file(path: Path, content: bytes) -> None: 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( diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py index b543f6645..f04e43ddc 100644 --- a/packages/data-designer-slurm/tests/services/test_profile_service.py +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import os import stat from pathlib import Path @@ -91,6 +92,91 @@ def test_profile_init_creates_deterministic_private_starter_without_side_effects 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" @@ -191,8 +277,8 @@ def test_profile_init_refuses_to_follow_dangling_destination_symlink(tmp_path: P @pytest.mark.parametrize( "profile_file", - (Path("profile.txt"), Path("missing/profile.yml")), - ids=("unsupported-suffix", "missing-parent"), + (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: diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 43053aba7..d55921b41 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -624,7 +624,7 @@ def test_image_add_rejects_versioned_vllm_as_client() -> None: def test_profile_init_creates_starter_and_emits_validation_command(tmp_path: Path) -> None: - profile_file = tmp_path / "profile.yml" + profile_file = tmp_path / "config" / "profile.yml" workspace = tmp_path / "workspace" result = CliRunner().invoke( @@ -649,6 +649,7 @@ def test_profile_init_creates_starter_and_emits_validation_command(tmp_path: Pat "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() From bd72a69ca465f0c6fb303d4d350d3301c4146a35 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 7 Oct 2026 20:52:32 -0600 Subject: [PATCH 40/62] fix(slurm): expose vLLM logs and honor JSON output (#1000) * fix(slurm): show vLLM startup log path Signed-off-by: Nabin Mulepati * fix(slurm): honor JSON output with newer Typer Signed-off-by: Nabin Mulepati * fix(slurm): refresh serving log guidance across executions Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/cli.py | 8 ++++- .../data_designer/slurm/services/wiring.py | 17 ++++++++-- .../tests/services/test_wiring.py | 31 +++++++++++++++++++ .../data-designer-slurm/tests/test_cli.py | 10 ++++++ 4 files changed, 63 insertions(+), 3 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 94cc15b0e..719df3400 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -6,6 +6,7 @@ import re import sys from collections.abc import Callable, Mapping +from contextvars import ContextVar from enum import Enum from functools import partial from pathlib import Path @@ -56,6 +57,9 @@ class _OutputFormat(str, Enum): 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", @@ -79,6 +83,8 @@ def slurm_callback( ) -> 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") @@ -417,7 +423,7 @@ def _output_format() -> _OutputFormat: if isinstance(context.obj, _OutputFormat): return context.obj context = context.parent - return _OutputFormat.HUMAN + return _ACTIVE_OUTPUT_FORMAT.get() def _human_lines(value: object, *, indent: int = 0) -> list[str]: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index e34e70060..f063f2325 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -420,7 +420,7 @@ def execute( 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 logs: " + "After allocation, follow shard 0 Slurm logs: " + shlex.join( ( "tail", @@ -430,8 +430,21 @@ def execute( ) ) ) + 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 log filenames with its task index.") + 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", diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 01213c32e..5bd7ce4d0 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -579,6 +579,37 @@ def test_execute_reports_safe_submission_and_log_guidance( 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( diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index d55921b41..7c8947361 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -157,6 +157,16 @@ def list(self) -> tuple[()]: assert result.stdout == "(none)\n" +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( From ba695e922fd5630c9784f33cfd73e4f0c1683c58 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:30:26 -0300 Subject: [PATCH 41/62] fix(slurm): address retry, startup, and proxy failure paths (#1005) * fix(slurm): align native retry and user-site policy Signed-off-by: Andre Manoel * fix(slurm): preserve delayed replica startup budget Signed-off-by: Andre Manoel * fix(slurm): cancel proxy requests and catch timeouts Signed-off-by: Andre Manoel * fix(slurm): isolate user-site path inspection Signed-off-by: Andre Manoel * test(slurm): guard proxy and readiness compatibility Signed-off-by: Andre Manoel --------- Signed-off-by: Andre Manoel --- .github/workflows/ci.yml | 17 +- .../data_designer/slurm/client/environment.py | 47 +++-- .../src/data_designer/slurm/client/runtime.py | 3 +- .../data_designer/slurm/launcher/renderer.py | 2 +- .../src/data_designer/slurm/runtime/proxy.py | 7 +- .../slurm/runtime/server_manifest.py | 7 +- .../tests/client/test_environment.py | 7 +- .../tests/client/test_runtime.py | 183 ++++++++++++++++++ .../tests/launcher/test_collection.py | 22 ++- .../tests/runtime/test_bootstrap.py | 105 ++++++++-- .../tests/runtime/test_proxy.py | 144 +++++++++++++- 11 files changed, 502 insertions(+), 42 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b36fb1411..350b25a8b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -130,9 +130,19 @@ jobs: pytest packages/data-designer/tests test-slurm-package: - name: Test Slurm package wheels + 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 @@ -142,16 +152,17 @@ jobs: uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: version: "latest" - python-version: "3.11" + python-version: ${{ matrix.python-version }} enable-cache: true - name: Install development dependencies run: make install-dev - name: Run package tests - run: .venv/bin/pytest packages/data-designer-slurm/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 # =========================================================================== 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 index 5144ac409..667d5e3c7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -303,28 +303,43 @@ def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ... def _active_distribution_search_paths() -> list[str]: - """Exclude submit-only ``PYTHONPATH`` entries from the client runtime inventory.""" + """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 } - if not configured_paths: - return sys.path.copy() - - site_paths = site.getsitepackages() - if site.ENABLE_USER_SITE: - user_site_path = site.getusersitepackages() - if isinstance(user_site_path, str): - site_paths.append(user_site_path) - else: - site_paths.extend(user_site_path) - standard_site_paths = {_normalize_import_path(entry) for entry in site_paths} - search_paths: list[str] = [] - for entry in sys.path: + 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 - search_paths.append(entry) - return search_paths + active_paths.append(entry) + return active_paths def _normalize_import_path(path: str) -> str: 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 index bfcd6774d..b6c4f6a63 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/runtime.py @@ -56,7 +56,8 @@ def inspect(self) -> ResolvedClientRuntime: missing = _REQUIRED_DISTRIBUTIONS.difference(names) if missing: raise ClientRuntimeInspectionError( - f"required client distributions are not installed: {', '.join(sorted(missing))}" + 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 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 index d933dcea8..9f3390c05 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -123,7 +123,7 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) 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)} -readonly DD_CLIENT_PYTHON={quote_shell_value(plan.client.runtime.python_executable)} +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)} 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 index 82de94e16..69003b0f5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/proxy.py @@ -222,7 +222,7 @@ async def _probe_backend(self, index: int, backend: _Backend) -> bool: ) as response: await response.read() return response.status == 200 - except (ClientError, TimeoutError): + except (ClientError, asyncio.TimeoutError): return False finally: self.pool.release(index) @@ -251,9 +251,9 @@ async def _forward(self, request: web.Request) -> web.StreamResponse: self.metrics.backend_requests[index] += 1 try: response = await self._request_backend(request, backend, body, index) - except (ClientError, TimeoutError) as error: + except (ClientError, asyncio.TimeoutError) as error: self.metrics.connection_failed += 1 - reason = "timeout" if isinstance(error, TimeoutError) else "connection_error" + reason = "timeout" if isinstance(error, asyncio.TimeoutError) else "connection_error" outcomes.append(reason) self.pool.release(index) excluded.add(index) @@ -421,6 +421,7 @@ def main(arguments: Sequence[str] | None = None) -> int: access_log=None, keepalive_timeout=_KEEPALIVE_TIMEOUT_SECONDS, shutdown_timeout=10.0, + handler_cancellation=True, ) return 0 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 index 43cc0c3d2..34973f1f3 100644 --- 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 @@ -80,7 +80,12 @@ def build_distributed_server_step( 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, + deadline_seconds=probe.deadline_seconds + + next( + process.launch_delay_seconds + for process in deployment.processes + if process.node_index == probe.node_index and process.http_port == probe.port + ), ) for probe in deployment.readiness_probes ) diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 167050459..920601dad 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -196,11 +196,16 @@ def distributions(**kwargs: object) -> tuple[()]: 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, *([user_site_path] if user_site_enabled else []), pth_path] + expected_paths = [interpreter_path, pth_path] assert calls == [{"path": expected_paths}] diff --git a/packages/data-designer-slurm/tests/client/test_runtime.py b/packages/data-designer-slurm/tests/client/test_runtime.py index 64827e4cc..787fc5adf 100644 --- a/packages/data-designer-slurm/tests/client/test_runtime.py +++ b/packages/data-designer-slurm/tests/client/test_runtime.py @@ -3,16 +3,199 @@ 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( diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 9f034ee54..f770275f4 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import shlex import subprocess from datetime import datetime, timezone from pathlib import Path @@ -95,6 +96,7 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( multi_node_plan: ResolvedSlurmRunPlan, + tmp_path: Path, ) -> None: retry = RetryPlan( schema_version=1, @@ -134,7 +136,25 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( assert "--container-image=" not in script assert "--container-mounts=" not in script assert "--container-env=" not in script - assert f'readonly DD_CLIENT_PYTHON="{multi_node_plan.client.runtime.python_executable}"' 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 diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 98187c753..853ec076d 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -3,6 +3,8 @@ from __future__ import annotations +import shlex +import subprocess from dataclasses import replace from pathlib import Path @@ -72,52 +74,80 @@ def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(ru assert "SLURM_TMPDIR" not in manifest.serialize_json() -@pytest.mark.parametrize(("tensor_parallel", "replicas"), ((1, 8), (2, 4), (4, 2))) -def test_one_node_multi_replica_manifest_uses_one_gpu_owned_worker( - runtime_case: RuntimeCase, tensor_parallel: int, replicas: int +@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})} + 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": "2s", + "distributed_init_timeout": "1s", + "lead_boot_standoff": "4s", + "rank_launch_stagger": f"{stagger}s", + } + ), + } ) topology = ResolvedTopology.derive( - node_count=1, gpus_per_node=plan.resolved_gpus_per_node, tensor_parallel=tensor_parallel, nodes_per_replica=1 + 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=0, port=18000 + index) + 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}) + 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=AllocationLayout(("compute-001",)), + 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"} + 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 = worker.nodes[0].processes + 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 * tensor_parallel, (index + 1) * tensor_parallel)) for index in range(replicas) + 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") == process.command[process.command.index("--host") :][:2] for process in processes + ("--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 "compute-001" not in endpoint.command - assert " ".join(endpoint.command).count("http://127.0.0.1:") == replicas + assert " ".join(endpoint.command).count("http://") == replicas + _assert_launch_readiness(manifest, runtime_case.workspace) @pytest.mark.parametrize("role", (RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.SERVER)) @@ -202,6 +232,14 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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": "2s", + "distributed_init_timeout": "1s", + "lead_boot_standoff": "4s", + "rank_launch_stagger": "2s", + } + ), } ), "model": f"/workspace/primary/models/model-{deployment_index}", @@ -255,6 +293,45 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( 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 + 2)) + cases.append(f"*://{probe.host}:{probe.port}/*) ((SECONDS >= {process.launch_delay_seconds + 1}));;") + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + command = f""" +source {shlex.quote(str(entrypoint))} +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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + assert completed.returncode == 0, completed.stderr + assert all(actual == expected for actual, expected in budgets) def test_pipeline_parallel_process_uses_multi_process_executor( diff --git a/packages/data-designer-slurm/tests/runtime/test_proxy.py b/packages/data-designer-slurm/tests/runtime/test_proxy.py index f3364d78f..a8afe7af5 100644 --- a/packages/data-designer-slurm/tests/runtime/test_proxy.py +++ b/packages/data-designer-slurm/tests/runtime/test_proxy.py @@ -10,7 +10,7 @@ from contextlib import asynccontextmanager import pytest -from aiohttp import ClientSession, web +from aiohttp import ClientSession, ClientTimeout, TCPConnector, web from aiohttp.test_utils import TestServer from data_designer.slurm.runtime import proxy as runtime_proxy @@ -41,6 +41,148 @@ def _backend(server: TestServer) -> _Backend: 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] = [] From b724d8615cdaf8c455d016957780460629fad56f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 12 Oct 2026 11:41:04 -0600 Subject: [PATCH 42/62] refactor(slurm): remove unused config flags and set client CPU default (#1012) * fix(slurm): clarify config defaults and inert fields Signed-off-by: Nabin Mulepati * docs(slurm): keep run config guidance concise Signed-off-by: Nabin Mulepati * refactor(slurm): remove unused run config flags Signed-off-by: Nabin Mulepati * docs(slurm): clarify benchmark model concurrency sweep Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/benchmarks.mdx | 5 +++-- .../latest/pages/slurm/run-configuration.mdx | 7 +++---- .../latest/pages/slurm/troubleshooting.mdx | 2 +- .../src/data_designer/slurm/config/__init__.py | 2 -- .../src/data_designer/slurm/config/builder.py | 5 +---- .../src/data_designer/slurm/config/run.py | 9 +-------- .../src/data_designer/slurm/planning/models.py | 1 - .../src/data_designer/slurm/planning/resolution.py | 1 - .../src/data_designer/slurm/planning/validation.py | 4 ---- .../tests/config/test_loading_builder.py | 6 ++++++ .../tests/contracts/golden/authored_run.json | 4 ---- .../contracts/golden/authored_run_single.json | 4 ---- .../tests/contracts/golden/multi_node_plan.json | 6 +----- .../tests/contracts/golden/single_node_plan.json | 6 +----- .../integration/golden/finalization_chain.json | 14 +++++++------- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 ++-- .../tests/state/test_state_golden_records.py | 2 +- 19 files changed, 29 insertions(+), 57 deletions(-) diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx index 059f46e6a..842ffbf73 100644 --- a/fern/versions/latest/pages/slurm/benchmarks.mdx +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -18,8 +18,7 @@ Create `benchmark.yaml` next to an existing `run.yaml`: schema_version: 1 name: generator-scaling base_run: run.yaml -model_aliases: - - generator +model_aliases: all concurrency_values: - 32 - 64 @@ -44,6 +43,8 @@ analysis: 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`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. `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)`, so it also depends on the selected profile and the base run's deployment topology. diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 6e938f12f..e08c48fba 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -15,7 +15,7 @@ invocation: model_concurrency: generator: 64 client: - cpus: 32 + cpus: 16 deployments: - model_alias: generator served_model_name: example/generator @@ -37,14 +37,13 @@ submission: output: format: parquet partitions: 1 - require_exact_record_count: false ``` `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, while `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. Output formats are `parquet`, `jsonl`, and `csv`. -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 32 and controls the client's CPU allocation, not its execution image. +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: @@ -180,7 +179,7 @@ slurm = ( dataset_name="generated", model_concurrency={"generator": 16}, ) - .with_client(cpus=32) + .with_client(cpus=16) .with_deployment( { "model_alias": "generator", diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index 78ccaefca..99d30e010 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -91,7 +91,7 @@ Do not edit persisted manifests. Conflicts preserve scheduler and provenance evi ## 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. Retry missing shards, then use `merge` with the original run directory. When `require_exact_record_count` is true, a short result remains a failure rather than being silently accepted. +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 short result cannot become a shard winner. Retry missing shards, then use `merge` with the original run directory. ## Scratch files are gone diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 3fca0ea0f..05f258600 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -62,7 +62,6 @@ DeploymentTopology, InputBindings, InvocationConfig, - InvocationDiagnostics, LocalStdioMCPProviderConfig, OutputConfig, RemoteMCPProviderConfig, @@ -99,7 +98,6 @@ "InputBindings", "InstalledDistribution", "InvocationConfig", - "InvocationDiagnostics", "LiteralEnvironmentBinding", "LocalStdioMCPProviderConfig", "OutputConfig", 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 index ccdecd384..b0f3a42a1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -23,7 +23,6 @@ DataDesignerSlurmConfig, InputBindings, InvocationConfig, - InvocationDiagnostics, MCPProviderConfig, OutputConfig, ServerDeploymentConfig, @@ -76,7 +75,6 @@ def with_invocation( input_bindings: InputBindings | Mapping[str, object] | None = None, mcp_providers: Sequence[MCPProviderConfig | Mapping[str, object]] = (), model_concurrency: Mapping[str, int] | None = None, - diagnostics: InvocationDiagnostics | Mapping[str, object] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set typed Data Designer invocation intent.""" self._invocation = _validate_model( @@ -89,7 +87,6 @@ def with_invocation( "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), - "diagnostics": {} if diagnostics is None else diagnostics, }, ) return self @@ -97,7 +94,7 @@ def with_invocation( def with_client( self, *, - cpus: int = 32, + cpus: int = 16, dependencies: ClientDependencies | Mapping[str, object] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set the separate zero-GPU Data Designer client declaration.""" 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 index a9d10855d..b3fe0a4fb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -51,7 +51,6 @@ "EnvironmentBinding", "InputBindings", "InvocationConfig", - "InvocationDiagnostics", "LiteralEnvironmentBinding", "LocalStdioMCPProviderConfig", "MCPProviderConfig", @@ -163,10 +162,6 @@ def validate_environment( ] -class InvocationDiagnostics(AuthoredConfig): - log_requests: bool = False - - class InvocationConfig(AuthoredConfig): num_records: PositiveInt dataset_name: Identifier @@ -175,7 +170,6 @@ class InvocationConfig(AuthoredConfig): input_bindings: InputBindings = Field(default_factory=InputBindings) mcp_providers: list[MCPProviderConfig] = Field(default_factory=list) model_concurrency: dict[str, PositiveInt] = Field(default_factory=dict) - diagnostics: InvocationDiagnostics = Field(default_factory=InvocationDiagnostics) @field_validator("run_config") @classmethod @@ -257,7 +251,7 @@ def validate_source(self) -> ClientDependencies: class ClientConfig(AuthoredConfig): - cpus: PositiveInt = 32 + cpus: PositiveInt = 16 dependencies: ClientDependencies = Field(default_factory=ClientDependencies) @@ -339,7 +333,6 @@ class OutputConfig(AuthoredConfig): root: str | None = None format: Literal["parquet", "jsonl", "csv"] = "parquet" partitions: PositiveInt = 1 - require_exact_record_count: bool = False @field_validator("root") @classmethod 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 index bc26bb085..b7c5b6377 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -388,7 +388,6 @@ class ResolvedOutput(ContractValue): root: str format: Literal["parquet", "jsonl", "csv"] partitions: PositiveInt - require_exact_record_count: bool _root_is_absolute = field_validator("root")(validate_absolute_path) 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 index 3ecf5db90..c00b1760e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -289,7 +289,6 @@ def _materialize_output(authored: DataDesignerSlurmConfig, run_root: str) -> Res root=authored.output.root or posixpath.join(run_root, "output"), format=authored.output.format, partitions=authored.output.partitions, - require_exact_record_count=authored.output.require_exact_record_count, ) 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 index dc0fa3245..71d761be6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -94,10 +94,6 @@ def validate_resolved_plan( _require( plan.output.partitions == authored.output.partitions, "resolved output partitions do not match authored input" ) - _require( - plan.output.require_exact_record_count == authored.output.require_exact_record_count, - "resolved exact-record policy does not match authored input", - ) if authored.output.root is not None: _require(plan.output.root == authored.output.root, "resolved output root does not match authored input") diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 4238b402e..083188628 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -13,6 +13,7 @@ from data_designer.slurm.config import ( DEFAULT_PROFILE_FILE_NAME, PROFILE_FILE_ENVIRONMENT, + ClientConfig, DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, DataDesignerSlurmConfigBuilder, @@ -61,9 +62,14 @@ def test_builder_builds_without_file_discovery_or_serialization(tmp_path: Path) 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") diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index a0e1a6e70..5643f2099 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -112,9 +112,6 @@ ], "invocation": { "dataset_name": "training-data", - "diagnostics": { - "log_requests": false - }, "input_bindings": { "managed_assets_path": null, "seed_path": null @@ -134,7 +131,6 @@ "output": { "format": "parquet", "partitions": 1, - "require_exact_record_count": false, "root": null }, "schema_version": 1, 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 index f1e81db8d..b2e807822 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -61,9 +61,6 @@ ], "invocation": { "dataset_name": "single-node", - "diagnostics": { - "log_requests": false - }, "input_bindings": { "managed_assets_path": null, "seed_path": null @@ -80,7 +77,6 @@ "output": { "format": "parquet", "partitions": 1, - "require_exact_record_count": false, "root": null }, "schema_version": 1, 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 index 5595a228e..7371ee1c9 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-001/authored-config.json", - "sha256": "fac5dec4c06558f13e761426bc52699a06017d5ba34514af422eda040fe43137" + "sha256": "c172186cbb7cedbb783c743faad2f7af8e6018bf6600a7040b608b5fe9967dbe" }, "builder": { "authored_source": null, @@ -322,9 +322,6 @@ "invocation": { "authored": { "dataset_name": "training-data", - "diagnostics": { - "log_requests": false - }, "input_bindings": { "managed_assets_path": null, "seed_path": null @@ -367,7 +364,6 @@ "output": { "format": "parquet", "partitions": 1, - "require_exact_record_count": false, "root": "/workspace/primary/runs/run-001/output" }, "package_version": "0.9.2", 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 index 03893958b..42b491a5a 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "7eb39d9f28b7c073bc520c6596bdf7f6e2afc5e7fa7881945c5bcb34b19f9920" + "sha256": "902aee7cc4bd9219ff3eed831f7203bc3722fd0ed9cbc0d01135d5c2b795668b" }, "builder": { "authored_source": null, @@ -171,9 +171,6 @@ "invocation": { "authored": { "dataset_name": "single-node", - "diagnostics": { - "log_requests": false - }, "input_bindings": { "managed_assets_path": null, "seed_path": null @@ -213,7 +210,6 @@ "output": { "format": "parquet", "partitions": 1, - "require_exact_record_count": false, "root": "/workspace/primary/runs/run-single/output" }, "package_version": "0.9.2", diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index f34ed1d52..32c0e365e 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" + "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" + "sha256": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36", + "provenance_digest": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" + "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "7eb39d9f28b7c073bc520c6596bdf7f6e2afc5e7fa7881945c5bcb34b19f9920" + "sha256": "902aee7cc4bd9219ff3eed831f7203bc3722fd0ed9cbc0d01135d5c2b795668b" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" + "sha256": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "b53124255d1e34b5ed31154bd5aa1e3d60f75b550b3b3b8ce90938245c36d13f" + "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 33a4cef2b..072e43472 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -17,7 +17,7 @@ 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="3984ca2a246f948b98aec22323e2d6724383739a90db0666918c8aaf3b41a89c" +readonly DD_PLAN_SHA256="1e915d097915a885657be68ca70e2a158b90e003e468462291805d3b6c88517e" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 12b7f1b89..8294d0cb5 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -17,7 +17,7 @@ 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="fcaef6b348e75e6956e6661335ff91432e04345091a6e59ed07d537410b85d36" +readonly DD_PLAN_SHA256="82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index d5d52ec38..5a664210c 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="e50d3a25a6f1aa05bd4b14818c02809007347a49073c0af8fd453dc1f9b8183a", + expected_fixture_sha256="c312271fe41fa988f439d54f5a9056880875431746c0d7b9bbf470a60f1b2316", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="5d06a7a7bd3bb8be9a62fdb251b8956368d27c9d2d9c4ab7d2af418c9bff889a", + expected_fixture_sha256="dbc7ff47a2e9a27f90e2ddffd06df0a2b00be9781635ece189a551bcb27f3943", ) diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index f591869bc..785a5b713 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "182770715c75496f62ed196cb04f088638a01c2d7ea7504547371cfce74c2b04" +_COMPATIBILITY_FIXTURE_DIGEST = "4128e04c4c4131de2029ec32bd9b9b337b2f695f42b938a1cb7d6cfe07dd4bd1" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) From 1646120bbe36615afc8c8e4cb863a11d26067a56 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 12 Oct 2026 16:09:00 -0600 Subject: [PATCH 43/62] fix(slurm): leave host patterns empty by default (#1014) Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/getting-started.mdx | 7 +++---- fern/versions/latest/pages/slurm/profiles.mdx | 5 ++--- .../src/data_designer/slurm/services/profiles.py | 3 +-- .../tests/services/test_profile_service.py | 8 +++----- 4 files changed, 9 insertions(+), 14 deletions(-) diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index b732f25ab..1a41aab6c 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -62,11 +62,10 @@ data-designer slurm profile init \ --workspace-root "$DD_SLURM_WORKSPACE" \ --image-build-partition cpu \ --account research \ - --partition gpu \ - --host-pattern 'login*.example.test' + --partition gpu ``` -The command writes the profile with restrictive permissions and never overwrites an existing file. Before submitting work, review every generated site value. In particular, confirm the account, partitions, hostname pattern, 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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. 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. +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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. 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 @@ -74,7 +73,7 @@ The command writes the profile with restrictive permissions and never overwrites data-designer slurm profile validate --cluster primary ``` -Validation checks strict profile loading, cluster selection, workspace access, and scheduler facts. Keep `--cluster primary` explicit when running from a compute host whose hostname does not match the login-host pattern. See [Profiles and clusters](/slurm/profiles-and-clusters) for non-default files and multiple clusters. +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 diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index 1aabf8dbd..822ca1713 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -21,11 +21,10 @@ data-designer slurm profile init \ --image-build-partition cpu \ --cluster primary \ --account research \ - --partition gpu \ - --host-pattern 'login*.example.test' + --partition gpu ``` -The starter sets `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, host patterns, workspace, GPU settings, and image-build resources against site policy before validation or image registration. +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 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 index 652330d5c..7a9d207f7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/profiles.py @@ -119,14 +119,13 @@ def initialize( environ=self._environ, home_directory=self._home_directory, ) - selected_patterns = host_patterns or _normalize_hostnames(self._hostname_resolver()) 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=selected_patterns, + host_patterns=host_patterns, ) SlurmProfileCatalog.model_validate(payload, strict=True) _create_profile_file(path, _serialize_catalog(payload, suffix=path.suffix)) diff --git a/packages/data-designer-slurm/tests/services/test_profile_service.py b/packages/data-designer-slurm/tests/services/test_profile_service.py index f04e43ddc..8914146d4 100644 --- a/packages/data-designer-slurm/tests/services/test_profile_service.py +++ b/packages/data-designer-slurm/tests/services/test_profile_service.py @@ -194,10 +194,12 @@ def test_profile_init_and_auto_gpu_resolution_accept_multiple_partitions(tmp_pat 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"] @@ -226,7 +228,6 @@ def test_profile_init_uses_explicit_environment_then_home_path_precedence(tmp_pa from_home = create_slurm_profile_service( environ={}, home_directory=tmp_path, - hostname_resolver=lambda: ("Node.EXAMPLE.test", "node.example.test", "node-short"), ).initialize( workspace_root=tmp_path / "workspace", image_build_partition="cpu", @@ -236,10 +237,7 @@ def test_profile_init_uses_explicit_environment_then_home_path_precedence(tmp_pa 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 == [ - "node.example.test", - "node-short", - ] + 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: From f5214c1f821ee95703071ec0310dfdc693e87aa5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 12 Oct 2026 16:52:39 -0600 Subject: [PATCH 44/62] feat(slurm): improve image add progress and removal (#1016) * feat(slurm): confirm image removal and allow artifact deletion Show an explicit removal result in human output. Allow opt-in deletion only for unshared package-owned OCI SQSH artifacts, with safe failure reporting and documentation. Signed-off-by: Nabin Mulepati * feat(slurm): timestamp image job progress Signed-off-by: Nabin Mulepati * feat(slurm): clarify image add progress Report deterministic image job phases and reduce idle status noise. Omit blank followed-log lines and show elapsed time after registration. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/cli-reference.mdx | 2 +- fern/versions/latest/pages/slurm/images.mdx | 12 ++- .../latest/pages/slurm/operations.mdx | 2 +- .../src/data_designer/slurm/cli.py | 23 ++++- .../src/data_designer/slurm/images/errors.py | 4 + .../data_designer/slurm/images/lifecycle.py | 33 ++++++- .../data_designer/slurm/images/registry.py | 92 +++++++++++++++++- .../src/data_designer/slurm/images/service.py | 6 +- .../slurm/services/image_lifecycle.py | 79 +++++++++++++--- .../data_designer/slurm/services/images.py | 12 ++- .../data_designer/slurm/services/wiring.py | 20 +++- .../tests/images/test_lifecycle.py | 1 + .../tests/images/test_registry_store.py | 94 ++++++++++++++++++- .../tests/services/test_image_wiring.py | 89 +++++++++++++++--- .../tests/services/test_wiring.py | 18 ++++ .../data-designer-slurm/tests/test_cli.py | 61 ++++++++++-- 16 files changed, 484 insertions(+), 64 deletions(-) diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx index 3eeff803f..8aece4647 100644 --- a/fern/versions/latest/pages/slurm/cli-reference.mdx +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -11,7 +11,7 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | `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 without deleting its SQSH artifact. | +| `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. | diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index 481fadb19..68a961d1f 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -33,14 +33,14 @@ Custom OCI sources must use `repository@sha256:<64 lowercase hex characters>`. T 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 scheduler state every 30 seconds and follows the image job's standard output and error by default: +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 log stream is bounded to 1 MiB and written to standard error; 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. +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 @@ -79,6 +79,14 @@ Removing an alias does not delete its SQSH artifact: 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`, and only when no other registered alias points to the same file. Existing SQSH files registered from outside the workspace are never deleted by this command. Make sure no active run needs the artifact before removing it. + +```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 diff --git a/fern/versions/latest/pages/slurm/operations.mdx b/fern/versions/latest/pages/slurm/operations.mdx index 2dae74437..0c3858a8b 100644 --- a/fern/versions/latest/pages/slurm/operations.mdx +++ b/fern/versions/latest/pages/slurm/operations.mdx @@ -7,7 +7,7 @@ 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; do not parse them as the command result. +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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 719df3400..7d971b2f5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -7,6 +7,7 @@ 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 @@ -358,14 +359,22 @@ def image_remove_command( name: str = typer.Argument(...), 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 OCI SQSH artifact."), ) -> None: - """Unregister one alias without deleting its SQSH artifact.""" + """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), + lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).remove( + name, delete_artifact=delete_artifact + ), ) - _emit_result(result) + 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: @@ -471,7 +480,7 @@ def _emit_json(value: object, *, err: bool = False) -> None: def _emit_progress(message: str) -> None: try: if _progress_enabled(): - typer.echo(message, err=True) + typer.echo(f"[{_utc_timestamp()}] {message}", err=True) except Exception: # Progress is optional and must not change the submission outcome. pass @@ -479,7 +488,7 @@ def _emit_progress(message: str) -> None: def _emit_image_log(message: str) -> None: try: - typer.echo(message, err=True) + typer.echo(f"[{_utc_timestamp()}] {message}", err=True) except Exception: # Log display is optional and must not change the submission outcome. pass @@ -489,6 +498,10 @@ 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 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 index db5807e21..156ffe1cd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/errors.py @@ -18,6 +18,10 @@ 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.""" 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 index f64319136..2f8add476 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/lifecycle.py @@ -47,6 +47,7 @@ _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) @@ -176,7 +177,9 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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 @@ -209,6 +212,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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} @@ -260,6 +264,11 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: [[ "$(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 @@ -289,7 +298,8 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: fi export ENROOT_MAX_PROCESSORS="${{SLURM_CPUS_PER_TASK}}" -{source_block}if [[ ! -f "${{DD_IMAGE_SQSH}}" || -L "${{DD_IMAGE_SQSH}}" ]]; then +{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 @@ -320,6 +330,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: 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}}" \ @@ -358,6 +369,7 @@ def render_image_lifecycle_script(plan: ImageLifecyclePlan) -> str: printf '%s\\n' 'SQSH path changed during inspection' >&2 exit 74 fi +report_phase inspection_complete true """ @@ -491,6 +503,25 @@ def read_image_lifecycle_log( 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, 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 index be75f72fe..e1afbf52c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py @@ -22,6 +22,7 @@ from data_designer.slurm.config import ImageRef from data_designer.slurm.contracts import Identifier, validate_absolute_path from data_designer.slurm.images.errors import ( + ImageArtifactRemovalError, ImageConflictError, ImageNotFoundError, ImageRegistryError, @@ -178,9 +179,11 @@ def register( raise return image - def unregister(self, name: Identifier) -> RegisteredImage: - """Atomically remove an alias without deleting its underlying SQSH artifact.""" + 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): @@ -209,6 +212,91 @@ def unregister(self, name: Identifier) -> RegisteredImage: self._save(image_root_descriptor, snapshot, updated) return removed + def _unregister_with_artifact(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, + ), + 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 _load(self) -> ImageRegistryDocument: try: self._image_root.lstat() 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 index 65d091b3d..f1004f68a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/service.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/service.py @@ -232,9 +232,9 @@ def resolve_for_planning(self, reference: ImageRef, *, expected_kind: ImageKind) inspection=image.inspection, ) - def unregister(self, name: str) -> RegisteredImage: - """Remove an alias while leaving its underlying artifact untouched.""" - return self._registry.unregister(name) + 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: 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 index dd3b3f6eb..656501ea2 100644 --- 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 @@ -26,6 +26,7 @@ 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 @@ -50,11 +51,19 @@ 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): @@ -77,10 +86,11 @@ def __init__(self, prepared: PreparedImageLifecycleJob, job_id: int, report: Log self._unavailable_reported = False self._limit_reported = False - def drain(self, *, final: bool = False) -> None: + 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 + return False while self._source_bytes < _LOG_TOTAL_BYTES: received = False for stream in ("stdout", "stderr"): @@ -99,14 +109,14 @@ def drain(self, *, final: bool = False) -> None: if not self._unavailable_reported: self._emit("Image job logs are unavailable.") self._unavailable_reported = True - return + 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 + return self._displayed_bytes > before if not final or not received: break if final or self._source_bytes >= _LOG_TOTAL_BYTES: @@ -117,6 +127,7 @@ def drain(self, *, final: bool = False) -> None: 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) @@ -135,6 +146,8 @@ def _append(self, stream: str, content: bytes, *, final: bool = False) -> None: 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)}") @@ -212,7 +225,7 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: raise _unavailable("image lifecycle job cannot be prepared") from None try: - self._report_progress("Submitting image inspection job...") + 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: @@ -227,10 +240,15 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: follower = _JobLogFollower(prepared, receipt.job_id, self._logs) if self._logs is not None else None try: - self._report_progress("Image inspection job submitted; checking every 30 seconds.") - self._wait_for_success(receipt.job_id, follower) - self._report_progress("Inspection complete; verifying and registering image...") - return publish_completed_image_lifecycle(prepared, replace=replace) + 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, @@ -255,11 +273,20 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> RegisteredImage: self._cancel_and_cleanup(receipt.job_id, prepared, follower) raise - def _wait_for_success(self, job_id: int, follower: _JobLogFollower | None) -> None: + 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( @@ -267,14 +294,31 @@ def _wait_for_success(self, job_id: int, follower: _JobLogFollower | None) -> No observed_at=observed_at, previous={job_id: previous}, )[0] - self._report_progress(f"Image inspection job: {observation.state.value.replace('_', ' ')}.") - if follower is not None: - follower.drain() + 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: - self._report_progress("Waiting for successful job exit evidence...") 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") @@ -367,4 +411,11 @@ 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 index 531c468c5..84ab6defa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/images.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/images.py @@ -47,8 +47,8 @@ def list(self) -> tuple[RegisteredImage, ...]: def get(self, name: Identifier) -> RegisteredImage: """Return one registered image.""" - def remove(self, name: Identifier) -> RegisteredImage: - """Unregister one alias without deleting its SQSH.""" + def remove(self, name: Identifier, *, delete_artifact: bool = False) -> RegisteredImage: + """Unregister one alias and optionally delete its package-owned SQSH.""" class SlurmImageService: @@ -145,14 +145,16 @@ def get_image() -> RegisteredImage: return _invoke_service_backend(operation, get_image) - def remove(self, name: Identifier) -> RegisteredImage: - """Unregister one alias without deleting its SQSH artifact.""" + 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) + 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index f063f2325..85e59e379 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -42,7 +42,12 @@ resolve_profile, ) from data_designer.slurm.contracts import Identifier, ShardId -from data_designer.slurm.images.errors import ImageConflictError, ImageNotFoundError, SlurmImageError +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 @@ -768,8 +773,11 @@ def list(self) -> tuple[RegisteredImage, ...]: 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) -> RegisteredImage: - return self._invoke_registry(SlurmServiceOperation.REMOVE_IMAGE, lambda: self._verified.unregister(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: @@ -777,8 +785,10 @@ def _invoke_registry(operation: SlurmServiceOperation, call: Callable[[], _Resul return call() except ImageNotFoundError: raise SlurmServiceError(SlurmServiceErrorCode.NOT_FOUND, operation, "image is not registered") from None - except ImageConflictError: - raise SlurmServiceError(SlurmServiceErrorCode.CONFLICT, operation, "image registry conflict") 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" diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index ea27d6d70..d428e7445 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -220,6 +220,7 @@ def test_rendered_image_lifecycle_job_computes_digest_and_runs_inspection( 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" diff --git a/packages/data-designer-slurm/tests/images/test_registry_store.py b/packages/data-designer-slurm/tests/images/test_registry_store.py index db1aabb31..fbb36b0bb 100644 --- a/packages/data-designer-slurm/tests/images/test_registry_store.py +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import json import os import stat @@ -20,7 +21,12 @@ 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 ImageConflictError, ImageRegistryError, ImageVerificationError +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 @@ -129,6 +135,92 @@ def test_registry_allows_aliases_with_identical_facts_for_one_path(tmp_path: Pat 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() + + +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") diff --git a/packages/data-designer-slurm/tests/services/test_image_wiring.py b/packages/data-designer-slurm/tests/services/test_image_wiring.py index 0bb3965d9..9358e6418 100644 --- a/packages/data-designer-slurm/tests/services/test_image_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_image_wiring.py @@ -121,15 +121,16 @@ def query_accounting(self, selectors: object) -> tuple[SlurmAccountingEntry, ... class _TransitionLauncher(_Launcher): - def __init__(self, *, on_submit: Callable[[], None]) -> None: + 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 == 3: + elif self.queue_queries == 2 + self.running_checks: self.queue_state = SchedulerState.COMPLETED self.accounting_state = SchedulerState.COMPLETED return super().query_queue(selectors) @@ -206,7 +207,7 @@ def test_default_image_add_waits_for_completed_job_accounting_evidence(tmp_path: assert not _job_directory(workspace).exists() -def test_image_add_reports_progress_on_each_thirty_second_check(tmp_path: Path) -> None: +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" @@ -227,14 +228,12 @@ def test_image_add_reports_progress_on_each_thirty_second_check(tmp_path: Path) assert clock.sleep_calls == [30.0, 30.0] assert progress == [ - "Submitting image inspection job...", - "Image inspection job submitted; checking every 30 seconds.", - "Image inspection job: completed.", - "Waiting for successful job exit evidence...", - "Image inspection job: completed.", + "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...", - "Image inspection job: completed.", - "Inspection complete; verifying and registering image...", + "Registering image alias...", + "Image add complete: client (total elapsed 00:01:00).", ] @@ -243,7 +242,7 @@ def test_image_add_reports_pending_running_and_completed(tmp_path: Path) -> None source = tmp_path / "client.sqsh" content = b"client image" source.write_bytes(content) - launcher = _TransitionLauncher(on_submit=lambda: _write_inspection(workspace, 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( @@ -257,10 +256,70 @@ def test_image_add_reports_pending_running_and_completed(tmp_path: Path) -> None service.add(ImageBuildRequest(name="client", kind="client", source=source.as_posix())) - assert clock.sleep_calls == [30.0, 30.0] - assert "Image inspection job: pending." in progress - assert "Image inspection job: running." in progress - assert "Image inspection job: completed." in progress + 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: diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 5bd7ce4d0..381e7da88 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -1472,3 +1472,21 @@ def test_production_image_registry_operations( 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/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index 7c8947361..cd5a39ad6 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -4,10 +4,12 @@ from __future__ import annotations import json +import re from pathlib import Path 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 @@ -24,6 +26,15 @@ 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]] = [] @@ -157,6 +168,38 @@ def list(self) -> tuple[()]: assert result.stdout == "(none)\n" +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) @@ -235,7 +278,7 @@ def test_execute_reports_submission_in_interactive_terminal( assert result.exit_code == 0 assert json.loads(result.stdout)["state"] == "submitted" - assert result.stderr == "Preparing and submitting run...\n" + assert _timestamped_lines(result.stderr) == ["Preparing and submitting run..."] def test_progress_stream_failure_does_not_stop_submissions( @@ -251,7 +294,7 @@ def test_progress_stream_failure_does_not_stop_submissions( original_echo = cli_module.typer.echo def echo_with_broken_progress(message: str, *, err: bool = False) -> None: - if message.startswith("Preparing and submitting"): + if "Preparing and submitting" in message: raise BrokenPipeError("progress stream closed") original_echo(message, err=err) @@ -311,7 +354,7 @@ def create_benchmark_service(**kwargs: object) -> _BenchmarkService: 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 run_result.stderr == "Preparing benchmark launch...\nBenchmark case 1/1: submitted.\n" + assert _timestamped_lines(run_result.stderr) == ["Preparing benchmark launch...", "Benchmark case 1/1: submitted."] def test_retry_emits_deterministic_json_and_maps_task_ids(monkeypatch) -> None: @@ -369,7 +412,7 @@ def test_retry_force_skips_confirmation(monkeypatch) -> None: assert result.exit_code == 0 assert service.retry_calls == [("42", None, "if_possible", False)] - assert result.stderr == "Preparing and submitting retry...\n" + assert _timestamped_lines(result.stderr) == ["Preparing and submitting retry..."] def test_retry_decline_emits_stable_result(monkeypatch) -> None: @@ -430,7 +473,7 @@ def test_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypat "state": "submitted", } assert service.collection_calls == [(input_path, output_path, 2)] - assert result.stderr == "Preparing and submitting collection job...\n" + assert _timestamped_lines(result.stderr) == ["Preparing and submitting collection job..."] @pytest.mark.parametrize( @@ -528,7 +571,7 @@ def add(self, request: ImageBuildRequest, *, replace: bool) -> ImageBuildRequest def create_service(**kwargs: object) -> ImageService: progress = kwargs["progress"] assert callable(progress) - progress("Image inspection job: pending.") + progress("Image add job 42: pending (elapsed 00:00:00).") return ImageService() monkeypatch.setattr(cli_module, "create_slurm_image_service", create_service) @@ -541,11 +584,11 @@ def create_service(**kwargs: object) -> ImageService: 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 result.stderr.splitlines() == ( + assert _timestamped_lines(result.stderr) == ( [ "Resolving vLLM image tag...", "Preparing image registration...", - "Image inspection job: pending.", + "Image add job 42: pending (elapsed 00:00:00).", ] if interactive else [] @@ -573,7 +616,7 @@ def create_service(**kwargs: object) -> ImageService: assert result.exit_code == 0, result.output assert json.loads(result.stdout)["source"] == source - assert result.stderr == "[image stderr] inspection started\n" + assert _timestamped_lines(result.stderr) == ["[image stderr] inspection started"] @pytest.mark.parametrize("interactive", [False, True]) From fa717c6971ddfc7eeed424d442769e2a408781c3 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 13 Oct 2026 10:08:23 -0600 Subject: [PATCH 45/62] fix(slurm): simplify image details output (#1018) * fix(slurm): simplify image details output Signed-off-by: Nabin Mulepati * fix(slurm): clarify image add and list results Signed-off-by: Nabin Mulepati * fix(slurm): clarify image command help Signed-off-by: Nabin Mulepati * fix(slurm): sanitize image list runtime Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/images.mdx | 2 + .../src/data_designer/slurm/cli.py | 76 +++++++-- .../data-designer-slurm/tests/test_cli.py | 150 +++++++++++++++++- 3 files changed, 214 insertions(+), 14 deletions(-) diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index 68a961d1f..1d78c9853 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -60,6 +60,8 @@ 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 7d971b2f5..7ecd7a663 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -18,9 +18,15 @@ from pydantic import BaseModel, ValidationError from data_designer.slurm.cli_benchmark import create_benchmark_app -from data_designer.slurm.config import ImageBuildRequest, SlurmConfigLoadError, load_run_config +from data_designer.slurm.config import ( + ClientImageInspection, + ImageBuildRequest, + ServingImageInspection, + SlurmConfigLoadError, + load_run_config, +) from data_designer.slurm.contracts import canonical_json -from data_designer.slurm.images.records import validate_oci_source_for_lifecycle +from data_designer.slurm.images.records import RegisteredImage, validate_oci_source_for_lifecycle from data_designer.slurm.images.vllm_source import ( VllmSourceResolutionError, VllmTagNotFoundError, @@ -265,10 +271,10 @@ def profile_validate_command( @image_app.command("add") def image_add_command( - source: str = typer.Argument(..., help="Digest-pinned OCI, versioned official vLLM image, or SQSH"), - kind: str = typer.Option("serving", "--kind", help="Image role (defaults to serving)"), - name: str | None = typer.Option(None, "--name"), - replace: bool = typer.Option(False, "--replace"), + 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", @@ -322,7 +328,7 @@ def add() -> BaseModel: logs=_emit_image_log if (follow_logs if follow_logs is not None else _progress_enabled()) else None, ).add(request, replace=replace) - _emit_result(_invoke(operation, add)) + _emit_registered_image(_invoke(operation, add), registered=True) @image_app.command("ls") @@ -336,12 +342,17 @@ def image_list_command( operation, lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).list(), ) - _emit_value([image.model_dump(mode="json") for image in images]) + 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(...), + 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: @@ -351,15 +362,17 @@ def image_info_command( operation, lambda: create_slurm_image_service(profile_file=profile_file, cluster=cluster).get(name), ) - _emit_result(result) + _emit_registered_image(result) @image_app.command("rm") def image_remove_command( - name: str = typer.Argument(...), + 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 OCI SQSH artifact."), + 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 @@ -419,6 +432,45 @@ def _emit_result(result: BaseModel) -> None: _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) diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index cd5a39ad6..ef805700d 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -14,7 +14,13 @@ 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 +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, @@ -165,7 +171,147 @@ def list(self) -> tuple[()]: result = CliRunner().invoke(cli_module.create_cli(), ["image", "ls"]) assert result.exit_code == 0 - assert result.stdout == "(none)\n" + 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: From e3fc404eabf5ffcd9c7328d4397b12256335d6c2 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 13 Oct 2026 13:29:34 -0600 Subject: [PATCH 46/62] fix(slurm): improve preparation and allocation diagnostics (#1019) * fix(slurm): explain missing builder sources Signed-off-by: Nabin Mulepati * fix(slurm): show serving startup progress in allocation logs Signed-off-by: Nabin Mulepati * fix(slurm): point allocation logs to client generation progress Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../src/data_designer/slurm/cli.py | 7 ++++- .../src/data_designer/slurm/config/loading.py | 2 ++ .../data_designer/slurm/runtime/entrypoint.sh | 21 ++++++++++++++- .../data_designer/slurm/services/wiring.py | 19 ++++++++++++- .../tests/runtime/test_shell_runtime.py | 14 ++++++++++ .../tests/services/test_wiring.py | 27 +++++++++++++++++++ .../data-designer-slurm/tests/test_cli.py | 19 +++++++++++++ 7 files changed, 106 insertions(+), 3 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index 7ecd7a663..ed17cd262 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -35,6 +35,8 @@ ) from data_designer.slurm.security import redact_sensitive_diagnostic from data_designer.slurm.services import ( + SlurmRetryExecution, + SlurmRunExecution, SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation, @@ -429,7 +431,10 @@ def _fail(error: SlurmServiceError) -> NoReturn: def _emit_result(result: BaseModel) -> None: - _emit_value(result.model_dump(mode="json")) + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py index b4cad4be5..e95bab69c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/loading.py @@ -90,6 +90,8 @@ def load_builder_payload(path: str | Path) -> dict[str, JsonValue]: 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: 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 index 8974914aa..05835100d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -14,6 +14,10 @@ DD_RUNTIME_FINALIZED=0 DD_RUNTIME_PHASE=initialization DD_RUNTIME_RETRY_ARGUMENTS=() +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 @@ -66,16 +70,21 @@ dd_slurm_run_allocation() { 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 @@ -85,12 +94,19 @@ dd_slurm_run_allocation() { 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 @@ -99,6 +115,7 @@ dd_slurm_run_allocation() { dd_run_bound_control_phase succeed DD_RUNTIME_FINALIZED=1 DD_RUNTIME_PHASE=completed + dd_runtime_log 'Run completed.' } dd_run_bound_control_phase() { @@ -187,7 +204,7 @@ dd_verify_gpu_count() { dd_start_servers() { dd_read_step_ids "${DD_RUNTIME_MANIFEST}" server - local step_id deployment delay remaining + 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}" @@ -199,6 +216,8 @@ dd_start_servers() { 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 } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 85e59e379..8a87b4a69 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -204,7 +204,14 @@ def _prepare( self._images.resolve_for_planning(deployment.server.image, expected_kind=ImageKind.SERVING) for deployment in authored.deployments ) - builder_payload = _resolve_builder_payload(authored, source_root) + 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 dependencies = stack.enter_context( self._dependencies.resolve( authored.client.dependencies, @@ -958,6 +965,16 @@ 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 _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: " diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 62b6cd3ec..7fe9e6b28 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -87,6 +87,20 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( 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: diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 381e7da88..793ec60a8 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -432,6 +432,33 @@ def test_public_plan_resolves_builder_relative_to_explicit_source_root( 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 + + def test_production_wiring_submits_after_publisher_initialization( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index ef805700d..cc7e3c6a4 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -144,6 +144,25 @@ def test_default_output_is_readable_and_json_is_explicit( 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") From ede15bb14597213bb7447c8d84ffd886f9aeb433 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 13 Oct 2026 20:34:51 -0600 Subject: [PATCH 47/62] fix(slurm): choose Parquet partitions at merge (#1021) * fix(slurm): choose Parquet partitions at merge Signed-off-by: Nabin Mulepati * fix(slurm): preserve collection intent on retry Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/cli-reference.mdx | 3 + .../latest/pages/slurm/getting-started.mdx | 8 +- fern/versions/latest/pages/slurm/overview.mdx | 2 +- .../pages/slurm/retry-and-collection.mdx | 6 +- .../latest/pages/slurm/run-configuration.mdx | 13 +- .../data-designer-slurm/examples/run.yaml | 3 - .../data_designer/slurm/benchmark/compiler.py | 8 - .../src/data_designer/slurm/cli.py | 4 +- .../data_designer/slurm/config/__init__.py | 2 - .../src/data_designer/slurm/config/builder.py | 8 - .../src/data_designer/slurm/config/run.py | 15 +- .../data_designer/slurm/planning/__init__.py | 2 - .../data_designer/slurm/planning/compiler.py | 1 - .../data_designer/slurm/planning/models.py | 16 -- .../slurm/planning/resolution.py | 49 ------ .../slurm/planning/validation.py | 7 - .../slurm/services/retry_collection.py | 13 +- .../data_designer/slurm/state/collection.py | 11 +- .../slurm/state/collection_merge.py | 60 +------ .../slurm/state/collection_validation.py | 5 +- .../slurm/state/collection_worker.py | 3 +- .../data_designer/slurm/state/destinations.py | 12 +- .../tests/config/test_loading_builder.py | 3 +- .../tests/contracts/golden/authored_run.json | 5 - .../contracts/golden/authored_run_single.json | 5 - .../contracts/golden/multi_node_plan.json | 7 +- .../contracts/golden/single_node_plan.json | 7 +- .../tests/contracts/test_config_records.py | 10 ++ .../tests/contracts/test_planning_records.py | 21 --- .../golden/finalization_chain.json | 14 +- .../tests/launcher/test_collection.py | 12 +- .../tests/planning/test_compiler.py | 96 +--------- .../tests/services/test_services.py | 4 +- .../tests/services/test_wiring.py | 14 +- .../golden/rendered/multi_node.sbatch | 2 +- .../golden/rendered/single_node.sbatch | 2 +- .../slurm_test_fakes/test_rendered_scripts.py | 4 +- .../tests/state/test_retry_collection.py | 165 +++++++----------- .../tests/state/test_state_golden_records.py | 2 +- .../tests/state/test_store.py | 31 ---- 40 files changed, 164 insertions(+), 491 deletions(-) diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx index 8aece4647..ec159b271 100644 --- a/fern/versions/latest/pages/slurm/cli-reference.mdx +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -29,6 +29,9 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | `--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 the managed run directory containing shard winners. | +| `--output-path PATH` | `merge` | Select an authorized compute-visible collection directory. | +| `--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. | diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index 1a41aab6c..228faee00 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -134,12 +134,9 @@ deployments: tensor_parallel: 1 submission: time_limit: 01:00:00 -output: - format: parquet - partitions: 1 ``` -The repository also contains these 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. +The repository also contains these 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. ## 9. Dry-run the submission @@ -182,7 +179,10 @@ Use `status` to confirm every shard succeeded. To create one collected dataset a data-designer slurm merge \ --input-path /workspace/alice/data-designer-slurm/runs/ \ --output-path /workspace/alice/exports/greetings \ + --num-partitions 1 \ --cluster primary ``` +`--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/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index fba10bde7..742e6ed5e 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -28,7 +28,7 @@ Run Slurm commands from a site-approved submit environment with the cluster CLI ## Storage model -The configured `workspace_root` is durable and shared. It contains image metadata, immutable run inputs, Slurm logs, attempt state, winning results, collection state, and default output. +The configured `workspace_root` is durable and shared. It contains image metadata, immutable run inputs, Slurm logs, attempt state, winning results, and collection state. Collected output goes to the destination selected with `slurm merge --output-path`. 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. diff --git a/fern/versions/latest/pages/slurm/retry-and-collection.mdx b/fern/versions/latest/pages/slurm/retry-and-collection.mdx index 7816bc75a..3cc0f9517 100644 --- a/fern/versions/latest/pages/slurm/retry-and-collection.mdx +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -34,7 +34,7 @@ Retries create new attempts. The original attempt records remain durable, and ea ## Collect winners -Default successful output is already managed beneath the run directory. Use `merge` to submit a zero-GPU CPU job that collects shard winners to an explicit compute-visible destination: +Successful shards publish Parquet winners under the run directory. Use `merge` to submit a zero-GPU CPU job that collects them to an explicit compute-visible destination: ```bash data-designer slurm merge \ @@ -44,7 +44,7 @@ data-designer slurm merge \ `--input-path` must be exactly a managed run directory under the selected profile workspace. The output path must be absolute, writable, compute-visible, and authorized by the profile mounts. -If supplied, `--num-partitions` must equal the partition count persisted in the run: +`--num-partitions` selects the number of collected Parquet files. It defaults to 1 for a new collection and cannot exceed the run's record count: ```bash data-designer slurm merge \ @@ -53,4 +53,4 @@ data-designer slurm merge \ --num-partitions 20 ``` -Repeating collection returns or reuses managed state when safe. Conflicting destinations or partition counts are rejected. +Repeating the same merge command returns or reuses managed collection state when safe. If you omit `--num-partitions` on a repeat call, the existing collection keeps its count. Changing the destination or explicitly requesting a different partition count 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 index e08c48fba..3731619ef 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -34,14 +34,11 @@ array_tasks: submission: job_name: product-descriptions time_limit: 03:55:00 -output: - format: parquet - partitions: 1 ``` -`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, while `input_bindings.seed_path` and `managed_assets_path` accept absolute compute-visible paths. +`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. Output formats are `parquet`, `jsonl`, and `csv`. +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. @@ -93,13 +90,11 @@ invocation: array_tasks: count: 20 max_concurrent: 4 -output: - partitions: 20 ``` -`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. +`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. Array task count controls execution shards; it does not set the number of collected Parquet files. Use `slurm merge --num-partitions 20` if you want 20 output files. -Runs with more than one array task require Parquet output. They 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`. +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 diff --git a/packages/data-designer-slurm/examples/run.yaml b/packages/data-designer-slurm/examples/run.yaml index 32158ce8b..71abe7c55 100644 --- a/packages/data-designer-slurm/examples/run.yaml +++ b/packages/data-designer-slurm/examples/run.yaml @@ -21,6 +21,3 @@ deployments: tensor_parallel: 1 submission: time_limit: 01:00:00 -output: - format: parquet - partitions: 1 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 index 568533af9..a78e82364 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py @@ -6,7 +6,6 @@ from __future__ import annotations import math -import posixpath from pydantic import Field @@ -75,7 +74,6 @@ def compile( case_id, _compile_child( base_run, - benchmark_id=benchmark_id, case_id=case_id, concurrency=concurrency, selected_aliases=selected_aliases, @@ -99,7 +97,6 @@ def compile( def _compile_child( base_run: DataDesignerSlurmConfig, *, - benchmark_id: str, case_id: str, concurrency: int, selected_aliases: tuple[str, ...], @@ -131,17 +128,12 @@ def _compile_child( ) ) - output_root = base_run.output.root - output = base_run.output - if output_root is not None: - output = output.model_copy(update={"root": posixpath.join(output_root, benchmark_id, case_id)}) return DataDesignerSlurmConfig.model_validate( base_run.model_copy( update={ "name": _bounded_identifier(f"{base_run.name}-{case_id}"), "invocation": invocation, "deployments": deployments, - "output": output, } ).model_dump(mode="python") ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index ed17cd262..bf5d5cd9a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -213,7 +213,9 @@ def retry_command( def merge_command( input_path: Path = typer.Option(..., "--input-path", file_okay=False), output_path: Path = typer.Option(..., "--output-path", file_okay=False), - num_partitions: int | None = typer.Option(None, "--num-partitions", min=1), + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 05f258600..9842577f0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -63,7 +63,6 @@ InputBindings, InvocationConfig, LocalStdioMCPProviderConfig, - OutputConfig, RemoteMCPProviderConfig, ServerDeploymentConfig, SubmissionConfig, @@ -100,7 +99,6 @@ "InvocationConfig", "LiteralEnvironmentBinding", "LocalStdioMCPProviderConfig", - "OutputConfig", "ProfileSelectionSource", "PROFILE_FILE_ENVIRONMENT", "QueueBackpressureConfig", 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 index b0f3a42a1..21670bb5e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -24,7 +24,6 @@ InputBindings, InvocationConfig, MCPProviderConfig, - OutputConfig, ServerDeploymentConfig, SubmissionConfig, ) @@ -43,7 +42,6 @@ def __init__(self, builder: BuilderInput, *, name: str = "data-designer") -> Non self._deployments: list[ServerDeploymentConfig] = [] self._array_tasks = ArrayTasksConfig() self._submission = SubmissionConfig() - self._output = OutputConfig() @classmethod def from_config_builder( @@ -133,11 +131,6 @@ def with_submission(self, **values: object) -> DataDesignerSlurmConfigBuilder: self._submission = _validate_model(SubmissionConfig, values) return self - def with_output(self, **values: object) -> DataDesignerSlurmConfigBuilder: - """Set typed dataset output intent.""" - self._output = _validate_model(OutputConfig, values) - return self - def build(self) -> DataDesignerSlurmConfig: """Return the complete authored declaration without resolving ambient state.""" missing = [] @@ -159,7 +152,6 @@ def build(self) -> DataDesignerSlurmConfig: "deployments": self._deployments, "array_tasks": self._array_tasks, "submission": self._submission, - "output": self._output, }, ) 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 index b3fe0a4fb..cb125ad6e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -54,7 +54,6 @@ "LiteralEnvironmentBinding", "LocalStdioMCPProviderConfig", "MCPProviderConfig", - "OutputConfig", "QueueBackpressureConfig", "RemoteMCPProviderConfig", "SecretRef", @@ -174,7 +173,7 @@ class InvocationConfig(AuthoredConfig): @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) + 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) @@ -329,17 +328,6 @@ 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 OutputConfig(AuthoredConfig): - root: str | None = None - format: Literal["parquet", "jsonl", "csv"] = "parquet" - partitions: PositiveInt = 1 - - @field_validator("root") - @classmethod - def validate_root(cls, value: str | None) -> str | None: - return None if value is None else validate_absolute_path(value) - - class DataDesignerSlurmConfig(AuthoredConfig): """Complete portable intent for one Data Designer Slurm run.""" @@ -351,7 +339,6 @@ class DataDesignerSlurmConfig(AuthoredConfig): deployments: list[ServerDeploymentConfig] = Field(min_length=1) array_tasks: ArrayTasksConfig = Field(default_factory=ArrayTasksConfig) submission: SubmissionConfig = Field(default_factory=SubmissionConfig) - output: OutputConfig = Field(default_factory=OutputConfig) @model_validator(mode="after") def validate_run(self) -> DataDesignerSlurmConfig: 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 index c4f3d0067..d0f837c41 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/__init__.py @@ -18,7 +18,6 @@ ResolvedDeployment, ResolvedImage, ResolvedInvocation, - ResolvedOutput, ResolvedSlurmRunPlan, ResolvedSubmission, ResolvedTopology, @@ -38,7 +37,6 @@ "ResolvedDeployment", "ResolvedImage", "ResolvedInvocation", - "ResolvedOutput", "ResolvedSlurmRunPlan", "ResolvedSubmission", "ResolvedTopology", 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 index 2b296c39a..81fd40698 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -81,7 +81,6 @@ def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan array_tasks=effective.authored.array_tasks, shards=_compile_shards(effective), submission=effective.submission, - output=effective.output, container_mounts=tuple(effective.selected_profile.profile.container_mounts), runtime_bundle=effective.runtime_bundle, ) 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 index b7c5b6377..48f13f191 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -384,14 +384,6 @@ def validate_submission(self) -> ResolvedSubmission: return self -class ResolvedOutput(ContractValue): - root: str - format: Literal["parquet", "jsonl", "csv"] - partitions: PositiveInt - - _root_is_absolute = field_validator("root")(validate_absolute_path) - - class ResolvedSlurmRunPlan(ContractRecord): """Immutable allocation input consumed without ambient configuration.""" @@ -408,7 +400,6 @@ class ResolvedSlurmRunPlan(ContractRecord): array_tasks: ArrayTasksConfig shards: tuple[PlannedShard, ...] = Field(min_length=1) submission: ResolvedSubmission - output: ResolvedOutput container_mounts: tuple[ContainerMount, ...] = () runtime_bundle: ArtifactReference @@ -504,13 +495,6 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: ): raise ValueError("runtime bundle must be a content-addressed tar archive below the workspace runtime root") self._validate_shards(run_root) - if not is_path_below(self.output.root, profile.workspace_root): - raise ValueError("resolved output root must be below the selected workspace_root") - if paths_overlap(self.output.root, managed_assets_path): - raise ValueError("resolved output root must not overlap managed assets") - shards_root = posixpath.join(run_root, "shards") - if paths_overlap(self.output.root, shards_root): - raise ValueError("resolved output root must not overlap the run shard workspace") return self def _validate_shards(self, run_root: str) -> None: 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 index c00b1760e..f54063536 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -35,7 +35,6 @@ ResolvedDependencyLock, ResolvedImage, ResolvedInvocation, - ResolvedOutput, ResolvedSubmission, ) @@ -73,7 +72,6 @@ InputBindings, ResolvedBuilderInput, ResolvedInvocation, - ResolvedOutput, ResolvedSubmission, ) @@ -94,7 +92,6 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): deployment_images: tuple[ResolvedImage, ...] dependency_lock: ResolvedDependencyLock submission: ResolvedSubmission - output: ResolvedOutput runtime_bundle: ArtifactReference @@ -138,7 +135,6 @@ def resolve_slurm_config( deployment_images=deployment_images, dependency_lock=dependency_lock, submission=_materialize_submission(authored, selected_profile), - output=_materialize_output(authored, run_root), runtime_bundle=runtime_bundle, ) return validate_effective_slurm_config(effective) @@ -162,7 +158,6 @@ def validate_effective_slurm_config( authored = effective.authored profile = effective.selected_profile.profile workspace_root = profile.workspace_root - run_root = posixpath.join(workspace_root, "runs", effective.run_id) 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: @@ -181,9 +176,6 @@ def validate_effective_slurm_config( 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") - expected_output = _materialize_output(authored, run_root) - if effective.output != expected_output: - raise SlurmConfigResolutionError("resolved output does not match the authored output") _validate_dependency_resolution(authored, effective.client_runtime, effective.dependency_lock) _validate_resolved_images(authored, effective.deployment_images) @@ -191,14 +183,6 @@ def validate_effective_slurm_config( 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) - _validate_output_destination( - effective.output.root, - workspace_root, - run_root, - managed_assets_path=managed_assets_path, - ) - if effective.output.partitions > authored.invocation.num_records: - raise SlurmConfigResolutionError("output partitions must not exceed requested records") runtime_root = posixpath.join(workspace_root, "runtime") runtime_name = posixpath.basename(effective.runtime_bundle.path) if ( @@ -284,14 +268,6 @@ def _materialize_submission( ) -def _materialize_output(authored: DataDesignerSlurmConfig, run_root: str) -> ResolvedOutput: - return ResolvedOutput( - root=authored.output.root or posixpath.join(run_root, "output"), - format=authored.output.format, - partitions=authored.output.partitions, - ) - - 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( @@ -311,9 +287,6 @@ def _validate_sharding_constraints( ) -> None: if authored.array_tasks.count == 1: return - if authored.output.format != "parquet": - raise SlurmConfigResolutionError("multi-shard runs require parquet output") - 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) @@ -392,28 +365,6 @@ def _validate_resolved_image_identity(image: ResolvedImage) -> None: raise SlurmConfigResolutionError("resolved image path does not match the authored path") -def _validate_output_destination( - output_root: str, - workspace_root: str, - run_root: str, - *, - managed_assets_path: str, -) -> None: - if not is_path_below(output_root, workspace_root): - raise SlurmConfigResolutionError("output root must be below the selected workspace_root") - reserved = tuple(posixpath.join(workspace_root, name) for name in ("images", "runtime", "benchmarks")) - if any(paths_overlap(output_root, path) for path in reserved): - raise SlurmConfigResolutionError("output root must not overlap package-managed workspace state") - if paths_overlap(output_root, managed_assets_path): - raise SlurmConfigResolutionError("output root must not overlap managed assets") - runs_root = posixpath.join(workspace_root, "runs") - run_output_root = posixpath.join(run_root, "output") - if paths_overlap(output_root, runs_root) and not ( - output_root == run_output_root or is_path_below(output_root, run_output_root) - ): - raise SlurmConfigResolutionError("output root must not overlap another package-managed run") - - 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") 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 index 71d761be6..8a35033c2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -90,13 +90,6 @@ def validate_resolved_plan( ) _require(plan.submission.comment == authored.submission.comment, "resolved comment does not match authored input") - _require(plan.output.format == authored.output.format, "resolved output format does not match authored input") - _require( - plan.output.partitions == authored.output.partitions, "resolved output partitions do not match authored input" - ) - if authored.output.root is not None: - _require(plan.output.root == authored.output.root, "resolved output root does not match authored input") - _require( plan.client.dependency_lock.sha256 == dependency_lock.compute_sha256(), "dependency lock digest does not match the resolved plan", 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 index 46de9f076..ead9d14fc 100644 --- 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 @@ -27,8 +27,10 @@ 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) @@ -140,15 +142,15 @@ def collect( 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.output.partitions: + if num_partitions is not None and num_partitions > plan.invocation.authored.num_records: raise SlurmServiceError( SlurmServiceErrorCode.INVALID_REQUEST, operation, - "num_partitions must match the persisted run output partitions", + "num_partitions must not exceed requested records", ) - effective_partitions = plan.output.partitions status = SlurmCollectionCoordinator(self._workspace_root, run_id, self._launcher).submit( destination=destination, + num_partitions=num_partitions, submitted_at=self._clock(), ) if status.scheduler is None: @@ -157,13 +159,16 @@ def collect( 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=destination.as_posix(), - num_partitions=effective_partitions, + num_partitions=collection_plan.num_partitions, ) except SlurmServiceError: raise 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 index 9995b61c1..a44d72d00 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -87,6 +87,7 @@ 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.""" @@ -101,6 +102,8 @@ def submit( 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(), @@ -113,6 +116,11 @@ def submit( ) run, resolved_plan, _ = self._reader.load_context() resolved_destination = self._destinations.resolve(resolved_plan, destination) + effective_partitions = num_partitions + if effective_partitions is None: + effective_partitions = current_plan.num_partitions if current is not None else 1 + if effective_partitions > resolved_plan.invocation.authored.num_records: + raise StateConflictError("collection partitions must not exceed requested records") collection_plan = CollectionPlan( schema_version=1, collection_id=self._collections.get_next_collection_id(), @@ -122,7 +130,7 @@ def submit( planned_shards=self._inputs.get_winner_shards(), python_executable=self._python_executable, host_destination=resolved_destination.host_path, - num_partitions=resolved_plan.output.partitions, + num_partitions=effective_partitions, ) self._inputs.resolve(collection_plan) prepare_collection_destination( @@ -332,7 +340,6 @@ def _load_valid_result( plan, result, expected_records=resolved_plan.invocation.authored.num_records, - output_format=resolved_plan.output.format, ) 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") 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 index d7ec2803a..4392bc5b4 100644 --- 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 @@ -9,11 +9,11 @@ import os import stat from collections.abc import Generator -from contextlib import AbstractContextManager, contextmanager +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import BinaryIO, Protocol, TextIO +from typing import BinaryIO, Protocol import data_designer.lazy_heavy_imports as lazy from data_designer.slurm.filesystem import get_file_facts @@ -49,7 +49,6 @@ def close(self) -> None: ... @dataclass(frozen=True, slots=True) class _PartitionLayout: stage_path: Path - output_format: str schema: object record_counts: tuple[int, ...] @@ -59,12 +58,10 @@ class CollectionMerger: def __init__( self, - output_format: str, *, verifier: CandidateArtifactVerifier | None = None, completed_at: datetime | None = None, ) -> None: - self._output_format = output_format self._verifier = verifier if verifier is not None else CandidateArtifactVerifier() self._completed_at = completed_at @@ -83,7 +80,6 @@ def merge( raise OSError("collection inputs contain no records") layout = _PartitionLayout( stage_path=staged.path, - output_format=self._output_format, schema=first_batch.schema.remove_metadata(), record_counts=_partition_record_counts( sum(candidate.actual_records for candidate in candidates), @@ -133,10 +129,9 @@ def _write_partitions( cursor = _BatchCursor(first_batch, remaining_batches) outputs: list[CollectedOutputFile] = [] for partition_index, record_count in enumerate(layout.record_counts): - suffix = "jsonl" if layout.output_format == "jsonl" else layout.output_format - relative_path = f"part-{partition_index:05d}.{suffix}" + relative_path = f"part-{partition_index:05d}.parquet" output_path = layout.stage_path / relative_path - with _open_partition_writer(output_path, layout.output_format, layout.schema) as writer: + 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(): @@ -190,7 +185,7 @@ def _advance(self) -> None: class _BoundOutput: - def __init__(self, path: Path, output: BinaryIO | TextIO) -> None: + def __init__(self, path: Path, output: BinaryIO) -> None: self.path = path self._output = output @@ -219,39 +214,6 @@ def close(self) -> None: self._writer.close() -class _TextWriter: - def __init__(self, output: TextIO, output_format: str, bound_output: _BoundOutput) -> None: - self._output = output - self._bound_output = bound_output - self._format = output_format - self._is_first = True - - def write(self, batch: object) -> None: - frame = batch.to_pandas() - if self._format == "csv": - frame.to_csv(self._output, header=self._is_first, index=False) - else: - content = frame.to_json(orient="records", lines=True, force_ascii=False, date_format="iso") - if content: - self._output.write(content) - self._is_first = False - - def finish(self, relative_path: str, record_count: int) -> CollectedOutputFile: - return self._bound_output.describe(relative_path, record_count) - - -def _open_partition_writer( - output_path: Path, - output_format: str, - schema: object, -) -> AbstractContextManager[CollectionPartitionWriter]: - if output_format == "parquet": - return _open_parquet_writer(output_path, schema) - if output_format in {"csv", "jsonl"}: - return _open_text_writer(output_path, output_format) - raise ValueError(f"unsupported collection output format {output_format!r}") - - @contextmanager def _open_parquet_writer(output_path: Path, schema: object) -> Generator[CollectionPartitionWriter, None, None]: with output_path.open("x+b") as output: @@ -264,17 +226,7 @@ def _open_parquet_writer(output_path: Path, schema: object) -> Generator[Collect _sync_output(output) -@contextmanager -def _open_text_writer(output_path: Path, output_format: str) -> Generator[CollectionPartitionWriter, None, None]: - with output_path.open("x+", encoding="utf-8", newline="") as output: - os.fchmod(output.fileno(), 0o600) - try: - yield _TextWriter(output, output_format, _BoundOutput(output_path, output)) - finally: - _sync_output(output) - - -def _sync_output(output: BinaryIO | TextIO) -> None: +def _sync_output(output: BinaryIO) -> None: output.flush() os.fsync(output.fileno()) 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 index 369b4cf6d..3bf19dd91 100644 --- 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 @@ -34,6 +34,7 @@ def validate_collection_inputs( ) expected_records = resolved_plan.invocation.authored.num_records + _require(collection_plan.num_partitions <= expected_records, "collection partitions exceed requested records") actual_records = sum(candidate.actual_records for candidate in candidates) _require(actual_records == expected_records, "collection candidate rows do not match requested run rows") return candidates @@ -66,7 +67,6 @@ def validate_collection_result( result: CollectionResult, *, expected_records: int, - output_format: str, ) -> 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") @@ -78,8 +78,7 @@ def validate_collection_result( _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") - suffix = "jsonl" if output_format == "jsonl" else output_format - expected_paths = tuple(f"part-{index:05d}.{suffix}" for index in range(collection_plan.num_partitions)) + 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", 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 index fa2019f2a..8680166e0 100644 --- 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 @@ -101,7 +101,7 @@ def _execute_collection( derive_collection_staging_directory(plan), Path(destination.authorized_root), ) as staged: - merger = CollectionMerger(resolved_plan.output.format, completed_at=completed_at) + merger = CollectionMerger(completed_at=completed_at) result = merger.merge( plan, candidates, @@ -144,7 +144,6 @@ def _load_valid_result( plan, result, expected_records=resolved_plan.invocation.authored.num_records, - output_format=resolved_plan.output.format, ) 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") 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 index 746483d0e..5ad25342f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -32,12 +32,16 @@ def resolve( requested_destination: str | Path | None = None, ) -> CollectionDestination: """Return one authorized host output path.""" - raw_destination = plan.output.root if requested_destination is None else Path(requested_destination).as_posix() + workspace_root = plan.selected_profile.profile.workspace_root + raw_destination = ( + posixpath.join(workspace_root, "runs", plan.run_id, "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 - workspace_root = plan.selected_profile.profile.workspace_root _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)) @@ -57,10 +61,10 @@ def validate_persisted( 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") - if collection_plan.num_partitions != resolved_plan.output.partitions: - raise StateConflictError("collection partition count no longer matches resolved intent") return destination diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 083188628..618a091b1 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -97,7 +97,7 @@ def test_benchmark_loader_round_trips_strict_config( assert load_benchmark_config(path) == benchmark_config -def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: +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") @@ -108,7 +108,6 @@ def test_builder_rejects_unsupported_output_format(tmp_path: Path) -> None: ("with_invocation", {"num_records": 0, "dataset_name": "generated"}), ("with_array_tasks", {"count": 1, "max_concurrent": 2}), ("with_submission", {"time_limit": "invalid"}), - ("with_output", {"partitions": 0}), ], ) def test_builder_normalizes_invalid_authored_values( diff --git a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json index 5643f2099..6bcdc0505 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run.json @@ -128,11 +128,6 @@ } }, "name": "two-model-generation", - "output": { - "format": "parquet", - "partitions": 1, - "root": null - }, "schema_version": 1, "submission": { "account": null, 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 index b2e807822..2c0485901 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json +++ b/packages/data-designer-slurm/tests/contracts/golden/authored_run_single.json @@ -74,11 +74,6 @@ "run_config": {} }, "name": "single-node-generation", - "output": { - "format": "parquet", - "partitions": 1, - "root": null - }, "schema_version": 1, "submission": { "account": 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 index 7371ee1c9..2b912541e 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/multi_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-001/authored-config.json", - "sha256": "c172186cbb7cedbb783c743faad2f7af8e6018bf6600a7040b608b5fe9967dbe" + "sha256": "8c453e1064bf37c675b8ed62f1fe7b89234c7854e2c2f0955ae0c83938c73506" }, "builder": { "authored_source": null, @@ -361,11 +361,6 @@ "write_scheduler_events": false } }, - "output": { - "format": "parquet", - "partitions": 1, - "root": "/workspace/primary/runs/run-001/output" - }, "package_version": "0.9.2", "resolved_gpus_per_node": 8, "run_id": "run-001", 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 index 42b491a5a..5938747c5 100644 --- a/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json +++ b/packages/data-designer-slurm/tests/contracts/golden/single_node_plan.json @@ -5,7 +5,7 @@ }, "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "902aee7cc4bd9219ff3eed831f7203bc3722fd0ed9cbc0d01135d5c2b795668b" + "sha256": "bf9d45d571f5f288c74a3332cdd23576df60ce7bb3d934774f50bdce4b984fba" }, "builder": { "authored_source": null, @@ -207,11 +207,6 @@ "write_scheduler_events": false } }, - "output": { - "format": "parquet", - "partitions": 1, - "root": "/workspace/primary/runs/run-single/output" - }, "package_version": "0.9.2", "resolved_gpus_per_node": 8, "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index f74dc6a75..8f130d1e7 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -688,6 +688,16 @@ def test_run_validates_public_run_config_and_shard_count(authored_run: DataDesig 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 diff --git a/packages/data-designer-slurm/tests/contracts/test_planning_records.py b/packages/data-designer-slurm/tests/contracts/test_planning_records.py index 4d04b9488..6601b8db9 100644 --- a/packages/data-designer-slurm/tests/contracts/test_planning_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_planning_records.py @@ -131,7 +131,6 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu 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["output"].update(root="/outside/output"), lambda payload: payload.update(container_mounts=[]), lambda payload: payload["deployments"][0]["topology"].update(replica_count=2), lambda payload: payload["invocation"]["effective_input_bindings"].update( @@ -139,7 +138,6 @@ def test_resolved_builder_rejects_plaintext_secrets(multi_node_plan: ResolvedSlu ), 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"), - lambda payload: payload["output"].update(root="/workspace/primary/managed-assets"), ], ) def test_plan_rejects_invalid_boundaries(multi_node_plan: ResolvedSlurmRunPlan, mutator: object) -> None: @@ -207,25 +205,6 @@ def test_plan_rejects_otel_port_collision(multi_node_plan: ResolvedSlurmRunPlan) ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) -@pytest.mark.parametrize( - "output_root", - [ - "/workspace/primary/runs/run-001/shards/shard-00000/dataset", - "/workspace/primary/runs/run-001", - "/workspace/primary/runs", - ], -) -def test_plan_rejects_output_overlapping_shard_workspace( - multi_node_plan: ResolvedSlurmRunPlan, - output_root: str, -) -> None: - payload = multi_node_plan.model_dump(mode="json") - payload["output"]["root"] = output_root - - with pytest.raises(ValidationError, match="must not overlap"): - 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"][ diff --git a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json index 32c0e365e..eeb5661b1 100644 --- a/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json +++ b/packages/data-designer-slurm/tests/integration/golden/finalization_chain.json @@ -4,12 +4,12 @@ "attempt_ordinal": 1, "candidate_output": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" }, "created_at": "2026-08-19T12:00:02Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" + "sha256": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" }, "run_id": "run-single", "scheduler": { @@ -38,7 +38,7 @@ } ], "outcome": "complete", - "provenance_digest": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a", + "provenance_digest": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75", "requested_records": 8, "run_id": "run-single", "schema_version": 1, @@ -49,7 +49,7 @@ "attempt_id": "attempt-0001", "candidate_output_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" }, "completed_at": "2026-08-19T12:05:01Z", "dataset_path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/dataset", @@ -85,12 +85,12 @@ "run": { "authored_config": { "path": "/workspace/primary/runs/run-single/authored-config.json", - "sha256": "902aee7cc4bd9219ff3eed831f7203bc3722fd0ed9cbc0d01135d5c2b795668b" + "sha256": "bf9d45d571f5f288c74a3332cdd23576df60ce7bb3d934774f50bdce4b984fba" }, "created_at": "2026-08-19T12:00:00Z", "resolved_plan": { "path": "/workspace/primary/runs/run-single/resolved-plan.json", - "sha256": "82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" + "sha256": "8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" }, "run_id": "run-single", "schema_version": 1, @@ -118,7 +118,7 @@ "attempt_ordinal": 1, "candidate_manifest": { "path": "/workspace/primary/runs/run-single/shards/shard-00000/attempts/attempt-0001/output-manifest.json", - "sha256": "287150bc3b0d19e14c745fddc7822d9a1448e25e7aa6d356d1e5c3444fef5621" + "sha256": "0afc18907a46ce999d233872e461feee3c5ce8d1067c511f696250ed9101d8b6" }, "published_at": "2026-08-19T12:05:03Z", "run_id": "run-single", diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index f770275f4..231007815 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -66,7 +66,7 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives ), python_executable="/shared/data-designer/bin/python", host_destination=destination.host_path, - num_partitions=plan.output.partitions, + num_partitions=1, ) script = render_collection_script(plan, collection, destination) @@ -226,7 +226,7 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( ), python_executable="/shared/data-designer/bin/python", host_destination=requested, - num_partitions=plan.output.partitions, + num_partitions=1, ) assert destination.authorized_root == workspace_root @@ -258,7 +258,13 @@ def test_destination_rejects_package_managed_paths( def test_destination_allows_current_run_output_subtree(multi_node_plan: ResolvedSlurmRunPlan) -> None: - requested = (Path(multi_node_plan.output.root) / "collected").as_posix() + requested = ( + Path(multi_node_plan.selected_profile.profile.workspace_root) + / "runs" + / multi_node_plan.run_id + / "output" + / "collected" + ).as_posix() assert CollectionDestinationResolver().resolve(multi_node_plan, requested).host_path == requested diff --git a/packages/data-designer-slurm/tests/planning/test_compiler.py b/packages/data-designer-slurm/tests/planning/test_compiler.py index 3a26bba22..2d58d9908 100644 --- a/packages/data-designer-slurm/tests/planning/test_compiler.py +++ b/packages/data-designer-slurm/tests/planning/test_compiler.py @@ -115,7 +115,6 @@ def test_compiler_resolves_explicit_hostname_and_default_profile_selection( "hostname", "default", } - assert {plan.output.root for plan in plans} == {"/workspace/primary/runs/run-single/output"} 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" @@ -157,7 +156,6 @@ def test_auto_gpu_resolution_is_explicit_and_scheduler_free( ) assert plan.resolved_gpus_per_node == 8 - assert plan.output.root == "/workspace/lab/runs/run-single/output" assert plan.submission.account == "lab" @@ -230,46 +228,6 @@ def test_compiler_rejects_tensor_parallelism_that_does_not_divide_gpu_shape( SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) -@pytest.mark.parametrize( - "output_update", - [ - {"root": "/outside/output"}, - {"root": "/workspace/primary/images/output"}, - {"root": "/workspace/primary/runtime/output"}, - {"root": "/workspace/primary/managed-assets/output"}, - {"root": "/workspace/primary/runs/other-run/output"}, - {"root": "/workspace/primary/runs/run-single/shards"}, - {"partitions": 9}, - ], -) -def test_resolution_rejects_invalid_output_destinations( - authored_run_single: DataDesignerSlurmConfig, - dependency_lock_single: ResolvedDependencyLock, - single_node_plan: ResolvedSlurmRunPlan, - output_update: dict[str, object], -) -> None: - payload = authored_run_single.model_dump(mode="json") - payload["output"].update(output_update) - authored = DataDesignerSlurmConfig.model_validate(payload) - - with pytest.raises(SlurmConfigResolutionError, match="output"): - _resolve_fixture(authored, dependency_lock_single, single_node_plan) - - -def test_resolution_rejects_output_overlapping_explicit_managed_assets( - authored_run_single: DataDesignerSlurmConfig, - dependency_lock_single: ResolvedDependencyLock, - single_node_plan: ResolvedSlurmRunPlan, -) -> None: - payload = authored_run_single.model_dump(mode="json") - payload["invocation"]["input_bindings"]["managed_assets_path"] = "/workspace/primary/custom-assets" - payload["output"]["root"] = "/workspace/primary/custom-assets/output" - authored = DataDesignerSlurmConfig.model_validate(payload) - - with pytest.raises(SlurmConfigResolutionError, match="managed assets"): - _resolve_fixture(authored, dependency_lock_single, single_node_plan) - - @pytest.mark.parametrize( "managed_assets_path", [ @@ -299,45 +257,6 @@ def test_compiler_rejects_direct_effective_input_resolution_bypass( ) -> None: effective = _resolve_fixture(authored_run, dependency_lock, multi_node_plan) - jsonl_output = authored_run.output.model_copy(update={"format": "jsonl"}) - jsonl_authored = authored_run.model_copy(update={"output": jsonl_output}) - with pytest.raises(SlurmPlanCompilationError, match="parquet output"): - SlurmRunCompiler.compile( - effective.model_copy( - update={ - "authored": jsonl_authored, - "output": effective.output.model_copy(update={"format": "jsonl"}), - } - ) - ) - - other_output = "/workspace/primary/runs/other-run/output" - other_authored = authored_run.model_copy( - update={"output": authored_run.output.model_copy(update={"root": other_output})} - ) - with pytest.raises(SlurmPlanCompilationError, match="another package-managed run"): - SlurmRunCompiler.compile( - effective.model_copy( - update={ - "authored": other_authored, - "output": effective.output.model_copy(update={"root": other_output}), - } - ) - ) - - partitioned_authored = authored_run.model_copy( - update={"output": authored_run.output.model_copy(update={"partitions": 101})} - ) - with pytest.raises(SlurmPlanCompilationError, match="requested records"): - SlurmRunCompiler.compile( - effective.model_copy( - update={ - "authored": partitioned_authored, - "output": effective.output.model_copy(update={"partitions": 101}), - } - ) - ) - with pytest.raises(SlurmPlanCompilationError, match="runtime bundle"): SlurmRunCompiler.compile( effective.model_copy( @@ -376,10 +295,6 @@ def test_compiler_rejects_direct_effective_record_drift( update={"invocation": effective.invocation.model_copy(update={"effective_run_config": run_config})} ) ) - with pytest.raises(SlurmPlanCompilationError, match="resolved output"): - SlurmRunCompiler.compile( - effective.model_copy(update={"output": effective.output.model_copy(update={"format": "jsonl"})}) - ) bindings = effective.invocation.effective_input_bindings.model_copy( update={"managed_assets_path": "/workspace/other-assets"} ) @@ -548,7 +463,7 @@ def test_uneven_shards_assign_remainder_to_the_final_task( @pytest.mark.parametrize( - ("builder_update", "output_update", "message"), + ("builder_update", "message"), [ ( { @@ -557,7 +472,6 @@ def test_uneven_shards_assign_remainder_to_the_final_task( "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, } }, - {}, "shuffled seed", ), ( @@ -567,7 +481,6 @@ def test_uneven_shards_assign_remainder_to_the_final_task( "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, } }, - {}, "seed_path", ), ( @@ -578,7 +491,6 @@ def test_uneven_shards_assign_remainder_to_the_final_task( "source": {"path": "/datasets/seed.parquet", "seed_type": "local"}, } }, - {}, "selection strategies", ), ( @@ -592,10 +504,8 @@ def test_uneven_shards_assign_remainder_to_the_final_task( } ] }, - {}, "media output", ), - ({}, {"format": "jsonl"}, "parquet output"), ], ) def test_resolution_rejects_unshardable_big_iron_fields( @@ -603,12 +513,10 @@ def test_resolution_rejects_unshardable_big_iron_fields( dependency_lock: ResolvedDependencyLock, multi_node_plan: ResolvedSlurmRunPlan, builder_update: dict[str, object], - output_update: dict[str, object], message: str, ) -> None: payload = authored_run.model_dump(mode="json") payload["builder"]["inline"]["data_designer"].update(builder_update) - payload["output"].update(output_update) authored = DataDesignerSlurmConfig.model_validate(payload) with pytest.raises(SlurmConfigResolutionError, match=message): @@ -726,13 +634,11 @@ def test_single_shard_allows_non_collectable_big_iron_fields( }, } ) - payload["output"]["format"] = "jsonl" authored = DataDesignerSlurmConfig.model_validate(payload) plan = SlurmRunCompiler.compile(_resolve_fixture(authored, dependency_lock_single, single_node_plan)) assert plan.array_tasks.count == 1 - assert plan.output.format == "jsonl" @pytest.mark.parametrize("sourced", [False, True], ids=["inline", "sourced"]) diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index 659af84b1..bd83a74e1 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -308,7 +308,7 @@ def test_run_service_normalizes_collection_paths(tmp_path: Path, monkeypatch: py ) -def test_run_service_uses_persisted_collection_partitions_when_omitted(tmp_path: Path) -> None: +def test_run_service_defaults_collection_to_one_partition(tmp_path: Path) -> None: backend = Mock(spec=SlurmRunBackend) expected = SlurmCollectionExecution( run_id="run-0001", @@ -316,7 +316,7 @@ def test_run_service_uses_persisted_collection_partitions_when_omitted(tmp_path: state=CollectionState.SUBMITTED, job_id=44, output_path=(tmp_path / "collected").as_posix(), - num_partitions=2, + num_partitions=1, ) backend.collect.return_value = expected service = SlurmRunService(FakeRunPlanningBackend(()), FakeBatchScriptRenderer(()), backend) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 793ec60a8..ba9370e87 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -1124,7 +1124,7 @@ def test_collection_rejects_unmanaged_input_and_destination( 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=2) + 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 @@ -1155,19 +1155,25 @@ def test_collection_returns_active_and_completed_jobs( ) class Coordinator: - def submit(self, *, destination: Path, submitted_at: datetime) -> SimpleNamespace: - del destination, submitted_at + 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)), + ) output_path = tmp_path / "collected" - active = service.collect(tmp_path / "runs/run-wired", destination=output_path) + active = service.collect(tmp_path / "runs/run-wired", destination=output_path, num_partitions=3) completed = service.collect(tmp_path / "runs/run-wired", destination=output_path) 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 def test_auto_gpu_resolution_rejects_mixed_node_shapes( diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 072e43472..6d45099c0 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -17,7 +17,7 @@ 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="1e915d097915a885657be68ca70e2a158b90e003e468462291805d3b6c88517e" +readonly DD_PLAN_SHA256="cc5b1fd8d7f3ad664ee56098a02747fa69888f0979646f67c78796d46674044b" readonly DD_RUN_ROOT="/workspace/primary/runs/run-001" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch index 8294d0cb5..8fe6975ea 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/single_node.sbatch @@ -17,7 +17,7 @@ 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="82229085d7141ea3035512d6cf8c018a16c477ef0d461661c5560a9d8fd7202a" +readonly DD_PLAN_SHA256="8318ae41cfdb1fbbfcea21fe28cd89b09bcf73772f69340a7863d333b2b76a75" readonly DD_RUN_ROOT="/workspace/primary/runs/run-single" readonly DD_ATTEMPT_ORDINAL="0001" diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index 5a664210c..e6371cb38 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -19,12 +19,12 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( single_node_plan, "single_node.sbatch", - expected_fixture_sha256="c312271fe41fa988f439d54f5a9056880875431746c0d7b9bbf470a60f1b2316", + expected_fixture_sha256="cc24e90595495373d47f470afa877333d41dd5ccb85df57da0e9a3d9e3ac1838", ) _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="dbc7ff47a2e9a27f90e2ddffd06df0a2b00be9781635ece189a551bcb27f3943", + expected_fixture_sha256="742d7ebc5412fe266a13e047d472aa2a2950d5276cc015fd93771b890a121d27", ) diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 49f3bdf44..55c37af5a 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -3,7 +3,6 @@ from __future__ import annotations -import csv import errno import hashlib import json @@ -55,11 +54,11 @@ compute_candidate_schema_digest, ) from data_designer.slurm.state.attempt_identity import require_attempt_scheduler_identity -from data_designer.slurm.state.collection_filesystem import derive_collection_staging_directory from data_designer.slurm.state.collection_inputs import CollectionInputResolver 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 @@ -751,7 +750,7 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( assert submitted.state is CollectionState.SUBMITTED assert submitted.scheduler == 5101 with pytest.raises(StateConflictError, match="requested destination"): - coordinator.submit(destination=Path(case.plan.output.root).parent / "other") + 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 @@ -764,10 +763,10 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( assert "DD_CLIENT_IMAGE" not in script assert "--gpus" not in script assert "#SBATCH --gres" not in script - stale_stage = Path(case.plan.output.root).parent / submitted.staging_directory + stale_stage = _output_root(case.plan).parent / submitted.staging_directory stale_stage.mkdir(mode=0o700) (stale_stage / "partial").write_text("incomplete") - unrelated_stage = Path(case.plan.output.root).parent / f".dd-collection-{'f' * 32}.tmp" + 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( @@ -777,9 +776,9 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=11)) - destination = Path(case.plan.output.root) + destination = _output_root(case.plan) assert result.actual_records == case.plan.invocation.authored.num_records - assert len(result.files) == case.plan.output.partitions + 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) ) @@ -883,15 +882,8 @@ def test_collection_prepares_missing_authorized_destination_parents( authored_run: DataDesignerSlurmConfig, multi_node_plan: ResolvedSlurmRunPlan, ) -> None: - workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) - output = multi_node_plan.output.model_copy( - update={"root": (workspace_root / "new" / "nested" / "output").as_posix()} - ) - plan = ResolvedSlurmRunPlan.model_validate_json( - json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) - ) - case = _initialize_run(tmp_path, authored_run, plan) - destination = Path(case.plan.output.root) + 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),)) @@ -900,7 +892,7 @@ def test_collection_prepares_missing_authorized_destination_parents( case.workspace, case.plan.run_id, SlurmCommandClient(runner), - ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + ).submit(destination=destination, submitted_at=case.created_at + timedelta(minutes=10)) assert destination.parent.is_dir() assert not destination.exists() @@ -1080,7 +1072,7 @@ def accept_then_lose_receipt(script: str) -> object: "collection-0001", environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=17)) - assert not Path(case.plan.output.root).exists() + assert not _output_root(case.plan).exists() assert sum(Path(call[0]).name == "sbatch" for call in runner.calls) == 2 @@ -1096,13 +1088,14 @@ def test_collection_definite_submission_failure_settles_and_can_be_retried( coordinator = SlurmCollectionCoordinator(case.workspace, case.plan.run_id, SlurmCommandClient(runner)) with pytest.raises(SlurmStateError, match="cannot submit collection"): - coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + coordinator.submit(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 def test_collection_discards_trailing_journal_interrupted_before_submission( @@ -1165,7 +1158,7 @@ def fail_after_writes(*args: object, **kwargs: object) -> None: environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=11)) - destination = Path(case.plan.output.root) + 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) @@ -1212,7 +1205,7 @@ def replace_output_before_describe( environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=11)) - destination = Path(case.plan.output.root) + 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) @@ -1229,7 +1222,7 @@ def test_collection_refresh_cleans_only_its_exact_stage_after_oom( 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 = Path(case.plan.output.root).parent + parent = _output_root(case.plan).parent exact_stage = parent / submitted.staging_directory exact_stage.mkdir(mode=0o700) (exact_stage / "partial.parquet").write_text("incomplete") @@ -1243,10 +1236,10 @@ def test_collection_refresh_cleans_only_its_exact_stage_after_oom( assert refreshed.state is CollectionState.FAILED assert not exact_stage.exists() assert (unrelated_stage / "active").read_text() == "preserve" - assert not Path(case.plan.output.root).exists() + assert not _output_root(case.plan).exists() -def test_collection_worker_reauthorizes_persisted_partition_intent( +def test_collection_uses_merge_partition_count( tmp_path: Path, authored_run: DataDesignerSlurmConfig, multi_node_plan: ResolvedSlurmRunPlan, @@ -1254,33 +1247,48 @@ def test_collection_worker_reauthorizes_persisted_partition_intent( case = _initialize_run(tmp_path, authored_run, multi_node_plan) _publish_all_winners(case) runner = FakeSlurmRunner(jobs=(FakeSlurmJob(5101),)) - submitted = SlurmCollectionCoordinator( + coordinator = SlurmCollectionCoordinator( case.workspace, case.plan.run_id, SlurmCommandClient(runner), - ).submit(submitted_at=case.created_at + timedelta(minutes=10)) + ) + 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) - tampered = plan.model_copy(update={"num_partitions": plan.num_partitions + 1}) - storage.get_plan_path(plan.collection_id).write_text(tampered.serialize_json()) - storage.replace_status( - submitted.model_copy( - update={ - "collection_plan": storage.get_plan_reference(tampered), - "staging_directory": derive_collection_staging_directory(tampered), - } - ) - ) + assert plan.num_partitions == 3 + assert coordinator.submit() == submitted with pytest.raises(StateConflictError, match="partition count"): - SlurmCollectionWorker( - case.workspace, - case.plan.run_id, - submitted.collection_id, - environment={"SLURM_JOB_ID": "5101"}, - ).run(completed_at=case.created_at + timedelta(minutes=11)) + 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 requested 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 not Path(case.plan.output.root).exists() + 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( @@ -1299,7 +1307,7 @@ def test_succeeded_collection_rejects_modified_partition_bytes( submitted.collection_id, environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=11)) - output_path = Path(case.plan.output.root) / result.files[0].relative_path + 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"): @@ -1322,7 +1330,7 @@ def test_succeeded_collection_rejects_same_size_partition_mutation( submitted.collection_id, environment={"SLURM_JOB_ID": "5101"}, ).run(completed_at=case.created_at + timedelta(minutes=11)) - output_path = Path(case.plan.output.root) / result.files[0].relative_path + 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) @@ -1398,7 +1406,7 @@ def test_collection_refuses_destination_collision_before_submission( ) -> None: case = _initialize_run(tmp_path, authored_run, multi_node_plan) _publish_all_winners(case) - destination = Path(case.plan.output.root) + destination = _output_root(case.plan) destination.mkdir(mode=0o700) marker = destination / "existing.txt" marker.write_text("preserve") @@ -1429,7 +1437,7 @@ def test_collection_refuses_atomic_publication_collision_without_replacement( case.plan.run_id, SlurmCommandClient(runner), ).submit(submitted_at=case.created_at + timedelta(minutes=10)) - destination = Path(case.plan.output.root) + destination = _output_root(case.plan) original_rename = collection_filesystem._rename_without_overwrite def collide( @@ -1523,7 +1531,7 @@ def test_collection_worker_recovers_interrupted_private_reservation( case.plan.run_id, SlurmCommandClient(runner), ).submit(submitted_at=case.created_at + timedelta(minutes=10)) - destination = Path(case.plan.output.root) + 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: @@ -1604,7 +1612,7 @@ def test_collection_recovery_preserves_unrelated_empty_destination( case.plan.run_id, SlurmCommandClient(runner), ).submit(submitted_at=case.created_at + timedelta(minutes=10)) - destination = Path(case.plan.output.root) + destination = _output_root(case.plan) unrelated = destination.parent / "unrelated-reservation" unrelated.mkdir(mode=0o700) destination.mkdir(mode=0o700) @@ -1686,20 +1694,15 @@ def test_collection_detects_destination_parent_replacement_before_success( multi_node_plan: ResolvedSlurmRunPlan, monkeypatch: pytest.MonkeyPatch, ) -> None: - workspace_root = Path(multi_node_plan.selected_profile.profile.workspace_root) - output = multi_node_plan.output.model_copy(update={"root": (workspace_root / "exports" / "output").as_posix()}) - plan = ResolvedSlurmRunPlan.model_validate_json( - json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) - ) - case = _initialize_run(tmp_path, authored_run, plan) + 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 = Path(case.plan.output.root) + ).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 @@ -1772,7 +1775,7 @@ def publish_then_fail( 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 Path(case.plan.output.root).is_dir() + assert _output_root(case.plan).is_dir() def test_collection_worker_rejects_login_node_execution( @@ -1846,50 +1849,8 @@ def test_collection_worker_validates_location_before_constructing_state(tmp_path SlurmCollectionWorker(tmp_path, "../run", "collection-0001") -@pytest.mark.parametrize("output_format", ["csv", "jsonl"]) -def test_collection_exports_non_parquet_formats_in_deterministic_partitions( - tmp_path: Path, - authored_run: DataDesignerSlurmConfig, - multi_node_plan: ResolvedSlurmRunPlan, - output_format: str, -) -> None: - output = multi_node_plan.output.model_copy(update={"format": output_format, "partitions": 3}) - plan = ResolvedSlurmRunPlan.model_validate_json( - json.dumps(multi_node_plan.model_copy(update={"output": output}).model_dump(mode="json")) - ) - case = _initialize_run(tmp_path, authored_run, 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)) - - 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 len(result.files) == 3 - requested_records = case.plan.invocation.authored.num_records - floor_count = requested_records // 3 - assert tuple(file.record_count for file in result.files) == ( - floor_count, - floor_count, - requested_records - 2 * floor_count, - ) - values: list[int] = [] - for output_file in result.files: - path = Path(case.plan.output.root) / output_file.relative_path - if output_format == "csv": - with path.open(newline="") as source: - values.extend(int(row["record_id"]) for row in csv.DictReader(source)) - else: - values.extend(int(json.loads(line)["record_id"]) for line in path.read_text().splitlines()) - assert values == list(range(requested_records)) +def _output_root(plan: ResolvedSlurmRunPlan) -> Path: + return Path(plan.selected_profile.profile.workspace_root) / "runs" / plan.run_id / "output" def _initialize_run( diff --git a/packages/data-designer-slurm/tests/state/test_state_golden_records.py b/packages/data-designer-slurm/tests/state/test_state_golden_records.py index 785a5b713..05a8561d3 100644 --- a/packages/data-designer-slurm/tests/state/test_state_golden_records.py +++ b/packages/data-designer-slurm/tests/state/test_state_golden_records.py @@ -44,7 +44,7 @@ CONTRACT_GOLDEN_DIRECTORY / "authored_run_single.json", CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json", ) -_COMPATIBILITY_FIXTURE_DIGEST = "4128e04c4c4131de2029ec32bd9b9b337b2f695f42b938a1cb7d6cfe07dd4bd1" +_COMPATIBILITY_FIXTURE_DIGEST = "18dbde1e5a6d883d90d6392487b84750a5dd8133f44d75d425d240ba7feb7b6c" @pytest.mark.parametrize(("filename", "model"), GOLDEN_MODELS) diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 470aa1383..8d04c8e28 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1526,37 +1526,6 @@ def test_result_publication_rejects_a_successful_attempt_bound_to_another_candid ) -@pytest.mark.parametrize("export_format", ("jsonl", "csv")) -def test_finalization_keeps_candidate_parquet_separate_from_export_format( - tmp_path: Path, - authored_run_single: DataDesignerSlurmConfig, - single_node_plan: ResolvedSlurmRunPlan, - export_format: str, -) -> None: - authored = authored_run_single.model_copy( - update={"output": authored_run_single.output.model_copy(update={"format": export_format})} - ) - plan = single_node_plan.model_copy( - update={ - "authored_config": single_node_plan.authored_config.model_copy( - update={"sha256": authored.compute_sha256()} - ), - "output": single_node_plan.output.model_copy(update={"format": export_format}), - } - ) - case = _initialized_case(tmp_path, authored, 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.attempt_id == finalization.attempt.attempt_id - assert finalization.candidate.files[0].relative_path.endswith(".parquet") - - def test_candidate_verifier_leases_manifest_bounded_files_through_publication( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From f4b342408955cb0ddb374ec9aa8a8f2a106cb6e7 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 14 Oct 2026 10:40:29 -0600 Subject: [PATCH 48/62] fix(slurm): reuse final batch files for candidates (#1023) Avoid writing a duplicate Parquet export when the engine batches can be used directly as a bounded candidate. Keep the export path for schema drift or too many batch files. Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/client/execution.py | 78 +++++++++++------ .../tests/client/conftest.py | 27 +++++- .../tests/client/test_worker.py | 86 ++++++++++++++----- 3 files changed, 143 insertions(+), 48 deletions(-) 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 index 0cc43a577..dd3db08ff 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -52,13 +52,25 @@ 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 CandidateOutcome, CandidateOutputFile, CandidateOutputManifest +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 @@ -67,6 +79,10 @@ class CreationResults(Protocol): 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.""" @@ -542,8 +558,8 @@ def _finalize( results: CreationResults, creation: _ValidatedCreation, ) -> ClientResult: - dataset_path, exported_path = self._place_dataset(creation, results) - candidate = self._build_candidate_manifest(context, prepared, creation, dataset_path, exported_path) + 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 @@ -600,20 +616,32 @@ def _validate_creation_result( def _place_dataset( creation: _ValidatedCreation, results: CreationResults, - ) -> tuple[Path, Path | None]: + ) -> tuple[Path, tuple[Path, ...]]: dataset_path = creation.dataset_path - exported_path: Path | None = None - if creation.actual_records: - exported_path = results.export(dataset_path / "part-00000.parquet", format="parquet") + 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) - if exported_path is not None: - exported_path = creation.expected_path / exported_path.relative_to(dataset_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, exported_path + 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, @@ -621,24 +649,26 @@ def _build_candidate_manifest( prepared: PreparedClientEnvironment, creation: _ValidatedCreation, dataset_path: Path, - exported_path: Path | None, + output_paths: tuple[Path, ...], ) -> CandidateOutputManifest: files: tuple[CandidateOutputFile, ...] = () schema_digest = context.plan.builder.content_sha256 - if exported_path is not None: - metadata = lazy.pq.read_metadata(exported_path) - if metadata.num_rows != creation.actual_records: - raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "exported record count differs") - schema_digest = hashlib.sha256( - lazy.pq.read_schema(exported_path).remove_metadata().serialize().to_pybytes() - ).hexdigest() - files = ( + 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=exported_path.relative_to(dataset_path).as_posix(), - sha256=compute_file_sha256(exported_path, missing_code=ClientErrorCode.OUTPUT_INVALID), - byte_size=exported_path.stat().st_size, - record_count=creation.actual_records, - ), + 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( diff --git a/packages/data-designer-slurm/tests/client/conftest.py b/packages/data-designer-slurm/tests/client/conftest.py index 4f49bb6e8..6ce539561 100644 --- a/packages/data-designer-slurm/tests/client/conftest.py +++ b/packages/data-designer-slurm/tests/client/conftest.py @@ -8,6 +8,7 @@ import sys from dataclasses import dataclass from pathlib import Path +from types import SimpleNamespace import pyarrow as pa import pyarrow.parquet as pq @@ -50,6 +51,7 @@ def __init__( 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" @@ -65,12 +67,18 @@ def __init__( 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 @@ -97,9 +105,22 @@ def create( 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) - batch_path = dataset_path / "batch.parquet" - pq.write_table(pa.table({"value": tuple(range(actual))}), batch_path) - on_batch_complete(batch_path) + 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) diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index dd25a3c92..045ed7f5f 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -12,10 +12,9 @@ from types import SimpleNamespace from unittest.mock import Mock -import pyarrow as pa import pyarrow.parquet as pq import pytest -from conftest import ClientWorkerCase, FakeCreationResults, FakeDataDesigner +from conftest import ClientWorkerCase, FakeDataDesigner import data_designer.slurm.client.worker as worker_module from data_designer.config import ResumeMode @@ -33,6 +32,7 @@ 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: @@ -180,36 +180,81 @@ def test_run_persists_semantic_result_and_candidate( 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_run_normalizes_parquet_schema_metadata( +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, - monkeypatch: pytest.MonkeyPatch, + batch_count: int, ) -> None: - def export_with_metadata(results: FakeCreationResults, path: Path, *, format: str) -> Path: - assert format == "parquet" - schema = pa.schema( - [("value", pa.int64())], - metadata={b"producer": b"data-designer"}, - ) - table = pa.Table.from_arrays( - [pa.array(range(results.actual_num_records))], - schema=schema, + 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", ) - pq.write_table(table, path) - return path + 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() - monkeypatch.setattr(FakeCreationResults, "export", export_with_metadata) - worker = ClientWorker(data_designer_factory=FakeDataDesigner) + +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, @@ -220,8 +265,7 @@ def export_with_metadata(results: FakeCreationResults, path: Path, *, format: st 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/part-00000.parquet") - assert candidate.dataset_schema_digest == compute_candidate_schema_digest(schema) + assert tuple(file.relative_path for file in candidate.files) == ("part-00000.parquet",) @pytest.mark.parametrize( @@ -446,7 +490,7 @@ def test_if_possible_interruption_preserves_workspace_for_retry(client_worker_ca ) resume_path = Path(plan.shards[0].resume_workspace.path) - assert (resume_path / "batch.parquet").is_file() + 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" From 568ebb4ced4eb513e57bc046ded13305e78b5177 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 14 Oct 2026 10:46:59 -0600 Subject: [PATCH 49/62] fix(slurm): isolate Harmony tokenizer caches (#1024) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Give each node-local vLLM replica its own tokenizer cache so concurrent startup cannot read or remove another replica’s partially downloaded vocabulary. Preserve an explicitly configured cache and clean generated caches after the replicas stop. Signed-off-by: Nabin Mulepati --- .../slurm/runtime/node_worker.py | 16 +++++++- .../tests/runtime/test_node_worker.py | 39 ++++++++++++++++++- 2 files changed, 52 insertions(+), 3 deletions(-) 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 index d8e8a0eff..dc282cce6 100644 --- 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 @@ -14,6 +14,7 @@ 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, @@ -41,6 +42,7 @@ class NodeProcessSupervisor: 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 @@ -83,15 +85,25 @@ def cleanup(self) -> None: child.wait(timeout=_TERMINATION_GRACE_SECONDS) except (OSError, subprocess.SubprocessError): pass - self._cleanup_complete = all(child.poll() is not None for child in self._children) - if not self._cleanup_complete: + 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, diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py index 3809ef195..b5fb57b13 100644 --- a/packages/data-designer-slurm/tests/runtime/test_node_worker.py +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -67,10 +67,14 @@ def test_partial_startup_failure_cleans_already_started_lane(monkeypatch: pytest process = _FakeProcess(pid=41) calls = 0 signals: list[int] = [] + cache_paths: list[Path] = [] def start(*arguments: object, **keywords: object) -> _FakeProcess: nonlocal calls - del arguments, keywords + 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") @@ -99,16 +103,22 @@ def terminate(child: _FakeProcess, selected: int) -> None: 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: @@ -129,6 +139,33 @@ def terminate(child: _FakeProcess, selected: int) -> None: 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: From d856b6ea4c311eb35872d39879008ddf9618af46 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 14 Oct 2026 11:04:59 -0600 Subject: [PATCH 50/62] feat(slurm): merge runs by ID with default output (#1025) * feat(slurm): merge runs by ID with default output Signed-off-by: Nabin Mulepati * refactor(slurm): name default merge directory Signed-off-by: Nabin Mulepati * test(slurm): align merge destination coverage Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/cli-reference.mdx | 6 +-- .../latest/pages/slurm/getting-started.mdx | 10 ++--- fern/versions/latest/pages/slurm/overview.mdx | 2 +- .../pages/slurm/retry-and-collection.mdx | 23 +++++++---- .../latest/pages/slurm/run-configuration.mdx | 2 +- .../src/data_designer/slurm/cli.py | 15 +++++-- .../slurm/services/retry_collection.py | 11 +++-- .../src/data_designer/slurm/services/run.py | 28 ++++++++----- .../data_designer/slurm/services/wiring.py | 4 +- .../data_designer/slurm/state/collection.py | 11 ++++- .../data_designer/slurm/state/destinations.py | 4 +- .../tests/launcher/test_collection.py | 5 ++- .../tests/services/test_services.py | 17 ++++++++ .../tests/services/test_wiring.py | 9 +++-- .../tests/state/test_retry_collection.py | 16 +++++++- .../data-designer-slurm/tests/test_cli.py | 40 +++++++++++++++++-- 16 files changed, 149 insertions(+), 54 deletions(-) diff --git a/fern/versions/latest/pages/slurm/cli-reference.mdx b/fern/versions/latest/pages/slurm/cli-reference.mdx index ec159b271..ba2583ea6 100644 --- a/fern/versions/latest/pages/slurm/cli-reference.mdx +++ b/fern/versions/latest/pages/slurm/cli-reference.mdx @@ -16,7 +16,7 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | `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` | Submit a zero-GPU winner collection job for a managed run directory. | +| `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. | @@ -29,8 +29,8 @@ All commands accept `--profile-file PATH` and `--cluster NAME` where cluster sel | `--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 the managed run directory containing shard winners. | -| `--output-path PATH` | `merge` | Select an authorized compute-visible collection directory. | +| `--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. | diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index 228faee00..adda90328 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -173,16 +173,12 @@ Each successful shard publishes a canonical winner at: 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 an explicit destination, submit a managed collection job: +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 \ - --input-path /workspace/alice/data-designer-slurm/runs/ \ - --output-path /workspace/alice/exports/greetings \ - --num-partitions 1 \ - --cluster primary +data-designer slurm merge ``` -`--num-partitions` controls the collected Parquet file count and defaults to 1 for a new collection. +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/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index 742e6ed5e..e6f64deac 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -28,7 +28,7 @@ Run Slurm commands from a site-approved submit environment with the cluster CLI ## 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. Collected output goes to the destination selected with `slurm merge --output-path`. +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. diff --git a/fern/versions/latest/pages/slurm/retry-and-collection.mdx b/fern/versions/latest/pages/slurm/retry-and-collection.mdx index 3cc0f9517..ee7c60d08 100644 --- a/fern/versions/latest/pages/slurm/retry-and-collection.mdx +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -34,7 +34,19 @@ Retries create new attempts. The original attempt records remain durable, and ea ## Collect winners -Successful shards publish Parquet winners under the run directory. Use `merge` to submit a zero-GPU CPU job that collects them to an explicit compute-visible destination: +Successful shards publish Parquet winners under the run directory. Use `merge` to submit a zero-GPU CPU job that collects them under `/runs//merge-output`: + +```bash +data-designer slurm merge +``` + +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 \ @@ -42,15 +54,12 @@ data-designer slurm merge \ --output-path /workspace/alice/exports/final-dataset ``` -`--input-path` must be exactly a managed run directory under the selected profile workspace. The output path must be absolute, writable, compute-visible, and authorized by the profile mounts. +`--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 run's record count: ```bash -data-designer slurm merge \ - --input-path /workspace/alice/data-designer-slurm/runs/ \ - --output-path /workspace/alice/exports/final-dataset \ - --num-partitions 20 +data-designer slurm merge --num-partitions 20 ``` -Repeating the same merge command returns or reuses managed collection state when safe. If you omit `--num-partitions` on a repeat call, the existing collection keeps its count. Changing the destination or explicitly requesting a different partition count for an active collection is a conflict. +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 index 3731619ef..a2019b3dd 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -92,7 +92,7 @@ array_tasks: max_concurrent: 4 ``` -`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. Array task count controls execution shards; it does not set the number of collected Parquet files. Use `slurm merge --num-partitions 20` if you want 20 output files. +`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. Array task count controls execution shards; it does not set the number of collected Parquet files. Use `slurm merge RUN_ID --num-partitions 20` if you want 20 output files. 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`. diff --git a/packages/data-designer-slurm/src/data_designer/slurm/cli.py b/packages/data-designer-slurm/src/data_designer/slurm/cli.py index bf5d5cd9a..60546f5ae 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli.py @@ -211,8 +211,9 @@ def retry_command( @app.command("merge") def merge_command( - input_path: Path = typer.Option(..., "--input-path", file_okay=False), - output_path: Path = typer.Option(..., "--output-path", file_okay=False), + 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)" ), @@ -221,11 +222,19 @@ def merge_command( ) -> 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( - input_path, + run_id if run_id is not None else input_path, destination=output_path, num_partitions=num_partitions, ), 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 index ead9d14fc..85c12a932 100644 --- 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 @@ -128,14 +128,17 @@ def retry( def collect( self, - input_path: Path, + input_path: Identifier | Path, *, - destination: Path, + destination: Path | None, num_partitions: int | None, ) -> SlurmCollectionExecution: """Submit or recover one winner-driven collection.""" operation = SlurmServiceOperation.COLLECT_RUN - run_id = self._resolve_run_input_path(input_path, operation) + 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: @@ -167,7 +170,7 @@ def collect( collection_id=status.collection_id, state=status.state, job_id=status.scheduler, - output_path=destination.as_posix(), + output_path=collection_plan.host_destination, num_partitions=collection_plan.num_partitions, ) except SlurmServiceError: 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 index e241b3f4f..4898c2820 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/run.py @@ -84,9 +84,9 @@ def retry( def collect( self, - input_path: Path, + input_path: Identifier | Path, *, - destination: Path, + destination: Path | None, num_partitions: int | None, ) -> SlurmCollectionExecution: """Submit or recover one winner-driven collection.""" @@ -261,17 +261,23 @@ def collect( self, input_path: str | Path, *, - destination: str | Path, + destination: str | Path | None = None, num_partitions: int | None = None, ) -> SlurmCollectionExecution: - """Submit or recover collection for one managed run directory.""" + """Submit or recover collection by run ID or managed run directory.""" operation = SlurmServiceOperation.COLLECT_RUN - if not isinstance(input_path, str | Path) or not isinstance(destination, str | Path): - raise _make_invalid_request_error(operation, "input_path and destination must be paths") + 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 = Path(input_path).expanduser().resolve() - normalized_destination = Path(destination).expanduser().resolve() + 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( @@ -320,9 +326,9 @@ def _retry_run( def _collect_run( backend: SlurmRunBackend, - input_path: Path, + input_path: Identifier | Path, *, - destination: Path, + destination: Path | None, num_partitions: int | None, ) -> SlurmCollectionExecution: result = backend.collect( @@ -332,7 +338,7 @@ def _collect_run( ) if not isinstance(result, SlurmCollectionExecution): raise TypeError("run backend returned an invalid collection result") - if result.output_path != destination.as_posix() or ( + 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") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 8a87b4a69..21cdd184b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -737,9 +737,9 @@ def retry( def collect( self, - input_path: Path, + input_path: Identifier | Path, *, - destination: Path, + destination: Path | None, num_partitions: int | None, ) -> SlurmCollectionExecution: return self._retry_collection.collect( 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 index a44d72d00..7f8d32e52 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -115,7 +115,10 @@ def submit( Path(resolved_destination.authorized_root), ) run, resolved_plan, _ = self._reader.load_context() - resolved_destination = self._destinations.resolve(resolved_plan, destination) + 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 @@ -303,7 +306,11 @@ def _validate_existing_destination( 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 not None and Path(requested_destination).as_posix() != resolved.host_path: + 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: 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 index 5ad25342f..d06bfefcd 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/destinations.py @@ -34,7 +34,7 @@ def resolve( """Return one authorized host output path.""" workspace_root = plan.selected_profile.profile.workspace_root raw_destination = ( - posixpath.join(workspace_root, "runs", plan.run_id, "output") + posixpath.join(workspace_root, "runs", plan.run_id, "merge-output") if requested_destination is None else Path(requested_destination).as_posix() ) @@ -78,7 +78,7 @@ def _validate_destination_overlap(plan: ResolvedSlurmRunPlan, host_path: str) -> 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, "output") + 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) ): diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 231007815..dd4d574a9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -243,6 +243,7 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( ("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( @@ -257,12 +258,12 @@ def test_destination_rejects_package_managed_paths( CollectionDestinationResolver().resolve(multi_node_plan, requested) -def test_destination_allows_current_run_output_subtree(multi_node_plan: ResolvedSlurmRunPlan) -> None: +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 - / "output" + / "merge-output" / "collected" ).as_posix() diff --git a/packages/data-designer-slurm/tests/services/test_services.py b/packages/data-designer-slurm/tests/services/test_services.py index bd83a74e1..9f6c5e972 100644 --- a/packages/data-designer-slurm/tests/services/test_services.py +++ b/packages/data-designer-slurm/tests/services/test_services.py @@ -329,6 +329,23 @@ def test_run_service_defaults_collection_to_one_partition(tmp_path: Path) -> Non ) +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)) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index ba9370e87..62fbf1313 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -1163,17 +1163,20 @@ def submit(self, *, destination: Path, num_partitions: int | None, submitted_at: monkeypatch.setattr( retry_collection_module, "CollectionStorage", - lambda *_: SimpleNamespace(read_plan=lambda _: SimpleNamespace(num_partitions=3)), + lambda *_: SimpleNamespace( + read_plan=lambda _: SimpleNamespace(num_partitions=3, host_destination=output_path.as_posix()) + ), ) output_path = tmp_path / "collected" - active = service.collect(tmp_path / "runs/run-wired", destination=output_path, num_partitions=3) - completed = service.collect(tmp_path / "runs/run-wired", destination=output_path) + 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( diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 55c37af5a..18aa3977b 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -896,6 +896,14 @@ def test_collection_prepares_missing_authorized_destination_parents( 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, @@ -1086,9 +1094,12 @@ def test_collection_definite_submission_failure_settles_and_can_be_retried( 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(num_partitions=3, submitted_at=case.created_at + timedelta(minutes=10)) + 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 @@ -1096,6 +1107,7 @@ def test_collection_definite_submission_failure_settles_and_can_be_retried( 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( @@ -1850,7 +1862,7 @@ def test_collection_worker_validates_location_before_constructing_state(tmp_path def _output_root(plan: ResolvedSlurmRunPlan) -> Path: - return Path(plan.selected_profile.profile.workspace_root) / "runs" / plan.run_id / "output" + return Path(plan.selected_profile.profile.workspace_root) / "runs" / plan.run_id / "merge-output" def _initialize_run( diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index cc7e3c6a4..b6f8a8ad5 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -45,7 +45,7 @@ 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[Path, Path, int | None]] = [] + self.collection_calls: list[tuple[str | Path, Path | None, int | None]] = [] def execute( self, @@ -86,9 +86,9 @@ def retry( def collect( self, - input_path: Path, + input_path: str | Path, *, - destination: Path, + destination: Path | None, num_partitions: int | None, ) -> SlurmCollectionExecution: self.collection_calls.append((input_path, destination, num_partitions)) @@ -97,7 +97,7 @@ def collect( collection_id="collection-0001", state=CollectionState.SUBMITTED, job_id=43, - output_path=destination.resolve().as_posix(), + 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, ) @@ -641,6 +641,38 @@ def test_merge_emits_collection_job_and_forwards_paths(tmp_path: Path, monkeypat 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"), [ From 07ba2d0bb74341a320c1acedf939090833f9ba95 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 14 Oct 2026 13:40:46 -0600 Subject: [PATCH 51/62] feat(slurm): accept nonempty partial shard outputs (#1026) * fix(slurm): read only committed attempt results Ignore client-owned partial result files when an attempt has no bound candidate. Failed shards can then be reconciled and reported without treating their incomplete output as a finalized result. Signed-off-by: Nabin Mulepati * accept nonempty partial Slurm shards Signed-off-by: Nabin Mulepati * fix(slurm): analyze nonempty partial benchmarks Use actual winning rows for benchmark throughput and capacity estimates. Update the benchmark result contract and Slurm docs for partial winners. Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../pages/slurm/retry-and-collection.mdx | 4 +- .../latest/pages/slurm/troubleshooting.mdx | 2 +- .../data_designer/slurm/benchmark/analysis.py | 7 -- .../data_designer/slurm/benchmark/observer.py | 3 +- .../data_designer/slurm/benchmark/records.py | 4 +- .../data_designer/slurm/runtime/controller.py | 10 +-- .../data_designer/slurm/runtime/entrypoint.py | 4 +- .../data_designer/slurm/runtime/records.py | 8 +- .../data_designer/slurm/state/collection.py | 7 +- .../slurm/state/collection_validation.py | 8 +- .../slurm/state/collection_worker.py | 4 +- .../data_designer/slurm/state/finalization.py | 2 +- .../slurm/state/finalization_validation.py | 18 ++-- .../src/data_designer/slurm/state/outputs.py | 6 +- .../src/data_designer/slurm/state/reader.py | 6 +- .../src/data_designer/slurm/state/results.py | 2 +- .../tests/benchmark/test_workflow.py | 44 ++++++++++ .../tests/contracts/test_shared_records.py | 3 +- .../test_integration_validation.py | 6 +- .../tests/runtime/test_controller.py | 41 +++++++-- .../tests/runtime/test_entrypoint.py | 2 +- .../tests/state/test_observer.py | 84 +++++++++++++++++++ .../tests/state/test_records.py | 4 +- .../tests/state/test_retry_collection.py | 54 +++++++++--- .../tests/state/test_store.py | 2 +- .../tests/state/test_validation.py | 5 +- 26 files changed, 262 insertions(+), 78 deletions(-) diff --git a/fern/versions/latest/pages/slurm/retry-and-collection.mdx b/fern/versions/latest/pages/slurm/retry-and-collection.mdx index ee7c60d08..894f13850 100644 --- a/fern/versions/latest/pages/slurm/retry-and-collection.mdx +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -34,7 +34,7 @@ Retries create new attempts. The original attempt records remain durable, and ea ## Collect winners -Successful shards publish Parquet winners under the run directory. Use `merge` to submit a zero-GPU CPU job that collects them under `/runs//merge-output`: +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 @@ -56,7 +56,7 @@ data-designer slurm merge \ `--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 run's record count: +`--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 diff --git a/fern/versions/latest/pages/slurm/troubleshooting.mdx b/fern/versions/latest/pages/slurm/troubleshooting.mdx index 99d30e010..e1a7476ca 100644 --- a/fern/versions/latest/pages/slurm/troubleshooting.mdx +++ b/fern/versions/latest/pages/slurm/troubleshooting.mdx @@ -91,7 +91,7 @@ Do not edit persisted manifests. Conflicts preserve scheduler and provenance evi ## 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 short result cannot become a shard winner. Retry missing shards, then use `merge` with the original run directory. +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 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 index 28eaec966..9b9dc0d33 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -187,13 +187,6 @@ def _analyze_case( boot_seconds = measurements.boot_seconds generation_seconds = measurements.generation_seconds wall_seconds = measurements.wall_seconds - requested_records = observed.authored_config.invocation.num_records - if actual_records != requested_records: - return _update_result( - result, - outcome=BenchmarkOutcome.INCOMPLETE, - actual_records=actual_records, - ) gpus_per_job = result.gpus_per_job if gpus_per_job is None: raise ValueError("resolved benchmark case has no GPU count") 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 index 3ae540698..0ed6b35db 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/observer.py @@ -191,7 +191,8 @@ def _measure_records(records: tuple[_MeasurementRecord, ...]) -> BenchmarkMeasur return None client_result, candidate = result if ( - client_result.outcome is not ClientOutcome.COMPLETE + 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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py index 40d98b232..c626266ca 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/records.py @@ -100,8 +100,8 @@ def validate_metrics(self) -> BenchmarkCaseResult: 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 != self.requested_records: - raise ValueError("successful benchmark cases require the requested record count") + 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: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 625050f2e..c1b987f9c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -21,7 +21,7 @@ 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_complete_client_candidate +from data_designer.slurm.runtime.records import load_nonempty_client_candidate from data_designer.slurm.runtime.steps import ( ClientStepBuilder, build_endpoint_steps, @@ -125,7 +125,7 @@ class _RuntimeTopology: class OneNodeAllocationController: - """Execute one planned shard attempt and finalize only complete candidates.""" + """Execute one planned shard attempt and finalize nonempty candidates.""" def __init__( self, @@ -362,7 +362,7 @@ def _run_client( 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_complete_client_candidate(self._context, self._attempt) + 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, @@ -370,10 +370,10 @@ def _run_client( ) 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 - the record contract requires this for complete results + if candidate_reference is None: # pragma: no cover - nonempty results require this reference raise SlurmRuntimeError( SlurmRuntimeErrorCode.FINALIZATION_FAILED, - "complete client result has no candidate reference", + "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) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index 67d74fd15..c6672b8c6 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -25,7 +25,7 @@ SystemAllocationPreflight, validate_allocation_layout, ) -from data_designer.slurm.runtime.records import load_complete_client_candidate +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, @@ -185,7 +185,7 @@ def _client( ) if return_code != 0: raise SlurmRuntimeError(SlurmRuntimeErrorCode.CLIENT_FAILED, "client generation failed") - client_result, candidate = load_complete_client_candidate( + client_result, candidate = load_nonempty_client_candidate( context, context.attempt, attempt_directory=arguments.attempt_dir, 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 index e61f48017..53119456e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/records.py @@ -24,16 +24,16 @@ _RecordT = TypeVar("_RecordT", bound=ContractRecord) -def load_complete_client_candidate( +def load_nonempty_client_candidate( context: AllocationContext, attempt: AttemptManifest, *, attempt_directory: Path | None = None, ) -> tuple[ClientResult, CandidateOutputManifest]: - """Load a complete semantic client result and its digest-bound candidate.""" + """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 is not ClientOutcome.COMPLETE: + 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}", @@ -72,4 +72,4 @@ def _read_record(directory: Path, name: str, record_type: type[_RecordT]) -> _Re ) from error -__all__ = ["load_complete_client_candidate"] +__all__ = ["load_nonempty_client_candidate"] 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 index 7f8d32e52..707b88456 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -122,8 +122,6 @@ def submit( effective_partitions = num_partitions if effective_partitions is None: effective_partitions = current_plan.num_partitions if current is not None else 1 - if effective_partitions > resolved_plan.invocation.authored.num_records: - raise StateConflictError("collection partitions must not exceed requested records") collection_plan = CollectionPlan( schema_version=1, collection_id=self._collections.get_next_collection_id(), @@ -340,13 +338,12 @@ def _load_valid_result( plan: CollectionPlan, status: CollectionStatus | None = None, ) -> CollectionResult: - self._inputs.resolve(plan) + _, candidates = self._inputs.resolve(plan) result = self._collections.read_result(plan) - resolved_plan = self._reader.load_resolved_plan() validated = validate_collection_result( plan, result, - expected_records=resolved_plan.invocation.authored.num_records, + 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") 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 index 3bf19dd91..2c63ac4d0 100644 --- 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 @@ -19,12 +19,12 @@ def validate_collection_inputs( collection_plan: CollectionPlan, candidates: tuple[CandidateOutputManifest, ...], ) -> tuple[CandidateOutputManifest, ...]: - """Require one compatible, complete candidate for every planned shard.""" + """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 not complete") + _require(all(candidate.winner_eligible for candidate in candidates), "collection candidate is empty") schema_digests = {candidate.dataset_schema_digest for candidate in candidates} _require(len(schema_digests) == 1, "collection candidates have incompatible schemas") @@ -33,10 +33,8 @@ def validate_collection_inputs( "collection candidate provenance does not match the resolved plan", ) - expected_records = resolved_plan.invocation.authored.num_records - _require(collection_plan.num_partitions <= expected_records, "collection partitions exceed requested records") actual_records = sum(candidate.actual_records for candidate in candidates) - _require(actual_records == expected_records, "collection candidate rows do not match requested run rows") + _require(collection_plan.num_partitions <= actual_records, "collection partitions exceed available records") return candidates 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 index 8680166e0..0fc4cd39e 100644 --- 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 @@ -139,11 +139,11 @@ def _load_valid_result( status: CollectionStatus | None = None, ) -> CollectionResult: result = self._collections.read_result(plan) - resolved_plan = self._reader.load_resolved_plan() + _, candidates = self._inputs.resolve(plan) validated = validate_collection_result( plan, result, - expected_records=resolved_plan.invocation.authored.num_records, + 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") 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 index 38fb75138..3d18e035d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/finalization.py @@ -357,7 +357,7 @@ def _build_winner( ) -> ShardWinner: reference = client_result.candidate_output_manifest if reference is None: - raise StateContractError("complete client result has no candidate reference") + raise StateContractError("nonempty client result has no candidate reference") return ShardWinner( schema_version=1, run_id=run.run_id, 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 index c5b7b6bfa..e47cc1f57 100644 --- 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 @@ -69,8 +69,11 @@ def _validate_result_outcomes( attempt.state in {AttemptLifecycleState.RUNNING, AttemptLifecycleState.SUCCEEDED}, "attempt results require a running or successful attempt", ) - _require(client_result.outcome is ClientOutcome.COMPLETE, "only complete client results may be finalized") - _require(candidate.winner_eligible, "only complete candidate outputs may be finalized") + _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: @@ -124,9 +127,10 @@ def _validate_output_contract( "client and candidate requested records must match the planned shard", ) _require( - client_result.actual_records == candidate.actual_records == planned_shard.requested_records, - "client and candidate actual records must complete the planned shard", + 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), @@ -170,7 +174,7 @@ def _validate_result_reference( "output-manifest.json", ) reference = client_result.candidate_output_manifest - _require(reference is not None, "complete client result has no candidate manifest reference") + _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") @@ -181,14 +185,14 @@ def _validate_winner_reference( winner: ShardWinner, ) -> None: reference = client_result.candidate_output_manifest - _require(reference is not None, "complete client result has no candidate manifest reference") + _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, "complete client result has no candidate manifest reference") + _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", 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 index c83ec6089..de8c558bb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -67,8 +67,10 @@ class CandidateOutputManifest(StateRecord): @property def winner_eligible(self) -> bool: - """Whether this complete candidate may be published as the shard winner.""" - return self.outcome is CandidateOutcome.COMPLETE + """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: 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 index a20828838..eeb23e66c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/reader.py @@ -184,7 +184,11 @@ def load_optional_attempt_result( shard: ShardManifest, attempt: AttemptManifest, ) -> tuple[ClientResult, CandidateOutputManifest] | None: - """Load and validate a complete producer result pair when present.""" + """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: 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 index 1e890cb05..5c15c568a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/results.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/results.py @@ -75,7 +75,7 @@ def _validate_record_location( def _bind_candidate_reference(attempt: AttemptManifest, client_result: ClientResult) -> AttemptManifest: reference = client_result.candidate_output_manifest if reference is None: - raise StateContractError("complete client result has no candidate manifest reference") + 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}) diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index 10f31d5ac..bcbacb4ea 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -455,6 +455,50 @@ def test_analysis_keeps_measured_infeasible_case_successful( 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, diff --git a/packages/data-designer-slurm/tests/contracts/test_shared_records.py b/packages/data-designer-slurm/tests/contracts/test_shared_records.py index e35ba4f98..83d725e1b 100644 --- a/packages/data-designer-slurm/tests/contracts/test_shared_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_shared_records.py @@ -248,7 +248,8 @@ def test_successful_benchmark_case_requires_metrics() -> None: @pytest.mark.parametrize( ("field", "value"), [ - ("actual_records", 999), + ("actual_records", 0), + ("actual_records", 1001), ("generation_seconds", 0), ("wall_seconds", 0), ("rows_per_second", 0), diff --git a/packages/data-designer-slurm/tests/integration/test_integration_validation.py b/packages/data-designer-slurm/tests/integration/test_integration_validation.py index 9f81f0226..dcc4be0ce 100644 --- a/packages/data-designer-slurm/tests/integration/test_integration_validation.py +++ b/packages/data-designer-slurm/tests/integration/test_integration_validation.py @@ -383,11 +383,11 @@ def test_client_candidate_validation_supports_preterminal_runtime_state(records: @pytest.mark.parametrize( ("mutation", "message"), [ - ("partial_result", "complete client results"), - ("failed_result", "complete client results"), + ("partial_result", "actual records must match"), + ("failed_result", "nonempty client results"), ("failed_attempt", "successful attempts"), ("wrong_terminal_classification", "successfully classified"), - ("ineligible_candidate", "complete candidate outputs"), + ("ineligible_candidate", "nonempty candidate outputs"), ("stale_winner", "winner attempt_id"), ("digest_mismatch", "digest"), ("dataset_path_mismatch", "dataset path"), diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index b8bf75f5f..afd19cdea 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -639,7 +639,7 @@ def write_future_result() -> None: assert state.attempt.state is AttemptLifecycleState.FAILED -def test_partial_client_result_is_classified_before_candidate_loading(runtime_case: RuntimeCase) -> None: +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)) @@ -658,10 +658,32 @@ def test_partial_client_result_is_classified_before_candidate_loading(runtime_ca with pytest.raises(SlurmRuntimeError) as raised: controller.run() - assert raised.value.code is SlurmRuntimeErrorCode.CLIENT_FAILED + 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) @@ -685,13 +707,14 @@ def test_stale_candidate_from_an_earlier_generation_cannot_succeed(runtime_case: assert state.attempt.state is AttemptLifecycleState.FAILED -def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: +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) - table = lazy.pa.table({"record_id": range(requested)}) + 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( @@ -703,14 +726,14 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: created_at=clock.now(), dataset_path=dataset_path, requested_records=requested, - actual_records=requested, - outcome=CandidateOutcome.COMPLETE, + 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=requested, + record_count=actual_records, ), ), dataset_schema_digest=compute_candidate_schema_digest(table.schema), @@ -726,8 +749,8 @@ def _write_complete_result(runtime_case: RuntimeCase, clock: FakeClock) -> None: attempt_id=context.attempt.attempt_id, completed_at=clock.now(), requested_records=requested, - actual_records=requested, - outcome=ClientOutcome.COMPLETE, + 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, diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index b8963513a..1d068cc6c 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -99,7 +99,7 @@ def load_candidate(*args: object, attempt_directory: Path | None = None) -> None assert attempt_directory == container_attempt_directory raise _InjectedFailure - monkeypatch.setattr(entrypoint, "load_complete_client_candidate", load_candidate) + monkeypatch.setattr(entrypoint, "load_nonempty_client_candidate", load_candidate) client = entrypoint._parse_arguments( _phase_arguments("client", runtime_case, attempt_directory=container_attempt_directory) ) diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 4781fdccf..8524973ac 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -515,6 +515,90 @@ def test_refresh_uses_terminal_accounting_over_stale_active_queue_state( 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, diff --git a/packages/data-designer-slurm/tests/state/test_records.py b/packages/data-designer-slurm/tests/state/test_records.py index 2eac315bf..b1c19e851 100644 --- a/packages/data-designer-slurm/tests/state/test_records.py +++ b/packages/data-designer-slurm/tests/state/test_records.py @@ -206,14 +206,14 @@ def test_pending_and_stopped_deployments_reject_ready_backends( AttemptReadiness.model_validate_json(json.dumps(payload)) -def test_partial_candidate_is_never_winner_eligible() -> None: +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 not partial_candidate.winner_eligible + assert partial_candidate.winner_eligible def test_candidate_output_rejects_count_and_file_mismatches() -> None: diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 18aa3977b..8a54514a6 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -789,6 +789,32 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( 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_requires_stable_plan_provenance_across_distinct_shards( tmp_path: Path, authored_run: DataDesignerSlurmConfig, @@ -1280,7 +1306,7 @@ def test_collection_uses_merge_partition_count( 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 requested records"): + with pytest.raises(StateContractError, match="partitions exceed available records"): validate_collection_inputs(resolved_plan, excessive, candidates) result = SlurmCollectionWorker( @@ -1942,7 +1968,7 @@ def _submitted_attempt( ) -def _publish_all_winners(case: _RunCase) -> None: +def _publish_all_winners(case: _RunCase, *, shortfall: int = 0) -> 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) @@ -1954,7 +1980,7 @@ def _publish_all_winners(case: _RunCase) -> None: ) 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) + _publish_candidate(case, shard, running, dataset_path, shortfall=shortfall) case.writer.finalize_winner( shard.shard_id, running.attempt_id, @@ -1962,9 +1988,17 @@ def _publish_all_winners(case: _RunCase) -> None: ) -def _publish_candidate(case: _RunCase, shard: ShardManifest, running: AttemptManifest, dataset_path: Path) -> None: +def _publish_candidate( + case: _RunCase, + shard: ShardManifest, + running: AttemptManifest, + dataset_path: Path, + *, + shortfall: int = 0, +) -> None: output_path = dataset_path / "part-00000.parquet" - values = range(shard.record_range.start_index, shard.record_range.end_index_exclusive) + 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}) lazy.pq.write_table(table, output_path) output_path.chmod(0o644) @@ -1981,14 +2015,14 @@ def _publish_candidate(case: _RunCase, shard: ShardManifest, running: AttemptMan created_at=running.updated_at + timedelta(minutes=1), dataset_path=dataset_path.as_posix(), requested_records=shard.record_range.record_count, - actual_records=shard.record_range.record_count, - outcome=CandidateOutcome.COMPLETE, + 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=shard.record_range.record_count, + record_count=actual_records, ), ), dataset_schema_digest=compute_candidate_schema_digest(table.schema), @@ -2002,8 +2036,8 @@ def _publish_candidate(case: _RunCase, shard: ShardManifest, running: AttemptMan attempt_id=running.attempt_id, completed_at=running.updated_at + timedelta(minutes=2), requested_records=shard.record_range.record_count, - actual_records=shard.record_range.record_count, - outcome=ClientOutcome.COMPLETE, + 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, diff --git a/packages/data-designer-slurm/tests/state/test_store.py b/packages/data-designer-slurm/tests/state/test_store.py index 8d04c8e28..9a9d902b9 100644 --- a/packages/data-designer-slurm/tests/state/test_store.py +++ b/packages/data-designer-slurm/tests/state/test_store.py @@ -1625,7 +1625,7 @@ def test_candidate_schema_digest_ignores_arrow_metadata() -> None: assert compute_candidate_schema_digest(annotated) == compute_candidate_schema_digest(schema) -def test_finalization_rejects_partial_client_results_and_candidate_file_drift( +def test_finalization_rejects_mismatched_partial_result_and_candidate_file_drift( tmp_path: Path, authored_run_single: DataDesignerSlurmConfig, single_node_plan: ResolvedSlurmRunPlan, diff --git a/packages/data-designer-slurm/tests/state/test_validation.py b/packages/data-designer-slurm/tests/state/test_validation.py index 8817b4339..6a254aedb 100644 --- a/packages/data-designer-slurm/tests/state/test_validation.py +++ b/packages/data-designer-slurm/tests/state/test_validation.py @@ -200,7 +200,7 @@ def test_attempt_transitions_preserve_identity_and_terminal_state() -> None: ) -def test_failed_partial_and_existing_winners_are_rejected() -> None: +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") @@ -230,8 +230,7 @@ def test_failed_partial_and_existing_winners_are_rejected() -> None: ) partial_attempt = _validated_copy(attempt, candidate_output=partial_reference) partial_winner = _validated_copy(winner, candidate_manifest=partial_reference) - with pytest.raises(StateContractError, match="winner policy"): - validate_shard_winner(run, shard, partial_attempt, partial, partial_winner) + assert validate_shard_winner(run, shard, partial_attempt, partial, partial_winner) == partial_winner with pytest.raises(StateContractError, match="immutable"): validate_shard_winner( From 71657149060f56388a95ee4d0c6907b46e6202ed Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 14 Oct 2026 14:55:08 -0600 Subject: [PATCH 52/62] fix(slurm): merge reordered nested shard schemas (#1028) * fix(slurm): merge reordered nested shard schemas Signed-off-by: Nabin Mulepati * fix(slurm): preserve saved collection plans Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../pages/slurm/retry-and-collection.mdx | 2 + .../slurm/launcher/collection.py | 8 +- .../data_designer/slurm/state/collection.py | 11 ++- .../slurm/state/collection_inputs.py | 5 + .../slurm/state/collection_merge.py | 20 +++- .../slurm/state/collection_schema.py | 77 +++++++++++++++ .../slurm/state/collection_validation.py | 2 - .../src/data_designer/slurm/state/outputs.py | 4 + .../tests/launcher/test_collection.py | 21 +++- .../tests/state/test_retry_collection.py | 96 ++++++++++++++++++- 10 files changed, 230 insertions(+), 16 deletions(-) create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/state/collection_schema.py diff --git a/fern/versions/latest/pages/slurm/retry-and-collection.mdx b/fern/versions/latest/pages/slurm/retry-and-collection.mdx index 894f13850..c3ba3dc15 100644 --- a/fern/versions/latest/pages/slurm/retry-and-collection.mdx +++ b/fern/versions/latest/pages/slurm/retry-and-collection.mdx @@ -40,6 +40,8 @@ Shards with at least one generated record can publish Parquet winners, even when 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 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 index 3ca48ade0..79cac21d3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -47,6 +47,8 @@ def render_collection_script( ) ) 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 @@ -57,9 +59,9 @@ def render_collection_script( 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(resolved_plan.package_version)} -readonly DD_RUNTIME_ARCHIVE={quote_shell_value(resolved_plan.runtime_bundle.path)} -readonly DD_RUNTIME_SHA256={quote_shell_value(resolved_plan.runtime_bundle.sha256)} +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)} 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 index 707b88456..85739aa54 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -7,6 +7,7 @@ import sys from datetime import datetime, timezone +from importlib.metadata import version from pathlib import Path from typing import Protocol @@ -17,6 +18,8 @@ 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, @@ -122,14 +125,18 @@ def submit( 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=self._inputs.get_winner_shards(), + 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, ) @@ -156,7 +163,7 @@ def submit( return self._submit_prepared(collection_plan, prepared, script, timestamp) except (StateConflictError, StateCorruptionError, SlurmStateError): raise - except (OSError, ValidationError, ValueError) as error: + except (OSError, SlurmRuntimeError, ValidationError, ValueError) as error: raise SlurmStateError(f"cannot submit collection for run {self._run_id!r}") from error def refresh( 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 index 764cf1ac3..75319d0e6 100644 --- 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 @@ -7,6 +7,8 @@ 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 @@ -23,6 +25,7 @@ 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.""" @@ -75,6 +78,8 @@ def resolve( 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) 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 index 4392bc5b4..aa2d2cda2 100644 --- 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 @@ -20,6 +20,7 @@ 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 @@ -73,6 +74,11 @@ def merge( ) -> 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) @@ -80,7 +86,7 @@ def merge( raise OSError("collection inputs contain no records") layout = _PartitionLayout( stage_path=staged.path, - schema=first_batch.schema.remove_metadata(), + 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, @@ -195,13 +201,17 @@ def describe(self, relative_path: str, record_count: int) -> CollectedOutputFile class _ParquetWriter: - def __init__(self, writer: _ArrowBatchWriter, output: _BoundOutput) -> None: + 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: - self._writer.write_batch(batch.replace_schema_metadata(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() @@ -218,7 +228,9 @@ def close(self) -> None: 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)) + partition_writer = _ParquetWriter( + lazy.pq.ParquetWriter(output, schema), _BoundOutput(output_path, output), schema + ) try: yield partition_writer finally: 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_validation.py b/packages/data-designer-slurm/src/data_designer/slurm/state/collection_validation.py index 2c63ac4d0..e98146db2 100644 --- 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 @@ -26,8 +26,6 @@ def validate_collection_inputs( _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") - schema_digests = {candidate.dataset_schema_digest for candidate in candidates} - _require(len(schema_digests) == 1, "collection candidates have incompatible schemas") _require( all(candidate.provenance_digest == resolved_plan.compute_sha256() for candidate in candidates), "collection candidate provenance does not match the resolved plan", 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 index de8c558bb..f1232753e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/outputs.py @@ -165,6 +165,8 @@ class CollectionPlan(StateRecord): 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 @@ -179,6 +181,8 @@ def submission_job_name(self) -> Identifier: @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)): diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index dd4d574a9..23d5012a9 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -10,6 +10,7 @@ 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 @@ -65,6 +66,8 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives ), ), 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, ) @@ -79,9 +82,9 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives '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.2"' in script - assert f'readonly DD_RUNTIME_ARCHIVE="{plan.runtime_bundle.path}"' in script - assert f'readonly DD_RUNTIME_SHA256="{plan.runtime_bundle.sha256}"' 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 @@ -93,6 +96,16 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives 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, @@ -225,6 +238,8 @@ def test_destination_reauthorizes_explicit_path_through_workspace_mapping( ), ), 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, ) diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 8a54514a6..15910b9f5 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -55,6 +55,7 @@ ) 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 @@ -758,6 +759,12 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( 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 @@ -815,6 +822,64 @@ def test_collection_merges_nonempty_partial_shard_winners( 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, @@ -1968,7 +2033,12 @@ def _submitted_attempt( ) -def _publish_all_winners(case: _RunCase, *, shortfall: int = 0) -> None: +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) @@ -1980,7 +2050,15 @@ def _publish_all_winners(case: _RunCase, *, shortfall: int = 0) -> None: ) 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) + _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, @@ -1995,11 +2073,25 @@ def _publish_candidate( 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() From 1ad7f68de87ae09766c196b36619dd0710e3edfe Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 15 Oct 2026 11:09:02 -0600 Subject: [PATCH 53/62] fix(slurm): benchmark one array task per case (#1037) * fix(slurm): benchmark one array task per case Signed-off-by: Nabin Mulepati * fix(slurm): analyze saved benchmark manifests Signed-off-by: Nabin Mulepati * fix(slurm): validate benchmark versions lazily Signed-off-by: Nabin Mulepati * docs: clarify saved benchmark rates Signed-off-by: Nabin Mulepati * feat(slurm): default benchmark records adaptively Signed-off-by: Nabin Mulepati * docs: fix Slurm example references Signed-off-by: Nabin Mulepati * docs: clarify Slurm topologies and shards Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/benchmarks.mdx | 32 ++---- .../latest/pages/slurm/getting-started.mdx | 28 +++-- .../latest/pages/slurm/operations.mdx | 2 + fern/versions/latest/pages/slurm/profiles.mdx | 2 + .../latest/pages/slurm/run-configuration.mdx | 108 ++++++++++++++++-- .../data_designer/slurm/benchmark/analysis.py | 21 ++-- .../data_designer/slurm/benchmark/compiler.py | 44 ++++--- .../data_designer/slurm/config/benchmark.py | 10 +- .../benchmark/test_benchmark_compiler.py | 24 ++++ .../tests/benchmark/test_workflow.py | 58 ++++++++++ 10 files changed, 254 insertions(+), 75 deletions(-) diff --git a/fern/versions/latest/pages/slurm/benchmarks.mdx b/fern/versions/latest/pages/slurm/benchmarks.mdx index 842ffbf73..b5db16406 100644 --- a/fern/versions/latest/pages/slurm/benchmarks.mdx +++ b/fern/versions/latest/pages/slurm/benchmarks.mdx @@ -4,12 +4,6 @@ Use a benchmark to compare concurrency and deployment topology while keeping eac Complete [Getting Started](/slurm/getting-started) first so the profile, images, and base run are ready. -Install the Slurm extension: - -```bash -pip install "data-designer[slurm]" -``` - ## Define cases Create `benchmark.yaml` next to an existing `run.yaml`: @@ -33,11 +27,6 @@ deployment_cases: generator: nodes: 2 nodes_per_replica: 2 -record_policy: - type: adaptive - base_records: 1000 - max_records: 5000 - records_per_concurrency: 1.0 analysis: target_total_records: 1000000 target_runtime: 4h @@ -45,21 +34,21 @@ analysis: `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`. An adaptive record policy uses `ceil(concurrency * records_per_concurrency)`, bounded by `base_records` and `max_records`. +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)`, so it also depends on the selected profile and the base run's deployment topology. +`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 \ - --profile-file ~/.data-designer-slurm-profile.yml \ - --cluster primary +data-designer slurm benchmark run benchmark.yaml ``` -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/`. +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. 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. +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 @@ -77,7 +66,7 @@ data-designer slurm benchmark analyze --refresh-state `--refresh-state` performs one fresh-process scheduler reconciliation per child before analysis. Add `--fail-if-incomplete` when automation should return a conflict after the report is persisted if any case is not successful. -Reports retain every manifest case in order. Outcomes include `pending`, `accounting_lag`, `succeeded`, `failed`, `incomplete`, `missing`, `stale`, and `scheduler_inconsistent`. Successful cases can be infeasible when boot time consumes the runtime budget. Boot and wall timing begin when execution starts inside each allocation, so queue wait does not affect topology comparisons. For array runs, `rows_per_second` sums the independently measured shard rates, while target jobs and GPU hours count individual task allocations. `generation_seconds` is the target runtime remaining after the slowest shard boot. Only successful feasible cases participate in Pareto, minimum-job, and minimum-GPU-hour recommendations. +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 @@ -89,10 +78,7 @@ from data_designer.slurm.services import create_slurm_benchmark_service benchmark_file = Path("benchmark.yaml").resolve() config = load_benchmark_config(benchmark_file) -service = create_slurm_benchmark_service( - profile_file="~/.data-designer-slurm-profile.yml", - cluster="primary", -) +service = create_slurm_benchmark_service() manifest = service.run(config, source_root=benchmark_file.parent) report = service.analyze(manifest.benchmark_id, refresh_state=True) diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index adda90328..b4ec3815f 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -65,12 +65,12 @@ data-designer slurm profile init \ --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. For this one-GPU walkthrough on a shared GRES cluster, set `gpus_per_node: 1` when your site permits partial-node allocations. 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. +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 --cluster primary +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. @@ -81,8 +81,7 @@ Validation checks strict profile loading, cluster selection, workspace access, a data-designer slurm image add \ vllm/vllm-openai:v0.22.0 \ --kind serving \ - --name vllm \ - --cluster primary + --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). @@ -98,13 +97,18 @@ data_designer: 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 greeting. + prompt: Write one short {{ tone }} greeting. model_alias: generator ``` -The model alias in this file must match a deployment in the Slurm run config. +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 @@ -136,12 +140,12 @@ submission: time_limit: 01:00:00 ``` -The repository also contains these 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. +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 --cluster primary +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`. @@ -149,16 +153,16 @@ Dry-run resolves the selected profile, records the active shared Python runtime, ## 10. Submit the run ```bash -data-designer slurm execute run.yaml --cluster primary +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 --cluster primary`. +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 --cluster primary -data-designer slurm cancel --cluster primary +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. diff --git a/fern/versions/latest/pages/slurm/operations.mdx b/fern/versions/latest/pages/slurm/operations.mdx index 0c3858a8b..a965b3131 100644 --- a/fern/versions/latest/pages/slurm/operations.mdx +++ b/fern/versions/latest/pages/slurm/operations.mdx @@ -29,6 +29,8 @@ In an interactive terminal, `execute` reports when it invokes `sbatch` with the 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 diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index 822ca1713..1e369ade2 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -87,6 +87,8 @@ The CLI also accepts `--partition gpu-a,gpu-b` during profile initialization. Sl `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. +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. diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index a2019b3dd..22f85767a 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -14,8 +14,6 @@ invocation: dataset_name: product-descriptions model_concurrency: generator: 64 -client: - cpus: 16 deployments: - model_alias: generator served_model_name: example/generator @@ -79,26 +77,37 @@ server: ## 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. 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 records into independent shards and limit how many allocations run at once: +Split the requested total into independent shards and cap the number of allocations that can run at once: ```yaml invocation: - num_records: 100000 - dataset_name: large-run + num_records: 1000 + dataset_name: sharded-run array_tasks: - count: 20 - max_concurrent: 4 + count: 5 + max_concurrent: 2 ``` -`count` cannot exceed `num_records`, and `max_concurrent` cannot exceed `count`. Status, retry, and collection retain shard identities and winners. Array task count controls execution shards; it does not set the number of collected Parquet files. Use `slurm merge RUN_ID --num-partitions 20` if you want 20 output files. +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 keep one vLLM replica per node: +Allocate multiple nodes and use all eight GPUs per node for one vLLM replica on each node: ```yaml resources: @@ -108,6 +117,16 @@ topology: 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: @@ -120,7 +139,31 @@ topology: nodes_per_replica: 2 ``` -`tensor_parallel` must divide the GPUs available on each node. With the eight GPUs per node in the profile example, `tensor_parallel: 8` and `nodes_per_replica: 2` use 16 GPUs per replica. `nodes_per_replica` must divide `resources.nodes`. Multi-node expert parallelism is not supported in v1. +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 @@ -148,6 +191,51 @@ invocation: judge: 32 ``` +### 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 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 index 9b9dc0d33..9fb56ba48 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -274,19 +274,20 @@ def _validate_and_get_requested_records( base_run: DataDesignerSlurmConfig, manifest: BenchmarkManifest, ) -> tuple[int, ...]: - try: - compiled = BenchmarkCompiler.compile(config, base_run) - except ValueError as error: - raise BenchmarkManifestMismatchError("benchmark configuration cannot reproduce its children") from error - expected = tuple( - (case.case_id, case.child_run_id, case.child_run_config.compute_sha256()) for case in compiled.cases - ) actual = tuple( (child.case_id, child.child_run_id, child.child_authored_config.sha256) for child in manifest.children ) - if manifest.benchmark_id != compiled.benchmark_id or actual != expected: - raise BenchmarkManifestMismatchError("benchmark manifest does not match deterministic child expansion") - return tuple(case.child_run_config.invocation.num_records for case in compiled.cases) + 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: 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 index a78e82364..4f6f7404c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/compiler.py @@ -11,12 +11,15 @@ 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.""" @@ -40,6 +43,15 @@ class BenchmarkCompiler: 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") @@ -58,12 +70,13 @@ def compile( f"benchmark deployment case references unknown aliases: {', '.join(sorted(unknown_deployments))}" ) - identity_digest = compute_canonical_json_sha256( - { - "benchmark": config.model_dump(mode="json"), - "base_run": base_run.model_dump(mode="json"), - } - ) + 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: @@ -79,6 +92,7 @@ def compile( selected_aliases=selected_aliases, deployment_case=deployment_case, requested_records=resolve_requested_records(config, concurrency), + single_task=not legacy, ), ) ) @@ -102,6 +116,7 @@ def _compile_child( 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)) @@ -128,15 +143,14 @@ def _compile_child( ) ) - return DataDesignerSlurmConfig.model_validate( - base_run.model_copy( - update={ - "name": _bounded_identifier(f"{base_run.name}-{case_id}"), - "invocation": invocation, - "deployments": deployments, - } - ).model_dump(mode="python") - ) + 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: 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 index 8fb92a574..10a2386eb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/benchmark.py @@ -62,10 +62,10 @@ class FixedRecordPolicy(AuthoredConfig): class AdaptiveRecordPolicy(AuthoredConfig): - type: Literal["adaptive"] - base_records: PositiveInt - max_records: PositiveInt - records_per_concurrency: PositiveFloat + 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: @@ -91,7 +91,7 @@ class DataDesignerSlurmBenchmarkConfig(AuthoredConfig): 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 + record_policy: BenchmarkRecordPolicy = Field(default_factory=AdaptiveRecordPolicy) analysis: BenchmarkAnalysisTargets @model_validator(mode="after") diff --git a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py index 907acd5ea..3c1c2a3b1 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py +++ b/packages/data-designer-slurm/tests/benchmark/test_benchmark_compiler.py @@ -27,6 +27,8 @@ def test_compiler_expands_stable_ordered_ordinary_runs( ) 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, @@ -75,3 +77,25 @@ def test_adaptive_record_policy_is_bounded( 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 index bcbacb4ea..c31b68b4e 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -29,6 +29,7 @@ BenchmarkBaseRun, DataDesignerSlurmBenchmarkConfig, DataDesignerSlurmConfig, + FixedRecordPolicy, ) from data_designer.slurm.contracts import ArtifactReference from data_designer.slurm.planning import ResolvedSlurmRunPlan @@ -529,6 +530,63 @@ def test_analysis_rejects_manifest_rewritten_away_from_compiled_children( ) +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, From f7eebabc9174ea697254fd2d06d536cacf89c2b0 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:48:56 -0300 Subject: [PATCH 54/62] feat(slurm): colocate distinct models with GPU quotas (#1034) Add opt-in shared-node GPU quotas and explicit GPU offsets for colocating distinct models, including multiple models on one GPU. Coordinate readiness before overlapping launches and support typed vLLM IO processor plugins. Signed-off-by: Andre Manoel --- .../latest/pages/slurm/run-configuration.mdx | 54 +++- .../examples/colocated-models.yaml | 42 +++ .../src/data_designer/slurm/_errors.py | 17 ++ .../data_designer/slurm/benchmark/analysis.py | 9 +- .../src/data_designer/slurm/config/run.py | 39 +++ .../src/data_designer/slurm/config/vllm.py | 4 +- .../data_designer/slurm/planning/compiler.py | 23 +- .../data_designer/slurm/planning/models.py | 68 ++++- .../slurm/planning/resolution.py | 12 + .../data_designer/slurm/runtime/controller.py | 5 + .../slurm/runtime/distributed.py | 59 +++- .../data_designer/slurm/runtime/manifest.py | 3 + .../data_designer/slurm/runtime/node_spec.py | 80 +++++- .../slurm/runtime/node_worker.py | 37 ++- .../slurm/runtime/plan_reader.sh | 4 +- .../slurm/runtime/server_manifest.py | 91 ++++-- .../slurm/runtime/step_runner.sh | 20 +- .../src/data_designer/slurm/runtime/steps.py | 7 + .../data_designer/slurm/serving/deployment.py | 3 +- .../data_designer/slurm/serving/resolver.py | 4 +- .../src/data_designer/slurm/serving/vllm.py | 1 + .../tests/benchmark/test_workflow.py | 70 ++++- .../data-designer-slurm/tests/conftest.py | 46 +++ .../tests/contracts/test_config_records.py | 14 + .../tests/planning/test_colocation.py | 248 +++++++++++++++++ .../tests/runtime/test_controller.py | 26 +- .../tests/runtime/test_entrypoint.py | 17 +- .../tests/runtime/test_node_worker.py | 195 ++++++++++++- .../tests/runtime/test_shared_node_runtime.py | 261 ++++++++++++++++++ .../tests/runtime/test_steps.py | 24 ++ .../tests/serving/test_resolver.py | 12 + .../tests/serving/test_serving_records.py | 13 +- plans/870/model-colocation.md | 25 ++ 33 files changed, 1451 insertions(+), 82 deletions(-) create mode 100644 packages/data-designer-slurm/examples/colocated-models.yaml create mode 100644 packages/data-designer-slurm/tests/planning/test_colocation.py create mode 100644 packages/data-designer-slurm/tests/runtime/test_shared_node_runtime.py create mode 100644 plans/870/model-colocation.md diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 22f85767a..0b8186340 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -73,11 +73,13 @@ server: `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. 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. +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 @@ -191,6 +193,56 @@ invocation: 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: 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/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py index 584bdf62b..ef4cdbbea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -27,6 +27,23 @@ "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" ), 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 index 9fb56ba48..9a7241ffc 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/analysis.py @@ -244,8 +244,8 @@ def _resolved_facts( authored: DataDesignerSlurmConfig, plan: ResolvedSlurmRunPlan, ) -> BenchmarkCaseResult: - nodes = sum(len(deployment.node_indices) for deployment in plan.deployments) - gpus = sum(len(deployment.node_indices) * deployment.gpus_per_node for deployment in plan.deployments) + 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( [ { @@ -254,6 +254,11 @@ def _resolved_facts( "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 ] 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 index cb125ad6e..dbc4c0fde 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -15,6 +15,7 @@ from pydantic import ( Field, JsonValue, + NonNegativeInt, PositiveInt, StringConstraints, field_validator, @@ -256,6 +257,19 @@ class ClientConfig(AuthoredConfig): 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): @@ -288,6 +302,13 @@ def validate_served_model_name(cls, value: str | None) -> str | None: @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: @@ -347,6 +368,24 @@ def validate_run(self) -> DataDesignerSlurmConfig: 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( 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 index d440881d5..2db28d086 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/vllm.py @@ -18,7 +18,7 @@ validate_vllm_readiness_path, ) from data_designer.slurm.contracts import AuthoredConfig -from data_designer.slurm.types import Duration, EnvironmentName, NonNegativeDuration +from data_designer.slurm.types import Duration, EnvironmentName, Identifier, NonNegativeDuration __all__ = [ "QueueBackpressureConfig", @@ -53,6 +53,7 @@ class VllmServerConfig(AuthoredConfig): 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. @@ -66,6 +67,7 @@ class VllmServerConfig(AuthoredConfig): 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) 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 index 81fd40698..f107a9e2e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -24,6 +24,7 @@ ResolvedDeployment, ResolvedSlurmRunPlan, ResolvedTopology, + derive_deployment_placement, ) from data_designer.slurm.planning.resolution import ( EffectiveDataDesignerSlurmConfig, @@ -109,21 +110,27 @@ def _compile_deployments( if len(effective.authored.deployments) > _PORT_RANGE_SIZE: raise SlurmPlanCompilationError("deployment count exceeds the compiler-owned logical endpoint port range") resolved: list[ResolvedDeployment] = [] - next_node_index = 0 + 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=effective.resolved_gpus_per_node, + 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}" - node_indices = tuple(range(next_node_index, next_node_index + authored.resources.nodes)) - ports = _compile_deployment_ports(deployment_id, node_indices, topology) + 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, @@ -131,12 +138,12 @@ def _compile_deployments( served_model_name=authored.served_model_name or authored.model, image=image, node_indices=node_indices, - gpus_per_node=effective.resolved_gpus_per_node, + gpus_per_node=quota, + gpu_start_index=gpu_start_index, topology=topology, ports=ports, ) ) - next_node_index += authored.resources.nodes return tuple(resolved) @@ -144,6 +151,8 @@ 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] = [] @@ -156,7 +165,7 @@ def _compile_deployment_ports( name=f"{deployment_id}-http-{replica_index:05d}", role="http", node_index=head, - port=_HTTP_PORT + lane_index, + port=_HTTP_PORT + http_offset + lane_index, ) ) if topology.nodes_per_replica > 1: 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 index 48f13f191..10ee11ee4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -4,6 +4,7 @@ from __future__ import annotations import posixpath +from collections.abc import Sequence from typing import Annotated, Literal from urllib.parse import urlsplit @@ -277,6 +278,7 @@ class ResolvedDeployment(ContractValue): 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, ...] = () @@ -422,8 +424,35 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: 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") - if any(deployment.gpus_per_node != self.resolved_gpus_per_node for deployment in self.deployments): - raise ValueError("every deployment must use the resolved profile GPU count") + 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") @@ -437,9 +466,6 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: if set(aliases) != set(self.builder.model_aliases): raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") - node_indices = tuple(index for deployment in self.deployments for index in deployment.node_indices) - if node_indices != tuple(range(len(node_indices))): - raise ValueError("deployment nodes must be disjoint and contiguous in authored order") 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 @@ -542,3 +568,35 @@ def _validate_shards(self, run_root: str) -> None: 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 index f54063536..5f735142d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -36,6 +36,7 @@ ResolvedImage, ResolvedInvocation, ResolvedSubmission, + derive_deployment_placement, ) __all__: list[str] = [] @@ -179,6 +180,17 @@ def validate_effective_slurm_config( _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 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index c1b987f9c..91330e39e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -145,6 +145,11 @@ def __init__( ) -> 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 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 index 1aabb4b6b..6d55c5893 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/distributed.py @@ -6,7 +6,7 @@ from __future__ import annotations from data_designer.slurm.planning import ResolvedSlurmRunPlan -from data_designer.slurm.runtime.node_spec import NodeProcessSpec, NodeSpec, NodeWorkerSpec +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 @@ -18,36 +18,65 @@ 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=_get_node_ports(deployment, node_index), + ports=tuple(sorted(port for member in deployments for port in _get_node_ports(member, node_index))), processes=tuple( - NodeProcessSpec( - process_id=process.process_id, - command=( - build_vllm_command(deployment, process, plan) - if len(layout.node_hosts) == 1 - else build_vllm_process_command(deployment, process, plan, layout) + 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 ), - gpu_indices=tuple(process.gpu_indices), - launch_delay_seconds=process.launch_delay_seconds, + key=lambda process: process.launch_delay_seconds, ) - for process in deployment.processes - if process.node_index == node_index ), ) for node_index in deployment.node_indices ) return NodeWorkerSpec( schema_version=1, - resolved_gpus_per_node=deployment.gpus_per_node, - required_model_path=model if deployment.model.startswith("/") else None, + 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, ) @@ -82,6 +111,8 @@ def build_vllm_process_command( ) 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 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 index 5569f38c2..60da12f7d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -47,6 +47,7 @@ class RuntimeStepSpec(ContractValue): 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) @@ -100,6 +101,8 @@ def _validate_placement(self) -> None: 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: 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 index b225328dd..3c9d49576 100644 --- 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 @@ -8,7 +8,7 @@ import base64 import binascii import json -from dataclasses import dataclass +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 @@ -16,6 +16,26 @@ _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.""" @@ -24,6 +44,7 @@ class NodeProcessSpec: command: tuple[str, ...] gpu_indices: tuple[int, ...] launch_delay_seconds: int + readiness: NodeReadinessSpec | None = None def __post_init__(self) -> None: if ( @@ -47,6 +68,8 @@ def __post_init__(self) -> None: 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) @@ -70,6 +93,10 @@ def __post_init__(self) -> None: 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) @@ -80,8 +107,12 @@ class NodeWorkerSpec: 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: @@ -90,6 +121,12 @@ def __post_init__(self) -> None: 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 @@ -104,7 +141,10 @@ def __post_init__(self) -> None: 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): - raise ValueError("node processes must not share 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: @@ -124,6 +164,7 @@ def encode_node_worker_spec(spec: NodeWorkerSpec) -> str: "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 ], @@ -131,6 +172,10 @@ def encode_node_worker_spec(spec: NodeWorkerSpec) -> str: 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") @@ -152,13 +197,20 @@ def decode_node_worker_spec(encoded: str) -> NodeWorkerSpec: def _parse_worker_spec(payload: object) -> NodeWorkerSpec: - root = _require_mapping(payload, {"schema_version", "resolved_gpus_per_node", "required_model_path", "nodes"}) + 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)), ) @@ -173,12 +225,25 @@ def _parse_node(payload: object) -> NodeSpec: def _parse_process(payload: object) -> NodeProcessSpec: - value = _require_mapping(payload, {"process_id", "command", "gpu_indices", "launch_delay_seconds"}) + 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"]), ) @@ -212,8 +277,15 @@ def _require_integer(payload: object) -> int: 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", 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 index dc282cce6..56186b79d 100644 --- 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 @@ -19,6 +19,7 @@ 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, ) @@ -31,6 +32,7 @@ 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 @@ -56,11 +58,19 @@ 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 process in node.processes: - self._wait_for_launch(process.launch_delay_seconds, started_at) + 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() @@ -113,13 +123,32 @@ def _start_process(self, process: NodeProcessSpec, visible_gpus: tuple[str, ...] close_fds=True, ) - def _wait_for_launch(self, delay_seconds: int, started_at: float) -> None: + 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: @@ -144,6 +173,8 @@ def main(arguments: Sequence[str] | None = None) -> int: _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) 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 index bf39dcbf6..45fdad59f 100644 --- 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 @@ -147,7 +147,8 @@ dd_read_step() { .stdout_path, "\u0000", .stderr_path, "\u0000", (.launch_delay_seconds | tostring), "\u0000", - (.kill_on_bad_exit | tostring), "\u0000" + (.kill_on_bad_exit | tostring), "\u0000", + (.gpu_binding // "mask"), "\u0000" ' "${manifest}" ) DD_STEP_EXECUTION=${DD_STEP_FIELDS[0]} @@ -157,6 +158,7 @@ dd_read_step() { 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}" 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 index 34973f1f3..f986a49cc 100644 --- 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 @@ -33,30 +33,53 @@ def build_server_steps( layout: AllocationLayout, ) -> tuple[RuntimeStepSpec, ...]: """Build remote preflight entries followed by all serving entries.""" - preflight_steps = tuple( - build_distributed_server_step( - deployment, - context, - runtime_container_root, - log_directory, - layout, - operation="preflight", - ) - for deployment in deployments - if _requires_server_preflight(deployment, context) - ) - serving_steps = tuple( - step - for deployment in deployments - for step in _build_deployment_serving_steps( - deployment, - context, - runtime_container_root, - log_directory, - layout, - ) - ) - return preflight_steps + serving_steps + 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( @@ -67,13 +90,14 @@ def build_distributed_server_step( 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) + worker_spec = build_node_worker_spec(deployment, context.plan, layout, colocated_deployments=colocated_deployments) readiness = ( tuple( RuntimeProbeSpec( @@ -83,16 +107,24 @@ def build_distributed_server_step( deadline_seconds=probe.deadline_seconds + next( process.launch_delay_seconds - for process in deployment.processes + for process in member.processes if process.node_index == probe.node_index and process.http_port == probe.port ), ) - for probe in deployment.readiness_probes + 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, @@ -106,7 +138,10 @@ def build_distributed_server_step( encode_node_worker_spec(worker_spec), ), cpus=context.plan.client.authored.cpus, - gpu_indices=tuple(range(deployment.gpus_per_node)), + 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, 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 index e4de7b973..c4ecbaaec 100644 --- 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 @@ -89,14 +89,18 @@ dd_build_srun_command() { fi if ((${#DD_STEP_GPU_INDICES[@]})); then if [[ ${DD_GPU_REQUEST_MODE} == gres ]]; 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+=( - "--gpus-per-task=${#DD_STEP_GPU_INDICES[@]}" - "--gpu-bind=mask_gpu:${gpu_mask}" - ) + 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[@]}" 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 index c6f2be674..b1b8f8a3c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -232,6 +232,11 @@ def build_vllm_steps( 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 @@ -402,6 +407,8 @@ def build_vllm_command( ) 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 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 index 51dc57f5a..673405d4e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/deployment.py @@ -37,6 +37,7 @@ class ResolvedVllmServerDeployment(ContractValue): 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) @@ -139,7 +140,7 @@ def _validate_replica(self, endpoint: ResolvedBackendEndpoint, probe: ResolvedRe for process in self.processes if process.deployment_replica_index == endpoint.deployment_replica_index ) - expected_gpu_start = expected_replica_index_in_node_group * self.topology.tensor_parallel + 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 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 index dd85211e7..95f038a37 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py @@ -122,7 +122,7 @@ def _resolve_vllm_server( replica_index_in_node_group=replica_index_in_node_group, deployment_replica_index=deployment_replica_index, ) - gpu_start = replica_index_in_node_group * tensor_parallel + 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): @@ -159,6 +159,7 @@ def _resolve_vllm_server( 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), @@ -167,6 +168,7 @@ def _resolve_vllm_server( 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), 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 index 73f3e085a..cb2bc52ea 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/vllm.py @@ -36,6 +36,7 @@ class ResolvedVllmLaunchPolicy(ContractValue): 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) diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index c31b68b4e..dcc6a77c0 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -32,7 +32,9 @@ FixedRecordPolicy, ) from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.planning import ResolvedSlurmRunPlan +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, @@ -43,6 +45,72 @@ 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 diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index ed14a4820..ec7950ebe 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -24,10 +24,14 @@ 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, @@ -81,6 +85,48 @@ 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()) diff --git a/packages/data-designer-slurm/tests/contracts/test_config_records.py b/packages/data-designer-slurm/tests/contracts/test_config_records.py index 8f130d1e7..33120368b 100644 --- a/packages/data-designer-slurm/tests/contracts/test_config_records.py +++ b/packages/data-designer-slurm/tests/contracts/test_config_records.py @@ -225,6 +225,20 @@ def test_vllm_rejects_nonpositive_retry_after_seconds() -> None: 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", [ 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/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index afd19cdea..ceb4a1e73 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -10,13 +10,14 @@ from typing import Callable import pytest -from conftest import FakeClientStepBuilder, FakePreflight, FakeStateStore, RuntimeCase +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 @@ -43,6 +44,29 @@ _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 diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 1d068cc6c..67e90aaf2 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -10,14 +10,14 @@ import os import sys import zipfile -from dataclasses import dataclass +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 +from conftest import FakeStateStore, RuntimeCase, relocate_plan from packaging.tags import interpreter_name, interpreter_version import data_designer.slurm.runtime.entrypoint as entrypoint @@ -161,10 +161,15 @@ def test_client_phase_installs_and_runs_plugin_in_default_worker_process( 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" @@ -184,10 +189,18 @@ def test_control_phases_record_running_ready_and_failed( 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( diff --git a/packages/data-designer-slurm/tests/runtime/test_node_worker.py b/packages/data-designer-slurm/tests/runtime/test_node_worker.py index b5fb57b13..10c1314e6 100644 --- a/packages/data-designer-slurm/tests/runtime/test_node_worker.py +++ b/packages/data-designer-slurm/tests/runtime/test_node_worker.py @@ -3,17 +3,24 @@ 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 +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, @@ -22,6 +29,192 @@ 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 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_steps.py b/packages/data-designer-slurm/tests/runtime/test_steps.py index 91a8f2657..7b9cbb3b5 100644 --- a/packages/data-designer-slurm/tests/runtime/test_steps.py +++ b/packages/data-designer-slurm/tests/runtime/test_steps.py @@ -18,8 +18,10 @@ 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, @@ -141,6 +143,28 @@ def test_vllm_command_maps_absolute_model_path(runtime_case: RuntimeCase) -> Non 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) diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index 28155e9fc..26532ca01 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -63,6 +63,18 @@ def test_multi_node_replica_resolution_matches_golden(multi_replica_plan: Resolv 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: diff --git a/packages/data-designer-slurm/tests/serving/test_serving_records.py b/packages/data-designer-slurm/tests/serving/test_serving_records.py index a79086ca9..3340709d7 100644 --- a/packages/data-designer-slurm/tests/serving/test_serving_records.py +++ b/packages/data-designer-slurm/tests/serving/test_serving_records.py @@ -12,7 +12,18 @@ 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 ResolvedVllmProcess +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( 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. From f522646d3e38da2fba347039b412a4dfc8c1dc73 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Thu, 15 Oct 2026 15:24:53 -0600 Subject: [PATCH 55/62] fix: recognize locked dependency licenses --- dependency-license-policy.toml | 7 +++++++ scripts/check_dependency_licenses.py | 1 + tests_e2e/test_dependency_license_check.py | 20 ++++++++++++++++++++ 3 files changed, 28 insertions(+) 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/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/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={}) From 98f19ef8fdb1aa4e1dc52c287f3470a48c14fe37 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 16 Oct 2026 10:49:35 -0300 Subject: [PATCH 56/62] fix(slurm): deduplicate runtime inventory paths Signed-off-by: Andre Manoel --- .../data_designer/slurm/client/environment.py | 3 +- .../tests/client/test_environment.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) 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 index 667d5e3c7..958fb4ab0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -338,7 +338,8 @@ def _active_distribution_search_paths() -> list[str]: continue if normalized_entry in configured_paths and normalized_entry not in standard_site_paths: continue - active_paths.append(entry) + if normalized_entry not in active_paths: + active_paths.append(normalized_entry) return active_paths diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 920601dad..9c16e0767 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -7,6 +7,7 @@ import importlib import json import os +import shutil import subprocess import sys from pathlib import Path @@ -209,6 +210,37 @@ def distributions(**kwargs: object) -> tuple[()]: 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) == () From 21b75957c6db039cabca7994ee7d7abbbea93485 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Fri, 16 Oct 2026 10:56:44 -0300 Subject: [PATCH 57/62] fix(slurm): sync checked examples with documentation Signed-off-by: Andre Manoel --- packages/data-designer-slurm/examples/benchmark.yaml | 8 +------- packages/data-designer-slurm/examples/builder.yaml | 7 ++++++- .../data-designer-slurm/examples/profile-catalog.yaml | 1 - scripts/test_slurm_package_install.py | 3 ++- 4 files changed, 9 insertions(+), 10 deletions(-) diff --git a/packages/data-designer-slurm/examples/benchmark.yaml b/packages/data-designer-slurm/examples/benchmark.yaml index 120120654..f3608b15e 100644 --- a/packages/data-designer-slurm/examples/benchmark.yaml +++ b/packages/data-designer-slurm/examples/benchmark.yaml @@ -1,8 +1,7 @@ schema_version: 1 name: generator-scaling base_run: run.yaml -model_aliases: - - generator +model_aliases: all concurrency_values: - 32 - 64 @@ -17,11 +16,6 @@ deployment_cases: generator: nodes: 2 nodes_per_replica: 2 -record_policy: - type: adaptive - base_records: 1000 - max_records: 5000 - records_per_concurrency: 1.0 analysis: target_total_records: 1000000 target_runtime: 4h diff --git a/packages/data-designer-slurm/examples/builder.yaml b/packages/data-designer-slurm/examples/builder.yaml index 9e2ca35d8..f2cf333b0 100644 --- a/packages/data-designer-slurm/examples/builder.yaml +++ b/packages/data-designer-slurm/examples/builder.yaml @@ -4,7 +4,12 @@ data_designer: 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 greeting. + prompt: Write one short {{ tone }} greeting. model_alias: generator diff --git a/packages/data-designer-slurm/examples/profile-catalog.yaml b/packages/data-designer-slurm/examples/profile-catalog.yaml index 77ce34607..add258b4a 100644 --- a/packages/data-designer-slurm/examples/profile-catalog.yaml +++ b/packages/data-designer-slurm/examples/profile-catalog.yaml @@ -9,7 +9,6 @@ clusters: account: research partition: gpu mem_per_gpu: 80G - bin_path: /opt/slurm/bin gpus_per_node: 8 workspace_root: /workspace/alice/data-designer-slurm image_build: diff --git a/scripts/test_slurm_package_install.py b/scripts/test_slurm_package_install.py index a6809e178..ce2c8b178 100644 --- a/scripts/test_slurm_package_install.py +++ b/scripts/test_slurm_package_install.py @@ -297,7 +297,8 @@ def verify_install(python: Path, version: str, *, slurm: bool, cwd: Path) -> Non assert StateRecordRange is ContractRecordRange assert StateResumeWorkspace is ContractResumeWorkspace examples = Path("examples") -assert load_builder_payload(examples / "builder.yaml")["data_designer"]["columns"][0]["name"] == "greeting" +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" From 86afd6976ba4ad3cb119cdf45fa281eacc74b8d4 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 16 Oct 2026 14:59:49 -0600 Subject: [PATCH 58/62] fix(slurm): harden run lifecycle and diagnostics (#1044) * fix(slurm): harden run lifecycle and diagnostics Validate built-in builder settings before allocation while deferring plugin and compute-local inputs. Protect image artifacts referenced by run plans and verify reused client overlays. Correct readiness, progress, benchmark, and multi-node behaviors found during branch review. Add regression coverage and update Slurm docs. Signed-off-by: Nabin Mulepati * fix(slurm): address review findings Poll readiness across all started servers so a slow deployment cannot consume another deployment timeout. Protect images referenced by custom run IDs and retry transient batch metadata failures in progress reporting. Signed-off-by: Nabin Mulepati * fix(slurm): narrow artifact locks and validate plans Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- fern/versions/latest/pages/slurm/images.mdx | 2 +- fern/versions/latest/pages/slurm/profiles.mdx | 2 + .../slurm/benchmark/execution.py | 21 +- .../src/data_designer/slurm/cli_benchmark.py | 17 +- .../data_designer/slurm/client/environment.py | 64 +++- .../data_designer/slurm/client/execution.py | 28 +- .../src/data_designer/slurm/filesystem.py | 3 +- .../data_designer/slurm/images/filesystem.py | 7 +- .../data_designer/slurm/images/registry.py | 206 ++++++++--- .../data_designer/slurm/launcher/renderer.py | 11 +- .../data_designer/slurm/runtime/entrypoint.sh | 44 ++- .../data_designer/slurm/services/wiring.py | 321 ++++++++++++------ .../src/data_designer/slurm/state/observer.py | 15 +- .../tests/benchmark/test_workflow.py | 10 +- .../tests/client/test_environment.py | 38 +++ .../tests/client/test_worker.py | 82 +++++ .../tests/images/test_registry_store.py | 157 +++++++++ .../tests/launcher/test_collection.py | 10 +- .../tests/runtime/test_bootstrap.py | 113 +++++- .../tests/services/test_wiring.py | 162 +++++++++ .../golden/rendered/multi_node.sbatch | 1 + .../slurm_test_fakes/test_rendered_scripts.py | 2 +- .../tests/state/test_observer.py | 48 +++ .../data-designer-slurm/tests/test_cli.py | 27 ++ 24 files changed, 1201 insertions(+), 190 deletions(-) diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index 1d78c9853..1bab843a0 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -81,7 +81,7 @@ Removing an alias does not delete its SQSH artifact: 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`, and only when no other registered alias points to the same file. Existing SQSH files registered from outside the workspace are never deleted by this command. Make sure no active run needs the artifact before removing it. +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 diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index 1e369ade2..30ad8d69e 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -87,6 +87,8 @@ The CLI also accepts `--partition gpu-a,gpu-b` during profile initialization. Sl `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. 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 index 5bbe370f2..7e85ea97b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/benchmark/execution.py @@ -34,6 +34,7 @@ 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, @@ -116,6 +117,7 @@ def run( ) from None failures: list[SlurmServiceErrorCode] = [] + first_failure: str | None = None submitted = 0 reused = 0 total = len(compiled.cases) @@ -142,19 +144,32 @@ def run( self._report_progress(f"Benchmark case {index}/{total}: submitted.") except SlurmServiceError as error: failures.append(error.code) - self._report_progress(f"Benchmark case {index}/{total}: submission failed; continuing.") + 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) - self._report_progress(f"Benchmark case {index}/{total}: submission failed; continuing.") + 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, - f"{len(failures)} of {len(compiled.cases)} benchmark child runs could not be submitted", + message[:509] + "..." if len(message) > 512 else message, ) return manifest 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 index b01cc7667..5e8d84c21 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/cli_benchmark.py @@ -12,9 +12,9 @@ import typer from pydantic import BaseModel -from data_designer.slurm.config import load_benchmark_config +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 SlurmServiceOperation +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation _Invoke = Callable[[SlurmServiceOperation, Callable[[], BaseModel]], BaseModel] _Emit = Callable[[BaseModel], None] @@ -43,7 +43,18 @@ def _analyze_benchmark( refresh: bool, fail_if_incomplete: bool, ) -> BaseModel: - benchmark_id = Path(benchmark.rstrip("/")).name + 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, 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 index 958fb4ab0..e588cd11c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -21,12 +21,18 @@ 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 +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) @@ -239,13 +245,18 @@ def _install_overlay( 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") - if not wheels: - return ClientInstallerOutcome.NOT_REQUIRED try: self._command_runner( ( @@ -269,9 +280,56 @@ def _install_overlay( 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: 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 index dd3db08ff..ff9b4edf5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -121,6 +121,10 @@ def __init__( 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, @@ -157,10 +161,30 @@ def best_effort( logger.warning("Could not persist client %s progress", phase.value) def on_batch_complete(self, path: Path) -> None: - generated_records = sum(lazy.pq.read_metadata(item).num_rows for item in path.parent.glob("*.parquet")) + 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(generated_records, self._context.shard.requested_records), + completed_records=min(self._completed_records, self._context.shard.requested_records), ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py b/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py index 99d8e7f11..8dfddb188 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/filesystem.py @@ -24,12 +24,13 @@ def acquire_restrictive_file_lock( 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_EX) + 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 ( 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 index 5baa2c690..c57b285f3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/filesystem.py @@ -27,8 +27,10 @@ @contextmanager -def acquire_file_lock(directory_descriptor: int, name: str, display_path: Path) -> Iterator[None]: - """Acquire one exclusive advisory lock without reopening its parent path.""" +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( @@ -37,6 +39,7 @@ def acquire_file_lock(directory_descriptor: int, name: str, display_path: Path) name, display_path, resource_name="registry", + shared=shared, ) ) except OSError as error: 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 index e1afbf52c..841e9fa9e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/images/registry.py @@ -21,6 +21,7 @@ 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, @@ -40,6 +41,7 @@ _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) @@ -71,6 +73,19 @@ 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 @@ -213,6 +228,10 @@ def unregister(self, name: Identifier, *, delete_artifact: bool = False) -> Regi 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" @@ -233,70 +252,157 @@ def _unregister_with_artifact(self, name: Identifier) -> RegisteredImage: target_lock_name, self._lock_directory / target_lock_name, ), - acquire_file_lock( + ): + # 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) + ): + 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 ( + 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 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) + 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() 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 index 9f3390c05..8b57f5149 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -31,7 +31,7 @@ def render_generation_attempt_script(plan: ResolvedSlurmRunPlan, *, attempt_ordi error_path=posixpath.join(run_root, f"slurm-attempt-{attempt}-%A_%a.err"), ) ) - if plan.selected_profile.profile.gpu_request_mode == "visible": + if _requires_exclusive_nodes(plan): directive_text = f"{directive_text}\n#SBATCH --exclusive" return f"""#!/usr/bin/env bash {directive_text} @@ -106,6 +106,8 @@ def render_generation_retry_script(plan: ResolvedSlurmRunPlan, retry: RetryPlan) 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 @@ -249,3 +251,10 @@ def _build_generation_directives( 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/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh index 05835100d..9904931c7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -13,6 +13,7 @@ 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" @@ -27,6 +28,7 @@ dd_slurm_run_allocation() { 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" @@ -213,6 +215,7 @@ dd_start_servers() { 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 @@ -234,6 +237,7 @@ 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 @@ -242,26 +246,42 @@ dd_start_endpoints() { dd_wait_for_role_readiness() { local role=$1 - local step_id index host port path deadline_seconds deadline + 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 - host=${DD_STEP_PROBE_FIELDS[index]} - port=${DD_STEP_PROBE_FIELDS[index + 1]} - path=${DD_STEP_PROBE_FIELDS[index + 2]} - deadline_seconds=${DD_STEP_PROBE_FIELDS[index + 3]} - deadline=$((SECONDS + deadline_seconds)) - until curl --fail --silent --max-time 1 "http://${host}:${port}${path}" >/dev/null 2>&1; do - dd_require_running - ((SECONDS < deadline)) || { - printf 'runtime step %q readiness timed out\n' "${step_id}" >&2 + 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 } - dd_sleep 0.5 - done + probe_ready[probe]=1 + pending=$((pending - 1)) + fi done + ((pending == 0)) || dd_sleep 0.5 done } diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 21cdd184b..2d90c069b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -16,8 +16,9 @@ from typing import Literal, Protocol, TypeVar from uuid import uuid4 -from pydantic import JsonValue +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 ( @@ -200,10 +201,6 @@ def _prepare( workspace_root = self._profile.profile.workspace_root run_root = Path(workspace_root) / "runs" / run_id client_runtime = self._runtime_inspector.inspect() - deployment_images = tuple( - self._images.resolve_for_planning(deployment.server.image, expected_kind=ImageKind.SERVING) - for deployment in authored.deployments - ) try: builder_payload = _resolve_builder_payload(authored, source_root) except SlurmConfigLoadError as error: @@ -212,6 +209,7 @@ def _prepare( 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, @@ -221,6 +219,18 @@ def _prepare( ) ) 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, @@ -349,121 +359,121 @@ def execute( prepared, force=False, ) - self._report_progress("Rendered batch script; invoking sbatch --parsable --hold (script via stdin).") + 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(publisher, plan, receipt.job_id) + except Exception as error: try: - receipt = self._launcher.submit_script( - prepared.batch_script, - hold=True, - export_environment=export_environment, - ) - except SlurmSubmissionError as error: + self._launcher.cancel(receipt.job_id) + except Exception: raise SlurmServiceError( SlurmServiceErrorCode.UNAVAILABLE, SlurmServiceOperation.EXECUTE_RUN, - _safe_submission_error(error), - ) from None - except SlurmLauncherError: + f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", + ) from error + try: + self._record_submission_failure(publisher, plan, receipt.job_id) + except SlurmServiceError as state_error: raise SlurmServiceError( - SlurmServiceErrorCode.UNAVAILABLE, + SlurmServiceErrorCode.INTERNAL, SlurmServiceOperation.EXECUTE_RUN, - "Slurm submission is unavailable", - ) from None + 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._record_submission(publisher, 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 + self._launcher.cancel(receipt.job_id) + except Exception: + pass + else: try: self._record_submission_failure(publisher, 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: + except BaseException: pass - else: - try: - self._record_submission_failure(publisher, plan, receipt.job_id) - except BaseException: - pass - raise + raise + try: + self._launcher.release(receipt.job_id) + except SlurmLauncherError as error: 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(publisher, 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 + 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 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(), - ) + f"held Slurm job {receipt.job_id} could not be released or cancelled", + ) from error + try: + self._record_submission_failure(publisher, 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 {} +" - ) + ) + 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( - "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, + "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: @@ -975,6 +985,113 @@ def _safe_builder_load_error(error: SlurmConfigLoadError) -> str: 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", + } + unknown_column = isinstance(columns, list) and any( + isinstance(column, dict) + and isinstance(column.get("column_type"), str) + and column["column_type"] not in builtin_columns + for column in columns + ) + unknown_processor = isinstance(processors, list) and any( + isinstance(processor, dict) + and isinstance(processor.get("processor_type"), str) + and processor["processor_type"] not in {"drop_columns", "schema_transform"} + 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 ( + isinstance(column, dict) + and isinstance(column.get("column_type"), str) + and column["column_type"] not in builtin_columns + ) + ] + if plugin_dependency and unknown_processor: + data_designer["processors"] = [ + processor + for processor in processors + if not ( + isinstance(processor, dict) + and isinstance(processor.get("processor_type"), str) + and processor["processor_type"] not in {"drop_columns", "schema_transform"} + ) + ] + 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: " diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py index aa459a6de..c0f8a4f58 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/observer.py @@ -127,7 +127,16 @@ def _observe(self, *, observed_at: datetime | None, persist: bool) -> RunStatus: ) for shard in shards ) - return self._compose_run_status(run, timestamp, shard_statuses) + 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, @@ -242,7 +251,9 @@ def _build_attempt_status( attempt, readiness, scheduler, - current_time=scheduler.observed_at, + 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: diff --git a/packages/data-designer-slurm/tests/benchmark/test_workflow.py b/packages/data-designer-slurm/tests/benchmark/test_workflow.py index dcc6a77c0..6fb48843d 100644 --- a/packages/data-designer-slurm/tests/benchmark/test_workflow.py +++ b/packages/data-designer-slurm/tests/benchmark/test_workflow.py @@ -420,10 +420,16 @@ def test_partial_submission_attempts_all_children_and_resumes_only_missing( progress=progress.append, ) - with pytest.raises(SlurmServiceError, match="1 of 2"): + with pytest.raises(SlurmServiceError, match="1 of 2") as failure: backend.run(config, source_root=tmp_path, force=True) assert len(calls) == 2 - assert "Benchmark case 1/2: submission failed; continuing." in progress + 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() diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index 9c16e0767..7146a5c10 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -444,6 +444,44 @@ def inventory(path: Path | None) -> tuple[InstalledDistribution, ...]: 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, diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 045ed7f5f..a3bb85ffe 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -16,6 +16,7 @@ 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 @@ -186,6 +187,87 @@ def test_run_persists_semantic_result_and_candidate( 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"}) diff --git a/packages/data-designer-slurm/tests/images/test_registry_store.py b/packages/data-designer-slurm/tests/images/test_registry_store.py index fbb36b0bb..2568dd656 100644 --- a/packages/data-designer-slurm/tests/images/test_registry_store.py +++ b/packages/data-designer-slurm/tests/images/test_registry_store.py @@ -151,6 +151,163 @@ def test_unregister_can_delete_unshared_imported_artifact(tmp_path: Path) -> Non 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" diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 23d5012a9..861bc1731 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -181,9 +181,14 @@ def test_retry_renderer_waits_for_persisted_attempt_before_starting_runtime( assert subprocess.run(("bash", "-n"), input=script, text=True, check=False).returncode == 0 -def test_retry_renderer_uses_profile_slurm_bin_path(multi_node_plan: ResolvedSlurmRunPlan) -> None: +@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}) + 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, @@ -208,6 +213,7 @@ def test_retry_renderer_uses_profile_slurm_bin_path(multi_node_plan: ResolvedSlu 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 diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 853ec076d..a4ea426a6 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -88,7 +88,7 @@ def test_multi_replica_manifest_uses_one_gpu_owned_worker( "resources": placement.authored.resources.model_copy(update={"nodes": nodes}), "server": placement.authored.server.model_copy( update={ - "startup_timeout": "2s", + "startup_timeout": "10s", "distributed_init_timeout": "1s", "lead_boot_standoff": "4s", "rank_launch_stagger": f"{stagger}s", @@ -234,7 +234,7 @@ def test_bootstrap_manifest_composes_multi_node_workers_and_remote_endpoints( "served_model_name": deployment.served_model_name, "server": deployment.authored.server.model_copy( update={ - "startup_timeout": "2s", + "startup_timeout": "10s", "distributed_init_timeout": "1s", "lead_boot_standoff": "4s", "rank_launch_stagger": "2s", @@ -314,7 +314,7 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path 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 + 2)) + budgets.append((probe.deadline_seconds, process.launch_delay_seconds + 10)) cases.append(f"*://{probe.host}:{probe.port}/*) ((SECONDS >= {process.launch_delay_seconds + 1}));;") entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" command = f""" @@ -334,6 +334,113 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path assert all(actual == expected for actual, expected in budgets) +def test_server_readiness_uses_one_deadline_for_all_probes() -> None: + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + command = f""" +source {shlex.quote(str(entrypoint))} +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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + + assert completed.returncode == 0 + assert "timeout-at:4" in completed.stdout + assert "readiness timed out" in completed.stderr + + +def test_server_readiness_rejects_late_success_after_earlier_probe() -> None: + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + command = f""" +source {shlex.quote(str(entrypoint))} +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() {{ :; }} +curl() {{ SECONDS=5; return 0; }} +if dd_wait_for_role_readiness server; then + echo unexpected-success +else + echo "timeout-at:$SECONDS" +fi +""" + completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + + assert completed.returncode == 0 + assert "timeout-at:5" in completed.stdout + assert "readiness timed out" in completed.stderr + + +def test_server_readiness_rejects_late_success_for_single_probe() -> None: + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + command = f""" +source {shlex.quote(str(entrypoint))} +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); }} +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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + + 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: + entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" + command = f""" +source {shlex.quote(str(entrypoint))} +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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + + 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, diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 62fbf1313..d4a5a02d8 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -5,6 +5,7 @@ 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 @@ -459,6 +460,129 @@ def test_execute_reports_missing_builder_source( 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_columns( + 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"} + ] + 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, @@ -491,6 +615,44 @@ def test_production_wiring_submits_after_publisher_initialization( 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, diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch index 6d45099c0..7a2687367 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/golden/rendered/multi_node.sbatch @@ -10,6 +10,7 @@ #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 diff --git a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py index e6371cb38..d1bf2461e 100644 --- a/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py +++ b/packages/data-designer-slurm/tests/slurm_test_fakes/test_rendered_scripts.py @@ -24,7 +24,7 @@ def test_rendered_script_fixtures_are_pinned_and_bound_to_canonical_plans( _assert_script_matches_plan( multi_node_plan, "multi_node.sbatch", - expected_fixture_sha256="742d7ebc5412fe266a13e047d472aa2a2950d5276cc015fd93771b890a121d27", + expected_fixture_sha256="c185ed76ac5e95e778d2ef8f442c046aa0fb982ef318a4404a9f5124c1a6484d", ) diff --git a/packages/data-designer-slurm/tests/state/test_observer.py b/packages/data-designer-slurm/tests/state/test_observer.py index 8524973ac..20436a0b9 100644 --- a/packages/data-designer-slurm/tests/state/test_observer.py +++ b/packages/data-designer-slurm/tests/state/test_observer.py @@ -828,6 +828,54 @@ def query_accounting(self, selectors: Sequence[SchedulerJobIdentity]) -> tuple[S 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, diff --git a/packages/data-designer-slurm/tests/test_cli.py b/packages/data-designer-slurm/tests/test_cli.py index b6f8a8ad5..9b491f6fb 100644 --- a/packages/data-designer-slurm/tests/test_cli.py +++ b/packages/data-designer-slurm/tests/test_cli.py @@ -6,6 +6,7 @@ import json import re from pathlib import Path +from types import SimpleNamespace import pytest from click.testing import CliRunner @@ -502,6 +503,11 @@ def create_benchmark_service(**kwargs: object) -> _BenchmarkService: 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( @@ -522,6 +528,27 @@ def create_benchmark_service(**kwargs: object) -> _BenchmarkService: 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) From 07a0267b8149cbfc85976f8978c68f8612313303 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 16 Oct 2026 15:38:09 -0600 Subject: [PATCH 59/62] refactor(slurm): simplify preflight and run preparation (#1047) * refactor(slurm): simplify preflight and readiness tests Signed-off-by: Nabin Mulepati * test(slurm): reuse readiness command helper Signed-off-by: Nabin Mulepati * refactor(slurm): simplify run preparation ownership Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/services/wiring.py | 67 +++++--------- .../tests/runtime/test_bootstrap.py | 90 ++++++++----------- .../tests/services/test_wiring.py | 3 +- 3 files changed, 60 insertions(+), 100 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 2d90c069b..06f421c52 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -175,12 +175,7 @@ def prepare( ) -> Iterator[_PreparedRun]: stack = ExitStack() try: - prepared = self._prepare(authored, source_root=source_root, operation=operation, stack=stack) - except BaseException: - stack.close() - raise - try: - yield prepared + yield self._prepare(authored, source_root=source_root, operation=operation, stack=stack) finally: stack.close() @@ -352,13 +347,7 @@ def execute( ) self._materialize_default_managed_assets(config, plan) export_environment = self._build_export_environment(config) - publisher = self._publisher - self._initialize_run( - publisher, - config, - prepared, - force=False, - ) + self._initialize_run(config, prepared) self._report_progress("Rendered batch script; invoking sbatch --parsable --hold (script via stdin).") try: receipt = self._launcher.submit_script( @@ -379,7 +368,7 @@ def execute( "Slurm submission is unavailable", ) from None try: - self._record_submission(publisher, plan, receipt.job_id) + self._record_submission(plan, receipt.job_id) except Exception as error: try: self._launcher.cancel(receipt.job_id) @@ -390,7 +379,7 @@ def execute( f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", ) from error try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -405,7 +394,7 @@ def execute( pass else: try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except BaseException: pass raise @@ -421,7 +410,7 @@ def execute( f"held Slurm job {receipt.job_id} could not be released or cancelled", ) from error try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -532,19 +521,16 @@ def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str def _initialize_run( self, - publisher: SlurmRunArtifactPublisher, config: DataDesignerSlurmConfig, prepared: _PreparedRun, - *, - force: bool, ) -> None: try: - publisher.initialize( + self._publisher.initialize( config, prepared.plan, prepared.dependencies, prepared.builder_payload, - force=force, + force=False, ) except StateNotFoundError: raise SlurmServiceError( @@ -567,12 +553,11 @@ def _initialize_run( def _record_submission( self, - publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan, job_id: int, ) -> None: try: - publisher.record_submission(plan, job_id, submitted_at=self._clock()) + self._publisher.record_submission(plan, job_id, submitted_at=self._clock()) except StateNotFoundError: raise SlurmServiceError( SlurmServiceErrorCode.NOT_FOUND, @@ -594,12 +579,11 @@ def _record_submission( def _record_submission_failure( self, - publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan, job_id: int, ) -> None: try: - publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) + self._publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) except (StateNotFoundError, StateConflictError, SlurmStateError): raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -1020,17 +1004,18 @@ def _validate_builder_for_submission( "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( - isinstance(column, dict) - and isinstance(column.get("column_type"), str) - and column["column_type"] not in builtin_columns - for column in columns + is_unrecognized_type(column, "column_type", builtin_columns) for column in columns ) unknown_processor = isinstance(processors, list) and any( - isinstance(processor, dict) - and isinstance(processor.get("processor_type"), str) - and processor["processor_type"] not in {"drop_columns", "schema_transform"} - for processor in processors + is_unrecognized_type(processor, "processor_type", builtin_processors) for processor in processors ) unknown_seed_source = ( isinstance(seed_source, dict) @@ -1050,23 +1035,13 @@ def _validate_builder_for_submission( data_designer = dict(data_designer) if plugin_dependency and unknown_column: data_designer["columns"] = [ - column - for column in columns - if not ( - isinstance(column, dict) - and isinstance(column.get("column_type"), str) - and column["column_type"] not in builtin_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 ( - isinstance(processor, dict) - and isinstance(processor.get("processor_type"), str) - and processor["processor_type"] not in {"drop_columns", "schema_transform"} - ) + 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"] = { diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index a4ea426a6..5c92d5a38 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -316,9 +316,7 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path ) budgets.append((probe.deadline_seconds, process.launch_delay_seconds + 10)) cases.append(f"*://{probe.host}:{probe.port}/*) ((SECONDS >= {process.launch_delay_seconds + 1}));;") - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" command = f""" -source {shlex.quote(str(entrypoint))} DD_RUNTIME_MANIFEST={shlex.quote(str(manifest_path))} SECONDS=0 dd_start_step() {{ DD_LAST_PID=1; }} @@ -329,70 +327,58 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path dd_start_servers dd_wait_for_role_readiness server """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0, completed.stderr assert all(actual == expected for actual, expected in budgets) -def test_server_readiness_uses_one_deadline_for_all_probes() -> None: +def _run_readiness_command(body: str) -> subprocess.CompletedProcess[str]: entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} -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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) - - assert completed.returncode == 0 - assert "timeout-at:4" in completed.stdout - assert "readiness timed out" in completed.stderr + return subprocess.run( + ("bash", "-c", f"source {shlex.quote(str(entrypoint))}\n{body}"), + capture_output=True, + text=True, + timeout=10, + ) -def test_server_readiness_rejects_late_success_after_earlier_probe() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} +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() {{ :; }} -curl() {{ SECONDS=5; return 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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0 - assert "timeout-at:5" in completed.stdout + assert "timeout-at:4" in completed.stdout assert "readiness timed out" in completed.stderr -def test_server_readiness_rejects_late_success_for_single_probe() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" +@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""" -source {shlex.quote(str(entrypoint))} 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); }} +dd_read_step() {{ DD_STEP_PROBE_FIELDS=({probe_fields}); }} dd_require_running() {{ :; }} dd_sleep() {{ :; }} curl() {{ SECONDS=5; return 0; }} @@ -402,7 +388,7 @@ def test_server_readiness_rejects_late_success_for_single_probe() -> None: echo "timeout-at:$SECONDS" fi """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0 assert "timeout-at:5" in completed.stdout @@ -410,32 +396,30 @@ def test_server_readiness_rejects_late_success_for_single_probe() -> None: def test_server_readiness_checks_later_steps_before_earlier_steps_finish() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} + 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() {{ +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 +} +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 = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0, completed.stderr assert "ready-at:110" in completed.stdout diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index d4a5a02d8..042152a50 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -491,7 +491,7 @@ def test_execute_rejects_invalid_builtin_builder_before_submission( assert not launcher.submissions -def test_submit_validation_defers_overlay_plugin_columns( +def test_submit_validation_defers_overlay_plugin_components( authored_run_single: DataDesignerSlurmConfig, ) -> None: payload = authored_run_single.model_dump(mode="json") @@ -499,6 +499,7 @@ def test_submit_validation_defers_overlay_plugin_columns( 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 From f5074a3a1fdd3d86705a6a0d890a2d7059634642 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Wed, 21 Oct 2026 10:54:27 -0600 Subject: [PATCH 60/62] fix(slurm): tag submitted jobs for telemetry (#1060) * fix(slurm): tag submitted jobs for telemetry Signed-off-by: Nabin Mulepati * fix(slurm): honor user deployment type during submission Signed-off-by: Nabin Mulepati * feat(telemetry): recognize internal Slurm deployments Signed-off-by: Nabin Mulepati * fix(slurm): preserve deployment tags across job operations Signed-off-by: Nabin Mulepati * fix(slurm): align fallback tag across initial and retry jobs Signed-off-by: Nabin Mulepati --------- Signed-off-by: Nabin Mulepati --- .../data_designer/engine/models/telemetry.py | 2 + .../tests/engine/models/test_telemetry.py | 22 +++++++-- .../data_designer/slurm/launcher/client.py | 13 ++--- .../slurm/services/retry_collection.py | 17 +++++-- .../data_designer/slurm/services/wiring.py | 9 ++++ .../data_designer/slurm/state/collection.py | 13 ++++- .../src/data_designer/slurm/state/retry.py | 14 ++++-- .../tests/images/test_lifecycle.py | 15 ++++-- .../tests/launcher/test_client.py | 41 ++++++++++++---- .../tests/services/test_wiring.py | 48 ++++++++++++++++++- .../tests/state/test_retry_collection.py | 13 ++++- 11 files changed, 173 insertions(+), 34 deletions(-) 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/src/data_designer/slurm/launcher/client.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py index e5f3fa6ad..a872daa19 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/client.py @@ -288,17 +288,18 @@ def _run( return stdout -def _format_export_environment(environment: Mapping[str, str] | None) -> tuple[str, Mapping[str, str] | None]: - if not environment: - return "--export=NIL", None - names = tuple(sorted(environment)) +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 = environment[name] + value = exported[name] if type(value) is not str or "\0" in value: raise ValueError("exported environment values must be strings without NUL") - return f"--export={','.join(names)}", environment + return "--export=ALL", exported def _format_selectors(selectors: Sequence[SchedulerJobIdentity]) -> str: 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 index 85c12a932..51d01b98d 100644 --- 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 @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from pathlib import Path from typing import Literal @@ -38,10 +38,17 @@ class RunRetryCollectionBackend: """Adapt public run operations to persisted retry and collection capabilities.""" - def __init__(self, workspace_root: str, launcher: SlurmCommandClient, clock: Callable[[], datetime]) -> None: + 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, @@ -56,7 +63,7 @@ def retry( 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) + 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) @@ -151,7 +158,9 @@ def collect( operation, "num_partitions must not exceed requested records", ) - status = SlurmCollectionCoordinator(self._workspace_root, run_id, self._launcher).submit( + status = SlurmCollectionCoordinator( + self._workspace_root, run_id, self._launcher, self._export_environment + ).submit( destination=destination, num_partitions=num_partitions, submitted_at=self._clock(), diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 06f421c52..a29a7674c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -322,6 +322,7 @@ def __init__( selected_profile.profile.workspace_root, launcher, clock, + source_environment, ) def execute( @@ -492,6 +493,14 @@ def _materialize_default_managed_assets( 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: 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 index 85739aa54..3028107c5 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/collection.py @@ -6,6 +6,7 @@ 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 @@ -52,7 +53,9 @@ class CollectionScheduler(SchedulerObservationClient, Protocol): """Scheduler operations required by collection submission and refresh.""" - def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + def submit_script( + self, script: str, *, export_environment: Mapping[str, str] | None = None + ) -> SlurmJobSubmissionReceipt: """Submit one rendered CPU collection job.""" ... @@ -74,9 +77,11 @@ def __init__( 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) @@ -242,7 +247,11 @@ def _submit_prepared( submitted_at: datetime, ) -> CollectionStatus: try: - receipt = self._scheduler.submit_script(script) + 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( 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 index a8b9188cd..8faef34c3 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Iterator, Sequence +from collections.abc import Iterator, Mapping, Sequence from contextlib import ExitStack, contextmanager from datetime import datetime, timezone from pathlib import Path @@ -48,7 +48,9 @@ class RetryScheduler(SchedulerObservationClient, Protocol): """Scheduler operations required by fresh-process retry.""" - def submit_script(self, script: str) -> SlurmJobSubmissionReceipt: + def submit_script( + self, script: str, *, export_environment: Mapping[str, str] | None = None + ) -> SlurmJobSubmissionReceipt: """Submit one rendered retry array.""" ... @@ -70,9 +72,11 @@ def __init__( 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) @@ -337,7 +341,11 @@ def _persist_prepared_retry(self, plan: RetryPlan, timestamp: datetime) -> None: def _submit(self, plan: RetryPlan, script: str, timestamp: datetime) -> SlurmJobSubmissionReceipt: try: - receipt = self._scheduler.submit_script(script) + 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) diff --git a/packages/data-designer-slurm/tests/images/test_lifecycle.py b/packages/data-designer-slurm/tests/images/test_lifecycle.py index d428e7445..8216b4830 100644 --- a/packages/data-designer-slurm/tests/images/test_lifecycle.py +++ b/packages/data-designer-slurm/tests/images/test_lifecycle.py @@ -9,7 +9,7 @@ import shutil import stat import subprocess -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from pathlib import Path @@ -513,7 +513,7 @@ def test_submit_prepared_image_lifecycle_uses_isolated_sbatch_client(tmp_path: P ( "sbatch", "--parsable", - "--export=NIL", + "--export=ALL", ) ] assert runner.inputs == [render_image_lifecycle_script(prepared.plan)] @@ -575,7 +575,10 @@ def test_submit_rejects_modified_script_even_when_artifact_digest_is_rebound(tmp submit_prepared_image_lifecycle(rebound, SlurmCommandClient(FakeSlurmRunner())) -def test_submit_uses_verified_script_bytes_when_path_changes_during_submission(tmp_path: Path) -> None: +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"), @@ -587,7 +590,8 @@ def test_submit_uses_verified_script_bytes_when_path_changes_during_submission(t receipt = submit_prepared_image_lifecycle(prepared, SlurmCommandClient(runner)) assert receipt.job_id == 5101 - assert runner.command == ("sbatch", "--parsable", "--export=NIL") + 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" @@ -1683,15 +1687,18 @@ 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/launcher/test_client.py b/packages/data-designer-slurm/tests/launcher/test_client.py index 7beb3225e..49f124289 100644 --- a/packages/data-designer-slurm/tests/launcher/test_client.py +++ b/packages/data-designer-slurm/tests/launcher/test_client.py @@ -24,7 +24,7 @@ def test_client_submits_and_observes_one_managed_array(fake_slurm_runner: FakeSl 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=NIL", "/workspace/run.sbatch"), + ("sbatch", "--parsable", "--export=ALL", "/workspace/run.sbatch"), ("squeue", "--noheader", "--array", "--format=%i|%T", "--jobs=4101"), ] @@ -37,7 +37,7 @@ def test_client_submits_verified_script_text_through_standard_input() -> None: submission = client.submit_script(script) assert submission.job_id == 5101 - assert runner.calls == [("sbatch", "--parsable", "--export=NIL")] + assert runner.calls == [("sbatch", "--parsable", "--export=ALL")] assert runner.inputs == [script] @@ -50,29 +50,54 @@ def test_client_holds_submission_until_initial_state_is_recorded() -> None: client.release(submission.job_id) assert runner.calls == [ - ("sbatch", "--parsable", "--hold", "--export=NIL"), + ("sbatch", "--parsable", "--hold", "--export=ALL"), ("scontrol", "release", "5101"), ] -def test_client_exports_only_explicit_environment_names_without_values_in_argv() -> None: +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", "SLURM_EXPORT_ENV": "ALL"}, + export_environment={"HF_TOKEN": "secret-value", "NEMO_DEPLOYMENT_TYPE": "library", "SLURM_EXPORT_ENV": "ALL"}, ) assert runner.command == ( "sbatch", "--parsable", - "--export=HF_TOKEN,SLURM_EXPORT_ENV", + "--export=ALL", ) - assert runner.environment == {"HF_TOKEN": "secret-value", "SLURM_EXPORT_ENV": "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)) @@ -744,7 +769,7 @@ def test_script_path_is_one_argument_vector_token(fake_slurm_runner: FakeSlurmRu assert fake_slurm_runner.calls[0] == ( "sbatch", "--parsable", - "--export=NIL", + "--export=ALL", "/workspace/run; touch injected.sbatch", ) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 042152a50..5da340517 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -589,7 +589,9 @@ def test_production_wiring_submits_after_publisher_initialization( 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() @@ -612,7 +614,7 @@ def test_production_wiring_submits_after_publisher_initialization( assert len(launcher.submissions) == 1 assert launcher.held_submissions == [True] assert launcher.releases == [42] - assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL"}] + assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": "slurm"}] assert (tmp_path / "managed-assets").is_dir() @@ -680,7 +682,33 @@ def test_production_wiring_exports_referenced_secrets_to_the_allocation( service.execute(authored, source_root=tmp_path) - assert launcher.exported_environments == [{"SLURM_EXPORT_ENV": "ALL", "SOURCE_TOKEN": "secret-value"}] + 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( @@ -1025,12 +1053,23 @@ def test_status_does_not_expire_requeue_window_from_another_attempt_clock( 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}), @@ -1046,6 +1085,7 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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" @@ -1079,6 +1119,10 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( 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( diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 15910b9f5..555ed2488 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -15,6 +15,7 @@ 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 @@ -744,12 +745,20 @@ def test_collection_submits_cpu_job_and_publishes_ordered_winners_atomically( 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)) + scheduler = SlurmCommandClient(runner) + coordinator = SlurmCollectionCoordinator( + case.workspace, + case.plan.run_id, + scheduler, + {"NEMO_DEPLOYMENT_TYPE": "slurm-nvidia-internal"}, + ) - submitted = coordinator.submit(submitted_at=case.created_at + timedelta(minutes=10)) + 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]) From 9ba05f5451f99df65dbf3835916139e23ad523e4 Mon Sep 17 00:00:00 2001 From: Andre Manoel <165937436+andreatnvidia@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:49:49 -0300 Subject: [PATCH 61/62] fix(slurm): speed up allocation artifact checks (#1071) Signed-off-by: Andre Manoel --- .../data_designer/slurm/runtime/entrypoint.py | 1 + .../data_designer/slurm/runtime/entrypoint.sh | 11 ----- .../slurm/runtime/plan_reader.sh | 20 -------- .../data_designer/slurm/runtime/preflight.py | 15 ++++-- .../tests/runtime/test_entrypoint.py | 23 ++++++++++ .../tests/runtime/test_preflight.py | 46 +++++++++++++++++++ 6 files changed, 81 insertions(+), 35 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index c6672b8c6..ce4897d65 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -101,6 +101,7 @@ def _prepare(arguments: argparse.Namespace, environment: Mapping[str, str]) -> N 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) 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 index 9904931c7..edac5241e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -149,17 +149,6 @@ dd_verify_host_context() { [[ ${SLURM_NODEID:-} == "${DD_CLIENT_NODE_INDEX}" ]] dd_resolve_allocation_hosts dd_verify_gpu_count - dd_read_artifacts "${DD_PLAN_PATH}" "${SLURM_ARRAY_TASK_ID}" - local index path digest actual - for ((index = 0; index < ${#DD_ARTIFACT_FIELDS[@]}; index += 2)); do - path=${DD_ARTIFACT_FIELDS[index]} - digest=${DD_ARTIFACT_FIELDS[index + 1]} - actual=$(sha256sum < "${path}") - [[ ${actual%% *} == "${digest}" ]] || { - printf 'allocation artifact digest mismatch: %s\n' "${path}" >&2 - return 65 - } - done } dd_resolve_allocation_hosts() { 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 index 45fdad59f..f1c78a462 100644 --- 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 @@ -64,26 +64,6 @@ dd_read_container_path() { ' "${plan}") } -dd_read_artifacts() { - local plan=$1 - local task_id=$2 - dd_read_null_values DD_ARTIFACT_FIELDS < <( - jq -j --argjson task_id "${task_id}" ' - [ - .runtime_bundle, - .client.dependency_lock, - (.deployments[] | {path: .image.path, sha256: .image.sha256}), - .builder.source, - (.shards[] | select(.array_task_index == $task_id) | .input_partition) - ] - | map(select(. != null)) - | unique_by([.path, .sha256]) - | .[] - | .path, "\u0000", .sha256, "\u0000" - ' "${plan}" - ) -} - dd_read_plan_secret_names() { local plan=$1 dd_read_null_values DD_REQUIRED_SECRET_NAMES < <( 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 index 480759039..1086a9f6d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/preflight.py @@ -93,7 +93,7 @@ def verify(self, context: AllocationContext, environment: Mapping[str, str]) -> attempt_directory = context.attempt_directory self.verify_attempt_directory(attempt_directory) self.verify_client_runtime(context) - self._verify_artifacts(context) + self.verify_artifacts(context) self.verify_ports(context, environment) except SlurmRuntimeError: raise @@ -161,7 +161,8 @@ def verify_attempt_directory(attempt_directory: Path) -> None: ) @staticmethod - def _verify_artifacts(context: AllocationContext) -> None: + def verify_artifacts(context: AllocationContext) -> None: + """Verify unique launch artifacts against their planned SHA-256 digests.""" plan = context.plan references = [ ArtifactReference( @@ -181,8 +182,14 @@ def _verify_artifacts(context: AllocationContext) -> None: ] references.extend(reference for reference in optional_references if reference is not None) unique_references = {(reference.path, reference.sha256): reference for reference in references} - for reference in unique_references.values(): - _verify_artifact(reference) + 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: diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 67e90aaf2..59c34c1f7 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -236,6 +236,28 @@ def test_succeed_phase_stops_runtime_and_finalizes_winner( 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, @@ -276,6 +298,7 @@ def _patch_runtime_context( ) -> 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") diff --git a/packages/data-designer-slurm/tests/runtime/test_preflight.py b/packages/data-designer-slurm/tests/runtime/test_preflight.py index 4c95f24c5..b9dcf5bfe 100644 --- a/packages/data-designer-slurm/tests/runtime/test_preflight.py +++ b/packages/data-designer-slurm/tests/runtime/test_preflight.py @@ -124,6 +124,52 @@ def test_artifact_verification_checks_exact_bytes_and_rejects_symlink(tmp_path: _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, From 4cb13d917d01054d444c6d6e99f5b1ad5bb33e48 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 23 Oct 2026 20:58:07 -0600 Subject: [PATCH 62/62] fix(slurm): remove vLLM release allowlist (#1086) Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/getting-started.mdx | 6 ++-- fern/versions/latest/pages/slurm/images.mdx | 10 ++++--- .../data_designer/slurm/serving/resolver.py | 21 +------------ .../tests/serving/test_resolver.py | 30 +++++++++---------- 4 files changed, 25 insertions(+), 42 deletions(-) diff --git a/fern/versions/latest/pages/slurm/getting-started.mdx b/fern/versions/latest/pages/slurm/getting-started.mdx index b4ec3815f..b631345a8 100644 --- a/fern/versions/latest/pages/slurm/getting-started.mdx +++ b/fern/versions/latest/pages/slurm/getting-started.mdx @@ -8,10 +8,10 @@ Obtain these site-specific values from your Slurm or platform administrator befo - 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 +- An official vLLM release tag, a digest-qualified custom serving OCI image, or an absolute SQSH path containing a vLLM installation compatible with your model - 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. +The site values and `example/generator` model below are placeholders. The official vLLM tag is a real image reference; select a release that supports your model and that your site can access. ## 1. Create and activate an environment @@ -84,7 +84,7 @@ data-designer slurm image add \ --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). +Replace the tag with a release that supports your model and is 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 diff --git a/fern/versions/latest/pages/slurm/images.mdx b/fern/versions/latest/pages/slurm/images.mdx index 1bab843a0..332813799 100644 --- a/fern/versions/latest/pages/slurm/images.mdx +++ b/fern/versions/latest/pages/slurm/images.mdx @@ -4,7 +4,7 @@ Data Designer Slurm images run vLLM model servers. The Data Designer client runs ## 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: +For vLLM, the simplest path is to select a 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 \ @@ -12,9 +12,11 @@ data-designer slurm image add vllm/vllm-openai:v0.22.0 \ --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. +Replace `v0.22.0` with a release that supports your model and 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. +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 vLLM executable and all runtime dependencies for the selected model. + +Slurm records the inspected vLLM version as provenance and does not restrict deployments to specific release lines. Custom and nightly builds can be registered through a digest-qualified OCI reference or an existing SQSH path. Image inspection and startup readiness checks still apply; accepting a version does not establish compatibility with the selected model, launch arguments, hardware, or topology. Validate your image with a small generation run before scaling up. 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. @@ -93,4 +95,4 @@ Published OCI imports are content-addressed. Identical imports can share an arti ## 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. +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 vLLM executable; the inspected runtime version is recorded without a release-line restriction. 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 index 95f038a37..0ae7d305c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/serving/resolver.py @@ -5,8 +5,6 @@ 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 @@ -27,11 +25,6 @@ ) 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.""" @@ -41,7 +34,7 @@ def resolve_vllm_server( plan: ResolvedSlurmRunPlan, deployment_id: Identifier, ) -> ResolvedVllmServerDeployment: - """Resolve one vLLM deployment admitted by the tested runtime matrix.""" + """Resolve one vLLM deployment from its inspected serving image.""" try: return _resolve_vllm_server(plan, deployment_id) except ValidationError as error: @@ -74,7 +67,6 @@ def _resolve_vllm_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") @@ -213,14 +205,3 @@ def _calculate_launch_delay_seconds(server: VllmServerConfig, deployment_replica 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/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index 26532ca01..e98835574 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -122,22 +122,20 @@ def test_two_deployments_keep_images_and_endpoint_identities_isolated( 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"), + ( + "0.20.9", + "0.21.0", + "0.22.0+vendor.1", + "0.22.0rc1", + "0.23.0", + "0.30.0", + "0.1.dev20073+g8e685d198", + "vendor-vllm-build", + ), ) -def test_resolution_rejects_unsupported_runtime_versions( +def test_resolution_preserves_inspected_runtime_version_as_provenance( single_node_plan: ResolvedSlurmRunPlan, runtime_version: str, ) -> None: @@ -146,8 +144,10 @@ def test_resolution_rejects_unsupported_runtime_versions( 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) + resolved = _resolve(plan, placement.deployment_id) + + assert resolved.image.inspection_facts.runtime_version == runtime_version + assert ResolvedVllmServerDeployment.model_validate_json(resolved.model_dump_json()) == resolved def test_resolution_rejects_image_inspection_mismatch(