feat: label sampled tokens with the policy step and log per-token policy age (vLLM 0.32 weight_versions)

src/prime_rl/inference/vllm/server.py

@@ -56,12 +56,15 @@ async def resume(request: Request):
 @router.post("/update_weights")
 async def update_weights(request: Request):
     data = await request.json()
-    await engine_client(request).collective_rpc("update_weights_from_path", args=(data.get("weight_dir"),))
+    client = engine_client(request)
+    await client.collective_rpc("update_weights_from_path", args=(data.get("weight_dir"),))
+    # Label the tokens sampled from now on (``weight_versions`` spans on /inference/v1/generate).
+    await client.update_weight_version(str(data["weight_version"]))
     return {"status": "ok"}
 
 
 @router.post("/load_lora_adapter")
-async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: Request):
+async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: Request, weight_version: str):
     """Wrapper around vLLM's /v1/load_lora_adapter.
 
     prime-rl reloads a fixed-name adapter with fresh weights every step (the path
@@ -89,6 +92,7 @@ async def load_lora_adapter(lora_request: LoadLoRAAdapterRequest, raw_request: R
         stored.load_inplace = False
     if isinstance(response, ErrorResponse):
         return JSONResponse(content=response.model_dump(), status_code=response.error.code)
+    await engine_client(raw_request).update_weight_version(weight_version)
     return {"status": "ok"}
 
 

src/prime_rl/orchestrator/clients.py

@@ -202,7 +202,7 @@ async def update_weights(
             transfer_times = await _post_all(
                 self.clients,
                 "/update_weights",
-                json={"weight_dir": weight_dir_posix},
+                json={"weight_dir": weight_dir_posix, "weight_version": step},
                 timeout_s=UPDATE_WEIGHTS_TIMEOUT_S,
             )
         finally:
@@ -475,7 +475,7 @@ def _is_retryable_lora_error(exception: BaseException) -> bool:
 LORA_LOAD_TOTAL_TIMEOUT_S = 120.0
 
 
-async def load_lora_adapter(admin_plane: AdminPlane, lora_name: str, lora_path: Path) -> None:
+async def load_lora_adapter(admin_plane: AdminPlane, lora_name: str, lora_path: Path, step: int) -> None:
     """Make a HTTP post request to the vLLM server to load a LoRA adapter.
 
     Uses our wrapper around vLLM's /v1/load_lora_adapter. The prefix cache is not reset
@@ -499,6 +499,7 @@ async def _load_lora_adapter(admin_client: AsyncClient) -> None:
         response = await admin_client.post(
             "/load_lora_adapter",
             json={"lora_name": lora_name, "lora_path": lora_path_posix},
+            params={"weight_version": step},
             timeout=httpx.Timeout(connect=10.0, read=LORA_LOAD_READ_TIMEOUT_S, write=60.0, pool=10.0),
         )
         response.raise_for_status()

src/prime_rl/orchestrator/orchestrator.py

@@ -68,6 +68,7 @@
     intercept_vf_logging,
     release_and_trim,
     set_default_executor,
+    token_policy_ages,
     trim_process_memory,
 )
 from prime_rl.orchestrator.watcher import WeightWatcher
@@ -728,6 +729,15 @@ def shipped_metrics() -> dict[str, float]:
                     "off_policy/in_queue/mean": sum(in_queue) / len(in_queue),
                     "off_policy/in_queue/max": float(max(in_queue)),
                 }
+            # Per-token age from the engines' weight-version labels: tokens sampled
+            # after mid-rollout weight updates are younger than their episode.
+            ages = [token_policy_ages(record.trace, step) for record in effective.records]
+            num_labelled = sum(count for _, count, _ in ages)
+            if num_labelled:
+                metrics |= {
+                    "off_policy/token/mean": sum(total for total, _, _ in ages) / num_labelled,
+                    "off_policy/token/max": float(max(oldest for _, _, oldest in ages)),
+                }
             for env_name, env_pool in effective.by_env().items():
                 metrics[f"batch/{env_name}"] = env_pool.num_traces / effective.num_traces
             return metrics

src/prime_rl/orchestrator/utils.py

@@ -2,9 +2,11 @@
 import ctypes
 import logging
 import math
+from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
+import numpy as np
 import verifiers.v1 as vf
 
 from prime_rl.utils.logger import InterceptHandler, get_logger, setup_logger
@@ -44,6 +46,22 @@ def min_fresh_version(step: int, max_off_policy_steps: int) -> int:
     return (step - 1) - max_off_policy_steps
 
 
+def node_policy_versions(node: vf.MessageNode, fallback: int) -> Iterator[tuple[int, int, int]]:
+    """``(start, end, version)`` runs over a node's sampled tokens, indexed into
+    ``node.token_ids``: the policy step whose weights sampled them, from the
+    engine's ``weight_versions`` labels. Tokens without a step label (engines
+    whose updates bypass ``/update_weights``, e.g. Dynamo) get ``fallback``."""
+    sampled = int(np.count_nonzero(node.mask))
+    offset = len(node.mask) - sampled
+    if not node.weight_versions:
+        if sampled:
+            yield offset, offset + sampled, fallback
+        return
+    for span in node.weight_versions:
+        version = int(span.version) if span.version.isdigit() else fallback
+        yield offset + span.start, offset + span.end, version
+
+
 def episode_staleness(episode: vf.Episode[Any, Any, Any], training_step: int) -> tuple[int, int, int]:
     """``(total, in_flight, in_queue)`` staleness of one train episode when
     consumed by batch ``training_step``: the version the batch trains on
@@ -60,6 +78,24 @@ def episode_staleness(episode: vf.Episode[Any, Any, Any], training_step: int) ->
     return total, in_flight, in_queue
 
 
+def token_policy_ages(trace: vf.Trace, training_step: int) -> tuple[int, int, int]:
+    """``(sum, count, max)`` of per-token policy age over the trace's sampled
+    tokens: the version batch ``training_step`` trains on (v{step-1}) minus the
+    policy step whose weights sampled the token, from the engines'
+    ``weight_versions`` labels. Unlabelled tokens (``"default"``: Dynamo
+    engines, which update weights outside ``/update_weights``) are not counted."""
+    total = count = oldest = 0
+    for node in trace.nodes:
+        for start, end, version in node_policy_versions(node, fallback=-1):
+            if version < 0:
+                continue
+            age = max(0, (training_step - 1) - version)
+            total += age * (end - start)
+            count += end - start
+            oldest = max(oldest, age)
+    return total, count, oldest
+
+
 def intercept_vf_logging(logger: str = "verifiers", level: str = "DEBUG", prefix: str | None = None):
     """Intercepts verifiers logging and routes through prime-rl logger with optional prefix."""
     vf_logger = logging.getLogger(logger)

src/prime_rl/transports/weights/filesystem.py

@@ -69,6 +69,6 @@ async def receive(self, step: int) -> None:
         self._ack(step)
         await wait_for_path(weights_dir / FINISHED_MARKER)
         if (weights_dir / "adapter_config.json").exists():
-            await load_lora_adapter(self.admin_plane, self.model_name, weights_dir)
+            await load_lora_adapter(self.admin_plane, self.model_name, weights_dir, step)
         else:
             await self.admin_plane.update_weights(weights_dir, transport="filesystem", step=step)

tests/unit/orchestrator/test_clients.py

@@ -48,11 +48,12 @@ def test_load_lora_adapter_succeeds_on_first_attempt():
     mock_response.raise_for_status = MagicMock()
     mock_client.post.return_value = mock_response
 
-    asyncio.run(load_lora_adapter(admin_plane, "test-lora", Path("/test/path")))
+    asyncio.run(load_lora_adapter(admin_plane, "test-lora", Path("/test/path"), 3))
 
     mock_client.post.assert_called_once_with(
         "/load_lora_adapter",
         json={"lora_name": "test-lora", "lora_path": "/test/path"},
+        params={"weight_version": 3},
         timeout=httpx.Timeout(connect=10.0, read=30.0, write=60.0, pool=10.0),
     )
 

tests/unit/orchestrator/test_token_policy_ages.py

@@ -0,0 +1,23 @@
+from types import SimpleNamespace
+
+import numpy as np
+from verifiers.v1.types import WeightVersionSpan
+
+from prime_rl.orchestrator.utils import token_policy_ages
+
+
+def test_token_policy_ages_counts_labelled_tokens_only():
+    spans = [
+        WeightVersionSpan(version="default", start=0, end=4),
+        WeightVersionSpan(version="7", start=4, end=10),
+        WeightVersionSpan(version="9", start=10, end=12),
+    ]
+    trace = SimpleNamespace(
+        nodes=[
+            SimpleNamespace(mask=np.array([False, True]), weight_versions=None),
+            SimpleNamespace(mask=np.array([False] + [True] * 12), weight_versions=spans),
+        ]
+    )
+
+    # Batch 11 trains on v10: 6 tokens of age 3 and 2 tokens of age 1.
+    assert token_policy_ages(trace, 11) == (20, 8, 3)