From b98a2bf6ce660eb773064243cd556ddd5f267fa5 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 1 Jul 2026 02:55:53 +0000 Subject: [PATCH 01/29] Add infra-failure/truncation filters; allow zero loss_mask for masked samples Co-Authored-By: Claude Opus 4.8 --- miles/backends/megatron_utils/actor.py | 4 +- .../filter_hub/dynamic_sampling_filters.py | 43 +++++++++++++++++++ miles/rollout/filter_hub/rollout_filters.py | 27 +++++++++++- 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index c138ea87460..5bb4600c57a 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -411,7 +411,9 @@ def _basic_batch_summary(): mask_sum = _sum_float(mask) if mask_sum <= 0: - _add_error(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}") + # Warning-only: rollout sample filters zero a sample's loss_mask + # (remove_sample) to drop it from the gradient; the reducer clamps. + _add_warning(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}") if mask_sum > resp: # Warning-only: float/weighted masks can legitimately have sum > resp. _add_warning( diff --git a/miles/rollout/filter_hub/dynamic_sampling_filters.py b/miles/rollout/filter_hub/dynamic_sampling_filters.py index f39ff8c9c13..0c0c6233ae5 100644 --- a/miles/rollout/filter_hub/dynamic_sampling_filters.py +++ b/miles/rollout/filter_hub/dynamic_sampling_filters.py @@ -6,10 +6,34 @@ __all__ = [ "check_reward_nonzero_std", "check_no_aborted", + "check_no_infra_failures", "drop_zero_std_groups_and_extreme_pass_rate", "drop_truncated_or_extreme_pass_rate", ] +# Harbor exit_status values (from _extract_exit_status, after _TIMEOUT_EXCEPTION_MAP) +# for infra / non-policy failures: the rollout was severed by the environment, +# engine, or verifier rather than the policy, so it isn't valid training signal. +INFRA_FAILURE_EXIT_STATUSES = frozenset( + { + "AgentTimeout", + "AgentTimeoutError", + "HealthcheckError", + "_K8sInternalInfraError", + "Cancelled", + "RewardFileNotFoundError", + "AgentSetupTimeout", + "AgentSetupTimeoutError", + "SqsConsumerError", + "VerifierTimeout", + "VerifierTimeoutError", + "EnvStartTimeout", + "EnvironmentStartTimeoutError", + "TimeoutError", + "AddTestsDirError", + } +) + def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): rewards = [sample.get_reward_value(args) for sample in samples] @@ -46,6 +70,25 @@ def check_no_aborted(args, samples: list[Sample], **kwargs): return DynamicFilterOutput(keep=True) +def check_no_infra_failures(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput: + """Reject the group if any rollout aborted or failed for an infra reason. + + Superset of ``check_no_aborted``: also drops samples whose Harbor + ``exit_status`` is an infra / non-policy failure. Keyed on ``exit_status`` + because such rollouts may still record COMPLETED turns; aborted samples + (zero records, no exit_status) are caught by the status check. + + --dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_infra_failures + """ + for sample in _flatten_samples(samples): + if sample.status == Sample.Status.ABORTED: + return DynamicFilterOutput(keep=False, reason="group_has_aborted") + exit_status = (sample.metadata or {}).get("exit_status", "") + if exit_status in INFRA_FAILURE_EXIT_STATUSES: + return DynamicFilterOutput(keep=False, reason=f"group_has_{exit_status}") + return DynamicFilterOutput(keep=True) + + def drop_zero_std_groups_and_extreme_pass_rate(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput: """Filter groups with near-zero reward std or extreme mean rewards. For 0/1 rewards, mean reward is equivalent to pass rate --- so this function can be used to filter diff --git a/miles/rollout/filter_hub/rollout_filters.py b/miles/rollout/filter_hub/rollout_filters.py index 72282d4c64b..30be747d712 100644 --- a/miles/rollout/filter_hub/rollout_filters.py +++ b/miles/rollout/filter_hub/rollout_filters.py @@ -1,7 +1,32 @@ from miles.rollout.filter_hub.dynamic_sampling_filters import _flatten_samples from miles.utils.types import Sample -__all__ = ["mask_truncated", "mask_truncated_and_llm_judge_failed"] +__all__ = ["mask_truncated", "mask_truncated_and_llm_judge_failed", "mask_token_truncated"] + +# Harbor exit_status values for rollouts truncated by running out of output +# tokens or context. Real attempts (verifier reward is meaningful) but the +# trajectory shouldn't get gradient — masked, reward kept in the group baseline. +TRUNCATION_EXIT_STATUSES = frozenset( + { + "BadRequestError", + "ContextWindowExceededError", + "OutputLengthExceededError", + } +) + + +def mask_token_truncated(args, samples: list[Sample]) -> None: + """Mask token/context-truncated samples (zero gradient, reward kept in baseline). + + Keyed on Harbor ``exit_status`` (BadRequest/ContextWindow record COMPLETED + turns, so a status check misses them); also covers TRUNCATED status. + + --rollout-sample-filter-path miles.rollout.filter_hub.rollout_filters.mask_token_truncated + """ + for sample in _flatten_samples(samples): + exit_status = (sample.metadata or {}).get("exit_status", "") + if exit_status in TRUNCATION_EXIT_STATUSES or sample.status == Sample.Status.TRUNCATED: + sample.remove_sample = True def mask_truncated(args, samples: list[Sample]) -> None: From b0b6638c2afdd53706b3b869d8ad5c4a9981db9d Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 1 Jul 2026 18:00:41 +0000 Subject: [PATCH 02/29] Sample tito_session_mismatch instead of computing it on every record fetch Re-tokenizing the full trajectory on the session-server event loop blocks record retrieval at long context; gate it behind --tito-session-mismatch-sample-rate (default off). Co-Authored-By: Claude Opus 4.8 (cherry picked from commit 759ca3048468adad0f498c764efa9d76624bbb6e) --- miles/rollout/session/sessions.py | 22 +++++++++++++++------- miles/utils/arguments.py | 11 +++++++++++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index ae6c8663615..bc7ae7889fd 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -2,6 +2,7 @@ import logging import time import uuid +import zlib import orjson from fastapi import Request @@ -112,13 +113,20 @@ async def get_session(session_id: str): raise SessionNotFoundError(f"session not found: session_id={session_id}") return GetSessionResponse(session_id=session_id, records=[], metadata={}) metadata = {} - try: - mismatch = registry.compute_session_mismatch(session) - except TokenizationError: - logger.exception("Failed to compute tito_session_mismatch for session %s", session_id) - mismatch = None - if mismatch is not None: - metadata["tito_session_mismatch"] = mismatch + # Re-tokenizing the full trajectory to check tito_session_mismatch is a + # diagnostic canary, not a training-correctness guard (train tokens come + # from the inference response). It runs synchronously here, so at long + # context it blocks the event loop and starves record retrieval. Sample + # a deterministic fraction by session_id to keep the signal cheap. + sample_rate = getattr(args, "tito_session_mismatch_sample_rate", 0.0) + if sample_rate > 0.0 and (zlib.crc32(session_id.encode()) % 10000) < sample_rate * 10000: + try: + mismatch = registry.compute_session_mismatch(session) + except TokenizationError: + logger.exception("Failed to compute tito_session_mismatch for session %s", session_id) + mismatch = None + if mismatch is not None: + metadata["tito_session_mismatch"] = mismatch metadata["accumulated_token_ids"] = session.token_ids metadata["max_trim_tokens"] = registry.tito_tokenizer.max_trim_tokens return GetSessionResponse( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2d6e43e1384..9d696bdb3a0 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1623,6 +1623,17 @@ def add_session_arguments(parser): choices=[t.value for t in TITOTokenizerType], help="TITO tokenizer type for pretokenized prefix reuse. Controls how token IDs are computed for messages appended after the pretokenized prefix in multi-turn agentic sessions.", ) + parser.add_argument( + "--tito-session-mismatch-sample-rate", + type=float, + default=0.0, + help="Fraction of sessions (0.0-1.0) for which the session server re-tokenizes the full " + "trajectory to compute tito_session_mismatch. This is a diagnostic canary for chat-template/" + "tokenizer drift; it does not affect training tokens (those come from the inference response). " + "It runs synchronously on the record-fetch path, so at long context it blocks the event loop and " + "starves record retrieval. Default 0.0 (off). Set to 0.02 for a cheap drift signal, 1.0 for legacy " + "always-on behaviour.", + ) parser.add_argument( "--tito-allowed-append-roles", nargs="+", From db5d476a50e971d0f2de9fd0ea17447820997b2b Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 1 Jul 2026 23:55:12 +0000 Subject: [PATCH 03/29] Parse HTTP responses with orjson, stdlib fallback Co-Authored-By: Claude Opus 4.8 --- miles/utils/http_utils.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 0aaf7926594..0c9c7736615 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -9,6 +9,7 @@ import time import httpx +import orjson logger = logging.getLogger(__name__) @@ -196,9 +197,16 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non response = await getattr(client, action)(url, json=payload or {}, headers=headers) response.raise_for_status() try: - output = response.json() - except json.JSONDecodeError: - output = response.text + # orjson is ~10x faster than stdlib json on the multi-MB + # token-id payloads from session-record GETs; at rollout scale + # stdlib parse time alone saturates the single-threaded event + # loop. Fall back to stdlib for JSON orjson rejects (NaN/Inf). + output = orjson.loads(response.content) + except orjson.JSONDecodeError: + try: + output = response.json() + except json.JSONDecodeError: + output = response.text except Exception as e: retry_count += 1 From 4d8c7fc1975c895dcf9d2eb875501d76f97249d3 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Thu, 2 Jul 2026 03:23:55 +0000 Subject: [PATCH 04/29] Log session-server stats heartbeat at debug level Co-Authored-By: Claude Opus 4.8 --- miles/rollout/session/session_server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/rollout/session/session_server.py b/miles/rollout/session/session_server.py index 07b90417cae..4535434991f 100644 --- a/miles/rollout/session/session_server.py +++ b/miles/rollout/session/session_server.py @@ -136,7 +136,7 @@ async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): stats = get_worker_stats(worker_port) if stats is None: # Routes not wired yet (very early startup) — emit a sparse log. - logger.info( + logger.debug( "[session-server] stats worker_port=%s reqs_total=0 reqs_since_last=0 " "inflight_now=0 turns_completed=0", worker_port, @@ -150,7 +150,7 @@ async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): mi = _proc.memory_info() rss_mb = mi.rss / 1024.0 / 1024.0 vms_mb = mi.vms / 1024.0 / 1024.0 - logger.info( + logger.debug( "[session-server] stats worker_port=%s reqs_total=%d reqs_since_last=%d " "inflight_now=%d turns_completed=%d rss_mb=%.0f vms_mb=%.0f", worker_port, From 0bd6425bd710266d382d484c2c2745ca8ffdf25a Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 11:31:11 -0700 Subject: [PATCH 05/29] fixes --- .../generate_utils/openai_endpoint_utils.py | 458 +++++++++++++++--- 1 file changed, 403 insertions(+), 55 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index e03197d6a5c..d5c3c357c6a 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -3,59 +3,322 @@ """ import asyncio +import json import logging import random from argparse import Namespace from copy import deepcopy +from typing import Any + +import httpx from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response from miles.rollout.session.session_types import GetSessionResponse, SessionRecord -from miles.utils.http_utils import post from miles.utils.types import Sample logger = logging.getLogger(__name__) -_SESSION_REQUEST_TIMEOUT = 120 +_SESSION_REQUEST_TIMEOUT = 120.0 + +_HTTP_CONNECT_TIMEOUT = 10.0 +_HTTP_READ_TIMEOUT = 120.0 +_HTTP_WRITE_TIMEOUT = 30.0 +_HTTP_POOL_TIMEOUT = 10.0 + +_HEALTH_RETRIES = 2 +_CREATE_RETRIES = 10 +_COLLECT_RETRIES = 3 +_DELETE_RETRIES = 3 + +_BACKOFF_INITIAL_SECONDS = 1.0 +_BACKOFF_MAX_SECONDS = 10.0 class OpenAIEndpointTracer: - def __init__(self, router_url: str, session_id: str, session_server_instance_id: str | None = None): - self.router_url = router_url + def __init__( + self, + router_url: str, + session_id: str, + session_server_instance_id: str | None = None, + ): + self.router_url = router_url.rstrip("/") self.session_id = session_id - self.base_url = f"{router_url}/sessions/{session_id}" + self.base_url = f"{self.router_url}/sessions/{session_id}" self.session_server_instance_id = session_server_instance_id + @staticmethod + def _timeout() -> httpx.Timeout: + return httpx.Timeout( + timeout=_SESSION_REQUEST_TIMEOUT, + connect=_HTTP_CONNECT_TIMEOUT, + read=_HTTP_READ_TIMEOUT, + write=_HTTP_WRITE_TIMEOUT, + pool=_HTTP_POOL_TIMEOUT, + ) + + @staticmethod + def _response_size(response: httpx.Response) -> int: + try: + return len(response.content) + except Exception: + return -1 + + @staticmethod + def _backoff_seconds(attempt: int) -> float: + # attempt is 1-indexed. + delay = _BACKOFF_INITIAL_SECONDS * (2 ** (attempt - 1)) + return min(delay, _BACKOFF_MAX_SECONDS) + + @classmethod + async def _request( + cls, + method: str, + url: str, + *, + phase: str, + payload: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + max_retries: int = 3, + expect_json: bool = True, + ) -> Any: + method = method.upper() + last_exc: BaseException | None = None + + for attempt in range(1, max_retries + 1): + try: + logger.info( + "[session-client] request_start phase=%s method=%s url=%s attempt=%d/%d", + phase, + method, + url, + attempt, + max_retries, + ) + + async with httpx.AsyncClient(timeout=cls._timeout()) as client: + if method in {"GET", "DELETE"}: + response = await client.request(method, url, headers=headers) + else: + response = await client.request( + method, + url, + json=payload or {}, + headers=headers, + ) + + body_bytes = cls._response_size(response) + + logger.info( + "[session-client] response_received phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + phase, + method, + url, + response.status_code, + body_bytes, + attempt, + max_retries, + ) + + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + last_exc = exc + + logger.info( + "[session-client] http_status_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + phase, + method, + url, + status, + body_bytes, + attempt, + max_retries, + ) + + # Do not retry deterministic client errors. + # Retry 429 and 5xx because they can be transient. + if status != 429 and status < 500: + raise + + else: + if response.status_code == 204 or not response.content: + return None + + if not expect_json: + return response.text + + try: + return response.json() + except json.JSONDecodeError as exc: + logger.info( + "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + phase, + method, + url, + response.status_code, + body_bytes, + attempt, + max_retries, + ) + raise exc + + except asyncio.CancelledError: + logger.info( + "[session-client] request_cancelled phase=%s method=%s url=%s attempt=%d/%d", + phase, + method, + url, + attempt, + max_retries, + ) + raise + + except httpx.RemoteProtocolError as exc: + last_exc = exc + logger.info( + "[session-client] remote_protocol_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) + + except httpx.TimeoutException as exc: + last_exc = exc + logger.info( + "[session-client] timeout phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) + + except httpx.TransportError as exc: + last_exc = exc + logger.info( + "[session-client] transport_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) + + except Exception as exc: + logger.info( + "[session-client] unexpected_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) + raise + + if attempt < max_retries: + delay = cls._backoff_seconds(attempt) + logger.info( + "[session-client] request_retry_sleep phase=%s method=%s url=%s attempt=%d/%d sleep_s=%.1f", + phase, + method, + url, + attempt, + max_retries, + delay, + ) + await asyncio.sleep(delay) + + logger.info( + "[session-client] request_failed phase=%s method=%s url=%s attempts=%d last_error_type=%s last_error=%r", + phase, + method, + url, + max_retries, + type(last_exc).__name__ if last_exc is not None else None, + last_exc, + ) + + if last_exc is not None: + raise last_exc + + raise RuntimeError(f"request failed without exception: phase={phase} method={method} url={url}") + @staticmethod async def create(args: Namespace): - # Multi-backend dispatch (callee-side random pick): if the spawner - # published ``session_server_backends`` (N>=1 vanilla session-server - # processes co-located with the SGLang gateway), pick one uniformly - # at random and bind this trajectory to it for its lifetime. Falls - # back to the legacy single-server fields when no backend list is - # published. backends = getattr(args, "session_server_backends", None) if backends: - session_url = random.choice(backends) + session_url = random.choice(backends).rstrip("/") else: session_ip = getattr(args, "session_server_ip", None) session_port = getattr(args, "session_server_port", None) if not session_ip or not session_port: raise RuntimeError( - "session_server_ip/session_server_port are not set. " - "Pass --use-session-server to start the session server." + "session_server_ip/session_server_port are not set. Pass --use-session-server to start the session server." ) session_url = f"http://{session_ip}:{session_port}" + + logger.info("[session-client] create_start session_url=%s", session_url) + session_server_instance_id = None + try: - health = await post(f"{session_url}/health", {}, action="get") + health = await OpenAIEndpointTracer._request( + "GET", + f"{session_url}/health", + phase="health", + max_retries=_HEALTH_RETRIES, + ) + if isinstance(health, dict): session_server_instance_id = health.get("session_server_instance_id") if session_server_instance_id is not None: args.session_server_instance_id = session_server_instance_id - except Exception as e: - logger.warning("Failed to get session server health from %s: %s", session_url, e) - response = await post(f"{session_url}/sessions", {}, action="post") + + logger.info( + "[session-client] health_ok session_url=%s instance_id=%s", + session_url, + session_server_instance_id, + ) + + except Exception as exc: + logger.info( + "[session-client] health_failed session_url=%s error_type=%s error=%r", + session_url, + type(exc).__name__, + exc, + ) + + response = await OpenAIEndpointTracer._request( + "POST", + f"{session_url}/sessions", + phase="create_session", + payload={}, + max_retries=_CREATE_RETRIES, + ) + + if not isinstance(response, dict) or "session_id" not in response: + raise RuntimeError(f"invalid create session response from {session_url}: {response!r}") + session_id = response["session_id"] + + logger.info( + "[session-client] create_done session_url=%s session_id=%s instance_id=%s", + session_url, + session_id, + session_server_instance_id, + ) + return OpenAIEndpointTracer( router_url=session_url, session_id=session_id, @@ -63,41 +326,130 @@ async def create(args: Namespace): ) async def collect_records(self) -> tuple[list[SessionRecord], dict]: + records: list[SessionRecord] = [] + metadata: dict = {} + + logger.info( + "[session-client] collect_start session_id=%s url=%s", + self.session_id, + self.base_url, + ) + try: - response = await asyncio.wait_for( - post(self.base_url, {}, action="get"), - timeout=_SESSION_REQUEST_TIMEOUT, + response = await self._request( + "GET", + self.base_url, + phase="collect_records", + max_retries=_COLLECT_RETRIES, ) - except asyncio.TimeoutError: - logger.error( - f"Timed out waiting for session {self.session_id} records after {_SESSION_REQUEST_TIMEOUT}s " - f"(likely stale HTTP keepalive connection). Returning empty records." + + parsed = GetSessionResponse.model_validate(response) + records = parsed.records or [] + metadata = parsed.metadata or {} + + logger.info( + "[session-client] collect_done session_id=%s records=%d metadata_keys=%s", + self.session_id, + len(records), + sorted(metadata.keys()), + ) + + return records, metadata + + except httpx.RemoteProtocolError as exc: + logger.info( + "[session-client] collect_failed_remote_protocol session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, ) - # Still attempt to clean up the session. - try: - await asyncio.wait_for( - post(self.base_url, {}, action="delete"), - timeout=_SESSION_REQUEST_TIMEOUT, - ) - except Exception: - logger.warning(f"Failed to delete session {self.session_id} after timeout") return [], {} - except Exception as e: - logger.warning(f"Failed to get session {self.session_id} records: {e}") - raise - response = GetSessionResponse.model_validate(response) - records = response.records - metadata = response.metadata + + except httpx.TimeoutException as exc: + logger.info( + "[session-client] collect_failed_timeout session_id=%s url=%s timeout_s=%.1f error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + _SESSION_REQUEST_TIMEOUT, + type(exc).__name__, + exc, + ) + return [], {} + + except httpx.TransportError as exc: + logger.info( + "[session-client] collect_failed_transport session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} + + except Exception as exc: + logger.info( + "[session-client] collect_failed_unexpected session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} + + finally: + await self.delete_session() + + async def delete_session(self) -> None: + logger.info( + "[session-client] delete_start session_id=%s url=%s", + self.session_id, + self.base_url, + ) try: - await asyncio.wait_for( - post(self.base_url, {}, action="delete"), - timeout=_SESSION_REQUEST_TIMEOUT, + await self._request( + "DELETE", + self.base_url, + phase="delete_session", + max_retries=_DELETE_RETRIES, + expect_json=False, ) - except Exception as e: - logger.warning(f"Failed to delete session {self.session_id} after collecting records: {e}") - return (records or []), metadata + logger.info( + "[session-client] delete_done session_id=%s url=%s", + self.session_id, + self.base_url, + ) + + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + + if status == 404: + logger.info( + "[session-client] delete_not_found session_id=%s url=%s", + self.session_id, + self.base_url, + ) + return + + logger.info( + "[session-client] delete_failed_status session_id=%s url=%s status=%d error_type=%s error=%r", + self.session_id, + self.base_url, + status, + type(exc).__name__, + exc, + ) + + except Exception as exc: + logger.info( + "[session-client] delete_failed session_id=%s url=%s error_type=%s error=%r", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) def compute_samples_from_openai_records( @@ -148,12 +500,9 @@ def compute_samples_from_openai_records( # the next message delimiter) — strip them from the sample. trim_count = len(output_ids) - matched allowed = 0 if is_last else max_trim_tokens - assert trim_count <= allowed, ( - f"trim_count {trim_count} exceeds allowed={allowed} " - f"(is_last={is_last}, max_trim_tokens={max_trim_tokens}); " - f"output_ids[-3:]={output_ids[-3:]}, " - f"accumulated[cursor:cursor+3]={accumulated_token_ids[cursor:cursor+3]}" - ) + assert ( + trim_count <= allowed + ), f"trim_count {trim_count} exceeds allowed={allowed} (is_last={is_last}, max_trim_tokens={max_trim_tokens}); output_ids[-3:]={output_ids[-3:]}, accumulated[cursor:cursor+3]={accumulated_token_ids[cursor : cursor + 3]}" # Step 4: advance cursor past matched output to the next turn cursor += matched @@ -163,10 +512,9 @@ def compute_samples_from_openai_records( if accumulated_token_ids is not None: # Step 5: verify the entire accumulated sequence was consumed - assert cursor == len(accumulated_token_ids), ( - f"cursor {cursor} != len(accumulated_token_ids) {len(accumulated_token_ids)} " - f"after processing all {len(records)} records" - ) + assert cursor == len( + accumulated_token_ids + ), f"cursor {cursor} != len(accumulated_token_ids) {len(accumulated_token_ids)} after processing all {len(records)} records" return samples From b7d36982d08d4500a2d5020b0eeaa556f298e508 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 15:27:42 -0700 Subject: [PATCH 06/29] add jitter --- .../generate_utils/openai_endpoint_utils.py | 248 +++++++++--------- 1 file changed, 125 insertions(+), 123 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index d5c3c357c6a..32c310e9877 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -32,6 +32,7 @@ _BACKOFF_INITIAL_SECONDS = 1.0 _BACKOFF_MAX_SECONDS = 10.0 +_BACKOFF_JITTER_FRACTION = 0.2 class OpenAIEndpointTracer: @@ -66,8 +67,14 @@ def _response_size(response: httpx.Response) -> int: @staticmethod def _backoff_seconds(attempt: int) -> float: # attempt is 1-indexed. - delay = _BACKOFF_INITIAL_SECONDS * (2 ** (attempt - 1)) - return min(delay, _BACKOFF_MAX_SECONDS) + base_delay = _BACKOFF_INITIAL_SECONDS * (2 ** (attempt - 1)) + base_delay = min(base_delay, _BACKOFF_MAX_SECONDS) + + jitter = random.uniform( + 1.0 - _BACKOFF_JITTER_FRACTION, + 1.0 + _BACKOFF_JITTER_FRACTION, + ) + return min(base_delay * jitter, _BACKOFF_MAX_SECONDS) @classmethod async def _request( @@ -84,18 +91,18 @@ async def _request( method = method.upper() last_exc: BaseException | None = None - for attempt in range(1, max_retries + 1): - try: - logger.info( - "[session-client] request_start phase=%s method=%s url=%s attempt=%d/%d", - phase, - method, - url, - attempt, - max_retries, - ) + async with httpx.AsyncClient(timeout=cls._timeout()) as client: + for attempt in range(1, max_retries + 1): + try: + logger.info( + "[session-client] request_start phase=%s method=%s url=%s attempt=%d/%d", + phase, + method, + url, + attempt, + max_retries, + ) - async with httpx.AsyncClient(timeout=cls._timeout()) as client: if method in {"GET", "DELETE"}: response = await client.request(method, url, headers=headers) else: @@ -106,138 +113,136 @@ async def _request( headers=headers, ) - body_bytes = cls._response_size(response) - - logger.info( - "[session-client] response_received phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", - phase, - method, - url, - response.status_code, - body_bytes, - attempt, - max_retries, - ) - - try: - response.raise_for_status() - except httpx.HTTPStatusError as exc: - status = exc.response.status_code - last_exc = exc + body_bytes = cls._response_size(response) logger.info( - "[session-client] http_status_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + "[session-client] response_received phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", phase, method, url, - status, + response.status_code, body_bytes, attempt, max_retries, ) - # Do not retry deterministic client errors. - # Retry 429 and 5xx because they can be transient. - if status != 429 and status < 500: - raise - - else: - if response.status_code == 204 or not response.content: - return None - - if not expect_json: - return response.text - try: - return response.json() - except json.JSONDecodeError as exc: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + last_exc = exc + logger.info( - "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + "[session-client] http_status_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", phase, method, url, - response.status_code, + status, body_bytes, attempt, max_retries, ) - raise exc - except asyncio.CancelledError: - logger.info( - "[session-client] request_cancelled phase=%s method=%s url=%s attempt=%d/%d", - phase, - method, - url, - attempt, - max_retries, - ) - raise + if status != 429 and status < 500: + raise - except httpx.RemoteProtocolError as exc: - last_exc = exc - logger.info( - "[session-client] remote_protocol_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", - phase, - method, - url, - attempt, - max_retries, - type(exc).__name__, - exc, - ) + else: + if response.status_code == 204 or not response.content: + return None + + if not expect_json: + return response.text + + try: + return response.json() + except json.JSONDecodeError as exc: + logger.info( + "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + phase, + method, + url, + response.status_code, + body_bytes, + attempt, + max_retries, + ) + raise exc + + except asyncio.CancelledError: + logger.info( + "[session-client] request_cancelled phase=%s method=%s url=%s attempt=%d/%d", + phase, + method, + url, + attempt, + max_retries, + ) + raise - except httpx.TimeoutException as exc: - last_exc = exc - logger.info( - "[session-client] timeout phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", - phase, - method, - url, - attempt, - max_retries, - type(exc).__name__, - exc, - ) + except httpx.RemoteProtocolError as exc: + last_exc = exc + logger.info( + "[session-client] remote_protocol_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) - except httpx.TransportError as exc: - last_exc = exc - logger.info( - "[session-client] transport_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", - phase, - method, - url, - attempt, - max_retries, - type(exc).__name__, - exc, - ) + except httpx.TimeoutException as exc: + last_exc = exc + logger.info( + "[session-client] timeout phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) - except Exception as exc: - logger.info( - "[session-client] unexpected_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", - phase, - method, - url, - attempt, - max_retries, - type(exc).__name__, - exc, - ) - raise + except httpx.TransportError as exc: + last_exc = exc + logger.info( + "[session-client] transport_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) - if attempt < max_retries: - delay = cls._backoff_seconds(attempt) - logger.info( - "[session-client] request_retry_sleep phase=%s method=%s url=%s attempt=%d/%d sleep_s=%.1f", - phase, - method, - url, - attempt, - max_retries, - delay, - ) - await asyncio.sleep(delay) + except Exception as exc: + logger.info( + "[session-client] unexpected_error phase=%s method=%s url=%s attempt=%d/%d error_type=%s error=%r", + phase, + method, + url, + attempt, + max_retries, + type(exc).__name__, + exc, + ) + raise + + if attempt < max_retries: + delay = cls._backoff_seconds(attempt) + logger.info( + "[session-client] request_retry_sleep phase=%s method=%s url=%s attempt=%d/%d sleep_s=%.3f", + phase, + method, + url, + attempt, + max_retries, + delay, + ) + await asyncio.sleep(delay) logger.info( "[session-client] request_failed phase=%s method=%s url=%s attempts=%d last_error_type=%s last_error=%r", @@ -326,9 +331,6 @@ async def create(args: Namespace): ) async def collect_records(self) -> tuple[list[SessionRecord], dict]: - records: list[SessionRecord] = [] - metadata: dict = {} - logger.info( "[session-client] collect_start session_id=%s url=%s", self.session_id, From 2943912250bbd3fa5341b012d2c289d786a3a1a1 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 16:02:44 -0700 Subject: [PATCH 07/29] add _compact_session_record, save on bytes --- miles/rollout/session/sessions.py | 97 +++++++++++++++++++++++++------ 1 file changed, 80 insertions(+), 17 deletions(-) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index aee310ad567..4ca6abb324e 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -9,12 +9,7 @@ from starlette.responses import Response from miles.rollout.session.linear_trajectory import SessionRegistry -from miles.rollout.session.session_errors import ( - SessionError, - SessionNotFoundError, - TokenizationError, - UpstreamResponseError, -) +from miles.rollout.session.session_errors import SessionError, SessionNotFoundError, UpstreamResponseError from miles.rollout.session.session_types import GetSessionResponse, SessionRecord from miles.utils.chat_template_utils import get_tito_tokenizer from miles.utils.processing_utils import load_tokenizer @@ -104,6 +99,78 @@ async def create_session(): session_id = registry.create_session() return {"session_id": session_id} + def _compact_output_token_logprobs(output_token_logprobs): + """Keep only the fields consumed by the sample builder. + + Consumer uses: + output_log_probs = [item[0] for item in output_token_logprobs] + output_token_ids = [item[1] for item in output_token_logprobs] + + Some backends may include extra per-token fields; do not return them. + """ + if not output_token_logprobs: + return [] + + compact = [] + for item in output_token_logprobs: + compact.append([item[0], item[1]]) + return compact + + def _compact_session_record(record: SessionRecord) -> SessionRecord: + """Return only the record fields required by compute_samples_from_openai_records.""" + choice = record.response.get("choices", [{}])[0] + meta_info = choice.get("meta_info") or {} + + compact_meta_info = { + "output_token_logprobs": _compact_output_token_logprobs(meta_info.get("output_token_logprobs")), + } + + # Used by _compute_sample_from_openai_record. + if "weight_version" in meta_info: + compact_meta_info["weight_version"] = meta_info["weight_version"] + + # Preserve prefix-cache metadata if present. The consumer does: + # sample.prefix_cache_info.add(choice.get("meta_info", {})) + for key in ( + "cached_tokens", + "prefix_cache_hit", + "prefix_cache_len", + "prefix_cache_tokens", + "prefix_cache_start_len", + "prefix_cache_end_len", + ): + if key in meta_info: + compact_meta_info[key] = meta_info[key] + + # Preserve routed experts if present. get_rollout_topk_from_response() + # may read this from choice or meta_info depending on implementation. + if "routed_experts" in meta_info: + compact_meta_info["routed_experts"] = meta_info["routed_experts"] + + compact_choice = { + "prompt_token_ids": choice.get("prompt_token_ids", []), + "finish_reason": choice.get("finish_reason"), + "meta_info": compact_meta_info, + } + + if "routed_experts" in choice: + compact_choice["routed_experts"] = choice["routed_experts"] + + return SessionRecord( + timestamp=record.timestamp, + method=record.method, + path=record.path, + status_code=record.status_code, + request={ + # Only used for this assertion: + # request_input_ids == prompt_token_ids + "input_ids": record.request.get("input_ids"), + }, + response={ + "choices": [compact_choice], + }, + ) + @app.get("/sessions/{session_id}") async def get_session(session_id: str): session = registry.sessions.get(session_id) @@ -111,19 +178,15 @@ async def get_session(session_id: str): if registry.is_deleted(session_id): raise SessionNotFoundError(f"session not found: session_id={session_id}") return GetSessionResponse(session_id=session_id, records=[], metadata={}) - metadata = {} - try: - mismatch = registry.compute_session_mismatch(session) - except TokenizationError: - logger.exception("Failed to compute tito_session_mismatch for session %s", session_id) - mismatch = None - if mismatch is not None: - metadata["tito_session_mismatch"] = mismatch - metadata["accumulated_token_ids"] = session.token_ids - metadata["max_trim_tokens"] = registry.tito_tokenizer.max_trim_tokens + + metadata = { + "accumulated_token_ids": session.token_ids, + "max_trim_tokens": registry.tito_tokenizer.max_trim_tokens, + } + return GetSessionResponse( session_id=session_id, - records=session.records, + records=[_compact_session_record(record) for record in session.records], metadata=metadata, ) From 082d11afedac8639a6f9c25debe79db82f5ec708 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Thu, 2 Jul 2026 23:03:50 +0000 Subject: [PATCH 08/29] 24h session-client timeouts; lifecycle/heartbeat logs at INFO --- .../generate_utils/openai_endpoint_utils.py | 4 ++-- miles/rollout/session/session_server.py | 4 ++-- miles/rollout/session/sessions.py | 15 +++++++++++---- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index d5c3c357c6a..578ea5ddcd0 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -18,10 +18,10 @@ logger = logging.getLogger(__name__) -_SESSION_REQUEST_TIMEOUT = 120.0 +_SESSION_REQUEST_TIMEOUT = 24 * 60 * 60.0 _HTTP_CONNECT_TIMEOUT = 10.0 -_HTTP_READ_TIMEOUT = 120.0 +_HTTP_READ_TIMEOUT = 24 * 60 * 60.0 _HTTP_WRITE_TIMEOUT = 30.0 _HTTP_POOL_TIMEOUT = 10.0 diff --git a/miles/rollout/session/session_server.py b/miles/rollout/session/session_server.py index cd564ca4036..6e48c10b354 100644 --- a/miles/rollout/session/session_server.py +++ b/miles/rollout/session/session_server.py @@ -131,7 +131,7 @@ async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): stats = get_worker_stats(worker_port) if stats is None: # Routes not wired yet (very early startup) — emit a sparse log. - logger.debug( + logger.info( "[session-server] stats worker_port=%s reqs_total=0 reqs_since_last=0 " "inflight_now=0 turns_completed=0", worker_port, @@ -145,7 +145,7 @@ async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): mi = _proc.memory_info() rss_mb = mi.rss / 1024.0 / 1024.0 vms_mb = mi.vms / 1024.0 / 1024.0 - logger.debug( + logger.info( "[session-server] stats worker_port=%s reqs_total=%d reqs_since_last=%d " "inflight_now=%d turns_completed=%d rss_mb=%.0f vms_mb=%.0f", worker_port, diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 9e7152864bc..59fcc8dc8e7 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -85,14 +85,21 @@ async def health(): async def debug_request_logger(request: Request, call_next): client = request.client client_info = f"{client.host}:{client.port}" if client else "unknown" - logger.debug( - f"[session-server] REQUEST ARRIVED: {request.method} {request.url.path} from={client_info} inflight_chat={_inflight_chat['count']}" + # Session lifecycle requests (create/get/delete, ~2 per trajectory) log + # at INFO so per-request latency is visible without enabling DEBUG; + # per-turn chat traffic and health probes stay at DEBUG for volume. + path = request.url.path + level = logging.DEBUG if path.endswith("/chat/completions") or path == "/health" else logging.INFO + logger.log( + level, + f"[session-server] REQUEST ARRIVED: {request.method} {path} from={client_info} inflight_chat={_inflight_chat['count']}", ) t0 = time.time() response = await call_next(request) elapsed = time.time() - t0 - logger.debug( - f"[session-server] REQUEST DONE: {request.method} {request.url.path} status={response.status_code} elapsed={elapsed:.3f}s from={client_info}" + logger.log( + level, + f"[session-server] REQUEST DONE: {request.method} {path} status={response.status_code} elapsed={elapsed:.3f}s from={client_info}", ) return response From d71934cf6207584b0131b937c88977da0d85be16 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 17:12:48 -0700 Subject: [PATCH 09/29] concurrency --- .../generate_utils/openai_endpoint_utils.py | 136 +++++++++++------- 1 file changed, 81 insertions(+), 55 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 32c310e9877..7651e5a2f26 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -34,6 +34,9 @@ _BACKOFF_MAX_SECONDS = 10.0 _BACKOFF_JITTER_FRACTION = 0.2 +_COLLECT_RECORDS_CONCURRENCY = 256 +_COLLECT_RECORDS_SEMAPHORE = asyncio.Semaphore(_COLLECT_RECORDS_CONCURRENCY) + class OpenAIEndpointTracer: def __init__( @@ -332,75 +335,98 @@ async def create(args: Namespace): async def collect_records(self) -> tuple[list[SessionRecord], dict]: logger.info( - "[session-client] collect_start session_id=%s url=%s", + "[session-client] collect_wait session_id=%s url=%s concurrency_limit=%d", self.session_id, self.base_url, + _COLLECT_RECORDS_CONCURRENCY, ) - try: - response = await self._request( - "GET", + async with _COLLECT_RECORDS_SEMAPHORE: + logger.info( + "[session-client] collect_start session_id=%s url=%s", + self.session_id, self.base_url, - phase="collect_records", - max_retries=_COLLECT_RETRIES, ) - parsed = GetSessionResponse.model_validate(response) - records = parsed.records or [] - metadata = parsed.metadata or {} + try: + response = await self._request( + "GET", + self.base_url, + phase="collect_records", + max_retries=_COLLECT_RETRIES, + ) - logger.info( - "[session-client] collect_done session_id=%s records=%d metadata_keys=%s", - self.session_id, - len(records), - sorted(metadata.keys()), - ) + parsed = GetSessionResponse.model_validate(response) + records = parsed.records or [] + metadata = parsed.metadata or {} - return records, metadata + logger.info( + "[session-client] collect_done session_id=%s records=%d metadata_keys=%s", + self.session_id, + len(records), + sorted(metadata.keys()), + ) - except httpx.RemoteProtocolError as exc: - logger.info( - "[session-client] collect_failed_remote_protocol session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", - self.session_id, - self.base_url, - type(exc).__name__, - exc, - ) - return [], {} + return records, metadata - except httpx.TimeoutException as exc: - logger.info( - "[session-client] collect_failed_timeout session_id=%s url=%s timeout_s=%.1f error_type=%s error=%r returning_empty_records=True", - self.session_id, - self.base_url, - _SESSION_REQUEST_TIMEOUT, - type(exc).__name__, - exc, - ) - return [], {} + except httpx.ConnectTimeout as exc: + logger.info( + "[session-client] collect_failed_connect_timeout session_id=%s url=%s " + "connect_timeout_s=%.1f error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + _HTTP_CONNECT_TIMEOUT, + type(exc).__name__, + exc, + ) + return [], {} - except httpx.TransportError as exc: - logger.info( - "[session-client] collect_failed_transport session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", - self.session_id, - self.base_url, - type(exc).__name__, - exc, - ) - return [], {} + except httpx.RemoteProtocolError as exc: + logger.info( + "[session-client] collect_failed_remote_protocol session_id=%s url=%s " + "error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} - except Exception as exc: - logger.info( - "[session-client] collect_failed_unexpected session_id=%s url=%s error_type=%s error=%r returning_empty_records=True", - self.session_id, - self.base_url, - type(exc).__name__, - exc, - ) - return [], {} + except httpx.TimeoutException as exc: + logger.info( + "[session-client] collect_failed_timeout session_id=%s url=%s " + "error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} + + except httpx.TransportError as exc: + logger.info( + "[session-client] collect_failed_transport session_id=%s url=%s " + "error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} + + except Exception as exc: + logger.info( + "[session-client] collect_failed_unexpected session_id=%s url=%s " + "error_type=%s error=%r returning_empty_records=True", + self.session_id, + self.base_url, + type(exc).__name__, + exc, + ) + return [], {} - finally: - await self.delete_session() + finally: + await self.delete_session() async def delete_session(self) -> None: logger.info( From e3112f106373fa0b7c69cd624ef5607b6263e500 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 20:49:52 -0700 Subject: [PATCH 10/29] add session server actor and distribute --- miles/ray/rollout.py | 182 ++++++++++++++++++++++--------------------- 1 file changed, 94 insertions(+), 88 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index baec979375d..eab5c228812 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -341,6 +341,46 @@ def onload_kv(self): return ray.get(handles) if handles else [] +@ray.remote +class SessionServerActor: + def __init__(self, args, router_url: str, index: int): + configure_logger() + self.args = copy.copy(args) + self.router_url = router_url + self.index = index + self.process = None + + def start(self) -> str: + from miles.rollout.session.session_server import run_session_server + + ip = _wrap_ipv6(get_host_info()[1]) + port = find_available_port(random.randint(5000, 6000)) + + self.args.session_server_ip = ip + self.args.session_server_port = port + self.args.session_server_instance_id = f"{uuid.uuid4().hex}-{self.index}" + + self.process = multiprocessing.Process( + target=run_session_server, + args=(self.args, self.router_url), + ) + self.process.daemon = True + self.process.start() + + wait_for_server_ready(ip, port, self.process, timeout=30) + + addr = f"{ip}:{port}" + logger.info("Started session server actor %s at %s", self.index, addr) + return addr + + def stop(self): + if self.process is not None and self.process.is_alive(): + self.process.terminate() + self.process.join(timeout=5) + if self.process.is_alive(): + self.process.kill() + + # --------------------------------------------------------------------------- # RolloutManager # --------------------------------------------------------------------------- @@ -385,7 +425,7 @@ def __init__(self, args, pg): else: init_http_client(args) self.servers = start_rollout_servers(args, pg) - _start_session_servers(args) + self.session_server_actors = _start_session_servers(args) self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() self.rollout_id = -1 @@ -425,6 +465,11 @@ def dispose(self): self._metric_checker.dispose() for monitor in self._health_monitors: monitor.stop() + for actor in getattr(self, "session_server_actors", []): + try: + ray.get(actor.stop.remote(), timeout=10) + except Exception as e: + logger.warning("Failed to stop session server actor: %s", e) @property def server(self) -> RolloutServer | None: @@ -629,9 +674,7 @@ def _compute_dynamic_global_batch_size(self, num_samples: int) -> int: if dynamic_gbs != original_gbs or wasted > 0: logger.info( - f"Dynamic global_batch_size: {original_gbs} -> {dynamic_gbs} " - f"(num_samples={num_samples}, dp_size={dp_size}, " - f"num_steps=1, wasted={wasted})" + f"Dynamic global_batch_size: {original_gbs} -> {dynamic_gbs} (num_samples={num_samples}, dp_size={dp_size}, num_steps=1, wasted={wasted})" ) return dynamic_gbs @@ -789,6 +832,7 @@ def set_train_parallel_config(self, config: dict): def _split_train_data_by_dp(self, data, dp_size): """Split the train data by data parallel size, with per-DP imbalance diagnostics.""" import sys + import ray def _safe_len(x): @@ -924,10 +968,7 @@ def _stat(xs): dp_summaries.append(summary) logger.warning( - "ROLLOUT_DP_SHARD " - "dp=%s samples=%s token_sum=%s token_min=%s token_max=%s token_avg=%s " - "response_sum=%s response_min=%s response_max=%s response_avg=%s " - "loss_mask_sum=%s payload_mb_est=%.2f object_ref=%s partition=%s", + "ROLLOUT_DP_SHARD dp=%s samples=%s token_sum=%s token_min=%s token_max=%s token_avg=%s response_sum=%s response_min=%s response_max=%s response_avg=%s loss_mask_sum=%s payload_mb_est=%.2f object_ref=%s partition=%s", summary["dp_rank"], summary["num_samples"], summary["tokens"]["sum"], @@ -959,11 +1000,7 @@ def _ratio(xs): return round(mx / mn, 3) if mn else float("inf") logger.warning( - "ROLLOUT_DP_IMBALANCE " - "dp_size=%s total_samples=%s total_tokens=%s " - "sample_counts=%s sample_ratio=%s " - "token_sums=%s token_ratio=%s " - "payload_mbs=%s payload_ratio=%s", + "ROLLOUT_DP_IMBALANCE dp_size=%s total_samples=%s total_tokens=%s sample_counts=%s sample_ratio=%s token_sums=%s token_ratio=%s payload_mbs=%s payload_ratio=%s", dp_size, len(total_lengths), sum(total_lengths), @@ -1151,8 +1188,7 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool port = router_port if not is_port_available(port): raise RuntimeError( - f"Port {port} is already in use — a stale router process may still be running. " - f"Run 'pkill -9 python' to kill it, then retry." + f"Port {port} is already in use — a stale router process may still be running. Run 'pkill -9 python' to kill it, then retry." ) process = multiprocessing.Process( @@ -1229,8 +1265,7 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: if args.offload_rollout and not needs_offload: overrides.setdefault("enable_memory_saver", False) logger.info( - f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} " - f"(abs={group_abs_start}): needs_offload={needs_offload}" + f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} (abs={group_abs_start}): needs_offload={needs_offload}" ) group = ServerGroup( @@ -1298,84 +1333,57 @@ def _resolve_sglang_config(args) -> SglangConfig: # --------------------------------------------------------------------------- +def _normalize_session_server_addrs(addrs): + if isinstance(addrs, str): + addrs = [a.strip() for a in addrs.split(",") if a.strip()] + + return [addr if addr.startswith(("http://", "https://")) else f"http://{addr}" for addr in addrs] + + def _start_session_servers(args): - """Start N independent vanilla session-server processes on the SGLang gateway node. - - Each session-server is a standalone process bound to its own port; MSA clients - pick one at random per session via ``args.session_server_backends`` (callee-side - multi-backend dispatch). This replaces the single-server design with a multi- - backend design that relieves the per-process GIL when many concurrent sessions - saturate one process. - - ``args.session_server_count`` (default 1) controls N. For N=1 the behavior is - identical to the previous single-server path. ``session_server_ip`` / - ``session_server_port`` are preserved for back-compat (set to host_ip / first - chosen port respectively). - """ if not getattr(args, "use_session_server", False): - return + return [] + + if getattr(args, "session_server_addrs", None): + args.session_server_backends = _normalize_session_server_addrs(args.session_server_addrs) + logger.info("Using external session servers: %s", args.session_server_backends) + return [] hf_checkpoint = getattr(args, "hf_checkpoint", None) if not hf_checkpoint: raise ValueError("--use-session-server requires --hf-checkpoint to be set.") - if getattr(args, "session_server_ip", None) is None: - args.session_server_ip = args.sglang_router_ip - if getattr(args, "session_server_instance_id", None) is None: - args.session_server_instance_id = uuid.uuid4().hex - - host_ip = args.session_server_ip - count = max(1, int(getattr(args, "session_server_count", 1) or 1)) - router_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + session_server_count = max(1, int(getattr(args, "session_server_count", 1) or 1)) + + actors = [ + SessionServerActor.options( + num_cpus=0.1, + num_gpus=0, + scheduling_strategy="SPREAD", + ).remote(args, router_url, i) + for i in range(session_server_count) + ] + + session_server_addrs = ray.get([actor.start.remote() for actor in actors]) + args.session_server_backends = _normalize_session_server_addrs(session_server_addrs) + + first_addr = session_server_addrs[0] + if first_addr.startswith("http://"): + first_addr = first_addr[len("http://") :] + elif first_addr.startswith("https://"): + first_addr = first_addr[len("https://") :] + + first_host, first_port = first_addr.rsplit(":", 1) + args.session_server_ip = first_host + args.session_server_port = int(first_port) - from miles.rollout.session.session_server import run_session_server - - chosen_ports: list[int] = [] - processes: list[multiprocessing.Process] = [] - - # Pre-fill the explicit port (if any) as the first backend so back-compat - # callers landing on backend 0 hit a deterministic port. - explicit_port = getattr(args, "session_server_port", None) - if explicit_port is not None: - if not is_port_available(explicit_port): - raise RuntimeError( - f"Port {explicit_port} is already in use — a stale session server may still be running. " - f"Run 'pkill -9 python' to kill it, then retry." - ) - chosen_ports.append(explicit_port) - - while len(chosen_ports) < count: - port = find_available_port(random.randint(5000, 6000)) - # Avoid races: reject duplicates produced by the random seed. - if port in chosen_ports: - continue - chosen_ports.append(port) - - for port in chosen_ports: - # Each child sees its own port via a shallow-copied args namespace. - child_args = copy.copy(args) - child_args.session_server_port = port - child_args.session_server_ip = host_ip - process = multiprocessing.Process(target=run_session_server, args=(child_args, router_url)) - process.daemon = True - process.start() - processes.append(process) - - for port, process in zip(chosen_ports, processes, strict=True): - wait_for_server_ready(host_ip, port, process, timeout=30) - logger.info(f"Session server launched at {host_ip}:{port}") - - # Publish the backend set; back-compat fields point at backend 0. - args.session_server_backends = [f"http://{host_ip}:{port}" for port in chosen_ports] - args.session_server_ip = host_ip - args.session_server_port = chosen_ports[0] logger.info( - "Session-server multi-backend dispatch enabled: N=%d, backends=%s", - count, + "Started %s Ray session servers: %s", + len(session_server_addrs), args.session_server_backends, ) - return args.session_server_backends + return actors def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): @@ -1468,11 +1476,9 @@ def compute_metrics_from_samples(args, samples): if args.ci_test: for strict_type in ("special_token_count", "special_token_type", "non_assistant_text"): rate = log_dict.get(f"tito_session_mismatch_rate/{strict_type}", 0) - assert rate == 0, ( - f"tito_session_mismatch_rate/{strict_type}={rate:.4f} must be 0 — " - "this indicates a bug in the TITO algorithm or chat template. " - "Please check your tito model and chat template." - ) + assert ( + rate == 0 + ), f"tito_session_mismatch_rate/{strict_type}={rate:.4f} must be 0 — this indicates a bug in the TITO algorithm or chat template. Please check your tito model and chat template." # assistant_text mismatch is non-critical: assistant tokens are inherited # from the pretokenized prefix and may differ from canonical tokenization. # new top-level grouped keys: global From 4bf33bf00f450e49990f0add8dbebe6fc311ae16 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 06:21:22 +0000 Subject: [PATCH 11/29] Raise connect/write/pool timeouts to 600s; lift collect concurrency cap to 8192 --- miles/rollout/generate_utils/openai_endpoint_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 1a73ad88c4b..c823fc225a2 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -20,10 +20,10 @@ _SESSION_REQUEST_TIMEOUT = 24 * 60 * 60.0 -_HTTP_CONNECT_TIMEOUT = 10.0 +_HTTP_CONNECT_TIMEOUT = 600.0 _HTTP_READ_TIMEOUT = 24 * 60 * 60.0 -_HTTP_WRITE_TIMEOUT = 30.0 -_HTTP_POOL_TIMEOUT = 10.0 +_HTTP_WRITE_TIMEOUT = 600.0 +_HTTP_POOL_TIMEOUT = 600.0 _HEALTH_RETRIES = 2 _CREATE_RETRIES = 10 @@ -34,7 +34,7 @@ _BACKOFF_MAX_SECONDS = 10.0 _BACKOFF_JITTER_FRACTION = 0.2 -_COLLECT_RECORDS_CONCURRENCY = 256 +_COLLECT_RECORDS_CONCURRENCY = 8192 _COLLECT_RECORDS_SEMAPHORE = asyncio.Semaphore(_COLLECT_RECORDS_CONCURRENCY) From 17a4d6b6203cd95adf95e0aceebfd835d767682d Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 01:29:54 -0700 Subject: [PATCH 12/29] to_thread --- miles/rollout/generate_hub/agentic_tool_call.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 46f89bccba6..0b847852cfc 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -24,6 +24,7 @@ async def my_agent( """ import argparse +import asyncio import logging import time from collections.abc import Callable @@ -104,7 +105,8 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) logger.debug(f"{log_prefix} Computing samples from {len(records)} records...") - samples = compute_samples_from_openai_records( + samples = await asyncio.to_thread( + compute_samples_from_openai_records, input.args, input.sample, records, From a86619eddaed643dc1d7dfec76947557ea30217a Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 01:37:58 -0700 Subject: [PATCH 13/29] add a threadpool --- train_async.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/train_async.py b/train_async.py index e9e05a40629..cc46612b8bf 100644 --- a/train_async.py +++ b/train_async.py @@ -1,4 +1,5 @@ import asyncio +from concurrent.futures import ThreadPoolExecutor from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from miles.utils.arguments import parse_args @@ -10,6 +11,11 @@ # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). async def train(args): + loop = asyncio.get_running_loop() + + # Raise the default threadpool size used by asyncio.to_thread() + loop.set_default_executor(ThreadPoolExecutor(max_workers=256)) + assert not args.colocate, "Colocation is not supported for async training." configure_logger() # allocate the GPUs From 546ee1d86d75b5f4212c74dad44c150880e3cb5b Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 08:43:26 +0000 Subject: [PATCH 14/29] Shallow-copy samples in merge; parse session responses with orjson --- miles/rollout/generate_utils/openai_endpoint_utils.py | 6 +++--- miles/rollout/generate_utils/sample_utils.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index c823fc225a2..f64bcfdff8a 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -3,7 +3,6 @@ """ import asyncio -import json import logging import random from argparse import Namespace @@ -11,6 +10,7 @@ from typing import Any import httpx +import orjson from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response from miles.rollout.session.session_types import GetSessionResponse, SessionRecord @@ -157,8 +157,8 @@ async def _request( return response.text try: - return response.json() - except json.JSONDecodeError as exc: + return orjson.loads(response.content) + except orjson.JSONDecodeError as exc: logger.info( "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", phase, diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 7845731bed5..4d8bee2f525 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -1,4 +1,4 @@ -from copy import deepcopy +from copy import copy from dataclasses import fields from miles.utils.types import Sample @@ -30,7 +30,7 @@ def merge_samples(samples: list[Sample], tokenizer) -> Sample: def _merge_sample_pair(a: Sample, b: Sample, tokenizer) -> Sample: """Merge two samples generated from sibling inference engine calls.""" - a, b = deepcopy(a), deepcopy(b) + a, b = copy(a), copy(b) def _merge_equal_value(field): x = getattr(a, field) From 008cf75d096879c7fbd7c15f6b706d6939c7011b Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 01:43:30 -0700 Subject: [PATCH 15/29] more to_threads --- miles/rollout/generate_hub/agentic_tool_call.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 0b847852cfc..12fca3a7e91 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -142,7 +142,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) if not input.args.generate_multi_samples: - samples = merge_samples(samples, input.state.tokenizer) + samples = await asyncio.to_thread( + merge_samples, + samples, + input.state.tokenizer, + ) samples.metadata.update(session_metadata) else: samples[-1].metadata.update(session_metadata) From 5364026e167da80240957c2ad5bbb9a31770dd78 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 10:07:22 +0000 Subject: [PATCH 16/29] Shallow-copy input sample per record; only mutable fields get fresh copies --- miles/rollout/generate_utils/openai_endpoint_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index f64bcfdff8a..eab5033a8a3 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -6,7 +6,7 @@ import logging import random from argparse import Namespace -from copy import deepcopy +from copy import copy, deepcopy from typing import Any import httpx @@ -560,7 +560,10 @@ def _compute_sample_from_openai_record( output_token_ids = [item[1] for item in choice["meta_info"]["output_token_logprobs"]] output_log_probs = [item[0] for item in choice["meta_info"]["output_token_logprobs"]] - sample = deepcopy(input_sample) + sample = copy(input_sample) + sample.metadata = dict(input_sample.metadata) + sample.weight_versions = list(input_sample.weight_versions) + sample.prefix_cache_info = deepcopy(input_sample.prefix_cache_info) request_input_ids = record.request.get("input_ids") if request_input_ids is not None: assert ( From d4645b21e5c9d70def7538af7572cb1915d89429 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 20:31:20 +0000 Subject: [PATCH 17/29] Skip SSL context for plain-http session clients; parse responses off the event loop --- miles/rollout/generate_utils/openai_endpoint_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index eab5033a8a3..f8bfca57d69 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -94,7 +94,7 @@ async def _request( method = method.upper() last_exc: BaseException | None = None - async with httpx.AsyncClient(timeout=cls._timeout()) as client: + async with httpx.AsyncClient(timeout=cls._timeout(), verify=False) as client: for attempt in range(1, max_retries + 1): try: logger.info( @@ -157,7 +157,7 @@ async def _request( return response.text try: - return orjson.loads(response.content) + return await asyncio.to_thread(orjson.loads, response.content) except orjson.JSONDecodeError as exc: logger.info( "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", From c007b6fb4b583fcbaee269e9da93c4c86882ae7d Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 21:04:14 +0000 Subject: [PATCH 18/29] prefix_cache_info shallow copy --- miles/rollout/generate_utils/openai_endpoint_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index f8bfca57d69..27b7200c03f 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -6,7 +6,7 @@ import logging import random from argparse import Namespace -from copy import copy, deepcopy +from copy import copy from typing import Any import httpx @@ -563,7 +563,7 @@ def _compute_sample_from_openai_record( sample = copy(input_sample) sample.metadata = dict(input_sample.metadata) sample.weight_versions = list(input_sample.weight_versions) - sample.prefix_cache_info = deepcopy(input_sample.prefix_cache_info) + sample.prefix_cache_info = copy(input_sample.prefix_cache_info) request_input_ids = record.request.get("input_ids") if request_input_ids is not None: assert ( From 417e4207b23b2398aaf679751c6dc2f388148838 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 14:56:47 -0700 Subject: [PATCH 19/29] skip decoding Sample.response if possible --- .../generate_utils/openai_endpoint_utils.py | 45 ++++++++++++++----- miles/rollout/generate_utils/sample_utils.py | 42 +++++++++++++++-- 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 7651e5a2f26..83f8995b92d 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -511,10 +511,8 @@ def compute_samples_from_openai_records( trim_count = 0 if accumulated_token_ids is not None: - # Step 1: position cursor right after this turn's prompt cursor = len(prompt_ids) - # Step 2: greedily match output_ids against accumulated[cursor:] matched = 0 for j in range(len(output_ids)): idx = cursor + j @@ -523,23 +521,18 @@ def compute_samples_from_openai_records( else: break - # Step 3: unmatched trailing tokens were consumed by the next - # turn's template rendering (e.g. stop tokens that double as - # the next message delimiter) — strip them from the sample. trim_count = len(output_ids) - matched allowed = 0 if is_last else max_trim_tokens assert ( trim_count <= allowed ), f"trim_count {trim_count} exceeds allowed={allowed} (is_last={is_last}, max_trim_tokens={max_trim_tokens}); output_ids[-3:]={output_ids[-3:]}, accumulated[cursor:cursor+3]={accumulated_token_ids[cursor : cursor + 3]}" - # Step 4: advance cursor past matched output to the next turn cursor += matched sample = _compute_sample_from_openai_record(args, input_sample, record, tokenizer, trim_count) samples.append(sample) if accumulated_token_ids is not None: - # Step 5: verify the entire accumulated sequence was consumed assert cursor == len( accumulated_token_ids ), f"cursor {cursor} != len(accumulated_token_ids) {len(accumulated_token_ids)} after processing all {len(records)} records" @@ -569,13 +562,17 @@ def _compute_sample_from_openai_record( sample.tokens = prompt_token_ids + output_token_ids sample.rollout_log_probs = output_log_probs - sample.response = tokenizer.decode(output_token_ids) + sample.response = "" sample.response_length = len(output_token_ids) sample.loss_mask = [1] * len(output_token_ids) sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts") + if not hasattr(sample, "metadata") or sample.metadata is None: + sample.metadata = {} + sample.metadata["response_decoded"] = False + if trim_count > 0: - sample.strip_last_output_tokens(trim_count, tokenizer) + _strip_last_output_tokens_without_decode(sample, trim_count) # TODO unify with Sample.update_from_meta_info match choice["finish_reason"]: @@ -593,6 +590,29 @@ def _compute_sample_from_openai_record( return sample +def _strip_last_output_tokens_without_decode(sample: Sample, trim_count: int) -> None: + """Strip trailing output tokens without decoding sample.response.""" + if trim_count <= 0: + return + + assert ( + trim_count <= sample.response_length + ), f"trim_count {trim_count} exceeds response_length {sample.response_length}" + + prompt_len = len(sample.tokens) - sample.response_length + keep_tokens = sample.response_length - trim_count + + sample.tokens = sample.tokens[: prompt_len + keep_tokens] + sample.response_length = keep_tokens + + if sample.rollout_log_probs is not None: + sample.rollout_log_probs = sample.rollout_log_probs[:keep_tokens] + if sample.loss_mask is not None: + sample.loss_mask = sample.loss_mask[:keep_tokens] + if sample.rollout_routed_experts is not None: + sample.rollout_routed_experts = sample.rollout_routed_experts[: len(sample.tokens) - 1] + + def truncate_samples_by_total_tokens( samples: list[Sample], max_seq_len: int, @@ -627,8 +647,13 @@ def _truncate_sample_output(sample: Sample, keep_tokens: int, tokenizer) -> None kept_ids = sample.tokens[prompt_len : prompt_len + keep_tokens] sample.tokens = sample.tokens[:prompt_len] + kept_ids - sample.response = tokenizer.decode(kept_ids) + sample.response = "" sample.response_length = keep_tokens + + if not hasattr(sample, "metadata") or sample.metadata is None: + sample.metadata = {} + sample.metadata["response_decoded"] = False + if sample.rollout_log_probs is not None: sample.rollout_log_probs = sample.rollout_log_probs[:keep_tokens] if sample.loss_mask is not None: diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 7845731bed5..b909551e564 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -49,8 +49,11 @@ def _fill_defaults(sample: Sample): obs_len = len(b.tokens) - len(a.tokens) - b.response_length obs_tokens = b.tokens[len(a.tokens) : len(a.tokens) + obs_len] - # TODO: is this acceptable? - obs_text = tokenizer.decode(obs_tokens) + + response_decoded = _sample_response_decoded(a) and _sample_response_decoded(b) + obs_text = tokenizer.decode(obs_tokens) if response_decoded else "" + response = a.response + obs_text + b.response if response_decoded else "" + metadata = _merge_metadata(a.metadata, b.metadata, response_decoded=response_decoded) try: a.validate() @@ -70,7 +73,7 @@ def _fill_defaults(sample: Sample): tokens=b.tokens, multimodal_inputs=_merge_equal_value("multimodal_inputs"), multimodal_train_inputs=_merge_equal_value("multimodal_train_inputs"), - response=a.response + obs_text + b.response, + response=response, response_length=a.response_length + obs_len + b.response_length, label=_merge_equal_value("label"), reward=_merge_equal_value("reward"), @@ -80,7 +83,7 @@ def _fill_defaults(sample: Sample): rollout_routed_experts=b.rollout_routed_experts, remove_sample=_merge_equal_value("remove_sample"), status=b.status, - metadata=_merge_equal_value("metadata"), + metadata=metadata, generate_function_path=_merge_equal_value("generate_function_path"), train_metadata=_merge_equal_value("train_metadata"), session_id=_merge_equal_value("session_id"), @@ -93,6 +96,37 @@ def _fill_defaults(sample: Sample): raise +def _sample_response_decoded(sample: Sample) -> bool: + """Old samples without the metadata flag are treated as decoded.""" + metadata = getattr(sample, "metadata", None) + if metadata is None: + return True + return metadata.get("response_decoded", True) + + +def _merge_metadata( + a: dict | None, + b: dict | None, + *, + response_decoded: bool, +) -> dict: + a = {} if a is None else dict(a) + b = {} if b is None else dict(b) + + a_without_response_decoded = dict(a) + b_without_response_decoded = dict(b) + a_without_response_decoded.pop("response_decoded", None) + b_without_response_decoded.pop("response_decoded", None) + + assert ( + a_without_response_decoded == b_without_response_decoded + ), f"metadata mismatch: a.metadata={a}, b.metadata={b}" + + merged = dict(a_without_response_decoded) + merged["response_decoded"] = response_decoded + return merged + + def _merge_spec_info(a: Sample.SpecInfo, b: Sample.SpecInfo) -> Sample.SpecInfo: def _merge_plus_value(field): return getattr(a, field) + getattr(b, field) From 7f01a8192bfbbe75c9b98992adf4fc0c6f76fc1e Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 22:44:13 +0000 Subject: [PATCH 20/29] Tail-cancel on total remaining groups (unsubmitted + in flight), not just in-flight --- .../inference_rollout/inference_rollout_train.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py index acc4c03f2a9..d778a79dd50 100644 --- a/miles/rollout/inference_rollout/inference_rollout_train.py +++ b/miles/rollout/inference_rollout/inference_rollout_train.py @@ -244,11 +244,15 @@ async def generate_rollout_async( data.append(group) pbar.update(args.n_samples_per_prompt) - # One pending task == one group, so len(pendings) is the live group count; - # with oversampling off, drop the slow tail once tail_cancel_groups remain. - if args.disable_oversampling and 0 < len(pendings) <= args.tail_cancel_groups: + # Remaining work is unsubmitted groups plus in-flight tasks (one task == + # one group); with oversampling off, drop the slow tail once at most + # tail_cancel_groups remain. Counting only in-flight tasks would fire at + # the start of a rolling-start rollout, when few groups are submitted. + groups_left = (target_data_size - submitted) + len(pendings) + if args.disable_oversampling and 0 < groups_left <= args.tail_cancel_groups: logger.info( - f"[rollout] tail-cancel: {len(pendings)} groups pending <= " + f"[rollout] tail-cancel: {groups_left} groups left " + f"({len(pendings)} in flight, {target_data_size - submitted} unsubmitted) <= " f"{args.tail_cancel_groups}; cutting tail with {len(data)}/{target_data_size} " f"groups collected" ) From 3ba65ffa87cdeb349fb7f9a2040cea90db70f54a Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 22:46:04 +0000 Subject: [PATCH 21/29] drop comment --- miles/rollout/inference_rollout/inference_rollout_train.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py index d778a79dd50..2e0e95403b3 100644 --- a/miles/rollout/inference_rollout/inference_rollout_train.py +++ b/miles/rollout/inference_rollout/inference_rollout_train.py @@ -244,10 +244,6 @@ async def generate_rollout_async( data.append(group) pbar.update(args.n_samples_per_prompt) - # Remaining work is unsubmitted groups plus in-flight tasks (one task == - # one group); with oversampling off, drop the slow tail once at most - # tail_cancel_groups remain. Counting only in-flight tasks would fire at - # the start of a rolling-start rollout, when few groups are submitted. groups_left = (target_data_size - submitted) + len(pendings) if args.disable_oversampling and 0 < groups_left <= args.tail_cancel_groups: logger.info( From df8a83f76a1c55ae48ebcb93f5a0acb936dd9876 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 8 Jul 2026 17:24:47 -0400 Subject: [PATCH 22/29] fix session server port conflict --- miles/ray/rollout.py | 53 ++++++++++++++++++++++++++++++++------------ 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index eab5c228812..0d462d46cdf 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -354,24 +354,49 @@ def start(self) -> str: from miles.rollout.session.session_server import run_session_server ip = _wrap_ipv6(get_host_info()[1]) - port = find_available_port(random.randint(5000, 6000)) + base_port = 20000 + (self.index % 100) * 100 + last_exc = None - self.args.session_server_ip = ip - self.args.session_server_port = port - self.args.session_server_instance_id = f"{uuid.uuid4().hex}-{self.index}" + for attempt in range(10): + port = find_available_port(base_port + attempt * 100) - self.process = multiprocessing.Process( - target=run_session_server, - args=(self.args, self.router_url), - ) - self.process.daemon = True - self.process.start() + self.args.session_server_ip = ip + self.args.session_server_port = port + self.args.session_server_instance_id = f"{uuid.uuid4().hex}-{self.index}" + + self.process = multiprocessing.Process( + target=run_session_server, + args=(self.args, self.router_url), + ) + self.process.daemon = True + self.process.start() + + try: + wait_for_server_ready(ip, port, self.process, timeout=60) + addr = f"{ip}:{port}" + logger.info("Started session server actor %s at %s", self.index, addr) + return addr + except RuntimeError as exc: + last_exc = exc + logger.warning( + "Session server actor %s failed to start on %s:%s " "(attempt %s/10); retrying", + self.index, + ip, + port, + attempt + 1, + exc_info=True, + ) - wait_for_server_ready(ip, port, self.process, timeout=30) + if self.process is not None: + if self.process.is_alive(): + self.process.terminate() + self.process.join(timeout=5) + if self.process.is_alive(): + self.process.kill() + self.process.join(timeout=5) + self.process = None - addr = f"{ip}:{port}" - logger.info("Started session server actor %s at %s", self.index, addr) - return addr + raise RuntimeError(f"Failed to start session server actor {self.index} after 10 attempts") from last_exc def stop(self): if self.process is not None and self.process.is_alive(): From ba0e8d89339363e11b1848db520e6051bfaefb6b Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 8 Jul 2026 19:09:22 -0400 Subject: [PATCH 23/29] lower stats logging (cherry picked from commit dba65d3ee423b3e959e14875f605d7503a050182) --- miles/rollout/session/session_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/rollout/session/session_server.py b/miles/rollout/session/session_server.py index 6e48c10b354..8fe73d10a1b 100644 --- a/miles/rollout/session/session_server.py +++ b/miles/rollout/session/session_server.py @@ -111,7 +111,7 @@ def build_proxy_response(self, result: dict) -> Response: ) -async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): +async def _stats_logger_loop(worker_port, interval_seconds: float = 300.0): """Per-worker observability heartbeat. Logs counters maintained in ``sessions._worker_stats`` plus RSS/VMS From ab0d5658354fa0955b2f2f4d3ecf20ad35640267 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 01:13:25 -0400 Subject: [PATCH 24/29] use gzip compression for /sessions/sessonid endpoint --- .../generate_utils/openai_endpoint_utils.py | 51 ++++++++++++++--- miles/rollout/session/sessions.py | 55 +++++++++++++++---- 2 files changed, 88 insertions(+), 18 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 6d557f3f2ee..d555c44ed1f 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -5,6 +5,7 @@ import asyncio import logging import random +import time from argparse import Namespace from copy import copy from typing import Any @@ -62,9 +63,12 @@ def _timeout() -> httpx.Timeout: @staticmethod def _response_size(response: httpx.Response) -> int: + value = response.headers.get("content-length") + if value is None: + return -1 try: - return len(response.content) - except Exception: + return int(value) + except ValueError: return -1 @staticmethod @@ -106,25 +110,33 @@ async def _request( max_retries, ) + request_headers = {"Accept-Encoding": "gzip", **(headers or {})} + request_t0 = time.monotonic() + if method in {"GET", "DELETE"}: - response = await client.request(method, url, headers=headers) + response = await client.request(method, url, headers=request_headers) else: response = await client.request( method, url, json=payload or {}, - headers=headers, + headers=request_headers, ) body_bytes = cls._response_size(response) + decoded_bytes = len(response.content) logger.info( - "[session-client] response_received phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + "[session-client] response_received phase=%s method=%s url=%s status=%d " + "wire_bytes=%d decoded_bytes=%d content_encoding=%s request_s=%.3f attempt=%d/%d", phase, method, url, response.status_code, body_bytes, + decoded_bytes, + response.headers.get("content-encoding", "identity"), + time.monotonic() - request_t0, attempt, max_retries, ) @@ -156,21 +168,37 @@ async def _request( if not expect_json: return response.text + parse_t0 = time.monotonic() try: - return await asyncio.to_thread(orjson.loads, response.content) + parsed_json = await asyncio.to_thread(orjson.loads, response.content) except orjson.JSONDecodeError as exc: logger.info( - "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", + "[session-client] json_decode_error phase=%s method=%s url=%s status=%d " + "wire_bytes=%d decoded_bytes=%d attempt=%d/%d", phase, method, url, response.status_code, body_bytes, + decoded_bytes, attempt, max_retries, ) raise exc + logger.info( + "[session-client] json_parse_done phase=%s method=%s url=%s decoded_bytes=%d " + "json_parse_s=%.3f attempt=%d/%d", + phase, + method, + url, + decoded_bytes, + time.monotonic() - parse_t0, + attempt, + max_retries, + ) + return parsed_json + except asyncio.CancelledError: logger.info( "[session-client] request_cancelled phase=%s method=%s url=%s attempt=%d/%d", @@ -356,7 +384,14 @@ async def collect_records(self) -> tuple[list[SessionRecord], dict]: max_retries=_COLLECT_RETRIES, ) - parsed = GetSessionResponse.model_validate(response) + validate_t0 = time.monotonic() + parsed = await asyncio.to_thread(GetSessionResponse.model_validate, response) + logger.info( + "[session-client] collect_validate_done session_id=%s validate_s=%.3f", + self.session_id, + time.monotonic() - validate_t0, + ) + records = parsed.records or [] metadata = parsed.metadata or {} diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index acacf5bed44..ea97c1003ed 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -1,4 +1,5 @@ import asyncio +import gzip import logging import time import uuid @@ -179,23 +180,57 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: ) @app.get("/sessions/{session_id}") - async def get_session(session_id: str): + async def get_session(request: Request, session_id: str): + t0 = time.monotonic() + session = registry.sessions.get(session_id) if session is None: if registry.is_deleted(session_id): raise SessionNotFoundError(f"session not found: session_id={session_id}") - return GetSessionResponse(session_id=session_id, records=[], metadata={}) + payload = GetSessionResponse(session_id=session_id, records=[], metadata={}) + else: + metadata = { + "accumulated_token_ids": session.token_ids, + "max_trim_tokens": registry.tito_tokenizer.max_trim_tokens, + } + payload = GetSessionResponse( + session_id=session_id, + records=[_compact_session_record(record) for record in session.records], + metadata=metadata, + ) - metadata = { - "accumulated_token_ids": session.token_ids, - "max_trim_tokens": registry.tito_tokenizer.max_trim_tokens, - } + dump_t0 = time.monotonic() + body = await asyncio.to_thread(orjson.dumps, payload.model_dump(mode="json")) + dump_s = time.monotonic() - dump_t0 + + uncompressed_bytes = len(body) + headers = {"Vary": "Accept-Encoding"} + content_encoding = "identity" - return GetSessionResponse( - session_id=session_id, - records=[_compact_session_record(record) for record in session.records], - metadata=metadata, + gzip_t0 = time.monotonic() + if "gzip" in request.headers.get("accept-encoding", "").lower(): + body = await asyncio.to_thread(gzip.compress, body, compresslevel=6) + headers["Content-Encoding"] = "gzip" + content_encoding = "gzip" + gzip_s = time.monotonic() - gzip_t0 + + wire_bytes = len(body) + headers["Content-Length"] = str(wire_bytes) + + logger.info( + "[session-server] get_session_response session_id=%s records=%d " + "uncompressed_bytes=%d wire_bytes=%d content_encoding=%s " + "json_dump_s=%.3f compress_s=%.3f total_s=%.3f", + session_id, + len(payload.records or []), + uncompressed_bytes, + wire_bytes, + content_encoding, + dump_s, + gzip_s, + time.monotonic() - t0, ) + return Response(content=body, media_type="application/json", headers=headers) @app.delete("/sessions/{session_id}") async def delete_session(session_id: str): From de5ad804155b5ab0540a82deb9df87713b594511 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 14:01:27 -0400 Subject: [PATCH 25/29] Revert "use gzip compression for /sessions/sessonid endpoint" This reverts commit cf6de4283b07453a8abb5d4cd8a8ae3163dcae89. --- .../generate_utils/openai_endpoint_utils.py | 55 ++++--------------- miles/rollout/session/sessions.py | 55 ++++--------------- 2 files changed, 20 insertions(+), 90 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index d555c44ed1f..e79775f1ef8 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -3,15 +3,14 @@ """ import asyncio +import json import logging import random -import time from argparse import Namespace from copy import copy from typing import Any import httpx -import orjson from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response from miles.rollout.session.session_types import GetSessionResponse, SessionRecord @@ -63,12 +62,9 @@ def _timeout() -> httpx.Timeout: @staticmethod def _response_size(response: httpx.Response) -> int: - value = response.headers.get("content-length") - if value is None: - return -1 try: - return int(value) - except ValueError: + return len(response.content) + except Exception: return -1 @staticmethod @@ -110,33 +106,25 @@ async def _request( max_retries, ) - request_headers = {"Accept-Encoding": "gzip", **(headers or {})} - request_t0 = time.monotonic() - if method in {"GET", "DELETE"}: - response = await client.request(method, url, headers=request_headers) + response = await client.request(method, url, headers=headers) else: response = await client.request( method, url, json=payload or {}, - headers=request_headers, + headers=headers, ) body_bytes = cls._response_size(response) - decoded_bytes = len(response.content) logger.info( - "[session-client] response_received phase=%s method=%s url=%s status=%d " - "wire_bytes=%d decoded_bytes=%d content_encoding=%s request_s=%.3f attempt=%d/%d", + "[session-client] response_received phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", phase, method, url, response.status_code, body_bytes, - decoded_bytes, - response.headers.get("content-encoding", "identity"), - time.monotonic() - request_t0, attempt, max_retries, ) @@ -168,37 +156,21 @@ async def _request( if not expect_json: return response.text - parse_t0 = time.monotonic() try: - parsed_json = await asyncio.to_thread(orjson.loads, response.content) - except orjson.JSONDecodeError as exc: + return response.json() + except json.JSONDecodeError as exc: logger.info( - "[session-client] json_decode_error phase=%s method=%s url=%s status=%d " - "wire_bytes=%d decoded_bytes=%d attempt=%d/%d", + "[session-client] json_decode_error phase=%s method=%s url=%s status=%d bytes=%d attempt=%d/%d", phase, method, url, response.status_code, body_bytes, - decoded_bytes, attempt, max_retries, ) raise exc - logger.info( - "[session-client] json_parse_done phase=%s method=%s url=%s decoded_bytes=%d " - "json_parse_s=%.3f attempt=%d/%d", - phase, - method, - url, - decoded_bytes, - time.monotonic() - parse_t0, - attempt, - max_retries, - ) - return parsed_json - except asyncio.CancelledError: logger.info( "[session-client] request_cancelled phase=%s method=%s url=%s attempt=%d/%d", @@ -384,14 +356,7 @@ async def collect_records(self) -> tuple[list[SessionRecord], dict]: max_retries=_COLLECT_RETRIES, ) - validate_t0 = time.monotonic() - parsed = await asyncio.to_thread(GetSessionResponse.model_validate, response) - logger.info( - "[session-client] collect_validate_done session_id=%s validate_s=%.3f", - self.session_id, - time.monotonic() - validate_t0, - ) - + parsed = GetSessionResponse.model_validate(response) records = parsed.records or [] metadata = parsed.metadata or {} diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index ea97c1003ed..acacf5bed44 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -1,5 +1,4 @@ import asyncio -import gzip import logging import time import uuid @@ -180,57 +179,23 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: ) @app.get("/sessions/{session_id}") - async def get_session(request: Request, session_id: str): - t0 = time.monotonic() - + async def get_session(session_id: str): session = registry.sessions.get(session_id) if session is None: if registry.is_deleted(session_id): raise SessionNotFoundError(f"session not found: session_id={session_id}") - payload = GetSessionResponse(session_id=session_id, records=[], metadata={}) - else: - metadata = { - "accumulated_token_ids": session.token_ids, - "max_trim_tokens": registry.tito_tokenizer.max_trim_tokens, - } - payload = GetSessionResponse( - session_id=session_id, - records=[_compact_session_record(record) for record in session.records], - metadata=metadata, - ) - - dump_t0 = time.monotonic() - body = await asyncio.to_thread(orjson.dumps, payload.model_dump(mode="json")) - dump_s = time.monotonic() - dump_t0 - - uncompressed_bytes = len(body) - headers = {"Vary": "Accept-Encoding"} - content_encoding = "identity" + return GetSessionResponse(session_id=session_id, records=[], metadata={}) - gzip_t0 = time.monotonic() - if "gzip" in request.headers.get("accept-encoding", "").lower(): - body = await asyncio.to_thread(gzip.compress, body, compresslevel=6) - headers["Content-Encoding"] = "gzip" - content_encoding = "gzip" - gzip_s = time.monotonic() - gzip_t0 - - wire_bytes = len(body) - headers["Content-Length"] = str(wire_bytes) + metadata = { + "accumulated_token_ids": session.token_ids, + "max_trim_tokens": registry.tito_tokenizer.max_trim_tokens, + } - logger.info( - "[session-server] get_session_response session_id=%s records=%d " - "uncompressed_bytes=%d wire_bytes=%d content_encoding=%s " - "json_dump_s=%.3f compress_s=%.3f total_s=%.3f", - session_id, - len(payload.records or []), - uncompressed_bytes, - wire_bytes, - content_encoding, - dump_s, - gzip_s, - time.monotonic() - t0, + return GetSessionResponse( + session_id=session_id, + records=[_compact_session_record(record) for record in session.records], + metadata=metadata, ) - return Response(content=body, media_type="application/json", headers=headers) @app.delete("/sessions/{session_id}") async def delete_session(session_id: str): From a176bed87f77594a33438cd987f8477dba364401 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 18:31:59 -0400 Subject: [PATCH 26/29] session-server-side merge - first shot --- .../rollout/generate_hub/agentic_tool_call.py | 77 +++--- .../generate_utils/openai_endpoint_utils.py | 191 +++++++++++++- miles/rollout/session/session_types.py | 19 ++ miles/rollout/session/sessions.py | 244 +++++++++++++++++- 4 files changed, 482 insertions(+), 49 deletions(-) diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 12fca3a7e91..ba4655ce09e 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -4,8 +4,7 @@ The agent logic is fully encapsulated in a user-provided async function (--custom-agent-function-path). This generate function only handles: 1. TITO session tracing (OpenAIEndpointTracer) - 2. Converting session records to training samples - 3. Multi-turn merge + 2. Collecting one server-merged training sample Agent function contract: async def my_agent( @@ -24,7 +23,6 @@ async def my_agent( """ import argparse -import asyncio import logging import time from collections.abc import Callable @@ -37,10 +35,9 @@ async def my_agent( from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput from miles.rollout.generate_utils.openai_endpoint_utils import ( OpenAIEndpointTracer, - compute_samples_from_openai_records, + apply_merged_session_sample, truncate_samples_by_total_tokens, ) -from miles.rollout.generate_utils.sample_utils import drop_samples_after_first_non_completed, merge_samples from miles.utils.misc import load_function from miles.utils.types import Sample @@ -60,6 +57,10 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: ), f"Custom agent function {input.args.custom_agent_function_path} not found" max_seq_len = getattr(input.args, "max_seq_len", None) + assert not getattr(input.args, "generate_multi_samples", False), ( + "agentic_tool_call.generate with session-server merged collection produces one merged sample; " + "--generate-multi-samples is not supported." + ) metadata = input.sample.metadata metadata = {**metadata, "sample_idx": input.sample.index} @@ -94,43 +95,29 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: logger.warning(f"{log_prefix} Agent function failed: {e}", exc_info=True) finally: - logger.debug(f"{log_prefix} Calling collect_records...") - records, session_metadata = await tracer.collect_records() - logger.debug(f"{log_prefix} collect_records done: {len(records)} records") + logger.debug(f"{log_prefix} Calling collect_merged_sample...") + merged_sample, session_metadata = await tracer.collect_merged_sample() + if merged_sample is None: + logger.debug(f"{log_prefix} collect_merged_sample done: empty sample") + else: + logger.debug( + f"{log_prefix} collect_merged_sample done: tokens={len(merged_sample.tokens)} " + f"response_length={merged_sample.response_length} " + f"loss_mask={len(merged_sample.loss_mask)} " + f"rollout_log_probs={len(merged_sample.rollout_log_probs)} " + f"status={merged_sample.status}" + ) - if not records: + if merged_sample is None: logger.warning("No model calls recorded for sample") sample = deepcopy(input.sample) sample.status = Sample.Status.ABORTED - return GenerateFnOutput(samples=sample) - - logger.debug(f"{log_prefix} Computing samples from {len(records)} records...") - samples = await asyncio.to_thread( - compute_samples_from_openai_records, - input.args, - input.sample, - records, - input.state.tokenizer, - accumulated_token_ids=session_metadata.get("accumulated_token_ids"), - max_trim_tokens=session_metadata.get("max_trim_tokens", 0), - ) + return GenerateFnOutput(samples=[sample]) - logger.debug( - f"{log_prefix} compute_samples done: {len(samples)} samples, total_time={time.monotonic()-t_start:.1f}s" - ) - for s in samples: - s.metadata.update(agent_metadata or {}) - - # An aborted/length-limited turn invalidates everything generated after - # it; the agent may have kept going on truncated output (e.g. when the - # engine aborts in-flight requests at the end of a rollout step). - samples, num_dropped = drop_samples_after_first_non_completed(samples) - if num_dropped > 0: - logger.warning( - f"{log_prefix} Dropped {num_dropped} trailing turn(s) generated after a " f"{samples[-1].status.name} turn" - ) - for s in samples: - s.metadata["dropped_trailing_turns"] = num_dropped + sample = apply_merged_session_sample(input.args, input.sample, merged_sample) + sample.metadata.update(agent_metadata or {}) + sample.metadata.update(session_metadata) + samples = [sample] if max_seq_len is not None: samples = truncate_samples_by_total_tokens(samples, max_seq_len, input.state.tokenizer) @@ -139,17 +126,13 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: logger.warning("All samples truncated (prompt already exceeds max_seq_len)") sample = deepcopy(input.sample) sample.status = Sample.Status.ABORTED - return GenerateFnOutput(samples=sample) + return GenerateFnOutput(samples=[sample]) - if not input.args.generate_multi_samples: - samples = await asyncio.to_thread( - merge_samples, - samples, - input.state.tokenizer, - ) - samples.metadata.update(session_metadata) - else: - samples[-1].metadata.update(session_metadata) + logger.debug( + f"{log_prefix} server-merged sample ready: samples={len(samples)} " + f"tokens={len(samples[0].tokens)} response_length={samples[0].response_length} " + f"total_time={time.monotonic()-t_start:.1f}s" + ) return GenerateFnOutput(samples=samples) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index e79775f1ef8..caa34e68446 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -13,11 +13,35 @@ import httpx from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response -from miles.rollout.session.session_types import GetSessionResponse, SessionRecord +from miles.rollout.session.session_types import ( + GetMergedSessionResponse, + GetSessionResponse, + MergedSessionSample, + SessionRecord, +) from miles.utils.types import Sample logger = logging.getLogger(__name__) +def _encoded_routed_experts_bytes(routed_experts) -> int: + if routed_experts is None: + return 0 + if isinstance(routed_experts, str): + return len(routed_experts.encode("ascii")) + try: + return len(routed_experts) + except TypeError: + return -1 + + +def _expected_routed_experts_bytes(args: Namespace, num_tokens: int) -> int | None: + num_layers = getattr(args, "num_layers", None) + moe_router_topk = getattr(args, "moe_router_topk", None) + if num_layers is None or moe_router_topk is None: + return None + return max(num_tokens - 1, 0) * int(num_layers) * int(moe_router_topk) * 4 + + _SESSION_REQUEST_TIMEOUT = 24 * 60 * 60.0 _HTTP_CONNECT_TIMEOUT = 600.0 @@ -333,6 +357,118 @@ async def create(args: Namespace): session_server_instance_id=session_server_instance_id, ) + async def collect_merged_sample(self) -> tuple[MergedSessionSample | None, dict]: + merged_url = f"{self.base_url}/merged" + logger.info( + "[session-client] collect_merged_wait session_id=%s url=%s concurrency_limit=%d", + self.session_id, + merged_url, + _COLLECT_RECORDS_CONCURRENCY, + ) + + async with _COLLECT_RECORDS_SEMAPHORE: + logger.info( + "[session-client] collect_merged_start session_id=%s url=%s", + self.session_id, + merged_url, + ) + + try: + response = await self._request( + "GET", + merged_url, + phase="collect_merged_sample", + max_retries=_COLLECT_RETRIES, + ) + + parsed = GetMergedSessionResponse.model_validate(response) + sample = parsed.sample + metadata = parsed.metadata or {} + + if sample is None: + logger.info( + "[session-client] collect_merged_done session_id=%s empty=True metadata_keys=%s", + self.session_id, + sorted(metadata.keys()), + ) + else: + logger.info( + "[session-client] collect_merged_done session_id=%s empty=False tokens=%d " + "response_length=%d loss_mask=%d rollout_log_probs=%d weight_versions=%d " + "prefix_cache_meta_infos=%d routed_experts_encoded_bytes=%d status=%s metadata_keys=%s", + self.session_id, + len(sample.tokens), + sample.response_length, + len(sample.loss_mask), + len(sample.rollout_log_probs), + len(sample.weight_versions), + len(sample.prefix_cache_meta_infos), + _encoded_routed_experts_bytes(sample.rollout_routed_experts), + sample.status, + sorted(metadata.keys()), + ) + + return sample, metadata + + except httpx.ConnectTimeout as exc: + logger.info( + "[session-client] collect_merged_failed_connect_timeout session_id=%s url=%s " + "connect_timeout_s=%.1f error_type=%s error=%r returning_empty_sample=True", + self.session_id, + merged_url, + _HTTP_CONNECT_TIMEOUT, + type(exc).__name__, + exc, + ) + return None, {} + + except httpx.RemoteProtocolError as exc: + logger.info( + "[session-client] collect_merged_failed_remote_protocol session_id=%s url=%s " + "error_type=%s error=%r returning_empty_sample=True", + self.session_id, + merged_url, + type(exc).__name__, + exc, + ) + return None, {} + + except httpx.TimeoutException as exc: + logger.info( + "[session-client] collect_merged_failed_timeout session_id=%s url=%s " + "error_type=%s error=%r returning_empty_sample=True", + self.session_id, + merged_url, + type(exc).__name__, + exc, + ) + return None, {} + + except httpx.TransportError as exc: + logger.info( + "[session-client] collect_merged_failed_transport session_id=%s url=%s " + "error_type=%s error=%r returning_empty_sample=True", + self.session_id, + merged_url, + type(exc).__name__, + exc, + ) + return None, {} + + except Exception as exc: + logger.info( + "[session-client] collect_merged_failed_unexpected session_id=%s url=%s " + "error_type=%s error=%r returning_empty_sample=True", + self.session_id, + merged_url, + type(exc).__name__, + exc, + ) + return None, {} + + finally: + await self.delete_session() + async def collect_records(self) -> tuple[list[SessionRecord], dict]: logger.info( "[session-client] collect_wait session_id=%s url=%s concurrency_limit=%d", @@ -480,6 +616,59 @@ async def delete_session(self) -> None: ) +def apply_merged_session_sample( + args: Namespace, + input_sample: Sample, + merged: MergedSessionSample, +) -> Sample: + """Materialize one server-merged session payload into a local Sample.""" + sample = deepcopy(input_sample) + sample.tokens = merged.tokens + sample.response = merged.response + sample.response_length = merged.response_length + sample.loss_mask = merged.loss_mask + sample.rollout_log_probs = merged.rollout_log_probs + sample.metadata = {**(sample.metadata or {}), **(merged.metadata or {})} + sample.status = Sample.Status(merged.status) + sample.weight_versions.extend(merged.weight_versions) + + for meta_info in merged.prefix_cache_meta_infos: + sample.prefix_cache_info.add(meta_info) + + if merged.rollout_routed_experts is not None: + choice = {"meta_info": {"routed_experts": merged.rollout_routed_experts}} + sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts") + + routed_experts_shape = None + routed_experts_decoded_bytes = 0 + if sample.rollout_routed_experts is not None: + routed_experts_shape = tuple(sample.rollout_routed_experts.shape) + routed_experts_decoded_bytes = int(sample.rollout_routed_experts.nbytes) + + expected_routed_experts_bytes = _expected_routed_experts_bytes(args, len(sample.tokens)) + + logger.info( + "[session-client] apply_merged_sample tokens=%d response_length=%d loss_mask=%d " + "rollout_log_probs=%d weight_versions=%d prefix_cache_meta_infos=%d " + "routed_experts_encoded_bytes=%d routed_experts_decoded_bytes=%d " + "expected_routed_experts_bytes=%s routed_experts_shape=%s status=%s", + len(sample.tokens), + sample.response_length, + len(sample.loss_mask or []), + len(sample.rollout_log_probs or []), + len(sample.weight_versions), + len(merged.prefix_cache_meta_infos), + _encoded_routed_experts_bytes(merged.rollout_routed_experts), + routed_experts_decoded_bytes, + expected_routed_experts_bytes, + routed_experts_shape, + sample.status.value, + ) + + sample.validate() + return sample + + def compute_samples_from_openai_records( args: Namespace, input_sample: Sample, diff --git a/miles/rollout/session/session_types.py b/miles/rollout/session/session_types.py index 65489020932..bff5da91b9f 100644 --- a/miles/rollout/session/session_types.py +++ b/miles/rollout/session/session_types.py @@ -14,3 +14,22 @@ class GetSessionResponse(BaseModel): session_id: str records: list[SessionRecord] metadata: dict = Field(default_factory=dict) + + +class MergedSessionSample(BaseModel): + tokens: list[int] + response: str = "" + response_length: int + loss_mask: list[int] + rollout_log_probs: list[float] + status: str + metadata: dict = Field(default_factory=lambda: {"response_decoded": False}) + weight_versions: list[str] = Field(default_factory=list) + prefix_cache_meta_infos: list[dict] = Field(default_factory=list) + rollout_routed_experts: str | None = None + + +class GetMergedSessionResponse(BaseModel): + session_id: str + sample: MergedSessionSample | None = None + metadata: dict = Field(default_factory=dict) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index acacf5bed44..25d1870c555 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -4,13 +4,19 @@ import uuid import orjson +import pybase64 from fastapi import Request from fastapi.responses import JSONResponse from starlette.responses import Response from miles.rollout.session.linear_trajectory import SessionRegistry from miles.rollout.session.session_errors import SessionError, SessionNotFoundError, UpstreamResponseError -from miles.rollout.session.session_types import GetSessionResponse, SessionRecord +from miles.rollout.session.session_types import ( + GetMergedSessionResponse, + GetSessionResponse, + MergedSessionSample, + SessionRecord, +) from miles.utils.chat_template_utils import get_tito_tokenizer from miles.utils.processing_utils import load_tokenizer @@ -140,6 +146,7 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: # sample.prefix_cache_info.add(choice.get("meta_info", {})) for key in ( "cached_tokens", + "prompt_tokens", "prefix_cache_hit", "prefix_cache_len", "prefix_cache_tokens", @@ -178,6 +185,186 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: }, ) + def _status_from_finish_reason(finish_reason: str | None) -> str: + match finish_reason: + case "stop" | "tool_calls": + return "completed" + case "length": + return "truncated" + case "abort": + return "aborted" + case _: + return "completed" + + def _routed_experts_encoded_bytes(routed_experts) -> int: + if routed_experts is None: + return 0 + if isinstance(routed_experts, str): + return len(routed_experts.encode("ascii")) + try: + return len(routed_experts) + except TypeError: + return -1 + + def _routed_experts_decoded_bytes(routed_experts) -> int: + if not isinstance(routed_experts, str): + return 0 + try: + return len(pybase64.b64decode(routed_experts.encode("ascii"))) + except Exception: + return -1 + + def _expected_routed_experts_bytes(num_tokens: int) -> int | None: + num_layers = getattr(args, "num_layers", None) + moe_router_topk = getattr(args, "moe_router_topk", None) + if num_layers is None or moe_router_topk is None: + return None + return max(num_tokens - 1, 0) * int(num_layers) * int(moe_router_topk) * 4 + + def _merge_session_records_to_sample( + records: list[SessionRecord], + accumulated_token_ids: list[int], + max_trim_tokens: int, + ) -> MergedSessionSample | None: + if not records: + return None + + assert accumulated_token_ids, "cannot merge records without accumulated_token_ids" + + tokens = accumulated_token_ids + first_prompt_len: int | None = None + loss_mask_full = [0] * len(tokens) + rollout_log_probs_full = [0.0] * len(tokens) + weight_versions: list[str] = [] + prefix_cache_meta_infos: list[dict] = [] + routed_experts = None + prev_checkpoint: list[int] | None = None + final_cursor = 0 + final_status = "completed" + + for i, record in enumerate(records): + is_last = i == len(records) - 1 + assert "choices" in record.response and record.response["choices"], f"record {i} missing response.choices" + choice = record.response["choices"][0] + assert "prompt_token_ids" in choice, f"record {i} missing choice.prompt_token_ids" + prompt_ids = choice["prompt_token_ids"] + assert isinstance(prompt_ids, list), f"record {i} prompt_token_ids must be a list" + + request_input_ids = record.request.get("input_ids") + if request_input_ids is not None: + assert request_input_ids == prompt_ids, f"record {i}: request.input_ids must equal prompt_token_ids" + + meta_info = choice.get("meta_info") + assert isinstance(meta_info, dict), f"record {i} missing choice.meta_info" + output_pairs = meta_info.get("output_token_logprobs") + assert isinstance(output_pairs, list), f"record {i} missing meta_info.output_token_logprobs" + + for j, item in enumerate(output_pairs): + assert ( + isinstance(item, (list, tuple)) and len(item) >= 2 + ), f"record {i} output_token_logprobs[{j}] must contain [logprob, token_id]" + assert isinstance(item[1], int), f"record {i} output_token_logprobs[{j}][1] must be token_id int" + + output_ids = [item[1] for item in output_pairs] + output_log_probs = [item[0] for item in output_pairs] + + if first_prompt_len is None: + first_prompt_len = len(prompt_ids) + + if prev_checkpoint is not None: + obs_len = len(prompt_ids) - len(prev_checkpoint) + assert obs_len > 0, f"record {i}: obs_len must be > 0, got {obs_len}" + check_len = max(0, len(prev_checkpoint) - max_trim_tokens) + assert prompt_ids[:check_len] == prev_checkpoint[:check_len], ( + f"record {i}: prompt_ids are not compatible with previous checkpoint " + f"within max_trim_tokens={max_trim_tokens}" + ) + + cursor = len(prompt_ids) + matched = 0 + for j, output_id in enumerate(output_ids): + idx = cursor + j + if idx < len(tokens) and output_id == tokens[idx]: + matched += 1 + else: + break + + trim_count = len(output_ids) - matched + allowed = 0 if is_last else max_trim_tokens + assert trim_count <= allowed, ( + f"record {i}: trim_count={trim_count} exceeds allowed={allowed}; " + f"is_last={is_last}, max_trim_tokens={max_trim_tokens}" + ) + + for j in range(matched): + idx = cursor + j + assert idx < len(tokens), f"record {i}: output token index {idx} outside accumulated tokens" + loss_mask_full[idx] = 1 + rollout_log_probs_full[idx] = output_log_probs[j] + + if "weight_version" in meta_info: + weight_versions.append(str(meta_info["weight_version"])) + + prefix_cache_meta_infos.append( + {key: meta_info[key] for key in ("cached_tokens", "prompt_tokens") if key in meta_info} + ) + + routed_experts = meta_info.get("routed_experts", choice.get("routed_experts")) + final_status = _status_from_finish_reason(choice.get("finish_reason")) + if not is_last: + assert ( + final_status == "completed" + ), f"record {i}: only the final record may be non-completed, got {final_status}" + + prev_checkpoint = prompt_ids + output_ids[:matched] + final_cursor = max(final_cursor, len(prev_checkpoint)) + + logger.debug( + "[session-server] merged_record record_index=%d prompt_tokens=%d output_tokens=%d " + "matched_output_tokens=%d trim_count=%d routed_experts_encoded_bytes=%d " + "routed_experts_decoded_bytes=%d", + i, + len(prompt_ids), + len(output_ids), + matched, + trim_count, + _routed_experts_encoded_bytes(routed_experts), + _routed_experts_decoded_bytes(routed_experts), + ) + + assert final_cursor == len(tokens), f"merged cursor {final_cursor} != accumulated length {len(tokens)}" + assert first_prompt_len is not None, "first_prompt_len was not initialized" + + response_length = len(tokens) - first_prompt_len + assert response_length >= 0, f"response_length must be non-negative, got {response_length}" + assert len(loss_mask_full) == len(tokens), "loss_mask_full length must equal tokens length" + assert len(rollout_log_probs_full) == len(tokens), "rollout_log_probs_full length must equal tokens length" + assert len(loss_mask_full[first_prompt_len:]) == response_length, "loss_mask length must equal response_length" + assert ( + len(rollout_log_probs_full[first_prompt_len:]) == response_length + ), "rollout_log_probs length must equal response_length" + + routed_experts_decoded_bytes = _routed_experts_decoded_bytes(routed_experts) + expected_routed_experts_bytes = _expected_routed_experts_bytes(len(tokens)) + if routed_experts is not None and expected_routed_experts_bytes is not None: + assert routed_experts_decoded_bytes == expected_routed_experts_bytes, ( + f"routed_experts decoded bytes {routed_experts_decoded_bytes} != " + f"expected {expected_routed_experts_bytes}" + ) + + return MergedSessionSample( + tokens=tokens, + response="", + response_length=response_length, + loss_mask=loss_mask_full[first_prompt_len:], + rollout_log_probs=rollout_log_probs_full[first_prompt_len:], + status=final_status, + metadata={"response_decoded": False}, + weight_versions=weight_versions, + prefix_cache_meta_infos=prefix_cache_meta_infos, + rollout_routed_experts=routed_experts, + ) + @app.get("/sessions/{session_id}") async def get_session(session_id: str): session = registry.sessions.get(session_id) @@ -197,6 +384,61 @@ async def get_session(session_id: str): metadata=metadata, ) + @app.get("/sessions/{session_id}/merged") + async def get_merged_session(session_id: str): + session = registry.sessions.get(session_id) + if session is None: + if registry.is_deleted(session_id): + raise SessionNotFoundError(f"session not found: session_id={session_id}") + return GetMergedSessionResponse(session_id=session_id, sample=None, metadata={}) + + async with session.lock: + records = [_compact_session_record(record) for record in session.records] + accumulated_token_ids = list(session.token_ids) + max_trim_tokens = registry.tito_tokenizer.max_trim_tokens + + sample = await asyncio.to_thread( + _merge_session_records_to_sample, + records, + accumulated_token_ids, + max_trim_tokens, + ) + metadata = { + "records_merged": len(records), + "accumulated_token_count": len(accumulated_token_ids), + "max_trim_tokens": max_trim_tokens, + } + + if sample is None: + logger.info( + "[session-server] merged_collect_empty session_id=%s records=%d accumulated_tokens=%d", + session_id, + len(records), + len(accumulated_token_ids), + ) + else: + expected_routed_experts_bytes = _expected_routed_experts_bytes(len(sample.tokens)) + logger.info( + "[session-server] merged_collect_done session_id=%s records=%d tokens=%d response_length=%d " + "loss_mask=%d rollout_log_probs=%d weight_versions=%d prefix_cache_meta_infos=%d " + "routed_experts_encoded_bytes=%d routed_experts_decoded_bytes=%d " + "expected_routed_experts_bytes=%s status=%s", + session_id, + len(records), + len(sample.tokens), + sample.response_length, + len(sample.loss_mask), + len(sample.rollout_log_probs), + len(sample.weight_versions), + len(sample.prefix_cache_meta_infos), + _routed_experts_encoded_bytes(sample.rollout_routed_experts), + _routed_experts_decoded_bytes(sample.rollout_routed_experts), + expected_routed_experts_bytes, + sample.status, + ) + + return GetMergedSessionResponse(session_id=session_id, sample=sample, metadata=metadata) + @app.delete("/sessions/{session_id}") async def delete_session(session_id: str): session = registry.get_session(session_id) From 260f9c51c92349fc00a67aae8493740c641fed90 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 19:52:49 -0400 Subject: [PATCH 27/29] return sample instead of samples --- miles/rollout/generate_hub/agentic_tool_call.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index ba4655ce09e..d89351a3dd4 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -112,7 +112,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: logger.warning("No model calls recorded for sample") sample = deepcopy(input.sample) sample.status = Sample.Status.ABORTED - return GenerateFnOutput(samples=[sample]) + return GenerateFnOutput(samples=sample) sample = apply_merged_session_sample(input.args, input.sample, merged_sample) sample.metadata.update(agent_metadata or {}) @@ -126,14 +126,15 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: logger.warning("All samples truncated (prompt already exceeds max_seq_len)") sample = deepcopy(input.sample) sample.status = Sample.Status.ABORTED - return GenerateFnOutput(samples=[sample]) + return GenerateFnOutput(samples=sample) + sample = samples[0] logger.debug( - f"{log_prefix} server-merged sample ready: samples={len(samples)} " - f"tokens={len(samples[0].tokens)} response_length={samples[0].response_length} " + f"{log_prefix} server-merged sample ready: " + f"tokens={len(sample.tokens)} response_length={sample.response_length} " f"total_time={time.monotonic()-t_start:.1f}s" ) - return GenerateFnOutput(samples=samples) + return GenerateFnOutput(samples=sample) def build_agent_function_kwargs( From 74c518c9c2ee8dc5d3e48b76f2b2c5fcc66885e9 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 15 Jul 2026 15:50:35 -0400 Subject: [PATCH 28/29] introduce _keep_records_until_first_non_completed in the session server to match previous implementation --- miles/rollout/session/sessions.py | 34 ++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 25d1870c555..b59deb3baa8 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -196,6 +196,19 @@ def _status_from_finish_reason(finish_reason: str | None) -> str: case _: return "completed" + def _keep_records_until_first_non_completed(records: list[SessionRecord]) -> tuple[list[SessionRecord], int]: + """Match sample_utils.drop_samples_after_first_non_completed semantics. + + Keep records through the first non-completed turn and drop any later + records, because later turns were conditioned on incomplete output. + """ + for i, record in enumerate(records[:-1]): + choice = record.response.get("choices", [{}])[0] + status = _status_from_finish_reason(choice.get("finish_reason")) + if status != "completed": + return records[: i + 1], len(records) - i - 1 + return records, 0 + def _routed_experts_encoded_bytes(routed_experts) -> int: if routed_experts is None: return 0 @@ -394,17 +407,28 @@ async def get_merged_session(session_id: str): async with session.lock: records = [_compact_session_record(record) for record in session.records] - accumulated_token_ids = list(session.token_ids) + records_to_merge, dropped_records = _keep_records_until_first_non_completed(records) + if records_to_merge: + # Records and trajectory_token_ids are kept in lockstep by + # LinearTrajectory.append_record/update_pretokenized_state and + # rollback truncates both together. If we drop trailing records, + # use the checkpoint for the last kept record so the merged + # tokens match the old rollout-manager merge result. + accumulated_token_ids = list(session.trajectory_token_ids[len(records_to_merge) - 1]) + else: + accumulated_token_ids = [] max_trim_tokens = registry.tito_tokenizer.max_trim_tokens sample = await asyncio.to_thread( _merge_session_records_to_sample, - records, + records_to_merge, accumulated_token_ids, max_trim_tokens, ) metadata = { - "records_merged": len(records), + "records_total": len(records), + "records_merged": len(records_to_merge), + "records_dropped_after_first_non_completed": dropped_records, "accumulated_token_count": len(accumulated_token_ids), "max_trim_tokens": max_trim_tokens, } @@ -413,7 +437,7 @@ async def get_merged_session(session_id: str): logger.info( "[session-server] merged_collect_empty session_id=%s records=%d accumulated_tokens=%d", session_id, - len(records), + len(records_to_merge), len(accumulated_token_ids), ) else: @@ -424,7 +448,7 @@ async def get_merged_session(session_id: str): "routed_experts_encoded_bytes=%d routed_experts_decoded_bytes=%d " "expected_routed_experts_bytes=%s status=%s", session_id, - len(records), + len(records_to_merge), len(sample.tokens), sample.response_length, len(sample.loss_mask), From 14a7daf3a575c006a87720b28ba48abb71fa0115 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Sun, 26 Jul 2026 09:52:01 +0000 Subject: [PATCH 29/29] Send routed_experts_start_len per turn; merge per-turn deltas positionally --- miles/rollout/session/sessions.py | 570 +++++++++++++++++------------- 1 file changed, 316 insertions(+), 254 deletions(-) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index b59deb3baa8..e05e1783a41 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -40,6 +40,302 @@ def get_worker_stats(port): return _worker_stats.get(port if port is not None else "default") +def _status_from_finish_reason(finish_reason: str | None) -> str: + match finish_reason: + case "stop" | "tool_calls": + return "completed" + case "length": + return "truncated" + case "abort": + return "aborted" + case _: + return "completed" + + +def _routed_experts_encoded_bytes(routed_experts) -> int: + if routed_experts is None: + return 0 + if isinstance(routed_experts, str): + return len(routed_experts.encode("ascii")) + try: + return len(routed_experts) + except TypeError: + return -1 + + +def _routed_experts_decoded_bytes(routed_experts) -> int: + if not isinstance(routed_experts, str): + return 0 + try: + return len(pybase64.b64decode(routed_experts.encode("ascii"))) + except Exception: + return -1 + + +def _expected_routed_experts_bytes(args, num_tokens: int) -> int | None: + num_layers = getattr(args, "num_layers", None) + moe_router_topk = getattr(args, "moe_router_topk", None) + if num_layers is None or moe_router_topk is None: + return None + return max(num_tokens - 1, 0) * int(num_layers) * int(moe_router_topk) * 4 + + +def _merge_session_records_to_sample( + records: list[SessionRecord], + accumulated_token_ids: list[int], + max_trim_tokens: int, + args, +) -> MergedSessionSample | None: + if not records: + return None + + assert accumulated_token_ids, "cannot merge records without accumulated_token_ids" + + tokens = accumulated_token_ids + first_prompt_len: int | None = None + loss_mask_full = [0] * len(tokens) + rollout_log_probs_full = [0.0] * len(tokens) + weight_versions: list[str] = [] + prefix_cache_meta_infos: list[dict] = [] + # Full-sequence routed-experts bytes assembled positionally from the + # records. A record covers rows [routed_experts_start_len, its end); + # records without the field start at row 0, reducing to last-wins. + routed_experts_buf = bytearray() + routed_experts_row_bytes = _expected_routed_experts_bytes(args, 2) + prev_checkpoint: list[int] | None = None + final_cursor = 0 + final_status = "completed" + + for i, record in enumerate(records): + is_last = i == len(records) - 1 + assert "choices" in record.response and record.response["choices"], f"record {i} missing response.choices" + choice = record.response["choices"][0] + assert "prompt_token_ids" in choice, f"record {i} missing choice.prompt_token_ids" + prompt_ids = choice["prompt_token_ids"] + assert isinstance(prompt_ids, list), f"record {i} prompt_token_ids must be a list" + + request_input_ids = record.request.get("input_ids") + if request_input_ids is not None: + assert request_input_ids == prompt_ids, f"record {i}: request.input_ids must equal prompt_token_ids" + + meta_info = choice.get("meta_info") + assert isinstance(meta_info, dict), f"record {i} missing choice.meta_info" + output_pairs = meta_info.get("output_token_logprobs") + assert isinstance(output_pairs, list), f"record {i} missing meta_info.output_token_logprobs" + + for j, item in enumerate(output_pairs): + assert ( + isinstance(item, (list, tuple)) and len(item) >= 2 + ), f"record {i} output_token_logprobs[{j}] must contain [logprob, token_id]" + assert isinstance(item[1], int), f"record {i} output_token_logprobs[{j}][1] must be token_id int" + + output_ids = [item[1] for item in output_pairs] + output_log_probs = [item[0] for item in output_pairs] + + if first_prompt_len is None: + first_prompt_len = len(prompt_ids) + + if prev_checkpoint is not None: + obs_len = len(prompt_ids) - len(prev_checkpoint) + assert obs_len > 0, f"record {i}: obs_len must be > 0, got {obs_len}" + check_len = max(0, len(prev_checkpoint) - max_trim_tokens) + assert prompt_ids[:check_len] == prev_checkpoint[:check_len], ( + f"record {i}: prompt_ids are not compatible with previous checkpoint " + f"within max_trim_tokens={max_trim_tokens}" + ) + + cursor = len(prompt_ids) + matched = 0 + for j, output_id in enumerate(output_ids): + idx = cursor + j + if idx < len(tokens) and output_id == tokens[idx]: + matched += 1 + else: + break + + trim_count = len(output_ids) - matched + allowed = 0 if is_last else max_trim_tokens + assert trim_count <= allowed, ( + f"record {i}: trim_count={trim_count} exceeds allowed={allowed}; " + f"is_last={is_last}, max_trim_tokens={max_trim_tokens}" + ) + + for j in range(matched): + idx = cursor + j + assert idx < len(tokens), f"record {i}: output token index {idx} outside accumulated tokens" + loss_mask_full[idx] = 1 + rollout_log_probs_full[idx] = output_log_probs[j] + + if "weight_version" in meta_info: + weight_versions.append(str(meta_info["weight_version"])) + + prefix_cache_meta_infos.append( + {key: meta_info[key] for key in ("cached_tokens", "prompt_tokens") if key in meta_info} + ) + + record_experts = meta_info.get("routed_experts", choice.get("routed_experts")) + if record_experts is not None: + delta = ( + pybase64.b64decode(record_experts.encode("ascii")) + if isinstance(record_experts, str) + else bytes(record_experts) + ) + start_rows = int(record.request.get("routed_experts_start_len") or 0) + if start_rows: + assert routed_experts_row_bytes, ( + f"record {i}: routed_experts_start_len={start_rows} requires " + "num_layers and moe_router_topk args to compute the row size" + ) + start_bytes = start_rows * (routed_experts_row_bytes or 0) + assert start_bytes <= len(routed_experts_buf), ( + f"record {i}: routed_experts_start_len={start_rows} leaves a gap " + f"(buffer has {len(routed_experts_buf)} bytes, needs {start_bytes})" + ) + del routed_experts_buf[start_bytes:] + routed_experts_buf += delta + + final_status = _status_from_finish_reason(choice.get("finish_reason")) + if not is_last: + assert ( + final_status == "completed" + ), f"record {i}: only the final record may be non-completed, got {final_status}" + + prev_checkpoint = prompt_ids + output_ids[:matched] + final_cursor = max(final_cursor, len(prev_checkpoint)) + + logger.debug( + "[session-server] merged_record record_index=%d prompt_tokens=%d output_tokens=%d " + "matched_output_tokens=%d trim_count=%d routed_experts_encoded_bytes=%d " + "routed_experts_decoded_bytes=%d", + i, + len(prompt_ids), + len(output_ids), + matched, + trim_count, + _routed_experts_encoded_bytes(record_experts), + len(routed_experts_buf), + ) + + assert final_cursor == len(tokens), f"merged cursor {final_cursor} != accumulated length {len(tokens)}" + assert first_prompt_len is not None, "first_prompt_len was not initialized" + + response_length = len(tokens) - first_prompt_len + assert response_length >= 0, f"response_length must be non-negative, got {response_length}" + assert len(loss_mask_full) == len(tokens), "loss_mask_full length must equal tokens length" + assert len(rollout_log_probs_full) == len(tokens), "rollout_log_probs_full length must equal tokens length" + assert len(loss_mask_full[first_prompt_len:]) == response_length, "loss_mask length must equal response_length" + assert ( + len(rollout_log_probs_full[first_prompt_len:]) == response_length + ), "rollout_log_probs length must equal response_length" + + routed_experts = ( + pybase64.b64encode(bytes(routed_experts_buf)).decode("ascii") if routed_experts_buf else None + ) + routed_experts_decoded_bytes = len(routed_experts_buf) + expected_routed_experts_bytes = _expected_routed_experts_bytes(args, len(tokens)) + if routed_experts is not None and expected_routed_experts_bytes is not None: + assert routed_experts_decoded_bytes == expected_routed_experts_bytes, ( + f"routed_experts decoded bytes {routed_experts_decoded_bytes} != " + f"expected {expected_routed_experts_bytes}" + ) + + return MergedSessionSample( + tokens=tokens, + response="", + response_length=response_length, + loss_mask=loss_mask_full[first_prompt_len:], + rollout_log_probs=rollout_log_probs_full[first_prompt_len:], + status=final_status, + metadata={"response_decoded": False}, + weight_versions=weight_versions, + prefix_cache_meta_infos=prefix_cache_meta_infos, + rollout_routed_experts=routed_experts, + ) + + +def _compact_output_token_logprobs(output_token_logprobs): + """Keep only the fields consumed by the sample builder. + + Consumer uses: + output_log_probs = [item[0] for item in output_token_logprobs] + output_token_ids = [item[1] for item in output_token_logprobs] + + Some backends may include extra per-token fields; do not return them. + """ + if not output_token_logprobs: + return [] + + compact = [] + for item in output_token_logprobs: + compact.append([item[0], item[1]]) + return compact + + +def _compact_session_record(record: SessionRecord) -> SessionRecord: + """Return only the record fields required by compute_samples_from_openai_records.""" + choice = record.response.get("choices", [{}])[0] + meta_info = choice.get("meta_info") or {} + + compact_meta_info = { + "output_token_logprobs": _compact_output_token_logprobs(meta_info.get("output_token_logprobs")), + } + + # Used by _compute_sample_from_openai_record. + if "weight_version" in meta_info: + compact_meta_info["weight_version"] = meta_info["weight_version"] + + # Preserve prefix-cache metadata if present. The consumer does: + # sample.prefix_cache_info.add(choice.get("meta_info", {})) + for key in ( + "cached_tokens", + "prompt_tokens", + "prefix_cache_hit", + "prefix_cache_len", + "prefix_cache_tokens", + "prefix_cache_start_len", + "prefix_cache_end_len", + ): + if key in meta_info: + compact_meta_info[key] = meta_info[key] + + # Preserve routed experts if present. get_rollout_topk_from_response() + # may read this from choice or meta_info depending on implementation. + if "routed_experts" in meta_info: + compact_meta_info["routed_experts"] = meta_info["routed_experts"] + + compact_choice = { + "prompt_token_ids": choice.get("prompt_token_ids", []), + "finish_reason": choice.get("finish_reason"), + "meta_info": compact_meta_info, + } + + if "routed_experts" in choice: + compact_choice["routed_experts"] = choice["routed_experts"] + + return SessionRecord( + timestamp=record.timestamp, + method=record.method, + path=record.path, + status_code=record.status_code, + request={ + # Only used for this assertion: + # request_input_ids == prompt_token_ids + "input_ids": record.request.get("input_ids"), + # Row offset of this record's routed_experts delta; without it + # the merge places the delta at row 0 and the full-coverage + # assert rejects the sample. + **( + {"routed_experts_start_len": record.request["routed_experts_start_len"]} + if "routed_experts_start_len" in record.request + else {} + ), + }, + response={ + "choices": [compact_choice], + }, + ) + + def setup_session_routes(app, backend, args): hf_checkpoint = getattr(args, "hf_checkpoint", None) if not hf_checkpoint: @@ -112,90 +408,6 @@ async def create_session(): session_id = registry.create_session() return {"session_id": session_id} - def _compact_output_token_logprobs(output_token_logprobs): - """Keep only the fields consumed by the sample builder. - - Consumer uses: - output_log_probs = [item[0] for item in output_token_logprobs] - output_token_ids = [item[1] for item in output_token_logprobs] - - Some backends may include extra per-token fields; do not return them. - """ - if not output_token_logprobs: - return [] - - compact = [] - for item in output_token_logprobs: - compact.append([item[0], item[1]]) - return compact - - def _compact_session_record(record: SessionRecord) -> SessionRecord: - """Return only the record fields required by compute_samples_from_openai_records.""" - choice = record.response.get("choices", [{}])[0] - meta_info = choice.get("meta_info") or {} - - compact_meta_info = { - "output_token_logprobs": _compact_output_token_logprobs(meta_info.get("output_token_logprobs")), - } - - # Used by _compute_sample_from_openai_record. - if "weight_version" in meta_info: - compact_meta_info["weight_version"] = meta_info["weight_version"] - - # Preserve prefix-cache metadata if present. The consumer does: - # sample.prefix_cache_info.add(choice.get("meta_info", {})) - for key in ( - "cached_tokens", - "prompt_tokens", - "prefix_cache_hit", - "prefix_cache_len", - "prefix_cache_tokens", - "prefix_cache_start_len", - "prefix_cache_end_len", - ): - if key in meta_info: - compact_meta_info[key] = meta_info[key] - - # Preserve routed experts if present. get_rollout_topk_from_response() - # may read this from choice or meta_info depending on implementation. - if "routed_experts" in meta_info: - compact_meta_info["routed_experts"] = meta_info["routed_experts"] - - compact_choice = { - "prompt_token_ids": choice.get("prompt_token_ids", []), - "finish_reason": choice.get("finish_reason"), - "meta_info": compact_meta_info, - } - - if "routed_experts" in choice: - compact_choice["routed_experts"] = choice["routed_experts"] - - return SessionRecord( - timestamp=record.timestamp, - method=record.method, - path=record.path, - status_code=record.status_code, - request={ - # Only used for this assertion: - # request_input_ids == prompt_token_ids - "input_ids": record.request.get("input_ids"), - }, - response={ - "choices": [compact_choice], - }, - ) - - def _status_from_finish_reason(finish_reason: str | None) -> str: - match finish_reason: - case "stop" | "tool_calls": - return "completed" - case "length": - return "truncated" - case "abort": - return "aborted" - case _: - return "completed" - def _keep_records_until_first_non_completed(records: list[SessionRecord]) -> tuple[list[SessionRecord], int]: """Match sample_utils.drop_samples_after_first_non_completed semantics. @@ -209,175 +421,6 @@ def _keep_records_until_first_non_completed(records: list[SessionRecord]) -> tup return records[: i + 1], len(records) - i - 1 return records, 0 - def _routed_experts_encoded_bytes(routed_experts) -> int: - if routed_experts is None: - return 0 - if isinstance(routed_experts, str): - return len(routed_experts.encode("ascii")) - try: - return len(routed_experts) - except TypeError: - return -1 - - def _routed_experts_decoded_bytes(routed_experts) -> int: - if not isinstance(routed_experts, str): - return 0 - try: - return len(pybase64.b64decode(routed_experts.encode("ascii"))) - except Exception: - return -1 - - def _expected_routed_experts_bytes(num_tokens: int) -> int | None: - num_layers = getattr(args, "num_layers", None) - moe_router_topk = getattr(args, "moe_router_topk", None) - if num_layers is None or moe_router_topk is None: - return None - return max(num_tokens - 1, 0) * int(num_layers) * int(moe_router_topk) * 4 - - def _merge_session_records_to_sample( - records: list[SessionRecord], - accumulated_token_ids: list[int], - max_trim_tokens: int, - ) -> MergedSessionSample | None: - if not records: - return None - - assert accumulated_token_ids, "cannot merge records without accumulated_token_ids" - - tokens = accumulated_token_ids - first_prompt_len: int | None = None - loss_mask_full = [0] * len(tokens) - rollout_log_probs_full = [0.0] * len(tokens) - weight_versions: list[str] = [] - prefix_cache_meta_infos: list[dict] = [] - routed_experts = None - prev_checkpoint: list[int] | None = None - final_cursor = 0 - final_status = "completed" - - for i, record in enumerate(records): - is_last = i == len(records) - 1 - assert "choices" in record.response and record.response["choices"], f"record {i} missing response.choices" - choice = record.response["choices"][0] - assert "prompt_token_ids" in choice, f"record {i} missing choice.prompt_token_ids" - prompt_ids = choice["prompt_token_ids"] - assert isinstance(prompt_ids, list), f"record {i} prompt_token_ids must be a list" - - request_input_ids = record.request.get("input_ids") - if request_input_ids is not None: - assert request_input_ids == prompt_ids, f"record {i}: request.input_ids must equal prompt_token_ids" - - meta_info = choice.get("meta_info") - assert isinstance(meta_info, dict), f"record {i} missing choice.meta_info" - output_pairs = meta_info.get("output_token_logprobs") - assert isinstance(output_pairs, list), f"record {i} missing meta_info.output_token_logprobs" - - for j, item in enumerate(output_pairs): - assert ( - isinstance(item, (list, tuple)) and len(item) >= 2 - ), f"record {i} output_token_logprobs[{j}] must contain [logprob, token_id]" - assert isinstance(item[1], int), f"record {i} output_token_logprobs[{j}][1] must be token_id int" - - output_ids = [item[1] for item in output_pairs] - output_log_probs = [item[0] for item in output_pairs] - - if first_prompt_len is None: - first_prompt_len = len(prompt_ids) - - if prev_checkpoint is not None: - obs_len = len(prompt_ids) - len(prev_checkpoint) - assert obs_len > 0, f"record {i}: obs_len must be > 0, got {obs_len}" - check_len = max(0, len(prev_checkpoint) - max_trim_tokens) - assert prompt_ids[:check_len] == prev_checkpoint[:check_len], ( - f"record {i}: prompt_ids are not compatible with previous checkpoint " - f"within max_trim_tokens={max_trim_tokens}" - ) - - cursor = len(prompt_ids) - matched = 0 - for j, output_id in enumerate(output_ids): - idx = cursor + j - if idx < len(tokens) and output_id == tokens[idx]: - matched += 1 - else: - break - - trim_count = len(output_ids) - matched - allowed = 0 if is_last else max_trim_tokens - assert trim_count <= allowed, ( - f"record {i}: trim_count={trim_count} exceeds allowed={allowed}; " - f"is_last={is_last}, max_trim_tokens={max_trim_tokens}" - ) - - for j in range(matched): - idx = cursor + j - assert idx < len(tokens), f"record {i}: output token index {idx} outside accumulated tokens" - loss_mask_full[idx] = 1 - rollout_log_probs_full[idx] = output_log_probs[j] - - if "weight_version" in meta_info: - weight_versions.append(str(meta_info["weight_version"])) - - prefix_cache_meta_infos.append( - {key: meta_info[key] for key in ("cached_tokens", "prompt_tokens") if key in meta_info} - ) - - routed_experts = meta_info.get("routed_experts", choice.get("routed_experts")) - final_status = _status_from_finish_reason(choice.get("finish_reason")) - if not is_last: - assert ( - final_status == "completed" - ), f"record {i}: only the final record may be non-completed, got {final_status}" - - prev_checkpoint = prompt_ids + output_ids[:matched] - final_cursor = max(final_cursor, len(prev_checkpoint)) - - logger.debug( - "[session-server] merged_record record_index=%d prompt_tokens=%d output_tokens=%d " - "matched_output_tokens=%d trim_count=%d routed_experts_encoded_bytes=%d " - "routed_experts_decoded_bytes=%d", - i, - len(prompt_ids), - len(output_ids), - matched, - trim_count, - _routed_experts_encoded_bytes(routed_experts), - _routed_experts_decoded_bytes(routed_experts), - ) - - assert final_cursor == len(tokens), f"merged cursor {final_cursor} != accumulated length {len(tokens)}" - assert first_prompt_len is not None, "first_prompt_len was not initialized" - - response_length = len(tokens) - first_prompt_len - assert response_length >= 0, f"response_length must be non-negative, got {response_length}" - assert len(loss_mask_full) == len(tokens), "loss_mask_full length must equal tokens length" - assert len(rollout_log_probs_full) == len(tokens), "rollout_log_probs_full length must equal tokens length" - assert len(loss_mask_full[first_prompt_len:]) == response_length, "loss_mask length must equal response_length" - assert ( - len(rollout_log_probs_full[first_prompt_len:]) == response_length - ), "rollout_log_probs length must equal response_length" - - routed_experts_decoded_bytes = _routed_experts_decoded_bytes(routed_experts) - expected_routed_experts_bytes = _expected_routed_experts_bytes(len(tokens)) - if routed_experts is not None and expected_routed_experts_bytes is not None: - assert routed_experts_decoded_bytes == expected_routed_experts_bytes, ( - f"routed_experts decoded bytes {routed_experts_decoded_bytes} != " - f"expected {expected_routed_experts_bytes}" - ) - - return MergedSessionSample( - tokens=tokens, - response="", - response_length=response_length, - loss_mask=loss_mask_full[first_prompt_len:], - rollout_log_probs=rollout_log_probs_full[first_prompt_len:], - status=final_status, - metadata={"response_decoded": False}, - weight_versions=weight_versions, - prefix_cache_meta_infos=prefix_cache_meta_infos, - rollout_routed_experts=routed_experts, - ) - @app.get("/sessions/{session_id}") async def get_session(session_id: str): session = registry.sessions.get(session_id) @@ -424,6 +467,7 @@ async def get_merged_session(session_id: str): records_to_merge, accumulated_token_ids, max_trim_tokens, + args, ) metadata = { "records_total": len(records), @@ -441,7 +485,7 @@ async def get_merged_session(session_id: str): len(accumulated_token_ids), ) else: - expected_routed_experts_bytes = _expected_routed_experts_bytes(len(sample.tokens)) + expected_routed_experts_bytes = _expected_routed_experts_bytes(args, len(sample.tokens)) logger.info( "[session-server] merged_collect_done session_id=%s records=%d tokens=%d response_length=%d " "loss_mask=%d rollout_log_probs=%d weight_versions=%d prefix_cache_meta_infos=%d " @@ -581,6 +625,17 @@ async def chat_completions(request: Request, session_id: str): len(pretokenized["input_ids"]), ) + if request_body.get("return_routed_experts"): + # Rows [0, start_len) were returned by earlier turns and are + # stitched back in _merge_session_records_to_sample. The + # max_trim_tokens tail may be rewritten by the engine, and + # the checkpoint's final token gets its expert row only in + # this request's forward pass, so both stay in the window. + request_body["routed_experts_start_len"] = max( + 0, + len(session.token_ids) - registry.tito_tokenizer.max_trim_tokens - 1, + ) + body = orjson.dumps(request_body) expected_num_assistant = session.num_assistant @@ -723,6 +778,13 @@ async def chat_completions(request: Request, session_id: str): completion_tokens_emit = -1 _stats["turns_completed"] += 1 + # Strip routed_experts/token-id blobs (in meta_info + prompt_token_ids) + # from the agent's copy only; the trainer reads them from the recorded + # SessionRecord. Returning them every turn made harbor RSS grow unbounded. + result["response_body"] = orjson.dumps({**response, "choices": [ + {k: v for k, v in c.items() if k not in ("meta_info", "prompt_token_ids")} + for c in response.get("choices", []) + ]}) return backend.build_proxy_response(result) finally: _inflight_chat["count"] -= 1