pass the engine weight version from the trainer instead of polling the router

miles/backends/experimental/fsdp_utils/actor.py

@@ -589,6 +589,8 @@ def update_weights(self, info: "EnginesAndLock") -> None:  # type: ignore[overri
                 ray.get(self.rollout_manager.clear_updatable_has_new_engines.remote())
 
         self.weight_updater.update_weights()
+        if dist.get_rank() == 0:
+            ray.get(self.rollout_manager.set_weight_version.remote(self.weight_updater.weight_version))
 
         if self.args.ci_test and len(rollout_engines) > 0:
             engine = random.choice(rollout_engines)

miles/backends/megatron_utils/actor.py

@@ -770,6 +770,8 @@ def update_weights(self, info: "EnginesAndLock") -> None:
             print_memory("before update_weights")
             self.weight_updater.update_weights()
             print_memory("after update_weights")
+            if dist.get_rank() == 0:
+                ray.get(self.rollout_manager.set_weight_version.remote(self.weight_updater.weight_version))
 
             if is_multi_lora_enabled(self.args):
                 from miles.backends.megatron_utils.multi_lora_utils import commit_weight_push

miles/ray/rollout/rollout_manager.py

@@ -60,6 +60,8 @@ def __init__(self, args, pg):
 
         self.pg = pg
         self.args = args
+        # set by the training actor after each weight update
+        self.weight_version: int | None = None
         # TODO make args immutable
         init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}")
         object_store.init_instance(args, contribute_segment=False)
@@ -238,7 +240,9 @@ async def _get_rollout_data(self, rollout_id):
         else:
             if self.use_experimental_refactor:
                 data = await asyncio.to_thread(
-                    call_rollout_function, self.generate_rollout, RolloutFnTrainInput(rollout_id=rollout_id)
+                    call_rollout_function,
+                    self.generate_rollout,
+                    RolloutFnTrainInput(rollout_id=rollout_id, weight_version=self.weight_version),
                 )
             else:
                 data = await asyncio.to_thread(
@@ -372,6 +376,9 @@ async def check_weights(
             action=action, allow_quant_error=allow_quant_error, selector=selector, skip_list=skip_list
         )
 
+    def set_weight_version(self, weight_version: int):
+        self.weight_version = weight_version
+
     def set_train_parallel_config(self, config: dict):
         self.train_parallel_config = config
 

miles/rollout/base_types.py

@@ -30,6 +30,9 @@ def evaluation(self):
 # subclassing for different data in the future
 @dataclass(frozen=True)
 class RolloutFnTrainInput(RolloutFnBaseInput):
+    # engine weight version, None before the first weight update
+    weight_version: int | None = None
+
     @property
     def evaluation(self):
         return False

miles/rollout/fully_async_rollout.py

@@ -18,32 +18,28 @@
 
 import asyncio
 import logging
-import time
 from collections.abc import Iterator
 
-import httpx
-
 from miles.rollout.base_types import (
     RolloutFnConstructorInput,
     RolloutFnEvalInput,
     RolloutFnEvalOutput,
     RolloutFnInput,
     RolloutFnOutput,
+    RolloutFnTrainInput,
     RolloutFnTrainOutput,
 )
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
 from miles.rollout.inference_rollout.inference_rollout_common import GenerateState, generate_and_rm_group
 from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets
 from miles.rollout.submission_scheduler import make_submission_scheduler
-from miles.utils.http_utils import get
 from miles.utils.misc import load_function
 from miles.utils.types import Sample
 
 logger = logging.getLogger(__name__)
 
 OUTPUT_QUEUE_MAX_GROUPS = 1000
 NO_PROGRESS_WARN_SECS = 30.0
-WEIGHT_VERSION_QUERY_TIMEOUT_SECS = 2.0
 
 # A finished group is list[Sample], or list[list[Sample]] when a generate function
 # returns multiple samples per trajectory (e.g. multi-agent).
@@ -68,32 +64,6 @@ def group_oldest_weight_version(group: Group) -> int | None:
     return min(versions) if versions else None
 
 
-class _CachedWeightVersion:
-    """Throttled query of the current engine weight version via the router's /model_info."""
-
-    def __init__(self, ttl: float = 1.0):
-        self._ttl = ttl
-        self._value: int | None = None
-        self._last_query = float("-inf")
-
-    async def get(self, args) -> int | None:
-        # Throttles failures too: the drain queries once per group, and an unreachable
-        # router would otherwise cost every one of them the full timeout.
-        if (time.monotonic() - self._last_query) < self._ttl:
-            return self._value
-        url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/model_info"
-        try:
-            data = await asyncio.wait_for(get(url), timeout=WEIGHT_VERSION_QUERY_TIMEOUT_SECS)
-            self._value = int(data["weight_version"])
-        except (httpx.HTTPError, asyncio.TimeoutError) as e:
-            # Transient router unavailability; the staleness filter is best-effort.
-            logger.debug(f"Failed to query engine weight version: {e}")
-        finally:
-            # Stamped on completion, so a router slower than the TTL still gets throttled.
-            self._last_query = time.monotonic()
-        return self._value
-
-
 class FullyAsyncRolloutFn:
     """Continuous rollout generation decoupled from training steps.
 
@@ -111,7 +81,6 @@ def __init__(self, input: RolloutFnConstructorInput):
         self._scheduler = make_submission_scheduler(input.args, default="sample")
         self._dynamic_filter = load_function(input.args.dynamic_sampling_filter_path)
         self._sample_filter = load_function(input.args.rollout_sample_filter_path)
-        self._weight_version = _CachedWeightVersion()
         self._worker: asyncio.Task | None = None
         self._eval_prompt_dataset_cache: dict = {}
         self._producer_resumed = asyncio.Event()
@@ -125,7 +94,7 @@ async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput:
             self._output = asyncio.Queue(maxsize=OUTPUT_QUEUE_MAX_GROUPS)
             self._worker = asyncio.create_task(self._worker_loop())
             logger.info("Started fully-async rollout worker")
-        return await self._drain(input.rollout_id)
+        return await self._drain(input)
 
     async def _call_eval(self, input: RolloutFnEvalInput) -> RolloutFnOutput:
         if input.generate_state is not None:
@@ -208,7 +177,7 @@ async def _next_group(self) -> tuple[list[Sample], Group]:
             if not queue_get.done():
                 queue_get.cancel()
 
-    async def _drain(self, rollout_id: int) -> RolloutFnTrainOutput:
+    async def _drain(self, input: RolloutFnTrainInput) -> RolloutFnTrainOutput:
         args = self.args
         assert args.rollout_global_dataset
 
@@ -232,7 +201,7 @@ async def _drain(self, rollout_id: int) -> RolloutFnTrainOutput:
 
             if args.max_weight_staleness is not None:
                 oldest = group_oldest_weight_version(group)
-                current = await self._weight_version.get(args)
+                current = input.weight_version
                 if oldest is not None and current is not None:
                     staleness = current - oldest
                     staleness_values.append(staleness)

tests/fast/rollout/test_fully_async_rollout.py

@@ -7,7 +7,6 @@
 from collections import deque
 from dataclasses import replace
 
-import httpx
 import pytest
 
 import miles.rollout.fully_async_rollout as fully_async
@@ -220,13 +219,7 @@ def get_samples_with_fresh_versions(num_samples):
 
     fn = make_fn(monkeypatch, make_args(rollout_batch_size=1, max_weight_staleness=2), data_source)
 
-    class FakeWeightVersion:
-        async def get(self, args):
-            return 10
-
-    fn._weight_version = FakeWeightVersion()
-
-    output = await fn(RolloutFnTrainInput(rollout_id=0))
+    output = await fn(RolloutFnTrainInput(rollout_id=0, weight_version=10))
 
     assert data_source.recycled == [stale]
     assert output.metrics["rollout/fully_async/stale_groups_recycled"] == 1
@@ -362,27 +355,17 @@ def mark_first_of_each_group(args, data):
     assert [sample.remove_sample for sample in output.samples[0]] == [True, False]
 
 
-async def test_weight_version_throttles_failed_queries(monkeypatch):
-    """A drain queries once per group, so an unreachable router must not cost one timeout each."""
-    calls = []
-
-    async def unreachable_router(url):
-        calls.append(url)
-        raise httpx.ConnectError("router down")
-
-    monkeypatch.setattr(fully_async, "get", unreachable_router)
-    args = make_args()
+async def test_staleness_filter_off_before_the_first_weight_update(monkeypatch):
+    """weight_version is None until the trainer pushes weights; staleness is unknown, not zero."""
+    stale = make_group(1, weight_versions=["5"])
+    data_source = FakeDataSource(scripted=[stale])
+    fn = make_fn(monkeypatch, make_args(rollout_batch_size=1, max_weight_staleness=0), data_source)
 
-    throttled = fully_async._CachedWeightVersion(ttl=60.0)
-    assert await throttled.get(args) is None
-    assert await throttled.get(args) is None
-    assert len(calls) == 1
+    output = await fn(RolloutFnTrainInput(rollout_id=0))
 
-    calls.clear()
-    expired = fully_async._CachedWeightVersion(ttl=0.0)
-    assert await expired.get(args) is None
-    assert await expired.get(args) is None
-    assert len(calls) == 2
+    assert data_source.recycled == []
+    assert output.samples[0][0].group_index == 1
+    assert "rollout/fully_async/max_staleness" not in output.metrics
 
 
 async def test_worker_defaults_to_sample_granularity(monkeypatch):