perf(r3): collate routes through a container-aware pool, and refuse replay under VPP - #2078
perf(r3): collate routes through a container-aware pool, and refuse replay under VPP#2078dyurk-lila wants to merge 6 commits into
Conversation
Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteGenerateClient plus a RemoteGenerateResult dataclass, so routed-expert results can be obtained without constructing the full inference/control-plane client. RemoteInferenceClient now owns an internal RemoteGenerateClient and delegates session management, _post, and _generate_single to it. Endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are unchanged.
For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces incremental routed expert tracking and refactors the remote inference client by extracting a reusable RemoteGenerateClient class. It also adds a CPU topology utility to dynamically size thread pools based on cgroup quotas and process affinity, and parallelizes batch row filling. Feedback on these changes includes catching aiohttp.ClientConnectionError instead of specific exceptions to handle all transient connection failures during POST retries, explicitly resetting _generate_client to None during unpickling to prevent stale states, and catching OSError when querying CPU affinity to ensure compatibility with restricted environments.
| except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as exc: | ||
| last_exc = exc | ||
| logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}") | ||
| await asyncio.sleep(1) |
There was a problem hiding this comment.
Catching only ServerDisconnectedError and ClientOSError misses other common transient connection errors like ClientConnectorError (which is raised when a connection is refused, e.g., during a server restart or scale-up). Since ClientConnectorError inherits from ClientConnectionError but not ClientOSError, it will bypass this retry loop entirely. Catching aiohttp.ClientConnectionError instead ensures all connection-related issues are safely retried.
| except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as exc: | |
| last_exc = exc | |
| logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}") | |
| await asyncio.sleep(1) | |
| except aiohttp.ClientConnectionError as exc: | |
| last_exc = exc | |
| logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}") | |
| await asyncio.sleep(1) |
| def __setstate__(self, state: dict) -> None: | ||
| """Restore state after unpickling.""" | ||
| self.__dict__.update(state) | ||
| self._session = None | ||
| self._gen_sem = None | ||
| self._detok_sem = None | ||
| self._sem_loop = None |
There was a problem hiding this comment.
Explicitly reset _generate_client to None during unpickling to ensure a clean state.
| def __setstate__(self, state: dict) -> None: | |
| """Restore state after unpickling.""" | |
| self.__dict__.update(state) | |
| self._session = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None | |
| def __setstate__(self, state: dict) -> None: | |
| """Restore state after unpickling.""" | |
| self.__dict__.update(state) | |
| self._generate_client = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None |
| try: | ||
| affinity = len(os.sched_getaffinity(0)) | ||
| except AttributeError: | ||
| affinity = os.cpu_count() or 1 |
There was a problem hiding this comment.
In some restricted environments or container runtimes with strict seccomp profiles, calling os.sched_getaffinity(0) can raise an OSError (such as PermissionError). Catching both AttributeError and OSError ensures the fallback to os.cpu_count() is always safe.
| try: | |
| affinity = len(os.sched_getaffinity(0)) | |
| except AttributeError: | |
| affinity = os.cpu_count() or 1 | |
| try: | |
| affinity = len(os.sched_getaffinity(0)) | |
| except (AttributeError, OSError): | |
| affinity = os.cpu_count() or 1 |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit e080afe. Configure here.
| rollout_expert_indices_out = agent_loop_state.routed_expert_trace.finalize( | ||
| token_count=len(prompt_ids) + len(response_ids), | ||
| loss_mask=[0] * len(prompt_ids) + loss_mask, | ||
| ) |
There was a problem hiding this comment.
Synthetic EOS breaks route finalize
Medium Severity
With use_conversation_multi_turn=False, a synthetic EOS is appended with loss weight 1 after generation, then RoutedExpertTrace.finalize runs. That finalize rejects loss-active targets past the captured route prefix, so R3 agent loops raise instead of padding the uncaptured EOS the way make_router_padding_mask documents.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit e080afe. Configure here.
2 participants
Problem
Routed-expert replay arrives as one variable-length NumPy array per trajectory. Collating those arrays into the left-padded global training batch is dominated by allocating and first-touching a large route buffer. The previous serial path also rebuilt padding intermediates and rescanned route IDs even though the inference boundary had already selected a canonical integer dtype.
Separately, replay is unsafe with Megatron virtual pipeline parallelism. Each forward appends a microbatch to every local replay FIFO, while an interleaved model chunk consumes only its own entry during backward. The queues can therefore desynchronize without an immediate crash.
Implementation
The sender remains responsible for route compaction. Collation validates the accepted
uint8,int16, andint32dtypes and promotes a mixed batch to the widest input dtype without another value scan.Performance
The included packed-route collation benchmark compares identical outputs across serial and pooled implementations. On a 32-core host with
OMP_NUM_THREADS=1, a representative 40-layer global batch improved from 29.61 s to 3.35 s (median of five runs).The result depends on the controller running with one OpenMP thread, as it does under its one-CPU Ray allocation. With unrestricted library threads, the pool can oversubscribe the container and lose performance; the worker cap and cgroup-aware sizing make that dependency explicit.
Testing