What does this PR do?

Adds two optional driver-side budgets, trainer.step_timeout_s and trainer.weight_sync_timeout_s. A hung training step becomes a typed, picklable StepTimeoutError / WeightSyncTimeoutError that names the step, stage and operation.

TLDR: the deadline enters asyncio.timeout, so every await in the step (generation, inference-engine pause/resume and sleep/wake_up, the fully-async buffer wait, to_thread saves) is cancelled on expiry. The ray.gets behind a step's WorkerDispatch calls and PPORayActorGroup offload/backload go through deadline.ray_get. Both trainer keys default to None; control-plane HTTP requests have a separate 600 s default timeout.

Usage

trainer:
  step_timeout_s: 1800        # one step: generation through weight sync
  weight_sync_timeout_s: 900  # each weight sync, incl. the initial one; must be <= step_timeout_s

StepTimeoutError: step 12 exceeded its 1800s budget after 1800.0s (stage='train_critic_and_policy', operation='forward_backward')

• Checkpoint and HF saves made outside a step (the final saves, and the fully-async trainer's saves) each get their own step_timeout_s budget. Saves inside a sync-trainer step share the step's budget.
• Eval is not covered.
• Only the driver's wait is interrupted: a hung worker keeps running until teardown.
• SKYRL_INFERENCE_CONTROL_PLANE_TIMEOUT_S bounds each inference-server control request (default 600 s; <= 0 disables it). A request timeout raises picklable InferenceServerTimeoutError. It does not apply to data-plane generation requests.

How it works

skyrl/train/utils/deadline.py:

• step_deadline(global_step, budget_s, error_cls) sets a _Deadline in a ContextVar and wraps its body in asyncio.timeout(budget_s). It is a no-op for None. A nested deadline can only shorten its parent: if the parent expires first, the parent, and its error class, stay in effect.
• On expiry, the cancellation unwinds through the Timers and deadline.operation(...) scopes. Each records itself on the deadline, innermost first, and the deadline raises error_cls(step, stage, operation or "await", …). A TimeoutError from the body that isn't this deadline's passes through unchanged.
• ray_get(refs, operation) calls ray.get(refs, timeout=remaining()) and turns GetTimeoutError into the deadline's error.
• stage(name) names the current stage, and operation(name) the call awaited within it: RemoteInferenceClient._call_all_servers names its endpoint and the sync trainer names generate, so a timeout during a weight sync reads e.g. stage='sync_weights', operation='/pause'. Timer.__enter__/__aenter__ enters it with the timer's message, so every existing Timer("…") / timed_phase("…") names its stage with no new bookkeeping. asyncio.to_thread copies the context, so ray_get in worker threads sees both the deadline and the stage.
• StepTimeoutError(RuntimeError) passes all five fields to super().__init__, so default pickling round-trips them. It is deliberately not a TimeoutError, so except TimeoutError retry paths (e.g. RemoteInferenceClient's) can't swallow it.
• WeightSyncTimeoutError subclasses it.

cleanup_preserving_primary's failure-path cleanup (the resume_generation / wake_up in save_weights_for_sampler) runs under asyncio.timeout(deadline.remaining()). It is skipped, with a warning, once the deadline has expired, and a cleanup that hangs past it adds a note to the primary exception. Without a deadline it is unbounded, as before. asyncio.timeout cancels only once, so without this a wedged engine would hang the unwind after the deadline fired.

Trainer scopes: the step block (async with self._step_deadline(), Timer("step", …)), each weight sync, including the initial one (_weight_sync_deadline()), and each save outside a step. The synchronous, fully-async and example async trainers all enter these scopes.

RemoteInferenceClient uses an aiohttp.ClientTimeout for each control-plane request, including LoRA load/unload. This bounds calls made outside a training step too.

validate_cfg:

• rejects <= 0 values and weight_sync_timeout_s > step_timeout_s;
• warns when a budget is below SKYRL_WORKER_NCCL_TIMEOUT_IN_S, because the driver would then give up before the NCCL watchdog reports which collective is stuck.

The keys' docstrings render in the config API reference. The troubleshooting page gets a "Step timeouts" section.

Alternatives considered:

• Per-call timeout= arguments: they would have to be threaded through every dispatch signature, and they can't express "whatever is left of this step".
• A watchdog thread that kills the process: that is the os._exit pattern the previous PR removes.

Test plan

tests/train/test_deadline.py:

• Nested deadlines keep the shorter budget and its error class; ray_get uses the remaining budget and the enclosing Timer stage.
• A real Ray actor blocked in backload_to_gpu times out through the actor-group call site; both error classes survive pickling.
• An awaited operation is cancelled on expiry with its stage and operation, including the "await" fallback; an inner weight-sync deadline raises WeightSyncTimeoutError.
• Unrelated TimeoutErrors pass through. After expiry, cleanup is skipped; a cleanup that hangs after a body error is cut off and adds a note to the original error.
• The example async trainer's blocked generation-buffer wait times out and cancels its generator task.

tests/train/test_config.py: parameterized rejection of invalid values and the NCCL warning.

test_remote_inference_client.py: a real mock-server control request that sleeps past a short budget raises a picklable InferenceServerTimeoutError.

Validation (whole stack)

Current CPU revision: PR61 trainer tests: 19 passed. PR63 focused stack tests: 322 passed. Ruff passes on all changed Python files. A full-suite rerun stopped around the existing Ray dispatch tests; the standalone dispatch test also timed out after 120 s. A prior stack revision passed the full CPU command (RAY_ADDRESS=local uv run --isolated --extra skyrl-train --extra dev pytest -q tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm"): 2020 passed, 32 skipped, 6 deselected.

Prior GPU end-to-end validation (earlier stack revision; 1× H100 node, 4 GPUs, Qwen2.5-0.5B-Instruct on GSM8K, FSDP + vLLM 0.30). A small driver runs the trainer as a Ray task and inspects the exception ray.get raises on the driver. Fault injection wraps generator.generate.

Before this stack, the injected-failure rows end in os._exit(1) / sys.exit(1), so the driver sees the worker process die.

Prior GPU CI subset (-m "not (integrations or megatron or mooncake)"), run on an earlier stack revision and on unmodified main in the same environment:

test_weight_sync.py fails the same 4 tests on main, for the same two reasons on both sides: AttributeError: 'types.SimpleNamespace' object has no attribute 'speculative_config', and a vLLM engine-core init failure. test_update_weights_rdt hits one reason on each side.

🤖 Generated with Claude Code