feat: count the GPUs, the VRAM and the other resources of the machine

[stephenc]

Implements section 2 of design-decisions.md. Closes part of #10.

What it does

The scheduler counted two things, the cores and the memory, and it held named locks. A machine with four GPUs had no way to say --gpu 1. This adds one new concept to the scheduler — a pool — and gives the command line the two names that an agent must get right on the first try.

qex submit --cpu 4 --mem 16GB --gpu 1 --vram 20GB -- uv run train.py
qex submit --cpu 8 --mem 32GB --gpu 2 -- uv run train.py    # 2 whole devices
qex submit --claim net=1 -- ./download.sh
qex run --lock target -- cargo test                         # unchanged

[[pool]]
name    = "gpu"
size    = "vram"
devices = ["24GB", "24GB", "24GB", "24GB"]
env     = "CUDA_VISIBLE_DEVICES"

[[pool]]
name  = "net"
count = 4

The model, in one paragraph

A pool is a name, a total, and — when qex must say which one — a list of devices with a capacity for each. --lock NAME is a pool of one unit, and a name that the configuration does not declare is a pool of one unit, so a lock still needs no configuration. --gpu and --vram are fixed names on the command line for a claim on the pool gpu; the scheduler holds no special case for that name and sees one claim map and one arithmetic. A plain pool admits a job when the free units cover the claim. An indexed pool admits a job when enough devices each have the claimed quantity free — the capacity of the devices is never added together, because four devices of 24GB are not 96GB for one job. qex reads no driver: the devices come from the configuration, so a machine with no CUDA and no driver library schedules GPU claims correctly. When the job starts, the coordinator chooses the devices with the most free capacity first (lowest index for a tie), inside the same lock hold that moves the job to starting, and writes the choice to status.json as assigned — an assignment is a result and not a request, so it never goes in spec.json. The supervisor then writes CUDA_VISIBLE_DEVICES, QEX_GPU_DEVICES and QEX_GPU_VRAM for the job, and those values replace the captured environment. A coordinator that starts again reads assigned back and rebuilds the occupancy of every pool.

The three traps, and how each is covered

• Peer needs #[serde(default)]. Both new fields have it, with a comment that says why, and a_peer_file_of_an_earlier_version_still_parses_and_still_counts parses a real 0.7.1 record and asserts that its cores and memory still count.
• locks: Vec<String> stays on the wire. No claim is ever written into locks, and no lock is ever written into claims. effective_claims does the conversion inside the coordinator, after the capability test. required_by gives locks for a lock and pools for a claim, and a_lock_does_not_need_the_pools_capability proves that a coordinator with locks and no pools still accepts --lock.
• New Info fields are Option. pools: Option<Vec<PoolReport>>. qex info prints unknown and names the coordinator version when the value is None, and never 0 or none.
• Stage mirrors JobFile. gpu, vram and claims went into spec::Resources, which Stage already reuses and stage_spec already copies whole. A pipeline stage therefore cannot drop a claim, and there is no third place to forget.

Decisions where the document left something open

1. An undeclared gpu pool is an error, and an undeclared net is a lock. The document gives both a "there is no pool gpu" refusal and an "an undeclared name is a lock" rule. They conflict for --gpu 1. The rule I used: a name that the configuration does not declare is a lock of one unit, except the built-in alias name gpu. --gpu promises a device index and an environment variable; a silent lock would give the job neither, and the job would then use a card that qex is not accounting for. One small special case, in pool_check, with the reasoning in a comment.
2. A generic way to claim a size. The document gives --vram for gpu and --claim NAME=N for everything else, which leaves no way to claim a quantity on a device of another indexed pool. I added --claim NAME=N:SIZE and the job-file table form tpu = { count = 2, size = "8GB" }. That makes PoolClaim.size reachable generically, so --vram really is only a friendly name and not the only path to the feature.
3. Assignment.size stays None for a whole-device claim. Writing the device capacity there looks equivalent and is not: the devices of a pool can differ in size, and any single number would leave part of the largest device free for a second job. Held::add resolves None to the capacity of each device it was given.
4. The record sorts the device indices; the choice does not. The selection uses free capacity, but the list written to status.json and to CUDA_VISIBLE_DEVICES is sorted, so two equal assignments give one text and 2,3 reads the way a person expects.
5. The CUDA_VISIBLE_DEVICES conflict is refused for an explicit value only. The check reads the job file [env] and --env, and not the captured shell environment. Refusing a captured value would stop anyone who exports the variable in a login file from submitting any GPU job; the supervisor replaces the value for that job, which is the correct result.
6. A [[pool]] with neither count nor devices is refused. The document does not say. Such an entry declares no size, and inventing one would be a size that no person chose.
7. A claim reserves capacity only when it is not all-or-nothing. Per §1.2/§4.1: a claim is all-or-nothing when the pool has one unit and the claim takes the whole of it. Those go through the lock path and continue, so a lock still does not park the queue. A counted claim goes through admit. An indexed pool of one device with a partial --vram claim is divisible, so it is counted and not a lock.

I did not find anything in section 2 that I think is wrong.

What I measured

• The full e2e suite takes 38s at --test-threads=2 on this machine, unchanged from main.
• admit gained one pool test. It reads no /proc and no driver; it walks the claim map of the job, which is empty for a job with no claim, so a queue with no pools does the same work as before.
• peers::claims is now read once per admit call rather than once per admit call — unchanged — and the result is reused for the cores, the memory and the pools instead of being read twice.
• No changing number went into any blocked_reason, so no reason text rewrites status.json on every tick.

What I tested

Unit (202 pass, up from 172). sched: a machine with no GPU admits a GPU claim from the configuration; VRAM is never summed; a claim above the pool total can never start; an undeclared name is a lock and two units of it are not; a GPU claim with no gpu pool names the configuration; VRAM with no device is refused; a size on a plain pool is refused; two jobs get different devices and a third waits; a whole-device claim leaves no room for a second job; a device holds three 8GB jobs on 24GB and no fourth; the most free capacity comes first; a device another user holds is not given again; a lock becomes a pool of one unit; a lock in an old record still counts; a counted pool admits until full. config: the documented example parses (now with both pools); count with devices is refused; cpu/mem cannot be a pool name; two entries with one name are refused; an entry with no size is refused. spec: the aliases become one claim; no --vram takes the whole device; [defaults] vram applies; a job file gives the same claims; an explicit CUDA_VISIBLE_DEVICES is refused and a captured one is not; --claim parsing. capabilities: a lock does not need pools; a claim is refused by a coordinator without pools. peers: an 0.7.1 record still parses and still counts; a record carries the pools and the devices.

End-to-end (81 pass, up from 71), all on a machine with no GPU.

• a_machine_with_no_gpu_schedules_a_gpu_claim_from_the_configuration — runs a job that prints CUDA_VISIBLE_DEVICES, QEX_GPU_DEVICES and QEX_GPU_VRAM, and checks the record as well.
• two_gpu_jobs_get_different_devices_and_a_third_waits
• a_vram_claim_above_the_largest_device_is_refused_at_the_submission — the message says "never start", and no record is left.
• a_gpu_claim_above_the_pool_total_is_refused_whatever_the_oversized_policy
• a_gpu_assignment_survives_a_coordinator_that_stops_and_starts — takes the pid from qex info --no-start --json, never a process-list search.
• a_lock_needs_no_configuration_and_still_excludes — and a job with no lock still passes it.
• a_counted_pool_lets_n_jobs_operate_and_makes_the_next_one_wait
• an_undeclared_pool_is_a_lock_and_two_units_of_it_are_refused
• a_job_that_sets_the_device_variable_itself_is_refused
• info_reports_the_pools_and_their_devices

Also updated

docs/reference.md (a new pools section, and the job-file example at line 180 no longer sets CUDA_VISIBLE_DEVICES), src/help.rs (the submit options, the job-file fields, the pipeline stage fields, the config topic and the resources topic — which now contains the sentence "qex does not add the VRAM of the devices together"), src/schema.rs (gpu, vram and claims on the job and pipeline schemas; locks, claims and assigned on the status schema), and Cargo.toml to 0.8.0.

A note

Another agent is changing the head-of-line rule in src/sched.rs on a separate branch. This branch works from main and does not implement that design. A merge conflict in choose is expected.

🤖 Generated with Claude Code

https://claude.ai/code/session_01KNvagiMEU3myn8EXGaGEM9