From 04be9de767b8ba16f502463c213d93b1921ec8cc Mon Sep 17 00:00:00 2001 From: dkn16 Date: Thu, 16 Jul 2026 02:27:16 -0700 Subject: [PATCH] Implement asynchronous Slurm job support --- docs/api/async_jobs.md | 26 + docs/api/cli.md | 17 +- docs/api/index.md | 6 +- docs/api/resources.md | 11 + docs/api/site_registry.md | 40 +- docs/api/snakefile.md | 18 +- docs/api/targets.md | 10 +- docs/architecture.md | 28 +- docs/cli/cancel.md | 23 + docs/cli/index.md | 1 + docs/cli/run.md | 59 +- docs/cli/status.md | 22 +- docs/contributing/hpc-sites.md | 9 +- docs/hpc/containers.md | 2 +- docs/hpc/site-registry.md | 10 +- docs/user/cluster.md | 96 ++-- docs/user/install.md | 10 +- docs/user/troubleshooting.md | 3 + src/lightcone/cli/commands.py | 194 ++++++- src/lightcone/engine/async_jobs.py | 749 ++++++++++++++++++++++++++ src/lightcone/engine/resources.py | 171 ++++++ src/lightcone/engine/site_registry.py | 45 +- src/lightcone/engine/snakefile.py | 10 +- tests/conftest.py | 9 + tests/test_async_jobs.py | 276 ++++++++++ tests/test_cli.py | 134 ++++- tests/test_resources.py | 61 +++ tests/test_site_registry.py | 10 + tests/test_snakefile.py | 29 + zensical.toml | 3 + 30 files changed, 1943 insertions(+), 139 deletions(-) create mode 100644 docs/api/async_jobs.md create mode 100644 docs/api/resources.md create mode 100644 docs/cli/cancel.md create mode 100644 src/lightcone/engine/async_jobs.py create mode 100644 src/lightcone/engine/resources.py create mode 100644 tests/test_async_jobs.py create mode 100644 tests/test_resources.py diff --git a/docs/api/async_jobs.md b/docs/api/async_jobs.md new file mode 100644 index 00000000..f7ceff4b --- /dev/null +++ b/docs/api/async_jobs.md @@ -0,0 +1,26 @@ +# lightcone.engine.async_jobs + +Coarse-grained SLURM submission, polling, and cancellation for +`lc run --async`. + +Source: `src/lightcone/engine/async_jobs.py`. + +The module resolves a requested ASTRA sub-DAG, aggregates its recipe +resources, selects a site policy, and renders one sbatch script. The script +re-enters ordinary `lc run`, so the async layer never duplicates recipe, +container, Dask, manifest, or validation logic. + +Important entry points: + +- `resolve_subdag_outputs()` — requested materializable outputs plus upstream + recipe dependencies. +- `aggregate_job_resources()` — element-wise maximum node shape and padded + serial walltime. +- `select_slurm_policy()` — deterministic site profile to `shared` or + `regular`. +- `submit_job()` — render, `sbatch --parsable`, and persist a `JobRecord`. +- `refresh_job_records()` — batch-query `squeue`, then `sacct`. +- `cancel_job()` — resolve an active record, call `scancel`, and update it. + +Project records live in `.lightcone/jobs/.json`; logs live under the +resolved scratch root. diff --git a/docs/api/cli.md b/docs/api/cli.md index d0634b3d..3ab2b12e 100644 --- a/docs/api/cli.md +++ b/docs/api/cli.md @@ -1,7 +1,8 @@ # lightcone.cli.commands -The Click surface. Defined in `src/lightcone/cli/commands.py`. Six -public commands: `init`, `run`, `status`, `verify`, `build`, `setup`. +The Click surface. Defined in `src/lightcone/cli/commands.py`. Public +commands include `init`, `run`, `status`, `cancel`, `verify`, `build`, and +`export`. The user-facing reference is in [CLI Overview](../cli/index.md). This page is a tour of the module internals. @@ -14,10 +15,7 @@ page is a tour of the module internals. @click.pass_context def main(ctx: click.Context) -> None: ctx.ensure_object(dict) - if ctx.invoked_subcommand in ("setup", "init", "eval"): - return - if not _config_path().exists(): - # print friendly error, sys.exit(1) + _ensure_global_config() ``` `main` is exposed as `lightcone.cli.main` (re-exported from @@ -44,8 +42,8 @@ it directly. To add a new tier, edit this dict and update the ### `_config_path() → Path` -Returns `~/.lightcone/config.yaml`. Used by the `main` group's -auto-init check and by `setup`. +Returns `~/.lightcone/config.yaml`. The main group creates it on first use +with container and SLURM defaults. ### `_project_root(start: Path | None = None) → Path` @@ -76,6 +74,9 @@ Map a status literal to the Rich-formatted display label: | `stale` | `[yellow]✸ stale[/yellow]` | | `missing` | `[red]✗ miss[/red]` | | `alias` | `[dim]→ alias[/dim]` | +| `queued` | `[blue]◷ queued[/blue]` | +| `running` | `[cyan]▶ running[/cyan]` | +| `failed` | `[red]✗ failed[/red]` | ## Boilerplate text diff --git a/docs/api/index.md b/docs/api/index.md index 491843b1..0ecb3cbe 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -7,17 +7,19 @@ is a thin Click wrapper around these modules. | Module | Role | |--------|------| -| [`lightcone.cli.commands`](cli.md) | Click CLI: `init`, `run`, `build`, `status`, `verify`, `setup`. | +| [`lightcone.cli.commands`](cli.md) | Click CLI: `init`, `run`, `build`, `status`, `cancel`, `verify`, `export`. | | [`lightcone.engine.manifest`](manifest.md) | Per-output `.lightcone-manifest.json` write/read; `code_version`, `sha256_dir`. The integrity layer. | | [`lightcone.engine.snakefile`](snakefile.md) | Generate `.lightcone/Snakefile` and `snakefile-config.json` from `astra.yaml`. | | [`lightcone.engine.container`](container.md) | Runtime detection, content-addressed image tags, `wrap_recipe`. | | [`lightcone.engine.dask_cluster`](dask_cluster.md) | Cluster lifecycle for `lc run` (local / SLURM / external). | +| [`lightcone.engine.resources`](resources.md) | Canonical ASTRA resource parsing and Snakemake-name mapping. | +| [`lightcone.engine.async_jobs`](async_jobs.md) | Coarse SLURM submission, polling, records, and cancellation. | | [`lightcone.engine.status`](status.md) | Manifest-driven status walker. | | [`lightcone.engine.verify`](verify.md) | Recompute hashes; walk the input chain. | | [`lightcone.engine.tree`](tree.md) | Sub-analysis tree helpers — outputs, decisions, `from:` resolution. | | [`lightcone.engine.validation`](validation.md) | Post-recipe sanity checks (empty dir, all-NaN columns, …). | | [`snakemake_executor_plugin_dask`](dask_executor.md) | Snakemake executor plugin → `dask.distributed`. | -| `lightcone.engine.site_registry` | Vestigial — no active code path imports it. See [api/site_registry](site_registry.md). | +| [`lightcone.engine.site_registry`](site_registry.md) | Site detection, scratch/runtime defaults, and async SLURM policies. | ## Common entry points diff --git a/docs/api/resources.md b/docs/api/resources.md new file mode 100644 index 00000000..d976b99a --- /dev/null +++ b/docs/api/resources.md @@ -0,0 +1,11 @@ +# lightcone.engine.resources + +Canonical parsing for ASTRA `Recipe.resources`. + +Source: `src/lightcone/engine/resources.py`. + +`parse_recipe_resources()` normalizes portable ASTRA values to CPUs, MB, +GPUs, and seconds. `RecipeResources.snakemake()` maps those to +`cpus_per_task`, `mem_mb`, `gpus_per_task`, and minute-based `runtime`. +Synchronous Snakefile generation and asynchronous SLURM sizing both use this +module, preventing the two paths from interpreting a recipe differently. diff --git a/docs/api/site_registry.md b/docs/api/site_registry.md index 09429604..53c38b8a 100644 --- a/docs/api/site_registry.md +++ b/docs/api/site_registry.md @@ -1,21 +1,18 @@ # lightcone.engine.site_registry -> **Status: orphaned.** Nothing in the active code path imports this -> module. It's residue from the (now-removed) target system. Keep -> reading if you want to know what's still there; otherwise skip to -> [api/container](container.md) and [api/dask_cluster](dask_cluster.md) -> for what actually drives execution today. +Site detection and deterministic per-site execution defaults. The active +container, scratch, and asynchronous SLURM paths all consume this registry. Source: `src/lightcone/engine/site_registry.py`. -## What's still in the file +## Public data and helpers The module exposes: - `SITE_DEFAULTS` — a dict mapping site keys (`"perlmutter"`, `"local"`) to a structured defaults dict (display name, hostname patterns, - suggested QoS / constraint / time-limit options, scratch deny paths, - container runtime). + suggested QoS / constraint / time-limit options, scratch paths, container + runtime, and `async_slurm` policies). - `detect_site(hostname_or_name) → str | None` - `get_site_defaults(site_key) → dict | None` - `list_known_sites() → list[tuple[str, str]]` @@ -24,25 +21,10 @@ The module exposes: The functions still work on their own; they just don't have a caller inside lightcone-cli right now. -## What's actually used today +## Async SLURM policy -The Perlmutter scratch deny rules used to be merged into -`.claude/settings.json` automatically when a non-local target was -configured. With the target system gone, the equivalent rules are -hard-coded inline in `PERMISSION_TIERS` (see -`src/lightcone/cli/commands.py`): - -```python -"Edit(//scratch/**)", -"Edit(//pscratch/**)", -``` - -If you want richer per-site rules without rebuilding the target -system, point `lc init`'s `_install_claude_plugin` at -`get_site_scratch_deny_rules(detect_site(socket.gethostname()))` and -merge the result into the deny list. Two lines of code. - -## Recommendation - -Either delete this module (no callers) or revive it for the use case -above. Leaving it as-is encourages the drift this audit is fighting. +Perlmutter declares CPU and GPU node shapes, fractional shared limits, and +the 48-hour `shared`/`regular` caps under `async_slurm`. The submission +engine selects the CPU profile when no GPU is declared and the GPU profile +otherwise. Adding another site is a table extension rather than a new branch +in the renderer. diff --git a/docs/api/snakefile.md b/docs/api/snakefile.md index ed8f3b43..8fc813a6 100644 --- a/docs/api/snakefile.md +++ b/docs/api/snakefile.md @@ -42,20 +42,20 @@ rule : output: data=directory(""), manifest="/.lightcone-manifest.json", + threads: + resources: + cpus_per_task=, + mem_mb=, + gpus_per_task=, + runtime=, params: cfg=lambda wc: CFG[""][wc.universe], run: - shell('printf "▶ [%s]\\n" "{wildcards.universe}" >&2') - shell(params.cfg["shell_command"]) - write_manifest( - output_dir=Path(output.data), - inputs={"": Path(input.), ...}, - cfg=params.cfg, - ) - for _w in validate_output(Path(output.data), params.cfg.get("output_type"), params.cfg["output_id"]): - print(f"\033[33m⚠\033[0m {_w}", file=sys.stderr) + run_rule(...) ``` +Unset optional memory/GPU/time values are omitted. CPU defaults to one. + ## `cfg` content Per-`(rule_key, universe)` entry written into diff --git a/docs/api/targets.md b/docs/api/targets.md index 4f29d61b..28beafa1 100644 --- a/docs/api/targets.md +++ b/docs/api/targets.md @@ -1,11 +1,15 @@ # lightcone.engine.targets (removed) -The target configuration module is gone. The only remaining global config -is `~/.lightcone/config.yaml`, which today carries one key: +The target configuration module is gone. The remaining global config is +`~/.lightcone/config.yaml`: ```yaml container: runtime: auto # auto | docker | podman | podman-hpc | none +slurm: + account: null + time_padding: 1.5 ``` -It is read by [`lightcone.engine.container.load_runtime`](container.md). +It is read by [`lightcone.engine.container.load_runtime`](container.md) and +[`lightcone.engine.async_jobs`](async_jobs.md). diff --git a/docs/architecture.md b/docs/architecture.md index 504fbb7b..d60a351b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -143,8 +143,9 @@ For each declared recipe input: - **`alias`** — output declared without a recipe; materialized only as a side effect of an upstream. -`status` reads only manifests. No Snakemake import, no `.snakemake/` -directory required, works on a fresh clone or frozen archive. +Recorded async jobs can temporarily overlay `queued`, `running`, `failed`, +`cancelled`, or `unknown`. `status` still does not import Snakemake or require +`.snakemake/`; it polls Slurm only when project job records exist. --- @@ -163,12 +164,25 @@ across three branches: one `dask worker` per node across the allocation. Workers advertise the node's resources via Dask abstract resources (`cpus`, `memory`, `gpus`). The Snakemake executor plugin maps per-rule - `cpus_per_task` / `mem_mb` / `gpus_per_task` to per-task constraints. + ASTRA's `cpus` / `memory` / `gpus` values are normalized to + `cpus_per_task` / `mem_mb` / `gpus_per_task` per-task constraints. 3. Neither → `LocalCluster()` sized to the local machine. The scheduler is always in-process so its lifetime equals the run's lifetime: no service to manage, no orphaned schedulers. +### Asynchronous allocations + +Module: [`lightcone.engine.async_jobs`](api/async_jobs.md). + +`lc run --async` resolves one requested universe's sub-DAG, aggregates +portable ASTRA resources, selects a site policy, and submits one coarse +sbatch job. The rendered script activates the submitting Python environment +and invokes plain `lc run`; once running, it follows branch 2 above and keeps +all container and provenance behavior identical to synchronous execution. +Job records are local caches; Slurm owns job state and manifests own output +state. + ### The Snakemake executor Module: [`snakemake_executor_plugin_dask`](api/dask_executor.md). @@ -270,18 +284,20 @@ warnings. src/lightcone/ # PEP 420 namespace package — NO __init__.py ├── cli/ # Click surface │ ├── __init__.py # exposes main() -│ ├── commands.py # init, run, status, verify, build, export +│ ├── commands.py # init, run, status, cancel, verify, build, export │ └── plugin.py # plugin source-dir discovery ├── engine/ # execution substrate │ ├── manifest.py # write_manifest, sha256_dir, code_version │ ├── snakefile.py # generate .lightcone/Snakefile from astra.yaml │ ├── container.py # docker/podman/podman-hpc build + recipe wrap │ ├── dask_cluster.py # cluster lifecycle (local/SLURM/external) +│ ├── resources.py # ASTRA resource normalization +│ ├── async_jobs.py # SLURM submit/status/cancel │ ├── status.py # manifest-driven status walker (no Snakemake) │ ├── verify.py # recompute hashes, walk the chain │ ├── tree.py # sub-analysis tree helpers │ ├── validation.py # post-recipe output sanity checks -│ └── site_registry.py # vestigial; not imported by active code +│ └── site_registry.py # site runtime, scratch, and SLURM policy └── eval/ # evaluation harness for the agent loop ├── cli.py harness.py sandbox.py graders.py build.py report.py models.py @@ -344,7 +360,7 @@ each rule to a Dask scheduler. | `.lightcone/Snakefile` | Project (generated) | Auto-generated by `lc run`. Don't edit. | | `.lightcone/snakefile-config.json` | Project (generated) | Per-`(rule, universe)` config. | | `.lightcone/lightcone.yaml` | Project | Tiny scratchpad — currently writes only `target: local`. Not consumed by today's code. | -| `~/.lightcone/config.yaml` | User | `container.runtime`. | +| `~/.lightcone/config.yaml` | User | `container.runtime`, `slurm.account`, async walltime padding. | | `.claude/settings.json` | Project | Claude Code permissions. | The `dagster.yaml` and `~/.lightcone/targets/*.yaml` files referenced in diff --git a/docs/cli/cancel.md b/docs/cli/cancel.md new file mode 100644 index 00000000..80ebd257 --- /dev/null +++ b/docs/cli/cancel.md @@ -0,0 +1,23 @@ +# lc cancel + +Cancel a recorded asynchronous SLURM job. + +## Synopsis + +```text +lc cancel +``` + +`lc cancel` refreshes recorded state, resolves the reference to one active +job, calls `scancel`, and stores `CANCELLED` in the job record. An output id +may refer to either an explicitly requested target or one of its resolved +upstream dependencies. If it matches multiple active jobs, cancel by job id +to disambiguate. + +```bash +lc cancel 1234567 +lc cancel heavy_fit +``` + +Only jobs recorded under the current project's `.lightcone/jobs/` directory +are managed. diff --git a/docs/cli/index.md b/docs/cli/index.md index 33ffba0b..aab5bfd0 100644 --- a/docs/cli/index.md +++ b/docs/cli/index.md @@ -19,6 +19,7 @@ Code skills, with the CLI as the durable, scriptable backstop. | [`lc run`](run.md) | Generate the Snakefile and dispatch through Snakemake + Dask. | | [`lc build`](build.md) | Build container images declared in `astra.yaml`. | | [`lc status`](status.md) | Manifest-driven status report. No Snakemake import needed. | +| [`lc cancel`](cancel.md) | Cancel a recorded asynchronous SLURM job. | | [`lc verify`](verify.md) | Recompute hashes, walk the input chain, surface tampering. | | [`lc export`](export.md) | Emit interoperable bundles (Workflow Run RO-Crate) for publication. | diff --git a/docs/cli/run.md b/docs/cli/run.md index 3e311be9..3414b6da 100644 --- a/docs/cli/run.md +++ b/docs/cli/run.md @@ -21,6 +21,8 @@ everything (Snakemake's `rule all`). | `--rerun-triggers TRIGGERS` | `code,input,mtime,params` | Comma-separated rerun triggers (forwarded to Snakemake). | | `--force`, `-f` | off | `--force` when targets are named, `--forceall` otherwise. | | `--verbose`, `-v` | off | Show the underlying Snakemake / executor chatter and the spawned `snakemake` invocation. | +| `--async` | off | Submit one coarse SLURM job per selected universe instead of executing immediately. Perlmutter only in v1. | +| `--account NAME` | `slurm.account` in `~/.lightcone/config.yaml` | Override the SLURM account for this async submission. Requires `--async`. | ## What happens, step by step @@ -43,6 +45,11 @@ everything (Snakemake's `rule all`). chatter so the output reads as lightcone's, not Snakemake's. Real error content always passes through. +With `--async`, steps 3–8 happen later inside the batch allocation. The +submission process instead resolves the requested sub-DAG, validates its +resources, chooses `shared` or `regular`, renders an sbatch script, submits +it, and writes `.lightcone/jobs/.json`. + ## Output qualification When the same `output_id` appears in multiple sub-analyses, you must @@ -69,8 +76,51 @@ lc run accuracy precision --universe baseline # several lc run --jobs 4 --verbose # parallel, with stack noise lc run --force --universe baseline # rebuild everything lc run --rerun-triggers params,input # tighter staleness +lc run --async accuracy -u baseline # submit one queued batch job +``` + +## Asynchronous SLURM jobs + +First set the account once: + +```yaml +# ~/.lightcone/config.yaml +container: + runtime: podman-hpc +slurm: + account: m1234 + time_padding: 1.5 ``` +Every recipe in the selected output's resolved sub-DAG must declare a +`time_limit`. Other resources default to one CPU, no requested memory, and +no GPU: + +```yaml +recipe: + command: python scripts/fit.py --output {output} + resources: + cpus: 16 + memory: 64GB + gpus: 1 + time_limit: 2h +``` + +The allocation shape is the element-wise maximum across those recipes. The +walltime is their serial sum multiplied by `slurm.time_padding`. On +Perlmutter, a job fitting the shared CPU/GPU profile uses `shared`; anything +larger uses `regular`, up to the 48-hour cap. A recipe larger than one node or +a padded walltime beyond the cap is rejected. + +Submission is allowed from both login and compute nodes. If `-u` is omitted, +Lightcone submits one job per discovered universe. The generated script +activates the environment containing the current Python executable and runs +plain `lc run ...` inside the allocation. It therefore uses the same Dask, +Snakemake, validation, manifest, and container path as a synchronous run, +including `podman-hpc`; it never installs packages on a compute node. + +Use `lc status` to poll and [`lc cancel`](cancel.md) to cancel a recorded job. + ## Inside SLURM ```bash @@ -80,9 +130,12 @@ lc run --universe baseline -j 16 `lc run` detects `SLURM_JOB_ID`, binds the Dask scheduler to the driver's hostname, and launches one `dask worker` per node via `srun`. -Workers advertise `cpus`, `memory`, and `gpus` resources. Per-rule -resource hints (`cpus_per_task`, `mem_mb`, `gpus_per_task`) constrain -which workers can pick up which jobs. +Workers advertise `cpus`, `memory`, and `gpus` resources. ASTRA recipe +resources are translated to Snakemake's `cpus_per_task`, `mem_mb`, +`gpus_per_task`, and `runtime` names, constraining which workers can pick up +which jobs. Before starting Dask, Lightcone checks that every selected rule +fits one worker in the current allocation; an impossible shape fails with a +larger-allocation / `--async` hint instead of waiting indefinitely. ## Provenance gotcha diff --git a/docs/cli/status.md b/docs/cli/status.md index 7985c3a5..f7cacbd2 100644 --- a/docs/cli/status.md +++ b/docs/cli/status.md @@ -1,7 +1,7 @@ # lc status -Manifest-driven status report for every output declared in -`astra.yaml`. +Manifest-driven status report for every output declared in `astra.yaml`, +annotated with recorded asynchronous SLURM jobs. ## Synopsis @@ -26,6 +26,8 @@ Universe baseline ✸ stale precision ✗ miss recall → alias inference + ◷ queued simulation (job 1234567, regular) + ▶ running fit (job 1234568, shared) ``` Statuses (defined in `lightcone.engine.status.StatusLiteral`): @@ -37,11 +39,20 @@ Statuses (defined in `lightcone.engine.status.StatusLiteral`): | `missing` | No manifest at the expected output path. | Never built, or the directory was deleted. | | `alias` | The output has no `recipe:` of its own — it's just a name pointing at a sibling output (typical for ASTRA "promoted" outputs from sub-analyses). | Status is implicitly determined by the upstream. | +When `.lightcone/jobs/*.json` records exist, `lc status` batch-queries +`squeue` and `sacct`. Affected outputs can additionally show `queued`, +`running`, `failed`, `cancelled`, or `unknown`, with the job id, QoS, and +failure log where relevant. Completed jobs fall back to the manifest status, +because manifests—not scheduler history—remain the source of truth for +outputs. JSON output includes a nullable `job` object beside each output's +materialization `status`. + ## Why it doesn't import Snakemake -`lc status` reads only the per-output `.lightcone-manifest.json` files -and recomputes `code_version` against the current spec. It never -imports Snakemake or touches `.snakemake/`. That makes it usable on: +`lc status` reads per-output manifests and small async job records. It never +imports Snakemake or touches `.snakemake/`; scheduler polling happens only +when records exist. If Slurm commands are unavailable, cached states are +preserved. That makes it usable on: - A fresh clone before any `lc run`. - A frozen archive copied off a cluster. @@ -63,5 +74,6 @@ lc status --json # machine-readable JSON output - [`lc verify`](verify.md) — recomputes data hashes too (slower; catches tampering and broken chains). +- [`lc cancel`](cancel.md) — cancel a queued or running recorded job. - [api/status](../api/status.md) — the Python API. - [api/manifest](../api/manifest.md) — the manifest schema. diff --git a/docs/contributing/hpc-sites.md b/docs/contributing/hpc-sites.md index 94df4d40..2fd0918b 100644 --- a/docs/contributing/hpc-sites.md +++ b/docs/contributing/hpc-sites.md @@ -1,8 +1,6 @@ -# Adding an HPC Site (deprecated) +# Adding an HPC Site -There is no longer a meaningful concept of "adding an HPC site" — the -target system that this used to feed is gone. The `site_registry` module -is still present in the source tree but unused. See +Site-specific defaults live in `engine/site_registry.py`. See [api/site_registry](../api/site_registry.md). If you want lightcone-cli to behave well on a new cluster, what you @@ -19,3 +17,6 @@ actually need is: tiers hard-code Perlmutter scratch deny rules; if you add a new site, update `PERMISSION_TIERS` in `src/lightcone/cli/commands.py`. +4. **An async SLURM policy.** Add CPU/GPU node shapes, shared thresholds, + constraints, and QoS walltime caps under `async_slurm`. Keep policy as + data so the renderer remains site-independent. diff --git a/docs/hpc/containers.md b/docs/hpc/containers.md index 786dc86d..bd57966a 100644 --- a/docs/hpc/containers.md +++ b/docs/hpc/containers.md @@ -12,7 +12,7 @@ without a registry. ```bash # On a login node, with podman-hpc on PATH: -lc setup # writes ~/.lightcone/config.yaml +lc status # creates ~/.lightcone/config.yaml $EDITOR ~/.lightcone/config.yaml # set container.runtime: podman-hpc lc build # builds + migrates each image ``` diff --git a/docs/hpc/site-registry.md b/docs/hpc/site-registry.md index 49dd39e7..8ebe5749 100644 --- a/docs/hpc/site-registry.md +++ b/docs/hpc/site-registry.md @@ -1,6 +1,6 @@ -# Site Registry (orphaned) +# Site Registry -The `lightcone.engine.site_registry` module still exists but is not imported -by any active code path. It carries Perlmutter scheduler defaults that used -to feed the wizard for the (now removed) target system. See -[api/site_registry](../api/site_registry.md) for the current state. +`lightcone.engine.site_registry` is the active source for site detection, +preferred container runtime, scratch placement, and deterministic async +SLURM policy. See [api/site_registry](../api/site_registry.md) for the data +shape and Perlmutter profiles. diff --git a/docs/user/cluster.md b/docs/user/cluster.md index ffa73f11..5b3dab38 100644 --- a/docs/user/cluster.md +++ b/docs/user/cluster.md @@ -1,9 +1,8 @@ # Running on a Cluster When local laptop time isn't enough, you can take the same project to -a SLURM HPC system. There's no separate configuration to learn — the -same `lc run` command works inside an allocation, just with more -hardware to spread across. +a SLURM HPC system. The same execution path supports interactive work +inside an allocation and queued work submitted with `lc run --async`. ## The big picture @@ -19,6 +18,10 @@ hardware to spread across. You don't pick — `lc run` detects which case applies. The only thing you do differently on a cluster is request the nodes. +For longer work, `lc run --async` is an outer submission layer. It asks +SLURM for an allocation and then executes ordinary `lc run` inside it, so +recipe behavior and provenance stay identical. + ## Pre-flight: pick the right container runtime On most HPC sites, docker isn't available on compute nodes. Most @@ -32,8 +35,15 @@ $EDITOR ~/.lightcone/config.yaml ```yaml container: runtime: podman-hpc +slurm: + account: m1234 + time_padding: 1.5 ``` +`slurm.account` is required the first time you use `--async`; alternatively, +pass `--account m1234` for one submission. `time_padding` multiplies the +serial sum of recipe time estimates and defaults to `1.5`. + Then build and migrate the images for your project: ```bash @@ -54,45 +64,26 @@ provenance warning.) ## A typical SLURM workflow -### 1. Get an allocation +### 1. Choose interactive or queued execution + +For short, agent-in-the-loop work, get an interactive allocation and run +synchronously: ```bash -salloc -N 4 -t 02:00:00 -C gpu # interactive -# or -sbatch run.sbatch # batch +salloc -A m1234 -N 1 -t 02:00:00 -C gpu -q interactive +lc run heavy_fit -u baseline ``` -`run.sbatch` looks like: - -=== "Generic" - ```bash - #!/bin/bash - #SBATCH -N 4 - #SBATCH -t 02:00:00 - #SBATCH -C gpu - - cd $HOME/my-analysis - source .venv/bin/activate - lc run -j 16 - ``` +For long or unattended work on Perlmutter, submit it from either a login or +compute node: -=== "NERSC Perlmutter" - ```bash - #!/bin/bash - #SBATCH -A - #SBATCH -q regular - #SBATCH -C gpu - #SBATCH -N 4 - #SBATCH -t 04:00:00 - - cd $SCRATCH/your-analysis - - # make `lc` available — pick the line that matches your install: - export PATH=$HOME/.local/bin:$PATH # uv tool install - # source ~/.conda/envs/your-env-name/bin/activate # conda env +```bash +lc run --async heavy_fit -u baseline +lc status +``` - lc run -j 16 - ``` +Lightcone chooses `shared` or `regular`, writes the script under +`.lightcone/jobs/`, stores logs in scratch, and records the returned job id. ### 2. `lc run` inside the allocation @@ -105,7 +96,7 @@ Once `SLURM_JOB_ID` is set in your environment, `lc run` does the rest: per-recipe `resources:` constraints land on workers that can hold them. -### 3. Per-recipe resource hints +### 3. Per-recipe resources Add resource hints in your `astra.yaml` recipe blocks: @@ -114,16 +105,18 @@ outputs: - id: heavy_fit type: metric recipe: - command: python scripts/fit.py --output {output[0]} + command: python scripts/fit.py --output {output} resources: - cpus_per_task: 32 - mem_mb: 64000 - gpus_per_task: 1 + cpus: 16 + memory: 64GB + gpus: 1 + time_limit: 2h ``` -The Snakemake-via-Dask executor maps these to per-task resource -requests, so a rule that needs a GPU only schedules on nodes that -advertise one. +Lightcone maps these portable ASTRA names to Snakemake/Dask resource +requests for synchronous execution and uses the same declaration to size +asynchronous jobs. Every recipe in an asynchronously submitted sub-DAG needs +`time_limit`; the other fields have conservative zero/one defaults. ## Interactive: agent-driven runs @@ -141,21 +134,22 @@ claude # or whichever agent CLI you prefer Everything the agent triggers (`lc run`, scripts, etc.) now executes on the allocated node. When you're done iterating and want a -hands-off sweep of all universes, submit `lc run` as a batch job -instead (the sbatch template above). +hands-off work, use `lc run --async`. This is also valid while the agent is +still on the compute node; `sbatch` queues an independent allocation. ## What about login-node-only operations? -Build images, dry-run, look at status — all fine on a login node -without an allocation: +Build images, submit queued work, and look at status are all fine on a login +node without an allocation: ```bash lc build # build images (uses podman-hpc on login node) -lc status # offline; reads only manifests +lc run --async heavy_fit # submits; does not execute on the login node +lc status # manifests + on-demand Slurm polling ``` -The actual `lc run` should happen inside an allocation, since that's -where the worker nodes are. +Plain `lc run` must happen inside an allocation. On Perlmutter login nodes it +fails with an allocation example and a hint to try `--async`. ## External Dask schedulers diff --git a/docs/user/install.md b/docs/user/install.md index 9ea81ff2..17460c49 100644 --- a/docs/user/install.md +++ b/docs/user/install.md @@ -133,16 +133,22 @@ Get a confirmation of the proper installation by running ## 3. Global configuration `~/.lightcone/config.yaml` is created automatically the first time you -run any `lc` command. No manual setup step is needed. The file starts +run an `lc` subcommand. No manual setup step is needed. The file starts as: ```yaml container: runtime: auto +slurm: + account: null + time_padding: 1.5 ``` `auto` detects whichever of `podman`, `docker`, or `podman-hpc` is on -your PATH (and skips docker if its daemon isn't running). Feel free to pin the runtime later by editing this file directly. +your PATH (and skips docker if its daemon isn't running). Feel free to pin the +runtime later by editing this file directly. Before the first +`lc run --async`, replace `slurm.account: null` with the SLURM project to +charge; synchronous and local work do not require it. ## 4. Agentic CLI diff --git a/docs/user/troubleshooting.md b/docs/user/troubleshooting.md index b7978b3a..3ad10baa 100644 --- a/docs/user/troubleshooting.md +++ b/docs/user/troubleshooting.md @@ -14,6 +14,9 @@ mkdir -p ~/.lightcone cat > ~/.lightcone/config.yaml <<'EOF' container: runtime: auto +slurm: + account: null + time_padding: 1.5 EOF ``` diff --git a/src/lightcone/cli/commands.py b/src/lightcone/cli/commands.py index 259f8428..624ff477 100644 --- a/src/lightcone/cli/commands.py +++ b/src/lightcone/cli/commands.py @@ -10,6 +10,7 @@ MyST report template). - ``lc run`` — generate Snakefile and run snakemake. - ``lc status`` — manifest-driven status walk (no Snakemake needed). +- ``lc cancel`` — cancel a recorded asynchronous SLURM job. - ``lc verify`` — recompute hashes and validate the provenance chain. - ``lc build`` — build containers from Containerfiles. @@ -97,6 +98,10 @@ def _ensure_global_config() -> None: # its daemon is unreachable); set explicitly to pin. ``none`` # disables containerization entirely. "container": {"runtime": "auto"}, + # Account is intentionally unset: allocations are site/user + # specific and must never be guessed. Async walltime is the + # serial sub-DAG estimate multiplied by this safety padding. + "slurm": {"account": None, "time_padding": 1.5}, } ) ) @@ -511,10 +516,67 @@ def _abort_on_perlmutter_login() -> None: "run inside a SLURM allocation.\n" " Start one with, e.g.:\n" " salloc -N 1 -C gpu -q interactive -t 1:00:00 -A \n" - " then re-run `lc run` from inside." + " then re-run `lc run` from inside.\n" + " If the recipes declare resources.time_limit, submit them instead with:\n" + " lc run --async --account " ) +def _abort_if_sync_resources_exceed_allocation( + project: Path, + outputs: tuple[str, ...], +) -> None: + """Fail before Dask waits forever on an unschedulable SLURM rule.""" + if "SLURM_JOB_ID" not in os.environ or os.environ.get("DASK_SCHEDULER_ADDRESS"): + return + + from lightcone.engine.async_jobs import AsyncJobError, resolve_subdag_outputs + from lightcone.engine.dask_cluster import _detect_node_shape + from lightcone.engine.resources import ResourceValueError, parse_recipe_resources + + try: + _, subdag = resolve_subdag_outputs(project, outputs) + except AsyncJobError as exc: + raise click.ClickException(str(exc)) from exc + + shape = _detect_node_shape() + available_memory_mb = shape.mem_bytes // 1_000_000 + shortfalls: list[str] = [] + for tree_output in subdag: + qualified = ( + f"{tree_output.analysis_id}.{tree_output.output_id}" + if tree_output.analysis_id + else tree_output.output_id + ) + try: + resources = parse_recipe_resources( + tree_output.output_def.get("recipe") or {}, + label=f"output {qualified!r}", + ) + except ResourceValueError as exc: + raise click.ClickException(str(exc)) from exc + exceeded: list[str] = [] + if resources.cpus > shape.cpus: + exceeded.append(f"CPUs {resources.cpus} > {shape.cpus}") + if resources.memory_mb and available_memory_mb: + if resources.memory_mb > available_memory_mb: + exceeded.append( + f"memory {resources.memory_mb} MB > {available_memory_mb} MB" + ) + if resources.gpus > shape.gpus: + exceeded.append(f"GPUs {resources.gpus} > {shape.gpus}") + if exceeded: + shortfalls.append(f" - {qualified}: {', '.join(exceeded)}") + + if shortfalls: + raise click.ClickException( + "The current SLURM allocation cannot satisfy every selected recipe:\n" + + "\n".join(shortfalls) + + "\nRequest a larger interactive allocation or submit an independent " + "batch allocation with `lc run --async`." + ) + + @main.command() @click.argument("outputs", nargs=-1) @click.option("--universe", "-u", default=None, help="Universe to materialize") @@ -526,6 +588,17 @@ def _abort_on_perlmutter_login() -> None: ) @click.option("--force", "-f", is_flag=True, help="Force re-materialization") @click.option("--verbose", "-v", is_flag=True, help="Show full executor output") +@click.option( + "--async", + "asynchronous", + is_flag=True, + help="Submit one SLURM batch job per selected universe.", +) +@click.option( + "--account", + default=None, + help="SLURM account for --async (overrides ~/.lightcone/config.yaml).", +) def run( outputs: tuple[str, ...], universe: str | None, @@ -533,6 +606,8 @@ def run( rerun_triggers: str, force: bool, verbose: bool, + asynchronous: bool, + account: str | None, ) -> None: """Materialize outputs declared in astra.yaml. @@ -540,7 +615,42 @@ def run( workstation, srun-launched workers inside a SLURM allocation, or an existing scheduler if ``DASK_SCHEDULER_ADDRESS`` is set. """ + project = _project_root() + + from lightcone.engine.snakefile import discover_universes + + universes = [universe] if universe else discover_universes(project) + if asynchronous: + from lightcone.engine.async_jobs import AsyncJobError, format_slurm_time, submit_job + + try: + for universe_id in universes: + record = submit_job( + project, + output_ids=outputs, + universe=universe_id, + account_override=account, + jobs=jobs, + rerun_triggers=rerun_triggers, + force=force, + verbose=verbose, + ) + console.print( + f"[green]Submitted job {record.job_id}[/green] " + f"([cyan]{record.qos}[/cyan], universe " + f"[cyan]{record.universe}[/cyan], " + f"walltime {format_slurm_time(int(record.resources['time_limit_seconds']))}).\n" + f" Script: [cyan]{record.sbatch_path}[/cyan]\n" + f" Log: [cyan]{record.log_path}[/cyan]" + ) + except AsyncJobError as exc: + raise click.ClickException(str(exc)) from exc + return + + if account is not None: + raise click.ClickException("--account is only valid together with --async.") _abort_on_perlmutter_login() + _abort_if_sync_resources_exceed_allocation(project, outputs) from lightcone.engine.container import load_runtime from lightcone.engine.dask_cluster import cluster_for_run @@ -551,10 +661,7 @@ def run( prepare_run_dirs, resolve_scratch_root, ) - from lightcone.engine.snakefile import discover_universes, generate - - project = _project_root() - universes = [universe] if universe else discover_universes(project) + from lightcone.engine.snakefile import generate # Resolve scratch and prepare per-run directories before anything # else. Snakemake's ``.snakemake/`` is redirected via symlink so its @@ -799,12 +906,38 @@ def _target_for(project: Path, output_id: str, universe: str) -> str: help="Emit machine-readable JSON instead of a styled table.", ) def status(universe: str | None, as_json: bool) -> None: - """Report materialization status for every declared output.""" + """Report materialization and recorded async-job status.""" + from lightcone.engine.async_jobs import ( + JobRecord, + job_display_state, + latest_job_for_output, + refresh_job_records, + ) from lightcone.engine.snakefile import discover_universes from lightcone.engine.status import get_output_status project = _project_root() universes = [universe] if universe else discover_universes(project) + records = refresh_job_records(project) + + def job_for(status_item: object, universe_id: str) -> JobRecord | None: + analysis_id = getattr(status_item, "analysis_id") + output_id = getattr(status_item, "output_id") + qualified = f"{analysis_id}.{output_id}" if analysis_id else output_id + return latest_job_for_output( + records, output_id=qualified, universe=universe_id + ) + + def job_payload(record: JobRecord | None) -> dict[str, str] | None: + if record is None: + return None + return { + "job_id": record.job_id, + "state": job_display_state(record.last_state), + "slurm_state": record.last_state, + "qos": record.qos, + "log_path": record.log_path, + } if as_json: payload = { @@ -817,6 +950,7 @@ def status(universe: str | None, as_json: bool) -> None: "analysis_id": s.analysis_id, "status": s.status, "recipe_command": s.recipe_command, + "job": job_payload(job_for(s, u)), } for s in get_output_status(project, universe_id=u) ], @@ -830,9 +964,26 @@ def status(universe: str | None, as_json: bool) -> None: for u in universes: console.print(f"\n[bold]Universe[/bold] [cyan]{u}[/cyan]") for s in get_output_status(project, universe_id=u): - label = _status_label(s.status) + record = job_for(s, u) + state = job_display_state(record.last_state) if record else None + active_state = ( + state + if state in {"queued", "running", "failed", "cancelled", "unknown"} + else None + ) + label = _status_label(active_state or s.status) scope = f"[dim]{s.analysis_id}.[/dim]" if s.analysis_id else "" - console.print(f" {label} {scope}{s.output_id}") + detail = "" + if record and state in {"queued", "running"}: + detail = f" [dim](job {record.job_id}, {record.qos})[/dim]" + elif record and state == "failed": + detail = ( + f" [dim](job {record.job_id}, {record.last_state}; " + f"log: {record.log_path})[/dim]" + ) + elif record and state in {"cancelled", "unknown"}: + detail = f" [dim](job {record.job_id}, {record.last_state})[/dim]" + console.print(f" {label} {scope}{s.output_id}{detail}") _STATUS_STYLES = { @@ -840,6 +991,11 @@ def status(universe: str | None, as_json: bool) -> None: "stale": "[yellow]✸ stale[/yellow] ", "missing": "[red]✗ miss[/red] ", "alias": "[dim]→ alias[/dim] ", + "queued": "[blue]◷ queued[/blue] ", + "running": "[cyan]▶ running[/cyan]", + "failed": "[red]✗ failed[/red] ", + "cancelled": "[dim]■ cancel[/dim] ", + "unknown": "[yellow]? unknown[/yellow]", } @@ -847,6 +1003,28 @@ def _status_label(s: str) -> str: return _STATUS_STYLES.get(s, s) +# ============================================================================= +# lc cancel +# ============================================================================= + + +@main.command() +@click.argument("job_or_output") +def cancel(job_or_output: str) -> None: + """Cancel an async job by recorded job id or output id.""" + from lightcone.engine.async_jobs import AsyncJobError, cancel_job + + project = _project_root() + try: + record = cancel_job(project, job_or_output) + except AsyncJobError as exc: + raise click.ClickException(str(exc)) from exc + console.print( + f"[green]Cancelled job {record.job_id}[/green] " + f"([cyan]{record.qos}[/cyan], universe [cyan]{record.universe}[/cyan])." + ) + + # ============================================================================= # lc verify # ============================================================================= diff --git a/src/lightcone/engine/async_jobs.py b/src/lightcone/engine/async_jobs.py new file mode 100644 index 00000000..4e5f330d --- /dev/null +++ b/src/lightcone/engine/async_jobs.py @@ -0,0 +1,749 @@ +"""Submit, track, and cancel coarse-grained Lightcone SLURM jobs. + +The batch layer deliberately does not execute recipes itself. It resolves the +requested ASTRA sub-DAG, allocates one SLURM job from declared resources, and +re-enters the ordinary synchronous ``lc run`` command inside that allocation. +That keeps container wrapping, Dask dispatch, manifests, and validation identical +between synchronous and asynchronous runs. +""" +from __future__ import annotations + +import json +import math +import os +import re +import shlex +import shutil +import subprocess +import sys +from dataclasses import asdict, dataclass, replace +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Literal + +import yaml +from astra.helpers import load_yaml, resolve_analysis_tree + +from lightcone.engine.resources import RecipeResources, ResourceValueError, parse_recipe_resources +from lightcone.engine.scratch import project_hash, resolve_scratch_root +from lightcone.engine.site_registry import HostSite, detect_current_site +from lightcone.engine.tree import TreeOutput, collect_tree_outputs, find_upstream_output + + +class AsyncJobError(RuntimeError): + """An asynchronous job cannot be safely submitted or managed.""" + + +@dataclass(frozen=True) +class JobResources: + """Aggregate allocation requirements for a resolved sub-DAG.""" + + cpus: int + memory_mb: int + gpus: int + time_limit_seconds: int + rule_count: int + + +@dataclass(frozen=True) +class SlurmSelection: + """Site-policy result used to render an sbatch script.""" + + site: str + profile: str + constraint: str + qos: str + nodes: int = 1 + allocation_cpus: int | None = None + allocation_memory_mb: int | None = None + allocation_gpus: int | None = None + + +@dataclass(frozen=True) +class SlurmSettings: + """User-controlled SLURM submission settings.""" + + account: str + time_padding: float + + +@dataclass(frozen=True) +class JobRecord: + """Small local cache record for one submitted job.""" + + job_id: str + targets: list[str] + resolved_targets: list[str] + universe: str + qos: str + resources: dict[str, Any] + sbatch_path: str + log_path: str + submitted_at: str + last_state: str + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> JobRecord: + targets = [str(value) for value in data.get("targets", [])] + resolved = [str(value) for value in data.get("resolved_targets", targets)] + return cls( + job_id=str(data["job_id"]), + targets=targets, + resolved_targets=resolved, + universe=str(data["universe"]), + qos=str(data["qos"]), + resources=dict(data.get("resources") or {}), + sbatch_path=str(data["sbatch_path"]), + log_path=str(data["log_path"]), + submitted_at=str(data["submitted_at"]), + last_state=str(data.get("last_state") or "UNKNOWN"), + ) + + +JobDisplayState = Literal[ + "queued", "running", "completed", "failed", "cancelled", "unknown" +] + + +def _qualified_output_id(tree_output: TreeOutput) -> str: + if tree_output.analysis_id: + return f"{tree_output.analysis_id}.{tree_output.output_id}" + return tree_output.output_id + + +def _resolve_requested_output( + output_id: str, + materializable: list[TreeOutput], +) -> TreeOutput: + matches = [ + tree_output + for tree_output in materializable + if _qualified_output_id(tree_output) == output_id + or tree_output.output_id == output_id + ] + if not matches: + raise AsyncJobError( + f"Output {output_id!r} not found in astra.yaml or has no recipe." + ) + if len(matches) > 1: + choices = ", ".join(_qualified_output_id(match) for match in matches) + raise AsyncJobError( + f"Output {output_id!r} is ambiguous; qualify it as one of: {choices}" + ) + return matches[0] + + +def resolve_subdag_outputs( + project_path: Path, + output_ids: tuple[str, ...], +) -> tuple[list[str], list[TreeOutput]]: + """Resolve requested outputs and all recipe-producing dependencies.""" + project_path = project_path.resolve() + spec = resolve_analysis_tree(load_yaml(project_path / "astra.yaml"), project_path) + all_outputs = collect_tree_outputs(spec) + materializable = [ + tree_output + for tree_output in all_outputs + if tree_output.output_def.get("recipe") is not None + ] + selected = ( + [_resolve_requested_output(output_id, materializable) for output_id in output_ids] + if output_ids + else materializable + ) + + required: set[str] = set() + visiting: set[str] = set() + + def visit(tree_output: TreeOutput) -> None: + key = _qualified_output_id(tree_output) + if key in required: + return + if key in visiting: + raise AsyncJobError(f"Cycle detected while resolving the sub-DAG at {key!r}.") + visiting.add(key) + for input_id in tree_output.output_def.get("inputs") or []: + upstream = find_upstream_output(tree_output, input_id, all_outputs) + if upstream is not None: + visit(upstream) + visiting.remove(key) + required.add(key) + + for tree_output in selected: + visit(tree_output) + + requested = [_qualified_output_id(tree_output) for tree_output in selected] + resolved = [ + tree_output + for tree_output in materializable + if _qualified_output_id(tree_output) in required + ] + return requested, resolved + + +def aggregate_job_resources( + tree_outputs: list[TreeOutput], + *, + time_padding: float, +) -> JobResources: + """Aggregate a sub-DAG to one node shape and conservative walltime.""" + if not tree_outputs: + raise AsyncJobError("The selected sub-DAG contains no materializable recipes.") + if not math.isfinite(time_padding) or time_padding <= 0: + raise AsyncJobError("slurm.time_padding must be a finite number greater than zero.") + + parsed: list[RecipeResources] = [] + missing: list[str] = [] + for tree_output in tree_outputs: + label = _qualified_output_id(tree_output) + recipe = tree_output.output_def.get("recipe") or {} + resources = recipe.get("resources") or {} + if not resources.get("time_limit"): + missing.append(label) + continue + try: + parsed.append( + parse_recipe_resources( + recipe, + require_time_limit=True, + label=f"output {label!r}", + ) + ) + except ResourceValueError as exc: + raise AsyncJobError(str(exc)) from exc + + if missing: + names = "\n".join(f" - {name}" for name in missing) + raise AsyncJobError( + "Async submission requires resources.time_limit on every recipe " + "in the resolved sub-DAG. Missing:\n" + f"{names}\n" + "Run the Lightcone resource-estimation skill, update astra.yaml, " + "then retry `lc run --async`." + ) + + return JobResources( + cpus=max(resource.cpus for resource in parsed), + memory_mb=max(resource.memory_mb for resource in parsed), + gpus=max(resource.gpus for resource in parsed), + time_limit_seconds=math.ceil( + sum(resource.time_limit_seconds or 0 for resource in parsed) * time_padding + ), + rule_count=len(parsed), + ) + + +def select_slurm_policy( + resources: JobResources, + *, + site: HostSite | None = None, +) -> SlurmSelection: + """Map aggregate resources to the site's deterministic QoS policy.""" + current_site = site or detect_current_site() + policy = current_site.get("async_slurm") + if not current_site or not isinstance(policy, dict): + raise AsyncJobError( + "Automatic async submission is currently supported only on a " + "configured SLURM site (v1: NERSC Perlmutter)." + ) + + profile_name = "gpu" if resources.gpus else "cpu" + profiles = policy.get("profiles") or {} + profile = profiles.get(profile_name) + qos_policy = policy.get("qos") or {} + if not isinstance(profile, dict): + raise AsyncJobError(f"Site {current_site.key!r} has no {profile_name!r} profile.") + + limits = { + "CPUs": (resources.cpus, int(profile["node_cpus"])), + "memory MB": (resources.memory_mb, int(profile["node_memory_mb"])), + "GPUs": (resources.gpus, int(profile["node_gpus"])), + } + exceeded = [ + f"{name} {value} > {maximum}" + for name, (value, maximum) in limits.items() + if value > maximum + ] + if exceeded: + raise AsyncJobError( + "A single recipe exceeds one Perlmutter node (" + + ", ".join(exceeded) + + "). Multi-node recipe shapes are outside async v1; restructure the work." + ) + + shared_max = int(qos_policy["shared"]["max_time_seconds"]) + if profile_name == "gpu": + shared_gpus = int(profile["shared_gpus"]) + shared_cpus = resources.gpus * int(profile["shared_cpus_per_gpu"]) + shared_memory = resources.gpus * int(profile["shared_memory_mb_per_gpu"]) + fits_shared_shape = ( + 1 <= resources.gpus <= shared_gpus + and resources.cpus <= shared_cpus + and resources.memory_mb <= shared_memory + ) + else: + fits_shared_shape = ( + resources.cpus <= int(profile["shared_cpus"]) + and resources.memory_mb <= int(profile["shared_memory_mb"]) + ) + + if fits_shared_shape and resources.time_limit_seconds <= shared_max: + qos = "shared" + else: + regular_max = int(qos_policy["regular"]["max_time_seconds"]) + if resources.time_limit_seconds > regular_max: + raise AsyncJobError( + "Aggregated padded walltime " + f"({format_slurm_time(resources.time_limit_seconds)}) exceeds the " + f"regular QoS cap ({format_slurm_time(regular_max)}). " + "Restructure or split the work before submitting." + ) + qos = "regular" + + allocation_cpus = None + allocation_memory_mb = None + allocation_gpus = None + if qos == "shared": + if profile_name == "gpu": + allocation_cpus = shared_cpus + allocation_memory_mb = shared_memory + allocation_gpus = resources.gpus + else: + allocation_cpus = resources.cpus + allocation_memory_mb = resources.memory_mb or None + + return SlurmSelection( + site=str(current_site.key), + profile=profile_name, + constraint=str(profile["constraint"]), + qos=qos, + allocation_cpus=allocation_cpus, + allocation_memory_mb=allocation_memory_mb, + allocation_gpus=allocation_gpus, + ) + + +_ACCOUNT_RE = re.compile(r"^[A-Za-z0-9_.-]+$") + + +def load_slurm_settings(*, account_override: str | None = None) -> SlurmSettings: + """Load account and padding from ``~/.lightcone/config.yaml``.""" + config_path = Path.home() / ".lightcone" / "config.yaml" + try: + data = yaml.safe_load(config_path.read_text()) or {} + except (OSError, yaml.YAMLError) as exc: + raise AsyncJobError(f"Could not read {config_path}: {exc}") from exc + if not isinstance(data, dict): + raise AsyncJobError(f"{config_path} must contain a YAML mapping.") + slurm = data.get("slurm") or {} + if not isinstance(slurm, dict): + raise AsyncJobError(f"{config_path} slurm entry must be a mapping.") + + account_value = account_override or slurm.get("account") + account = str(account_value).strip() if account_value else "" + if not account: + raise AsyncJobError( + "No SLURM account configured. Set `slurm.account` in " + "~/.lightcone/config.yaml or pass `lc run --async --account `." + ) + if _ACCOUNT_RE.fullmatch(account) is None: + raise AsyncJobError(f"Invalid SLURM account {account!r}.") + + raw_padding = slurm.get("time_padding", 1.5) + if isinstance(raw_padding, bool): + raise AsyncJobError("slurm.time_padding must be a number greater than zero.") + try: + padding = float(raw_padding) + except (TypeError, ValueError) as exc: + raise AsyncJobError("slurm.time_padding must be a number greater than zero.") from exc + if not math.isfinite(padding) or padding <= 0: + raise AsyncJobError("slurm.time_padding must be a finite number greater than zero.") + return SlurmSettings(account=account, time_padding=padding) + + +def format_slurm_time(seconds: int) -> str: + """Format seconds as a SLURM-compatible ``HH:MM:SS`` duration.""" + hours, remainder = divmod(seconds, 3600) + minutes, secs = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{secs:02d}" + + +def _slug(value: str, *, maximum: int = 48) -> str: + slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-.") or "all" + return slug[:maximum] + + +def jobs_dir(project_path: Path) -> Path: + """Return the project-local async job-record directory.""" + return project_path / ".lightcone" / "jobs" + + +def _environment_commands() -> tuple[list[str], str]: + """Return environment setup lines and the same-environment ``lc`` path.""" + executable = Path(sys.executable).absolute() + bin_dir = executable.parent + activate = bin_dir / "activate" + lines: list[str] = [] + if activate.is_file(): + lines.append(f"source {shlex.quote(str(activate))}") + lines.append(f"export PATH={shlex.quote(str(bin_dir))}:\"$PATH\"") + sibling_lc = bin_dir / "lc" + lc_executable = sibling_lc if sibling_lc.is_file() else Path(shutil.which("lc") or "lc") + return lines, str(lc_executable) + + +def render_sbatch_script( + *, + project_path: Path, + account: str, + resources: JobResources, + selection: SlurmSelection, + log_template: Path, + job_name: str, + lc_args: list[str], +) -> str: + """Render one sbatch script which re-enters the synchronous run path.""" + setup_lines, lc_executable = _environment_commands() + command = shlex.join([lc_executable, *lc_args]) + lines = [ + "#!/bin/bash", + f"#SBATCH --job-name={_slug(job_name, maximum=64)}", + f"#SBATCH --account={account}", + f"#SBATCH --qos={selection.qos}", + f"#SBATCH --constraint={selection.constraint}", + f"#SBATCH --nodes={selection.nodes}", + f"#SBATCH --time={format_slurm_time(resources.time_limit_seconds)}", + f"#SBATCH --output={log_template}", + f"#SBATCH --error={log_template}", + ] + if selection.allocation_gpus: + lines.append(f"#SBATCH --gpus={selection.allocation_gpus}") + if selection.allocation_cpus: + lines.append(f"#SBATCH --cpus-per-task={selection.allocation_cpus}") + if selection.allocation_memory_mb: + lines.append(f"#SBATCH --mem={selection.allocation_memory_mb}M") + lines.extend( + [ + "", + "set -euo pipefail", + f"cd {shlex.quote(str(project_path))}", + # Never inherit an interactive/external scheduler into the batch + # job. The ordinary run path must create its scheduler inside + # this new allocation. + "unset DASK_SCHEDULER_ADDRESS", + *setup_lines, + f"exec {command}", + "", + ] + ) + return "\n".join(lines) + + +def _write_record(project_path: Path, record: JobRecord) -> None: + directory = jobs_dir(project_path) + directory.mkdir(parents=True, exist_ok=True) + destination = directory / f"{record.job_id}.json" + temporary = directory / f".{record.job_id}.{os.getpid()}.tmp" + temporary.write_text(json.dumps(asdict(record), indent=2, sort_keys=True) + "\n") + os.replace(temporary, destination) + + +def submit_job( + project_path: Path, + *, + output_ids: tuple[str, ...], + universe: str, + account_override: str | None = None, + jobs: int | None = None, + rerun_triggers: str = "code,input,mtime,params", + force: bool = False, + verbose: bool = False, +) -> JobRecord: + """Resolve, render, submit, and record one universe-scoped async job.""" + project_path = project_path.resolve() + settings = load_slurm_settings(account_override=account_override) + requested, subdag = resolve_subdag_outputs(project_path, output_ids) + resources = aggregate_job_resources(subdag, time_padding=settings.time_padding) + selection = select_slurm_policy(resources) + + directory = jobs_dir(project_path) + directory.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%S%fZ") + target_slug = _slug("-".join(requested) if output_ids else "all") + script_path = directory / f"{stamp}-{target_slug}.sbatch" + + log_directory = ( + resolve_scratch_root(project_path) + / ".lightcone" + / "jobs" + / project_hash(project_path) + ) + log_directory.mkdir(parents=True, exist_ok=True) + log_template = log_directory / f"{stamp}-{target_slug}-%j.log" + + lc_args = ["run", *requested, "--universe", universe] + if jobs is not None: + lc_args.extend(["--jobs", str(jobs)]) + if rerun_triggers != "code,input,mtime,params": + lc_args.extend(["--rerun-triggers", rerun_triggers]) + if force: + lc_args.append("--force") + if verbose: + lc_args.append("--verbose") + + script = render_sbatch_script( + project_path=project_path, + account=settings.account, + resources=resources, + selection=selection, + log_template=log_template, + job_name=f"lc-{target_slug}-{universe}", + lc_args=lc_args, + ) + script_path.write_text(script) + script_path.chmod(0o750) + + try: + result = subprocess.run( + ["sbatch", "--parsable", str(script_path)], + capture_output=True, + text=True, + check=False, + ) + except FileNotFoundError as exc: + raise AsyncJobError("`sbatch` was not found on PATH.") from exc + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "unknown error" + raise AsyncJobError(f"sbatch failed: {detail}\nRendered script: {script_path}") + match = re.match(r"\s*(\d+)", result.stdout) + if match is None: + raise AsyncJobError( + f"Could not parse job id from sbatch output {result.stdout.strip()!r}." + ) + job_id = match.group(1) + + record = JobRecord( + job_id=job_id, + targets=requested, + resolved_targets=[_qualified_output_id(tree_output) for tree_output in subdag], + universe=universe, + qos=selection.qos, + resources={ + **asdict(resources), + "constraint": selection.constraint, + "profile": selection.profile, + }, + sbatch_path=str(script_path), + log_path=str(log_template).replace("%j", job_id), + submitted_at=datetime.now(UTC).isoformat(), + last_state="PENDING", + ) + _write_record(project_path, record) + return record + + +def load_job_records(project_path: Path) -> list[JobRecord]: + """Load valid project job records, oldest first.""" + directory = jobs_dir(project_path) + if not directory.is_dir(): + return [] + records: list[JobRecord] = [] + for path in directory.glob("*.json"): + try: + data = json.loads(path.read_text()) + if isinstance(data, dict): + records.append(JobRecord.from_dict(data)) + except (OSError, ValueError, KeyError, TypeError): + continue + return sorted(records, key=lambda record: record.submitted_at) + + +def job_display_state(raw_state: str) -> JobDisplayState: + """Collapse SLURM's detailed states to Lightcone's status vocabulary.""" + state = raw_state.strip().upper().split(maxsplit=1)[0].rstrip("+") + if state in {"PENDING", "CONFIGURING", "REQUEUED", "REQUEUE_FED", "SUBMITTED"}: + return "queued" + if state in {"RUNNING", "COMPLETING", "STAGE_OUT", "RESIZING", "SUSPENDED"}: + return "running" + if state == "COMPLETED": + return "completed" + if state.startswith("CANCELLED"): + return "cancelled" + if state in { + "BOOT_FAIL", + "DEADLINE", + "FAILED", + "NODE_FAIL", + "OUT_OF_MEMORY", + "PREEMPTED", + "REVOKED", + "SPECIAL_EXIT", + "TIMEOUT", + }: + return "failed" + return "unknown" + + +def _parse_state_lines(output: str, requested: set[str]) -> dict[str, str]: + states: dict[str, str] = {} + for line in output.splitlines(): + pieces = [piece.strip() for piece in line.split("|", 1)] + if len(pieces) != 2: + continue + job_id, state = pieces + if job_id in requested and state: + states[job_id] = state + return states + + +def _scheduler_query(command: list[str]) -> subprocess.CompletedProcess[str] | None: + try: + return subprocess.run(command, capture_output=True, text=True, check=False) + except FileNotFoundError: + return None + + +def query_job_states(job_ids: list[str]) -> tuple[dict[str, str], bool]: + """Batch-query active jobs with squeue, then missing jobs with sacct. + + Returns ``(states, scheduler_available)``. An offline archive keeps + cached states instead of converting every record to ``unknown``. + """ + if not job_ids: + return {}, True + requested = set(job_ids) + states: dict[str, str] = {} + scheduler_available = False + + squeue = _scheduler_query( + ["squeue", "--noheader", "--jobs", ",".join(job_ids), "--format=%i|%T"] + ) + if squeue is not None and squeue.returncode == 0: + scheduler_available = True + states.update(_parse_state_lines(squeue.stdout, requested)) + + missing = requested - states.keys() + if missing: + sacct = _scheduler_query( + [ + "sacct", + "--noheader", + "--parsable2", + "--allocations", + "--jobs", + ",".join(sorted(missing)), + "--format=JobIDRaw,State", + ] + ) + if sacct is not None and sacct.returncode == 0: + scheduler_available = True + states.update(_parse_state_lines(sacct.stdout, missing)) + return states, scheduler_available + + +def refresh_job_records(project_path: Path) -> list[JobRecord]: + """Refresh non-terminal records from SLURM and persist changed states.""" + records = load_job_records(project_path) + active = [ + record + for record in records + if job_display_state(record.last_state) not in {"completed", "failed", "cancelled"} + ] + states, scheduler_available = query_job_states([record.job_id for record in active]) + if not scheduler_available: + return records + + refreshed: list[JobRecord] = [] + for record in records: + if record not in active: + refreshed.append(record) + continue + raw_state = states.get(record.job_id, "UNKNOWN") + updated = replace(record, last_state=raw_state) + if updated != record: + _write_record(project_path, updated) + refreshed.append(updated) + return refreshed + + +def latest_job_for_output( + records: list[JobRecord], + *, + output_id: str, + universe: str, +) -> JobRecord | None: + """Return the newest job whose resolved sub-DAG contains an output.""" + matches = [ + record + for record in records + if record.universe == universe + and output_id in (record.resolved_targets or record.targets) + ] + return max(matches, key=lambda record: record.submitted_at, default=None) + + +def _resolve_cancel_record(records: list[JobRecord], reference: str) -> JobRecord: + if reference.isdigit(): + matches = [record for record in records if record.job_id == reference] + else: + matches = [ + record + for record in records + if reference in (record.resolved_targets or record.targets) + or any( + target.rsplit(".", 1)[-1] == reference + for target in (record.resolved_targets or record.targets) + ) + ] + active = [ + record + for record in matches + if job_display_state(record.last_state) not in {"completed", "failed", "cancelled"} + ] + if not active: + raise AsyncJobError(f"No active recorded job matches {reference!r}.") + if len(active) > 1: + job_ids = ", ".join(record.job_id for record in active) + raise AsyncJobError( + f"{reference!r} matches multiple active jobs ({job_ids}); cancel by job id." + ) + return active[0] + + +def cancel_job(project_path: Path, reference: str) -> JobRecord: + """Cancel one recorded job by job id or unambiguous output id.""" + records = refresh_job_records(project_path) + record = _resolve_cancel_record(records, reference) + try: + result = subprocess.run( + ["scancel", record.job_id], capture_output=True, text=True, check=False + ) + except FileNotFoundError as exc: + raise AsyncJobError("`scancel` was not found on PATH.") from exc + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "unknown error" + raise AsyncJobError(f"scancel failed for job {record.job_id}: {detail}") + cancelled = replace(record, last_state="CANCELLED") + _write_record(project_path, cancelled) + return cancelled + + +__all__ = [ + "AsyncJobError", + "JobRecord", + "JobResources", + "SlurmSelection", + "aggregate_job_resources", + "cancel_job", + "format_slurm_time", + "job_display_state", + "latest_job_for_output", + "load_job_records", + "load_slurm_settings", + "query_job_states", + "refresh_job_records", + "render_sbatch_script", + "resolve_subdag_outputs", + "select_slurm_policy", + "submit_job", +] diff --git a/src/lightcone/engine/resources.py b/src/lightcone/engine/resources.py new file mode 100644 index 00000000..00c3d614 --- /dev/null +++ b/src/lightcone/engine/resources.py @@ -0,0 +1,171 @@ +"""Canonical ASTRA recipe-resource parsing. + +ASTRA records portable resource hints on ``Recipe.resources`` using the +``cpus`` / ``memory`` / ``gpus`` / ``time_limit`` vocabulary. Lightcone +normalizes those values once here so synchronous Snakemake rules and +asynchronous SLURM submissions cannot interpret them differently. +""" +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from typing import Any + + +class ResourceValueError(ValueError): + """A recipe resource value cannot be interpreted safely.""" + + +@dataclass(frozen=True) +class RecipeResources: + """Normalized resources for one recipe.""" + + cpus: int = 1 + memory_mb: int = 0 + gpus: int = 0 + time_limit_seconds: int | None = None + + def snakemake(self) -> dict[str, int]: + """Return the canonical Snakemake resource-name mapping. + + Snakemake's conventional ``runtime`` resource is expressed in + minutes. All other normalized values are already in the units + consumed by Lightcone's Dask executor. + """ + result = {"cpus_per_task": self.cpus} + if self.memory_mb: + result["mem_mb"] = self.memory_mb + if self.gpus: + result["gpus_per_task"] = self.gpus + if self.time_limit_seconds is not None: + result["runtime"] = math.ceil(self.time_limit_seconds / 60) + return result + + +_MEMORY_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([kmgt]?i?b)\s*$", re.IGNORECASE) +_MEMORY_MB = { + "b": 1 / 1_000_000, + "kb": 1 / 1_000, + "kib": 1024 / 1_000_000, + "mb": 1, + "mib": 1024**2 / 1_000_000, + "gb": 1_000, + "gib": 1024**3 / 1_000_000, + "tb": 1_000_000, + "tib": 1024**4 / 1_000_000, +} + + +def parse_memory_mb(value: str) -> int: + """Parse an ASTRA memory string (for example ``8GB``) into MB.""" + match = _MEMORY_RE.fullmatch(value) + if match is None: + raise ResourceValueError( + f"Invalid memory value {value!r}; use a value such as '8GB' or '512MB'." + ) + amount = float(match.group(1)) + return max(1, math.ceil(amount * _MEMORY_MB[match.group(2).lower()])) + + +_TIME_TOKEN_RE = re.compile(r"(\d+(?:\.\d+)?)([smhd])", re.IGNORECASE) +_TIME_MULTIPLIER = {"s": 1, "m": 60, "h": 3600, "d": 86400} + + +def parse_time_seconds(value: str) -> int: + """Parse an ASTRA duration into seconds. + + Supports compact durations (``30m``, ``2h``, ``1h30m``) and SLURM-style + ``HH:MM:SS`` / ``D-HH:MM:SS`` strings. + """ + raw = value.strip() + if not raw: + raise ResourceValueError("time_limit must not be empty.") + + if ":" in raw: + day_part = "0" + clock = raw + if "-" in raw: + day_part, clock = raw.split("-", 1) + pieces = clock.split(":") + if len(pieces) not in {2, 3} or not day_part.isdigit(): + raise ResourceValueError(f"Invalid time_limit {value!r}.") + try: + numbers = [int(piece) for piece in pieces] + except ValueError as exc: + raise ResourceValueError(f"Invalid time_limit {value!r}.") from exc + if len(numbers) == 2: + hours = 0 + minutes, seconds = numbers + else: + hours, minutes, seconds = numbers + if minutes >= 60 or seconds >= 60: + raise ResourceValueError(f"Invalid time_limit {value!r}.") + total = int(day_part) * 86400 + hours * 3600 + minutes * 60 + seconds + if total <= 0: + raise ResourceValueError("time_limit must be greater than zero.") + return total + + compact_total = 0.0 + position = 0 + for match in _TIME_TOKEN_RE.finditer(raw): + if match.start() != position: + raise ResourceValueError( + f"Invalid time_limit {value!r}; use a value such as '2h' or '30m'." + ) + compact_total += ( + float(match.group(1)) * _TIME_MULTIPLIER[match.group(2).lower()] + ) + position = match.end() + if position != len(raw) or compact_total <= 0: + raise ResourceValueError( + f"Invalid time_limit {value!r}; use a value such as '2h' or '30m'." + ) + return math.ceil(compact_total) + + +def parse_recipe_resources( + recipe: dict[str, Any], + *, + require_time_limit: bool = False, + label: str = "recipe", +) -> RecipeResources: + """Normalize one ASTRA recipe's resource declaration.""" + raw = recipe.get("resources") or {} + if not isinstance(raw, dict): + raise ResourceValueError(f"{label} resources must be a mapping.") + + cpus = raw.get("cpus") or 1 + gpus = raw.get("gpus") or 0 + if not isinstance(cpus, int) or isinstance(cpus, bool) or cpus < 1: + raise ResourceValueError(f"{label} resources.cpus must be an integer >= 1.") + if not isinstance(gpus, int) or isinstance(gpus, bool) or gpus < 0: + raise ResourceValueError(f"{label} resources.gpus must be an integer >= 0.") + + memory = raw.get("memory") + if memory is not None and not isinstance(memory, str): + raise ResourceValueError(f"{label} resources.memory must be a string such as '8GB'.") + memory_mb = parse_memory_mb(memory) if memory else 0 + + time_limit = raw.get("time_limit") + if require_time_limit and not time_limit: + raise ResourceValueError(f"{label} is missing resources.time_limit.") + if time_limit is not None and not isinstance(time_limit, str): + raise ResourceValueError(f"{label} resources.time_limit must be a string such as '2h'.") + time_limit_seconds = parse_time_seconds(time_limit) if time_limit else None + + return RecipeResources( + cpus=cpus, + memory_mb=memory_mb, + gpus=gpus, + time_limit_seconds=time_limit_seconds, + ) + + +__all__ = [ + "RecipeResources", + "ResourceValueError", + "parse_memory_mb", + "parse_recipe_resources", + "parse_time_seconds", +] diff --git a/src/lightcone/engine/site_registry.py b/src/lightcone/engine/site_registry.py index f4c00351..eb3e92f5 100644 --- a/src/lightcone/engine/site_registry.py +++ b/src/lightcone/engine/site_registry.py @@ -14,6 +14,7 @@ """ from __future__ import annotations +import os import socket from collections.abc import Mapping from dataclasses import dataclass, field @@ -33,6 +34,35 @@ "hostname": "perlmutter.nersc.gov", }, "container_runtime": "podman-hpc", + # Deterministic policy for ``lc run --async``. Profiles are + # separate because Perlmutter's CPU and GPU nodes have different + # shapes, and GPU shared jobs receive 16 CPU cores + 64 GB host + # memory per requested GPU. + "async_slurm": { + "qos": { + "shared": {"max_time_seconds": 48 * 60 * 60}, + "regular": {"max_time_seconds": 48 * 60 * 60}, + }, + "profiles": { + "cpu": { + "constraint": "cpu", + "node_cpus": 128, + "node_memory_mb": 512_000, + "node_gpus": 0, + "shared_cpus": 64, + "shared_memory_mb": 256_000, + }, + "gpu": { + "constraint": "gpu", + "node_cpus": 64, + "node_memory_mb": 256_000, + "node_gpus": 4, + "shared_gpus": 2, + "shared_cpus_per_gpu": 16, + "shared_memory_mb_per_gpu": 64_000, + }, + }, + }, # Where lightcone keeps its operational state (snakemake metadata, # dask spill, cross-node stdout locks). NERSC's $HOME and CFS are # mounted on compute via DVS, which silently swallows ``flock`` and @@ -161,7 +191,20 @@ def detect_current_site() -> HostSite: :class:`HostSite` (``key is None``) when the hostname matches no known site. """ - key = detect_site(socket.gethostname()) + # Perlmutter compute-node hostnames are commonly ``nidNNNNNN`` and do + # not identify the system. NERSC_HOST is set consistently on login and + # compute nodes; LIGHTCONE_SITE is a portable explicit override for a + # future site whose node names have the same problem. + candidates = ( + os.environ.get("LIGHTCONE_SITE"), + os.environ.get("NERSC_HOST"), + socket.gethostname(), + ) + key = None + for value in candidates: + if value and (detected := detect_site(value)): + key = detected + break if key is None: return _UNKNOWN_HOST_SITE return HostSite(key=key, defaults=get_site_defaults(key) or {}) diff --git a/src/lightcone/engine/snakefile.py b/src/lightcone/engine/snakefile.py index 87b20158..f1d6d509 100644 --- a/src/lightcone/engine/snakefile.py +++ b/src/lightcone/engine/snakefile.py @@ -31,6 +31,7 @@ from lightcone.engine.container import make_image_tag_resolver, wrap_recipe from lightcone.engine.manifest import code_version +from lightcone.engine.resources import parse_recipe_resources from lightcone.engine.tree import ( TreeOutput, collect_tree_outputs, @@ -252,7 +253,8 @@ def _render_snakefile( """Render the Snakefile string from rule descriptors. Each rule descriptor has ``name`` (Snakemake-safe), ``key`` (cfg - lookup), ``output_dir`` (wildcard pattern), and ``inputs`` (list of + lookup), ``output_dir`` (wildcard pattern), ``resources`` (canonical + Snakemake names), and ``inputs`` (list of ``(raw_id, safe_key, path_pattern)`` triples — both sibling outputs and analysis-level Inputs, in declaration order). External Input patterns are static source strings; sibling-output patterns carry a @@ -298,6 +300,10 @@ def _render_snakefile( lines.append(" output:") lines.append(f' data=directory("{r["output_dir"]}"),') lines.append(f' manifest="{r["output_dir"]}/.lightcone-manifest.json",') + lines.append(f' threads: {r["resources"]["cpus_per_task"]}') + lines.append(" resources:") + for resource_name, value in r["resources"].items(): + lines.append(f" {resource_name}={value},") lines.append(" params:") lines.append(f' cfg=lambda wc: CFG["{r["key"]}"][wc.universe],') lines.append(" run:") @@ -361,6 +367,7 @@ def generate( rule_key = _rule_key(to) rule_name = _rule_name(to) out_dir_pattern = _output_dir_pattern(to) + resources = parse_recipe_resources(recipe, label=f"output {rule_key!r}") # v0.0.7: declared upstream inputs live on the Output, not the # Recipe. Each ID resolves to either a sibling output (a @@ -392,6 +399,7 @@ def generate( "key": rule_key, "output_dir": out_dir_pattern, "inputs": rule_inputs, + "resources": resources.snakemake(), } ) diff --git a/tests/conftest.py b/tests/conftest.py index baba0659..1aec1db8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,12 @@ """Shared test fixtures for lightcone-cli tests.""" from __future__ import annotations + +import pytest + + +@pytest.fixture(autouse=True) +def _clear_site_environment(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep host-site tests independent of the machine running pytest.""" + monkeypatch.delenv("LIGHTCONE_SITE", raising=False) + monkeypatch.delenv("NERSC_HOST", raising=False) diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py new file mode 100644 index 00000000..3094c5f2 --- /dev/null +++ b/tests/test_async_jobs.py @@ -0,0 +1,276 @@ +"""Tests for coarse-grained asynchronous SLURM jobs.""" +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from typing import Any + +import pytest +import yaml + +import lightcone.engine.async_jobs as async_jobs +from lightcone.engine.async_jobs import ( + AsyncJobError, + JobRecord, + JobResources, + SlurmSelection, + aggregate_job_resources, + cancel_job, + job_display_state, + query_job_states, + render_sbatch_script, + resolve_subdag_outputs, + select_slurm_policy, + submit_job, +) +from lightcone.engine.site_registry import SITE_DEFAULTS, HostSite + + +def _write_project(path: Path, outputs: list[dict[str, Any]]) -> None: + path.mkdir(parents=True) + (path / "astra.yaml").write_text(yaml.safe_dump({"outputs": outputs})) + universes = path / "universes" + universes.mkdir() + (universes / "baseline.yaml").write_text("decisions: {}\n") + + +def _perlmutter() -> HostSite: + return HostSite(key="perlmutter", defaults=SITE_DEFAULTS["perlmutter"]) + + +def test_resolve_subdag_and_aggregate_resources(tmp_path: Path) -> None: + project = tmp_path / "project" + _write_project( + project, + [ + { + "id": "prepare", + "recipe": { + "command": "python prepare.py", + "resources": {"cpus": 2, "memory": "2GB", "time_limit": "10m"}, + }, + }, + { + "id": "fit", + "inputs": ["prepare"], + "recipe": { + "command": "python fit.py", + "resources": { + "cpus": 16, + "memory": "64GB", + "gpus": 1, + "time_limit": "2h", + }, + }, + }, + { + "id": "unrelated", + "recipe": { + "command": "echo no", + "resources": {"time_limit": "5m"}, + }, + }, + ], + ) + requested, subdag = resolve_subdag_outputs(project, ("fit",)) + assert requested == ["fit"] + assert [output.output_id for output in subdag] == ["prepare", "fit"] + + resources = aggregate_job_resources(subdag, time_padding=1.5) + assert resources == JobResources( + cpus=16, + memory_mb=64000, + gpus=1, + time_limit_seconds=11700, + rule_count=2, + ) + + +def test_async_requires_time_limit_on_every_dependency(tmp_path: Path) -> None: + project = tmp_path / "project" + _write_project( + project, + [ + {"id": "prepare", "recipe": {"command": "echo prepare"}}, + { + "id": "fit", + "inputs": ["prepare"], + "recipe": { + "command": "echo fit", + "resources": {"time_limit": "1h"}, + }, + }, + ], + ) + _, subdag = resolve_subdag_outputs(project, ("fit",)) + with pytest.raises(AsyncJobError, match="prepare"): + aggregate_job_resources(subdag, time_padding=1.5) + + +def test_policy_prefers_cpu_shared() -> None: + selection = select_slurm_policy( + JobResources(32, 128000, 0, 3600, 1), site=_perlmutter() + ) + assert selection.qos == "shared" + assert selection.profile == "cpu" + assert selection.allocation_cpus == 32 + assert selection.allocation_memory_mb == 128000 + + +def test_policy_uses_gpu_shared_allocation_ratio() -> None: + shared = select_slurm_policy( + JobResources(16, 64000, 1, 3600, 1), site=_perlmutter() + ) + regular = select_slurm_policy( + JobResources(32, 64000, 1, 3600, 1), site=_perlmutter() + ) + assert shared.qos == "shared" + assert shared.profile == "gpu" + assert regular.qos == "regular" + + +def test_policy_rejects_walltime_beyond_regular_cap() -> None: + with pytest.raises(AsyncJobError, match="exceeds the regular QoS cap"): + select_slurm_policy( + JobResources(1, 0, 0, 48 * 3600 + 1, 1), site=_perlmutter() + ) + + +def test_render_script_reenters_plain_lc_run(tmp_path: Path) -> None: + script = render_sbatch_script( + project_path=tmp_path, + account="m1234", + resources=JobResources(16, 64000, 1, 7200, 1), + selection=SlurmSelection( + "perlmutter", + "gpu", + "gpu", + "shared", + allocation_cpus=16, + allocation_memory_mb=64000, + allocation_gpus=1, + ), + log_template=tmp_path / "job-%j.log", + job_name="lc-fit-baseline", + lc_args=["run", "fit", "--universe", "baseline"], + ) + assert "#SBATCH --account=m1234" in script + assert "#SBATCH --qos=shared" in script + assert "#SBATCH --gpus=1" in script + assert "#SBATCH --cpus-per-task=16" in script + assert "#SBATCH --mem=64000M" in script + assert "unset DASK_SCHEDULER_ADDRESS" in script + assert " run fit --universe baseline" in script + assert "--async" not in script + assert "pip install" not in script + + +def test_submit_job_writes_script_and_record( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "project" + _write_project( + project, + [ + { + "id": "fit", + "recipe": { + "command": "python fit.py", + "resources": { + "cpus": 16, + "memory": "64GB", + "gpus": 1, + "time_limit": "2h", + }, + }, + } + ], + ) + home = tmp_path / "home" + (home / ".lightcone").mkdir(parents=True) + (home / ".lightcone" / "config.yaml").write_text( + "container:\n runtime: auto\nslurm:\n account: m1234\n time_padding: 1.5\n" + ) + monkeypatch.setattr(Path, "home", lambda: home) + monkeypatch.setattr(async_jobs, "detect_current_site", _perlmutter) + monkeypatch.setattr(async_jobs, "resolve_scratch_root", lambda _: tmp_path / "scratch") + + calls: list[list[str]] = [] + + def fake_run(command: list[str], **_: Any) -> subprocess.CompletedProcess[str]: + calls.append(command) + return subprocess.CompletedProcess(command, 0, stdout="12345;perlmutter\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + record = submit_job(project, output_ids=("fit",), universe="baseline") + + assert record.job_id == "12345" + assert record.qos == "shared" + assert calls == [["sbatch", "--parsable", record.sbatch_path]] + assert Path(record.sbatch_path).is_file() + assert Path(record.sbatch_path).read_text().count("lc") > 0 + saved = json.loads((project / ".lightcone" / "jobs" / "12345.json").read_text()) + assert saved["resolved_targets"] == ["fit"] + assert saved["resources"]["time_limit_seconds"] == 10800 + + +def test_query_job_states_uses_squeue_then_sacct( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fake_query(command: list[str]) -> subprocess.CompletedProcess[str]: + if command[0] == "squeue": + return subprocess.CompletedProcess(command, 0, stdout="11|RUNNING\n", stderr="") + return subprocess.CompletedProcess(command, 0, stdout="22|TIMEOUT\n", stderr="") + + monkeypatch.setattr(async_jobs, "_scheduler_query", fake_query) + states, available = query_job_states(["11", "22"]) + assert available is True + assert states == {"11": "RUNNING", "22": "TIMEOUT"} + + +@pytest.mark.parametrize( + ("raw", "display"), + [ + ("PENDING", "queued"), + ("RUNNING", "running"), + ("COMPLETED", "completed"), + ("TIMEOUT", "failed"), + ("CANCELLED by 1", "cancelled"), + ("mystery", "unknown"), + ], +) +def test_job_display_state(raw: str, display: str) -> None: + assert job_display_state(raw) == display + + +def test_cancel_job_updates_record(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + project = tmp_path / "project" + directory = project / ".lightcone" / "jobs" + directory.mkdir(parents=True) + record = JobRecord( + job_id="123", + targets=["fit"], + resolved_targets=["prepare", "fit"], + universe="baseline", + qos="regular", + resources={}, + sbatch_path="job.sbatch", + log_path="job.log", + submitted_at="2026-07-16T00:00:00+00:00", + last_state="PENDING", + ) + (directory / "123.json").write_text(json.dumps(record.__dict__)) + monkeypatch.setattr(async_jobs, "query_job_states", lambda _: ({"123": "PENDING"}, True)) + + calls: list[list[str]] = [] + + def fake_run(command: list[str], **_: Any) -> subprocess.CompletedProcess[str]: + calls.append(command) + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + cancelled = cancel_job(project, "prepare") + assert cancelled.last_state == "CANCELLED" + assert calls == [["scancel", "123"]] + assert json.loads((directory / "123.json").read_text())["last_state"] == "CANCELLED" diff --git a/tests/test_cli.py b/tests/test_cli.py index a5c03fc1..b86b0516 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -33,7 +33,7 @@ def _isolated_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: def test_help_lists_core_commands(runner: CliRunner) -> None: result = runner.invoke(main, ["--help"]) assert result.exit_code == 0 - for cmd in ("init", "run", "status", "verify", "build"): + for cmd in ("init", "run", "status", "verify", "build", "cancel"): assert cmd in result.output @@ -58,6 +58,8 @@ def test_first_invocation_auto_creates_global_config( assert result.exit_code == 0, result.output assert config.exists() assert "runtime: auto" in config.read_text() + assert "account: null" in config.read_text() + assert "time_padding: 1.5" in config.read_text() # ---- lc init -------------------------------------------------------------- @@ -232,3 +234,133 @@ def test_run_cmd_multiple_triggers_all_before_separator() -> None: for trigger in ("code", "input", "mtime", "params"): assert trigger in cmd, f"trigger '{trigger}' missing from cmd" assert cmd.index(trigger) < sep_idx, f"trigger '{trigger}' must come before '--'" + + +def test_sync_run_on_perlmutter_login_hints_async( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "project" + project.mkdir() + (project / "astra.yaml").write_text( + "outputs:\n - id: foo\n recipe:\n command: echo foo\n" + ) + monkeypatch.chdir(project) + monkeypatch.setenv("NERSC_HOST", "perlmutter") + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + monkeypatch.delenv("DASK_SCHEDULER_ADDRESS", raising=False) + result = runner.invoke(main, ["run", "foo"]) + assert result.exit_code != 0 + assert "lc run --async" in result.output + + +def test_sync_run_rejects_rule_larger_than_current_allocation( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "project" + project.mkdir() + (project / "astra.yaml").write_text( + "outputs:\n" + " - id: fit\n" + " recipe:\n" + " command: echo fit\n" + " resources:\n" + " cpus: 16\n" + " memory: 64GB\n" + " gpus: 2\n" + ) + monkeypatch.chdir(project) + monkeypatch.setenv("SLURM_JOB_ID", "999") + monkeypatch.setenv("SLURM_CPUS_ON_NODE", "32") + monkeypatch.setenv("SLURM_MEM_PER_NODE", "128000") + monkeypatch.setenv("SLURM_GPUS_ON_NODE", "1") + result = runner.invoke(main, ["run", "fit"]) + assert result.exit_code != 0 + assert "GPUs 2 > 1" in result.output + assert "lc run --async" in result.output + + +def test_async_run_submits_each_discovered_universe( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lightcone.engine import async_jobs + from lightcone.engine.async_jobs import JobRecord + + project = tmp_path / "project" + project.mkdir() + (project / "astra.yaml").write_text( + "outputs:\n" + " - id: foo\n" + " recipe:\n" + " command: echo foo\n" + " resources:\n" + " time_limit: 10m\n" + ) + universes = project / "universes" + universes.mkdir() + (universes / "baseline.yaml").write_text("decisions: {}\n") + (universes / "alternate.yaml").write_text("decisions: {}\n") + monkeypatch.chdir(project) + + submitted: list[str] = [] + + def fake_submit(project_path: Path, **kwargs: object) -> JobRecord: + universe = str(kwargs["universe"]) + submitted.append(universe) + return JobRecord( + job_id=str(len(submitted)), + targets=["foo"], + resolved_targets=["foo"], + universe=universe, + qos="shared", + resources={"time_limit_seconds": 900}, + sbatch_path=str(project_path / f"{universe}.sbatch"), + log_path=str(project_path / f"{universe}.log"), + submitted_at="2026-07-16T00:00:00+00:00", + last_state="PENDING", + ) + + monkeypatch.setattr(async_jobs, "submit_job", fake_submit) + result = runner.invoke(main, ["run", "--async", "foo", "--account", "m1234"]) + assert result.exit_code == 0, result.output + assert submitted == ["alternate", "baseline"] + assert result.output.count("Submitted job") == 2 + + +def test_status_json_includes_async_job( + runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import json + + from lightcone.engine import async_jobs + from lightcone.engine.async_jobs import JobRecord + + project = tmp_path / "project" + project.mkdir() + (project / "astra.yaml").write_text( + "outputs:\n - id: foo\n recipe:\n command: echo foo\n" + ) + monkeypatch.chdir(project) + record = JobRecord( + job_id="123", + targets=["foo"], + resolved_targets=["foo"], + universe="default", + qos="shared", + resources={}, + sbatch_path="job.sbatch", + log_path="job.log", + submitted_at="2026-07-16T00:00:00+00:00", + last_state="RUNNING", + ) + monkeypatch.setattr(async_jobs, "refresh_job_records", lambda _: [record]) + result = runner.invoke(main, ["status", "--json"]) + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + job = payload["universes"][0]["outputs"][0]["job"] + assert job == { + "job_id": "123", + "state": "running", + "slurm_state": "RUNNING", + "qos": "shared", + "log_path": "job.log", + } diff --git a/tests/test_resources.py b/tests/test_resources.py new file mode 100644 index 00000000..147cc0cd --- /dev/null +++ b/tests/test_resources.py @@ -0,0 +1,61 @@ +"""Tests for canonical ASTRA recipe-resource normalization.""" +from __future__ import annotations + +import pytest + +from lightcone.engine.resources import ( + ResourceValueError, + parse_memory_mb, + parse_recipe_resources, + parse_time_seconds, +) + + +@pytest.mark.parametrize( + ("value", "expected"), + [("512MB", 512), ("8GB", 8000), ("1.5GiB", 1611)], +) +def test_parse_memory_mb(value: str, expected: int) -> None: + assert parse_memory_mb(value) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [("30m", 1800), ("1h30m", 5400), ("02:15:00", 8100), ("1-02:00:00", 93600)], +) +def test_parse_time_seconds(value: str, expected: int) -> None: + assert parse_time_seconds(value) == expected + + +def test_parse_recipe_resources_maps_to_snakemake_names() -> None: + normalized = parse_recipe_resources( + { + "command": "python fit.py", + "resources": { + "cpus": 8, + "memory": "16GB", + "gpus": 2, + "time_limit": "1h1s", + }, + } + ) + assert normalized.snakemake() == { + "cpus_per_task": 8, + "mem_mb": 16000, + "gpus_per_task": 2, + "runtime": 61, + } + + +def test_parse_recipe_resources_requires_time_for_async() -> None: + with pytest.raises(ResourceValueError, match="missing resources.time_limit"): + parse_recipe_resources( + {"command": "echo"}, require_time_limit=True, label="output 'fit'" + ) + + +@pytest.mark.parametrize("value", ["eight GB", "8", "", "-1GB"]) +def test_parse_memory_rejects_ambiguous_values(value: str) -> None: + with pytest.raises(ResourceValueError): + parse_memory_mb(value) + diff --git a/tests/test_site_registry.py b/tests/test_site_registry.py index 61396e1a..cfe0d6c4 100644 --- a/tests/test_site_registry.py +++ b/tests/test_site_registry.py @@ -18,6 +18,8 @@ def fake_hostname(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], None]: """Return a setter that pins ``socket.gethostname`` for the test.""" def _set(name: str) -> None: + monkeypatch.delenv("LIGHTCONE_SITE", raising=False) + monkeypatch.delenv("NERSC_HOST", raising=False) monkeypatch.setattr(site_registry.socket, "gethostname", lambda: name) return _set @@ -68,6 +70,14 @@ def test_display_name_for_unknown_site(self) -> None: class TestDetectCurrentSite: + def test_nersc_host_detects_site_on_compute_node( + self, fake_hostname: Callable[[str], None], monkeypatch: pytest.MonkeyPatch + ) -> None: + fake_hostname("nid001234") + monkeypatch.setenv("NERSC_HOST", "perlmutter") + site = detect_current_site() + assert site.key == "perlmutter" + def test_known_host_returns_populated_site( self, fake_hostname: Callable[[str], None] ) -> None: diff --git a/tests/test_snakefile.py b/tests/test_snakefile.py index 62b4a065..59067cb8 100644 --- a/tests/test_snakefile.py +++ b/tests/test_snakefile.py @@ -42,6 +42,35 @@ def test_generate_simple_spec(tmp_path: Path) -> None: assert "results/{universe}/foo" in snake_text +def test_generate_maps_astra_recipe_resources_to_snakemake(tmp_path: Path) -> None: + _spec( + tmp_path, + { + "outputs": [ + { + "id": "fit", + "recipe": { + "command": "python fit.py", + "resources": { + "cpus": 8, + "memory": "16GB", + "gpus": 1, + "time_limit": "2h", + }, + }, + } + ] + }, + ) + snakefile, _ = generate(tmp_path, universes=["baseline"]) + text = snakefile.read_text() + assert " threads: 8" in text + assert " cpus_per_task=8," in text + assert " mem_mb=16000," in text + assert " gpus_per_task=1," in text + assert " runtime=120," in text + + def test_generate_universe_expansion(tmp_path: Path) -> None: _spec(tmp_path, {"outputs": [{"id": "foo", "recipe": {"command": "echo"}}]}) _, cfg_path = generate(tmp_path, universes=["u1", "u2"]) diff --git a/zensical.toml b/zensical.toml index 5206196a..27538930 100644 --- a/zensical.toml +++ b/zensical.toml @@ -28,6 +28,7 @@ nav = [ {"lc run" = "cli/run.md"}, {"lc build" = "cli/build.md"}, {"lc status" = "cli/status.md"}, + {"lc cancel" = "cli/cancel.md"}, {"lc verify" = "cli/verify.md"}, {"lc export" = "cli/export.md"}, ]}, @@ -42,6 +43,8 @@ nav = [ {"engine/tree" = "api/tree.md"}, {"engine/validation" = "api/validation.md"}, {"engine/dask_cluster" = "api/dask_cluster.md"}, + {"engine/resources" = "api/resources.md"}, + {"engine/async_jobs" = "api/async_jobs.md"}, {"snakemake_executor_plugin_dask" = "api/dask_executor.md"}, ]}, {"Skills" = [