Sync to vLLM 0.21.0

[CalvinXKY]

Purpose

This PR adapts vime’s Megatron→vLLM native weight sync to vLLM 0.21.x, as part of the image / dependency upgrade tracked in Issue #20 (vllm 0.21 + torch 2.11 stack).

vLLM 0.21 exposes a four-phase weight transfer control plane (vLLM #39212):

• POST /start_weight_update
• POST /update_weights
• POST /finish_weight_update

Trainer-side sends use NCCLTrainerSendWeightsArgs with NCCLWeightTransferEngine.trainer_send_weights(...).

In-process NCCL on the trainer (replacing the subprocess bridge) is already covered by #13; this PR builds on that and focuses on 0.21 HTTP/API alignment plus unit test coverage.

What’s included

vLLM 0.21 weight sync (update_weight_from_distributed.py)

• Wrap each full sync with rank-0 start_weight_update(is_checkpoint_format=True) → per-bucket /update_weights + trainer_send_weights(NCCLTrainerSendWeightsArgs(...)) → finish_weight_update, with Gloo barriers.
• Per-chunk update_info no longer carries legacy is_checkpoint_format (handled in start).

VLLMEngine (vllm_engine.py)

• Always call vLLM 0.21 start/finish/update routes (no legacy protocol branch).
• Shared _response_json_or_fallback(), unified weight-transfer HTTP timeout, _http_base() init guard.

Unit tests (new layout)

Introduce layered unit tests under tests/unit/ (pytest unit marker unchanged):

Path	Scope
tests/unit/backends/vllm_utils/test_vllm_engine.py	vLLM 0.21 HTTP four-phase endpoints, helpers (_normalize_vllm_wake_tags, _serialize_for_cli, init retry, IPv6 base URL, etc.) — 28 cases, no live vLLM server
tests/unit/backends/vllm_utils/test_arguments.py	Moved from tests/test_vllm_arguments.py (CLI / router args)
tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_distributed.py	Moved from tests/test_update_weight_from_distributed.py; covers in-process NCCL send, start/finish session, packed buckets
tests/unit/conftest.py	Ray stub for local collection without a cluster

Run:

pytest tests/unit/backends -m unit -v
# or
pytest tests/unit/backends/vllm_utils/test_vllm_engine.py -m unit -v

Test plan

• Unit: pytest tests/unit/backends -m unit -v
• E2E: bash run_scripts/qwen_4b.sh with --rollout-backend vllm --vllm-weight-sync-mode native
• vLLM server exposes /start_weight_update, /update_weights, /finish_weight_update (0.21+)

Environment / install: same as #3 and #20 (vllm 0.21 cu129 image work in #17).

Test result

• test_vllm_engine.py: 28 passed (HTTP mocks)
• test_update_weight_from_distributed.py: passed on training host (incl. start/finish session + Gloo mock)
• E2E (Qwen3-4B, 8×A100, 4 train + 4 vLLM): 6+ steps, stable sync; benefits vs bridge baseline documented in [Not maintained · vLLM 0.16.0 only · kept for reference] Fix trainer GPU0 memory skew: replace NcclBridge with in-process trainer_init #13

Known issues / follow-up

1. Requires vLLM ≥ 0.21 with four-phase routes and NCCLTrainerSendWeightsArgs (see [RFC] Docker image roadmap #20).
2. Design doc slime_vllm_backend_design_v1.md still mentions NcclBridge — update separately.
3. Docker / cu13 path continues under [RFC] Docker image roadmap #20 / dockerfile: cu12 base for vllm rollout (vllm/vllm-openai:v0.21.0-cu129) #17.

[gemini-code-assist]

Code Review

This pull request refactors the vLLM weight transfer mechanism by replacing the subprocess-based _NcclBridge with an in-process implementation using NCCLWeightTransferEngine. It introduces session management for weight updates and adds support for new HTTP endpoints in the VLLMEngine. Feedback was provided regarding the _response_json_or_fallback helper, noting that hardcoding a success flag in fallback paths could mask server errors or malformed responses.

[gemini-code-assist]

The _response_json_or_fallback helper hardcodes "ok": True in all fallback paths, including when JSON decoding fails or when the response body is not a dictionary. This is misleading because it masks potential logical errors or malformed responses from the vLLM server that didn't result in a non-2xx HTTP status code. If the server returns plain text indicating an error but with a 200 OK status, this function will report success to the caller.

Suggested change

	def _response_json_or_fallback(response: requests.Response) -> dict:
	"""Parse JSON body; on decode failure return a minimal dict with raw text."""
	try:
	body = response.json()
	if isinstance(body, dict):
	return body
	return {"ok": True, "data": body}
	except ValueError:
	return {"ok": True, "raw": response.text}
	def _response_json_or_fallback(response: requests.Response) -> dict:
	"""Parse JSON body; on decode failure return a minimal dict with raw text."""
	try:
	body = response.json()
	if isinstance(body, dict):
	return body
	return {"ok": False, "error": "Response is not a dictionary", "data": body}
	except ValueError:
	return {"ok": False, "error": "Invalid JSON response", "raw": response.text}

[CalvinXKY]

The qwen_4b running result：

[aoshen02]

/gemini review this

[gemini-code-assist]

Code Review

This pull request refactors the vLLM weight transfer mechanism to use an in-process NCCLWeightTransferEngine, replacing the previous subprocess-based _NcclBridge. It introduces explicit session management with start_weight_update and finish_weight_update calls and updates the VLLMEngine to support these new endpoints. Feedback suggests optimizing performance by moving local imports and caching environment variable lookups that occur within hot loops.

[gemini-code-assist]

The local imports for NCCLTrainerSendWeightsArgs and NCCLWeightTransferEngine are performed inside update_weights_from_distributed, which is called for every bucket of weights during the sync process. While sys.modules caching makes this relatively fast, performing these lookups in a hot loop adds unnecessary overhead. Consider moving these imports to the top of the file (using a try-except block if vllm is an optional dependency) or caching them at the module level.

[gemini-code-assist]

The _weight_transfer_http_timeout method performs multiple os.environ.get calls every time it is invoked. Since this is used in _post_vllm_update_weights_http, which is called per weight bucket, this results in redundant environment variable lookups. It is more efficient to resolve and cache this timeout value once during __init__ or the init method.

Suggested change

	def _weight_transfer_http_timeout(self) -> float:
	return float(
	os.environ.get(
	"SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC",
	os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"),
	)
	)
	def _weight_transfer_http_timeout(self) -> float:
	if not hasattr(self, "_cached_timeout"):
	self._cached_timeout = float(
	os.environ.get(
	"SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC",
	os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"),
	)
	)
	return self._cached_timeout

[aoshen02]

What about syncing after all weight bucket are sent?

[CalvinXKY]

Better！ Done：

After the entire sync round finishes, call torch.cuda.synchronize() to avoid synchronizing per bucket.

Works well：

[aoshen02]

Please refactor this in slime/backends/vllm_utils/arguments.py#L220:
help="Disable packed sync; use per-bucket NCCL via NcclBridge instead.",

[aoshen02]

move it to the top

[CalvinXKY]

Done

[CalvinXKY]

update：