From 0bd6425bd710266d382d484c2c2745ca8ffdf25a Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 11:31:11 -0700 Subject: [PATCH 01/17] 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 e03197d6a5..d5c3c357c6 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 02/17] 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 d5c3c357c6..32c310e987 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 03/17] 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 aee310ad56..4ca6abb324 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 d71934cf6207584b0131b937c88977da0d85be16 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 2 Jul 2026 17:12:48 -0700 Subject: [PATCH 04/17] 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 32c310e987..7651e5a2f2 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 05/17] 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 baec979375..eab5c22881 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 17a4d6b6203cd95adf95e0aceebfd835d767682d Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 01:29:54 -0700 Subject: [PATCH 06/17] 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 46f89bccba..0b847852cf 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 07/17] add a threadpool --- train_async.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/train_async.py b/train_async.py index e9e05a4062..cc46612b8b 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 008cf75d096879c7fbd7c15f6b706d6939c7011b Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 01:43:30 -0700 Subject: [PATCH 08/17] 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 0b847852cf..12fca3a7e9 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 417e4207b23b2398aaf679751c6dc2f388148838 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Fri, 3 Jul 2026 14:56:47 -0700 Subject: [PATCH 09/17] 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 7651e5a2f2..83f8995b92 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 7845731bed..b909551e56 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 07b5c6207cb912559becf600a70bfa3ffaf1f23a Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 8 Jul 2026 17:24:47 -0400 Subject: [PATCH 10/17] 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 eab5c22881..0d462d46cd 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 dba65d3ee423b3e959e14875f605d7503a050182 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 8 Jul 2026 19:09:22 -0400 Subject: [PATCH 11/17] lower stats logging --- 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 6e48c10b35..8fe73d10a1 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 cf6de4283b07453a8abb5d4cd8a8ae3163dcae89 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 01:13:25 -0400 Subject: [PATCH 12/17] use gzip compression for /sessions/sessonid endpoint --- .../generate_utils/openai_endpoint_utils.py | 55 +++++++++++++++---- miles/rollout/session/sessions.py | 55 +++++++++++++++---- 2 files changed, 90 insertions(+), 20 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 83f8995b92..4323590f28 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -3,14 +3,15 @@ """ import asyncio -import json import logging import random +import time from argparse import Namespace from copy import deepcopy 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 @@ -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 response.json() - except json.JSONDecodeError as exc: + 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 4ca6abb324..315f262942 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 @@ -172,23 +173,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 d8905371e1ab23bd76e0e653a5c1755a3f3a9b17 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 14:01:27 -0400 Subject: [PATCH 13/17] 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 4323590f28..83f8995b92 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 deepcopy 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 315f262942..4ca6abb324 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 @@ -173,57 +172,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 8b5dfb632b9f02a955b857041df803e163fbc4f7 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 18:31:59 -0400 Subject: [PATCH 14/17] session-server-side merge - first shot --- .../rollout/generate_hub/agentic_tool_call.py | 77 +++--- .../generate_utils/openai_endpoint_utils.py | 192 +++++++++++++- miles/rollout/session/session_types.py | 19 ++ miles/rollout/session/sessions.py | 244 +++++++++++++++++- 4 files changed, 483 insertions(+), 49 deletions(-) diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 12fca3a7e9..ba4655ce09 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 83f8995b92..645742afb3 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -13,11 +13,36 @@ 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 = 120.0 _HTTP_CONNECT_TIMEOUT = 10.0 @@ -333,6 +358,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 +617,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 6548902093..bff5da91b9 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 4ca6abb324..a6014c6aaa 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 @@ -133,6 +139,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", @@ -171,6 +178,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) @@ -190,6 +377,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 92f32a440631371c48f8bf3bdacb0027719d47d1 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Thu, 9 Jul 2026 19:52:49 -0400 Subject: [PATCH 15/17] 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 ba4655ce09..d89351a3dd 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 2096dd0e5b305326e834e09161b09897fad2ac57 Mon Sep 17 00:00:00 2001 From: Omkar Pangarkar Date: Wed, 15 Jul 2026 15:50:35 -0400 Subject: [PATCH 16/17] 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 a6014c6aaa..27dd94615d 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -189,6 +189,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 @@ -387,17 +400,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, } @@ -406,7 +430,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: @@ -417,7 +441,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 6f534053c2081ee7b0f60c4c25fcbb74d002dea5 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 15 Jul 2026 21:45:59 +0000 Subject: [PATCH 17/17] Make session client limits configurable --- .../generate_utils/openai_endpoint_utils.py | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 645742afb3..b2fd57ea7a 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 json import logging +import os import random from argparse import Namespace from copy import deepcopy @@ -43,12 +44,28 @@ def _expected_routed_experts_bytes(args: Namespace, num_tokens: int) -> int | No return max(num_tokens - 1, 0) * int(num_layers) * int(moe_router_topk) * 4 -_SESSION_REQUEST_TIMEOUT = 120.0 +def _positive_float_env(name: str, default: float) -> float: + raw_value = os.environ.get(name) + value = default if raw_value is None else float(raw_value) + if value <= 0: + raise ValueError(f"{name} must be greater than zero, got {value}") + return value -_HTTP_CONNECT_TIMEOUT = 10.0 -_HTTP_READ_TIMEOUT = 120.0 -_HTTP_WRITE_TIMEOUT = 30.0 -_HTTP_POOL_TIMEOUT = 10.0 + +def _positive_int_env(name: str, default: int) -> int: + raw_value = os.environ.get(name) + value = default if raw_value is None else int(raw_value) + if value <= 0: + raise ValueError(f"{name} must be greater than zero, got {value}") + return value + + +_SESSION_REQUEST_TIMEOUT = _positive_float_env("MILES_SESSION_REQUEST_TIMEOUT_SECONDS", 24 * 60 * 60.0) + +_HTTP_CONNECT_TIMEOUT = _positive_float_env("MILES_SESSION_HTTP_CONNECT_TIMEOUT_SECONDS", 600.0) +_HTTP_READ_TIMEOUT = _positive_float_env("MILES_SESSION_HTTP_READ_TIMEOUT_SECONDS", 24 * 60 * 60.0) +_HTTP_WRITE_TIMEOUT = _positive_float_env("MILES_SESSION_HTTP_WRITE_TIMEOUT_SECONDS", 600.0) +_HTTP_POOL_TIMEOUT = _positive_float_env("MILES_SESSION_HTTP_POOL_TIMEOUT_SECONDS", 600.0) _HEALTH_RETRIES = 2 _CREATE_RETRIES = 10 @@ -59,7 +76,7 @@ def _expected_routed_experts_bytes(args: Namespace, num_tokens: int) -> int | No _BACKOFF_MAX_SECONDS = 10.0 _BACKOFF_JITTER_FRACTION = 0.2 -_COLLECT_RECORDS_CONCURRENCY = 256 +_COLLECT_RECORDS_CONCURRENCY = _positive_int_env("MILES_SESSION_COLLECT_CONCURRENCY", 8192) _COLLECT_RECORDS_SEMAPHORE = asyncio.Semaphore(_COLLECT_RECORDS_CONCURRENCY)