diff --git a/docs/cli.md b/docs/cli.md index 55bdeea06..4a95fc291 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -41,6 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index dbcc1f2fa..154570ec6 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -26,6 +26,7 @@ class BackendInfo: supported_types: frozenset[AttnType] requires_flashinfer: bool = False requires_sgl_kernel: bool = False + requires_sm90: bool = False requires_sm100: bool = False # Allowed page sizes (None -> any). Config-time resolution coerces to the last # entry when the resolved page_size is not in the list. @@ -77,6 +78,7 @@ def create_fi_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL}), requires_sgl_kernel=True, + requires_sm90=True, ), ) def create_fa_backend(config: ModelConfig): diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 6c4e70602..c1cc7a7fb 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -13,6 +13,7 @@ from freetoken.kvcache.kv_host_offload import KVHostOffloader from freetoken.kvcache.linear_state_pool import LinearStatePool from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.tokenizer.detokenize import DecodeStatus @dataclass @@ -103,6 +104,7 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + stop_decode_status: DecodeStatus | None = None def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index fa61d9765..7947c416a 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -88,11 +88,17 @@ class EngineConfig: # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime # rebuild can pin an absolute window instead). swa_full_tokens_ratio: float = 0.2 + # Force one decode step after this many consecutive prefill steps, so a long chunked + # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first + # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE + # first"). At 8 the cost is ~1% of prefill wall time. + decode_interleave_every: int | None = None # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None distributed_timeout: float = 1800.0 # ranks reach the first collective minutes apart on a 100+ GiB offload load + distributed_port: int = 2333 use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None @@ -166,4 +172,4 @@ def max_forward_len(self) -> int: @property def distributed_addr(self) -> str: - return "tcp://127.0.0.1:2333" + return f"tcp://127.0.0.1:{self.distributed_port}" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index e5869c1a9..88c3b1fc8 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -5,6 +5,7 @@ import gc import math import os +import socket from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -128,6 +129,9 @@ def _backend_requirements_met(name: str) -> bool: return False if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available(): return False + # requires_sm90 also admits sm_100 (FA4). Hopper-only would reject B200. + if any(i.requires_sm90 for i in infos) and not (is_sm90_family() or is_sm100_family()): + return False if any(i.requires_sm100 for i in infos) and not is_sm100_family(): return False return True @@ -284,6 +288,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att "not installed. Install it with `pip install 'freetoken[sgl]'` (or " "'freetoken[accel]'), or use --attention-backend triton." ) + if info.requires_sm90 and not (is_sm90_family() or is_sm100_family()): + raise RuntimeError( + f"Attention backend {config.attention_backend!r} requires a compute capability " + "9.x GPU (FA3, sm_90 cubins) or 10.x GPU (FA4). Ada (sm_89) and consumer " + "Blackwell (sm_12x) are not in the published sgl_kernel flash_ops set. " + "Use --attention-backend fi (or triton) instead." + ) if info.requires_sm100 and not is_sm100_family(): raise RuntimeError( f"Attention backend {config.attention_backend!r} requires a compute capability " @@ -611,15 +622,40 @@ def __init__(self, config: EngineConfig): if config.prefill_warmup: self._warmup_prefill() - def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: - if config.tp_info.size == 1 or config.use_pynccl: - torch.distributed.init_process_group( - backend="gloo", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, + def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Store: + """The rendezvous store's C10d server ignores the host it's given and always listens + on every interface, so an unauthenticated TCPStore ends up reachable off-box even + when distributed_addr says 127.0.0.1. Rank 0 pre-binds the listening socket to + loopback itself and hands the fd to TCPStore (master_listen_fd) to force that.""" + timeout = timedelta(seconds=config.distributed_timeout) + if not config.tp_info.is_primary(): + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=False, timeout=timeout, multi_tenant=True, ) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", config.distributed_port)) + sock.listen(1024) + self._distributed_listen_addr = sock.getsockname() + # TCPStore takes ownership of the fd and closes it with the store; detach so the + # Python socket never closes that number again after it has been reused. + return torch.distributed.TCPStore( + "127.0.0.1", config.distributed_port, config.tp_info.size, + is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(), + ) + + def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: + store = self._make_distributed_store(config) + use_gloo = config.tp_info.size == 1 or config.use_pynccl + torch.distributed.init_process_group( + backend="gloo" if use_gloo else "nccl", + rank=config.tp_info.rank, + world_size=config.tp_info.size, + timeout=timedelta(seconds=config.distributed_timeout), + store=store, + ) + if use_gloo: tp_cpu_group = torch.distributed.group.WORLD assert tp_cpu_group is not None max_bytes = ( @@ -627,13 +663,6 @@ def _init_communication(self, config: EngineConfig) -> torch.distributed.Process ) enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes) else: - torch.distributed.init_process_group( - backend="nccl", - rank=config.tp_info.rank, - world_size=config.tp_info.size, - timeout=timedelta(seconds=config.distributed_timeout), - init_method=config.distributed_addr, - ) tp_cpu_group = torch.distributed.new_group(backend="gloo") assert tp_cpu_group is not None return tp_cpu_group diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py index cb4a7f90b..c09800d95 100644 --- a/python/freetoken/kvcache/hybrid_radix_cache.py +++ b/python/freetoken/kvcache/hybrid_radix_cache.py @@ -116,7 +116,10 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, node, prefix_len = self._walk(input_ids) if prefix_len != insert_len: new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], kv_indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), kv_indices[prefix_len:].clone() + ) new_node.set_parent(node) self.full_evictable += new_node.length node = new_node diff --git a/python/freetoken/kvcache/radix_cache.py b/python/freetoken/kvcache/radix_cache.py index 5437286d1..ec0bf75d3 100644 --- a/python/freetoken/kvcache/radix_cache.py +++ b/python/freetoken/kvcache/radix_cache.py @@ -167,7 +167,10 @@ def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> Inser node, prefix_len = self._tree_walk(input_ids) if prefix_len != insert_len: # NOTE: prefix_len < insert_len new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), indices[prefix_len:].clone() + ) new_node.set_parent(node) self.evictable_size += new_node.length node = new_node diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py index 260439ad6..d2cc89c26 100644 --- a/python/freetoken/kvcache/swa_radix_cache.py +++ b/python/freetoken/kvcache/swa_radix_cache.py @@ -215,7 +215,8 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, # then a live (non-tombstone) leaf for the in-window remainder (never a tombstone leaf -- # the clamp + the free_swa -page_size margin guarantee a live tail). if total < insert_len: - suffix_ids = input_ids[total:] + # Both suffix nodes can share this copy without retaining the request buffer. + suffix_ids = input_ids[total:].clone() suffix_kv = kv_indices[total:].clone() boundary = max(0, min(swa_evicted_seqlen, insert_len) - total) boundary = min(boundary, max(0, len(suffix_ids) - self.page_size)) diff --git a/python/freetoken/scheduler/interleave.py b/python/freetoken/scheduler/interleave.py new file mode 100644 index 000000000..76f293c8f --- /dev/null +++ b/python/freetoken/scheduler/interleave.py @@ -0,0 +1,73 @@ +"""Decode-interleave scheduling policy. + +Background (measured on production, 2026-09-15) +----------------------------------------------- +``Scheduler._schedule_next_batch`` picks:: + + batch = (prefill_manager.schedule_next_batch(budget) + or decode_manager.schedule_next_batch()) + +Prefill wins unconditionally. A long prompt is chunked at the window pool's +budget (384 tokens in production), so a 300k-token context is 797 consecutive +prefill steps -- and while they run, *decoding requests already in flight are +never scheduled*. Measured on the live service: + + four long cold prefills: 680/797/506/250 chunks, 267/317/190/93 s + prefill-burst share of a 15-minute window: 70% + +so an in-flight request waits up to ~5 minutes for its next token. The +scheduler already carries the marker for this:: + + # TODO: support other policies: e.g. DECODE first + +Reversing the order outright is wrong -- it starves prefill, and a request that +never prefills never starts. What is needed is *interleaving*: let prefill keep +priority, but force a decode step every N prefill steps so in-flight requests +make progress. The cost is small by construction: measured chunk cost is +0.388 s, and one decode step for a handful of requests is a fraction of that, +so every 8th step costs ~1% of prefill wall time while bounding the stall at +8 chunks (~3 s) instead of the whole burst (~300 s). + +This module holds the *policy* only -- a counter and a decision -- so it can be +unit-tested without a Scheduler, an engine, or a GPU. +""" + +from __future__ import annotations + + +class DecodeInterleavePolicy: + """Decide whether to force a decode step between prefill steps. + + ``every`` is the number of consecutive prefill steps allowed before a decode + step is forced. ``None`` or <= 0 disables the policy, reproducing the + historical prefill-first behaviour bit-for-bit. + + The counter only advances on prefill steps that were actually scheduled, so + a scheduler with nothing to prefill never accumulates credit, and the first + decode of a burst is not delayed. + """ + + def __init__(self, every: int | None = None) -> None: + self.every = int(every) if every else None + self._prefill_streak = 0 + + @property + def enabled(self) -> bool: + return self.every is not None and self.every > 0 + + def note_prefill(self) -> None: + """Record that a prefill step was scheduled in this slot.""" + self._prefill_streak += 1 + + def note_decode(self) -> None: + """Record that a decode step was scheduled in this slot (streak resets).""" + self._prefill_streak = 0 + + def wants_decode(self) -> bool: + """True when the streak has reached the threshold and a decode is due.""" + if not self.enabled: + return False + return self._prefill_streak >= self.every + + def reset(self) -> None: + self._prefill_streak = 0 diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 090e68b97..acb3cf05e 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -32,6 +32,7 @@ from .cache import CacheManager from .config import SchedulerConfig from .decode import DecodeManager +from .interleave import DecodeInterleavePolicy from .io import SchedulerIOMixin from .mm import cut_image_spans, plan_mm_batch from .prefill import ChunkedReq, PrefillManager @@ -92,6 +93,14 @@ def __init__(self, config: SchedulerConfig): ) or getattr(self.engine.kv_cache, "sliding_window_size", None), ) self.decode_manager = DecodeManager(config.page_size) + # Bound how long a run of prefill steps may starve in-flight decodes. A long + # prompt chunked at the window budget is hundreds of consecutive prefill steps; + # measured on an 8xRTX4090 DSV4 deployment, that left an already-decoding + # request unscheduled for 267-317 s. Unset reproduces the historical + # prefill-first order exactly. + self._interleave = DecodeInterleavePolicy( + getattr(config, "decode_interleave_every", None) + ) self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups) self.prefill_manager = PrefillManager( self.cache_manager, @@ -462,16 +471,38 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: self.send_result(reply) def _match_stop_str(self, req: Req) -> str | None: - """First stop string present in this request's generated tail, else None. Decodes - only a short suffix (bounded by the longest stop string's char length, so a stop of - N chars spans at most N tokens) to keep the per-step cost small.""" + """Match stops against the same incrementally decoded text as the frontend.""" + from freetoken.tokenizer.detokenize import DecodeStatus + stop_strs = req.sampling_params.stop_strs prompt_len = req.max_device_len - req.output_len - if len(req.input_ids) <= prompt_len: + end = len(req.input_ids) + # The frontend omits a terminal EOS, including at the output limit with ignore_eos. + # The limit is the delivered length, as in hit_length: under overlap device_len (and + # so can_decode) runs a step ahead of the host. + if ( + end > prompt_len + and end >= req.max_device_len + and int(req.input_ids[-1]) in self.eos_token_ids + ): + end -= 1 + if end <= prompt_len: return None + if req.stop_decode_status is None: + req.stop_decode_status = DecodeStatus( + decoded_ids=[], + decoded_str="", + read_offset=0, + surr_offset=0, + sent_offset=0, + skip_special=req.sampling_params.skip_special_tokens, + ) + state = req.stop_decode_status + state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist()) + tail = state.decode(self.tokenizer) max_chars = max(len(s) for s in stop_strs) - tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1)) - tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist()) + # Bound the text used for matching, retaining token context for incomplete characters. + state.decoded_str = state.decoded_str[-max_chars:] for s in stop_strs: if s in tail: return s @@ -916,11 +947,28 @@ def _gather_multimodal(self, batch: Batch) -> None: ) def _schedule_next_batch(self) -> ForwardInput | None: - # TODO: support other policies: e.g. DECODE first - batch = ( - self.prefill_manager.schedule_next_batch(self.prefill_budget) - or self.decode_manager.schedule_next_batch() - ) + # Prefill keeps priority -- a request that never prefills never starts -- but a + # long run of prefill chunks must not starve in-flight decodes. When the policy + # says a decode is due, take one if a decode is actually runnable; otherwise stay + # with prefill (never spend a slot idle just to keep a promise). + # getattr, not attribute access: the accounting tests drive this method on a + # partially built Scheduler, and interleaving must be an opt-in that a stripped + # object simply does not have rather than something that breaks it. A missing + # policy is the historical prefill-first order. + policy = getattr(self, "_interleave", None) + batch = None + if policy is not None and policy.wants_decode(): + batch = self.decode_manager.schedule_next_batch() + if batch is not None: + policy.note_decode() + if batch is None: + batch = self.prefill_manager.schedule_next_batch(self.prefill_budget) + if batch is not None and policy is not None: + policy.note_prefill() + if batch is None: + batch = self.decode_manager.schedule_next_batch() + if batch is not None and policy is not None: + policy.note_decode() if batch is None: return None forward_input = self._prepare_batch(batch) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 1a039402f..4eaa29682 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -57,6 +57,10 @@ # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor. _SHUTTING_DOWN = threading.Event() BACKEND_DEATH_EXIT_GRACE_S = 10.0 +# The longest a stream stays silent before it sends an SSE comment. A long prompt prefills for +# minutes before its first token, and clients time an idle body out (Node's fetch after 300 s). +STREAM_KEEPALIVE_SECONDS = 15.0 +_STREAM_KEEPALIVE = b": keep-alive\n\n" def get_global_state() -> FrontendManager: @@ -367,21 +371,36 @@ async def stream_generate(self, uid: int): async def stream_with_cancellation(self, generator, request: Request, uid): """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams - several generations into one response): a disconnect aborts every one.""" + several generations into one response): a disconnect aborts every one. While the + generator is silent (a long prefill), send an SSE comment every + ``STREAM_KEEPALIVE_SECONDS`` and check whether the client has gone.""" uids = list(uid) if isinstance(uid, (list, tuple)) else [uid] finished = False + chunks = generator.__aiter__() + pending: asyncio.Future | None = None try: - async for chunk in generator: - # detect if the client has disconnected + while True: + if pending is None: + pending = asyncio.ensure_future(chunks.__anext__()) + done, _ = await asyncio.wait({pending}, timeout=STREAM_KEEPALIVE_SECONDS) if await request.is_disconnected(): logger.info("Client disconnected for user %s", uids) raise asyncio.CancelledError + if not done: + yield _STREAM_KEEPALIVE + continue + try: + chunk = pending.result() + except StopAsyncIteration: + finished = True + return + pending = None yield chunk - finished = True finally: + if pending is not None and not pending.done(): + pending.cancel() # finally, not `except CancelledError`: the abort must go out however the - # stream dies -- the server cancelling the response task (the CancelledError - # lands on wait_for_ack's await, never on the poll above), a late close + # stream dies -- the server cancelling the response task, a late close # delivering GeneratorExit, or an exception out of the generator. Only a # stream that ran to completion leaves the engine with nothing to stop. if not finished: @@ -753,6 +772,30 @@ def _reasoning_geometry(state: Any) -> dict | None: return {"gears": list(gears), "default": default, "kwargs": kwargs} +def kv_pool_geometry(state: Any) -> tuple[int, int]: + """The allocated KV pool as ``(num_pages, page_size)``, most-recent-truth first: the last + rebuild's result, else the running UserReply snapshot, else the load-time ("meta", …) ack. + ``num_pages`` is 0 until one of those lands (still loading / older engine build). + + Split out of cache_geometry so /v1/models can clamp the advertised context to the pool + without recomputing the whole panel payload -- both callers must agree on the resolution + order or they would report different capacities for the same server. + """ + stats = getattr(state, "stats", None) + last = getattr(state, "last_rebuild", None) or {} + pools = getattr(state, "cache_pools", None) or {} + num_pages = int( + last.get("num_pages") + or getattr(stats, "kv_total_pages", 0) + or pools.get("num_pages", 0) + or 0 + ) + page_size = int( + pools.get("page_size", 0) or getattr(getattr(state, "config", None), "page_size", 1) or 1 + ) + return num_pages, page_size + + def cache_geometry(state: Any) -> dict: """Current cache geometry for the desktop cache panel. Each pool size resolves most-recent-truth first: the last rebuild's result, else the running UserReply snapshot @@ -768,11 +811,10 @@ def cache_geometry(state: Any) -> dict: config = state.config last = getattr(state, "last_rebuild", None) or {} pools = getattr(state, "cache_pools", None) or {} - num_pages = int(last.get("num_pages") or tr.kv_total_pages or pools.get("num_pages", 0) or 0) + num_pages, page_size = kv_pool_geometry(state) num_mamba_slots = int( last.get("mamba_slots") or tr.mamba_total_slots or pools.get("num_mamba_slots", 0) or 0 ) - page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) moe_cache_size = last.get("moe_cache_size") if moe_cache_size is None: moe_cache_size = int(pools.get("moe_cache_size", 0) or 0) or configured_moe_cache_size( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 9a463c2bf..900cf46a7 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -109,10 +109,6 @@ def backend_create_detokenizer_link(self) -> bool: def frontend_create_tokenizer_link(self) -> bool: return not self.share_tokenizer - @property - def distributed_addr(self) -> str: - return f"tcp://127.0.0.1:{self.server_port + 1}" - def _json_object(text: str) -> dict: try: @@ -169,6 +165,15 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _valid_port(value: str) -> int: + try: + n = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if not 0 <= n <= 65535: + raise argparse.ArgumentTypeError("must be between 0 and 65535") + return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: from freetoken.gpu_select import gpu_arg @@ -375,6 +380,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--dist-port", + type=_valid_port, + dest="distributed_port", + default=None, + help=( + "Port for the internal TP rendezvous store, loopback-only regardless of --host. " + "Defaults to --port + 1; override when that collides with another instance or " + "service on the same host." + ), + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -428,6 +445,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--decode-interleave-every", + type=int, + default=ServerArgs.decode_interleave_every, + help=( + "Force one decode step after this many consecutive prefill steps. A long " + "prompt is prefilled in chunks, so it occupies that many consecutive " + "scheduler steps; measured on an 8xRTX4090 DSV4 deployment, a 300k-token " + "context left an already-decoding request unscheduled for 267-317 s. " + "Unset keeps the historical prefill-first order (the scheduler's own TODO " + "names this policy)." + ), + ) + parser.add_argument( "--page-size", type=int, @@ -919,6 +950,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if entry: kwargs["quant_backend"] = entry + if kwargs["distributed_port"] is None: + kwargs["distributed_port"] = kwargs["server_port"] + 1 + if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 5c0e60f55..dd653486c 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -44,6 +44,250 @@ ) +# Qwen's decoder intentionally keeps protocol markers visible so reasoning and +# tool-call parsers can consume their (non-special-token) tags. At the semantic +# API waist, after those parsers, only transport/internal multimodal markers are +# suppressed. Grounding markers (object_ref/box/quad), think tags, tool tags, +# and FIM/repository markers are deliberately absent from this list. +_QWEN_SEMANTIC_HIDDEN_TOKENS = ( + "<|endoftext|>", + "<|im_start|>", + "<|im_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", + "<|audio_start|>", + "<|audio_end|>", + "<|audio_pad|>", + "", + "", + "", + "", +) + + +def _uses_qwen_semantic_protocol(state: Any) -> bool: + """Whether this server is configured for a Qwen semantic output dialect. + + Parser selection is the stable frontend-visible family signal available in + ``ServerArgs``. Either half is enough because users may explicitly disable + reasoning or omit tools for an individual request. + """ + config = state.config + return ( + getattr(config, "reasoning_parser", None) == "qwen3" + or getattr(config, "tool_call_parser", None) in {"qwen", "qwen25", "qwen3_coder"} + ) + + +class _SemanticSpecialTokenFilter: + """Request-local, chunk-safe suppression of semantic protocol markers. + + A proper marker prefix is retained until the next detokenizer delta, so a + marker split at any byte boundary cannot leak. A marker inside a quoted/code + literal is retained only once the matching closer arrives; an unclosed quote + is sanitized at end-of-stream. Structured events can be held inside either + ambiguity so their ordering relative to surrounding content stays exact. + """ + + def __init__(self, tokens: tuple[str, ...]) -> None: + self._tokens = tokens + self._quote: str | None = None + self._escaped = False + self._last_char = "" + self._candidate_text = "" + self._candidate_items: list[Any] = [] + self._candidate_quoted = False + self._literal_items: list[Any] | None = None + + @staticmethod + def _append_item(items: list[Any], item: Any) -> None: + if item == "": + return + if isinstance(item, str) and items and isinstance(items[-1], str): + items[-1] += item + else: + items.append(item) + + @classmethod + def _strip_items(cls, items: list[Any], tokens: tuple[str, ...]) -> list[Any]: + """Strip markers across text segments while preserving event positions.""" + out: list[Any] = [] + candidate_text = "" + candidate_items: list[Any] = [] + + def flush(*, keep_text: bool) -> None: + nonlocal candidate_text, candidate_items + for pending in candidate_items: + if keep_text or not isinstance(pending, str): + cls._append_item(out, pending) + candidate_text = "" + candidate_items = [] + + for item in items: + if not isinstance(item, str): + if candidate_text: + cls._append_item(candidate_items, item) + else: + cls._append_item(out, item) + continue + index = 0 + while index < len(item): + ch = item[index] + if not candidate_text: + if any(token.startswith(ch) for token in tokens): + candidate_text = ch + cls._append_item(candidate_items, ch) + else: + cls._append_item(out, ch) + index += 1 + continue + extended = candidate_text + ch + if extended in tokens: + candidate_text = extended + cls._append_item(candidate_items, ch) + flush(keep_text=False) + index += 1 + elif any(token.startswith(extended) for token in tokens): + candidate_text = extended + cls._append_item(candidate_items, ch) + index += 1 + else: + flush(keep_text=True) + flush(keep_text=True) + return out + + def _start_candidate(self, ch: str, *, quoted: bool) -> None: + self._candidate_text = ch + self._candidate_items = [ch] + self._candidate_quoted = quoted + + def _extend_candidate(self, ch: str) -> tuple[list[Any], bool]: + extended = self._candidate_text + ch + if extended in self._tokens: + self._candidate_text = extended + self._append_item(self._candidate_items, ch) + if self._candidate_quoted: + self._literal_items = self._candidate_items + out: list[Any] = [] + else: + out = [item for item in self._candidate_items if not isinstance(item, str)] + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + return out, True + if any(token.startswith(extended) for token in self._tokens): + self._candidate_text = extended + self._append_item(self._candidate_items, ch) + return [], True + out = self._candidate_items + self._last_char = self._candidate_text[-1] + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + return out, False + + def feed_items(self, text: str, *, final: bool = False) -> list[Any]: + if not self._tokens: + return [text] if text else [] + out: list[Any] = [] + index = 0 + while index < len(text): + ch = text[index] + if self._literal_items is not None: + self._append_item(self._literal_items, ch) + if self._escaped: + self._escaped = False + elif ch == "\\": + self._escaped = True + elif ch == self._quote: + out.extend(self._literal_items) + self._literal_items = None + self._quote = None + self._last_char = ch + index += 1 + continue + + if self._candidate_text: + emitted, consumed = self._extend_candidate(ch) + out.extend(emitted) + if consumed: + index += 1 + continue + + if self._quote is not None: + if any(token.startswith(ch) for token in self._tokens): + self._escaped = False + self._start_candidate(ch, quoted=True) + index += 1 + continue + if self._escaped: + self._escaped = False + elif ch == "\\": + self._escaped = True + elif ch == self._quote: + self._quote = None + self._append_item(out, ch) + self._last_char = ch + index += 1 + continue + + opens_quote = ch in {'"', "`", "“", "‘"} or ( + ch == "'" and not self._last_char.isalnum() + ) + if opens_quote: + self._quote = {"“": "”", "‘": "’"}.get(ch, ch) + self._append_item(out, ch) + elif any(token.startswith(ch) for token in self._tokens): + self._start_candidate(ch, quoted=False) + index += 1 + continue + else: + self._append_item(out, ch) + self._last_char = ch + index += 1 + + if final: + if self._literal_items is not None: + out.extend(self._strip_items(self._literal_items, self._tokens)) + self._literal_items = None + elif self._candidate_text: + out.extend(self._candidate_items) + self._candidate_text = "" + self._candidate_items = [] + self._candidate_quoted = False + self._quote = None + self._escaped = False + return out + + def feed(self, text: str, *, final: bool = False) -> str: + items = self.feed_items(text, final=final) + assert all(isinstance(item, str) for item in items) + return "".join(items) + + def boundary(self, item: Any) -> list[Any]: + if self._literal_items is not None: + self._append_item(self._literal_items, item) + return [] + if self._candidate_text: + self._append_item(self._candidate_items, item) + return [] + return [item] + + def finish(self) -> str: + return self.feed("", final=True) + + def finish_items(self) -> list[Any]: + return self.feed_items("", final=True) + + +def _make_semantic_special_token_filter(state: Any) -> _SemanticSpecialTokenFilter: + tokens = _QWEN_SEMANTIC_HIDDEN_TOKENS if _uses_qwen_semantic_protocol(state) else () + return _SemanticSpecialTokenFilter(tokens) + + class GenerationError(ValueError): """A request failed before producing output (surfaced via ``UserReply.error`` — e.g. a chat template the tokenizer cannot render, or a prompt that exceeds the KV budget). Each @@ -783,6 +1027,14 @@ async def _generate_events_impl( parse_tools = spec.parse_tools reasoning_parser = _make_reasoning_parser(spec, state) specials = _leaked_special_tokens(state) + reasoning_special_filter = _make_semantic_special_token_filter(state) + content_special_filter = _make_semantic_special_token_filter(state) + + def _filter_reasoning(piece: str) -> str: + return reasoning_special_filter.feed(strip_special_tokens(piece, specials)) + + def _filter_content(piece: str) -> str: + return content_special_filter.feed(strip_special_tokens(piece, specials)) def _content_delta(text: str) -> ContentDelta: nonlocal pending_logprobs @@ -818,6 +1070,28 @@ def _content_delta(text: str) -> ContentDelta: # chunk is a legitimate word gap. suppress_ws = False + def _semantic_content_events(items: list[Any]) -> list[GenEvent]: + nonlocal suppress_ws + out: list[GenEvent] = [] + for item in items: + if isinstance(item, str): + if item and not (item.strip() == "" and suppress_ws): + out.append(_content_delta(item)) + if item.strip(): + suppress_ws = False + else: + out.append(item) + if isinstance(item, ToolCallsDelta): + suppress_ws = True + return out + + def _filter_tool_text(piece: str) -> list[GenEvent]: + text = strip_special_tokens(piece, specials) + return _semantic_content_events(content_special_filter.feed_items(text)) + + def _filter_tool_boundary(event: GenEvent) -> list[GenEvent]: + return _semantic_content_events(content_special_filter.boundary(event)) + def _close_open_call() -> ToolCallsDelta | None: nonlocal open_call, calls_emitted if open_call is None: @@ -833,8 +1107,6 @@ def _close_open_call() -> ToolCallsDelta | None: ) open_call = None calls_emitted += 1 - nonlocal suppress_ws - suppress_ws = True return ToolCallsDelta([call]) def _route_tool_text(piece: str) -> list[GenEvent]: @@ -848,12 +1120,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]: # close it first so wire order matches generation order. done = _close_open_call() if done is not None: - out.append(done) - stripped = strip_special_tokens(payload, specials) - if stripped and not (stripped.strip() == "" and suppress_ws): - out.append(_content_delta(stripped)) - if stripped.strip(): - suppress_ws = False + out.extend(_filter_tool_boundary(done)) + out.extend(_filter_tool_text(payload)) continue for frag in payload: starts_new = frag.name is not None and ( @@ -862,7 +1130,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: if starts_new: done = _close_open_call() if done is not None: - out.append(done) + out.extend(_filter_tool_boundary(done)) if open_call is None or starts_new: open_call = { "detector_index": frag.tool_index, @@ -870,14 +1138,16 @@ def _route_tool_text(piece: str) -> list[GenEvent]: "params": "", "ordinal": calls_emitted, } - out.append(ToolCallStart( - tool_index=calls_emitted, name=frag.name, args_prefix_stable=frag_stable, - )) + out.extend(_filter_tool_boundary(ToolCallStart( + tool_index=calls_emitted, + name=frag.name, + args_prefix_stable=frag_stable, + ))) if frag.parameters: open_call["params"] += frag.parameters - out.append( + out.extend(_filter_tool_boundary( ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) - ) + )) return out engine_finish_reason: str | None = None @@ -896,7 +1166,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: if reasoning_parser is not None and content_delta: reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta) if reasoning_delta: - stripped_reasoning = strip_special_tokens(reasoning_delta, specials) + stripped_reasoning = _filter_reasoning(reasoning_delta) if stripped_reasoning: # a bare special token must not open a thinking block yield ReasoningDelta(stripped_reasoning) if content_delta: @@ -906,7 +1176,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: elif parse_tools: pending += content_delta else: - yield _content_delta(strip_special_tokens(content_delta, specials)) + stripped_content = _filter_content(content_delta) + if stripped_content: + yield _content_delta(stripped_content) if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) engine_matched_stop = getattr(ack, "matched_stop", None) @@ -917,7 +1189,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: if reasoning_parser is not None: flush_reasoning, flush_content = reasoning_parser.flush() if flush_reasoning: - stripped_reasoning = strip_special_tokens(flush_reasoning, specials) + stripped_reasoning = _filter_reasoning(flush_reasoning) if stripped_reasoning: yield ReasoningDelta(stripped_reasoning) if flush_content: @@ -927,7 +1199,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: elif parse_tools: pending += flush_content else: - yield _content_delta(strip_special_tokens(flush_content, specials)) + stripped_content = _filter_content(flush_content) + if stripped_content: + yield _content_delta(stripped_content) # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins. finish_reason = engine_finish_reason or "stop" @@ -941,35 +1215,60 @@ def _route_tool_text(piece: str) -> list[GenEvent]: for frag in tool_parser.finalize_stream(): if open_call is not None and frag.parameters: open_call["params"] += frag.parameters - yield ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) + for ev in _filter_tool_boundary( + ToolCallArgsDelta(tool_index=open_call["ordinal"], fragment=frag.parameters) + ): + yield ev done = _close_open_call() if done is not None: - yield done + for ev in _filter_tool_boundary(done): + yield ev for item in tool_parser.recover_truncated_call(): - yield ToolCallStart(tool_index=calls_emitted, name=item.name) - yield ToolCallsDelta( - [ToolCallItem(tool_index=calls_emitted, name=item.name, parameters=item.parameters)] - ) + recovered = [ + ToolCallStart(tool_index=calls_emitted, name=item.name), + ToolCallsDelta([ + ToolCallItem( + tool_index=calls_emitted, + name=item.name, + parameters=item.parameters, + ) + ]), + ] + for event in recovered: + for ev in _filter_tool_boundary(event): + yield ev calls_emitted += 1 residual = tool_parser.finish_stream() if residual: - stripped = strip_special_tokens(residual, specials) - if stripped and not (stripped.strip() == "" and suppress_ws): - yield _content_delta(stripped) + for ev in _filter_tool_text(residual): + yield ev if calls_emitted and finish_reason != "length": finish_reason = "tool_calls" else: parsed = _parse_tool_response(pending, spec, state) if parse_tools else None if parsed is not None: normal_text, tool_calls = parsed - normal_text = strip_special_tokens(normal_text, specials) + normal_text = _filter_content(normal_text) if normal_text: yield _content_delta(normal_text) yield ToolCallsDelta(tool_calls) if finish_reason != "length": finish_reason = "tool_calls" elif parse_tools and pending: - yield _content_delta(strip_special_tokens(pending, specials)) + stripped_pending = _filter_content(pending) + if stripped_pending: + yield _content_delta(stripped_pending) + + trailing_reasoning = reasoning_special_filter.finish() + if trailing_reasoning: + yield ReasoningDelta(trailing_reasoning) + if tool_parser is not None: + for ev in _semantic_content_events(content_special_filter.finish_items()): + yield ev + else: + trailing_content = content_special_filter.finish() + if trailing_content: + yield _content_delta(trailing_content) # Entries without a content delta are intentionally dropped in streaming mode. yield GenDone( @@ -1018,9 +1317,15 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer finish_reason = "tool_calls" specials = _leaked_special_tokens(state) + reasoning_special_filter = _make_semantic_special_token_filter(state) + content_special_filter = _make_semantic_special_token_filter(state) return GenResult( - reasoning=strip_special_tokens(reasoning_text, specials).strip(), - content=strip_special_tokens(content_text, specials), + reasoning=reasoning_special_filter.feed( + strip_special_tokens(reasoning_text, specials), final=True + ).strip(), + content=content_special_filter.feed( + strip_special_tokens(content_text, specials), final=True + ), tool_calls=tool_calls, finish_reason=finish_reason, prompt_tokens=prompt_tokens, diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f4..ef4c40cb3 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -52,7 +52,18 @@ def _run_tokenize_worker(detach: bool, **kwargs) -> None: _detach_process_group() from freetoken.tokenizer import tokenize_worker - tokenize_worker(**kwargs) + ack_queue = kwargs.get("ack_queue") + try: + tokenize_worker(**kwargs) + except Exception as exc: # noqa: BLE001 -- surface the reason, then let it propagate + # Same contract as _run_scheduler: without this, a tokenizer/detokenizer that dies + # before its ready ack reaches the parent only as a dead process, so the supervisor + # reports the generic "backend worker freetoken-detokenizer-0 exited during load" and + # the real cause (an unreadable tokenizer.json, a missing checkpoint file, a ZMQ bind + # clash) is lost. The traceback still prints and the process still exits non-zero. + if ack_queue is not None: + _report_startup_error(ack_queue, exc) + raise def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 9ce63060a..b09c480e1 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -1137,10 +1137,31 @@ def _served_model_name(state: Any) -> str: def _model_context_length(state: Any) -> int | None: - """The model ceiling, not `min(ceiling, KV budget)`: a rebuild moves the latter, and agents - read this once at startup.""" + """The usable context: ``min(model ceiling, allocated KV tokens)``. + + This is the clamp the engine already applies to its own ``max_seq_len`` (Engine.__init__, + and again in _refresh_seq_state after a rebuild), and therefore the limit the scheduler + admits against. The frontend process holds an unclamped ServerArgs copy, so publishing + ``config.max_seq_len`` here advertised the model's positional ceiling instead -- e.g. 262144 + against a 178176-token pool, so the engine and this route disagreed about one quantity. + + `ft launch` reads this to size each agent's context window (opencode's ``limit.context``, + codex's ``context_window``, ``CLAUDE_CODE_MAX_CONTEXT_TOKENS``), so overstating it stops + those agents compacting before the KV pool runs out. + """ try: # never 500 a metadata route: max_seq_len walks into the HF config on some builds value = int(state.config.max_seq_len) except Exception: # noqa: BLE001 return None - return value if value > 0 else None + if value <= 0: + return None + try: + # Local import: api_server imports this module, so a module-level one would cycle. + from .api_server import kv_pool_geometry + + num_pages, page_size = kv_pool_geometry(state) + kv_tokens = num_pages * page_size + except Exception: # noqa: BLE001 + kv_tokens = 0 + # 0 before the ("meta", …) ack lands; the ceiling is the best available answer until then. + return min(value, kv_tokens) if kv_tokens > 0 else value diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index 15bde4fa0..2eb3029df 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -74,6 +74,24 @@ class DecodeStatus: skip_special: bool = False # decode this request with skip_special_tokens think_end_at: int | None = None # generated tokens up to and incl. the reasoning end tag + def update(self, read_str: str, surr_str: str) -> str: + new_text = read_str[len(surr_str) :] + if len(new_text) > 0 and not new_text.endswith("�"): + self.decoded_str += new_text + self.surr_offset = self.read_offset + self.read_offset = len(self.decoded_ids) + return self.decoded_str + return self.decoded_str + find_printable_text(new_text) + + def decode(self, tokenizer: PreTrainedTokenizerBase) -> str: + skip = self.skip_special + return self.update( + tokenizer.decode(self.decoded_ids[self.surr_offset :], skip_special_tokens=skip), + tokenizer.decode( + self.decoded_ids[self.surr_offset : self.read_offset], skip_special_tokens=skip + ), + ) + class DetokenizeManager: def __init__( @@ -116,6 +134,10 @@ def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]: def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]: """Incremental text per message plus, for a finished message, its reasoning token count (0 otherwise or when the model emitted no reasoning end tag).""" + # Each message must advance its request's decode state before the next one. + if len({msg.uid for msg in msgs}) != len(msgs): + parts = [self.detokenize_with_meta([msg]) for msg in msgs] + return [p[0][0] for p in parts], [p[1][0] for p in parts] read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] skip: List[bool] = [] @@ -130,7 +152,9 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li skip_special=bool(getattr(msg, "skip_special_tokens", False)), ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): + if not ( + msg.finished and not msg.matched_stop and msg.next_token in self.eos_token_ids + ): s.decoded_ids.append(msg.next_token) if ( self.think_end_id is not None @@ -149,16 +173,7 @@ def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], Li reasoning_tokens: List[int] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] - new_text = read_str[len(surr_str) :] - # Streaming chunk: update the decode status - if len(new_text) > 0 and not new_text.endswith("�"): - output_str = s.decoded_str + new_text - s.decoded_str = output_str - s.surr_offset = s.read_offset - s.read_offset = len(s.decoded_ids) - else: - new_text = find_printable_text(new_text) - output_str = s.decoded_str + new_text + output_str = s.update(read_str, surr_str) prev_sent = s.sent_offset if msg.finished: diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py index 9271bd425..44f9f8777 100644 --- a/tests/engine/test_attention_backend_arch.py +++ b/tests/engine/test_attention_backend_arch.py @@ -110,6 +110,42 @@ def test_explicit_trtllm_allowed_on_sm100_family(monkeypatch): assert config.page_size in (16, 32, 64) +def test_explicit_fa_rejected_outside_sm90_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=8) + config = _engine_config(attention_backend="fa") + with pytest.raises(RuntimeError, match="9.x"): + _adjust_config(config) + + +def test_explicit_fa_allowed_on_sm90_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=9) + config = _engine_config(attention_backend="fa") + _adjust_config(config) + assert config.attention_backend == "fa" + + +def test_explicit_fa_allowed_on_sm100_family(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=10) + config = _engine_config(attention_backend="fa") + _adjust_config(config) + assert config.attention_backend == "fa" + + +def test_explicit_fa_rejected_on_consumer_blackwell(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch, major=12) + config = _engine_config(attention_backend="fa") + with pytest.raises(RuntimeError, match="9.x"): + _adjust_config(config) + + def test_fa_backend_rejects_cc12(monkeypatch): from freetoken.attention.fa import FlashAttentionBackend diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py new file mode 100644 index 000000000..4fc247376 --- /dev/null +++ b/tests/engine/test_distributed_store.py @@ -0,0 +1,113 @@ +"""Regression test for #301: the TP rendezvous store must not fall back to a wildcard +bind. torch's tcp:// rendezvous handler ignores the host it's given when binding the +master's listening socket and always listens on every interface, so +init_method="tcp://127.0.0.1:PORT" looks loopback-only but isn't -- see +Engine._make_distributed_store. This drives the real store-construction path across two +spawned local ranks (CPU/gloo) and checks the actual bind address, not just the URL. +""" + +from __future__ import annotations + +import gc +import multiprocessing as mp +import os +from datetime import timedelta +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +from freetoken.distributed import DistributedInfo +from freetoken.engine.engine import Engine + +WORLD_SIZE = 2 +TIMEOUT = timedelta(seconds=15) + + +def _config(rank: int, port: int) -> SimpleNamespace: + return SimpleNamespace( + tp_info=DistributedInfo(rank, WORLD_SIZE), + distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + + +def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None: + engine = Engine.__new__(Engine) # only _make_distributed_store is under test + store = engine._make_distributed_store(_config(rank, port)) + dist.init_process_group( + backend="gloo", rank=rank, world_size=WORLD_SIZE, timeout=TIMEOUT, store=store + ) + total = torch.tensor([float(rank)]) + dist.all_reduce(total) + bind_host = engine._distributed_listen_addr[0] if rank == 0 else None + result_q.put((rank, total.item(), bind_host)) + dist.destroy_process_group() + + +def test_two_ranks_rendezvous_with_loopback_only_store(): + ctx = mp.get_context("spawn") + port = 29511 + result_q = ctx.Queue() + procs = [ctx.Process(target=_run_rank, args=(rank, port, result_q)) for rank in range(WORLD_SIZE)] + for p in procs: + p.start() + + results = {} + for _ in procs: + rank, total, bind_host = result_q.get(timeout=30) + results[rank] = (total, bind_host) + for p in procs: + p.join(timeout=30) + assert p.exitcode == 0 + + # both ranks completed the same all_reduce over the store -- rendezvous worked + assert results[0][0] == results[1][0] == 1.0 # sum(0, 1) + # the master's listening socket was pre-bound to loopback, not the wildcard + # address the C10d TCPStore server binds to on its own + assert results[0][1] == "127.0.0.1" + + +def _teardown_rank(port: int, result_q: mp.Queue) -> None: + listen_fds = [] + real_store = dist.TCPStore + + def recording_store(*args, **kwargs): + listen_fds.append(kwargs.get("master_listen_fd")) + return real_store(*args, **kwargs) + + dist.TCPStore = recording_store + engine = Engine.__new__(Engine) + config = SimpleNamespace( + tp_info=DistributedInfo(0, 1), distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + store = engine._make_distributed_store(config) + dist.init_process_group(backend="gloo", rank=0, world_size=1, timeout=TIMEOUT, store=store) + dist.destroy_process_group() + del store + gc.collect() + # the listening fd number is reused by an unrelated file, as a later open() would + fd = listen_fds[0] + r, w = os.pipe() + os.dup2(r, fd) + del engine + gc.collect() + try: + os.fstat(fd) + result_q.put("open") + except OSError as exc: + result_q.put(f"closed: {exc}") + + +def test_the_engine_does_not_close_the_listen_fd_the_store_owns(): + # TCPStore takes ownership of master_listen_fd and closes it itself; an engine that also + # kept the Python socket closed whatever file reused that number when it was freed + ctx = mp.get_context("spawn") + result_q = ctx.Queue() + proc = ctx.Process(target=_teardown_rank, args=(29513, result_q)) + proc.start() + outcome = result_q.get(timeout=30) + proc.join(timeout=30) + assert proc.exitcode == 0 + assert outcome == "open" diff --git a/tests/scheduler/test_interleave.py b/tests/scheduler/test_interleave.py new file mode 100644 index 000000000..42a3f7221 --- /dev/null +++ b/tests/scheduler/test_interleave.py @@ -0,0 +1,346 @@ +"""Tests for DecodeInterleavePolicy and its wiring into Scheduler._schedule_next_batch. + +The production symptom this exists for (measured 2026-09-15): a 300k-token prompt +is chunked at the window pool budget (384 tok), so it becomes 797 consecutive +prefill steps, and an already-decoding request is not scheduled once during the +267-317 s that takes. Prefill must keep priority (a request that never prefills +never starts) but must yield to decode periodically. +""" + +from __future__ import annotations + +import pytest + +from freetoken.scheduler.interleave import DecodeInterleavePolicy + + +# --------------------------------------------------------------------------- # +# policy, in isolation +# --------------------------------------------------------------------------- # +def test_disabled_by_default_and_never_forces_decode(): + """Unset reproduces historical prefill-first behaviour exactly.""" + p = DecodeInterleavePolicy() + assert not p.enabled + for _ in range(10_000): + p.note_prefill() + assert not p.wants_decode() + + +@pytest.mark.parametrize("every", [1, 2, 8, 16]) +def test_forces_decode_after_exactly_every_prefills(every): + p = DecodeInterleavePolicy(every) + assert p.enabled + for i in range(1, every): + p.note_prefill() + assert not p.wants_decode(), f"fired early at {i} < {every}" + p.note_prefill() + assert p.wants_decode(), f"did not fire at {every}" + + +def test_decode_resets_the_streak(): + p = DecodeInterleavePolicy(4) + for _ in range(4): + p.note_prefill() + assert p.wants_decode() + p.note_decode() + assert not p.wants_decode(), "streak must reset after a decode step" + for _ in range(3): + p.note_prefill() + assert not p.wants_decode() + p.note_prefill() + assert p.wants_decode() + + +def test_every_one_interleaves_strictly(): + """every=1 is the extreme: decode after every single prefill step.""" + p = DecodeInterleavePolicy(1) + p.note_prefill() + assert p.wants_decode() + + +@pytest.mark.parametrize("bad", [0, -1, -100]) +def test_non_positive_disables(bad): + p = DecodeInterleavePolicy(bad) + assert not p.enabled + for _ in range(50): + p.note_prefill() + assert not p.wants_decode() + + +def test_reset_clears_credit(): + p = DecodeInterleavePolicy(8) + for _ in range(7): + p.note_prefill() + p.reset() + assert not p.wants_decode() + p.note_prefill() + assert not p.wants_decode() + + +# --------------------------------------------------------------------------- # +# the decision the scheduler actually makes +# --------------------------------------------------------------------------- # +def _decide(policy, has_prefill, has_decode, every=8): + """Mirror of the scheduler's new decision, as a pure function. + + prefill wins unless the policy says decode is due AND a decode is available. + """ + if has_prefill and not (policy.wants_decode() and has_decode): + policy.note_prefill() + return "prefill" + if has_decode: + policy.note_decode() + return "decode" + if has_prefill: + policy.note_prefill() + return "prefill" + return None + + +def test_burst_is_punctuated_by_decode_every_n(): + """797 prefill steps in a row -- the production shape -- must contain decode steps.""" + policy = DecodeInterleavePolicy(8) + seq = [] + for _ in range(797): + seq.append(_decide(policy, has_prefill=True, has_decode=True)) + n_dec = seq.count("decode") + # A decode step resets the streak, so a cycle is every+1 slots (8 prefill + 1 decode): + # 797 slots hold 88 full cycles, i.e. 88 decode steps. + assert n_dec == 797 // 9, f"expected {797 // 9} decode steps, got {n_dec}" + # and no run of prefills exceeds the threshold + run = 0 + worst = 0 + for s in seq: + run = run + 1 if s == "prefill" else 0 + worst = max(worst, run) + assert worst == 8, f"longest prefill run {worst} exceeds the threshold" + + +def test_disabled_policy_yields_the_historical_sequence(): + policy = DecodeInterleavePolicy() + seq = [_decide(policy, True, True) for _ in range(100)] + assert seq.count("decode") == 0, "disabled must never take a decode slot" + + +def test_decode_is_not_forced_when_none_is_runnable(): + """Nothing to decode -> prefill keeps the slot (no idle step).""" + policy = DecodeInterleavePolicy(2) + seq = [_decide(policy, has_prefill=True, has_decode=False) for _ in range(10)] + assert seq == ["prefill"] * 10 + + +def test_prefill_keeps_priority_below_the_threshold(): + """The first N-1 slots of a burst must still go to prefill.""" + policy = DecodeInterleavePolicy(8) + seq = [_decide(policy, True, True) for _ in range(7)] + assert seq.count("decode") == 0 + + +def test_pure_decode_phase_unaffected(): + """With nothing to prefill, every slot is decode (no policy interference).""" + policy = DecodeInterleavePolicy(8) + seq = [_decide(policy, has_prefill=False, has_decode=True) for _ in range(50)] + assert seq == ["decode"] * 50 + + +def test_scheduler_without_policy_keeps_prefill_first(): + """A Scheduler built without __init__ (as the accounting tests do) must not break. + + _schedule_next_batch reads the policy through getattr because the accounting tests + build a Scheduler this way; a stripped object keeps the historical prefill-first + order instead of raising. + """ + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + # no _interleave attribute at all + assert not hasattr(s, "_interleave") + assert Scheduler._schedule_next_batch(s) is prefill_batch, "must stay prefill-first" + + +def test_scheduler_with_policy_and_no_decode_stays_prefill(): + """When no decode is runnable the policy must not steal a slot.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + s._interleave = DecodeInterleavePolicy(1) # fires every step + for _ in range(5): + assert Scheduler._schedule_next_batch(s) is prefill_batch + + +def test_scheduler_with_policy_takes_decode_when_due(): + """Once the streak reaches the threshold, a runnable decode gets the slot.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + prefill_batch = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + decode_batch = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: prefill_batch) + s.decode_manager = SimpleNamespace(schedule_next_batch=lambda: decode_batch) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + s._interleave = DecodeInterleavePolicy(3) + got = [Scheduler._schedule_next_batch(s) for _ in range(4)] + assert got[0] is prefill_batch and got[1] is prefill_batch and got[2] is prefill_batch + assert got[3] is decode_batch, "4th slot must be the forced decode" + + +# --------------------------------------------------------------------------- # +# Truth table against the REAL scheduler (codex review: _decide is a model of +# the decision, not the decision -- it can drift from the code it describes). +# --------------------------------------------------------------------------- # +def _real_scheduler(every=None, prefill=True, decode=True): + """Build a Scheduler whose two managers are stubs with call counters.""" + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + from freetoken.scheduler.interleave import DecodeInterleavePolicy + + calls = {"prefill": 0, "decode": 0} + pf = SimpleNamespace(is_prefill=True, prompt_admissions=[]) + dc = SimpleNamespace(is_prefill=False, prompt_admissions=[]) + + def _pf(budget): + calls["prefill"] += 1 + return pf if prefill else None + + def _dc(): + calls["decode"] += 1 + return dc if decode else None + + s = Scheduler.__new__(Scheduler) + s.prefill_budget = 384 + s.prefill_manager = SimpleNamespace(schedule_next_batch=_pf) + s.decode_manager = SimpleNamespace(schedule_next_batch=_dc) + s._prepare_batch = lambda value: value + s.send_result = lambda messages: None + if every is not None: + s._interleave = DecodeInterleavePolicy(every) + return s, calls, pf, dc + + +@pytest.mark.parametrize("every,expected", [ + # (policy N, expected sequence of slot winners over 9 slots) + (None, ["p"] * 9), # disabled -> always prefill + # streak starts at 0, so slot 0 is prefill; then strict alternation + (1, ["p", "d", "p", "d", "p", "d", "p", "d", "p"]), + (3, ["p", "p", "p", "d", "p", "p", "p", "d", "p"]), + (8, ["p"] * 8 + ["d"]), +]) +def test_truth_table_against_real_scheduler(every, expected): + """The real _schedule_next_batch must produce exactly this slot sequence.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=every) + got = [] + for _ in range(9): + b = Scheduler._schedule_next_batch(s) + got.append("p" if b is pf else ("d" if b is dc else "?")) + assert got == expected, f"every={every}: got {got}, want {expected}" + + +def test_every_one_alternates_while_both_stay_runnable(): + """every=1 is the degenerate case: one decode between every prefill.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=1) + seq = [] + for _ in range(6): + b = Scheduler._schedule_next_batch(s) + seq.append("p" if b is pf else "d") + # slot 0 prefill (streak starts at 0), then alternation + assert seq == ["p", "d", "p", "d", "p", "d"] + + +def test_latch_releases_when_decode_becomes_available(): + """Due decode unavailable -> keep prefilling -> decode becomes available -> it fires. + + Exercises the subtle path: the policy must not lose its 'decode is due' state + while waiting, and must not fire twice once it does. + """ + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=2) + # phase 1: no decode runnable at all + s.decode_manager.schedule_next_batch = lambda: None + for _ in range(5): + assert Scheduler._schedule_next_batch(s) is pf + # phase 2: decode becomes available -> next slot must be decode + s.decode_manager.schedule_next_batch = lambda: dc + assert Scheduler._schedule_next_batch(s) is dc, "latch did not release" + + +def test_no_work_returns_none(): + from freetoken.scheduler.scheduler import Scheduler + + s, _, _, _ = _real_scheduler(every=8, prefill=False, decode=False) + assert Scheduler._schedule_next_batch(s) is None + + +def test_pure_decode_fallback_when_nothing_to_prefill(): + """With no prefill available every slot is decode (no policy interference).""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=2, prefill=False, decode=True) + got = [Scheduler._schedule_next_batch(s) for _ in range(5)] + assert all(b is dc for b in got) + + +def test_disabled_policy_object_not_just_missing_attribute(): + """codex: cover a policy that EXISTS but is disabled, not only an absent one.""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=None) + s._interleave = __import__( + "freetoken.scheduler.interleave", fromlist=["DecodeInterleavePolicy"] + ).DecodeInterleavePolicy(0) # present but disabled + got = [Scheduler._schedule_next_batch(s) for _ in range(6)] + assert all(b is pf for b in got), "a disabled policy must never take a slot" + + +def test_decode_manager_not_probed_when_policy_disabled(): + """When disabled, the decode manager must not be consulted at all (no extra call).""" + from freetoken.scheduler.scheduler import Scheduler + + s, calls, pf, dc = _real_scheduler(every=None) + Scheduler._schedule_next_batch(s) + assert calls["decode"] == 0, "disabled policy should not probe decode" + + +def test_config_to_scheduler_wiring(): + """codex: cover CLI/config -> Scheduler.__init__ wiring.""" + import inspect + + from freetoken.engine.config import EngineConfig + from freetoken.scheduler import scheduler as sched_mod + + assert "decode_interleave_every" in {f.name for f in __import__("dataclasses").fields(EngineConfig)} + src = inspect.getsource(sched_mod.Scheduler.__init__) + assert "_interleave" in src and "decode_interleave_every" in src, ( + "Scheduler.__init__ must build the policy from the config field" + ) + # and the CLI exposes it + from freetoken.server.args import ServerArgs + assert hasattr(ServerArgs, "decode_interleave_every") diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py new file mode 100644 index 000000000..e9a5c07da --- /dev/null +++ b/tests/scheduler/test_match_stop_str.py @@ -0,0 +1,86 @@ +"""Scheduler-side stop-string matching (Scheduler._match_stop_str) against the same +incrementally decoded text the frontend streams.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.core import Req, SamplingParams +from freetoken.scheduler.scheduler import Scheduler + +EOS, SPECIAL = 0, 1 + + +class _ByteTok: + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS, + 1 a special marker that skip_special_tokens drops.""" + + def decode(self, ids, skip_special_tokens=False): + def piece(t): + if t == EOS: + return b"" + if t == SPECIAL: + return b"" if skip_special_tokens else b"" + return bytes([t]) + + return b"".join(piece(t) for t in ids).decode("utf-8", errors="replace") + + +def _req(prompt: list[int], output_len: int, **sp) -> Req: + return Req( + input_ids=torch.tensor(prompt, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=output_len, + uid=1, + sampling_params=SamplingParams(**sp), + cache_handle=None, + ) + + +def _match(req: Req) -> str | None: + stub = SimpleNamespace(tokenizer=_ByteTok(), eos_token_ids={EOS}) + return Scheduler._match_stop_str(stub, req) + + +def _feed(req: Req, ids: list[int]) -> list[str | None]: + seen = [] + for t in ids: + req.append_host(torch.tensor([t], dtype=torch.int32)) + req.complete_one() + seen.append(_match(req)) + return seen + + +def test_a_stop_spanning_more_tokens_than_characters_is_matched(): + dna = list("\N{DNA DOUBLE HELIX}".encode()) # one character, four byte tokens + req = _req([65, 66], 32, stop_strs=["\N{DNA DOUBLE HELIX}"]) + text = list(b"Here is DNA ") + dna + list(b" and more") + seen = _feed(req, text) + first = next(i for i, m in enumerate(seen) if m is not None) + assert first == len(b"Here is DNA ") + len(dna) - 1 + + +def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): + # Overlap scheduling advances device_len ahead of the host: the budget is spent on the + # device (can_decode is False) while the host holds one token less than the limit. That + # EOS is not terminal, so the frontend decodes it and the scheduler must too. + req = _req([65], 3, stop_strs=["zz"], ignore_eos=True) + for t in (65, EOS): + req.append_host(torch.tensor([t], dtype=torch.int32)) + while req.can_decode: + req.complete_one() + assert req.input_ids.numel() < req.max_device_len + assert _match(req) is None + assert req.stop_decode_status.decoded_ids == [65, EOS] + + +def test_stops_match_the_text_decoded_with_the_requests_skip_special_tokens(): + # the frontend streams "ab" for a request that skips special tokens; the scheduler has + # to match the stop in that same text, not in "ab" + req = _req([65], 8, stop_strs=["ab"], skip_special_tokens=True) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, "ab"] + req = _req([65], 8, stop_strs=["ab"]) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, None] diff --git a/tests/server/test_dist_port_arg.py b/tests/server/test_dist_port_arg.py new file mode 100644 index 000000000..a49def7b6 --- /dev/null +++ b/tests/server/test_dist_port_arg.py @@ -0,0 +1,41 @@ +"""Regression test for #301: the TP rendezvous port must be independently configurable +from --port, so two instances on adjacent ports don't have one's API port collide with +the other's rendezvous port. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +ANON_PATH = "/models/anon" + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"} + + +def _parse(argv: list[str]): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + return parse_args(["--model", ANON_PATH, *argv]) + + +def test_dist_port_defaults_to_server_port_plus_one(): + args, _ = _parse(["--port", "8081"]) + assert args.distributed_port == 8082 + assert args.distributed_addr == "tcp://127.0.0.1:8082" + + +def test_dist_port_override_is_independent_of_server_port(): + args, _ = _parse(["--port", "8082", "--dist-port", "9000"]) + assert args.distributed_port == 9000 + assert args.distributed_addr == "tcp://127.0.0.1:9000" + + +def test_dist_port_rejects_out_of_range_value(): + with pytest.raises(SystemExit): + _parse(["--dist-port", "70000"]) diff --git a/tests/server/test_launch_worker_errors.py b/tests/server/test_launch_worker_errors.py new file mode 100644 index 000000000..8779016bd --- /dev/null +++ b/tests/server/test_launch_worker_errors.py @@ -0,0 +1,69 @@ +"""A tokenizer/detokenizer worker that dies during load must report WHY. + +Regression test for the asymmetry that made every tokenizer-side startup failure surface as +the generic "backend worker freetoken-detokenizer-0 exited during load": ``_run_scheduler`` +pushed an ("error", reason) ack before dying, ``_run_tokenize_worker`` did not. +""" + +from __future__ import annotations + +import queue + +import pytest + +from freetoken.server.launch import _run_tokenize_worker + + +def _drain(q: "queue.Queue") -> list: + out = [] + while True: + try: + out.append(q.get_nowait()) + except queue.Empty: + return out + + +def test_tokenize_worker_reports_the_real_reason_before_dying(monkeypatch): + """A raising tokenize_worker pushes ("error", reason) so the supervisor can surface it.""" + import freetoken.tokenizer as tok_mod + + def boom(**_kwargs): + raise FileNotFoundError("tokenizer.json is missing from the checkpoint") + + monkeypatch.setattr(tok_mod, "tokenize_worker", boom) + q: "queue.Queue" = queue.Queue() + + with pytest.raises(FileNotFoundError): + _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope") + + acks = _drain(q) + assert acks, "the dying worker pushed no ack at all" + kind, reason = acks[0] + assert kind == "error" + assert "FileNotFoundError" in reason + assert "tokenizer.json is missing" in reason + + +def test_tokenize_worker_without_ack_queue_still_propagates(monkeypatch): + """No ack_queue (direct call / test harness) must not turn the failure into an AttributeError.""" + import freetoken.tokenizer as tok_mod + + def boom(**_kwargs): + raise ValueError("bad tokenizer config") + + monkeypatch.setattr(tok_mod, "tokenize_worker", boom) + + with pytest.raises(ValueError, match="bad tokenizer config"): + _run_tokenize_worker(detach=False, tokenizer_path="/nope") + + +def test_tokenize_worker_success_pushes_no_error_ack(monkeypatch): + """The happy path must stay clean — no spurious ("error", …) on the ack queue.""" + import freetoken.tokenizer as tok_mod + + monkeypatch.setattr(tok_mod, "tokenize_worker", lambda **_k: None) + q: "queue.Queue" = queue.Queue() + + _run_tokenize_worker(detach=False, ack_queue=q, tokenizer_path="/nope") + + assert not [a for a in _drain(q) if isinstance(a, tuple) and a and a[0] == "error"] diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 7d5b3cd11..376cc2a7b 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -462,19 +462,67 @@ def test_models_route_returns_served_model_name(): assert card["max_model_len"] is None and card["context_length"] is None +def _models_card(state): + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + return TestClient(app).get("/v1/models").json()["data"][0] + + def test_models_route_publishes_the_model_context_length(): """`ft launch` reads this to size each agent's context window.""" state = FakeState([]) state.config.max_seq_len = 262144 - app = FastAPI() - register_openai_routes(app, lambda: state, lambda: {}) - card = TestClient(app).get("/v1/models").json()["data"][0] + card = _models_card(state) + # No pool geometry resolved yet (still loading, or pre-("meta", …) ack): the model + # ceiling is the only answer available. assert card["max_model_len"] == 262144 assert card["context_length"] == 262144 +def test_models_route_clamps_the_context_length_to_the_kv_pool(): + """The engine clamps its own max_seq_len to the allocated pool and the scheduler admits + against that, so publishing the unclamped ceiling here made `ft launch` size each agent's + compaction window past what the server can actually hold (#448).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 178176 + assert card["context_length"] == 178176 + + +def test_models_route_keeps_the_ceiling_when_the_pool_exceeds_it(): + """A pool larger than the model's positional ceiling does not extend the context.""" + state = FakeState([]) + state.config.max_seq_len = 32768 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + + card = _models_card(state) + + assert card["max_model_len"] == 32768 + assert card["context_length"] == 32768 + + +def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool(): + """A rebuild moves the pool; /v1/models must follow it rather than freeze the load-time + allocation (same most-recent-truth order /v1/cache/status reports).""" + state = FakeState([]) + state.config.max_seq_len = 262144 + state.config.page_size = 1 + state.cache_pools = {"num_pages": 178176, "page_size": 1} + state.last_rebuild = {"num_pages": 40000} + + card = _models_card(state) + + assert card["max_model_len"] == 40000 + + async def _collect(generator): return [chunk async for chunk in generator] diff --git a/tests/server/test_qwen_semantic_special_tokens.py b/tests/server/test_qwen_semantic_special_tokens.py new file mode 100644 index 000000000..11e35ccb3 --- /dev/null +++ b/tests/server/test_qwen_semantic_special_tokens.py @@ -0,0 +1,405 @@ +"""Qwen protocol markers are private to semantic chat-style API responses.""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from types import SimpleNamespace + +import pytest + +_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +_PY = os.path.join(_ROOT, "python") +if _PY not in sys.path: + sys.path.insert(0, _PY) + +from freetoken.llm.llm import LLM, RequestStatus # noqa: E402 +from freetoken.message import DetokenizeMsg, UserReply # noqa: E402 +from freetoken.server.api_server import FrontendManager # noqa: E402 +from freetoken.server.api_models import CompletionRequest # noqa: E402 +from freetoken.server.generation import ( # noqa: E402 + _QWEN_SEMANTIC_HIDDEN_TOKENS, + ContentDelta, + GenDone, + GenSpec, + ReasoningDelta, + ToolCallsDelta, + generate_events, + generate_full, +) +from freetoken.server.openai_api import handle_completion, stream_completion_chunks # noqa: E402 + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "echo", + "parameters": { + "type": "object", + "properties": {"value": {"type": "string"}}, + }, + }, + } +] + + +class FakeState: + def __init__( + self, + chunks: list[str], + *, + completion_tokens: list[int] | None = None, + reasoning_parser: str | None = "qwen3", + tool_call_parser: str = "qwen3_coder", + ) -> None: + self.config = SimpleNamespace( + reasoning_parser=reasoning_parser, + tool_call_parser=tool_call_parser, + served_model_name="qwen-unit", + ) + self._chunks = chunks + self._completion_tokens = completion_tokens or [1] * len(chunks) + self.sent = None + + def new_user(self) -> int: + return 42 + + async def send_one(self, msg) -> None: + self.sent = msg + + async def wait_for_ack(self, uid: int): + assert uid == 42 + for index, (chunk, token_delta) in enumerate(zip(self._chunks, self._completion_tokens)): + last = index == len(self._chunks) - 1 + yield UserReply( + uid=uid, + incremental_output=chunk, + finished=last, + finish_reason="stop" if last else None, + prompt_tokens_delta=5 if index == 0 else 0, + completion_tokens_delta=token_delta, + ) + + +def _spec(*, tools=None) -> GenSpec: + return GenSpec( + messages=[], + sampling_params=SimpleNamespace(), + chat_template_kwargs={"enable_thinking": False}, + template_tools=tools, + parser_tools=tools, + ) + + +async def _events(chunks: list[str], *, tools=None): + return [event async for event in generate_events(42, _spec(tools=tools), FakeState(chunks))] + + +def _visible(events) -> str: + return "".join(event.text for event in events if isinstance(event, ContentDelta)) + + +@pytest.mark.parametrize("marker", _QWEN_SEMANTIC_HIDDEN_TOKENS) +def test_every_hidden_marker_is_chunk_boundary_safe(marker: str): + text = f"before{marker}after" + events = asyncio.run(_events(list(text))) + + assert _visible(events) == "beforeafter" + assert isinstance(events[-1], GenDone) + assert events[-1].completion_tokens == len(text) + + +def test_stream_and_nonstream_match_while_preserving_semantic_markers_and_literals(): + grounding = ( + "<|object_ref_start|>cat<|object_ref_end|>" + "<|box_start|>(1,2),(3,4)<|box_end|>" + "<|quad_start|>(1,2),(3,4),(5,6),(7,8)<|quad_end|>" + ) + text = ( + "<|im_start|>plananswer" + "<|vision_start|><|image_pad|><|vision_end|>" + f"{grounding} literal \"<|im_start|>\" and `<|image_pad|>`<|im_end|>" + ) + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text], completion_tokens=[len(text)]))) + + reasoning = "".join(event.text for event in streamed if isinstance(event, ReasoningDelta)) + expected_content = f"answer{grounding} literal \"<|im_start|>\" and `<|image_pad|>`" + assert reasoning == full.reasoning == "plan" + assert _visible(streamed) == full.content == expected_content + assert full.completion_tokens == streamed[-1].completion_tokens == len(text) + + +def test_tool_and_think_markers_reach_their_parsers_before_filtering(): + block = ( + "" + 'literal <|im_start|>' + "" + ) + text = "choose tool<|im_start|>" + block + "<|im_end|>" + streamed = asyncio.run(_events(list(text), tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState([text]))) + + stream_calls = [ + call for event in streamed if isinstance(event, ToolCallsDelta) for call in event.calls + ] + assert "".join( + event.text for event in streamed if isinstance(event, ReasoningDelta) + ) == full.reasoning == "choose tool" + assert _visible(streamed) == full.content == "" + assert len(stream_calls) == len(full.tool_calls) == 1 + assert stream_calls[0].name == full.tool_calls[0].name == "echo" + assert json.loads(stream_calls[0].parameters) == json.loads(full.tool_calls[0].parameters) == { + "value": "literal <|im_start|>" + } + + +def test_filter_state_is_isolated_per_concurrent_request(): + async def run_pair(): + return await asyncio.gather( + _events(["left<|im_", "start|>right"]), + _events(['quoted "<|im_start|>"']), + ) + + first, second = asyncio.run(run_pair()) + assert _visible(first) == "leftright" + assert _visible(second) == 'quoted "<|im_start|>"' + + +@pytest.mark.parametrize("quote", ['"', "`", "'"]) +def test_unterminated_literal_cannot_disable_filtering(quote: str): + text = f"answer {quote}oops <|im_end|>" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == f"answer {quote}oops " + + +@pytest.mark.parametrize("slashes", ["\\", "\\\\"]) +def test_escaped_marker_opener_does_not_bypass_unterminated_quote_filter(slashes: str): + text = f'answer "oops {slashes}<|im_end|>' + expected = f'answer "oops {slashes}' + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == expected + + +@pytest.mark.parametrize(("opening", "closing"), [("“", "”"), ("‘", "’")]) +def test_typographic_quoted_marker_is_preserved(opening: str, closing: str): + text = f"literal {opening}<|im_start|>{closing}" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == text + + +def test_suppressed_marker_does_not_turn_a_possessive_into_an_open_quote(): + text = "foo<|im_end|>'s <|vision_pad|>done" + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + assert _visible(streamed) == full.content == "foo's done" + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("<<|im_end|>", "<"), + ("<|im_<|im_end|>", "<|im_"), + ("x<<|vision_pad|>y", "xok" + "" + ) + chunks = ["before<|im_", block, "start|>after"] + streamed = asyncio.run(_events(chunks, tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)]))) + + assert _visible(streamed) == full.content == "beforeafter" + assert len([event for event in streamed if isinstance(event, ToolCallsDelta)]) == 1 + assert len(full.tool_calls) == 1 + + +def test_disproved_marker_prefix_stays_before_intervening_tool_event(): + block = ( + "ok" + "" + ) + streamed = asyncio.run(_events(["pre<|im_", block, "Xpost"], tools=TOOLS)) + + kinds = [type(event).__name__ for event in streamed] + call_start_index = kinds.index("ToolCallStart") + call_end_index = kinds.index("ToolCallsDelta") + before = "".join( + event.text + for event in streamed[:call_start_index] + if isinstance(event, ContentDelta) + ) + after = "".join( + event.text + for event in streamed[call_end_index + 1 :] + if isinstance(event, ContentDelta) + ) + + assert before == "pre<|im_" + assert after == "Xpost" + + +def test_tool_event_flushes_prior_quoted_prose_before_call(): + block = ( + "ok" + "" + ) + chunks = [ + 'before "<|im_start|> quoted', + block, + ' tail <|vision_pad|>" after<|im_end|>', + ] + streamed = asyncio.run(_events(chunks, tools=TOOLS)) + full = asyncio.run(generate_full(42, _spec(tools=TOOLS), FakeState(["".join(chunks)]))) + + kinds = [type(event).__name__ for event in streamed] + call_start_index = kinds.index("ToolCallStart") + call_end_index = kinds.index("ToolCallsDelta") + prefix = "".join( + event.text + for event in streamed[:call_start_index] + if isinstance(event, ContentDelta) + ) + suffix = "".join( + event.text + for event in streamed[call_end_index + 1 :] + if isinstance(event, ContentDelta) + ) + assert prefix == 'before "<|im_start|> quoted' + assert suffix == ' tail <|vision_pad|>" after' + assert call_start_index < call_end_index + assert _visible(streamed) == full.content == ( + 'before "<|im_start|> quoted tail <|vision_pad|>" after' + ) + + +def test_reasoning_filter_hides_transport_marker_and_preserves_fim_tokens(): + text = ( + "plan<|vision_pad|>" + "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|><|im_end|>" + ) + streamed = asyncio.run(_events(list(text))) + full = asyncio.run(generate_full(42, _spec(), FakeState([text]))) + + reasoning = "".join( + event.text for event in streamed if isinstance(event, ReasoningDelta) + ) + expected_content = "<|fim_prefix|>left<|fim_suffix|>right<|fim_middle|>" + assert reasoning == full.reasoning == "plan" + assert _visible(streamed) == full.content == expected_content + + +def test_non_qwen_semantic_dialect_is_unchanged(): + text = '<|im_start|>literal "<|vision_pad|>"<|im_end|>' + state = FakeState( + [text], + reasoning_parser=None, + tool_call_parser="llama3", + ) + + result = asyncio.run(generate_full(42, _spec(), state)) + + assert result.content == text + + +def test_legacy_qwen_parser_alias_enables_semantic_filter_without_reasoning(): + text = "<|im_start|>visible<|vision_pad|><|im_end|>" + state = FakeState([text], reasoning_parser=None, tool_call_parser="qwen") + result = asyncio.run(generate_full(42, _spec(), state)) + + assert result.content == "visible" + + +def test_legacy_completions_paths_remain_raw_for_qwen(): + raw = '<|im_start|>raw "<|vision_pad|>"<|im_end|>' + + full = asyncio.run( + handle_completion( + CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8), + request=None, + state=FakeState([raw]), + model_sampling={}, + ) + ) + assert full["choices"][0]["text"] == raw + + async def stream_text() -> str: + chunks = [ + chunk async for chunk in stream_completion_chunks( + 42, + CompletionRequest(model="qwen-unit", prompt="continue", max_tokens=8, stream=True), + FakeState([raw]), + ) + ] + frames = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + frames.append(json.loads(line[6:])) + return "".join(frame["choices"][0]["text"] for frame in frames) + + assert asyncio.run(stream_text()) == raw + + +def test_raw_generate_stream_remains_unfiltered_for_qwen(): + raw = "<|im_start|>raw<|vision_pad|><|im_end|>" + + async def stream_text() -> str: + chunks = [ + chunk async for chunk in FrontendManager.stream_generate(FakeState([raw]), 42) + ] + frames = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + frames.append(json.loads(line[6:])) + return "".join(frame["text"] for frame in frames) + + assert asyncio.run(stream_text()) == raw + + +def test_offline_generation_keeps_exact_output_token_ids(): + offline = SimpleNamespace( + status_map={7: RequestStatus(uid=7, input_ids=[], output_ids=[])}, + eos_token_ids=set(), + ) + marker_token_ids = [248045, 248044, 248046] + + LLM.offline_send_result( + offline, + [ + DetokenizeMsg(uid=7, next_token=token_id, finished=False) + for token_id in marker_token_ids + ], + ) + + assert offline.status_map[7].output_ids == marker_token_ids diff --git a/tests/server/test_stream_keepalive.py b/tests/server/test_stream_keepalive.py new file mode 100644 index 000000000..e86d4365f --- /dev/null +++ b/tests/server/test_stream_keepalive.py @@ -0,0 +1,115 @@ +"""A stream that stays silent -- the prefill of a long prompt -- keeps its client, and notices when +the client leaves instead of prefilling for nobody.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.server import api_server +from freetoken.server.api_server import FrontendManager + +# Converts a stream that never ends into a failure; the streams below finish in milliseconds. +HANG_GUARD_SECONDS = 10.0 + + +class _Request: + def __init__(self) -> None: + self.gone = False + + async def is_disconnected(self) -> bool: + return self.gone + + +def _manager() -> SimpleNamespace: + async def abort_user(uid: int) -> None: + pass + + return SimpleNamespace(abort_user=abort_user) + + +def _is_sse_comment(chunk: bytes) -> bool: + return chunk.startswith(b":") and chunk.endswith(b"\n\n") + + +def test_a_silent_stream_sends_comments_until_its_next_chunk(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + released = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await released.wait() + yield b"data: last\n\n" + + async def consume() -> list[bytes]: + out = [] + request = _Request() + async for chunk in FrontendManager.stream_with_cancellation(_manager(), chunks(), request, 7): + out.append(chunk) + if _is_sse_comment(chunk): + released.set() + return out + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out[0] == b"data: first\n\n" and out[-1] == b"data: last\n\n" + assert len(out) > 2 and all(_is_sse_comment(chunk) for chunk in out[1:-1]) + + +def test_a_client_that_leaves_during_the_silence_aborts_its_request(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + prefill_done = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await prefill_done.wait() + yield b"data: never read\n\n" + + async def scenario() -> tuple[list[bytes], list[int]]: + aborted = asyncio.Event() + seen: list[int] = [] + + async def abort_user(uid: int) -> None: + seen.append(uid) + aborted.set() + + request = _Request() + out: list[bytes] = [] + + async def consume() -> None: + manager = SimpleNamespace(abort_user=abort_user) + async for chunk in FrontendManager.stream_with_cancellation(manager, chunks(), request, 9): + out.append(chunk) + request.gone = True + + consumer = asyncio.ensure_future(consume()) + # Only the wrapper noticing the departure can end this wait; the guard below fails it. + await aborted.wait() + with pytest.raises(asyncio.CancelledError): + await consumer + return out, seen + + out, seen = asyncio.run(asyncio.wait_for(scenario(), HANG_GUARD_SECONDS)) + assert out == [b"data: first\n\n"] + assert seen == [9] + assert not prefill_done.is_set() + + +def test_a_stream_that_never_goes_quiet_carries_no_comments(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS) + + async def chunks(): + for index in range(5): + yield f"data: {index}\n\n".encode() + + async def consume() -> list[bytes]: + return [ + chunk + async for chunk in FrontendManager.stream_with_cancellation( + _manager(), chunks(), _Request(), 3 + ) + ] + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out == [f"data: {index}\n\n".encode() for index in range(5)] diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py index 5617a1053..4717617b9 100644 --- a/tests/tokenizer/test_detokenize_extras.py +++ b/tests/tokenizer/test_detokenize_extras.py @@ -79,3 +79,19 @@ def test_reasoning_tokens_count_up_to_the_end_tag(): assert reasoning == 3 # b, c and the tag itself _, none = _run(manager, 2, [1, 2, EOS]) assert none == 0 + + +def test_a_batch_with_one_uid_twice_advances_its_state_in_order(): + # each message must see the decode state the previous one left, and the reasoning + # counts must stay aligned with the texts + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + msgs = [ + DetokenizeMsg(uid=1, next_token=1, finished=False), + DetokenizeMsg(uid=1, next_token=THINK_END, finished=False), + DetokenizeMsg(uid=2, next_token=3, finished=False), + DetokenizeMsg(uid=1, next_token=2, finished=True), + ] + texts, counts = manager.detokenize_with_meta(msgs) + assert texts == ["b", "", "d", "c"] + assert counts == [0, 0, 0, 2] + assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"]