Skip to content

Arrays#

An array runs one worker template many times, each copy with its own index. Use it for embarrassingly parallel work: a hyper-parameter sweep, one shard of a dataset per worker, a batch of evaluations. The copies are independent — none waits for another, and one failing does not stop the rest — and you can cap how many run at once so the array does not take every GPU of the workspace.

It is Astraeus' form of a Slurm job array.

How an array runs#

  • array.size workers are created, named <run>-0 … <run>-<size−1>.
  • Each worker has the template's whole ask: with gpu_requests.count: 1, every copy gets one GPU. Nothing is split.
  • Worker i gets ASTRAEUS_ARRAY_TASK_ID=i (from 0) and ASTRAEUS_ARRAY_SIZE=<size>. RANK and ASTRAEUS_REPLICA_INDEX hold the same index, and WORLD_SIZE the size.
  • Copies are placed independently, as room appears, and several may share a machine.
  • array.max_parallel caps how many hold machines at once. The next waiting copy says Array runs at most 10 workers at once.
  • Each copy restarts on its own under its restart_policy; the run is Completed when every copy completed, and Failed if any failed for good (3 of 100 workers failed; the first: Worker sweep-17 failed on gpu-02: Exited with code 1).

Run a parameter sweep#

100 learning-rate and batch-size combinations, one GPU each, at most 8 at a time, each with a one-hour limit.

sweep.py
import itertools, os

lrs = [1e-4, 3e-4, 1e-3, 3e-3, 1e-2]
batches = [16, 32, 64, 128]
seeds = range(5)
grid = list(itertools.product(lrs, batches, seeds))      # 100 combinations

i = int(os.environ["ASTRAEUS_ARRAY_TASK_ID"])
lr, batch, seed = grid[i]
train(lr=lr, batch=batch, seed=seed, out=f"/results/{i:03d}.json")
  1. New run: Name sweep, Image nvcr.io/nvidia/pytorch:24.08-py3, Command python, Arguments /data/code/sweep.py.
  2. GPUs 1, CPU cores 8, Memory 32Gi (with an array, CPU and memory are per worker).
  3. Drives: datasets:/data,results:/results.
  4. More options: Array size 100, At most in parallel 8, Time limit 1h.
  5. Start run. Start is forced to independent for an array.

astra astraeus run has no array flag; use Slurm's:

sweep.sbatch
#!/bin/bash
#SBATCH --array=0-99%8
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=1:00:00
#SBATCH --container-image=registry.example.com/ml/sweep:1
python /opt/sweep/sweep.py

$ astra slurm sbatch sweep.sbatch
Submitted batch job sweep

The script runs in the image, so the code must be in it (/opt/sweep here). Inside, read ASTRAEUS_ARRAY_TASK_ID: SLURM_ARRAY_TASK_ID is not set. astra cannot mount drives: for the results drive, use the API.

sweep.json
{
  "metadata": { "name": "sweep" },
  "spec": {
    "array": { "size": 100, "max_parallel": 8 },
    "task_template": {
      "image": "nvcr.io/nvidia/pytorch:24.08-py3",
      "command": "python",
      "args": ["/data/code/sweep.py"],
      "time_limit_seconds": 3600,
      "datavolume_refs": [
        { "name": "datasets", "mount_path": "/data", "mode": "ReadOnly" },
        { "name": "results", "mount_path": "/results", "mode": "ReadWrite" }
      ],
      "requested_resources": {
        "cpu_cores": 8,
        "memory_bytes": 34359738368,
        "gpu_requests": { "count": 1 }
      }
    }
  }
}

$ curl -sS -X POST "$API/jobs" -H "Authorization: Bearer $ASTRA_TOKEN" \
    -H 'Content-Type: application/json' --data-binary @sweep.json | jq .spec.array
{
  "size": 100,
  "max_parallel": 8
}

Give each copy a time limit: a run whose workers all state one can backfill, which on a busy cluster often starts copies sooner.

Follow an array#

The run's Workers tab lists every copy with its state, machine, GPUs and exit code; the readout says how many run (Workers 8 running of 100). Select Log on a row for that copy's output.

$ astra astraeus workers sweep
WORKER                        RANK   STATE       MACHINE               REASON
sweep-0                       0      Completed   gpu-01                Exited with code 0
sweep-1                       1      Running     gpu-02                Container running
sweep-9                       9      Pending     -                     Array runs at most 8 workers at once
$ astra astraeus logs sweep-17

$ curl -sS "$API/jobs/sweep" -H "Authorization: Bearer $ASTRA_TOKEN" | jq .task_counts
{
  "Completed": 41,
  "Pending": 51,
  "Running": 8
}

Limits#

Limit Value
array.size 1 to 10 000
array.max_parallel 0 (no cap) to size
Start policy Independent only: start: Gang or policy: Gang with an array is refused (An array's workers are independent: use policy Independent).
Worker groups Not with task_groups.
Lifetime Not with lifetime: Service.
Model engine Not with model.engine.
Indices Always 0 to size − 1, in steps of 1.

Troubleshooting#

Symptom Cause Fix
Every copy computes the same thing The script reads SLURM_ARRAY_TASK_ID, which is not set. Read ASTRAEUS_ARRAY_TASK_ID.
Only a few copies run though GPUs are free max_parallel, or the workspace's quota. Check the waiting copy's reason.
The run is Failed though most copies completed One copy failed for good. Its row on Workers has the exit code and log; run only the missing indices in a new array, mapping its index onto them.