Metadata-Version: 2.4
Name: neurarch-skypilot
Version: 0.1.1
Summary: A SkyPilot admin policy that sizes a training job's model before it takes a GPU: refuses a request that cannot hold it, right-sizes one that asks for far more, and enforces a cost cap.
Author: neurarch-ai
License: MIT
Project-URL: Homepage, https://neurarch.com/memory-calibration.html
Project-URL: Source, https://github.com/neurarch-ai/neurarch
Keywords: skypilot,admin-policy,gpu,admission-control,capacity-planning,training,oom
Classifier: Intended Audience :: System Administrators
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: skypilot
Requires-Dist: skypilot>=0.14; extra == "skypilot"
Dynamic: license-file

# neurarch-skypilot

A SkyPilot admin policy that sizes a training job's model when the job is
submitted, before a VM is provisioned. It refuses a request that cannot hold
the job, warns on (or rewrites) a request far larger than the job needs, and
refuses a run whose list-price cost is over a cap.

```bash
pip install neurarch-skypilot                   # on the host that applies the policy (PyPI)
echo "admin_policy: neurarch_skypilot.NeurarchPolicy" >> ~/.sky/config.yaml
sky launch examples/finetune-qwen3-8b.sky.yaml
```

For a team, set the same `admin_policy:` line in the API server's config
instead of each user's `~/.sky/config.yaml`, with the package installed on the
server. The policy needs the `neurarch` CLI on that host to size a Hugging Face
id offline (see [Where the number comes from](#where-the-number-comes-from)).

## What a refused job looks like

The example task asks for one A10G to fully fine-tune Qwen3-8B at batch 8:

```
sky.exceptions.UserRequestRejectedByPolicy: Neurarch: Qwen/Qwen3-8B needs 252.0 GB to train at
batch 8 per device, 2048 tokens (peak 219.1 GB x 1.15 headroom), and A10G:1 holds 24 GB. No batch
fits A10G:1: the weights, gradients and optimizer state alone need 189.9 GB. No single card holds
it; the smallest request that does is H200:2, sharded (FSDP or ZeRO-3) at the same global batch
of 8. Sized offline by the neurarch CLI from the model's structure; nothing was trained. Change
`accelerators:` and resubmit.
```

When some batch does fit, the message names the largest one and the smallest
single card that holds the job as asked:

```
Neurarch: Qwen/Qwen3-0.6B needs 34.4 GB to train at batch 16 per device, 2048 tokens (peak 29.9 GB
x 1.15 headroom), and T4:1 holds 16 GB. The largest batch that fits T4:1 is 1 per device. The
smallest card that holds this job is A100:1. ...
```

## Telling the policy about the job

The job's facts go in the task's `envs`, the same variables the training
script reads, so the policy and the script cannot size two different runs:

```yaml
resources:
  accelerators: A10G:1
envs:
  NEURARCH_MODEL: Qwen/Qwen3-8B   # HF id, or a path to a .neurarch.json from neurarch-trace
  NEURARCH_BATCH: "8"             # per device; default 32
  NEURARCH_SEQ_LEN: "2048"        # tokens per sample; default 2048, or the model's context if shorter
  NEURARCH_LORA_RANK: "16"        # omit for a full fine-tune
  NEURARCH_ATTENTION: sdpa        # sdpa | flash | eager (only eager stores attention scores)
  NEURARCH_HOURS: "6"             # expected wall hours; needed for the cost cap
```

Resource `labels` with the same names in lower case (`neurarch-model`,
`neurarch-batch`, ...) are read when the env is absent. A task that names no
model is not judged at all, and neither is a CPU-only task or SkyPilot's own
controller launch. A value the policy cannot read (`NEURARCH_BATCH: eight`) is
refused rather than defaulted.

## The verdicts

The request is `accelerators` x `num_nodes`. A multi-card request is sized as
one job sharded across all its cards: the peak at the global batch (per-device
batch x cards) against the summed memory, the same rule `npx neurarch rightsize`
uses. Need is the peak x 1.15 headroom.

| verdict | when | what happens |
|---|---|---|
| does not fit | need > requested memory | refused, naming the need, the card, the largest batch that fits it and the smallest request that holds the job |
| borderline | need within 15% of the requested memory | admitted, with a warning |
| over-requested | a much smaller request holds it (more than one size up from the smallest that does; single card first) | warning; with `mutate`, the accelerator request is rewritten to the cheapest request that holds it |
| over the cost cap | cards x hourly x `NEURARCH_HOURS` > `max_cost_usd` | refused, naming the cost and the rate |
| cannot size | no sizing source answered, or an accelerator with no capacity on file | admitted with a warning (fail open), or refused with `fail_closed` |

The over-requested rule is `classify` from `src/utils/rightsizeCore.ts`, ported;
a shared fixture pins the port to the original and the card table to its price
list. A multi-node request is never rewritten. With several accelerator
candidates (`any_of`, or a list), a candidate that cannot hold the job refuses
the request, or with `mutate` is dropped while at least one candidate remains.

Warnings go to SkyPilot's own logger (`sky.neurarch_skypilot`), so they print
in the terminal for a client-side policy and in the request log for a
server-side one.

## Settings

SkyPilot instantiates the policy with no arguments, so settings are class
attributes overridden by environment variables on the host that applies it.
To fix them in code, subclass (see [`examples/cluster_policy.py`](examples/cluster_policy.py)):

```python
class ClusterPolicy(NeurarchPolicy):
    max_cost_usd = 500
    mutate = True
```

| attribute | env var | default |
|---|---|---|
| `mutate` | `NEURARCH_POLICY_MUTATE` | `False` |
| `max_cost_usd` | `NEURARCH_POLICY_MAX_COST_USD` | none |
| `fail_closed` | `NEURARCH_POLICY_FAIL_CLOSED` | `False` |
| `sizing` | `NEURARCH_POLICY_SIZING` | `auto` (CLI, then hosted) |
| `cli` | `NEURARCH_POLICY_CLI` | `neurarch` on PATH, else `npx --no neurarch` |
| `plan_url` | `NEURARCH_POLICY_PLAN_URL` | `https://www.neurarch.com/api/v1/plan` |
| `prices` | `NEURARCH_POLICY_PRICES` (path to a JSON object, card to USD per GPU-hour) | list prices |
| `timeout_s` | `NEURARCH_POLICY_TIMEOUT_S` | 60 |
| | `NEURARCH_API_KEY` | none (the hosted plan allows 60 anonymous calls an hour) |

Default prices are on-demand list prices per GPU-hour from Modal, the same table
the Neurarch app and `rightsize` use. A team on reserved or negotiated rates sets
`prices`; the cost in a refusal says which rate it used.

## The repository's `.neurarch.yml`

When the task's `workdir` (or the API server's working directory, or
`NEURARCH_POLICY_FILE`) holds a `.neurarch.yml`, the verdicts above are handed
to `neurarch policy apply`, the engine the CLI uses, as findings
(`admission/request-does-not-fit`, `admission/cost-cap`,
`admission/over-requested`, `admission/borderline-fit`) with the facts a budget
line needs. A waiver scoped to the model, path or cluster (`NEURARCH_CLUSTER`)
turns a refusal into a warning that names its owner and expiry; a budget line
(`cards_allowed`, `max_cost_per_run_usd`, `max_params`) refuses a request that
would otherwise go through. The cluster's own `max_cost_usd` is never loosened
by a repo file. Waiver rules (30 days when no expiry, never forever) are the
CLI's, documented in its README. No CLI, a timeout or an unreadable file: the
verdicts stand and one line says the policy was not applied (fail open).

```
Neurarch: Qwen/Qwen3-8B needs 252.0 GB to train at batch 8 per device, ... Allowed: waived in .neurarch.yml by @ops until 2026-11-30: OOM test on purpose.
Neurarch: .neurarch.yml: H100 is not an allowed card (a100). Request one of a100.
```

## Where the number comes from

The training peak is weights, gradients and AdamW state plus activations,
computed from the model's structure (a Hugging Face `config.json`, or a traced
graph), calibrated on measured runs and published next to `accelerate
estimate-memory` at <https://neurarch.com/memory-calibration.html>. It is
arithmetic: nothing is trained, downloaded or run on a GPU.

1. **Offline, the `neurarch` CLI** (`neurarch rightsize <csv> --json`). Needs
   Node 18+ on the host. Reads `config.json` from huggingface.co (set
   `HF_TOKEN` for a gated repo) or a local graph; nothing is uploaded. Handles
   LoRA. `npx --no` is used so admission never fetches a package from the
   registry; point `NEURARCH_POLICY_CLI` at an installed copy
   (`node /path/to/neurarch/bin/neurarch.mjs`).
2. **Hosted, `POST /api/v1/plan`** with `seq_len` and `attention`. Takes a
   graph, so it serves a `.neurarch.json` path the policy host can read, not a
   bare Hugging Face id, and it does not size LoRA. The URL is `www.`, because
   the apex domain drops the `Authorization` header on its redirect.

Either source is reduced to `fixed + per_sample x batch` (the peak is linear in
batch), which is how the policy names the largest batch a card holds without a
second call. Answers are cached per job shape for the life of the process,
because SkyPilot applies a policy more than once per launch.

## Tests

```bash
cd packages/neurarch-skypilot && python3 -m unittest discover -s tests
```

Standard library only; SkyPilot is not needed (one test checks the exception
type and is skipped without it). `scripts/skypilot-parity.test.ts` in the main
repo runs the TypeScript side of the shared fixture under `npm test`.

## Not done yet

- **A server-side policy cannot read the user's local files**, so a
  `.neurarch.json` path only works client-side or on a path the server mounts.
  A Hugging Face id works in both places.
- Kueue (an AdmissionCheck controller) is not built. SLURM is a sibling
  package, [`neurarch-slurm`](../neurarch-slurm), a `cli_filter` plugin that
  calls this package's `judge` with SLURM wording.
- The policy does not record its verdicts anywhere. The refused-before-dispatch
  count at quoted prices is the metric, and today it lives only in the log.
- Fractional accelerators (`A10G:0.5`), TPUs and cards outside the table are
  not judged.

MIT licensed.
