[v5] Keep an explicit zero optimizer step out of the driver budget

src/twinkle_agentic/async_rl/streaming_driver.py

@@ -797,7 +797,8 @@ def _train_pull(self, records: List[BufferRecord], ready: List[BufferRecord],
             if group_counts is not None:
                 group_counts.pop(record.group_key, None)
         steps = self._consume_batch(records)
-        steps = int(steps) if steps else 1
+        # None means the consumer did not report a count. An explicit 0 is a real report.
+        steps = 1 if steps is None else int(steps)
         self._last_pull_steps = steps
         if steps > 0:
             current = self._current_version()
@@ -1153,16 +1154,18 @@ def __next__(self) -> Any:
             raise StopIteration
         return batch
 
-    def commit(self, optimizer_steps: int = 1) -> None:
+    def commit(self, optimizer_steps: Optional[int] = 1) -> None:
         """Release the driver thread blocked in the bridge, reporting the real optimizer steps just run.
 
         The recipe calls this after training the batch ``__next__`` handed it. The count becomes the
         bridge's consume return, which drives the publish cadence (twinkle publishes after consume returns,
-        so the trained weights -- not stale ones -- are what the next admission window samples from). A falsy
-        count is treated as one step, matching :meth:`_train_pull`. The same count also feeds the internal
-        ``max_steps`` budget, so a driver constructed with a numeric budget exits on its own.
+        so the trained weights -- not stale ones -- are what the next admission window samples from). Omitting
+        the count, or passing ``None``, means one step. An explicit ``0`` means this batch did not update the
+        optimizer and does not consume the ``max_steps`` budget. The same count also feeds that budget, so a
+        driver constructed with a numeric budget exits on its own.
         """
-        self._committed_steps = int(optimizer_steps) if optimizer_steps else 1
+        # None means the caller omitted the count. An explicit 0 does not consume budget.
+        self._committed_steps = 1 if optimizer_steps is None else int(optimizer_steps)
         self._budget_steps += self._committed_steps
         self._ack_event.set()
 

tests/twinkle_agentic/test_streaming_driver_budget.py

@@ -0,0 +1,116 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""A reported optimizer step of 0 must not be rewritten into a budget step."""
+
+import types
+
+from twinkle_agentic.async_rl.streaming_driver import BufferRecord, StreamingDriver
+
+
+def _driver(**overrides):
+    kwargs = dict(
+        prompt_stream=[],
+        max_staleness=0,
+        num_generations=1,
+        groups_per_partition=1,
+        buffer_depth=1,
+        submit=lambda prompt, index, policy: None,
+        poll=lambda handles: [],
+        collect=lambda handle: None,
+        assembly_ready=lambda ready: None,
+        consume=lambda records: 0,
+        max_steps=3,
+    )
+    kwargs.update(overrides)
+    return StreamingDriver(**kwargs)
+
+
+def test_commit_zero_does_not_exhaust_a_three_step_budget():
+    driver = _driver()
+    for reported in (0, 1, 1):
+        driver.commit(reported)
+    assert driver._committed_steps == 1
+    assert driver._budget_steps == 2
+    assert driver._default_reached_max() is False
+    driver.commit(1)
+    assert driver._budget_steps == 3
+    assert driver._default_reached_max() is True
+
+
+def test_omitted_commit_still_counts_as_one_step():
+    driver = _driver()
+    driver.commit()
+    driver.commit(None)
+    assert driver._budget_steps == 2
+    assert driver._committed_steps == 1
+
+
+def test_train_pull_keeps_explicit_zero_out_of_the_stats():
+    driver = _driver()
+    record = BufferRecord(sample=object(), version=0, group_key='g')
+    ready = [record]
+    driver._consume = lambda records: 0
+    driver._train_pull([record], ready, {'g': 1})
+    assert ready == []
+    assert driver._last_pull_steps == 0
+    assert driver.stats.optimizer_steps == 0
+    assert driver.stats.train_pulls == 0
+    assert driver.stats.consumed == 0
+
+
+def test_train_pull_missing_report_defaults_to_one_step():
+    driver = _driver()
+    record = BufferRecord(sample=object(), version=0, group_key='g')
+    ready = [record]
+    driver._consume = lambda records: None
+    driver._train_pull([record], ready, {'g': 1})
+    assert ready == []
+    assert driver._last_pull_steps == 1
+    assert driver.stats.optimizer_steps == 1
+    assert driver.stats.train_pulls == 1
+    assert driver.stats.consumed == 1
+
+
+def test_iterator_runs_until_three_real_updates():
+    """The disaggregated loop reports [0, 1, 1, 1]; only the real updates fill max_steps=3."""
+    active = {}
+    next_handle = 0
+
+    def submit(prompt, index, policy):
+        nonlocal next_handle
+        next_handle += 1
+        active[next_handle] = policy.version
+        return next_handle
+
+    def collect(handle):
+        active.pop(handle)
+        return types.SimpleNamespace()
+
+    driver = StreamingDriver(
+        prompt_stream=list(range(32)),
+        max_staleness=0,
+        num_generations=1,
+        groups_per_partition=1,
+        buffer_depth=4,
+        submit=submit,
+        poll=lambda handles: list(handles),
+        collect=collect,
+        group_pull=True,
+        need_rows=4,
+        version_control=True,
+        max_steps=3,
+        consume=None,
+    )
+    reported = iter((0, 1, 1, 1))
+    seen = []
+    try:
+        for batch in driver:
+            assert len(batch) == 4
+            steps = next(reported)
+            seen.append(steps)
+            driver.commit(steps)
+    finally:
+        driver.abort()
+    assert seen == [0, 1, 1, 1]
+    assert driver._budget_steps == 3
+    assert driver.stats.optimizer_steps == 3
+    assert driver.stats.consumed == 12