fix: stamp response with sampler version atomically with the client

tests/unit/test_tinker_responses_model.py

@@ -336,3 +336,83 @@ def reader() -> None:
 
     assert not bad_reads, f"Reader saw invalid values: {bad_reads[:5]}"
     assert read_count[0] > 100, "Reader thread did not run enough iterations"
+
+
+# ---------------------------------------------------------------------------
+# Test 7: atomic (client, version) snapshot — a hot-swap that lands mid-route
+# must NOT retroactively re-stamp the response with the new version.
+# ---------------------------------------------------------------------------
+
+
+def test_get_sampling_client_and_version_reads_atomically():
+    """The combined getter returns a consistent (client, version) snapshot."""
+    client_a = object()
+    m.set_sampling_client(client_a, version=7)  # type: ignore[arg-type]
+    got_client, got_version = m.get_sampling_client_and_version()
+    assert got_client is client_a
+    assert got_version == 7
+
+
+def test_route_stamps_version_of_client_that_served_across_hot_swap():
+    """A concurrent ``set_sampling_client(B, v+1)`` landing *inside* a
+    ``_route_to_tinker`` call must not corrupt the stamped ``sampler_version``.
+
+    The trainer hot-swaps the sampling client between RL steps
+    (``set_sampling_client(new_client, save_count + 1)``). The exact drift
+    diagnostic (``staleness = save_count - sampler_version``, the HF BF16
+    mismatch measurement) relies on the stamped version being the version of
+    the client that *actually produced the tokens*.
+
+    The race window is the span between reading the active client and reading
+    its version. In ``_route_to_tinker`` that span is purely synchronous
+    (render the prompt + build SamplingParams), so an interleaving swap can
+    only be injected there — an ``asyncio`` task cannot preempt it because
+    there is no ``await`` point. We simulate the trainer's swap landing in that
+    window by performing it from inside ``renderer.build_generation_prompt``,
+    which runs *between* the two reads.
+
+    Old two-lock code: ``client = get_sampling_client()`` reads A, then the
+    render triggers the swap to (B, 8), then a late ``get_sampler_version()``
+    reads 8 → the response is stamped 8 even though A (version 7) served it, so
+    the trainer under-counts staleness. New snapshot code captures (A, 7)
+    together up front → stamped 7. This test therefore FAILS on the old code
+    and PASSES on the fix.
+    """
+    client_a = _make_sampling_client(gen_tokens=[1], gen_logprobs=[-0.5])
+    client_b = _make_sampling_client(gen_tokens=[9], gen_logprobs=[-0.05])
+
+    swaps: list[int] = []
+
+    renderer = MagicMock(name="renderer")
+
+    def _build_and_swap(_prompt_messages):
+        # Simulate the trainer's hot-swap landing in the race window: this runs
+        # *between* the client read and the version read in the buggy code.
+        if not swaps:
+            swaps.append(1)
+            m.set_sampling_client(client_b, version=8)
+        return _StubModelInput([10, 20, 30])
+
+    renderer.build_generation_prompt.side_effect = _build_and_swap
+    renderer.parse_response.return_value = ({"role": "assistant", "content": "hi"}, True)
+    renderer.get_stop_sequences.return_value = ["<|eot|>"]
+
+    m.init("m", tokenizer=object(), renderer=renderer)
+    m.set_sampling_client(client_a, version=7)
+
+    out = asyncio.run(
+        m._route_to_tinker(
+            prompt_messages=[{"role": "user", "content": "hi"}],
+            max_tokens=8,
+            temperature=1.0,
+            stop=None,
+        )
+    )
+
+    # The client captured at the top of the call (A, version 7) must be the one
+    # that served the sample — not the mid-flight replacement B.
+    assert client_a.sample_async.await_count == 1
+    assert client_b.sample_async.await_count == 0
+    # And the stamped version must be A's version (7), NOT B's (8). This is the
+    # assertion that fails on the two-lock code.
+    assert out["sampler_version"] == 7

tinker_nemogym/tinker_responses_model.py

@@ -108,6 +108,25 @@ def get_sampler_version() -> int:
         return _current_sampler_version
 
 
+def get_sampling_client_and_version() -> tuple[tinker.SamplingClient | None, int]:
+    """Atomically read the active sampling client *and* its version together.
+
+    Returns ``(client, version)`` captured under a single acquisition of
+    :data:`_sampling_client_lock` so the two values are a consistent snapshot.
+
+    This is what :func:`_route_to_tinker` uses at the top of a call: reading the
+    client and its version in one lock guarantees the version we stamp onto the
+    response matches the client that actually produced the tokens. Reading them
+    with two separate lock acquisitions (``get_sampling_client()`` then a later
+    ``get_sampler_version()``) races an in-flight ``set_sampling_client(new,
+    version + 1)`` hot-swap: the swap can land between the two reads, stamping a
+    version that never produced the sampled tokens and making the trainer's
+    ``staleness = save_count - sampler_version`` diagnostic under-count drift.
+    """
+    with _sampling_client_lock:
+        return _current_sampling_client, _current_sampler_version
+
+
 def init(base_model: str, tokenizer: Any, renderer: Any) -> None:
     """One-time init — stash the tokenizer + renderer + base_model metadata.
 
@@ -169,7 +188,11 @@ async def _route_to_tinker(
         raise SamplerNotReadyError(
             "tinker_responses_model.init(base_model, tokenizer, renderer) has not been called"
         )
-    client = get_sampling_client()
+    # Capture the client AND its version in one lock so they're a consistent
+    # snapshot: the version we stamp below is guaranteed to be the version of
+    # the exact client that serves this call, even if a concurrent
+    # set_sampling_client(new, version + 1) hot-swap lands mid-flight.
+    client, sampler_version_at_call = get_sampling_client_and_version()
     if client is None:
         raise SamplerNotReadyError(
             "No active sampling client; trainer must call set_sampling_client() first"
@@ -221,11 +244,12 @@ async def _route_to_tinker(
         "seed": getattr(sampling_params, "seed", None),
     }
 
-    # 3) Sample. Use the snapshot from above (not _current_sampling_client) so
-    # a concurrent hot-swap mid-flight still gets a stable reference. We
-    # record the sampler_version BEFORE the sample call so downstream
-    # callers can tell "what weights produced these tokens".
-    sampler_version_at_call = get_sampler_version()
+    # 3) Sample. Use the client snapshot captured above (not
+    # _current_sampling_client) so a concurrent hot-swap mid-flight still gets a
+    # stable reference. ``sampler_version_at_call`` was captured in the SAME
+    # lock acquisition as ``client`` (see get_sampling_client_and_version), so
+    # the version we stamp is exactly the one that produced these tokens —
+    # downstream callers can tell "what weights produced these tokens".
     started_at = time.time()
     t0 = time.monotonic()
     result = await client.sample_async(
@@ -651,6 +675,7 @@ async def responses(
     "set_sampling_client",
     "get_sampling_client",
     "get_sampler_version",
+    "get_sampling_client_and_version",
     "init",
     "_route_to_tinker",
     "TinkerResponsesAPIModel",