diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index baec979375..0d462d46cd 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -341,6 +341,71 @@ 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]) + base_port = 20000 + (self.index % 100) * 100 + last_exc = None + + for attempt in range(10): + port = find_available_port(base_port + attempt * 100) + + 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, + ) + + 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 + + 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(): + self.process.terminate() + self.process.join(timeout=5) + if self.process.is_alive(): + self.process.kill() + + # --------------------------------------------------------------------------- # RolloutManager # --------------------------------------------------------------------------- @@ -385,7 +450,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 +490,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 +699,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 +857,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 +993,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 +1025,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 +1213,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 +1290,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 +1358,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 +1501,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 diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index 46f89bccba..d89351a3dd 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( @@ -36,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 @@ -59,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} @@ -93,42 +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 = 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), - ) - - 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,12 +128,13 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: sample.status = Sample.Status.ABORTED return GenerateFnOutput(samples=sample) - if not input.args.generate_multi_samples: - samples = merge_samples(samples, input.state.tokenizer) - samples.metadata.update(session_metadata) - else: - samples[-1].metadata.update(session_metadata) - return GenerateFnOutput(samples=samples) + sample = samples[0] + logger.debug( + 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=sample) def build_agent_function_kwargs( diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index e03197d6a5..b2fd57ea7a 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -3,101 +3,688 @@ """ import asyncio +import json import logging +import os 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.rollout.session.session_types import ( + GetMergedSessionResponse, + GetSessionResponse, + MergedSessionSample, + SessionRecord, +) from miles.utils.types import Sample logger = logging.getLogger(__name__) -_SESSION_REQUEST_TIMEOUT = 120 + +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 + + +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 + + +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 +_COLLECT_RETRIES = 3 +_DELETE_RETRIES = 3 + +_BACKOFF_INITIAL_SECONDS = 1.0 +_BACKOFF_MAX_SECONDS = 10.0 +_BACKOFF_JITTER_FRACTION = 0.2 + +_COLLECT_RECORDS_CONCURRENCY = _positive_int_env("MILES_SESSION_COLLECT_CONCURRENCY", 8192) +_COLLECT_RECORDS_SEMAPHORE = asyncio.Semaphore(_COLLECT_RECORDS_CONCURRENCY) 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. + 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( + 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 + + 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, + ) + + 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, + ) + + 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=%.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", + 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, session_server_instance_id=session_server_instance_id, ) - async def collect_records(self) -> tuple[list[SessionRecord], dict]: - try: - response = await asyncio.wait_for( - post(self.base_url, {}, action="get"), - timeout=_SESSION_REQUEST_TIMEOUT, + 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, ) - 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." + + 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", + self.session_id, + self.base_url, + _COLLECT_RECORDS_CONCURRENCY, + ) + + async with _COLLECT_RECORDS_SEMAPHORE: + logger.info( + "[session-client] collect_start session_id=%s url=%s", + self.session_id, + self.base_url, ) - # Still attempt to clean up the session. + try: - await asyncio.wait_for( - post(self.base_url, {}, action="delete"), - timeout=_SESSION_REQUEST_TIMEOUT, + response = await self._request( + "GET", + self.base_url, + phase="collect_records", + max_retries=_COLLECT_RETRIES, + ) + + 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()), ) - 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 + + return records, metadata + + 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.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 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() + + 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, + ) + + 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, ) - except Exception as e: - logger.warning(f"Failed to delete session {self.session_id} after collecting records: {e}") - return (records or []), metadata + +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( @@ -131,10 +718,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 @@ -143,30 +728,21 @@ 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} " - 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 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)} " - 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 @@ -193,13 +769,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"]: @@ -217,6 +797,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, @@ -251,8 +854,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) 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 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 aee310ad56..27dd94615d 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -4,18 +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, - TokenizationError, - UpstreamResponseError, +from miles.rollout.session.session_errors import SessionError, SessionNotFoundError, UpstreamResponseError +from miles.rollout.session.session_types import ( + GetMergedSessionResponse, + GetSessionResponse, + MergedSessionSample, + SessionRecord, ) -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 +105,272 @@ async def create_session(): session_id = registry.create_session() return {"session_id": session_id} + def _compact_output_token_logprobs(output_token_logprobs): + """Keep only the fields consumed by the sample builder. + + Consumer uses: + output_log_probs = [item[0] for item in output_token_logprobs] + output_token_ids = [item[1] for item in output_token_logprobs] + + Some backends may include extra per-token fields; do not return them. + """ + if not output_token_logprobs: + return [] + + compact = [] + for item in output_token_logprobs: + compact.append([item[0], item[1]]) + return compact + + def _compact_session_record(record: SessionRecord) -> SessionRecord: + """Return only the record fields required by compute_samples_from_openai_records.""" + choice = record.response.get("choices", [{}])[0] + meta_info = choice.get("meta_info") or {} + + compact_meta_info = { + "output_token_logprobs": _compact_output_token_logprobs(meta_info.get("output_token_logprobs")), + } + + # Used by _compute_sample_from_openai_record. + if "weight_version" in meta_info: + compact_meta_info["weight_version"] = meta_info["weight_version"] + + # Preserve prefix-cache metadata if present. The consumer does: + # sample.prefix_cache_info.add(choice.get("meta_info", {})) + for key in ( + "cached_tokens", + "prompt_tokens", + "prefix_cache_hit", + "prefix_cache_len", + "prefix_cache_tokens", + "prefix_cache_start_len", + "prefix_cache_end_len", + ): + if key in meta_info: + compact_meta_info[key] = meta_info[key] + + # Preserve routed experts if present. get_rollout_topk_from_response() + # may read this from choice or meta_info depending on implementation. + if "routed_experts" in meta_info: + compact_meta_info["routed_experts"] = meta_info["routed_experts"] + + compact_choice = { + "prompt_token_ids": choice.get("prompt_token_ids", []), + "finish_reason": choice.get("finish_reason"), + "meta_info": compact_meta_info, + } + + if "routed_experts" in choice: + compact_choice["routed_experts"] = choice["routed_experts"] + + return SessionRecord( + timestamp=record.timestamp, + method=record.method, + path=record.path, + status_code=record.status_code, + request={ + # Only used for this assertion: + # request_input_ids == prompt_token_ids + "input_ids": record.request.get("input_ids"), + }, + response={ + "choices": [compact_choice], + }, + ) + + def _status_from_finish_reason(finish_reason: str | None) -> str: + match finish_reason: + case "stop" | "tool_calls": + return "completed" + case "length": + return "truncated" + case "abort": + return "aborted" + case _: + return "completed" + + def _keep_records_until_first_non_completed(records: list[SessionRecord]) -> tuple[list[SessionRecord], int]: + """Match sample_utils.drop_samples_after_first_non_completed semantics. + + 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 + 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) @@ -111,22 +378,84 @@ 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, ) + @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] + 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_to_merge, + accumulated_token_ids, + max_trim_tokens, + ) + metadata = { + "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, + } + + if sample is None: + logger.info( + "[session-server] merged_collect_empty session_id=%s records=%d accumulated_tokens=%d", + session_id, + len(records_to_merge), + 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_to_merge), + 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) 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