diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index ae6c866361..bc7ae7889f 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -2,6 +2,7 @@ import logging import time import uuid +import zlib import orjson from fastapi import Request @@ -112,13 +113,20 @@ async def get_session(session_id: str): raise SessionNotFoundError(f"session not found: session_id={session_id}") return GetSessionResponse(session_id=session_id, records=[], metadata={}) metadata = {} - try: - mismatch = registry.compute_session_mismatch(session) - except TokenizationError: - logger.exception("Failed to compute tito_session_mismatch for session %s", session_id) - mismatch = None - if mismatch is not None: - metadata["tito_session_mismatch"] = mismatch + # Re-tokenizing the full trajectory to check tito_session_mismatch is a + # diagnostic canary, not a training-correctness guard (train tokens come + # from the inference response). It runs synchronously here, so at long + # context it blocks the event loop and starves record retrieval. Sample + # a deterministic fraction by session_id to keep the signal cheap. + sample_rate = getattr(args, "tito_session_mismatch_sample_rate", 0.0) + if sample_rate > 0.0 and (zlib.crc32(session_id.encode()) % 10000) < sample_rate * 10000: + try: + mismatch = registry.compute_session_mismatch(session) + except TokenizationError: + logger.exception("Failed to compute tito_session_mismatch for session %s", session_id) + mismatch = None + if mismatch is not None: + metadata["tito_session_mismatch"] = mismatch metadata["accumulated_token_ids"] = session.token_ids metadata["max_trim_tokens"] = registry.tito_tokenizer.max_trim_tokens return GetSessionResponse( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2d6e43e138..9d696bdb3a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1623,6 +1623,17 @@ def add_session_arguments(parser): choices=[t.value for t in TITOTokenizerType], help="TITO tokenizer type for pretokenized prefix reuse. Controls how token IDs are computed for messages appended after the pretokenized prefix in multi-turn agentic sessions.", ) + parser.add_argument( + "--tito-session-mismatch-sample-rate", + type=float, + default=0.0, + help="Fraction of sessions (0.0-1.0) for which the session server re-tokenizes the full " + "trajectory to compute tito_session_mismatch. This is a diagnostic canary for chat-template/" + "tokenizer drift; it does not affect training tokens (those come from the inference response). " + "It runs synchronously on the record-fetch path, so at long context it blocks the event loop and " + "starves record retrieval. Default 0.0 (off). Set to 0.02 for a cheap drift signal, 1.0 for legacy " + "always-on behaviour.", + ) parser.add_argument( "--tito-allowed-append-roles", nargs="+",