diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index ad6082eb71a..6687ba9841b 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -545,14 +545,16 @@ def _fetch(self, sample: dict) -> EagleFetchPayload | None: time.sleep(0.0002) agent.release_xfer_handle(h) hidden_states = view.clone() # copy out before /done so the gen check brackets the read - # /done frees the slot + reports valid; valid=False -> ring lapped us mid-read, bytes - # stale -> resample. A failed /done can't prove staleness, so default valid=True. + # /done frees the slot and reports whether the ring lapped us mid-read (stale bytes). + # Fail closed: a failed /done cannot prove the bytes are ours, and a resample is cheap. try: - valid = self._http_rdma.get( - f"http://{host}:{port}/done", params={"req_id": rid} - ).json()["valid"] - except Exception: - valid = True + resp = self._http_rdma.get(f"http://{host}:{port}/done", params={"req_id": rid}) + resp.raise_for_status() + except _TRANSIENT_FETCH_ERRORS as exc: + # Kept apart from the lap warning: these point at the sidecar, not ring pressure. + warn_rank_0(f"[streaming] /done failed for {sample['cid']} ({exc!r}); resampling") + return None + valid = resp.json()["valid"] if not valid: warn_rank_0(f"[streaming] slot lapped mid-read for {sample['cid']}; resampling") return None diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index 00b5a1f6ab5..90dd7ab345a 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -430,6 +430,90 @@ def test_lapped_slot_is_treated_as_miss(monkeypatch): ds[0] +def _first_done_raising_handler(seq, n_layers, hidden, calls): + """Sidecar handler whose first /done raises a transport error; ``calls`` logs the paths.""" + inner = _rdma_sidecar_handler(seq, n_layers, hidden) + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(request.url.path) + if request.url.path == "/done" and calls.count("/done") == 1: + raise httpx.ConnectError("simulated sidecar failure") + return inner(request) + + return handler + + +def test_failed_done_resamples_instead_of_returning_the_read(monkeypatch): + """The discard is a resample, not a hard failure: the next entry is fetched and returned.""" + seq, n_layers, hidden = 8, 3, 16 + calls: list[str] = [] + _mock_rdma( + monkeypatch, + _first_done_raising_handler(seq, n_layers, hidden, calls), + ) + + ds = EagleVllmStreamingDataset( + entries=[ + {"conversation_id": f"c-{i}", "messages": [{"role": "user", "content": "x"}]} + for i in range(2) + ], + tokenizer=_tokenizer_returning(seq), + config=EagleVllmStreamingConfig( + server_urls="http://mock:8000", + model="mock-model", + max_seq_len=seq, + fail_after_consecutive_skips=100, + ), + ) + + batch = ds[0] + assert batch["base_model_hidden_states"].shape == (seq, hidden) + # Two prompts posted: the first read was thrown away, the second is what came back. + assert calls.count("/v1/completions") == 2 + assert calls.count("/done") == 2 + + +def _done_answering(seq, n_layers, hidden, status, body): + """Sidecar handler whose /done answers ``status`` with JSON ``body``.""" + inner = _rdma_sidecar_handler(seq, n_layers, hidden) + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/done": + return httpx.Response(status, json=body) + return inner(request) + + return handler + + +def _one_entry_dataset(seq): + return EagleVllmStreamingDataset( + entries=[{"conversation_id": "c-0", "messages": [{"role": "user", "content": "x"}]}], + tokenizer=_tokenizer_returning(seq), + config=EagleVllmStreamingConfig( + server_urls="http://mock:8000", + model="mock-model", + max_seq_len=seq, + fail_after_consecutive_skips=100, + ), + ) + + +def test_done_server_error_is_a_miss_whatever_it_says(monkeypatch): + """A non-2xx /done is a transport failure: the read is discarded, its body not trusted.""" + seq, n_layers, hidden = 8, 3, 16 + _mock_rdma(monkeypatch, _done_answering(seq, n_layers, hidden, 503, {"valid": True})) + with pytest.raises(RuntimeError, match="no fetchable sample"): + _one_entry_dataset(seq)[0] + + +def test_done_without_valid_is_a_contract_error(monkeypatch): + """A 200 /done without ``valid`` breaks the sidecar contract: it raises, not a miss.""" + seq, n_layers, hidden = 8, 3, 16 + _mock_rdma(monkeypatch, _done_answering(seq, n_layers, hidden, 200, {"freed": "c-0"})) + with pytest.raises(KeyError, match="valid"): + _one_entry_dataset(seq)[0] + + def test_oversize_server_response_raises(monkeypatch): """If the server captured more tokens than max_seq_len (its connector max_tokens > our recv buffer), reading would silently truncate the slice; fail loud instead so the