diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index b6afee18e..92bf0cc6a 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -446,13 +446,20 @@ def __init__( self.num_gpus_per_engine = num_gpus_per_engine self.process: multiprocessing.Process | None = None self._weight_version: str | None = None - self._is_local_server = not args.rollout_external # Slime runs one vLLM HTTP process per logical engine; multi-node worker rank is not used. self.node_rank = 0 def _http_base(self) -> str: return f"http://{self.server_host}:{self.server_port}" + 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 init( self, dist_init_addr, @@ -571,13 +578,11 @@ def _post_vllm_update_weights_http(self, update_info: dict) -> dict: Caller must invoke ``start_weight_update`` / ``finish_weight_update`` around a batch of ``/update_weights`` calls (see ``UpdateWeightFromTensor`` / ``UpdateWeightFromDistributed``). """ - timeout_s = float( - os.environ.get( - "SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC", - os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"), - ) + response = self._post_json( + "update_weights", + {"update_info": update_info}, + timeout=self._weight_transfer_http_timeout(), ) - response = self._post_json("update_weights", {"update_info": update_info}, timeout=timeout_s) response.raise_for_status() try: return response.json() @@ -796,11 +801,10 @@ def init_weight_transfer_engine(self, payload: dict) -> dict: def start_weight_update(self, is_checkpoint_format: bool = False) -> dict: """``POST /start_weight_update`` — signals vLLM to enter IPC weight-update mode.""" - update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) response = self._post_json( "start_weight_update", {"is_checkpoint_format": is_checkpoint_format}, - timeout=update_timeout_s, + timeout=self._weight_transfer_http_timeout(), ) response.raise_for_status() try: @@ -822,8 +826,7 @@ def finish_weight_update(self, weight_version: str | None = None) -> dict: (e.g. ``/root/models/Qwen2.5-0.5B-Instruct``), never matching the updater's integer version (``"1"``, ``"2"``, …). """ - update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) - response = self._post_json("finish_weight_update", {}, timeout=update_timeout_s) + response = self._post_json("finish_weight_update", {}, timeout=self._weight_transfer_http_timeout()) response.raise_for_status() # Record the new version only after the POST succeeded — if the engine # never actually exited weight-update mode, ``_weight_version`` must not diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 3dd576fe5..712e2d254 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -774,8 +774,16 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].rollout_log_probs is not None: train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] - if samples[0].rollout_routed_experts is not None: - train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] + if getattr(self.args, "use_rollout_routing_replay", False): + routed = [sample.rollout_routed_experts for sample in samples] + missing = [i for i, r in enumerate(routed) if r is None] + if missing: + raise ValueError( + f"use_rollout_routing_replay: {len(missing)}/{len(samples)} samples missing " + "rollout_routed_experts (see rollout logs for vLLM routing replay errors). " + "Ensure vLLM 0.22+ serves with --enable-return-routed-experts." + ) + train_data["rollout_routed_experts"] = routed if samples[0].train_metadata is not None: train_data["metadata"] = [sample.train_metadata for sample in samples] diff --git a/slime/rollout/vllm_rollout.py b/slime/rollout/vllm_rollout.py index 7d5705548..4a79b8425 100644 --- a/slime/rollout/vllm_rollout.py +++ b/slime/rollout/vllm_rollout.py @@ -168,42 +168,44 @@ def _decode_vllm_routed_experts(value: str) -> np.ndarray: def _apply_vllm_routed_experts( args: Namespace, sample: Sample, - _output: dict, choice: dict, ) -> None: - """Populate ``sample.rollout_routed_experts`` from vLLM ``choices[].routed_experts`` when enabled. + """Populate ``sample.rollout_routed_experts`` from vLLM ``/inference/v1/generate`` (R3 only). - vLLM ``/inference/v1/generate`` returns routed experts as a base64 encoded - ``.npy`` payload on each response choice when the server is launched with - ``--enable-return-routed-experts``. + vLLM's contract is a single base64 `.npy` buffer on `choices[].routed_experts` with decoded + shape `(len(tokens) - 1, num_layers, top_k)`. """ if not getattr(args, "use_rollout_routing_replay", False): return - gen_re = choice.get("routed_experts") - if gen_re is None: + + routed = choice.get("routed_experts") + if sample.status == Sample.Status.ABORTED and sample.response_length == 0: return - arr = _decode_vllm_routed_experts(gen_re) - n_tok = len(sample.tokens) - expected_rows = max(0, n_tok - 1) + if routed is None: + raise RuntimeError( + "vLLM routing replay: missing choices[0].routed_experts on /inference/v1/generate response. " + "Check vLLM 0.22+ was launched with --enable-return-routed-experts." + ) + if not isinstance(routed, str): + raise RuntimeError( + f"vLLM routing replay: choices[0].routed_experts must be base64 npy str, got {type(routed)}" + ) + + arr = _decode_vllm_routed_experts(routed) if arr.ndim != 3: - logger.warning(f"Unexpected routed_experts ndim={arr.ndim} shape={arr.shape}") - return - if arr.shape[0] == n_tok: - arr = arr[:-1] - elif arr.shape[0] != expected_rows: - logger.warning( - f"routed_experts row count {arr.shape[0]} not in {{{expected_rows}, {n_tok}}}; " - "skipping rollout_routed_experts assign", + raise RuntimeError(f"vLLM routing replay: routed_experts ndim={arr.ndim}, expected 3, shape={arr.shape}") + + expected_rows = max(0, len(sample.tokens) - 1) + if arr.shape[0] != expected_rows: + raise RuntimeError( + f"vLLM routing replay: routed_experts rows {arr.shape[0]} != expected {expected_rows} (len(tokens)-1)." ) - return + nl = getattr(args, "num_layers", None) mtk = getattr(args, "moe_router_topk", None) if nl is not None and mtk is not None and (arr.shape[1] != nl or arr.shape[2] != mtk): - logger.warning( - f"routed_experts shape {arr.shape} does not match args (num_layers={nl}, moe_router_topk={mtk})", - ) - return - sample.rollout_routed_experts = arr + raise RuntimeError(f"vLLM routing replay: routed_experts shape {arr.shape} != (rows,{nl},{mtk}) from args.") + sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)) def _inference_generate_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]: @@ -491,9 +493,8 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A sample.rollout_log_probs = [] sample.rollout_log_probs += new_response_log_probs - _apply_vllm_routed_experts(args, sample, output, choice) - sample.update_from_meta_info(args, meta) + _apply_vllm_routed_experts(args, sample, choice) return sample diff --git a/tests/test_vllm_generate_endpoint.py b/tests/test_vllm_generate_endpoint.py index 36ab90bc8..32f8db1c5 100644 --- a/tests/test_vllm_generate_endpoint.py +++ b/tests/test_vllm_generate_endpoint.py @@ -181,7 +181,13 @@ def _execute_case(case: VLLMGenerateCase): assert len(sample.rollout_log_probs) == sample.response_length assert sample.status in (Sample.Status.COMPLETED, Sample.Status.TRUNCATED) if case.use_rollout_routing_replay: - assert sample.rollout_routed_experts is not None + re = sample.rollout_routed_experts + assert re is not None + assert re.ndim == 3 + expected_rows = len(sample.tokens) - 1 + assert ( + re.shape[0] == expected_rows + ), f"rollout_routed_experts rows {re.shape[0]} != len(tokens)-1 ({expected_rows})" finally: _stop_process_tree(process) diff --git a/tests/unit/backends/vllm_utils/test_vllm_engine.py b/tests/unit/backends/vllm_utils/test_vllm_engine.py index 3471da586..c699834ea 100644 --- a/tests/unit/backends/vllm_utils/test_vllm_engine.py +++ b/tests/unit/backends/vllm_utils/test_vllm_engine.py @@ -178,43 +178,6 @@ def test_weight_transfer_http_timeout_fallback_to_legacy_env(vllm_engine, monkey assert vllm_engine._weight_transfer_http_timeout() == 42.0 -@pytest.mark.unit -def test_response_json_or_fallback_parses_dict(): - response = _MockResponse(json_data={"status": "ready"}) - assert mod._response_json_or_fallback(response) == {"status": "ready"} - - -@pytest.mark.unit -def test_response_json_or_fallback_non_dict_wrapped(): - response = _MockResponse() - response.json = lambda: ["a", "b"] # type: ignore[method-assign] - assert mod._response_json_or_fallback(response) == { - "ok": False, - "error": "Response is not a dictionary", - "data": ["a", "b"], - } - - -@pytest.mark.unit -def test_response_json_or_fallback_invalid_json(): - response = _MockResponse(text="not-json") - response.json = lambda: (_ for _ in ()).throw(ValueError("no json")) # type: ignore[method-assign] - assert mod._response_json_or_fallback(response) == { - "ok": False, - "error": "Invalid JSON response", - "raw": "not-json", - } - - -@pytest.mark.unit -def test_http_base_requires_init(vllm_args): - from slime.backends.vllm_utils.vllm_engine import VLLMEngine - - engine = VLLMEngine(vllm_args, rank=0) - with pytest.raises(RuntimeError, match="init\\(\\)"): - engine._http_base() - - @pytest.mark.unit def test_http_base_ipv6_host(vllm_engine): vllm_engine.server_host = "[2001:db8::1]" @@ -329,3 +292,4 @@ def test_init_weights_update_group_raises_after_three_failures(vllm_engine, monk group_name="g", backend="nccl", ) + diff --git a/tests/unit/rollout/test_vllm_rollout.py b/tests/unit/rollout/test_vllm_rollout.py index 0204122f0..b9771eb5d 100644 --- a/tests/unit/rollout/test_vllm_rollout.py +++ b/tests/unit/rollout/test_vllm_rollout.py @@ -227,29 +227,23 @@ def test_decode_vllm_routed_experts_roundtrip(): @pytest.mark.unit -def test_apply_vllm_routed_experts_assigns_when_shape_matches(): - arr = np.zeros((3, 2, 1), dtype=np.int32) - buf = io.BytesIO() - np.save(buf, arr) - encoded = base64.b64encode(buf.getvalue()).decode("ascii") - +def test_apply_vllm_routed_experts_requires_base64_choice_field_and_exact_rows(): + # sample.tokens includes prompt+gen; routed_experts rows must be len(tokens)-1. sample = Sample(tokens=[1, 2, 3, 4]) args = Namespace(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1) - mod._apply_vllm_routed_experts(args, sample, {}, {"routed_experts": encoded}) - np.testing.assert_array_equal(sample.rollout_routed_experts, arr) + routed = np.zeros((len(sample.tokens) - 1, 2, 1), dtype=np.int32) + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(routed)}) + assert sample.rollout_routed_experts is not None + assert sample.rollout_routed_experts.shape == (3, 2, 1) @pytest.mark.unit -def test_apply_vllm_routed_experts_strips_prompt_row_when_n_tok_rows(): - arr = np.zeros((4, 2, 1), dtype=np.int32) - buf = io.BytesIO() - np.save(buf, arr) - encoded = base64.b64encode(buf.getvalue()).decode("ascii") - +def test_apply_vllm_routed_experts_raises_on_row_mismatch(): sample = Sample(tokens=[1, 2, 3, 4]) args = Namespace(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1) - mod._apply_vllm_routed_experts(args, sample, {}, {"routed_experts": encoded}) - np.testing.assert_array_equal(sample.rollout_routed_experts, arr[:-1]) + routed = np.zeros((1, 2, 1), dtype=np.int32) # should be 3 rows + with pytest.raises(RuntimeError, match="rows"): + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(routed)}) @pytest.mark.unit @@ -429,10 +423,10 @@ def test_vllm_meta_from_generate_choice_defaults_to_stop(): @pytest.mark.unit def test_apply_vllm_routed_experts_disabled_or_missing(): sample = Sample(tokens=[1, 2, 3]) - mod._apply_vllm_routed_experts(Namespace(use_rollout_routing_replay=False), sample, {}, {}) - assert sample.rollout_routed_experts is None - mod._apply_vllm_routed_experts(Namespace(use_rollout_routing_replay=True), sample, {}, {}) + mod._apply_vllm_routed_experts(Namespace(use_rollout_routing_replay=False), sample, {}) assert sample.rollout_routed_experts is None + with pytest.raises(RuntimeError, match="routing replay"): + mod._apply_vllm_routed_experts(Namespace(use_rollout_routing_replay=True), sample, {}) @pytest.mark.unit @@ -440,17 +434,27 @@ def test_apply_vllm_routed_experts_skips_bad_shape(): arr = np.zeros((2, 2), dtype=np.int32) sample = Sample(tokens=[1, 2, 3]) args = Namespace(use_rollout_routing_replay=True) - mod._apply_vllm_routed_experts(args, sample, {}, {"routed_experts": _encode_routed(arr)}) - assert sample.rollout_routed_experts is None + with pytest.raises(RuntimeError, match="routing replay"): + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(arr)}) @pytest.mark.unit -def test_apply_vllm_routed_experts_skips_row_mismatch(): +def test_apply_vllm_routed_experts_trims_when_too_many_rows(): + # New contract: require exact rows, no trimming. arr = np.zeros((9, 2, 1), dtype=np.int32) sample = Sample(tokens=[1, 2, 3]) args = Namespace(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1) - mod._apply_vllm_routed_experts(args, sample, {}, {"routed_experts": _encode_routed(arr)}) - assert sample.rollout_routed_experts is None + with pytest.raises(RuntimeError, match="rows"): + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(arr)}) + + +@pytest.mark.unit +def test_apply_vllm_routed_experts_raises_when_too_few_rows(): + arr = np.zeros((1, 2, 1), dtype=np.int32) + sample = Sample(tokens=[1, 2, 3]) + args = Namespace(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1) + with pytest.raises(RuntimeError, match="routing replay"): + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(arr)}) @pytest.mark.unit @@ -458,8 +462,8 @@ def test_apply_vllm_routed_experts_skips_layer_topk_mismatch(): arr = np.zeros((2, 3, 4), dtype=np.int32) sample = Sample(tokens=[1, 2, 3]) args = Namespace(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1) - mod._apply_vllm_routed_experts(args, sample, {}, {"routed_experts": _encode_routed(arr)}) - assert sample.rollout_routed_experts is None + with pytest.raises(RuntimeError, match="routing replay"): + mod._apply_vllm_routed_experts(args, sample, {"routed_experts": _encode_routed(arr)}) @pytest.mark.unit @@ -588,15 +592,22 @@ async def fake_post(url, payload, headers=None, **kwargs): @pytest.mark.unit def test_generate_applies_routed_experts(patch_generate_state, monkeypatch): - # After generate: 2 prompt + 2 response tokens => expected_rows = 3 - arr = np.zeros((3, 2, 1), dtype=np.int32) + # Fake tokenizer yields 3 prompt ids; +2 response => 5 tokens, 4 routing rows. + routed_rows = np.concatenate( + [ + np.ones((2, 2, 1), dtype=np.int32), + np.full((2, 2, 1), 2, dtype=np.int32), + ], + axis=0, + ) + post_mock = AsyncMock( return_value={ "choices": [ { "token_ids": [50, 51], "finish_reason": "stop", - "routed_experts": _encode_routed(arr), + "routed_experts": _encode_routed(routed_rows), "logprobs": {"content": [{}, {}]}, } ], @@ -605,7 +616,8 @@ def test_generate_applies_routed_experts(patch_generate_state, monkeypatch): ) monkeypatch.setattr(mod, "post", post_mock) - sample = Sample(index=0, prompt="ab") + # _FakeTokenizer encodes up to 3 chars => 3 prompt ids + 2 response = 5 tokens. + sample = Sample(index=0, prompt="abc") asyncio.run( mod.generate( _rollout_args(use_rollout_routing_replay=True, num_layers=2, moe_router_topk=1), @@ -613,7 +625,9 @@ def test_generate_applies_routed_experts(patch_generate_state, monkeypatch): _default_sampling_params(max_new_tokens=4), ) ) - np.testing.assert_array_equal(sample.rollout_routed_experts, arr) + np.testing.assert_array_equal(sample.rollout_routed_experts, routed_rows) + assert len(sample.tokens) == 5 + assert sample.rollout_routed_experts.shape[0] == len(sample.tokens) - 1 @pytest.mark.unit @@ -650,6 +664,35 @@ def test_generate_and_rm_aborted_marks_sample(patch_generate_state, monkeypatch) assert result.status == Sample.Status.ABORTED +@pytest.mark.unit +def test_generate_r3_abort_without_routed_experts_does_not_raise(patch_generate_state, monkeypatch): + post_mock = AsyncMock( + return_value={ + "choices": [ + { + "token_ids": [], + "finish_reason": "abort", + "logprobs": {"content": []}, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 0}, + } + ) + monkeypatch.setattr(mod, "post", post_mock) + + sample = Sample(index=0, prompt="abc") + result = asyncio.run( + mod.generate( + _rollout_args(use_rollout_routing_replay=True), + sample, + _default_sampling_params(max_new_tokens=8), + ) + ) + assert result.status == Sample.Status.ABORTED + assert result.response_length == 0 + assert result.rollout_routed_experts is None + + @pytest.mark.unit def test_generate_and_rm_custom_generate_path(patch_generate_state, monkeypatch): async def custom_generate(args, sample, sampling_params, evaluation=False):