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.

pip install ./packages/neurarch-skypilot        # on the host that applies the policy
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).

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:

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.

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):

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

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.

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

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) and SLURM (job_submit.lua) are not built; this is the SkyPilot one only.
• 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.
• Not published to PyPI.

MIT licensed.