feat(client): accept partial replay supports listed by the server

renderers/client.py

@@ -199,6 +199,14 @@ def _parse_compact_logprobs(
     if not np.isfinite(values).all() or (values > 1e-6).any():
         raise MalformedGenerateResponseError("Invalid compact logprob values")
     samples = packed.get("samples") is True
+    # Replayed supports wider than the server's width return only part of their ids, with
+    # probabilities still normalized over the whole support.
+    partial_rows = packed.get("partial_rows", [])
+    if not isinstance(partial_rows, list) or any(
+        type(row) is not int or not 0 <= row < len(completion_ids) for row in partial_rows
+    ):
+        raise MalformedGenerateResponseError("Invalid compact logprob partial rows")
+    partial_rows = set(partial_rows)
     sampled = []
     top_ids: list[list[int]] = []
     top_logprobs: list[list[float]] = []
@@ -234,6 +242,7 @@ def _parse_compact_logprobs(
                 row[token_id],
                 k,
                 supports[index] if supports is not None else None,
+                partial=index in partial_rows,
             )
             top_ids.append(selected_ids)
             top_logprobs.append(selected_logprobs)
@@ -360,8 +369,10 @@ def _select_sampler_head(
     sampled_logprob: float,
     k: int,
     support: list[int] | None = None,
+    partial: bool = False,
 ) -> tuple[list[int], list[float]]:
-    """Keep an action-independent head, or the complete recorded replay support."""
+    """Keep an action-independent head, or the recorded replay support. A ``partial``
+    support holds part of the ids, so its probabilities may sum to less than one."""
     if sampled_id in head and not math.isclose(
         head[sampled_id], sampled_logprob, rel_tol=1e-6, abs_tol=1e-5
     ):
@@ -380,7 +391,7 @@ def _select_sampler_head(
                 "Missing sampling-support probabilities"
             )
         ranked = [(i, p) for i, p in ranked if i in support_set]
-        if not math.isclose(
+        if not partial and not math.isclose(
             math.fsum(math.exp(p) for _, p in ranked), 1.0, rel_tol=0, abs_tol=1e-4
         ):
             raise MalformedGenerateResponseError(

tests/test_client.py

@@ -76,6 +76,9 @@ def _parse_test_head(choice, sampled, k, compact):
         {"data": "?"},
         {"data": None},
         {"data": ""},
+        {"partial_rows": "0"},
+        {"partial_rows": [1]},
+        {"partial_rows": [True]},
     ],
 )
 def test_compact_logprobs_reject_malformed_payload(update):
@@ -194,12 +197,12 @@ def test_top_logprobs_head_is_independent_of_sampled_action(
 
 @pytest.mark.parametrize("compact", [False, True])
 @pytest.mark.parametrize("packed", [False, True])
-@pytest.mark.parametrize("error", [None, "missing", "mass", "mask"])
+@pytest.mark.parametrize("error", [None, "missing", "mass", "mask", "partial"])
 def test_top_logprobs_preserve_complete_replay_support(compact, packed, error):
     import math
 
     probabilities = {1: 0.6, 2: 0.2, 3: 0.2}
-    if error == "mass":
+    if error in ("mass", "partial"):
         probabilities[1] = 0.5
     entry = {
         "token": "token_id:3",
@@ -222,7 +225,14 @@ def test_top_logprobs_preserve_complete_replay_support(compact, packed, error):
             for name, values in (("ids", support), ("counts", [len(support)]))
         }
     choice = {"logprobs": {"content": [entry]}, "sampling_mask": mask}
-    if error is not None:
+    if error == "partial" and compact:
+        # Only the server declares a support partial: its mass may then fall short of one.
+        choice = _compact_choice(choice)
+        choice["compact_logprobs"]["partial_rows"] = [0]
+        _, ids, logprobs = _parse_compact_logprobs(choice, [3], 2)
+        assert ids == [[1, 2, 3]]
+        np.testing.assert_allclose(np.exp(logprobs), [[0.5, 0.2, 0.2]])
+    elif error is not None:
         with pytest.raises(MalformedGenerateResponseError):
             _parse_test_head(choice, [3], 2, compact)
     else: