Fix weight update race condition between trainer cleanup and orchestrator

[cdreetz]

current bug in prime-rl when having excessive event loop lag, it will cause the cleanup of old broadcast/step_n that wont happen immediately, then inference will check step_n and see its there, then the delete takes place and the following load fails because step_n is missing by the time it tries to load

The trainer's maybe_clean deletes broadcast checkpoint directories on a timer, but the orchestrator may not have loaded them yet (its event loop can be blocked for 70+ seconds during rollout generation). This caused FileNotFoundError on the inference server and crashed the run.

Fix:

• Orchestrator writes a LOADED_STEP marker after each successful weight update so the trainer knows which checkpoints have been consumed.
• Trainer only cleans broadcast directories the orchestrator has moved past (candidate_step < loaded_step).
• Orchestrator catches update_weights failures and retries on the next poll as defense-in-depth.
• Fix DefaultModelLoader import for vllm >=0.16 (module renamed from loader to default_loader).

---

Note

Medium Risk
Touches orchestrator↔trainer synchronization around checkpoint lifecycle; mistakes could still lead to missing checkpoints or disk bloat. Also downgrades core GPU/ML dependencies (Torch/Triton/NVIDIA libs), which can impact runtime stability and compatibility.

Overview
Prevents a race where the trainer deletes broadcast step_* directories before the orchestrator/inference pool has finished loading them.

The orchestrator now atomically writes a LOADED_STEP marker after a successful update_weights, and the trainer’s maybe_clean only deletes step_* directories up to min(candidate_step, loaded_step-1) (sweeping all eligible historical steps while respecting interval_to_keep).

Separately, uv.lock is updated to pin older versions of torch/torchaudio/torchvision and Linux triton, along with a nvidia-nvshmem-cu12 downgrade.

^{Written by Cursor Bugbot for commit 221cf70. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

[mikasenghaas]

i think it should be fine? im a little lost with our ckpt logic (ie which ckpt goes where) since multi tenant but overall looks sensible

[mikasenghaas]

😭

[mikasenghaas]

this will break if we ever rename get_step_path but ig we wont haha so ig we're good?