| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """OpenAI-compatible proxy for E+PD and E+P+D disaggregation. |
| 5 | |
| 6 | Use static server URLs, or --dynamic-registration for launcher-managed HTTP |
| 7 | registration of E, P/PD and D instances. Both modes route OpenAI-compatible |
| 8 | ``/v1/chat/completions`` requests through encoder and inference clusters and |
| 9 | share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming |
| 10 | forwarding. |
| 11 | |
| 12 | For MM input we: |
| 13 | 1. Extract *every* image/audio/video item. |
| 14 | 2. Send concurrent encoder requests with all text removed, |
| 15 | grouping images assigned to the same encoder. |
| 16 | 3. Wait for all of them to succeed. |
| 17 | 4. Forward the *original* request to a decode server. |
| 18 | """ |
| 19 | |
| 20 | from __future__ import annotations |
| 21 | |
| 22 | import argparse |
| 23 | import asyncio |
| 24 | import enum |
| 25 | import hashlib |
| 26 | import itertools |
| 27 | import json |
| 28 | import logging |
| 29 | import os |
| 30 | import random |
| 31 | import secrets |
| 32 | import time |
| 33 | import uuid |
| 34 | from collections.abc import AsyncIterator |
| 35 | from contextlib import asynccontextmanager, suppress |
| 36 | from dataclasses import dataclass, field |
| 37 | from typing import Any |
| 38 | |
| 39 | import aiohttp |
| 40 | import msgspec |
| 41 | import uvicorn |
| 42 | from fastapi import Depends, FastAPI, Header, HTTPException, Request |
| 43 | from fastapi.responses import JSONResponse, Response, StreamingResponse |
| 44 | from pydantic import AnyHttpUrl, BaseModel, Field, model_validator |
| 45 | |
| 46 | ############################################################################### |
| 47 | # FastAPI app & global state |
| 48 | ############################################################################### |
| 49 | |
| 50 | logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") |
| 51 | logger = logging.getLogger("proxy") |
| 52 | |
| 53 | DEFAULT_PROBE_INTERVAL = 5.0 |
| 54 | DEFAULT_PROBE_TIMEOUT = 2.0 |
| 55 | DEFAULT_FAIL_THRESHOLD = 3 |
| 56 | # Stop probing an instance that has been down this long. 0 probes forever, |
| 57 | # which is what a cluster that restarts instances in place wants. |
| 58 | DEFAULT_EVICTED_TTL = 900.0 |
| 59 | |
| 60 | |
| 61 | class InstanceRole(str, enum.Enum): |
| 62 | ENCODE = "encode" |
| 63 | PREFILL = "prefill" |
| 64 | DECODE = "decode" |
| 65 | PREFILL_DECODE = "prefill_decode" |
| 66 | |
| 67 | |
| 68 | @dataclass |
| 69 | class InstanceRecord: |
| 70 | """One registered instance. |
| 71 | |
| 72 | Attributes: |
| 73 | role: Which stage this instance serves. |
| 74 | url: Base OpenAI-compatible URL, e.g. ``http://host:8000``. |
| 75 | ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica, |
| 76 | supplied by the launcher from the consumer's fixed port config. |
| 77 | dp_size: Data-parallel replicas behind `url`, so the proxy can pick |
| 78 | a replica and name the same one to both halves of a request. |
| 79 | |
| 80 | """ |
| 81 | |
| 82 | role: InstanceRole |
| 83 | url: str |
| 84 | ec_zmq_addrs: list[str] = field(default_factory=list) |
| 85 | dp_size: int = 1 |
| 86 | |
| 87 | |
| 88 | class InstanceRegistry: |
| 89 | def __init__( |
| 90 | self, |
| 91 | probe_interval: float = DEFAULT_PROBE_INTERVAL, |
| 92 | probe_timeout: float = DEFAULT_PROBE_TIMEOUT, |
| 93 | fail_threshold: int = DEFAULT_FAIL_THRESHOLD, |
| 94 | evicted_ttl: float = DEFAULT_EVICTED_TTL, |
| 95 | ): |
| 96 | self._probe_interval = probe_interval |
| 97 | self._probe_timeout = probe_timeout |
| 98 | self._fail_threshold = fail_threshold |
| 99 | self._evicted_ttl = evicted_ttl |
| 100 | |
| 101 | self._live: dict[str, InstanceRecord] = {} |
| 102 | self._evicted: dict[str, InstanceRecord] = {} |
| 103 | self._evicted_since: dict[str, float] = {} |
| 104 | self._fail_counts: dict[str, int] = {} |
| 105 | # One cursor per role, only ever incremented. Rebuilding it whenever |
| 106 | # the roster changes -- what an `itertools.cycle` over a mutable list |
| 107 | # forces -- restarts every fan-out at the first instance and hot-spots |
| 108 | # it after each registration. |
| 109 | self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole} |
| 110 | self._replica_cursors: dict[str, int] = {} |
| 111 | self._probe_task: asyncio.Task | None = None |
| 112 | |
| 113 | def register(self, record: InstanceRecord) -> bool: |
| 114 | """Add or refresh an instance without overriding its health status.""" |
| 115 | roles = { |
| 116 | other.role |
| 117 | for other in itertools.chain(self._live.values(), self._evicted.values()) |
| 118 | } |
| 119 | split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE} |
| 120 | if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or ( |
| 121 | record.role in split_roles and InstanceRole.PREFILL_DECODE in roles |
| 122 | ): |
| 123 | raise ValueError("Cannot mix prefill_decode with standalone prefill/decode") |
| 124 | key = record.url |
| 125 | previous = self._live.get(key) or self._evicted.get(key) |
| 126 | if previous is not None: |
| 127 | if previous == record: |
| 128 | return False |
| 129 | if previous.role is not record.role: |
| 130 | raise ValueError("Unregister the instance before changing its role") |
| 131 | target = self._live if key in self._live else self._evicted |
| 132 | target[key] = record |
| 133 | return False |
| 134 | self._live[key] = record |
| 135 | logger.info("Registered instance %s: %s", record.role.value, record.url) |
| 136 | return True |
| 137 | |
| 138 | def unregister(self, url: str) -> bool: |
| 139 | """Drop an instance for good, so a probe cannot bring it back.""" |
| 140 | key = url |
| 141 | if key not in self._live and key not in self._evicted: |
| 142 | return False |
| 143 | self._live.pop(key, None) |
| 144 | self._evicted.pop(key, None) |
| 145 | self._evicted_since.pop(key, None) |
| 146 | self._fail_counts.pop(key, None) |
| 147 | self._replica_cursors.pop(key, None) |
| 148 | logger.info("Unregistered instance: %s", url) |
| 149 | return True |
| 150 | |
| 151 | def instances(self, role: InstanceRole) -> list[InstanceRecord]: |
| 152 | return [record for record in self._live.values() if record.role is role] |
| 153 | |
| 154 | def urls(self, role: InstanceRole) -> list[str]: |
| 155 | return [record.url for record in self.instances(role)] |
| 156 | |
| 157 | def pick(self, role: InstanceRole) -> InstanceRecord | None: |
| 158 | """Take the next instance of `role` in round-robin order.""" |
| 159 | picked = self.pick_many(role, 1) |
| 160 | return picked[0] if picked else None |
| 161 | |
| 162 | def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]: |
| 163 | """Take `count` instances, continuing the rotation across calls. |
| 164 | |
| 165 | A multimodal request fans out one encoder request per item, so the |
| 166 | assignment has to be contiguous with the previous request's rather |
| 167 | than restart at the first instance every time. |
| 168 | """ |
| 169 | alive = self.instances(role) |
| 170 | if not alive or count <= 0: |
| 171 | return [] |
| 172 | start = self._cursors[role] |
| 173 | self._cursors[role] = start + count |
| 174 | return [alive[(start + offset) % len(alive)] for offset in range(count)] |
| 175 | |
| 176 | def next_replica(self, record: InstanceRecord) -> int: |
| 177 | """Take the next data-parallel replica of `record`, round-robin. |
| 178 | |
| 179 | The encoder pushes to one replica's receive channel, so the request |
| 180 | has to run on that same replica; the caller names the rank to both |
| 181 | halves. |
| 182 | """ |
| 183 | if record.dp_size <= 1: |
| 184 | return 0 |
| 185 | cursor = self._replica_cursors.get(record.url, 0) |
| 186 | self._replica_cursors[record.url] = cursor + 1 |
| 187 | return cursor % record.dp_size |
| 188 | |
| 189 | def status(self) -> dict[str, Any]: |
| 190 | return { |
| 191 | role.value: { |
| 192 | "live": [record.url for record in self.instances(role)], |
| 193 | "evicted": [ |
| 194 | record.url |
| 195 | for record in self._evicted.values() |
| 196 | if record.role is role |
| 197 | ], |
| 198 | } |
| 199 | for role in InstanceRole |
| 200 | } |
| 201 | |
| 202 | def start_probing(self) -> None: |
| 203 | if self._probe_task is None and self._probe_interval > 0: |
| 204 | self._probe_task = asyncio.create_task(self._probe_loop()) |
| 205 | |
| 206 | async def stop_probing(self) -> None: |
| 207 | if self._probe_task is None: |
| 208 | return |
| 209 | self._probe_task.cancel() |
| 210 | with suppress(asyncio.CancelledError): |
| 211 | await self._probe_task |
| 212 | self._probe_task = None |
| 213 | |
| 214 | async def _probe_loop(self) -> None: |
| 215 | timeout = aiohttp.ClientTimeout(total=self._probe_timeout) |
| 216 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 217 | while True: |
| 218 | await asyncio.sleep(self._probe_interval) |
| 219 | try: |
| 220 | await self._probe_once(session) |
| 221 | except Exception: |
| 222 | logger.exception("EPD registry probe round failed") |
| 223 | |
| 224 | async def _probe_once(self, session: aiohttp.ClientSession) -> None: |
| 225 | targets = list(self._live.values()) + list(self._evicted.values()) |
| 226 | if not targets: |
| 227 | return |
| 228 | results = await asyncio.gather( |
| 229 | *(self._probe(session, record.url) for record in targets), |
| 230 | return_exceptions=True, |
| 231 | ) |
| 232 | now = time.monotonic() |
| 233 | for record, healthy in zip(targets, results): |
| 234 | current = self._live.get(record.url) or self._evicted.get(record.url) |
| 235 | if current is not record: |
| 236 | continue |
| 237 | if healthy is True: |
| 238 | self._on_probe_success(record) |
| 239 | else: |
| 240 | self._on_probe_failure(record, now) |
| 241 | self._drop_expired(now) |
| 242 | |
| 243 | async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool: |
| 244 | async with session.get(f"{url}/health") as resp: |
| 245 | return resp.status == 200 |
| 246 | |
| 247 | def _on_probe_success(self, record: InstanceRecord) -> None: |
| 248 | key = record.url |
| 249 | self._fail_counts.pop(key, None) |
| 250 | if key in self._evicted: |
| 251 | self._evicted.pop(key, None) |
| 252 | self._evicted_since.pop(key, None) |
| 253 | self._live[key] = record |
| 254 | logger.info( |
| 255 | "Instance %s (%s) is healthy again; routing resumed", |
| 256 | record.url, |
| 257 | record.role.value, |
| 258 | ) |
| 259 | |
| 260 | def _on_probe_failure(self, record: InstanceRecord, now: float) -> None: |
| 261 | key = record.url |
| 262 | if key not in self._live: |
| 263 | # Either already evicted, or unregistered while this probe was in |
| 264 | # flight. A round snapshots its targets and then awaits, so a |
| 265 | # removal inside that window would otherwise be undone by a result |
| 266 | # describing a registry that no longer exists. |
| 267 | return |
| 268 | failures = self._fail_counts.get(key, 0) + 1 |
| 269 | self._fail_counts[key] = failures |
| 270 | if failures < self._fail_threshold: |
| 271 | return |
| 272 | self._live.pop(key, None) |
| 273 | self._evicted[key] = record |
| 274 | self._evicted_since[key] = now |
| 275 | logger.warning( |
| 276 | "Instance %s (%s) failed %d consecutive probes; stopped routing " |
| 277 | "to it. It rejoins on its own once it responds again.", |
| 278 | record.url, |
| 279 | record.role.value, |
| 280 | failures, |
| 281 | ) |
| 282 | |
| 283 | def _drop_expired(self, now: float) -> None: |
| 284 | if self._evicted_ttl <= 0: |
| 285 | return |
| 286 | for key, since in list(self._evicted_since.items()): |
| 287 | if now - since < self._evicted_ttl: |
| 288 | continue |
| 289 | record = self._evicted.pop(key, None) |
| 290 | self._evicted_since.pop(key, None) |
| 291 | self._fail_counts.pop(key, None) |
| 292 | self._replica_cursors.pop(key, None) |
| 293 | if record is not None: |
| 294 | logger.warning("Instance %s stayed down; forgetting it", record.url) |
| 295 | |
| 296 | |
| 297 | class InstanceRegistration(BaseModel): |
| 298 | role: InstanceRole |
| 299 | url: AnyHttpUrl |
| 300 | ec_zmq_addrs: list[str] = Field(default_factory=list) |
| 301 | dp_size: int = Field(default=1, ge=1) |
| 302 | |
| 303 | @model_validator(mode="after") |
| 304 | def validate_consumer_addresses(self): |
| 305 | if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size: |
| 306 | raise ValueError("Provide one Mooncake control address per DP replica") |
| 307 | return self |
| 308 | |
| 309 | |
| 310 | def require_admin_key(x_api_key: str = Header(default="")) -> None: |
| 311 | expected = os.getenv("ADMIN_API_KEY", "") |
| 312 | if not expected or not secrets.compare_digest(x_api_key, expected): |
| 313 | raise HTTPException(403, "Invalid admin API key") |
| 314 | |
| 315 | |
| 316 | @dataclass |
| 317 | class EPDProxyConfig: |
| 318 | probe_interval: float = 5.0 |
| 319 | probe_timeout: float = 2.0 |
| 320 | fail_threshold: int = 3 |
| 321 | evicted_ttl: float = 900.0 |
| 322 | |
| 323 | |
| 324 | @dataclass |
| 325 | class _Route: |
| 326 | """The instances one request was assigned to. |
| 327 | |
| 328 | Attributes: |
| 329 | consumer: The stage that receives the embedding, if any stage does. |
| 330 | consumer_zmq: The receive address named to the encoders. |
| 331 | dp_rank: Which replica of `consumer` was named, so the request can be |
| 332 | pinned to it. |
| 333 | |
| 334 | """ |
| 335 | |
| 336 | encoder_urls: list[str] |
| 337 | prefill: InstanceRecord | None |
| 338 | decode: InstanceRecord |
| 339 | consumer: InstanceRecord | None = None |
| 340 | consumer_zmq: str | None = None |
| 341 | dp_rank: int | None = None |
| 342 | |
| 343 | |
| 344 | class EPDProxy: |
| 345 | def __init__(self, registry: InstanceRegistry): |
| 346 | self.registry = registry |
| 347 | |
| 348 | # ---------------------------------------------------------------- # |
| 349 | # Routing # |
| 350 | # ---------------------------------------------------------------- # |
| 351 | def route(self, num_items: int) -> _Route: |
| 352 | decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick( |
| 353 | InstanceRole.DECODE |
| 354 | ) |
| 355 | if decode is None: |
| 356 | raise HTTPException( |
| 357 | status_code=503, detail="No decode instance is registered" |
| 358 | ) |
| 359 | prefill = ( |
| 360 | self.registry.pick(InstanceRole.PREFILL) |
| 361 | if decode.role is InstanceRole.DECODE |
| 362 | else None |
| 363 | ) |
| 364 | if decode.role is InstanceRole.DECODE and prefill is None: |
| 365 | raise HTTPException( |
| 366 | status_code=503, detail="No prefill instance is registered" |
| 367 | ) |
| 368 | encoders = self.registry.urls(InstanceRole.ENCODE) |
| 369 | if num_items and not encoders: |
| 370 | raise HTTPException( |
| 371 | status_code=503, detail="No encode instance is registered" |
| 372 | ) |
| 373 | route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode) |
| 374 | self._name_consumer(route) |
| 375 | return route |
| 376 | |
| 377 | def _name_consumer(self, route: _Route) -> None: |
| 378 | """Pin the EC consumer replica and, for push connectors, its endpoint.""" |
| 379 | candidate = route.prefill or route.decode |
| 380 | route.consumer = candidate |
| 381 | rank = self.registry.next_replica(candidate) |
| 382 | route.dp_rank = rank if candidate.dp_size > 1 else None |
| 383 | if candidate.ec_zmq_addrs: |
| 384 | route.consumer_zmq = candidate.ec_zmq_addrs[rank] |
| 385 | |
| 386 | |
| 387 | app = FastAPI() |
| 388 | encode_session: aiohttp.ClientSession | None = None |
| 389 | prefill_session: aiohttp.ClientSession | None = None |
| 390 | decode_session: aiohttp.ClientSession | None = None |
| 391 | |
| 392 | # Cursor for round-robin encoder assignment, shared across requests so the |
| 393 | # fan-out doesn't restart from e_urls[0] every time. |
| 394 | encoder_rr_idx = 0 |
| 395 | encoder_rr_lock = asyncio.Lock() |
| 396 | |
| 397 | ############################################################################### |
| 398 | # Utils |
| 399 | ############################################################################### |
| 400 | |
| 401 | |
| 402 | MM_TYPES = {"image_url", "audio_url", "input_audio", "video_url"} |
| 403 | |
| 404 | # The embeds content type each MM item is rewritten to once the encoder has |
| 405 | # published its embedding out of band. |
| 406 | EMBEDS_TYPES = { |
| 407 | "image_url": "image_embeds", |
| 408 | "audio_url": "audio_embeds", |
| 409 | "input_audio": "audio_embeds", |
| 410 | "video_url": "video_embeds", |
| 411 | } |
| 412 | |
| 413 | |
| 414 | def encoder_rr_assignment( |
| 415 | e_urls: list[str], start: int, count: int |
| 416 | ) -> tuple[list[str], int]: |
| 417 | """Assign `count` items to encoder URLs starting from cursor `start`. |
| 418 | |
| 419 | Returns the per-item URL list and the cursor value the next call should |
| 420 | start from, so the assignment is contiguous across calls instead of |
| 421 | restarting at e_urls[0] every time. |
| 422 | """ |
| 423 | urls = [e_urls[(start + i) % len(e_urls)] for i in range(count)] |
| 424 | next_start = (start + count) % len(e_urls) |
| 425 | return urls, next_start |
| 426 | |
| 427 | |
| 428 | def validate_ec_consumer_routing( |
| 429 | prefill_urls: list[str], consumer_addrs: list[str] |
| 430 | ) -> None: |
| 431 | """Reject the topology whose EC destination cannot be routed safely.""" |
| 432 | if prefill_urls and consumer_addrs: |
| 433 | raise ValueError( |
| 434 | "Mooncake EC consumer routing supports E+PD only; disable independent " |
| 435 | "prefill or omit --ec-consumer-zmq-addrs." |
| 436 | ) |
| 437 | |
| 438 | |
| 439 | # Diagnostic switch: forward the original request to the decoder so the |
| 440 | # only difference from the rewrite path is the rewrite itself. |
| 441 | NO_REWRITE = False |
| 442 | # Maximum images per encoder subrequest; 0 leaves batches unlimited. |
| 443 | ENCODER_MAX_BATCH_SIZE = int(os.getenv("ENCODER_MAX_BATCH_SIZE", "0")) |
| 444 | if ENCODER_MAX_BATCH_SIZE < 0: |
| 445 | raise ValueError("ENCODER_MAX_BATCH_SIZE must be non-negative") |
| 446 | |
| 447 | # Decode-side retries for a retryable internal error (`finish_reason="error"`, |
| 448 | # e.g. an encoder embedding the connector could not deliver). Re-issuing runs |
| 449 | # the encode again, which produces a fresh transfer. |
| 450 | DECODE_RETRIES = 1 |
| 451 | |
| 452 | |
| 453 | # Grid metadata reported by the encoder instance, keyed by item index. |
| 454 | # Empty when the encoder did not report any (then nothing is rewritten). |
| 455 | def content_uuid(item: dict) -> str: |
| 456 | """Cache key for a multimodal item, derived from its content. |
| 457 | |
| 458 | Must be content-derived, not request-derived: the EC cache is keyed by this |
| 459 | value, so a per-request key (a request id, say) would make every request a |
| 460 | miss and throw away cross-request reuse of already-encoded media -- while |
| 461 | the unmodified path, which hashes the content, would keep it. That asymmetry |
| 462 | silently biases any comparison between the two. |
| 463 | """ |
| 464 | url = ( |
| 465 | item.get("image_url") or item.get("audio_url") or item.get("video_url") or {} |
| 466 | ).get("url") or "" |
| 467 | payload = url or json.dumps(item, sort_keys=True) |
| 468 | return hashlib.sha256(payload.encode()).hexdigest() |
| 469 | |
| 470 | |
| 471 | def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict: |
| 472 | """Replace each media item with a metadata-only reference for the decoder. |
| 473 | |
| 474 | The decoder does not need the pixels: the encoder instance already produced |
| 475 | the embedding and published it through the EC connector under the same uuid. |
| 476 | Sending only the grid lets the decoder size the placeholder range without |
| 477 | re-running the media transform. |
| 478 | |
| 479 | `item_meta` holds what the encoder reported for each item (its cache key and |
| 480 | the grid its processor actually produced), so the grid is never re-derived |
| 481 | here -- a second derivation could disagree with the encoder's. |
| 482 | """ |
| 483 | mm_items = extract_mm_items(req_data) |
| 484 | # Audio processors may consume placeholders from other modalities too. |
| 485 | # Keep the whole request raw if any item cannot be rewritten safely. |
| 486 | control_fields = {"mm_hash", "ec_mm_hash", "transfer_id"} |
| 487 | raw_audio_request = any( |
| 488 | item["type"] in {"audio_url", "input_audio"} for item in mm_items |
| 489 | ) and any( |
| 490 | not item_meta.get(i, {}).get("mm_hash") |
| 491 | or not (item_meta.get(i, {}).keys() - control_fields) |
| 492 | for i in range(len(mm_items)) |
| 493 | ) |
| 494 | rewritten = 0 |
| 495 | transfer_items = [] |
| 496 | idx = 0 |
| 497 | new_messages = [] |
| 498 | for msg in req_data.get("messages", []): |
| 499 | content = msg.get("content") |
| 500 | if not isinstance(content, list): |
| 501 | new_messages.append(msg) |
| 502 | continue |
| 503 | new_content = [] |
| 504 | for item in content: |
| 505 | if item.get("type") not in MM_TYPES: |
| 506 | new_content.append(item) |
| 507 | continue |
| 508 | meta = dict(item_meta.get(idx) or {}) |
| 509 | idx += 1 |
| 510 | item_uuid = meta.pop("mm_hash", None) |
| 511 | ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid |
| 512 | transfer_id = meta.pop("transfer_id", None) |
| 513 | if raw_audio_request: |
| 514 | new_content.append( |
| 515 | { |
| 516 | **item, |
| 517 | "uuid": item_uuid or item.get("uuid") or content_uuid(item), |
| 518 | } |
| 519 | ) |
| 520 | if transfer_id is not None: |
| 521 | transfer_items.append( |
| 522 | {"mm_hash": ec_mm_hash, "transfer_id": transfer_id} |
| 523 | ) |
| 524 | continue |
| 525 | # Whatever keys the encoder reported are the metadata its model |
| 526 | # declared as needed to size the placeholder range; the proxy does |
| 527 | # not need to know their names. |
| 528 | # Downstream stacks per item; keep the existing per-item vector shape. |
| 529 | metadata = { |
| 530 | k: [ |
| 531 | x |
| 532 | for item in v |
| 533 | for x in (item if isinstance(item, list) else [item]) |
| 534 | ] |
| 535 | for k, v in meta.items() |
| 536 | } |
| 537 | if not metadata or not item_uuid: |
| 538 | # Nothing to size the placeholder range with. A processor cache |
| 539 | # hit is not a cause on its own: with the default `lru` type the |
| 540 | # engine restores the item before the scheduler reports it. It |
| 541 | # goes missing when the encode request failed, or under |
| 542 | # `--mm-processor-cache-type shm`, where a hit replaces the item |
| 543 | # with its shared-memory address and only the worker restores |
| 544 | # it. Send the media so the decoder can derive the grid itself. |
| 545 | new_content.append(item) |
| 546 | continue |
| 547 | embeds_type = EMBEDS_TYPES[item["type"]] |
| 548 | new_content.append( |
| 549 | {"type": embeds_type, embeds_type: metadata, "uuid": item_uuid} |
| 550 | ) |
| 551 | if transfer_id is not None: |
| 552 | transfer_items.append( |
| 553 | {"mm_hash": ec_mm_hash, "transfer_id": transfer_id} |
| 554 | ) |
| 555 | rewritten += 1 |
| 556 | new_messages.append({**msg, "content": new_content}) |
| 557 | |
| 558 | if not rewritten and not raw_audio_request: |
| 559 | return req_data |
| 560 | if rewritten: |
| 561 | logger.info("Rewrote %d media item(s) as metadata references", rewritten) |
| 562 | rewritten_request = {**req_data, "messages": new_messages} |
| 563 | if transfer_items: |
| 564 | ec_transfer_params = dict(req_data.get("ec_transfer_params") or {}) |
| 565 | ec_transfer_params["ec_items"] = transfer_items |
| 566 | rewritten_request["ec_transfer_params"] = ec_transfer_params |
| 567 | return rewritten_request |
| 568 | |
| 569 | |
| 570 | def extract_mm_items(request_data: dict) -> list[dict]: |
| 571 | """Return *all* image/audio/video items that appear anywhere in `messages`. |
| 572 | |
| 573 | Each returned dict looks like: |
| 574 | { "type": "image_url", "image_url": {...} } |
| 575 | """ |
| 576 | items: list[dict] = [] |
| 577 | for msg in request_data.get("messages", []): |
| 578 | content = msg.get("content") |
| 579 | if not isinstance(content, list): |
| 580 | continue |
| 581 | |
| 582 | for item in content: |
| 583 | if item.get("type") in MM_TYPES: |
| 584 | items.append(item) |
| 585 | return items |
| 586 | |
| 587 | |
| 588 | async def fanout_encoder_primer( |
| 589 | orig_request: dict, |
| 590 | e_urls: list[str], |
| 591 | req_id: str, |
| 592 | consumer_zmq: str | None = None, |
| 593 | ) -> tuple[dict[int, dict], dict[str, Any]]: |
| 594 | """1. Group images by encoder, retaining per-item round-robin assignment. |
| 595 | 2. Send them concurrently to the encode cluster. |
| 596 | 3. Raise if any of them fails. |
| 597 | |
| 598 | Returns, per item index, the metadata the encoder reported in |
| 599 | `ec_transfer_params`: its EC cache key and the grid its processor produced. |
| 600 | The proxy still supplies the uuid so both sides key the cache the same way; |
| 601 | the grid can only come from the encoder, which is the side that computed it. |
| 602 | |
| 603 | Also returns the connector handles to put on the decode body, as a fresh |
| 604 | mapping. `orig_request` is left untouched so a retry re-encodes from the |
| 605 | original request instead of carrying the previous attempt's handles. |
| 606 | """ |
| 607 | logger.info("[%s] Processing multimodal items...", req_id) |
| 608 | |
| 609 | mm_items = extract_mm_items(orig_request) |
| 610 | if not mm_items: |
| 611 | logger.info("[%s] No multimodal items, skipping encoder", req_id) |
| 612 | return {}, {} # nothing to do |
| 613 | |
| 614 | logger.info("[%s] got %d multimodal items...", req_id, len(mm_items)) |
| 615 | |
| 616 | tasks = [] |
| 617 | item_uuids: dict[int, str] = {} |
| 618 | item_transfer_ids: dict[int, str] = {} |
| 619 | item_meta: dict[int, dict] = {} |
| 620 | transfer_items: dict[int, dict[str, str]] = {} |
| 621 | ec_params: dict[str, Any] = {} |
| 622 | |
| 623 | # Round-robin over encode servers to distribute load a bit. The cursor |
| 624 | # persists across requests so fan-out doesn't restart at e_urls[0] every |
| 625 | # time (which would hot-spot the first encoder for single-item requests). |
| 626 | global encoder_rr_idx |
| 627 | async with encoder_rr_lock: |
| 628 | url_cycle, encoder_rr_idx = encoder_rr_assignment( |
| 629 | e_urls, encoder_rr_idx, len(mm_items) |
| 630 | ) |
| 631 | |
| 632 | groups: list[tuple[str, list[int]]] = [] |
| 633 | image_groups: dict[str, list[int]] = {} |
| 634 | for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)): |
| 635 | if item["type"] == "image_url": |
| 636 | indices = image_groups.get(target_url) |
| 637 | if indices is None or ( |
| 638 | ENCODER_MAX_BATCH_SIZE and len(indices) >= ENCODER_MAX_BATCH_SIZE |
| 639 | ): |
| 640 | indices = [] |
| 641 | image_groups[target_url] = indices |
| 642 | groups.append((target_url, indices)) |
| 643 | indices.append(idx) |
| 644 | else: |
| 645 | groups.append((target_url, [idx])) |
| 646 | |
| 647 | for target_url, indices in groups: |
| 648 | # Derive a *child* request id: <parent>:<index>:<random-short> |
| 649 | child_req_id = f"{req_id}:{indices[0]}:{uuid.uuid4().hex[:6]}" |
| 650 | headers = {"x-request-id": child_req_id, "Content-Type": "application/json"} |
| 651 | |
| 652 | # With --no-rewrite the decoder still receives the raw image and derives |
| 653 | # the cache key by hashing it, so the encoder must do the same -- passing |
| 654 | # a uuid here would make the two disagree and silently defeat the EC |
| 655 | # transfer, leaving the decoder to encode the image itself. |
| 656 | content = [] |
| 657 | for idx in indices: |
| 658 | item = mm_items[idx] |
| 659 | item_uuid = None if NO_REWRITE else (item.get("uuid") or content_uuid(item)) |
| 660 | if item_uuid is not None: |
| 661 | item_uuids[idx] = item_uuid |
| 662 | item_transfer_ids[idx] = uuid.uuid4().hex |
| 663 | content.append(item if item_uuid is None else {**item, "uuid": item_uuid}) |
| 664 | |
| 665 | encoder_req = { |
| 666 | "model": orig_request.get("model"), |
| 667 | "messages": [ |
| 668 | { |
| 669 | "role": "user", |
| 670 | "content": content, |
| 671 | }, |
| 672 | ], |
| 673 | # No max_tokens cap: the encoder instance never samples, it finishes |
| 674 | # once the prompt is encoded and its embeddings are published. |
| 675 | "stream": False, |
| 676 | } |
| 677 | for key in ( |
| 678 | "mm_processor_kwargs", |
| 679 | "media_io_kwargs", |
| 680 | "priority", |
| 681 | "session_id", |
| 682 | ): |
| 683 | if key in orig_request: |
| 684 | encoder_req[key] = orig_request[key] |
| 685 | if consumer_zmq is not None: |
| 686 | # The engine may rehash UUIDs with processing options. Match by |
| 687 | # position instead; batches contain only images in input order. |
| 688 | encoder_req["ec_transfer_params"] = { |
| 689 | "consumer_zmq": consumer_zmq, |
| 690 | "ec_items": [ |
| 691 | {"transfer_id": item_transfer_ids[idx]} for idx in indices |
| 692 | ], |
| 693 | } |
| 694 | tasks.append( |
| 695 | encode_session.post( |
| 696 | f"{target_url}/v1/chat/completions", |
| 697 | data=msgspec.json.encode(encoder_req), |
| 698 | headers=headers, |
| 699 | ) |
| 700 | ) |
| 701 | |
| 702 | results = await asyncio.gather(*tasks, return_exceptions=True) |
| 703 | |
| 704 | # Fail fast if any sub-request failed |
| 705 | for (_, indices), r in zip(groups, results): |
| 706 | idx = indices[0] |
| 707 | if isinstance(r, Exception): |
| 708 | logger.error( |
| 709 | "[%s] Encoder request #%d raised exception: %s", |
| 710 | req_id, |
| 711 | idx, |
| 712 | r, |
| 713 | exc_info=r, |
| 714 | ) |
| 715 | raise HTTPException( |
| 716 | status_code=502, detail=f"Encoder request failed: {str(r)}" |
| 717 | ) |
| 718 | if r.status != 200: |
| 719 | try: |
| 720 | detail = await r.text() |
| 721 | except Exception: |
| 722 | detail = "<unable to read body>" |
| 723 | logger.error( |
| 724 | "[%s] Encoder request #%d returned status %s: %s", |
| 725 | req_id, |
| 726 | idx, |
| 727 | r.status, |
| 728 | detail, |
| 729 | ) |
| 730 | raise HTTPException( |
| 731 | status_code=r.status, |
| 732 | detail=f"Encoder request failed: {detail}", |
| 733 | ) |
| 734 | |
| 735 | # The encoder reports each mm_hash's metadata (e.g. the grid) here, |
| 736 | # keyed by the same uuid this proxy assigned above. |
| 737 | try: |
| 738 | params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {} |
| 739 | except Exception: |
| 740 | logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx) |
| 741 | params = {} |
| 742 | if params: |
| 743 | by_index = {} |
| 744 | for mm_hash, reported in params.items(): |
| 745 | if len(indices) == 1: |
| 746 | break |
| 747 | for local_index in reported.get("item_indices", []): |
| 748 | if type(local_index) is not int or not 0 <= local_index < len( |
| 749 | indices |
| 750 | ): |
| 751 | raise HTTPException(502, "Invalid encoder item index") |
| 752 | if local_index in by_index: |
| 753 | raise HTTPException(502, "Duplicate encoder item index") |
| 754 | by_index[local_index] = (mm_hash, reported) |
| 755 | for local_index, idx in enumerate(indices): |
| 756 | matched = by_index.get(local_index) |
| 757 | if matched is None: |
| 758 | # Compatibility with encoders without position metadata. |
| 759 | mm_hash = item_uuids.get(idx) |
| 760 | if mm_hash in params: |
| 761 | matched = (mm_hash, params[mm_hash]) |
| 762 | elif len(indices) == 1 and len(params) == 1: |
| 763 | matched = next(iter(params.items())) |
| 764 | elif ( |
| 765 | mm_items[idx]["type"] == "video_url" |
| 766 | and (orig_request.get("mm_processor_kwargs") or {}).get( |
| 767 | "use_audio_in_video" |
| 768 | ) |
| 769 | and len(params) == 2 |
| 770 | and consumer_zmq is None |
| 771 | and all( |
| 772 | entry.keys() <= {"metadata", "item_indices"} |
| 773 | for entry in params.values() |
| 774 | ) |
| 775 | ): |
| 776 | # Keep raw video for its audio/video features when no transfer |
| 777 | # handles need to be forwarded. |
| 778 | continue |
| 779 | else: |
| 780 | raise HTTPException(502, "Encoder metadata cannot be matched") |
| 781 | ec_mm_hash, reported = matched |
| 782 | metadata = reported.get("metadata") or {} |
| 783 | item_meta[idx] = { |
| 784 | **metadata, |
| 785 | "mm_hash": item_uuids.get(idx, ec_mm_hash), |
| 786 | "ec_mm_hash": ec_mm_hash, |
| 787 | } |
| 788 | if idx in item_transfer_ids: |
| 789 | item_meta[idx]["transfer_id"] = item_transfer_ids[idx] |
| 790 | # Whatever the encoder reported alongside `metadata` is the |
| 791 | # connector's own handle on the published embedding (for NIXL, |
| 792 | # peer_host/peer_port/size_bytes). The decoder's connector |
| 793 | # looks it up by mm_hash on the request, so carry it through. |
| 794 | ec_params[ec_mm_hash] = reported |
| 795 | if NO_REWRITE and consumer_zmq is not None: |
| 796 | transfer_items[idx] = { |
| 797 | "mm_hash": ec_mm_hash, |
| 798 | "transfer_id": item_transfer_ids[idx], |
| 799 | } |
| 800 | |
| 801 | if transfer_items: |
| 802 | ec_params["ec_items"] = [transfer_items[idx] for idx in sorted(transfer_items)] |
| 803 | |
| 804 | logger.info( |
| 805 | "[%s] All %d encoder requests completed successfully", req_id, len(groups) |
| 806 | ) |
| 807 | return item_meta, ec_params |
| 808 | |
| 809 | |
| 810 | async def maybe_prefill( |
| 811 | req_data: dict, |
| 812 | p_url: str, |
| 813 | req_id: str, |
| 814 | dp_rank: int | None = None, |
| 815 | ) -> dict: |
| 816 | """- Do prefill-only task if p_url exist; |
| 817 | - Return a new body carrying kv transfer params (for nixl connector) |
| 818 | - Else, skip and return the original request data for decode |
| 819 | |
| 820 | `req_data` is never mutated: a decode retry re-enters this function with the |
| 821 | same body, and one attempt's `remote_block_ids` must not reach the next. |
| 822 | """ |
| 823 | if p_url: |
| 824 | logger.info("[%s] Processing through prefill: %s", req_id, p_url) |
| 825 | |
| 826 | prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank) |
| 827 | # for nixl connector to facilitate kv transfer... |
| 828 | prefill_response_json = msgspec.json.decode(await prefill_response.read()) |
| 829 | kv_transfer_params = prefill_response_json.get("kv_transfer_params", {}) |
| 830 | if kv_transfer_params: |
| 831 | return {**req_data, "kv_transfer_params": kv_transfer_params} |
| 832 | |
| 833 | return req_data |
| 834 | |
| 835 | |
| 836 | async def process_prefill_stage( |
| 837 | req_data: dict, |
| 838 | p_url: str, |
| 839 | req_id: str, |
| 840 | dp_rank: int | None = None, |
| 841 | ) -> dict: |
| 842 | """Process request through Prefill stage and return kv_transfer_params""" |
| 843 | logger.info("[%s] Sending prefill request to: %s", req_id, p_url) |
| 844 | |
| 845 | prefill_request = req_data.copy() |
| 846 | prefill_request["kv_transfer_params"] = { |
| 847 | "do_remote_decode": True, |
| 848 | "do_remote_prefill": False, |
| 849 | "remote_engine_id": None, |
| 850 | "remote_block_ids": None, |
| 851 | "remote_host": None, |
| 852 | "remote_port": None, |
| 853 | } |
| 854 | prefill_request["stream"] = False |
| 855 | prefill_request["max_tokens"] = 1 |
| 856 | if "max_completion_tokens" in prefill_request: |
| 857 | prefill_request["max_completion_tokens"] = 1 |
| 858 | if "stream_options" in prefill_request: |
| 859 | del prefill_request["stream_options"] |
| 860 | |
| 861 | headers = {"x-request-id": req_id, "Content-Type": "application/json"} |
| 862 | if dp_rank is not None: |
| 863 | headers["X-data-parallel-rank"] = str(dp_rank) |
| 864 | try: |
| 865 | prefill_response = await prefill_session.post( |
| 866 | f"{p_url}/v1/chat/completions", |
| 867 | data=msgspec.json.encode(prefill_request), |
| 868 | headers=headers, |
| 869 | ) |
| 870 | prefill_response.raise_for_status() |
| 871 | |
| 872 | if prefill_response.status != 200: |
| 873 | error_text = await prefill_response.text() |
| 874 | logger.error( |
| 875 | "[%s] Prefill request failed with status %d: %s", |
| 876 | req_id, |
| 877 | prefill_response.status, |
| 878 | error_text, |
| 879 | ) |
| 880 | raise HTTPException( |
| 881 | status_code=prefill_response.status, |
| 882 | detail={"error": "Prefill request failed", "message": error_text}, |
| 883 | ) |
| 884 | logger.info("[%s] Prefill request completed successfully", req_id) |
| 885 | |
| 886 | return prefill_response |
| 887 | |
| 888 | except Exception as e: |
| 889 | logger.error("Prefill processing failed: %s", str(e)) |
| 890 | raise HTTPException( |
| 891 | status_code=500, |
| 892 | detail={"error": "Prefill processing error", "message": str(e)}, |
| 893 | ) from e |
| 894 | |
| 895 | |
| 896 | ############################################################################### |
| 897 | # Middleware for request/response logging |
| 898 | ############################################################################### |
| 899 | |
| 900 | |
| 901 | async def log_requests(request: Request, call_next): |
| 902 | """Middleware to log all incoming requests and responses""" |
| 903 | req_id = request.headers.get("x-request-id", str(uuid.uuid4())) |
| 904 | |
| 905 | # Log incoming request |
| 906 | logger.info( |
| 907 | ">>> [%s] %s %s from %s", |
| 908 | req_id, |
| 909 | request.method, |
| 910 | request.url.path, |
| 911 | request.client.host if request.client else "unknown", |
| 912 | ) |
| 913 | |
| 914 | try: |
| 915 | # Process request |
| 916 | response = await call_next(request) |
| 917 | |
| 918 | # Log response |
| 919 | logger.info( |
| 920 | "<<< [%s] %s %s completed with status %d", |
| 921 | req_id, |
| 922 | request.method, |
| 923 | request.url.path, |
| 924 | response.status_code, |
| 925 | ) |
| 926 | |
| 927 | return response |
| 928 | except Exception as e: |
| 929 | # Log errors |
| 930 | logger.exception( |
| 931 | "!!! [%s] %s %s failed with error: %s", |
| 932 | req_id, |
| 933 | request.method, |
| 934 | request.url.path, |
| 935 | str(e), |
| 936 | ) |
| 937 | raise |
| 938 | |
| 939 | |
| 940 | ############################################################################### |
| 941 | # FastAPI lifecycle |
| 942 | ############################################################################### |
| 943 | |
| 944 | |
| 945 | @app.on_event("startup") |
| 946 | async def on_startup() -> None: |
| 947 | global encode_session, prefill_session, decode_session |
| 948 | timeout = aiohttp.ClientTimeout(total=100_000) |
| 949 | # vLLM closes an idle keep-alive connection after |
| 950 | # VLLM_HTTP_TIMEOUT_KEEP_ALIVE seconds (5 by default), while aiohttp keeps |
| 951 | # pooling it for 15. Reusing one it has already closed fails the request |
| 952 | # with ServerDisconnectedError, and the server logs nothing at all: it |
| 953 | # closed the socket before the request arrived. Retire ours first. |
| 954 | server_keep_alive = float(os.getenv("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5")) |
| 955 | connector = aiohttp.TCPConnector( |
| 956 | limit=0, |
| 957 | **( |
| 958 | {"keepalive_timeout": server_keep_alive / 2} |
| 959 | if server_keep_alive > 0 |
| 960 | else {"force_close": True} |
| 961 | ), |
| 962 | ) |
| 963 | encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) |
| 964 | prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) |
| 965 | decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) |
| 966 | |
| 967 | |
| 968 | @app.on_event("shutdown") |
| 969 | async def on_shutdown() -> None: |
| 970 | global encode_session, prefill_session, decode_session |
| 971 | if encode_session: |
| 972 | await encode_session.close() |
| 973 | if prefill_session: |
| 974 | await prefill_session.close() |
| 975 | if decode_session: |
| 976 | await decode_session.close() |
| 977 | |
| 978 | |
| 979 | ############################################################################### |
| 980 | # Core forwarding |
| 981 | ############################################################################### |
| 982 | |
| 983 | |
| 984 | async def prepare_for_decode( |
| 985 | req_data: dict, |
| 986 | req_id: str, |
| 987 | e_urls: list[str], |
| 988 | p_url: str, |
| 989 | consumer_zmq: str | None, |
| 990 | prefill_dp_rank: int | None = None, |
| 991 | ) -> tuple[dict, float, float]: |
| 992 | """Encode, rewrite and prefill, returning the body to send to decode. |
| 993 | |
| 994 | `req_data` is left untouched so a retry starts from the original media |
| 995 | rather than from a body whose images are already metadata references. |
| 996 | """ |
| 997 | _t0 = time.perf_counter() |
| 998 | item_meta, ec_params = await fanout_encoder_primer( |
| 999 | req_data, e_urls, req_id, consumer_zmq |
| 1000 | ) |
| 1001 | _t1 = time.perf_counter() |
| 1002 | prepared = req_data if NO_REWRITE else rewrite_for_decode(req_data, item_meta) |
| 1003 | if ec_params: |
| 1004 | # A fresh body every time: `rewrite_for_decode` hands back `req_data` |
| 1005 | # itself when it rewrote nothing, and this attempt's handles must not |
| 1006 | # outlive it into a retry. |
| 1007 | handles = dict(prepared.get("ec_transfer_params") or {}) |
| 1008 | handles.update(ec_params) |
| 1009 | prepared = {**prepared, "ec_transfer_params": handles} |
| 1010 | _t2 = time.perf_counter() |
| 1011 | prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank) |
| 1012 | return prepared, _t1 - _t0, _t2 - _t1 |
| 1013 | |
| 1014 | |
| 1015 | async def forward_non_stream( |
| 1016 | req_data: dict, |
| 1017 | req_id: str, |
| 1018 | e_urls: list[str], |
| 1019 | p_url: str, |
| 1020 | d_url: str, |
| 1021 | consumer_zmq: str | None, |
| 1022 | dp_rank: int | None = None, |
| 1023 | prefill_dp_rank: int | None = None, |
| 1024 | ) -> Response: |
| 1025 | try: |
| 1026 | for attempt in range(DECODE_RETRIES + 1): |
| 1027 | _t0 = time.perf_counter() |
| 1028 | prepared, encode_s, rewrite_s = await prepare_for_decode( |
| 1029 | req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank |
| 1030 | ) |
| 1031 | _t2 = time.perf_counter() |
| 1032 | |
| 1033 | logger.info("[%s] Forwarding to decode: %s", req_id, d_url) |
| 1034 | headers = {"x-request-id": req_id, "Content-Type": "application/json"} |
| 1035 | if dp_rank is not None: |
| 1036 | headers["X-data-parallel-rank"] = str(dp_rank) |
| 1037 | |
| 1038 | async with decode_session.post( |
| 1039 | f"{d_url}/v1/chat/completions", |
| 1040 | data=msgspec.json.encode(prepared), |
| 1041 | headers=headers, |
| 1042 | ) as resp: |
| 1043 | if resp.status >= 400: |
| 1044 | detail = await resp.text() |
| 1045 | # 500 is the decoder's retryable internal error, which |
| 1046 | # includes an encoder embedding it could not obtain. Redoing |
| 1047 | # the encode publishes the item again. |
| 1048 | if resp.status == 500 and attempt < DECODE_RETRIES: |
| 1049 | logger.warning( |
| 1050 | "[%s] Decode returned 500, re-encoding and retrying " |
| 1051 | "(attempt %d/%d): %s", |
| 1052 | req_id, |
| 1053 | attempt + 1, |
| 1054 | DECODE_RETRIES, |
| 1055 | detail[:200], |
| 1056 | ) |
| 1057 | continue |
| 1058 | logger.error( |
| 1059 | "[%s] Decode request returned status %s: %s", |
| 1060 | req_id, |
| 1061 | resp.status, |
| 1062 | detail, |
| 1063 | ) |
| 1064 | raise HTTPException(status_code=resp.status, detail=detail) |
| 1065 | out = await resp.read() |
| 1066 | _t3 = time.perf_counter() |
| 1067 | logger.info( |
| 1068 | "STAGE %s encode=%.1f rewrite=%.1f decode=%.1f total=%.1f " |
| 1069 | "attempt=%d", |
| 1070 | "no-rewrite" if NO_REWRITE else "rewrite", |
| 1071 | encode_s * 1e3, |
| 1072 | rewrite_s * 1e3, |
| 1073 | (_t3 - _t2) * 1e3, |
| 1074 | (_t3 - _t0) * 1e3, |
| 1075 | attempt, |
| 1076 | ) |
| 1077 | return Response( |
| 1078 | content=out, |
| 1079 | status_code=resp.status, |
| 1080 | headers={ |
| 1081 | "Content-Type": resp.headers.get( |
| 1082 | "Content-Type", "application/json" |
| 1083 | ) |
| 1084 | }, |
| 1085 | ) |
| 1086 | raise HTTPException(status_code=500, detail="Decode failed after re-encoding") |
| 1087 | |
| 1088 | except HTTPException: |
| 1089 | raise |
| 1090 | except Exception as e: |
| 1091 | logger.exception("[%s] Error in forward_non_stream: %s", req_id, str(e)) |
| 1092 | raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") from e |
| 1093 | |
| 1094 | |
| 1095 | async def forward_stream( |
| 1096 | req_data: dict, |
| 1097 | req_id: str, |
| 1098 | e_urls: list[str], |
| 1099 | p_url: str, |
| 1100 | d_url: str, |
| 1101 | consumer_zmq: str | None, |
| 1102 | dp_rank: int | None = None, |
| 1103 | prefill_dp_rank: int | None = None, |
| 1104 | ) -> AsyncIterator[bytes]: |
| 1105 | try: |
| 1106 | for attempt in range(DECODE_RETRIES + 1): |
| 1107 | _t0 = time.perf_counter() |
| 1108 | prepared, encode_s, rewrite_s = await prepare_for_decode( |
| 1109 | req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank |
| 1110 | ) |
| 1111 | _t2 = time.perf_counter() |
| 1112 | |
| 1113 | logger.info("[%s] Starting streaming from decode: %s", req_id, d_url) |
| 1114 | headers = {"x-request-id": req_id, "Content-Type": "application/json"} |
| 1115 | if dp_rank is not None: |
| 1116 | headers["X-data-parallel-rank"] = str(dp_rank) |
| 1117 | |
| 1118 | _first = None |
| 1119 | async with decode_session.post( |
| 1120 | f"{d_url}/v1/chat/completions", |
| 1121 | data=msgspec.json.encode(prepared), |
| 1122 | headers=headers, |
| 1123 | ) as resp: |
| 1124 | # Retry only before the first chunk: once anything reached the |
| 1125 | # client the response cannot be replaced. |
| 1126 | if resp.status == 500 and attempt < DECODE_RETRIES: |
| 1127 | detail = await resp.text() |
| 1128 | logger.warning( |
| 1129 | "[%s] Decode returned 500 before streaming, re-encoding " |
| 1130 | "and retrying (attempt %d/%d): %s", |
| 1131 | req_id, |
| 1132 | attempt + 1, |
| 1133 | DECODE_RETRIES, |
| 1134 | detail[:200], |
| 1135 | ) |
| 1136 | continue |
| 1137 | resp.raise_for_status() |
| 1138 | async for chunk in resp.content.iter_any(): |
| 1139 | if chunk: |
| 1140 | if _first is None: |
| 1141 | _first = time.perf_counter() |
| 1142 | yield chunk |
| 1143 | _t3 = time.perf_counter() |
| 1144 | |
| 1145 | logger.info( |
| 1146 | "STAGE %s encode=%.1f rewrite=%.2f decode_ttfb=%.1f " |
| 1147 | "decode_total=%.1f attempt=%d", |
| 1148 | "no-rewrite" if NO_REWRITE else "rewrite", |
| 1149 | encode_s * 1e3, |
| 1150 | rewrite_s * 1e3, |
| 1151 | ((_first or _t3) - _t2) * 1e3, |
| 1152 | (_t3 - _t2) * 1e3, |
| 1153 | attempt, |
| 1154 | ) |
| 1155 | logger.info("[%s] Streaming completed", req_id) |
| 1156 | return |
| 1157 | |
| 1158 | except HTTPException: |
| 1159 | logger.exception("[%s] HTTPException in forward_stream", req_id) |
| 1160 | raise |
| 1161 | except Exception as e: |
| 1162 | logger.exception("[%s] Error in forward_stream: %s", req_id, str(e)) |
| 1163 | raise HTTPException( |
| 1164 | status_code=500, detail=f"Proxy streaming error: {str(e)}" |
| 1165 | ) from e |
| 1166 | |
| 1167 | |
| 1168 | ############################################################################### |
| 1169 | # Public routes |
| 1170 | ############################################################################### |
| 1171 | |
| 1172 | |
| 1173 | @app.post("/v1/chat/completions") |
| 1174 | async def chat_completions(request: Request): |
| 1175 | try: |
| 1176 | req_data = msgspec.json.decode(await request.body()) |
| 1177 | req_id = request.headers.get("x-request-id", str(uuid.uuid4())) |
| 1178 | |
| 1179 | e_urls = app.state.e_urls # we want the full list for fan-out |
| 1180 | p_url = random.choice(app.state.p_urls) if app.state.p_urls else None |
| 1181 | decode_index = random.randrange(len(app.state.d_urls)) |
| 1182 | d_url = app.state.d_urls[decode_index] |
| 1183 | dp_size = app.state.ec_consumer_dp_size |
| 1184 | # Round-robin the replica, then name it to both halves: the decoder |
| 1185 | # honours the rank header instead of its own balancer, and the encoder |
| 1186 | # pushes to that replica's control channel. Choosing once here means a |
| 1187 | # decode retry re-encodes to the same replica. |
| 1188 | dp_rank = next(app.state.replica_counter) % dp_size if dp_size > 1 else None |
| 1189 | ec_index = decode_index * dp_size + (dp_rank or 0) |
| 1190 | consumer_zmq = app.state.d_ec_urls[ec_index] if app.state.d_ec_urls else None |
| 1191 | |
| 1192 | is_streaming = req_data.get("stream", False) |
| 1193 | |
| 1194 | if is_streaming: |
| 1195 | return StreamingResponse( |
| 1196 | forward_stream( |
| 1197 | req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank |
| 1198 | ), |
| 1199 | media_type="text/event-stream", |
| 1200 | ) |
| 1201 | return await forward_non_stream( |
| 1202 | req_data, req_id, e_urls, p_url, d_url, consumer_zmq, dp_rank |
| 1203 | ) |
| 1204 | |
| 1205 | except HTTPException: |
| 1206 | raise |
| 1207 | except Exception as e: |
| 1208 | logger.exception("Error in chat_completions endpoint: %s", str(e)) |
| 1209 | raise HTTPException( |
| 1210 | status_code=500, detail=f"Request processing error: {str(e)}" |
| 1211 | ) from e |
| 1212 | |
| 1213 | |
| 1214 | @app.get("/v1/models") |
| 1215 | async def list_models(): |
| 1216 | async with decode_session.get(f"{app.state.d_urls[0]}/v1/models") as resp: |
| 1217 | resp.raise_for_status() |
| 1218 | return await resp.json() |
| 1219 | |
| 1220 | |
| 1221 | @app.get("/health") |
| 1222 | async def health_check(): |
| 1223 | async def healthy(urls): |
| 1224 | if not urls: |
| 1225 | return "empty" |
| 1226 | for u in urls: |
| 1227 | try: |
| 1228 | async with encode_session.get(f"{u}/health") as resp: |
| 1229 | resp.raise_for_status() |
| 1230 | except Exception: |
| 1231 | return "unhealthy" |
| 1232 | return "healthy" |
| 1233 | |
| 1234 | e_status, p_status, d_status = await asyncio.gather( |
| 1235 | healthy(app.state.e_urls), healthy(app.state.p_urls), healthy(app.state.d_urls) |
| 1236 | ) |
| 1237 | |
| 1238 | overall_healthy = all( |
| 1239 | status != "unhealthy" for status in (e_status, p_status, d_status) |
| 1240 | ) |
| 1241 | |
| 1242 | status_code = 200 if overall_healthy else 503 |
| 1243 | |
| 1244 | return JSONResponse( |
| 1245 | { |
| 1246 | "proxy": "healthy", |
| 1247 | "encode_cluster": e_status, |
| 1248 | "prefill_cluster": p_status, |
| 1249 | "decode_cluster": d_status, |
| 1250 | }, |
| 1251 | status_code=status_code, |
| 1252 | ) |
| 1253 | |
| 1254 | |
| 1255 | ############################################################################### |
| 1256 | # Simple profiler fan-out (unchanged except for sessions) |
| 1257 | ############################################################################### |
| 1258 | |
| 1259 | |
| 1260 | async def _post_if_available( |
| 1261 | session: aiohttp.ClientSession, |
| 1262 | url: str, |
| 1263 | payload: dict, |
| 1264 | headers: dict, |
| 1265 | ) -> dict | None: |
| 1266 | """POST `payload` to `url`. |
| 1267 | |
| 1268 | Returns |
| 1269 | ------- |
| 1270 | • The decoded JSON body on success (2xx) |
| 1271 | • None if the endpoint does not exist (404) |
| 1272 | • Raises for anything else. |
| 1273 | |
| 1274 | """ |
| 1275 | try: |
| 1276 | resp = await session.post(url, json=payload, headers=headers) |
| 1277 | if resp.status == 404: # profiling disabled on that server |
| 1278 | logger.warning("Profiling endpoint missing on %s", url) |
| 1279 | return None |
| 1280 | resp.raise_for_status() |
| 1281 | return await resp.json(content_type=None) |
| 1282 | except aiohttp.ClientResponseError as exc: |
| 1283 | # Pass 404 through the branch above, re-raise everything else |
| 1284 | if exc.status == 404: |
| 1285 | logger.warning("Profiling endpoint missing on %s", url) |
| 1286 | return None |
| 1287 | raise |
| 1288 | except Exception: |
| 1289 | # Network errors etc.: propagate |
| 1290 | raise |
| 1291 | |
| 1292 | |
| 1293 | async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str): |
| 1294 | """Fire & forget to both clusters, tolerate 404.""" |
| 1295 | headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} |
| 1296 | |
| 1297 | encode_task = _post_if_available( |
| 1298 | encode_session, f"{e_url}/{cmd}_profile", payload, headers |
| 1299 | ) |
| 1300 | prefill_task = ( |
| 1301 | _post_if_available(prefill_session, f"{p_url}/{cmd}_profile", payload, headers) |
| 1302 | if p_url is not None |
| 1303 | else asyncio.sleep(0) |
| 1304 | ) |
| 1305 | decode_task = _post_if_available( |
| 1306 | decode_session, f"{d_url}/{cmd}_profile", payload, headers |
| 1307 | ) |
| 1308 | |
| 1309 | encode_res, prefill_res, decode_res = await asyncio.gather( |
| 1310 | encode_task, prefill_task, decode_task |
| 1311 | ) |
| 1312 | |
| 1313 | # If *all* clusters said “I don’t have that route”, surface an error |
| 1314 | if encode_res is prefill_res is decode_res is None: |
| 1315 | raise HTTPException( |
| 1316 | status_code=503, |
| 1317 | detail="Profiling endpoints are disabled on all clusters", |
| 1318 | ) |
| 1319 | |
| 1320 | return { |
| 1321 | "encode": encode_res, # may be None |
| 1322 | "prefill": prefill_res, # may be None |
| 1323 | "decode": decode_res, # may be None |
| 1324 | } |
| 1325 | |
| 1326 | |
| 1327 | @app.post("/start_profile") |
| 1328 | async def start_profile(request: Request): |
| 1329 | body = await request.json() |
| 1330 | # TODO: handle multi urls properly |
| 1331 | e_url = random.choice(app.state.e_urls) |
| 1332 | p_url = random.choice(app.state.p_urls) if app.state.p_urls else None |
| 1333 | d_url = random.choice(app.state.d_urls) |
| 1334 | return await _profile_cmd("start", body, e_url, p_url, d_url) |
| 1335 | |
| 1336 | |
| 1337 | @app.post("/stop_profile") |
| 1338 | async def stop_profile(request: Request): |
| 1339 | body = await request.json() |
| 1340 | # TODO: handle multi urls properly |
| 1341 | e_url = random.choice(app.state.e_urls) |
| 1342 | p_url = random.choice(app.state.p_urls) if app.state.p_urls else None |
| 1343 | d_url = random.choice(app.state.d_urls) |
| 1344 | return await _profile_cmd("stop", body, e_url, p_url, d_url) |
| 1345 | |
| 1346 | |
| 1347 | def build_app(config: EPDProxyConfig | None = None) -> FastAPI: |
| 1348 | config = config or EPDProxyConfig() |
| 1349 | registry = InstanceRegistry( |
| 1350 | probe_interval=config.probe_interval, |
| 1351 | probe_timeout=config.probe_timeout, |
| 1352 | fail_threshold=config.fail_threshold, |
| 1353 | evicted_ttl=config.evicted_ttl, |
| 1354 | ) |
| 1355 | proxy = EPDProxy(registry) |
| 1356 | |
| 1357 | @asynccontextmanager |
| 1358 | async def lifespan(app: FastAPI): |
| 1359 | await on_startup() |
| 1360 | try: |
| 1361 | registry.start_probing() |
| 1362 | yield |
| 1363 | finally: |
| 1364 | await registry.stop_probing() |
| 1365 | await on_shutdown() |
| 1366 | |
| 1367 | app = FastAPI(lifespan=lifespan) |
| 1368 | app.state.proxy = proxy |
| 1369 | app.state.registry = registry |
| 1370 | |
| 1371 | @app.get("/instances") |
| 1372 | async def list_instances(): |
| 1373 | return registry.status() |
| 1374 | |
| 1375 | @app.post("/instances", dependencies=[Depends(require_admin_key)]) |
| 1376 | async def register_instance(body: InstanceRegistration): |
| 1377 | if body.ec_zmq_addrs and body.role not in ( |
| 1378 | InstanceRole.PREFILL, |
| 1379 | InstanceRole.PREFILL_DECODE, |
| 1380 | ): |
| 1381 | raise HTTPException(400, "Only EC consumers accept Mooncake addresses") |
| 1382 | try: |
| 1383 | created = registry.register( |
| 1384 | InstanceRecord( |
| 1385 | body.role, |
| 1386 | str(body.url).rstrip("/"), |
| 1387 | body.ec_zmq_addrs, |
| 1388 | body.dp_size, |
| 1389 | ) |
| 1390 | ) |
| 1391 | except ValueError as error: |
| 1392 | raise HTTPException(409, str(error)) from error |
| 1393 | return {"registered": created} |
| 1394 | |
| 1395 | @app.delete("/instances", dependencies=[Depends(require_admin_key)]) |
| 1396 | async def unregister_instance(url: AnyHttpUrl): |
| 1397 | return {"removed": registry.unregister(str(url).rstrip("/"))} |
| 1398 | |
| 1399 | @app.post("/v1/chat/completions") |
| 1400 | async def chat_completions(request: Request): |
| 1401 | req_data = msgspec.json.decode(await request.body()) |
| 1402 | req_id = request.headers.get("x-request-id", str(uuid.uuid4())) |
| 1403 | route = proxy.route(len(extract_mm_items(req_data))) |
| 1404 | args = ( |
| 1405 | req_data, |
| 1406 | req_id, |
| 1407 | route.encoder_urls, |
| 1408 | route.prefill.url if route.prefill else None, |
| 1409 | route.decode.url, |
| 1410 | route.consumer_zmq, |
| 1411 | ) |
| 1412 | ranks = { |
| 1413 | "dp_rank": route.dp_rank if route.prefill is None else None, |
| 1414 | "prefill_dp_rank": route.dp_rank if route.prefill is not None else None, |
| 1415 | } |
| 1416 | if req_data.get("stream", False): |
| 1417 | return StreamingResponse( |
| 1418 | forward_stream(*args, **ranks), media_type="text/event-stream" |
| 1419 | ) |
| 1420 | return await forward_non_stream(*args, **ranks) |
| 1421 | |
| 1422 | @app.get("/v1/models") |
| 1423 | async def list_models(): |
| 1424 | decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick( |
| 1425 | InstanceRole.DECODE |
| 1426 | ) |
| 1427 | if decode is None: |
| 1428 | raise HTTPException( |
| 1429 | status_code=503, detail="No decode instance is registered" |
| 1430 | ) |
| 1431 | async with decode_session.get(f"{decode.url}/v1/models") as resp: |
| 1432 | resp.raise_for_status() |
| 1433 | return await resp.json() |
| 1434 | |
| 1435 | @app.get("/health") |
| 1436 | async def health(): |
| 1437 | status = registry.status() |
| 1438 | # An empty roster is not unhealthy: the proxy is meant to come up |
| 1439 | # before anything registers with it. |
| 1440 | return JSONResponse({"proxy": "healthy", "instances": status}) |
| 1441 | |
| 1442 | @app.post("/start_profile") |
| 1443 | async def start_profile(request: Request): |
| 1444 | return await _profile(registry, "start", await request.json()) |
| 1445 | |
| 1446 | @app.post("/stop_profile") |
| 1447 | async def stop_profile(request: Request): |
| 1448 | return await _profile(registry, "stop", await request.json()) |
| 1449 | |
| 1450 | return app |
| 1451 | |
| 1452 | |
| 1453 | async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict: |
| 1454 | headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} |
| 1455 | targets = {role.value: registry.pick(role) for role in InstanceRole} |
| 1456 | results = await asyncio.gather( |
| 1457 | *( |
| 1458 | _post_if_available( |
| 1459 | decode_session, f"{record.url}/{cmd}_profile", payload, headers |
| 1460 | ) |
| 1461 | for record in targets.values() |
| 1462 | if record is not None |
| 1463 | ) |
| 1464 | ) |
| 1465 | reachable = [result for result in results if result is not None] |
| 1466 | if not reachable: |
| 1467 | raise HTTPException( |
| 1468 | status_code=503, |
| 1469 | detail="Profiling endpoints are disabled on every instance", |
| 1470 | ) |
| 1471 | live = [name for name, record in targets.items() if record is not None] |
| 1472 | return dict(zip(live, results)) |
| 1473 | |
| 1474 | |
| 1475 | if __name__ == "__main__": |
| 1476 | parser = argparse.ArgumentParser() |
| 1477 | parser.add_argument("--host", default="0.0.0.0") |
| 1478 | parser.add_argument("--port", type=int, default=8000) |
| 1479 | parser.add_argument( |
| 1480 | "--dynamic-registration", |
| 1481 | action="store_true", |
| 1482 | help="Enable launcher-managed HTTP registration for E+PD or E+P+D.", |
| 1483 | ) |
| 1484 | parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL) |
| 1485 | parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT) |
| 1486 | parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD) |
| 1487 | parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL) |
| 1488 | parser.add_argument( |
| 1489 | "--log-requests", |
| 1490 | action="store_true", |
| 1491 | help=( |
| 1492 | "Log every request in and out, and raise the log level to DEBUG. " |
| 1493 | "Off by default: the proxy is on the request path." |
| 1494 | ), |
| 1495 | ) |
| 1496 | parser.add_argument( |
| 1497 | "--no-rewrite", |
| 1498 | action="store_true", |
| 1499 | help="Forward images to the decoder unchanged (for stage-timing A/B).", |
| 1500 | ) |
| 1501 | parser.add_argument( |
| 1502 | "--encode-servers-urls", |
| 1503 | default="", |
| 1504 | help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")', |
| 1505 | ) |
| 1506 | parser.add_argument( |
| 1507 | "--prefill-servers-urls", |
| 1508 | default="none", |
| 1509 | help=( |
| 1510 | 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ' |
| 1511 | 'to enable E->P->D, set "disable" or "none" to enable E->PD' |
| 1512 | ), |
| 1513 | ) |
| 1514 | parser.add_argument( |
| 1515 | "--decode-servers-urls", |
| 1516 | default="", |
| 1517 | help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")', |
| 1518 | ) |
| 1519 | parser.add_argument( |
| 1520 | "--decode-retries", |
| 1521 | type=int, |
| 1522 | default=1, |
| 1523 | help=( |
| 1524 | "Re-encode and re-send when decode returns 500, which is its " |
| 1525 | "retryable internal error (an undeliverable encoder embedding " |
| 1526 | "among them). 0 disables." |
| 1527 | ), |
| 1528 | ) |
| 1529 | parser.add_argument( |
| 1530 | "--ec-consumer-zmq-addrs", |
| 1531 | default="", |
| 1532 | help=( |
| 1533 | "Comma-separated Mooncake EC consumer control addresses, aligned " |
| 1534 | "with --decode-servers-urls. Required for Mooncake EC consumers and " |
| 1535 | "supported only in E+PD mode. With --ec-consumer-dp-size > 1, list " |
| 1536 | "each server's replicas consecutively: s0r0,s0r1,s1r0,s1r1." |
| 1537 | ), |
| 1538 | ) |
| 1539 | parser.add_argument( |
| 1540 | "--ec-consumer-dp-size", |
| 1541 | type=int, |
| 1542 | default=1, |
| 1543 | help=( |
| 1544 | "Data-parallel replicas per EC consumer. The proxy picks a replica " |
| 1545 | "round-robin and names it to both halves of the request, because an " |
| 1546 | "encoder push has to land where the request will run." |
| 1547 | ), |
| 1548 | ) |
| 1549 | |
| 1550 | args = parser.parse_args() |
| 1551 | if args.fail_threshold < 1: |
| 1552 | parser.error("--fail-threshold must be at least 1") |
| 1553 | if args.log_requests: |
| 1554 | logging.getLogger().setLevel(logging.DEBUG) |
| 1555 | app.middleware("http")(log_requests) |
| 1556 | NO_REWRITE = args.no_rewrite |
| 1557 | DECODE_RETRIES = max(0, args.decode_retries) |
| 1558 | if args.dynamic_registration: |
| 1559 | if not os.getenv("ADMIN_API_KEY"): |
| 1560 | parser.error("--dynamic-registration requires ADMIN_API_KEY") |
| 1561 | if ( |
| 1562 | args.encode_servers_urls |
| 1563 | or args.prefill_servers_urls.lower() not in ("disable", "none", "") |
| 1564 | or args.decode_servers_urls |
| 1565 | or args.ec_consumer_zmq_addrs |
| 1566 | ): |
| 1567 | parser.error( |
| 1568 | "With --dynamic-registration, instances register through /instances" |
| 1569 | ) |
| 1570 | app = build_app( |
| 1571 | EPDProxyConfig( |
| 1572 | probe_interval=args.probe_interval, |
| 1573 | probe_timeout=args.probe_timeout, |
| 1574 | fail_threshold=args.fail_threshold, |
| 1575 | evicted_ttl=args.evicted_ttl, |
| 1576 | ) |
| 1577 | ) |
| 1578 | if args.log_requests: |
| 1579 | app.middleware("http")(log_requests) |
| 1580 | else: |
| 1581 | if not args.encode_servers_urls or not args.decode_servers_urls: |
| 1582 | parser.error( |
| 1583 | "Static routing requires --encode-servers-urls " |
| 1584 | "and --decode-servers-urls" |
| 1585 | ) |
| 1586 | app.state.e_urls = [ |
| 1587 | u.strip() for u in args.encode_servers_urls.split(",") if u.strip() |
| 1588 | ] |
| 1589 | app.state.d_urls = [ |
| 1590 | u.strip() for u in args.decode_servers_urls.split(",") if u.strip() |
| 1591 | ] |
| 1592 | app.state.d_ec_urls = [ |
| 1593 | u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() |
| 1594 | ] |
| 1595 | if args.ec_consumer_dp_size < 1: |
| 1596 | parser.error("--ec-consumer-dp-size must be at least 1") |
| 1597 | app.state.ec_consumer_dp_size = args.ec_consumer_dp_size |
| 1598 | app.state.replica_counter = itertools.count() |
| 1599 | expected = len(app.state.d_urls) * args.ec_consumer_dp_size |
| 1600 | if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: |
| 1601 | parser.error( |
| 1602 | "--ec-consumer-zmq-addrs must contain one address per consumer " |
| 1603 | f"replica: expected {expected} " |
| 1604 | f"({len(app.state.d_urls)} servers x " |
| 1605 | f"{args.ec_consumer_dp_size} replicas), " |
| 1606 | f"got {len(app.state.d_ec_urls)}" |
| 1607 | ) |
| 1608 | # handle prefill instances |
| 1609 | if args.prefill_servers_urls.lower() in ("disable", "none", ""): |
| 1610 | app.state.p_urls = [] |
| 1611 | logger.info("Disaggregated prefill disabled. Running E + PD...") |
| 1612 | else: |
| 1613 | app.state.p_urls = [ |
| 1614 | u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() |
| 1615 | ] |
| 1616 | logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") |
| 1617 | try: |
| 1618 | validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) |
| 1619 | except ValueError as exc: |
| 1620 | parser.error(str(exc)) |
| 1621 | |
| 1622 | logger.info("Proxy listening on %s:%s", args.host, args.port) |
| 1623 | logger.info("Encode servers: %s", app.state.e_urls) |
| 1624 | logger.info("Prefill instances %s", app.state.p_urls) |
| 1625 | logger.info("Decode servers: %s", app.state.d_urls) |
| 1626 | if app.state.ec_consumer_dp_size > 1: |
| 1627 | logger.info( |
| 1628 | "EC consumer replicas per server: %d (control addresses: %s)", |
| 1629 | app.state.ec_consumer_dp_size, |
| 1630 | app.state.d_ec_urls, |
| 1631 | ) |
| 1632 | |
| 1633 | uvicorn.run( |
| 1634 | app, |
| 1635 | host=args.host, |
| 1636 | port=args.port, |
| 1637 | log_level="info", |
| 1638 | loop="uvloop", |
| 1639 | access_log=True, |
| 1640 | ) |