diff --git a/.agents/skills/cosmos-automl/SKILL.md b/.agents/skills/cosmos-automl/SKILL.md index aa077ca56..1f65c40d6 100644 --- a/.agents/skills/cosmos-automl/SKILL.md +++ b/.agents/skills/cosmos-automl/SKILL.md @@ -6,10 +6,31 @@ description: Optimize Cosmos Framework reasoner SFT hyperparameters or evaluatio # Cosmos Framework AutoML Read [docs/automl.md](../../../docs/automl.md) before planning a search. It owns -the installation/release prerequisites, manifest contract, result semantics +the installation prerequisites, manifest contract, result semantics and current limits. Use the native `cosmos-automl` CLI; do not call a TAO runner -or introduce model/runtime selection. The engine is an optional wheel, not a -source-checkout dependency. +or introduce model/runtime selection. Install the optional engine from the +`packages/automl-core` subdirectory of `NVIDIA-TAO/tao-automl`; no TAO application, +SDK, persistent source checkout, or publicly published wheel is required. + +Before a new installation, ask: **"Use AutoML core from `tao-automl/main`, or +do you want a development branch from that repository? If so, which branch?"** +Use `main` unless the user selects another branch; if they already supplied a +branch, confirm that selection without asking again. Follow the +[source-install instructions](../../../docs/automl.md#installation-from-source), +including the named-branch override. Do not substitute a development branch +when the selected source is unavailable. Record the resolved commit, installed +package version and `engine_build`; a branch name alone is not a build identity. +For an existing run, inspect and preserve its original environment/build instead +of reinstalling the latest branch head. Source selection does not authorize +upgrading a live controller or migrating a stopped experiment. + +Follow the [fresh-machine checklist](../../../docs/automl.md#fresh-machine-reproducibility-checklist): +record Framework/core commits, the actual installed inventory, data/model +identities and reviewed recipe/budgets. Preflight the evaluator and checkpoint +export as well as training. Prefer local POSIX workspaces for local runs; shared +storage needs concurrent atomic-read/replacement and locking checks. Report +dependency warnings accurately; a successful training smoke is not a clean +full-environment dependency check. Reuse **cosmos3-setup** for environment preparation and **cosmos3-post-training** for a prepared reasoner recipe and dataset/checkpoint @@ -59,11 +80,19 @@ Keep this one skill entrypoint for both workflows. - Supervise the native CPU allocation before the brain starts; trial retries cannot handle bootstrap/import/prolog failures. Follow the SLURM guide's [startup contract](../../../docs/automl-slurm.md#before-the-search-controller-starts): + when startup recovery is required, explicitly configure `AUTOML_STARTUP_HELPER`, + its registered attempt base and the original deadline. Leaving the helper unset + preserves the direct-CLI launch path, but does not supply bootstrap validation + or pre-controller recovery; do not claim those checks ran. For the supervised path, arm the cross-node observer automatically before submission, preserve startup receipts, and mark the boundary before any algorithm/API/job fan-out. Only confirmed pre-fan-out infrastructure/observer failures may use approved CPU startup retries. Unknown submission outcomes block replacement. A missing experiment ledger is not success; report unlaunched modes as `NOT_STARTED`. + Use the guide's persistent observer service, verify `ready` before submission, + and register every acknowledged CPU job with its remaining allocation cap. + One-shot ticks or chat polling do not cover later requeue generations. The + watcher never replaces unknown submissions or extends the original deadline. - If the CPU allocation can expire first, review the SLURM guide's opt-in `--controller-max-allocations`, allocation seconds and grace in both plan and run. CPU handoff requeues the same CPU job and adopts existing GPU jobs; it @@ -131,8 +160,16 @@ Keep this one skill entrypoint for both workflows. A `scheduler_unavailable` observation is not a failed trial: the controller retries read-only SLURM queries within the original deadline. Do not wrap a transient status error in blanket cancellation or restart a paid request. + Timed-out read-only preflight checks also wait within the original limits; + inspect `_scheduler_check.json`. A fully sealed, unsubmitted candidate can + continue across handoff; an attempted submission still requires reconciliation. If cancellation is unconfirmed, reconcile native jobs and worker receipts; neither a stale ledger nor a stop marker proves that GPUs have been released. + A local venv controller can exit while its worker continues. Reconcile both + before restarting; never upgrade its environment while workers are active. + Same-host venv cold migration is a separately approved, terminal-worker-only + operation preserving results, inventory, remaining budgets and deadline; see + [local recovery](../../../docs/automl.md#results-monitoring-and-recovery). - Hand off `best.json`, the winning native TOML, DCP and held-out metrics, or report that no trial succeeded. Do not call validation improvement final test performance or describe fixture checks as GPU end-to-end verification. @@ -145,8 +182,9 @@ searches also passed bounded GPU smokes. Six of nine full-data numerical/populat searches completed; Bayesian, Hyperband and DEHB did not complete their approved budgets. Live infrastructure recovery and Hyperband-ES curve pruning remain unverified; smoke success is not production-scale or model-quality validation. -Keep the core at `1.0.0.dev0` during the port; local builds are not releases. -Use wheel hashes and `engine_build` to distinguish builds, not version bumps. +Use the core package version required by Framework's extras (currently `0.1.0`); +source installs and local builds are not releases. Use resolved Git commits, +wheel hashes when applicable, and `engine_build` to distinguish builds, not version bumps. Preserve exact old wheels/manifests for paused searches; changed builds fail resume identity checks even when package versions match. Route ordinary training/evaluation to the existing skills, and do not infer diff --git a/.agents/skills/cosmos3-setup/SKILL.md b/.agents/skills/cosmos3-setup/SKILL.md index e1160f0c9..6795c440e 100644 --- a/.agents/skills/cosmos3-setup/SKILL.md +++ b/.agents/skills/cosmos3-setup/SKILL.md @@ -51,6 +51,11 @@ The canonical setup reference is `docs/setup.md`. The README (`README.md` § Set - **CUDA compatibility**: compare the host driver with the selected CUDA build's requirements. `nvidia-smi` reports driver capability, not the installed toolkit; its CUDA version need not equal `torch.version.cuda`. - **`HF_HOME`**: controls where checkpoints are cached (default: `~/.cache/huggingface`). Set this if disk space is tight or you want a shared cache. - **Conflicting env vars**: stale `HF_TOKEN` or `HUGGING_FACE_HUB_TOKEN` env vars can silently override CLI auth. Check with `printenv | grep HF_`. +- **AutoML is a separate source install**: `uv sync --all-extras` does not install + the core engine. For AutoML setup, follow `docs/automl.md#installation-from-source` + and the `cosmos-automl` skill: ask whether to use `NVIDIA-TAO/tao-automl`'s `main` + branch or a user-selected development branch before a new installation. Keep + the recorded source revision for existing experiments; do not upgrade on resume. ## Related skills diff --git a/.claude/skills/cosmos3-setup/SKILL.md b/.claude/skills/cosmos3-setup/SKILL.md index 646ccc629..10c45a89e 100644 --- a/.claude/skills/cosmos3-setup/SKILL.md +++ b/.claude/skills/cosmos3-setup/SKILL.md @@ -51,6 +51,11 @@ The canonical setup reference is `docs/setup.md`. The README (`README.md` § Set - **CUDA version alignment**: the major CUDA version from `nvidia-smi` must match `torch.version.cuda`. Mismatches cause cryptic shared-library errors. - **`HF_HOME`**: controls where checkpoints are cached (default: `~/.cache/huggingface`). Set this if disk space is tight or you want a shared cache. - **Conflicting env vars**: stale `HF_TOKEN` or `HUGGING_FACE_HUB_TOKEN` env vars can silently override CLI auth. Check with `printenv | grep HF_`. +- **AutoML is a separate source install**: `uv sync --all-extras` does not install + the core engine. For AutoML setup, follow `docs/automl.md#installation-from-source` + and the `cosmos-automl` skill: ask whether to use `NVIDIA-TAO/tao-automl`'s `main` + branch or a user-selected development branch before a new installation. Keep + the recorded source revision for existing experiments; do not upgrade on resume. ## Related skills diff --git a/.github/workflows/automl-helpers.yml b/.github/workflows/automl-helpers.yml new file mode 100644 index 000000000..974b434a8 --- /dev/null +++ b/.github/workflows/automl-helpers.yml @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +name: AutoML helpers +on: + pull_request: + push: + branches: [main] + paths: + - 'tools/cosmos_workflows/**' + - 'examples/automl/**' + - '.github/workflows/automl-helpers.yml' + - '.pytest.toml' +permissions: + contents: read +jobs: + cpu-regression: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: '3.12' + - run: python -m pip install pytest jsonschema + - name: Test execution helpers without GPU dependencies or scheduler access + run: >- + python -m pytest -o addopts= --confcutdir=tools/cosmos_workflows + -o 'python_files=test_*.py *_test.py' tools/cosmos_workflows diff --git a/README.md b/README.md index 156006bb7..6d3898dfa 100644 --- a/README.md +++ b/README.md @@ -143,14 +143,14 @@ invokes automatically when the user's request matches its description. Skills live in [`.agents/skills/`](./.agents/skills) (canonical) and are mirrored under [`.claude/skills/`](./.claude/skills) for Claude Code: -| Skill | When it activates | -| -------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | -| [`cosmos3-setup`](./.agents/skills/cosmos3-setup/SKILL.md) | Installation, environment setup, checkpoint downloads, Docker. | -| [`cosmos3-codebase-nav`](./.agents/skills/cosmos3-codebase-nav/SKILL.md) | "Where is X" / "where do I change parameter Y" questions across `cosmos_framework/`. | -| [`cosmos3-inference`](./.agents/skills/cosmos3-inference/SKILL.md) | Running offline or online inference, parallelism, sampling parameters. | -| [`cosmos3-post-training`](./.agents/skills/cosmos3-post-training/SKILL.md) | SFT post-training end-to-end: data prep, DCP conversion, launch, export. | -| [`cosmos-automl`](./.agents/skills/cosmos-automl/SKILL.md) | Reasoner hyperparameter search and GEPA prompt optimization; [setup and release status](./docs/automl.md). | -| [`cosmos3-env-troubleshoot`](./.agents/skills/cosmos3-env-troubleshoot/SKILL.md) | Diagnosing install/runtime errors (ImportError, CUDA, Docker, checkpoint failures). | +| Skill | When it activates | +| -------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | +| [`cosmos3-setup`](./.agents/skills/cosmos3-setup/SKILL.md) | Installation, environment setup, checkpoint downloads, Docker. | +| [`cosmos3-codebase-nav`](./.agents/skills/cosmos3-codebase-nav/SKILL.md) | "Where is X" / "where do I change parameter Y" questions across `cosmos_framework/`. | +| [`cosmos3-inference`](./.agents/skills/cosmos3-inference/SKILL.md) | Running offline or online inference, parallelism, sampling parameters. | +| [`cosmos3-post-training`](./.agents/skills/cosmos3-post-training/SKILL.md) | SFT post-training end-to-end: data prep, DCP conversion, launch, export. | +| [`cosmos-automl`](./.agents/skills/cosmos-automl/SKILL.md) | Reasoner hyperparameter search and GEPA prompt optimization; [main/development-branch setup](./docs/automl.md#installation-from-source). | +| [`cosmos3-env-troubleshoot`](./.agents/skills/cosmos3-env-troubleshoot/SKILL.md) | Diagnosing install/runtime errors (ImportError, CUDA, Docker, checkpoint failures). | See [`AGENTS.md`](./AGENTS.md) for the canonical repo map that agents load first; the skills above are referenced from there. diff --git a/cosmos_framework/scripts/automl.py b/cosmos_framework/scripts/automl.py index bc8e948d0..20da293f0 100644 --- a/cosmos_framework/scripts/automl.py +++ b/cosmos_framework/scripts/automl.py @@ -12,10 +12,14 @@ def main() -> None: if exc.name not in {"automl_core", "automl_core.cli"}: raise raise SystemExit( - "AutoML requires the separately supplied nvidia-automl-core 1.0.0.dev0 wheel. " - "Install that wheel into this Python environment first, then install " - "cosmos-framework[automl,reasoner-eval]. See docs/automl.md for the local-wheel " - "installation workflow (including uv run --no-sync)." + "AutoML requires nvidia-automl-core==0.1.0. Install packages/automl-core from " + "NVIDIA-TAO/tao-automl (main by default, or your selected development branch) " + "into this Python environment, then install cosmos-framework[automl,reasoner-eval]. " + "Git access to the selected repository/revision is required; if unavailable, " + "use a reviewed compatible core wheel supplied by your project maintainer. " + "Record the resolved commit and engine_build, and reuse them for existing runs. " + "See docs/automl.md#installation-from-source for source-install commands, " + "compatible versions, and uv run --no-sync. No published wheel is required." ) from exc if len(sys.argv) > 1 and sys.argv[1] == "prompts": from cosmos_framework.automl.prompts import load diff --git a/docs/automl-algorithms.md b/docs/automl-algorithms.md index ebc6b3e99..d21c886d6 100644 --- a/docs/automl-algorithms.md +++ b/docs/automl-algorithms.md @@ -1,6 +1,7 @@ # AutoML algorithm contracts -Install the local core 1.0.0.dev0 wheel as described in [AutoML](automl.md). Choose +Install the core from `tao-automl/main` or a selected development branch as +described in [AutoML setup](automl.md#installation-from-source). Choose `search.algorithm` in the same native reasoner manifest; there are no competing launchers or target selectors. All modes use explicit objective direction, declared parameter domains, recorded native jobs and bounded sequential execution. diff --git a/docs/automl-platforms.md b/docs/automl-platforms.md index ce4a180e5..d8a9bc2e1 100644 --- a/docs/automl-platforms.md +++ b/docs/automl-platforms.md @@ -1,6 +1,6 @@ # AutoML on Kubernetes, Brev and Lepton -These are execution adapters in the optional `nvidia-automl-core` wheel, used by +These are execution adapters in the optional `nvidia-automl-core` package, used by the existing `cosmos-automl` and `cosmos-automl prompts` commands. They do not call a TAO runner, require a source checkout, or introduce another search loop. All twelve SFT modes and GEPA use the same workload, scoring and promotion diff --git a/docs/automl-prompts.md b/docs/automl-prompts.md index 4be31fa96..288064a2a 100644 --- a/docs/automl-prompts.md +++ b/docs/automl-prompts.md @@ -6,15 +6,24 @@ weights and is not an SFT search algorithm. The reusable implementation lives in `automl_core`; Framework supplies native annotations, evaluation and metrics. Users need no TAO source checkout or SDK. -Install the supplied core 1.0.0.dev0 wheel and Framework's `automl-prompts` extra: +Follow [source installation](automl.md#installation-from-source), using `main` by +default or an explicitly selected development branch. For a pip-managed, +prepared Framework environment without a previous core installation: ```bash -python -m pip install '/path/to/nvidia_automl_core-1.0.0.dev0-py3-none-any.whl[prompts]' -python -m pip install -e '.[automl-prompts,reasoner-eval]' +AUTOML_CORE_REF=main # Or your development branch in NVIDIA-TAO/tao-automl. +AUTOML_CORE_SOURCE="nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" +python -m pip install "$AUTOML_CORE_SOURCE" -e '.[automl-prompts,reasoner-eval]' +python -m pip check cosmos-automl prompts --help ``` -The controller extra pins `gepa==0.1.2`. The workload image needs the matching +For a `uv`-managed environment, use the main guide's post-sync source install +and `uv run --no-sync cosmos-automl prompts --help`; `--all-extras` includes +Framework's GEPA dependency. Preserve the resolved source commit on resume. + +The Framework controller extra pins `gepa==0.1.2` directly; it does not depend on +the selected core branch declaring a `prompts` extra. The workload image needs the matching Framework prompt evaluator and its normal GPU inference dependencies, not GEPA or an LLM credential. Existing training-only images are not implicitly upgraded. diff --git a/docs/automl-slurm.md b/docs/automl-slurm.md index 4efba450f..397babb97 100644 --- a/docs/automl-slurm.md +++ b/docs/automl-slurm.md @@ -1,6 +1,6 @@ # AutoML on SLURM -The local core wheel provides a single-node SLURM executor. The controller +The core package provides a single-node SLURM executor. The controller runs on a CPU node **on the cluster**, with `sbatch`, `squeue`, `sacct`, `scontrol`, `scancel`, `srun`, a compatible Python environment, and shared storage. It submits one GPU allocation per candidate, waits, collects the @@ -51,17 +51,23 @@ allocation identity and GPU visibility belong to the scheduler. Forward `NVIDIA_DRIVER_CAPABILITIES` by name after exporting `compute,utility,video` for CUDA video decoding; credentials, if needed, are names only as well. -The CPU controller needs the Framework and core wheels plus their planning +The CPU controller needs the Framework and core packages plus their planning dependencies, not GPU training packages. Verify its Python version/imports, native scheduler CLIs, and shared interpreter path on compute nodes. Do not assume the system Python or the interpreter inside the training image is a suitable host-side controller environment. -For an isolated planning-only controller, install the core wheel with its -dependencies, then the Framework wheel with `--no-deps`: +For an isolated planning-only controller, install the core from the selected +source revision with its dependencies, then the Framework wheel with `--no-deps`. +Use `main` by default or the user's explicit development branch, as described in +[source installation](automl.md#installation-from-source). Resolve and record the +commit during preparation, never reinstall from a moving branch at job startup: ```bash -python -m pip install /shared/wheels/nvidia_automl_core-1.0.0.dev0-py3-none-any.whl +AUTOML_CORE_REF=main # Or the selected development branch / recorded commit SHA. +python -m pip install \ + "nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" \ + 'nvidia-automl-core==0.1.0' python -m pip install --no-deps /shared/wheels/cosmos_framework-1.2.2-py3-none-any.whl python -m cosmos_framework.scripts.automl --help ``` @@ -73,6 +79,9 @@ dependencies before installing Framework and exercise the installed adapter's `plan` command outside both source checkouts. Never use this limited environment for training; the prepared image supplies that stack. Point Python installer caches at shared storage when the login home has a small quota. +For GEPA controllers, also install `gepa==0.1.2` before verifying imports; the +planning-only `--no-deps` Framework install does not install its extras. A wheel +built from the recorded core revision remains an alternative for offline nodes. ## Private node history and exclusions @@ -119,7 +128,7 @@ native termination/retry-budget safeguards. GPU exclusions do not change the separately submitted CPU controller's placement. `SBATCH_EXCLUDE` is not a substitute: the executor filters scheduler environment overrides. -This requires a new core controller build, still `1.0.0.dev0`; the training +This requires a core controller build with node-exclusion support; the training image need not change. Do not upgrade a live controller or mutate its sealed manifest. Existing searches keep their original placement policy; select the new build and exclusions when preparing a fresh experiment. @@ -171,6 +180,40 @@ other workloads instead of claiming these checks cover them. ### Before the search controller starts +Cross-node startup probing is **opt-in** in +[slurm_controller.sbatch](../examples/automl/slurm_controller.sbatch): leave +`AUTOML_STARTUP_HELPER` unset for the original direct-CLI launch path. It needs +only `AUTOML_CONTROLLER_PYTHON`, `AUTOML_MANIFEST`, `AUTOML_WORKSPACE`, and +`AUTOML_CONTROLLER_MAX_ALLOCATIONS`; no startup observer, base directory, or +separate deadline variable is required. The search/controller limits still apply, +but this mode does not provide cross-node bootstrap validation or automatic +recovery for failures before the controller starts. Monitor native CPU state +through the site's existing submission service. + +The allocation cap **includes the initial allocation**. Native restart generation +is zero-based: cap `3` allows generations `0`, `1`, and `2`; generation `3` or +higher exits `2` before a probe or controller runs. Both values must be decimal +integers (cap positive, generation nonnegative); zero padding such as `003` and +`0002` is accepted without octal interpretation. Missing `SLURM_RESTART_COUNT` +means generation `0`. Malformed values now fail early rather than reaching the +CLI. The shell compares normalized copies but passes the original cap argument +unchanged. With the helper unset and `AUTOML_ACTION` unset, valid launches keep +the original CLI argv and direct `exec` signal delivery. + +This is the same count/generation boundary as the core's `controller.Session.run` +check, not a replacement for it: the shell reads `SLURM_RESTART_COUNT`, while the +core independently reads native `scontrol` `Restarts` and checks ownership, +resource shape, ledger continuity, and remaining time. Passing the shell check +does not guarantee core acceptance. A stale or inconsistent scheduler environment +must be reconciled, not used to bypass either guard. + +To enable the supervised path below, additionally set `AUTOML_STARTUP_HELPER` +to the sealed shared helper path, `AUTOML_STARTUP_BASE` to the registered shared +attempt base, and `AUTOML_DEADLINE` to the original absolute Unix deadline. +Once opted in, missing settings or an absent observer fail closed; there is no +fallback to an unprobed launch. Use this path when startup recovery is part of +the reviewed launch policy. + Trial recovery and CPU handoff run inside Python. They cannot recover a failed batch prolog, missing interpreter, failed environment import or bootstrap that never reaches the search controller. The submitting process must therefore @@ -224,6 +267,56 @@ and `USR1` to exit while bootstrap is running; never ignore an early signal. Use a fresh attempt directory for each native job ID and `SLURM_RESTART_COUNT` generation. Validate the generation against the approved allocation cap before probing. The external observer must watch each generation automatically. +The batch template diagnoses invalid restart/allocation values and exhausted +caps before installing signal traps or running helpers. Bootstrap signals and +helper failures exit with a diagnostic and status 2, not `EX_TEMPFAIL` (75). +They do not request native requeue; the external supervisor still owns terminal +reconciliation and approved startup retries. Site-enforced restarts can still +occur, but an out-of-budget generation cannot probe or start the controller. + +#### Persistent observer service + +Use [slurm_controller_watch.py](../tools/cosmos_workflows/scripts/slurm_controller_watch.py) +on an approved always-on host with the same shared storage and native SLURM +commands, **outside the CPU allocations it observes**. It automatically selects +`/-g` for each owned running job. Scheduler outages +retry read-only; the registry, deadline and observations survive process restarts. +The watcher never submits, requeues, cancels, or authorizes replacement jobs. + +Before submitting CPU controllers: + +```bash +python3 /shared/tools/slurm_controller_watch.py init \ + --state-dir /shared/run/observer --deadline +# Review/install slurm_observer.service on the approved service host first. +systemctl --user start slurm_observer.service +python3 /shared/tools/slurm_controller_watch.py ready --state-dir /shared/run/observer +``` + +Replace paths in the [user-service template](../examples/automl/slurm_observer.service). +It restarts the watcher after unexpected exits; the watcher exits normally at the +original deadline. The user manager must survive logout and the whole run. User +lingering, login-node residency and service-host availability are site prerequisites, +not settings to enable implicitly. Occasional `tick` calls or `nohup` alone are not +a supervised service. `ready` requires a fresh heartbeat **and** a held observer +lock. Stage the watcher and its sibling startup helper together; both are stdlib-only. + +After an acknowledged CPU submission, register its exact ID: + +```bash +python3 /shared/tools/slurm_controller_watch.py register \ + --state-dir /shared/run/observer --job-id \ + --attempt-base /shared/run/startup --max-allocations +``` + +Use that attempt base and original deadline in the reviewed +[CPU batch template](../examples/automl/slurm_controller.sbatch). Register every +acknowledged controller, including cold continuations; native requeues need no +manual registration. A lost submission reply still needs the submitter's native-ID +reconciliation, not a guessed ID or replacement. Registered entries and the deadline +are immutable. The submitting service still owns terminal monitoring and startup +retry policy; this watcher supplies the cross-node handshake, not another search +controller. Restart it on the same host; moving hosts mid-handshake fails closed. Probe success only means the storage check passed; it does not prove that the brain started. Preserve `fanout` before exec and use the native CLI's @@ -257,7 +350,7 @@ The controller's own allocation must leave time for shutdown before its hard wall limit; configure its terminal signals and timeout in the launch envelope. Keep the original engine wheel for existing experiments: changing the engine build changes its identity even while the development version stays fixed. -See [development version and build identity](automl.md#development-version-and-build-identity). +See [source revision, package version and build identity](automl.md#source-revision-package-version-and-build-identity). ### Explicit cold recovery migration @@ -297,8 +390,13 @@ resource shape, wall time and grace, but set `--controller-max-allocations` to n more than the migration's **remaining** allocation count. Prior allocations are deducted, and bypassing the CPU session is rejected. The original absolute deadline still applies, including controller downtime. Ordinary identity checks are not -disabled. This narrow operation does not support GEPA, other executors, arbitrary -engine upgrades or chained migrations. +disabled. A previously migrated, stopped SFT search requires explicit +`--continue-migration` with its **current exact wheel**. Each transaction retains +its predecessor in `migrations/recovery-v.json`, validates CPU history against +the preceding remaining envelope, and deducts allocations again; it never restores +the initial cap. Unacknowledged native handoffs still require reconciliation, +including CPU generations that failed before entering Python. This narrow operation +does not support GEPA, other executors or arbitrary engine upgrades. ### Bounded CPU-controller handoff @@ -328,6 +426,10 @@ in-flight LLM request plus at least 60 seconds. Prepare [slurm_controller.sbatch](../examples/automl/slurm_controller.sbatch), replace its site placeholders, and set the documented environment variables. +`AUTOML_WORKFLOW` is `sft` (default) or `prompts`; `AUTOML_ACTION` is `run` +(default) or `resume` for a reviewed continuation. Other actions are rejected +before any probe or approved CLI invocation. Startup probing remains optional +as described [above](#before-the-search-controller-starts). The example starts one experiment per CPU job, using `exec` so the Python process receives `--signal=B:USR1@300`. Review before `sbatch`; the example is not a ready-to-launch site configuration. It is not a supervisor for an @@ -397,6 +499,27 @@ cancellation markers or automatically submit a replacement. ### Submission reconciliation +Timed-out read-only preparation checks are separate from uncertain submissions. +The controller retries cluster/node/partition checks, batch syntax checks and +`sbatch --test-only`, backing off up to 300 seconds within the original deadline +and CPU handoff envelope. `_scheduler_check.json` identifies the operation, +timeout and `WAITING`/`COMPLETE` state without saving argv, stderr or credentials. +Nonzero exits, invalid configuration, ownership conflicts and record-hook errors +are not silently classified as transient scheduler timeouts. + +Before submission, the core seals the complete staged attempt's files. A handoff +during these checks resumes **that reserved candidate** after checking the seal, +open launch fence, owner, absence of worker/stop receipts and unchanged deadline. +It does not repeat a proposal or consume a training retry. Older or incomplete +staging cannot be replayed automatically. The preflight clock is durable even +before the first ledger is created. The CPU session also retries timed-out native +identity reads before search startup; unverified ownership still stops safely +when its bounded preparation time expires. + +Once `submission_attempted` is recorded, only the uncertain-submission protocol +below applies. Failures after the ledger opens now persist an explicit terminal +error and cleanup observation instead of leaving a stale `RUNNING` ledger. + `execution.slurm.submission_reconcile_seconds` defaults to **120** (1–900). A timeout, transport/exit error or invalid `sbatch` acknowledgement keeps the candidate reserved while the controller queries both queue and accounting. @@ -541,9 +664,9 @@ retries, count conservatively toward `max_evaluated_examples`. Reflection API ca are never retried by this mechanism. Automatic infrastructure retries currently require SLURM; local venv/Docker failure classification is not inferred. -These changes require new exact local controller wheels, still core -`1.0.0.dev0`; they do not upgrade already-running controllers or old experiment -identities automatically. Do not edit a live manifest/ledger or replace its installed wheel to +These changes require a controller build with bounded recovery support; selecting +a new source revision does not upgrade already-running controllers or old experiment +identities automatically. Do not edit a live manifest/ledger or replace its installed package to enable recovery. The worker is staged from the controller; the Framework changes declare existing native status files, not new training-image entrypoints. CPU fault-injection checks are not live cluster recovery validation. diff --git a/docs/automl.md b/docs/automl.md index 6c644e758..3534eda1c 100644 --- a/docs/automl.md +++ b/docs/automl.md @@ -21,7 +21,7 @@ implementation coverage is not equivalent to live GPU validation. | Execution platform | Core support | Framework evidence / limits | | ------------------ | ------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------- | | Docker | Implemented | Local Docker WTS GPU training/evaluation exercised. | -| Virtualenv | Implemented | CPU lifecycle verified; not equivalent to the recorded Framework GPU matrix. | +| Virtualenv | Implemented | Fresh-environment WTS/H100 training and evaluation passed after explicit controller recovery; see the bounded verification below. | | SLURM | Implemented | On-cluster CPU controller submits and monitors separate single-node GPU allocations. Recorded bounded SFT/GEPA matrix results are below. | | Kubernetes | Implemented; live validation pending | Single-pod Jobs, explicit context/PVCs, native deadline and UID-safe suspension. | | Brev | Implemented; live validation pending | Docker on an existing instance, explicit daemon identity and shared/controller-local instance storage; no provisioning or teardown. | @@ -41,75 +41,183 @@ For SLURM CPU allocations shorter than an experiment, the opt-in preserves active jobs and the original deadline. This is not an implementation of local-to-remote orchestration or a cloud controller service. -## Installation and release status +## Installation from source + +The default source is **`main` in [NVIDIA-TAO/tao-automl](https://github.com/NVIDIA-TAO/tao-automl)**, +specifically `packages/automl-core`. A named development branch is an explicit +override, not a fallback. Installers fetch and build this independent package; +users do not need a pre-existing TAO checkout, TAO application/SDK, or a publicly +published wheel. Git, repository access, and access to the package's build/runtime +dependencies are required. Configure Git authentication separately if needed; +do not embed credentials in the source URL. + +Framework requires **`nvidia-automl-core==0.1.0`**. The source repository is public, +but your environment must still be able to reach it and read the selected revision. +If access is unavailable, obtain a reviewed `nvidia_automl_core-0.1.0-*.whl` from +your project maintainer and install it into the controller environment instead. +Record its SHA256 and originating source commit; do not substitute the TAO +application package. Neither a `main` nor a development-branch name is an immutable +build: [record the resolved commit and engine hash](#source-revision-package-version-and-build-identity) +after a new installation, and pin that full commit (or retain the exact wheel) +for reproduction and existing experiments. Source access failures are not permission +to change branches or replace an existing controller environment. + +The repository contains two distributions: its root package is +`nvidia-tao-automl` (the TAO application), while `packages/automl-core` builds +`nvidia-automl-core` (the engine imported by Framework as `automl_core`). The +[PyPI TAO application](https://pypi.org/project/nvidia-tao-automl/) is not a +replacement for this core package. Preserve the selected source's declared +package name and version; do not rename it or substitute the TAO application's +release number. Framework's extras must match the core's distribution metadata. + +The following commands run from a prepared Framework checkout. Source selection +is for a **new environment or new experiment**. Do not update the controller of +a running or resumable experiment just because a branch has advanced. +For a reproducible run or deployment, use the [commit-pinned installation](#source-revision-package-version-and-build-identity) +instead of the moving-branch examples below. + +### Default: main -The independently built `nvidia-automl-core` 1.0.0.dev0 wheel must be supplied by the -maintainer: **this change does not publish it to an index**. Install the wheel -first, then install the extra into your already-prepared Framework environment: +```bash +AUTOML_CORE_REF=main +AUTOML_CORE_SOURCE="nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" + +# Prepare Framework first; choose the CUDA group appropriate for your machine. +uv sync --all-extras --group=cu128-train +uv pip install --no-config --python .venv/bin/python \ + --reinstall-package nvidia-automl-core \ + "$AUTOML_CORE_SOURCE" 'nvidia-automl-core==0.1.0' +uv pip check --python .venv/bin/python +uv run --no-sync cosmos-automl --help +``` + +The version constraint matches the committed core package metadata and Framework's +AutoML extras; a branch declaring an incompatible version fails resolution rather +than weakening that requirement. `--reinstall-package` refreshes this selected +package even when another commit has the same package version. It does not +request reinstallation of the entire training stack. + +Framework's `uv` configuration excludes **only `nvidia-automl-core`** from managed +resolution so ordinary setup/CI does not require a source installation or select +a development branch for you. `--no-config` makes the explicit installation above +honor the source requirement instead of that exclusion. Use `uv run --no-sync` +or `.venv/bin/cosmos-automl` afterward. An exact `uv sync` removes the separately +installed core; reinstall the **recorded commit**, not a newer branch head, if +repairing an existing experiment's environment. + +### Install from a named development branch + +Set the actual branch name in the same repository, including any `/` separators: ```bash -python -m pip install /path/to/nvidia_automl_core-1.0.0.dev0-py3-none-any.whl -python -m pip install -e '.[automl,reasoner-eval]' -cosmos-automl --help +AUTOML_CORE_REF='codex/repository-neutral-automl' # Replace with your development branch. +AUTOML_CORE_SOURCE="nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" +uv pip install --no-config --python .venv/bin/python \ + --reinstall-package nvidia-automl-core \ + "$AUTOML_CORE_SOURCE" 'nvidia-automl-core==0.1.0' +uv pip check --python .venv/bin/python +uv run --no-sync cosmos-automl --help ``` -The core is maintained in `tao-automl/packages/automl-core` and built independently; -that is a maintainer source location, not a consumer installation prerequisite. -Current delivery is **local wheels only**. To install a supplied Framework wheel -instead of an editable checkout, install the core wheel above, then: +The branch must contain `packages/automl-core/pyproject.toml`. A nonexistent or +inaccessible branch, missing package, or incompatible version is an installation +error; do not silently switch to another branch. Skill-driven setup asks whether +to use `main` or a user-supplied development branch before a new installation. +Manual CLI setup uses the variable above and does not prompt interactively. + +### Pip-only installation + +In a prepared environment without a previous core installation, use the same +`AUTOML_CORE_REF` / `AUTOML_CORE_SOURCE` selection (`main` by default): ```bash -python -m pip install '/path/to/cosmos_framework-1.2.2-py3-none-any.whl[automl,reasoner-eval]' +AUTOML_CORE_REF=main # Or the explicitly selected development branch. +AUTOML_CORE_SOURCE="nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" +python -m pip install "$AUTOML_CORE_SOURCE" -e '.[automl,reasoner-eval]' python -m pip check cosmos-automl --help ``` -If Framework 1.2.2 is already installed from a different build, resolve the extras -above, then replace just that package with `pip install --force-reinstall --no-deps -/path/to/cosmos_framework-1.2.2-py3-none-any.whl`. The wheelhouse is not an offline -mirror of third-party dependencies. +For GEPA, select `.[automl-prompts,reasoner-eval]` instead. That Framework extra +declares `gepa==0.1.2` directly, so it also works with source branches that do not +declare the core's optional `prompts` extra. A supplied Framework wheel may +replace the editable requirement with `/path/to/cosmos_framework.whl[automl,reasoner-eval]` +(using its actual filename). When switching same-version core builds, use a new +environment or the targeted `uv --reinstall-package` command above; do not assume +an already-satisfied version means the selected source was installed. -A bare index install cannot find the unpublished core. The repository's `uv` -configuration explicitly excludes **only `nvidia-automl-core`** from its managed -resolution while delivery remains local-wheel-only. This lets ordinary Framework -setup and CI resolve without an additional source checkout; it does **not** -install AutoML. -The wheel metadata still declares the core requirement for the AutoML extras. -A distributable core location and removal of this development exception remain -a release gate. +### Source revision, package version and build identity -For a `uv`-managed Framework checkout, sync first, then install the supplied core -wheel into that environment (include `[prompts]` for GEPA): +`main` remains the default for a new exploratory installation. For a reproducible +run or deployment, **resolve the selected branch to a full commit before installing** +and retain that commit as the installation input, rather than reinstalling from +the branch name later. After the `main` or development-branch selection above: ```bash -uv sync --all-extras --group=cu128-train -uv pip install --no-config --python .venv/bin/python '/path/to/nvidia_automl_core-1.0.0.dev0-py3-none-any.whl[prompts]' +AUTOML_CORE_COMMIT=$(git ls-remote --exit-code \ + https://github.com/NVIDIA-TAO/tao-automl.git "refs/heads/${AUTOML_CORE_REF}" | cut -f1) +if [[ ! "$AUTOML_CORE_COMMIT" =~ ^[0-9a-f]{40}$ ]]; then + echo 'Could not resolve the selected AutoML branch; stop installation.' >&2 + exit 1 +fi +AUTOML_CORE_REF=$AUTOML_CORE_COMMIT +AUTOML_CORE_SOURCE="nvidia-automl-core @ git+https://github.com/NVIDIA-TAO/tao-automl.git@${AUTOML_CORE_REF}#subdirectory=packages/automl-core" +uv pip install --no-config --python .venv/bin/python \ + --reinstall-package nvidia-automl-core \ + "$AUTOML_CORE_SOURCE" 'nvidia-automl-core==0.1.0' uv pip check --python .venv/bin/python -uv run --no-sync cosmos-automl --help ``` -The install command's `--no-config` bypasses the project's development exclusion -for this explicitly supplied wheel. Use `uv run --no-sync` or the environment's -executable directly for AutoML. -An exact `uv sync` removes separately installed packages; reinstall the supplied -wheel afterward. No source checkout or machine-specific wheel path is recorded -in the project's lock file. - -### Development version and build identity - -AutoML is in its initial **v1 development phase**, not a released package series. -Keep `nvidia-automl-core` at **`1.0.0.dev0`** throughout the port. Bug fixes and local -wheel builds do not increment the patch version or development counter. Change -the package version only as part of an explicitly approved release; building a -local validation wheel is not a release. Cosmos Framework's existing package -version remains **1.2.2**, independent of the AutoML development label. - -Store each build in a separate immutable wheelhouse directory and record its -wheel SHA256 and source revision. Do not overwrite an earlier same-named wheel. -When intentionally selecting another same-version wheel for a new environment -or run, install the exact supplied artifact; an existing installation may require -`pip install --force-reinstall --no-deps /path/to/the/wheel.whl` after resolving -dependencies. A package version alone does not identify a development build. +For an already recorded full commit, set `AUTOML_CORE_REF` directly to that SHA; +do not resolve a moving branch again. Pip-only installs can use the same pinned +`AUTOML_CORE_SOURCE`. After installation, verify and record the resolved commit +and engine hash (the commit must match the requested pin): + +```bash +python - <<'PY' +import json +from importlib.metadata import distribution +from automl_core import build_fingerprint + +package = distribution("nvidia-automl-core") +origin = json.loads(package.read_text("direct_url.json") or "{}") +vcs = origin.get("vcs_info", {}) +print("package_version:", package.version) +print("requested_revision:", vcs.get("requested_revision")) +print("source_commit:", vcs.get("commit_id")) +print("engine_build:", build_fingerprint()) +PY +``` + +Use the target environment's Python (for example `.venv/bin/python`). Put the +resolved core commit in the run's provenance. For a source reinstall, set +`AUTOML_CORE_REF` to that full commit SHA and reconstruct `AUTOML_CORE_SOURCE`; +the same install commands then reproduce the core source selection. A Git commit +does not freeze third-party dependencies, datasets, or training images; preserve +those separately. For offline deployment or cold migration, retain an immutable +wheel built from the same commit and record its SHA256. Source installation does +not remove the migration protocol's requirement for the exact original wheel. + +The current core **distribution metadata is `0.1.0`**; Framework remains `1.2.2`. +Installing a branch or building a local wheel is not a release and does not require +a version bump. A package version alone does not identify a development build. + +This replaces Framework's former `==1.0.0.dev0` requirement; an environment with +that local development wheel does **not** satisfy the new extras. For a new +experiment, prepare a separate environment with the selected `0.1.0` source build. +Do not downgrade a live or resumable controller in place, relabel old wheels, or +rewrite its ledger. Keep the old environment and exact artifacts for ordinary +resume; use [reviewed cold migration](automl-slurm.md#explicit-cold-recovery-migration) +only where the migration protocol accepts the stopped experiment. That protocol +does not support arbitrary version changes or GEPA migration. The old pin is no +longer supported for new installations through these extras; this is not a claim +that every historical `1.0.0.dev0` artifact is API-compatible or formally retired +by its publisher. + +Compatibility evidence includes the [two-candidate H100 SFT smoke](#fresh-venv-wts-verification-2026-10-19) +with core `0.1.0`, native training, export, and evaluation. It required explicit +controller recovery and validates the recorded builds, not every future commit +that retains the same version, nor a new GEPA GPU qualification. New SFT and GEPA identities include `engine_build`, a SHA256 fingerprint of the installed core's relative Python file names and contents. Identical source and @@ -140,6 +248,41 @@ in the image or export it and include its name in `execution.environment`. Also verify that the non-root container has a valid user identity (for example, forward `USER` and `LOGNAME`) and writable per-trial caches. +### Fresh-machine reproducibility checklist + +For a new run, preserve a small provenance bundle alongside the reviewed manifest: + +- Framework commit, its `uv.lock`, selected CUDA group, Python version and full + installed package inventory. For a non-editable training installation, use + `uv sync --frozen --all-extras --group=cu128-train --no-editable` (or the + matching CUDA group), then install the selected core as above. Record the core + commit and `engine_build`, plus wheel hashes or the immutable container digest. +- The actual training/evaluation TOMLs, seed, search domains and trial/time caps. + Preserve dataset split/content revisions and model checkpoint identity; + recipe paths and record counts alone do not reproduce their contents. +- Required environment variable **names**, decoder/export prerequisites and + session-local library setup. Keep credentials out of the bundle. For CUDA 12, + consult the [TorchCodec NPP setup](setup.md#torchcodec-cannot-find-cuda-12-npp) + when the host does not expose matching libraries. Declare required exported + variable names, such as `LD_PRELOAD`, in `execution.environment` too. +- A writable workspace with capacity for every trial's checkpoints and exports. + Prefer local POSIX storage for local venv runs. Shared storage must pass + concurrent atomic-replacement/read checks as well as locking and fsync checks; + a single successful write is not sufficient. + +Run the install/dependency checks, `cosmos-automl --help` and CPU-only `plan` +before approving a small GPU run. Verify training, checkpoint export and the +selected evaluator, not just imports. Keep the run environment unchanged until +its controller and workers are terminal. Neither the optional engine nor private +datasets/checkpoints are bundled into Framework; a new user supplies assets and +paths, not another developer's checkout or credentials. + +The recorded Python 3.13 CUDA training installation reports a `pip check` warning: +`jupyter-compare-view==0.2.4` declares Python `<3.13`. This notebook-only dependency +was unused by the bounded training/evaluation smoke below; its metadata conflict +is not resolved by installing AutoML. Do not describe the full environment's +dependency check as passing or silently change its Python/dependency pins. + ## Prepare and review Start with [reasoner.yaml](../examples/automl/reasoner.yaml), @@ -294,6 +437,26 @@ the extra allocation caps in `plan`. An `attempt_id` can differ from the logical trial `id`; `attempts` preserves failed executions and their native identities. Unconfirmed cancellation blocks replacement rather than creating duplicate work. +For shared-storage receipt/ledger reads, the core retries only transient `EIO`, +`ESTALE` and `ETIMEDOUT`, reopening up to five times with 0.75 seconds of total +backoff. It does not replay training, submissions or paid requests. Persistent +errors and malformed receipts still stop safely. A dead CPU controller can leave +a GPU worker running; check both native process/job identity and worker receipts +before interpreting a ledger's last `RUNNING` state or starting a replacement. + +Ordinary resume requires the original build. A stopped same-host venv SFT search +can use `cosmos-automl migrate --manifest --workspace +--from-wheel ` to review a narrow recovery-build continuation; +`--approve` applies it, and a separately approved `resume` continues the search. +All old workers and command processes must first be terminal on their recorded +host/boot. Completed-but-unrecorded results must be reconciled with the original +build before changing it. Preserve the interpreter, package inventory, data, +completed results, remaining trials and original absolute deadline. Never install +a new wheel into an environment used by active workers, edit sealed identities, +or use a fresh experiment to disguise a retry. This migration does not support +GEPA, Docker or cloud runs; see the [SLURM migration contract](automl-slurm.md#explicit-cold-recovery-migration) +for scheduler-owned continuations. + ## Scope and checks The shared engine ports the original eleven algorithm/controller modes and @@ -312,7 +475,7 @@ by scheduler-query timeouts during evaluation. Hyperband-ES curve pruning and completed GPU searches driven by those structured controllers remain verification gaps; these checks do not establish model quality or validate a production search. -Run adapter tests with an installed core wheel (no GPU imports required): +Run adapter tests with the selected core installed (no GPU imports required): ```bash python -m pytest -o addopts='' --confcutdir=tests/automl tests/automl @@ -831,3 +994,54 @@ does not require a training-image rebuild; no existing cluster controller was upgraded. The earlier GPU results used the earlier core hash and do not validate this new build. The separate GEPA progress-file gap above remains open, as do the core distribution/dependency review and coordinated merge gates. + +### Controller startup and restart assets + +The checkout now supplies the [persistent observer service](automl-slurm.md#persistent-observer-service) +and a CPU batch template that probes every native restart generation before +starting the controller. An initial one-shot handshake cannot cover a later +requeue. The observer must be independently supervised on an approved host; it +does not submit or replace jobs, authorize retries, or extend the deadline. + +The core also retries timed-out read-only scheduler preparation checks and +continues a fully sealed, unsubmitted candidate across a bounded CPU handoff. +An attempted submission still uses the separate reconciliation/fencing protocol. +These controller paths were checked with CPU subprocesses and a simulated +scheduler; this is not live cross-node SLURM recovery or service-deployment proof. + +### Fresh-venv WTS verification (2026-10-19) + +A fresh, non-editable Python 3.13 CUDA 12.8 environment ran native Framework +AutoML on Cosmos3-Nano-VLM with 64 WTS training records and 16 disjoint held-out +questions. Two random-search candidates each completed five LoRA steps on eight +H100 GPUs, DCP-to-HF export and single-GPU evaluation. Both scored **8/16 (50%)**. +The approved limits were 30 minutes per trial and two hours for the search, with +no automatic training retries or paid LLM calls. + +The initial source-install receipts identify Framework commit +`def247af66f22abb45d25c3138bf13d0edcb4496` and core commit +`bd21cbd3a6aa24e3277c0d232dab1614f935c03e`, with core distribution `0.1.0`. +These are historical build-specific results, not a GPU rerun of later PR heads. + +This was a **pass after explicit controller recovery**, not an uninterrupted +run. A transient shared-drive `EIO` stopped the CPU controller during the first +trial's evaluation while its worker continued. The original build collected that +completed result; a reviewed cold migration then resumed the corrected build. +Both trials finished within the original deadline, with no repeated training, +additional trial slots or budget reset. The resumed controller exited normally. + +The corrected core's `engine_build` is +`cdfb7731512119c1b27db079fb34cda24b4b3506a7a8beef225894a375aaeca2`; +its distribution version remains **0.1.0**. A separate real shared-drive stress +probe recovered 86 transient read errors without escaping a logical read, and a +two-launch CPU dummy workflow verified stopped-controller continuation, result +preservation, unchanged deadlines and rejection of increased budgets. Existing +core tests passed; no new repository unit tests or package-version bumps were +introduced. + +The notebook dependency warning above remains. Transformers also emitted a +tokenizer-regex warning while loading Qwen3-VL; tokenizer files/settings were not +changed. This bounded random-search result establishes training/export/evaluation +execution after recovery, **not** full-data accuracy, all-algorithm qualification, +or automatic recovery from persistent storage failures. Private asset paths and +run receipts are not part of the public installation instructions. diff --git a/docs/setup.md b/docs/setup.md index 0c2149d3d..cc83fd0fd 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -68,6 +68,13 @@ cd cosmos-framework The two supported install paths are the recommended base image and the Docker container. For other paths (standalone venv, custom torch/cuda) see [Advanced](#advanced). +For AutoML, also follow [installation from source](automl.md#installation-from-source) +after preparing Framework. The core is installed from `NVIDIA-TAO/tao-automl`'s +`main` branch by default; the guide includes an explicit named-development-branch +override. A public wheel release is not required. Ordinary `uv sync --all-extras` +does not install the core; use the documented post-sync install and `--no-sync` +invocation. Keep the installed revision unchanged when resuming an experiment. +
Quickstart: From the Recommended Base Image ### Quickstart: From the Recommended Base Image diff --git a/examples/automl/slurm_controller.sbatch b/examples/automl/slurm_controller.sbatch index 1f0943624..cb1ac60cd 100755 --- a/examples/automl/slurm_controller.sbatch +++ b/examples/automl/slurm_controller.sbatch @@ -20,6 +20,30 @@ set -euo pipefail : "${AUTOML_MANIFEST:?Set the sealed, absolute shared manifest path}" : "${AUTOML_WORKSPACE:?Set the absolute shared experiment workspace}" : "${AUTOML_CONTROLLER_MAX_ALLOCATIONS:?Set the reviewed CPU allocation cap}" +generation=${SLURM_RESTART_COUNT:-0} +if [[ ! "$generation" =~ ^[0-9]+$ ]]; then + echo 'SLURM_RESTART_COUNT must be a nonnegative decimal integer' >&2 + exit 2 +fi +if [[ ! "$AUTOML_CONTROLLER_MAX_ALLOCATIONS" =~ ^[0-9]+$ || ! "$AUTOML_CONTROLLER_MAX_ALLOCATIONS" =~ [1-9] ]]; then + echo 'AUTOML_CONTROLLER_MAX_ALLOCATIONS must be a positive decimal integer' >&2 + exit 2 +fi +# Normalize native zero padding; compare decimal strings without octal parsing +# or arithmetic overflow. No helper or signal trap runs before this cap check. +generation=${generation#"${generation%%[!0]*}"} +generation=${generation:-0} +allocation_cap=$AUTOML_CONTROLLER_MAX_ALLOCATIONS +allocation_cap=${allocation_cap#"${allocation_cap%%[!0]*}"} +# The cap counts the initial allocation: cap 3 admits generations 0, 1 and 2. +# Preserve the original cap argument for the CLI's decimal int parser below. +if [[ ${#generation} -gt ${#allocation_cap} ]] || { + [[ ${#generation} -eq ${#allocation_cap} ]] && + [[ "$generation" == "$allocation_cap" || "$generation" > "$allocation_cap" ]] +}; then + echo "Controller allocation cap exhausted: generation $generation, approved count $allocation_cap" >&2 + exit 2 +fi # One experiment owns this CPU job. The verified successor automatically resumes # its ledger; do not clear outputs or switch manifests on SLURM_RESTART_COUNT. @@ -29,11 +53,58 @@ case "${AUTOML_WORKFLOW:-sft}" in prompts) workflow=(prompts) ;; *) echo 'AUTOML_WORKFLOW must be sft or prompts' >&2; exit 2 ;; esac +action=${AUTOML_ACTION:-run} +case "$action" in + run|resume) ;; + *) echo 'AUTOML_ACTION must be run or resume' >&2; exit 2 ;; +esac + +# Optional cross-node startup supervision. Unset helper preserves the original +# direct-CLI launch path; it does not provide pre-controller startup recovery. +if [[ -n "${AUTOML_STARTUP_HELPER:-}" ]]; then + : "${AUTOML_STARTUP_BASE:?Set the shared startup directory registered with the observer}" + : "${AUTOML_DEADLINE:?Set the original absolute experiment deadline}" + : "${SLURM_JOB_ID:?The startup probe requires a native SLURM job ID}" + + stop_startup() { + trap - TERM HUP INT USR1 + echo 'Controller startup interrupted; no automatic requeue requested' >&2 + if [[ -n "${startup_pid:-}" ]]; then + kill -TERM "$startup_pid" 2>/dev/null || true + fi + exit 2 + } + run_startup_helper() { + # A trapped signal interrupts wait promptly, unlike a foreground Python + # call. Forward termination to the probe; SLURM owns batch-child cleanup. + "$AUTOML_CONTROLLER_PYTHON" "$AUTOML_STARTUP_HELPER" "$@" & + startup_pid=$! + if wait "$startup_pid"; then + startup_pid= + else + startup_status=$? + startup_pid= + echo "Controller startup $1 failed (exit $startup_status); controller not started" >&2 + # Do not leak EX_TEMPFAIL (75) into site automatic-requeue policies. + # The external supervisor owns native reconciliation and retry caps. + exit 2 + fi + } + trap stop_startup TERM HUP INT USR1 + # The observer must be ready before sbatch and this job registered afterwards. + # Every requeue gets fresh receipts; do not wrap the handoff CLI in run. + export AUTOML_STARTUP_ATTEMPT="$AUTOML_STARTUP_BASE/${SLURM_JOB_ID}-g${generation}" + run_startup_helper probe \ + --attempt-dir "$AUTOML_STARTUP_ATTEMPT" --deadline "$AUTOML_DEADLINE" \ + --probe-seconds 600 --minimum-run-seconds 600 + run_startup_helper fanout --attempt-dir "$AUTOML_STARTUP_ATTEMPT" + trap - TERM HUP INT USR1 +fi # exec makes Python the batch process receiving B:USR1. Do not add an EXIT trap # that cancels trials or removes credentials needed by a successor allocation. exec "$AUTOML_CONTROLLER_PYTHON" -m cosmos_framework.scripts.automl \ - "${workflow[@]}" run --manifest "$AUTOML_MANIFEST" \ + "${workflow[@]}" "$action" --manifest "$AUTOML_MANIFEST" \ --workspace "$AUTOML_WORKSPACE" --approve \ --controller-max-allocations "$AUTOML_CONTROLLER_MAX_ALLOCATIONS" \ --controller-allocation-seconds 3600 --controller-grace-seconds 300 diff --git a/examples/automl/slurm_observer.service b/examples/automl/slurm_observer.service new file mode 100644 index 000000000..5dc345276 --- /dev/null +++ b/examples/automl/slurm_observer.service @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# User-local template for an approved always-on, shared-storage SLURM service host. +# Replace ALL paths; install as a user service only after site/launch review. +# Enabling user lingering or using a login node needs site approval separately. +[Unit] +Description=Bounded AutoML cross-node startup observer + +[Service] +Type=simple +ExecStart=/usr/bin/python3 /shared/tools/slurm_controller_watch.py watch --state-dir /shared/run/observer +Restart=on-failure +RestartSec=5 +TimeoutStopSec=150 +UMask=0077 + +[Install] +WantedBy=default.target diff --git a/pyproject.toml b/pyproject.toml index 59f46737f..386421dba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,8 +59,8 @@ dependencies = [ ] [project.optional-dependencies] -automl = ["nvidia-automl-core==1.0.0.dev0"] -automl-prompts = ["nvidia-automl-core[prompts]==1.0.0.dev0"] +automl = ["nvidia-automl-core==0.1.0"] +automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"] reasoner-eval = [ "attrs", "evaluate", @@ -349,11 +349,11 @@ repository = "https://github.com/NVIDIA/cosmos-framework" # activates this project's venv, whose `uv` then runs that sync. required-version = "==0.11.29" constraint-dependencies = ["pydantic-extra-types==2.11.1", "uv==0.11.29"] -# Development delivery of the optional AutoML engine is by independently built -# local wheel, not an index. Keep its published extra requirements above intact, -# but leave that engine out of the repository's uv-managed environment. Install -# the supplied wheel explicitly after sync; use uv run --no-sync for AutoML. -# Remove this exception once the core has an approved distribution location. +# Install the optional AutoML engine from tao-automl/main (or an explicitly +# selected development branch) after sync; see docs/automl.md. Keep ordinary +# Framework setup independent of that source installation. The extras above +# still enforce package compatibility. Use uv run --no-sync after installing +# the core so exact sync does not remove the selected source build. exclude-dependencies = ["nvidia-automl-core"] conflicts = [ # Bimanual FK training and the LIBERO simulator require separate MuJoCo runtimes. diff --git a/tools/cosmos_workflows/README.md b/tools/cosmos_workflows/README.md index ac4a8260d..b7125fb95 100644 --- a/tools/cosmos_workflows/README.md +++ b/tools/cosmos_workflows/README.md @@ -12,6 +12,10 @@ These checkout-local assets support [AutoML on SLURM](../../docs/automl-slurm.md supplies the cross-node lock observer and pre-fan-out marker. The submitting service must schedule observation automatically and reconcile native terminal state before any bounded startup retry; see the SLURM guide's startup contract. +- `scripts/slurm_controller_watch.py` persistently observes registered CPU jobs + and each native restart generation. Run it as an independently supervised, + deadline-bounded service; its `ready` check gates submission. It never launches + or replaces jobs. See `examples/automl/slurm_observer.service` at repository root. The search controller is `cosmos-automl`, not a workflow wrapper in this directory. Stage these assets from this repository, seal their hashes in the manifest, and @@ -21,3 +25,19 @@ of truth for whether jobs are running. The helpers retain their original Apache-2.0 license and attribution in [LICENSE](LICENSE) and [NOTICE](NOTICE). Annotation, serving, quantization, data conversion, and other platform workflows are not part of this focused AutoML change. + +## CPU-only regression checks + +From the repository root, using a test environment with `pytest` and `jsonschema`: + +```bash +python -m pytest -o addopts= --confcutdir=tools/cosmos_workflows \ + -o 'python_files=test_*.py *_test.py' tools/cosmos_workflows +``` + +The AutoML helpers CI job runs this suite without installing the GPU stack. +Observer tests cover registry identity, readiness/locking, native state and +generation checks, and scheduler/storage errors. Batch-process tests check +legacy argv/exec, decimal allocation boundaries, and opt-in bootstrap failures. +Scheduler calls and controller children are mocked; these are not live SLURM, +cross-node filesystem, or GPU-training verification. diff --git a/tools/cosmos_workflows/scripts/slurm_controller_watch.py b/tools/cosmos_workflows/scripts/slurm_controller_watch.py new file mode 100755 index 000000000..f3c01a416 --- /dev/null +++ b/tools/cosmos_workflows/scripts/slurm_controller_watch.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Persistent, restartable startup observer. Never submits, requeues or cancels. + +Run on an approved service host with shared storage and native SLURM tools, +outside the CPU allocations it observes. Supervise this process independently +of an interactive shell/chat (see the accompanying user-service template). +""" + +import argparse +import fcntl +import math +import os +import re +import signal +import subprocess +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from slurm_controller_startup import atomic, observe, read + + +def registry(root): + value = read(root / "registry.json") + if value.get("uid") != os.getuid() or not math.isfinite(value.get("deadline", float("inf"))): + raise ValueError("Observer registry requires an owner and finite absolute deadline") + if not isinstance(value.get("jobs"), dict): + raise ValueError("Invalid observer job registry") + for job, entry in value["jobs"].items(): + if ( + not re.fullmatch(r"[1-9][0-9]*", job) + or not Path(entry["attempt_base"]).is_absolute() + or type(entry["max_allocations"]) is not int + or entry["max_allocations"] < 1 + ): + raise ValueError("Invalid observer job identity or allocation cap") + return value + + +def ready(root, max_age=90): + value = registry(root) + state = read(root / "watch.json") + if ( + state.get("state") != "WATCHING" + or state.get("uid") != value["uid"] + or state.get("deadline") != value["deadline"] + or not 0 <= time.time() - state.get("updated_at", 0) < max_age + or time.time() >= value["deadline"] + ): + raise ValueError("Persistent startup observer is not ready; do not submit controllers") + # A recent heartbeat alone is insufficient after the observer has exited. + with (root / "watch.lock").open("a") as handle: + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + return state + raise ValueError("Startup observer is not running") + + +def register(root, job_id, attempt_base, max_allocations): + """Append an acknowledged native ID; do not adopt unknown submissions here.""" + ready(root) + with (root / "registry.lock").open("a") as handle: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + value = registry(root) + entry = {"attempt_base": str(attempt_base.absolute()), "max_allocations": max_allocations} + if not re.fullmatch(r"[1-9][0-9]*", job_id) or max_allocations < 1: + raise ValueError("Native ID and positive allocation cap required") + if job_id in value["jobs"] and value["jobs"][job_id] != entry: + raise ValueError("Observer identity cannot be replaced") + value["jobs"][job_id] = entry + atomic(root / "registry.json", value) + + +def tick(job_id, entry): + try: + result = subprocess.run( + ["scontrol", "show", "job", "-o", job_id], + capture_output=True, + text=True, + timeout=30, + check=True, + ) + fields = dict(re.findall(r"(?:^|\s)(\w+)=([^\s]+)", result.stdout)) + if fields.get("JobId") != job_id or not fields.get("UserId", "").endswith(f"({os.getuid()})"): + raise ValueError("Native observer ownership mismatch") + generation = int(fields["Restarts"]) + if not 0 <= generation < entry["max_allocations"]: + raise ValueError("Native restart exceeds the approved allocation cap") + if fields.get("JobState") != "RUNNING": + return {"state": "OBSERVING", "native_state": fields.get("JobState"), "generation": generation} + attempt = Path(entry["attempt_base"]) / f"{job_id}-g{generation}" + return {"state": "OBSERVING", "generation": generation, **observe(attempt, job_id)} + except (OSError, subprocess.SubprocessError) as exc: + # Scheduler/shared-storage outages do not kill all observers or authorize + # a replacement. Do not persist exception text, stderr or environment. + return {"state": "WAITING", "error_type": type(exc).__name__, "errno": getattr(exc, "errno", None)} + except (ValueError, KeyError, TypeError) as exc: + return {"state": "BLOCKED", "error_type": type(exc).__name__} + + +def watch(root, poll_seconds=2): + value = registry(root) + saved = read(root / "watch.json") + if saved and (saved.get("uid") != value["uid"] or saved.get("deadline") != value["deadline"]): + raise ValueError("Observer identity/deadline changed") + known = saved.get("registered", {}) + next_check = {} + observations = saved.get("observations", {}) + stopping = False + + def stop(*unused): + nonlocal stopping + stopping = True + + for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP): + signal.signal(sig, stop) + with (root / "watch.lock").open("a") as handle, ThreadPoolExecutor(max_workers=4) as pool: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + while not stopping and time.time() < value["deadline"]: + current = registry(root) + if current["deadline"] != value["deadline"] or any(current["jobs"].get(k) != v for k, v in known.items()): + raise ValueError("Registered observer identities/deadline must be immutable") + known = current["jobs"] + futures = { + job: pool.submit(tick, job, entry) + for job, entry in known.items() + if time.time() >= next_check.get(job, 0) + } + for job, future in futures.items(): + observations[job] = future.result() + observations[job]["observed_at"] = time.time() + # Fast only during the handshake; do not hammer the scheduler + # for hours after the brain starts or while it is queued. + phase = observations[job].get("phase") + delay = poll_seconds if phase in {"NOT_ENTERED", "LOCKED", "CLAIMED"} else 30 + next_check[job] = time.time() + delay + saved = dict( + uid=value["uid"], + deadline=value["deadline"], + registered=known, + state="WATCHING", + updated_at=time.time(), + observations=observations, + ) + atomic(root / "watch.json", saved) + wake = min(value["deadline"], time.time() + poll_seconds) + while not stopping and time.time() < wake: + time.sleep(max(0, min(0.2, wake - time.time()))) + saved.update(state="STOPPED" if stopping else "EXPIRED", updated_at=time.time()) + atomic(root / "watch.json", saved) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("verb", choices=("init", "watch", "ready", "register")) + parser.add_argument("--state-dir", required=True, type=Path) + parser.add_argument("--deadline", type=float) + parser.add_argument("--job-id") + parser.add_argument("--attempt-base", type=Path) + parser.add_argument("--max-allocations", type=int) + options = parser.parse_args() + root = options.state_dir.absolute() + if options.verb == "init": + if options.deadline is None or not math.isfinite(options.deadline) or options.deadline <= time.time(): + parser.error("init needs a future finite absolute --deadline") + root.mkdir(parents=True, exist_ok=False) + atomic(root / "registry.json", dict(uid=os.getuid(), deadline=options.deadline, jobs={})) + elif options.verb == "watch": + watch(root) + elif options.verb == "ready": + ready(root) + else: + if not options.job_id or options.attempt_base is None or options.max_allocations is None: + parser.error("register needs --job-id, --attempt-base and --max-allocations") + register(root, options.job_id, options.attempt_base, options.max_allocations) + + +if __name__ == "__main__": + main() diff --git a/tools/cosmos_workflows/scripts/tests/slurm_controller_batch_test.py b/tools/cosmos_workflows/scripts/tests/slurm_controller_batch_test.py new file mode 100644 index 000000000..329b123f7 --- /dev/null +++ b/tools/cosmos_workflows/scripts/tests/slurm_controller_batch_test.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execute the real batch shell with a recording child, never a controller.""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +BATCH = Path(__file__).resolve().parents[4] / "examples/automl/slurm_controller.sbatch" + + +@pytest.fixture +def launch(tmp_path): + interpreter = tmp_path / "record-python" + interpreter.write_text( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "with open(os.environ['CALLS'], 'a') as stream:\n" + " stream.write(json.dumps({'argv': sys.argv[1:], 'pid': os.getpid(),\n" + " 'attempt': os.environ.get('AUTOML_STARTUP_ATTEMPT')}) + '\\n')\n" + "if sys.argv[1] != '-m':\n" + " sys.exit(int(os.environ.get('HELPER_EXIT', '0')))\n" + ) + interpreter.chmod(0o755) + count = 0 + + def run(**overrides): + nonlocal count + count += 1 + calls = tmp_path / f"calls-{count}.jsonl" + env = { + "PATH": os.environ["PATH"], + "CALLS": str(calls), + "AUTOML_CONTROLLER_PYTHON": str(interpreter), + "AUTOML_MANIFEST": "/shared/path with spaces/search.yaml", + "AUTOML_WORKSPACE": "/shared/runs", + "AUTOML_CONTROLLER_MAX_ALLOCATIONS": "3", + **overrides, + } + process = subprocess.Popen(["bash", str(BATCH)], env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + try: + _, stderr = process.communicate(timeout=5) + finally: + if process.poll() is None: + process.kill() + process.wait() + entries = [json.loads(line) for line in calls.read_text().splitlines()] if calls.exists() else [] + return process.returncode, stderr.decode(), entries, process.pid + + return run + + +def expected_cli(workflow="sft", action="run", cap="3"): + # Original #288 argv for default action, including exact argument ordering. + return [ + "-m", + "cosmos_framework.scripts.automl", + *(["prompts"] if workflow == "prompts" else []), + action, + "--manifest", + "/shared/path with spaces/search.yaml", + "--workspace", + "/shared/runs", + "--approve", + "--controller-max-allocations", + cap, + "--controller-allocation-seconds", + "3600", + "--controller-grace-seconds", + "300", + ] + + +@pytest.mark.parametrize("workflow", ["sft", "prompts"]) +@pytest.mark.parametrize("settings", [{}, {"AUTOML_STARTUP_HELPER": ""}, {"AUTOML_STARTUP_BASE": "/ignored"}]) +def test_unsupervised_default_preserves_original_argv_and_exec(launch, workflow, settings): + code, stderr, entries, pid = launch(AUTOML_WORKFLOW=workflow, **settings) + assert (code, stderr) == (0, "") + assert entries == [{"argv": expected_cli(workflow), "pid": pid, "attempt": None}] + + +@pytest.mark.parametrize("cap", ["1", "3", "0003", "10", "9" * 25]) +def test_allocation_count_boundary_includes_initial_allocation(launch, cap): + # Same mathematical boundary as automl_core.controller.Session.run: + # generation >= max_allocations is exhausted; generation starts at zero. + for generation in [0, int(cap) - 1, int(cap), int(cap) + 1]: + code, stderr, entries, _ = launch(AUTOML_CONTROLLER_MAX_ALLOCATIONS=cap, SLURM_RESTART_COUNT=str(generation)) + if generation < int(cap): + assert (code, stderr) == (0, "") + assert entries[0]["argv"] == expected_cli(cap=cap) + else: + assert code == 2 and "allocation cap exhausted" in stderr + assert entries == [] + + +@pytest.mark.parametrize("cap", ["0", "000", "-1", "1+1", "x"]) +def test_invalid_cap_fails_before_child(launch, cap): + code, stderr, entries, _ = launch(AUTOML_CONTROLLER_MAX_ALLOCATIONS=cap) + assert code == 2 and "positive decimal integer" in stderr + assert entries == [] + + +@pytest.mark.parametrize("generation", ["-1", "1+1", "x"]) +def test_invalid_generation_fails_before_child(launch, generation): + code, stderr, entries, _ = launch(SLURM_RESTART_COUNT=generation) + assert code == 2 and "nonnegative decimal integer" in stderr + assert entries == [] + + +@pytest.fixture +def supervised(): + return { + "AUTOML_STARTUP_HELPER": "/shared/probe.py", + "AUTOML_STARTUP_BASE": "/shared/startup", + "AUTOML_DEADLINE": "2000000000", + "SLURM_JOB_ID": "123", + } + + +def test_supervised_restart_normalizes_generation_but_preserves_cap_argv(launch, supervised): + code, stderr, entries, pid = launch( + **supervised, SLURM_RESTART_COUNT="0002", AUTOML_CONTROLLER_MAX_ALLOCATIONS="003" + ) + assert (code, stderr) == (0, "") + assert [entry["argv"][1] for entry in entries[:2]] == ["probe", "fanout"] + assert all(entry["attempt"] == "/shared/startup/123-g2" for entry in entries) + assert entries[-1]["argv"] == expected_cli(cap="003") + assert entries[-1]["pid"] == pid + + +def test_exhausted_cap_does_not_start_probe(launch, supervised): + code, stderr, entries, _ = launch(**supervised, SLURM_RESTART_COUNT="3") + assert code == 2 and "allocation cap exhausted" in stderr + assert entries == [] + + +def test_probe_failure_does_not_start_cli_or_leak_requeue_exit(launch, supervised): + code, stderr, entries, _ = launch(**supervised, HELPER_EXIT="75") + assert code == 2 and "probe failed (exit 75)" in stderr + assert len(entries) == 1 and entries[0]["argv"][1] == "probe" + + +@pytest.mark.parametrize("workflow", ["sft", "prompts"]) +def test_explicit_resume_preserves_exec(launch, workflow): + code, stderr, entries, pid = launch(AUTOML_ACTION="resume", AUTOML_WORKFLOW=workflow) + assert (code, stderr) == (0, "") + assert entries == [{"argv": expected_cli(workflow, "resume"), "pid": pid, "attempt": None}] + + +@pytest.mark.parametrize("action", ["plan", "status", "cancel", "migrate", "run;true"]) +def test_nonlaunch_action_cannot_reach_helper(launch, supervised, action): + code, stderr, entries, _ = launch(**supervised, AUTOML_ACTION=action) + assert code == 2 and "AUTOML_ACTION must be run or resume" in stderr + assert entries == [] diff --git a/tools/cosmos_workflows/scripts/tests/slurm_controller_watch_test.py b/tools/cosmos_workflows/scripts/tests/slurm_controller_watch_test.py new file mode 100644 index 000000000..4c6dffb7f --- /dev/null +++ b/tools/cosmos_workflows/scripts/tests/slurm_controller_watch_test.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only observer contracts; no native scheduler or remote filesystem calls.""" + +import errno +import fcntl +import json +import os +import subprocess +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import slurm_controller_watch as observer # noqa: E402 + + +@pytest.fixture +def registry_data(tmp_path): + return { + "uid": os.getuid(), + "deadline": 2000.0, + "jobs": {"123": {"attempt_base": str(tmp_path / "attempts"), "max_allocations": 3}}, + } + + +def save_registry(root, value): + (root / "registry.json").write_text(json.dumps(value)) + + +def test_registry_round_trip(tmp_path, registry_data): + save_registry(tmp_path, registry_data) + assert observer.registry(tmp_path) == registry_data + + +@pytest.mark.parametrize( + "changes", + [ + {"uid": -1}, + {"deadline": float("inf")}, + {"deadline": float("nan")}, + {"jobs": []}, + ], +) +def test_registry_rejects_invalid_identity(tmp_path, registry_data, changes): + save_registry(tmp_path, registry_data | changes) + with pytest.raises(ValueError): + observer.registry(tmp_path) + + +@pytest.mark.parametrize("job_id", ["0", "0123", "-1", "123;true", ""]) +def test_registry_rejects_invalid_native_id(tmp_path, registry_data, job_id): + registry_data["jobs"] = {job_id: registry_data["jobs"]["123"]} + save_registry(tmp_path, registry_data) + with pytest.raises(ValueError, match="identity or allocation cap"): + observer.registry(tmp_path) + + +@pytest.mark.parametrize( + "changes", + [{"attempt_base": "relative"}, {"max_allocations": 0}, {"max_allocations": True}, {"max_allocations": 1.5}], +) +def test_registry_rejects_invalid_job_contract(tmp_path, registry_data, changes): + registry_data["jobs"]["123"].update(changes) + save_registry(tmp_path, registry_data) + with pytest.raises(ValueError, match="identity or allocation cap"): + observer.registry(tmp_path) + + +@pytest.fixture +def live_observer(tmp_path, registry_data, monkeypatch): + monkeypatch.setattr(observer.time, "time", lambda: 1000.0) + save_registry(tmp_path, registry_data) + state = {"state": "WATCHING", "uid": os.getuid(), "deadline": 2000.0, "updated_at": 999.0} + observer.atomic(tmp_path / "watch.json", state) + with (tmp_path / "watch.lock").open("a") as handle: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + yield state + + +def test_ready_requires_live_lock(tmp_path, registry_data, monkeypatch): + monkeypatch.setattr(observer.time, "time", lambda: 1000.0) + save_registry(tmp_path, registry_data) + observer.atomic( + tmp_path / "watch.json", + {"state": "WATCHING", "uid": os.getuid(), "deadline": 2000.0, "updated_at": 999.0}, + ) + with pytest.raises(ValueError, match="not running"): + observer.ready(tmp_path) + + +def test_ready_accepts_live_owner(tmp_path, live_observer): + assert observer.ready(tmp_path) == live_observer + + +@pytest.mark.parametrize( + "changes", + [ + {"state": "STOPPED"}, + {"uid": -1}, + {"deadline": 2001.0}, + {"updated_at": 910.0}, + {"updated_at": 1001.0}, + ], +) +def test_ready_rejects_stale_or_conflicting_receipt(tmp_path, live_observer, changes): + observer.atomic(tmp_path / "watch.json", live_observer | changes) + with pytest.raises(ValueError, match="not ready"): + observer.ready(tmp_path) + + +def test_register_is_idempotent_and_cannot_replace_identity(tmp_path, live_observer): + base = tmp_path / "another-attempt" + observer.register(tmp_path, "124", base, 2) + saved = (tmp_path / "registry.json").read_bytes() + observer.register(tmp_path, "124", base, 2) + assert (tmp_path / "registry.json").read_bytes() == saved + for new_base, cap in [(base, 3), (tmp_path / "other", 2)]: + with pytest.raises(ValueError, match="cannot be replaced"): + observer.register(tmp_path, "124", new_base, cap) + assert (tmp_path / "registry.json").read_bytes() == saved + + +@pytest.fixture +def native(monkeypatch): + # Any attempted submit/requeue/cancel would fail the exact-argv assertion. + fields = {"JobId": "123", "UserId": f"user({os.getuid()})", "Restarts": "0", "JobState": "RUNNING"} + + def query(argv, **kwargs): + assert argv == ["scontrol", "show", "job", "-o", "123"] + assert kwargs == {"capture_output": True, "text": True, "timeout": 30, "check": True} + return subprocess.CompletedProcess(argv, 0, " ".join(f"{k}={v}" for k, v in fields.items())) + + monkeypatch.setattr(observer.subprocess, "run", query) + probe = Mock(return_value={"phase": "NOT_ENTERED"}) + monkeypatch.setattr(observer, "observe", probe) + return fields, probe + + +@pytest.mark.parametrize("generation", [0, 1, 2]) +def test_tick_observes_each_allowed_generation(registry_data, native, generation): + fields, probe = native + fields["Restarts"] = str(generation) + entry = registry_data["jobs"]["123"] + assert observer.tick("123", entry) == {"state": "OBSERVING", "generation": generation, "phase": "NOT_ENTERED"} + probe.assert_called_once_with(Path(entry["attempt_base"]) / f"123-g{generation}", "123") + + +@pytest.mark.parametrize("state", ["PENDING", "COMPLETED", "FAILED", "CANCELLED"]) +def test_tick_never_probes_nonrunning_jobs(registry_data, native, state): + fields, probe = native + fields["JobState"] = state + assert observer.tick("123", registry_data["jobs"]["123"]) == { + "state": "OBSERVING", + "native_state": state, + "generation": 0, + } + probe.assert_not_called() + + +@pytest.mark.parametrize( + "changes", + [ + {"JobId": "124"}, + {"UserId": "other(-1)"}, + {"Restarts": "-1"}, + {"Restarts": "3"}, + {"Restarts": "999999999999999999999999999"}, + {"Restarts": "invalid"}, + ], +) +def test_tick_blocks_identity_or_generation_mismatch(registry_data, native, changes): + fields, probe = native + fields.update(changes) + assert observer.tick("123", registry_data["jobs"]["123"]) == {"state": "BLOCKED", "error_type": "ValueError"} + probe.assert_not_called() + + +def test_tick_blocks_missing_native_restart_field(registry_data, native): + fields, probe = native + del fields["Restarts"] + assert observer.tick("123", registry_data["jobs"]["123"]) == {"state": "BLOCKED", "error_type": "KeyError"} + probe.assert_not_called() + + +@pytest.mark.parametrize( + "failure", + [ + OSError(errno.EIO, "private detail"), + subprocess.TimeoutExpired("private command", 30, stderr="private detail"), + subprocess.CalledProcessError(1, "private command", stderr="private detail"), + ], +) +def test_tick_waits_on_scheduler_errors_without_recording_details(registry_data, native, monkeypatch, failure): + _, probe = native + monkeypatch.setattr(observer.subprocess, "run", Mock(side_effect=failure)) + assert observer.tick("123", registry_data["jobs"]["123"]) == { + "state": "WAITING", + "error_type": type(failure).__name__, + "errno": getattr(failure, "errno", None), + } + probe.assert_not_called() + + +def test_tick_waits_on_shared_storage_error(registry_data, native): + _, probe = native + probe.side_effect = OSError(errno.EIO, "private path") + assert observer.tick("123", registry_data["jobs"]["123"]) == { + "state": "WAITING", + "error_type": "OSError", + "errno": errno.EIO, + } diff --git a/uv.lock b/uv.lock index 8d15db1d7..733fd19ce 100644 --- a/uv.lock +++ b/uv.lock @@ -2378,6 +2378,9 @@ dependencies = [ ] [package.optional-dependencies] +automl-prompts = [ + { name = "gepa" }, +] guardrail = [ { name = "better-profanity" }, { name = "nltk" }, @@ -2708,6 +2711,7 @@ requires-dist = [ { name = "ftfy", marker = "extra == 'train'" }, { name = "futureproof", marker = "extra == 'train'" }, { name = "fvcore", marker = "extra == 'train'" }, + { name = "gepa", marker = "extra == 'automl-prompts'", specifier = "==0.1.2" }, { name = "glfw", marker = "extra == 'train'" }, { name = "gradio", marker = "extra == 'serve'" }, { name = "h5py", marker = "extra == 'train'" }, @@ -2738,8 +2742,8 @@ requires-dist = [ { name = "multi-storage-client", extras = ["boto3", "google-cloud-storage", "fsspec", "observability-otel", "vault"], marker = "extra == 'train'", specifier = "==0.44.0" }, { name = "ninja", marker = "extra == 'train'" }, { name = "nltk", marker = "extra == 'guardrail'" }, - { name = "nvidia-automl-core", marker = "extra == 'automl'", specifier = "==1.0.0.dev0" }, - { name = "nvidia-automl-core", extras = ["prompts"], marker = "extra == 'automl-prompts'", specifier = "==1.0.0.dev0" }, + { name = "nvidia-automl-core", marker = "extra == 'automl'", specifier = "==0.1.0" }, + { name = "nvidia-automl-core", marker = "extra == 'automl-prompts'", specifier = "==0.1.0" }, { name = "nvidia-cudnn-frontend" }, { name = "nvidia-dali-cuda120", marker = "extra == 'train'" }, { name = "nvidia-ml-py" }, @@ -5011,6 +5015,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1d/33/f1c6a276de27b7d7339a34749cc33fa87f077f921969c47185d34a887ae2/gast-0.7.0-py3-none-any.whl", hash = "sha256:99cbf1365633a74099f69c59bd650476b96baa5ef196fec88032b00b31ba36f7", size = 22966, upload-time = "2025-11-29T15:30:03.983Z" }, ] +[[package]] +name = "gepa" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c6/2d/9521c2b94bc4c95a81bfb483278cd3fbd1c4261f228638cb7844c84293da/gepa-0.1.2.tar.gz", hash = "sha256:1587e7081d3e321d7739785f00ba0bc27509dcc7ca4e75546d0210e1c292410c", size = 317051, upload-time = "2026-07-14T12:07:00.216Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/81/33e455d043926f95356f19f43b6e371a895708e067726071f0b5b1b05205/gepa-0.1.2-py3-none-any.whl", hash = "sha256:9ccf1e8353ec6d5d606a5aabd7a9ed75a86be5024371584646abc7d532321eb3", size = 273524, upload-time = "2026-07-14T12:06:58.865Z" }, +] + [[package]] name = "gguf" version = "0.18.0"