fix(slurm): make SLURM mode work with VRAM-based GPU sizing

[nilsmechtel]

Problem

SLURM mode could not run GPU applications that use VRAM-based sizing.

Ray Serve replicas that declare gpu_memory_mb are packed onto GPUs by real memory: the builder requests a VRAM_MB custom resource and a tiny num_gpus value that only binds a device. Four independent defects broke that path end-to-end on a SLURM cluster.

1. SLURM workers never advertised VRAM_MB. The worker job started Ray with a hardcoded --resources string carrying only the job id, so ray.cluster_resources() on a SLURM cluster contained no VRAM_MB at all. Every GPU app fell back to a whole-GPU reservation, and any app explicitly requesting VRAM_MB stayed pending forever.

2. The autoscaler ignored VRAM-only demand. _check_scale_up read pending GPU demand from required_resources["GPU"]. A VRAM-packed replica reports GPU: 0.01, which is < 1, so the autoscaler submitted a CPU-only SLURM job — or nothing — and the replica never scheduled.

3. The builder locked in whole-GPU reservations on a scaled-to-zero cluster. GPU sizing is resolved once at deploy time by asking the proxy actor what GPUs the cluster has. An autoscaling SLURM cluster is normally at zero workers when an app is deployed, so the answer is "no GPU visible", and the builder chose a whole-GPU reservation. That reservation is never revisited after a worker joins, so VRAM packing could not take effect on the mode that needs it most.

4. The admission gate counted capacity in flight as consumed headroom. _check_resources first looks for a registered Ray node with room, then falls back to asking whether another SLURM worker could be spawned — comparing the total number of worker jobs against max_workers. A job that has been submitted but has not yet registered as a Ray node satisfies neither test: it is invisible to the node scan and it consumes headroom in the job count. That state is the norm rather than an edge case, because a deployment on a cold cluster triggers its own scale-up: the app-introspection Ray task creates pending demand before the resource check runs. With max_workers=1 the check therefore saw one job and zero nodes and rejected the deployment with "Insufficient resources" every time. Larger max_workers narrows the window but does not close it.

Separately, SlurmWorkers.close_all() returned early whenever no Ray node was registered. SLURM jobs that were still queued, or had started but not yet joined Ray, therefore survived the worker's shutdown and burned allocation until they hit their time limit.

Solution

Advertise real per-node VRAM. The generated sbatch script queries nvidia-smi on the allocated node and appends "VRAM_MB": <total> to the Ray resource string. Detection happens inside the job because only the allocated node knows its own GPU size. It is gated to single-GPU workers, matching the existing invariant for the head node: on a multi-GPU node the node-level VRAM_MB is a sum Ray can satisfy across devices, which would let a replica claim more memory than any single GPU has. If detection fails the job logs a notice and starts without VRAM_MB, falling back to fractional sizing.

Treat a VRAM request as GPU demand. _check_scale_up now raises num_gpus to 1 when pending resources include VRAM_MB with num_gpus < 1. The CPU divisor is guarded against a zero pending-CPU value that could raise ZeroDivisionError.

Request VRAM up front on elastic clusters. AppBuilder takes a new elastic_gpu_cluster flag, set by AppsManager for slurm mode. When the cluster is elastic and no GPU is visible yet, the replica requests VRAM_MB rather than a whole GPU, so it stays pending until the autoscaler provides a node that advertises it. Clusters with a known GPU are unaffected, and gpu_memory_mb=-1 still books a whole device.

Admit deployments backed by capacity in flight. Worker jobs with no corresponding Ray node are no longer counted against max_workers; they admit the deployment when the worker defaults fit the request, and the replica then waits for the node in the normal pending state.

Keep fractional GPU demand visible. _sum_resources truncated num_gpus to an int, so the epsilon reservation used to bind a device summed to zero and hid the app's GPU demand entirely. num_gpus now stays a float.

Always sweep SLURM on shutdown. close_all() drains registered nodes, awaits those graceful stops so shutdown does not race process exit, and then cancels any remaining jobs regardless of whether nodes were found.

Behavioural changes

• SLURM worker nodes now report a VRAM_MB resource. Existing deployments that reserve whole GPUs are unaffected; apps declaring gpu_memory_mb will start packing by memory.
• A GPU app deployed against a scaled-to-zero SLURM cluster now triggers a GPU worker allocation instead of hanging, allocating a CPU-only worker, or being rejected as unschedulable.
• Stopping a SLURM-mode worker no longer leaves queued or newly started jobs behind.

No configuration or manifest changes are required.

Test plan

Validated on an NSC Berzelius A100 cluster (SLURM, apptainer, min_workers=0, max_workers=1, -C thin) against the published dev image ghcr.io/aicell-lab/bioengine-worker:0.16.5-dev1, built from this branch. The worker booted from that image with no source bind-mount, so both the head and — via runtime_env.py_modules — every replica ran the image's own bioengine. The test application requests gpu_memory_mb and pulls torch==2.5.1 through runtime_env.pip, i.e. the same transport a real BioEngine app uses. Every GPU app shipped in this repo requests a whole device, so a purpose-built app was needed to exercise the VRAM packing path at all.

Cold start, from a cluster with zero worker nodes, driven entirely through Hypha (upload_app -> deploy_app -> RPC call):

• The deployment was admitted rather than rejected, which is the path defect 4 broke: at admission time the cluster had no registered node and one worker job in flight.
• The autoscaler logged "Scaling up with 1 GPU(s) and 16 CPU(s) due to pending resources" and submitted a GPU job.
• The job's stdout shows the advertisement: Ray resources: {"slurm_job_id:17464167": 1, "VRAM_MB": 40960}.
• The replica packed by memory rather than booking the device: with the reservation held, Ray Serve reported VRAM_MB: 32960.0 still available on the node, i.e. 40960 minus the requested 8000.
• The application reached RUNNING after 181 s, and a method call over Hypha RPC ran a real fp16 8192³ matmul on the allocated device: {'gpu': 'NVIDIA A100-SXM4-40GB', 'gpu_total_mb': 40440, 'visible_devices': '0', 'matmul_tflops': 250.2, 'peak_alloc_mb': 520} — confirming the epsilon num_gpus reservation binds a usable CUDA device.
• stop_app and delete_app left no applications behind. Shutting the worker down cancelled the still-running job (Cancelling 1 job(s): ['17464167']), completed in 73.3 s, and left squeue empty.

Two further runs over the same branch source covered the remaining paths: an in-place redeploy of the same application_id at gpu_memory_mb=-1 reached RUNNING on the whole-GPU path at 236.4 TFLOPS, so both reservation shapes work on one cluster; and a pure-Ray run confirmed scale-down, where deleting the application let the node go idle and the autoscaler cancelled the job after 102.2 s against a 90 s threshold.

Generated sbatch scripts were also inspected for num_gpus of 0, 1 and 2 to confirm the VRAM block is emitted only for single-GPU workers.

Unit tests need no cluster. tests/test_gpu_sizing.py (9 cases) covers the sizing helpers directly: VRAM-advertised packing, the fraction fallback, the elastic-cluster path, the static-cluster whole-GPU fallback, the gpu_memory_mb=-1 sentinel, disable_gpu, idempotent re-derivation, the oversized-request error, and fractional num_gpus summation. tests/test_slurm_resource_admission.py (4 cases) drives _check_resources against a stub cluster: empty cluster, in-flight job counted as capacity, registered node with room, and a full cluster correctly rejecting an oversized request. Reverting the manager.py change makes the in-flight case fail and leaves the other three passing.

Files changed

File	Change
bioengine/cluster/slurm_workers.py	Detect and advertise per-node VRAM_MB in the sbatch script; treat VRAM_MB demand as requiring a GPU in _check_scale_up; guard the CPU divisor; always sweep SLURM jobs in close_all() and await graceful stops
bioengine/apps/builder.py	New elastic_gpu_cluster flag; merge it into the cluster GPU sizing result; request VRAM_MB when an elastic cluster has no GPU visible yet; keep num_gpus a float in _sum_resources
bioengine/apps/manager.py	Pass elastic_gpu_cluster=self.ray_cluster.mode == "slurm" to AppBuilder; treat unregistered SLURM worker jobs as capacity in flight in _check_resources
tests/test_gpu_sizing.py	New unit tests for the GPU sizing helpers
tests/test_slurm_resource_admission.py	New unit tests for the SLURM headroom check
pyproject.toml, bioengine/_version.py	Version bump to 0.16.5