fix(speculative): fail closed when the streaming /done check cannot be reached

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

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