feat: add disk-mode weight update flow to gateway

[garrett4wade]

Description

Add disk-mode weight updates to the experimental weight-update gateway so training workers can save checkpoints to shared storage and inference workers can load from disk. The update flow now consistently pauses and resumes inference around transfer operations, including error paths, and extends controller/gateway contracts for mode-specific fields.

Related Issue

Fixes #(issue)

Type of Change

• 🐛 Bug fix
• ✨ New feature
• 💥 Breaking change
• 📝 Documentation update
• ♻️ Refactoring
• ⚡ Performance improvement
• ✅ Test coverage improvement

Checklist

• I have read the Contributing Guide
• Pre-commit hooks pass (pre-commit run --all-files)
• Relevant tests pass; new tests added for new functionality
• Documentation updated (if applicable; built with ./docs/build_all.sh)
• Branch is up to date with main
• Self-reviewed via /review-pr command
• This PR was created by a coding agent via /create-pr
• This PR is a breaking change

Breaking Change Details (if applicable):

N/A

Additional Context

Risk areas:

• Disk path correctness and shared storage permissions in deployment.
• Inference pause/continue endpoint behavior under transient failures.

Test commands run:

• uv run pytest tests/experimental/weight_update/test_wu_controller.py
• uv run pytest tests/experimental/weight_update/test_disk_integration.py -k "not multi_gpu and not slow and not sglang"

Skipped suites:

• 3 multi-GPU/slow/sglang-marked cases in test_disk_integration.py (environment-dependent E2E coverage).

[gemini-code-assist]

Code Review

This pull request introduces a disk-based weight update mechanism, allowing model weights to be transferred via shared storage as an alternative to the AWEX method. The changes include updates to the API structures, controllers, and the gateway to handle disk-specific parameters such as save paths and LoRA adapter names. The gateway now manages the update lifecycle by pausing inference, triggering saves from training workers, loading weights into inference workers, and resuming generation. Feedback focuses on optimizing the save process for distributed training to avoid redundant operations, validating mandatory fields for disk mode during the connection phase, and improving error logging with tracebacks.

[gemini-code-assist]

Calling the /save endpoint on all training workers simultaneously can be problematic if the training engine (e.g., FSDP) performs a collective consolidated save. In such cases, the save operation should typically be triggered on only one rank (usually rank 0) to avoid redundant operations or race conditions when writing to shared storage. If the workers are expected to save shards, they would likely need distinct paths.

        # Trigger save on the first training worker (typically rank 0).
        # For distributed training like FSDP, a single trigger is usually sufficient
        # for a consolidated checkpoint save to shared storage.
        await _post_json(
            session,
            f"{pair_info.train_worker_urls[0]}/save",
            timeout_s,
            json_data=save_payload,
        )

[gemini-code-assist]

When mode is set to "disk", the save_path is required for the weight update flow to function correctly. It should be validated during the connection phase to return a clear error to the client if it's missing, rather than failing later during the transfer operation.

Suggested change

	if body.mode == "disk":
	pair_info = PairInfo(
	pair_name=pair_name,
	train_worker_urls=train_urls,
	inference_worker_urls=inference_urls,
	mode="disk",
	save_path=body.save_path,
	use_lora=body.use_lora,
	lora_name=body.lora_name,
	)
	if body.mode == "disk":
	if not body.save_path:
	return JSONResponse(
	status_code=400,
	content={"error": "save_path is required for disk mode"},
	)
	pair_info = PairInfo(
	pair_name=pair_name,
	train_worker_urls=train_urls,
	inference_worker_urls=inference_urls,
	mode="disk",
	save_path=body.save_path,
	use_lora=body.use_lora,
	lora_name=body.lora_name,
	)

[gemini-code-assist]

Including the exception traceback in the log message will greatly assist in debugging weight update failures.

Suggested change

	logger.error(
	"Weight update failed for pair '%s': %s",
	pair_info.pair_name,
	e,
	)
	logger.error(
	"Weight update failed for pair '%s'",
	pair_info.pair_name,
	exc_info=True,
	)

[rchardx]

LGTM