prain: draft rows from the batch layout

2026-10-27 13:52:25 +01:00
parent cf9a347792
commit 1cf3c18a82
2 changed files with 76 additions and 0 deletions
@@ -0,0 +1,42 @@
+import unittest
+from types import SimpleNamespace
+
+from prain import drafts
+from tests.helpers import TmpCase  # noqa: F401  (every test module imports the helpers)
+
+
+def new_request(req_id, xargs):
+    return SimpleNamespace(req_id=req_id, sampling_params=SimpleNamespace(extra_args=xargs))
+
+
+def step(new=(), finished=()):
+    return SimpleNamespace(scheduled_new_reqs=list(new), finished_req_ids=set(finished))
+
+
+class Drafts(unittest.TestCase):
+    def setUp(self):
+        drafts._DRAFTS.clear()
+        drafts.clear()
+
+    def test_draft_rows_follow_the_batch_layout(self):
+        # a: 3 prompt tokens (full); b: 2 prompt tokens (draft)
+        drafts.note_step(step([new_request("a", None), new_request("b", {"prain": "draft"})]),
+                         ["a", "b"], [0, 3, 5])
+        self.assertEqual(drafts.spans(), [(3, 5)])
+        # next step: both decode one token each, in a new order
+        drafts.note_step(step(), ["b", "a"], [0, 1, 2])
+        self.assertEqual(drafts.spans(), [(0, 1)])
+        drafts.clear()
+        self.assertIsNone(drafts.spans())
+
+    def test_finished_drafts_are_forgotten(self):
+        drafts.note_step(step([new_request("b", {"prain": "draft"})]), ["b"], [0, 4])
+        drafts.note_step(step(finished=["b"]), [], [0])
+        self.assertIsNone(drafts.spans())
+        self.assertEqual(drafts._DRAFTS, set())
+
+    def test_requests_without_the_flag_are_full(self):
+        drafts.note_step(step([new_request("a", {"other": 1}),
+                               SimpleNamespace(req_id="c", sampling_params=None)]),
+                         ["a", "c"], [0, 1, 2])
+        self.assertIsNone(drafts.spans())