已合并
支持Nodemanager原生拉起引擎，解除Engine Server层依赖 #698
tobking创建于 23 天前
支持Nodemanager原生拉起引擎，解除Engine Server层依赖 #698
已合并
共 127 个文件变更+7168-3884
| @@ -38,6 +38,11 @@ Endpoint (ABC) → MgmtEndpoint / InferEndpoint — HTTP server lifecycl | |||
| 38 | InferEndpoint.get_lifespan() → VLLMEndpoint / SGLangEndpoint — engine-specific startup | 38 | InferEndpoint.get_lifespan() → VLLMEndpoint / SGLangEndpoint — engine-specific startup |
| 39 | ``` | 39 | ``` |
| 40 | 40 | ||
| 41 | +The transitional EngineServer reuses the CLI configuration adapters owned by | ||
| 42 | +`motor/node_manager/core/services/native_engine/backends/`. Process launch and supervision for the | ||
| 43 | +new direct-native path remain in `NativeEngineService`; the classes in this document describe only | ||
| 44 | +the legacy in-process EngineServer path. | ||
| 45 | + | ||
| 41 | ### Config → CLI Args Pipeline | 46 | ### Config → CLI Args Pipeline |
| 42 | 47 | ||
| 43 | 1. `ConfigFactory` looks up the engine's `IConfig` in an explicit `_ENGINE_CONFIG_MAP` dict (`vllm` / `sglang` → module paths); `parse()` runs the full `initialize()` → `convert()` → `validate()` pipeline | 48 | 1. `ConfigFactory` looks up the engine's `IConfig` in an explicit `_ENGINE_CONFIG_MAP` dict (`vllm` / `sglang` → module paths); `parse()` runs the full `initialize()` → `convert()` → `validate()` pipeline |
| @@ -176,7 +181,7 @@ Layer 4: SimInference (proactive health) | |||
| 176 | | File | Role | | 181 | | File | Role | |
| 177 | |------|------| | 182 | |------|------| |
| 178 | | `motor/engine_server/cli/main.py` | Entry point: CLI arg parsing, factory wiring, start Mgmt+Infer endpoints | | 183 | | `motor/engine_server/cli/main.py` | Entry point: CLI arg parsing, factory wiring, start Mgmt+Infer endpoints | |
| 179 | -| `motor/engine_server/core/config.py` | `IConfig` ABC: `initialize()`, `validate()`, `convert()`, `get_args()` | | 184 | +| `motor/node_manager/core/services/native_engine/backends/base.py` | Shared `IConfig` ABC used by the transitional EngineServer and native backends | |
| 180 | | `motor/engine_server/core/engine.py` | `Engine` ABC: `launch()`, `shutdown()` | | 185 | | `motor/engine_server/core/engine.py` | `Engine` ABC: `launch()`, `shutdown()` | |
| 181 | | `motor/engine_server/core/endpoint.py` | `Endpoint` ABC: `run()`, uvicorn lifecycle | | 186 | | `motor/engine_server/core/endpoint.py` | `Endpoint` ABC: `run()`, uvicorn lifecycle | |
| 182 | | `motor/engine_server/core/infer_endpoint.py` | `InferEndpoint`: FastAPI app, uvicorn, route registration, lifespan, dispatch adapter wiring, TLS | | 187 | | `motor/engine_server/core/infer_endpoint.py` | `InferEndpoint`: FastAPI app, uvicorn, route registration, lifespan, dispatch adapter wiring, TLS | |
| @@ -186,12 +191,13 @@ Layer 4: SimInference (proactive health) | |||
| 186 | | `motor/engine_server/core/dispatch_adapter/` | Dispatch adapter subsystem: `base.py` (534), `vllm_adapter.py` (340), `normalization.py` (220), `sglang_adapter.py` (49), `factory.py` (23) | | 191 | | `motor/engine_server/core/dispatch_adapter/` | Dispatch adapter subsystem: `base.py` (534), `vllm_adapter.py` (340), `normalization.py` (220), `sglang_adapter.py` (49), `factory.py` (23) | |
| 187 | | `motor/engine_server/core/snapshot_sentinel.py` | `SnapshotSentinel` thread: checkpoint wait + suspend/resume driving | | 192 | | `motor/engine_server/core/snapshot_sentinel.py` | `SnapshotSentinel` thread: checkpoint wait + suspend/resume driving | |
| 188 | | `motor/engine_server/core/snapshot_monitor.py` | `SnapshotMonitor`: suspend/unlock/resume completion states | | 193 | | `motor/engine_server/core/snapshot_monitor.py` | `SnapshotMonitor`: suspend/unlock/resume completion states | |
| 189 | -| `motor/engine_server/core/vllm/vllm_config.py` | `VLLMConfig`: field mapping, DP address, KV transfer, D2D config | | 194 | +| `motor/node_manager/core/services/native_engine/backends/vllm/config.py` | `VLLMConfig`: field mapping, DP address, KV transfer, D2D config | |
| 190 | | `motor/engine_server/core/vllm/vllm_engine.py` | `VLLMEngine`: `AsyncEngineArgs.from_cli_args` + `AsyncLLM.from_vllm_config`, headless PCP follower | | 195 | | `motor/engine_server/core/vllm/vllm_engine.py` | `VLLMEngine`: `AsyncEngineArgs.from_cli_args` + `AsyncLLM.from_vllm_config`, headless PCP follower | |
| 191 | | `motor/engine_server/core/vllm/vllm_endpoint.py` | `VLLMEndpoint`: `_vllm_lifespan` context manager, route init | | 196 | | `motor/engine_server/core/vllm/vllm_endpoint.py` | `VLLMEndpoint`: `_vllm_lifespan` context manager, route init | |
| 192 | | `motor/engine_server/core/vllm/vllm_openai_compat.py` | OpenAI-compat shims (model lists, request mapping) for the vLLM backend | | 197 | | `motor/engine_server/core/vllm/vllm_openai_compat.py` | OpenAI-compat shims (model lists, request mapping) for the vLLM backend | |
| 193 | -| `motor/engine_server/core/sglang/` | `SGLangConfig`, `SGLangEngine`, `SGLangEndpoint` | | 198 | +| `motor/node_manager/core/services/native_engine/backends/sglang/config.py` | `SGLangConfig`: native CLI conversion and validation | |
| 194 | -| `motor/engine_server/factory/config_factory.py` | `ConfigFactory`: explicit `_ENGINE_CONFIG_MAP` + `parse()` (initialize→convert→validate) | | 199 | +| `motor/engine_server/core/sglang/` | Transitional `SGLangEngine` and `SGLangEndpoint` | |
| 200 | +| `motor/node_manager/core/services/native_engine/config_factory.py` | `ConfigFactory`: explicit `_ENGINE_CONFIG_MAP` + `parse()` (initialize→convert→validate) | | ||
| 195 | | `motor/engine_server/factory/endpoint_factory.py` | InferEndpoint loading by engine name | | 201 | | `motor/engine_server/factory/endpoint_factory.py` | InferEndpoint loading by engine name | |
| 196 | | `motor/config/endpoint.py` | `EndpointConfig`: engine_type, host, port, mgmt_port, role, dp_rank, master_dp_ip, node_rank, d2d_peer_ips, snapshot_metadata, deploy_config; `HealthCheckConfig`: timeout, retry attempts, npu_usage_threshold, enable_virtual_inference, max_failure_count | | 202 | | `motor/config/endpoint.py` | `EndpointConfig`: engine_type, host, port, mgmt_port, role, dp_rank, master_dp_ip, node_rank, d2d_peer_ips, snapshot_metadata, deploy_config; `HealthCheckConfig`: timeout, retry attempts, npu_usage_threshold, enable_virtual_inference, max_failure_count | |
| 197 | 203 | ||
| @@ -201,11 +207,12 @@ Layer 4: SimInference (proactive health) | |||
| 201 | 207 | ||
| 202 | To add support for a new engine (e.g., "trtllm"): | 208 | To add support for a new engine (e.g., "trtllm"): |
| 203 | 209 | ||
| 204 | -1. Create `motor/engine_server/core/{engine}/` directory | 210 | +1. Create `motor/node_manager/core/services/native_engine/backends/{engine}/config.py` |
| 205 | -2. Implement `{engine}_config.py`: subclass `IConfig`, implement `initialize()`/`validate()`/`convert()` | 211 | +2. Implement the engine config adapter as an `IConfig` subclass and register it in `ConfigFactory._ENGINE_CONFIG_MAP` |
| 206 | -3. Implement `{engine}_engine.py`: subclass `Engine`, wrap engine client creation in `launch()` | 212 | +3. For transitional EngineServer support, create `motor/engine_server/core/{engine}/` |
| 207 | -4. Implement `{engine}_endpoint.py`: subclass `InferEndpoint`, implement `get_lifespan()` + `init_request_handlers()` | 213 | +4. Implement `{engine}_engine.py`: subclass `Engine`, wrap engine client creation in `launch()` |
| 208 | -5. Register engine name → module path in `ConfigFactory._ENGINE_CONFIG_MAP` and `EndpointFactory` | 214 | +5. Implement `{engine}_endpoint.py`: subclass `InferEndpoint`, implement `get_lifespan()` + `init_request_handlers()` |
| 215 | +6. Register the transitional endpoint in `EndpointFactory` | ||
| 209 | 216 | ||
| 210 | ### Other Rules | 217 | ### Other Rules |
| 211 | 218 | ||
| @@ -214,7 +221,7 @@ To add support for a new engine (e.g., "trtllm"): | |||
| 214 | - **PD disaggregation**: KV transfer config is set up during `IConfig.initialize()` based on the endpoint's role (prefill=producer, decode=consumer) and the selected connector (Mooncake/Multi/UCM/AscendStore). | 221 | - **PD disaggregation**: KV transfer config is set up during `IConfig.initialize()` based on the endpoint's role (prefill=producer, decode=consumer) and the selected connector (Mooncake/Multi/UCM/AscendStore). |
| 215 | - **Device pinning**: handled by NodeManager via `ASCEND_RT_VISIBLE_DEVICES` env var. EngineServer must NOT manage device assignment — it inherits visibility from the parent process. | 222 | - **Device pinning**: handled by NodeManager via `ASCEND_RT_VISIBLE_DEVICES` env var. EngineServer must NOT manage device assignment — it inherits visibility from the parent process. |
| 216 | - **Prometheus multiprocess**: `PROMETHEUS_MULTIPROC_DIR` must be set in `main()` before any metric import. Each engine_server writes its own metrics file to this directory. | 223 | - **Prometheus multiprocess**: `PROMETHEUS_MULTIPROC_DIR` must be set in `main()` before any metric import. Each engine_server writes its own metrics file to this directory. |
| 217 | -- **Port ranges**: ports are validated in `NodeManager`'s `EngineService._check_params()` (business_port in [1024, 65535]) before the engine process is spawned. | 224 | +- **Port ranges**: `EndpointConfig.validate()` validates business and management ports before the engine process is spawned. |
| 218 | 225 | ||
| 219 | ## Testing | 226 | ## Testing |
| 220 | 227 | ||
| @@ -1,238 +1,273 @@ | |||
| 1 | # NodeManager Module — Architecture & Implementation | 1 | # NodeManager Module — Architecture & Implementation |
| 2 | 2 | ||
| 3 | -## Architecture: Sidecar Daemon | 3 | +## Architecture: Native Runtime Sidecar |
| 4 | 4 | ||
| 5 | -One NodeManager process per NPU pod/container. Its job is to manage the lifecycle of `engine_server` subprocesses on the local machine and report health to the Controller. | 5 | +One NodeManager process runs in each inference pod. It receives lifecycle commands from |
| 6 | +Controller, starts one native vLLM or SGLang process per endpoint, supervises the complete | ||
| 7 | +process group, probes the native `/health` endpoint, and reports endpoint status through | ||
| 8 | +Controller heartbeats. | ||
| 6 | 9 | ||
| 7 | ``` text | 10 | ``` text |
| 8 | -NodeManager Process (sidecar, one per pod) — NodeManager(Application) main loop | 11 | +NodeManager (Application) |
| 9 | │ | 12 | │ |
| 10 | ├── NodeManagerAPI (FastAPI thread) | 13 | ├── NodeManagerAPI (FastAPI thread) |
| 11 | -│ POST /node-manager/start — spawn engines with StartCmdMsg | 14 | +│ POST /node-manager/start — validate StartCmdMsg and launch native engines |
| 12 | -│ POST /node-manager/stop — kill engines | 15 | +│ POST /node-manager/stop — stop native process groups |
| 13 | -│ POST /node-manager/pause — pause endpoints (snapshot/upgrade flow) | 16 | +│ POST /node-manager/pause — mark endpoints PAUSED for PreStop |
| 14 | -│ POST /node-manager/resume — resume endpoints | 17 | +│ POST /node-manager/resume — restore PAUSED endpoints |
| 15 | -│ GET /node-manager/status — current engine states | 18 | +│ GET /node-manager/status — report endpoint readiness |
| 16 | -│ GET /readiness — k8s readiness probe | 19 | +│ GET /readiness — Kubernetes readiness |
| 17 | -│ (TLS via mgmt_tls_config) | ||
| 18 | │ | 20 | │ |
| 19 | -├── EngineManager (ThreadSafeSingleton) | 21 | +├── Daemon |
| 20 | -│ Registration protocol: POST /controller/register (with retry) | 22 | +│ service registry for engine and optional KV-store services |
| 21 | -│ Ranktable file writing: saves ranktable JSON for engine RPC | 23 | +│ 5-second service monitor → health_check() |
| 22 | │ | 24 | │ |
| 23 | -├── Daemon (ThreadSafeSingleton) | 25 | +├── NativeEngineService |
| 24 | -│ Service registry orchestration (core/services/registry.py): | 26 | +│ builds LaunchContext, selects Native Engine Backend, delegates lifecycle to ProcessSupervisor |
| 25 | -│ "engine" services → EngineService (subprocess.Popen, device pinning) | ||
| 26 | -│ "kv-store" backends → memcache/lifecycle services | ||
| 27 | -│ Device pinning via ASCEND_RT_VISIBLE_DEVICES env var | ||
| 28 | -│ SIGKILL on stop (no graceful shutdown — engines are stateless) | ||
| 29 | -│ 5s process monitor thread → svc.health_check() for every service | ||
| 30 | │ | 27 | │ |
| 31 | -├── HeartbeatManager (ThreadSafeSingleton) | 28 | +├── ProcessSupervisor |
| 32 | -│ Two daemon threads: | 29 | +│ subprocess.Popen(start_new_session=True) |
| 33 | -│ _engine_server_status_thread — poll each engine GET /status every interval | 30 | +│ owns RuntimeProcess records, process groups and native health probes |
| 34 | -│ _heartbeat_report_thread — POST /controller/heartbeat every interval | ||
| 35 | -│ Fault detection: 120s grace period + 5 consecutive abnormal reports → suicide | ||
| 36 | │ | 31 | │ |
| 37 | -└── FaultReporter (ThreadSafeSingleton) | 32 | +├── HeartbeatManager |
| 38 | - HTTP poll GET {business_port}/fault_tolerance/status per engine | 33 | +│ polls ProcessSupervisor every second |
| 39 | - (vLLM FT REST API) → report_software_fault to Controller | 34 | +│ reports Controller heartbeat at configured interval |
| 35 | +│ preserves STARTING/STOPPING/PAUSED semantics and suicide threshold | ||
| 36 | +│ | ||
| 37 | +└── FaultReporter | ||
| 38 | + optional GET {business_port}/fault_tolerance/status polling for engine software faults | ||
| 40 | ``` | 39 | ``` |
| 41 | 40 | ||
| 42 | -All core components use `ThreadSafeSingleton` — `__new__` + `threading.Lock` ensures one instance per process with thread-safe lazy initialization. | 41 | +`motor/node_manager/core/services/native_engine/` is the native engine service boundary. Its |
| 42 | +engine backends are stateless converters: they turn a validated `LaunchContext` into an immutable | ||
| 43 | +`LaunchSpec` containing a `CommandSpec` and a `ProbeSpec`. The Coordinator and Controller do not | ||
| 44 | +depend on these node-local runtime states. | ||
| 43 | 45 | ||
| 44 | ## Complete Lifecycle | 46 | ## Complete Lifecycle |
| 45 | 47 | ||
| 46 | -### Phase 1: Startup & Registration | 48 | +### Phase 1: Startup and Registration |
| 47 | 49 | ||
| 48 | ``` text | 50 | ``` text |
| 49 | - | 51 | +main.py |
| 50 | -1. main.py is a thin wrapper (41 lines): load config → port setup → NodeManager(config).run() | 52 | + → NodeManagerConfig.from_json() |
| 51 | - | 53 | + → port allocation / configuration validation |
| 52 | - The real orchestration lives in NodeManager(Application) (motor/node_manager/node_manager.py) | 54 | + → NodeManager(Application).run() |
| 53 | - and the Application base class (motor/common/app/application.py): signal handlers, | 55 | + → init_modules(): Daemon, NodeManagerAPI, EngineManager, HeartbeatManager |
| 54 | - daemon tick loop, module init/start, graceful shutdown. There is no separate | ||
| 55 | - suicide_procedure() function — suicide is driven by the tick loop (see Phase 4). | ||
| 56 | - | ||
| 57 | -2. NodeManager.init_modules() — registration order: | ||
| 58 | - Daemon → NodeManagerAPI → [EngineManager → HeartbeatManager] | ||
| 59 | - | ||
| 60 | - EngineManager/HeartbeatManager are registered ONLY when daemon.has_engine | ||
| 61 | - is True. A KV-only pod (kv_cache_store_config mode="separated") registers | ||
| 62 | - just Daemon + NodeManagerAPI. | ||
| 63 | - | ||
| 64 | -3. EngineManager._register() [background thread] | ||
| 65 | - | ||
| 66 | - Loop: | ||
| 67 | - wait_until_api_ready(timeout=30.0) # NodeManagerAPI must be serving | ||
| 68 | - POST /controller/register (with instance metadata, capabilities) | ||
| 69 | - → 200: registration accepted, break | ||
| 70 | - → non-200: retry indefinitely with exponential backoff (2, 4, 8, 16, 32s, capped at 32s) | ||
| 71 | - Registration failure does not send SIGTERM / suicide the process. | ||
| 72 | - | ||
| 73 | -4. Main loop (Application.run()) blocks until stop_event | ||
| 74 | - - SIGTERM/SIGINT → cleanup and exit | ||
| 75 | - - stdin EOF → cleanup and exit | ||
| 76 | - - each daemon tick also checks the HeartbeatManager suicide flag | ||
| 77 | - | ||
| 78 | ``` | 56 | ``` |
| 79 | 57 | ||
| 80 | -### Phase 2: Receiving Start Command | 58 | +`EngineManager` and `HeartbeatManager` are initialized only when the active service registry |
| 59 | +contains an engine. A KV-only pod (`kv_cache_store_config.mode="separated"`) starts the KV service | ||
| 60 | +and NodeManager API without native engine registration or endpoint heartbeats. | ||
| 61 | + | ||
| 62 | +`EngineManager` registers the pod with Controller in a background loop. It waits for the API-ready | ||
| 63 | +event, posts `/controller/register`, and retries indefinitely with exponential backoff starting at | ||
| 64 | +2 seconds and capped at 32 seconds. Registration failure does not terminate NodeManager; normal | ||
| 65 | +shutdown and heartbeat-triggered recovery remain controlled by the application lifecycle. | ||
| 66 | + | ||
| 67 | +### Phase 2: Start Command and Native Launch | ||
| 68 | + | ||
| 69 | +Controller sends `POST /node-manager/start` with `StartCmdMsg`: | ||
| 81 | 70 | ||
| 82 | ``` text | 71 | ``` text |
| 83 | -Controller → POST /node-manager/start (StartCmdMsg) | 72 | +{ |
| 84 | - { | 73 | + instance_id, job_name, role, |
| 85 | - instance_id: int, | 74 | + endpoints: [{id, ip, business_port, mgmt_port, dp_rank, headless, ...}], |
| 86 | - endpoints: [Endpoint], # host, business_port, mgmt_port, dp_rank, role | 75 | + master_dp_ip, node_rank, d2d_peer_ips, ranktable |
| 87 | - ranktable: {...}, # for distributed RPC | 76 | +} |
| 88 | - job_name: str, | ||
| 89 | - role: str, # prefill / decode / union | ||
| 90 | - master_dp_ip: str, | ||
| 91 | - d2d_peer_ips: [str], | ||
| 92 | - node_rank: int, | ||
| 93 | - } | ||
| 94 | - | ||
| 95 | -1. EngineManager.parse_start_cmd(msg): | ||
| 96 | - - Validates all fields (endpoint IPs, ports in range) | ||
| 97 | - - Stores instance_id, endpoints, role, master_dp_ip, d2d_peer_ips, node_rank | ||
| 98 | - - Writes ranktable JSON to file (for engine RPC initialization) | ||
| 99 | - | ||
| 100 | -2. Daemon.pull_engine(pd_role_info, endpoints, instance_id, ...): | ||
| 101 | - - Phase 1: run PreparableService.prepare() on KV-store services | ||
| 102 | - - Phase 2: EngineService.pull() — for each endpoint: | ||
| 103 | - | ||
| 104 | - - Device pinning (_calc_visible_device_ids, core/services/engine.py): | ||
| 105 | - start_device = i * local_world_size % device_num | ||
| 106 | - visible = local_world_size consecutive devices from start_device | ||
| 107 | - (wraps around when past the end) | ||
| 108 | - ASCEND_RT_VISIBLE_DEVICES is set ONLY when enable_multi_endpoints | ||
| 109 | - is enabled; the logic moved from Daemon into EngineService. | ||
| 110 | - | ||
| 111 | - - subprocess.Popen (shell=False, with env): | ||
| 112 | - engine_server --dp-rank <id> --instance-id <id> --role <role> | ||
| 113 | - --host <ip> --port <business_port> --mgmt-port <mgmt_port> | ||
| 114 | - --master-dp-ip <ip> --node-rank <rank> --config-path <path> | ||
| 115 | - optional flags appended when configured: | ||
| 116 | - --snapshot-metadata / --kv-port / --dp-rpc-port / | ||
| 117 | - --lookup-rpc-port / --d2d-peer-ips | ||
| 118 | - - Track PID in self.engine_pids list | ||
| 119 | - | ||
| 120 | -3. HeartbeatManager.start() — takes NO arguments: | ||
| 121 | - - Starts _engine_server_status_thread | ||
| 122 | - - Starts _heartbeat_report_thread | ||
| 123 | - - instance/role/endpoint fields are set afterwards via | ||
| 124 | - update_endpoint(StartCmdMsg) | ||
| 125 | - | ||
| 126 | ``` | 77 | ``` |
| 127 | 78 | ||
| 128 | -### Phase 3: Steady State | 79 | +The API parses and validates the message, then `Daemon.pull_engine()` runs preparable KV-store |
| 80 | +services before asking `NativeEngineService.pull()` to launch each endpoint. | ||
| 81 | + | ||
| 82 | +For each endpoint, `NativeEngineService`: | ||
| 83 | + | ||
| 84 | +1. Builds an immutable `LaunchContext` with role, rank, host, ports, distributed setup, D2D peers, | ||
| 85 | + environment, and the endpoint's `headless` flag. | ||
| 86 | +2. Selects `VllmBackend` or `SGLangBackend` from the configured engine type. | ||
| 87 | +3. Rebuilds and validates the role-specific `EndpointConfig` and native engine configuration. | ||
| 88 | +4. Creates a `CommandSpec` and a `ProbeSpec`. | ||
| 89 | +5. Calls `ProcessSupervisor.start()`. | ||
| 90 | + | ||
| 91 | +The native commands are: | ||
| 129 | 92 | ||
| 130 | ``` text | 93 | ``` text |
| 131 | -Every heartbeat_interval — _engine_server_status_thread: | 94 | +vllm serve <native vLLM arguments> |
| 132 | - For each endpoint: | 95 | +python3 -m sglang.launch_server <native SGLang arguments> |
| 133 | - GET http://{ip}:{mgmt_port}/status | ||
| 134 | - → parse EndpointStatus (initial / normal / abnormal / paused / wait2start) | ||
| 135 | - → update self._endpoints[i].status | ||
| 136 | - | ||
| 137 | -Every heartbeat_interval (default 3s) — _heartbeat_report_thread: | ||
| 138 | - POST /controller/heartbeat | ||
| 139 | - body: HeartbeatMsg { | ||
| 140 | - job_name, ins_id, ip, | ||
| 141 | - status: {endpoint_id: endpoint_status} # dict keyed by endpoint_id | ||
| 142 | - } | ||
| 143 | ``` | 96 | ``` |
| 144 | 97 | ||
| 145 | -### Phase 4: Fault Detection & Suicide | 98 | +No `engine_server` process is inserted by the Native Runtime path. |
| 99 | + | ||
| 100 | +Device pinning is applied by `NativeEngineService` when `enable_multi_endpoints` is enabled: | ||
| 101 | +`ASCEND_RT_VISIBLE_DEVICES` contains the endpoint's calculated local device range. `POD_IP` is | ||
| 102 | +used as the default `VLLM_HOST_IP` when it is not already set. When the Mooncake IPv6 experiment is | ||
| 103 | +enabled, `MC_USE_IPV6=1` is supplied unless the environment already defines it. | ||
| 104 | + | ||
| 105 | +### Native Engine Backend Contract | ||
| 106 | + | ||
| 107 | +| Type | Responsibility | | ||
| 108 | +|------|----------------| | ||
| 109 | +| `LaunchContext` | Immutable normalized input for one endpoint launch | | ||
| 110 | +| `CommandSpec` | Immutable argv, environment and optional working directory | | ||
| 111 | +| `ProbeSpec` | Native health path, request timeout, startup timeout, retry limit, TLS and headless mode | | ||
| 112 | +| `LaunchSpec` | Pair of `CommandSpec` and `ProbeSpec` | | ||
| 113 | +| `VllmBackend` | Builds `vllm serve`; Native P/D only accepts the supported HANDOFF profile | | ||
| 114 | +| `SGLangBackend` | Builds `python3 -m sglang.launch_server`; encode role is rejected | | ||
| 115 | + | ||
| 116 | +The backend configuration adapters flatten engine-specific JSON into native CLI arguments. They | ||
| 117 | +must call the engine configuration converter and validator before returning a launch specification. | ||
| 118 | +`None` values are omitted from CLI arguments. SGLang receives `enable-metrics=true` because | ||
| 119 | +Coordinator metrics and PreStop drain depend on the native metrics endpoint. | ||
| 120 | + | ||
| 121 | +For SGLang P/D roles, the backend sets `disaggregation-mode=prefill|decode`; union uses | ||
| 122 | +`disaggregation-mode=null`. Multi-node configurations validate `nnodes` and require | ||
| 123 | +`master_dp_ip`, which is formatted with the shared IPv4/IPv6 address helper. | ||
| 124 | + | ||
| 125 | +### Probe and Runtime State | ||
| 126 | + | ||
| 127 | +`ProbeSpec.max_attempts` is the configured | ||
| 128 | +`health_check_config.health_collector_timeout_retry_attempts` value. It counts the first request | ||
| 129 | +and retries only `requests` timeout failures (including exceptions wrapped by `SafeHTTPSClient`). | ||
| 130 | +HTTP status failures, TLS errors, connection errors and other exceptions are not retried. | ||
| 131 | + | ||
| 132 | +`ProcessSupervisor.state()` follows this state model: | ||
| 146 | 133 | ||
| 147 | ``` text | 134 | ``` text |
| 148 | -Grace period: 120s hardcoded from engine start | 135 | +start → STARTING |
| 149 | - (engines need time to load models — don't kill them during warmup) | 136 | + ├─ headless process alive → RUNNING |
| 150 | - | 137 | + ├─ /health success → READY |
| 151 | -After grace period: | 138 | + ├─ process exits → STOPPED |
| 152 | - If any endpoint.status == ABNORMAL: | 139 | + └─ startup timeout expires while probe fails → UNHEALTHY |
| 153 | - _consecutive_abnormal_count += 1 | ||
| 154 | - Else: | ||
| 155 | - _consecutive_abnormal_count = 0 | ||
| 156 | - | ||
| 157 | - If _consecutive_abnormal_count >= 5 (5 consecutive reports): | ||
| 158 | - _should_suicide = True | ||
| 159 | - → NodeManager._on_daemon_tick() detects the flag → stop_event.set() | ||
| 160 | - → main loop exits → shutdown() gracefully stops all modules | ||
| 161 | - (Daemon.stop SIGKILLs engine subprocesses — no os._exit(-1)) | ||
| 162 | - → run() returns exit code -1 | ||
| 163 | - → Kubernetes restarts the pod → fresh registration | ||
| 164 | - | ||
| 165 | -HTTP 503 from Controller: | ||
| 166 | - → Controller restarted → HeartbeatManager._reregister() | ||
| 167 | - → POST /controller/reregister (ReregisterMsg) — single attempt, no backoff; | ||
| 168 | - a later heartbeat exception triggers the next retry | ||
| 169 | ``` | 140 | ``` |
| 170 | 141 | ||
| 142 | +During `startup_timeout` a failed probe keeps `STARTING`; this prevents slow model loading from | ||
| 143 | +being reported as a fault. A headless process is only proven alive, not independently ready. The | ||
| 144 | +heartbeat layer maps `RUNNING` to `WAIT2START` and `READY` to `NORMAL`. | ||
| 145 | + | ||
| 146 | +`ProcessSupervisor` protects state updates with a lock and performs HTTP probes outside the lock. | ||
| 147 | +Each launch uses `start_new_session=True` on POSIX, caches the process-group ID, and treats the | ||
| 148 | +process group as the lifecycle unit. A stopped or dead record is removed exactly once; dead | ||
| 149 | +launchers trigger cleanup of the cached group so workers are not leaked. | ||
| 150 | + | ||
| 151 | +### Phase 3: Heartbeat and Status Reporting | ||
| 152 | + | ||
| 153 | +`HeartbeatManager` snapshots endpoint records under `_endpoint_lock`, probes each native endpoint | ||
| 154 | +through `Daemon.get_engine_runtime_state()`, then commits results only if the endpoint generation | ||
| 155 | +has not changed. HTTP probing is therefore not performed while holding the endpoint lock. | ||
| 156 | + | ||
| 157 | +Status mapping: | ||
| 158 | + | ||
| 159 | +| Native runtime state | Controller endpoint status | | ||
| 160 | +|----------------------|----------------------------| | ||
| 161 | +| `STARTING` / `STOPPING` | Preserve `INITIAL` or the last status | | ||
| 162 | +| `RUNNING` | `WAIT2START` | | ||
| 163 | +| `READY` | `NORMAL` | | ||
| 164 | +| `UNHEALTHY` / `STOPPED` | `ABNORMAL` | | ||
| 165 | +| manual `PAUSED` | Preserve `PAUSED` | | ||
| 166 | + | ||
| 167 | +The heartbeat body is `HeartbeatMsg(job_name, ins_id, ip, status)` where `status` is keyed by | ||
| 168 | +endpoint ID. Controller readiness requires routable endpoints to be `NORMAL`; headless members | ||
| 169 | +must at least be alive and reported as `WAIT2START`. | ||
| 170 | + | ||
| 171 | +`POST /node-manager/pause` marks all endpoints `PAUSED` and returns native metrics URLs for | ||
| 172 | +non-headless endpoints. `resume` changes only `PAUSED` records back to `NORMAL`. | ||
| 173 | + | ||
| 174 | +### Phase 4: Fault Detection and Recovery | ||
| 175 | + | ||
| 176 | +`Daemon` calls each service's `health_check()` every 5 seconds. If `NativeEngineService` observes a dead | ||
| 177 | +native launcher, it removes the record, cleans the process group, and requests one Pod-level | ||
| 178 | +recovery through `SIGTERM` when `motor_restart_engine` is enabled. `_recovery_requested` prevents | ||
| 179 | +duplicate recovery signals until a successful new pull resets it. | ||
| 180 | + | ||
| 181 | +`HeartbeatManager` counts consecutive successful heartbeat reports containing `ABNORMAL` endpoints. | ||
| 182 | +After five consecutive abnormal reports it sets the suicide flag. The main application tick sees | ||
| 183 | +the flag, stops modules, and exits for platform rescheduling. | ||
| 184 | + | ||
| 185 | +Controller HTTP 503 triggers the existing re-registration path. Heartbeat HTTP requests use a | ||
| 186 | +shared long-lived client with bounded retry behavior in `ControllerApiClient`. | ||
| 187 | + | ||
| 171 | ### Phase 5: Shutdown | 188 | ### Phase 5: Shutdown |
| 172 | 189 | ||
| 173 | -``` text | 190 | +Shutdown stops modules in reverse initialization order. `Daemon.stop()` asks each service to stop; |
| 174 | -NodeManager.shutdown() — stop modules in reverse registration order: | 191 | +`NativeEngineService` marks records `STOPPING`, sends SIGTERM to each native process group, and waits up to |
| 175 | - HeartbeatManager.stop() → join threads | 192 | +the configured grace period for the launcher. It then checks the cached PGID: a group still present after |
| 176 | - EngineManager.stop() → stop registration thread | 193 | +the launcher exits is force-killed too, so surviving workers cannot retain NPU, port or memory resources. |
| 177 | - NodeManagerAPI.stop() → shutdown FastAPI | 194 | +Records are removed after cleanup so concurrent status reads cannot report a stopped process as a fresh endpoint. |
| 178 | - Daemon.stop() → stop 5s monitor thread, then stop each service | ||
| 179 | - in reverse registration order (SIGKILL engine PIDs) | ||
| 180 | - (there is no NodeManagerConfig.stop() — the config object has no lifecycle) | ||
| 181 | -``` | ||
| 182 | 195 | ||
| 183 | -### Snapshot Restore | 196 | +### Snapshot Boundary |
| 184 | 197 | ||
| 185 | -When `is_restored_from_host_side_snapshot()` returns True: | 198 | +The Native Runtime path currently does not support container snapshot suspend/resume or native |
| 199 | +engine restore. `snapshot_config.enable_snapshot=true` is rejected during NodeManager config | ||
| 200 | +validation, and `snapshot_metadata_path` is not consumed by Native Runtime. | ||
| 186 | 201 | ||
| 187 | -- `_checkpoint_done_inspect_retry_count` tracks checkpoint readiness polling | 202 | +Legacy snapshot helpers and EngineServer compatibility code remain in the repository for the |
| 188 | -- `_register_after_restore()` — single POST /controller/register, no backoff; | 203 | +transition period, but they are not a supported Native Runtime launch path. Do not add new native |
| 189 | - on failure the next heartbeat-report exception triggers it again | 204 | +runtime behavior that depends on `engine_server` management endpoints or snapshot metadata until a |
| 190 | -- `_start_after_restore()` does NOT exist — after restore the Controller | 205 | +separate runtime contract is defined. |
| 191 | - re-issues a fresh StartCmdMsg via /node-manager/start, and the engine | ||
| 192 | - resumes through the snapshot routes (see engine-server.md) | ||
| 193 | 206 | ||
| 194 | ## Key Files | 207 | ## Key Files |
| 195 | 208 | ||
| 196 | | File | Role | | 209 | | File | Role | |
| 197 | |------|------| | 210 | |------|------| |
| 198 | -| `motor/node_manager/main.py` | Thin wrapper: config → port setup → `NodeManager(config).run()` | | 211 | +| `motor/node_manager/main.py` | Thin process entrypoint: config, port allocation and `NodeManager.run()` | |
| 199 | -| `motor/node_manager/node_manager.py` | `NodeManager(Application)`: init_modules, daemon tick, suicide detection, shutdown | | 212 | +| `motor/node_manager/node_manager.py` | Application lifecycle, module initialization, suicide handling and shutdown | |
| 200 | -| `motor/node_manager/api_server/node_manager_api.py` | FastAPI: `/node-manager/start`, `/stop`, `/pause`, `/resume`, `/status`, `/readiness` (TLS via `mgmt_tls_config`) | | 213 | +| `motor/node_manager/api_server/node_manager_api.py` | FastAPI start/stop/pause/resume/status/readiness APIs | |
| 201 | -| `motor/node_manager/core/engine_manager.py` | Registration thread with exponential backoff + ranktable file writing + StartCmdMsg handling | | 214 | +| `motor/node_manager/core/daemon.py` | Service discovery, preparable-service ordering and process monitor | |
| 202 | -| `motor/node_manager/core/daemon.py` | Service registry orchestration (engine + KV-store services), 5s process monitor thread | | 215 | +| `motor/node_manager/core/services/native_engine/service.py` | LaunchContext creation, device pinning, native launch and recovery request | |
| 203 | -| `motor/node_manager/core/services/engine.py` | `EngineService`: subprocess.Popen, device pinning (`_calc_visible_device_ids`), CLI args, `_check_params` (business_port range), PID-death → SIGTERM self | | 216 | +| `motor/node_manager/core/services/native_engine/models.py` | LaunchContext, CommandSpec, ProbeSpec, LaunchSpec and RuntimeState | |
| 204 | -| `motor/node_manager/core/services/registry.py` | `@register_service` decorator + `_MODULE_MAP`; discovers active services by pod profile | | 217 | +| `motor/node_manager/core/services/native_engine/factory.py` | Selects the stateless vLLM/SGLang backend by engine type | |
| 205 | -| `motor/node_manager/core/services/memcache/` | KV-store (memcache) service implementation | | 218 | +| `motor/node_manager/core/services/native_engine/config_factory.py` | Lazily loads engine-specific CLI configuration adapters | |
| 206 | -| `motor/node_manager/core/heartbeat_manager.py` | Two daemon threads: status polling + heartbeat reporting, fault detection state machine, reregister | | 219 | +| `motor/node_manager/core/services/native_engine/backends/` | vLLM/SGLang command construction, configuration conversion and validation | |
| 207 | -| `motor/node_manager/core/fault_reporter.py` | HTTP poll of engine `GET /fault_tolerance/status` (business port) → `report_software_fault` to Controller | | 220 | +| `motor/node_manager/core/services/native_engine/supervisor.py` | Process groups, bounded native health probes and runtime state ownership | |
| 208 | -| `motor/node_manager/api_client/controller_api_client.py` | Sync HTTP client to Controller: `/register`, `/reregister`, `/heartbeat` | | 221 | +| `motor/node_manager/core/services/registry.py` | Service registration and backend discovery | |
| 209 | -| `motor/node_manager/api_client/engine_server_api_client.py` | Sync HTTP client: `GET /status` on engine's mgmt port | | 222 | +| `motor/node_manager/core/services/memcache/` | Optional KV-store service implementation | |
| 210 | -| `motor/config/node_manager.py` | `NodeManagerConfig`: BasicConfig, APIConfig, EndpointConfig, SnapshotConfig, SingleContainerConfig, PortAllocatorConfig | | 223 | +| `motor/node_manager/core/heartbeat_manager.py` | Native state polling, status mapping, heartbeat and suicide threshold | |
| 224 | +| `motor/node_manager/core/engine_manager.py` | Controller registration, StartCmdMsg validation, ranktable and legacy transition hooks | | ||
| 225 | +| `motor/node_manager/api_client/controller_api_client.py` | Controller register, reregister and heartbeat HTTP client | | ||
| 226 | +| `motor/node_manager/core/fault_reporter.py` | Optional native engine software-fault polling | | ||
| 227 | +| `motor/config/node_manager.py` | NodeManager schema, endpoint derivation, ports and snapshot validation | | ||
| 211 | 228 | ||
| 212 | -## Port Allocation | 229 | +## Port and Address Rules |
| 213 | 230 | ||
| 214 | -- `service_ports[i] = base_port + i * 2` (even ports — inference API) | 231 | +- Inference service ports and management ports are allocated from the NodeManager port allocator; |
| 215 | -- `mgmt_ports[i] = base_port + i * 2 + 1` (odd ports — management API) | 232 | + native runtime health probes use the endpoint's `business_port`. |
| 216 | -- Ports validated in range [1024, 65535] | 233 | +- Ports are validated before launch. The shared address helpers bracket IPv6 literals when building |
| 217 | -- Endpoint count: `min(dp_size, device_num // local_world_size)` | 234 | + URLs and distributed addresses. |
| 218 | -- When `port_allocator_config.enable`, runtime probing (motor/common/utils/port_allocator.py, | 235 | +- `master_dp_ip`, D2D peer addresses and endpoint hosts are passed through the common address |
| 219 | - `apply_node_manager_ports`) probes the host (`allocate_auto` with scan_range + | 236 | + formatting utilities; do not concatenate IPv6 host and port strings manually. |
| 220 | - probe_timeout) and re-allocates any busy port: node_manager_port, each | 237 | +- In multi-endpoint mode, device visibility is calculated from `local_world_size`; single-container |
| 221 | - service/mgmt port, and in single-container mode kv_port, lookup_rpc_port, | 238 | + offsets are applied after local device selection. |
| 222 | - dp_rpc_port. | ||
| 223 | 239 | ||
| 224 | ## Development Rules | 240 | ## Development Rules |
| 225 | 241 | ||
| 226 | -- **New core components** → use `ThreadSafeSingleton` pattern (`__new__` + double-checked locking with `threading.Lock`). | 242 | +- Keep native process lifecycle behind `Daemon` → `NativeEngineService` → `ProcessSupervisor`; do not call |
| 227 | -- **Module lifecycle**: init-then-start pattern — construct in `init_modules()`, start via API command, stop in `shutdown()`. | 243 | + `subprocess.Popen` from API handlers or native engine backends. |
| 228 | -- **Fault thresholds**: keep configurable constants: grace period (120s, hardcoded in heartbeat_manager), consecutive abnormal count (5), heartbeat interval (default 3s) — do not hardcode deeper in the code. | 244 | +- Native engine backends remain stateless and return immutable launch specifications. |
| 229 | -- **Engine process management** → always through Daemon → EngineService; never `subprocess.Popen` elsewhere. Daemon tracks all PIDs for cleanup. | 245 | +- Add a regression test for every new state transition, CLI mapping, probe policy or cleanup path. |
| 230 | -- **Device pinning**: `ASCEND_RT_VISIBLE_DEVICES` env var set per engine subprocess by `EngineService._calc_visible_device_ids` (only when `enable_multi_endpoints`). | 246 | +- Keep health probing outside endpoint locks; commit results under lock with generation checks. |
| 231 | -- **New engine/KV-store backends** → add a service module with `@register_service` and an entry in the registry's `_MODULE_MAP`; Daemon stays unchanged. | 247 | +- Health retries must remain bounded and timeout-only. Do not retry explicit HTTP rejection, TLS or |
| 232 | -- **Snapshot support**: check `is_restored_from_host_side_snapshot()` before registration; use longer timeouts during restore. | 248 | + configuration errors. |
| 249 | +- Preserve the distinction between process liveness (`RUNNING`/`WAIT2START`) and service readiness | ||
| 250 | + (`READY`/`NORMAL`). | ||
| 251 | +- New engine backends should add a backend/config package and tests without coupling `Daemon` to the | ||
| 252 | + engine-specific CLI. | ||
| 253 | +- Native Runtime snapshot support is intentionally disabled until its lifecycle contract is defined. | ||
| 233 | 254 | ||
| 234 | ## Testing | 255 | ## Testing |
| 235 | 256 | ||
| 236 | ```bash | 257 | ```bash |
| 258 | +# Focused backend and supervisor tests | ||
| 259 | +bash tests/run_tests.sh --serial tests/node_manager/core/services/native_engine/ | ||
| 260 | + | ||
| 261 | +# NodeManager module tests | ||
| 237 | bash tests/run_tests.sh tests/node_manager/ | 262 | bash tests/run_tests.sh tests/node_manager/ |
| 238 | ``` | 263 | ``` |
| 264 | + | ||
| 265 | +Important test areas: | ||
| 266 | + | ||
| 267 | +- `tests/node_manager/core/services/native_engine/test_supervisor.py`: process-group ownership, state races, bounded | ||
| 268 | + timeout retries, headless liveness and cleanup. | ||
| 269 | +- `tests/node_manager/core/services/native_engine/test_backends.py`: vLLM/SGLang command construction and ProbeSpec | ||
| 270 | + mapping. | ||
| 271 | +- `tests/node_manager/core/test_heartbeat_manager.py`: generation-safe status mapping and heartbeat | ||
| 272 | + behavior. | ||
| 273 | +- `tests/node_manager/core/services/native_engine/test_service.py`: recovery latch and process-death handling. | ||
| @@ -4,13 +4,17 @@ | |||
| 4 | 4 | ||
| 5 | 在本仓库中，**Engine Server** 指可执行入口 **`engine_server`**（`setup.py` 中 entry point：`engine_server = motor.engine_server.cli.main:main`），实现该功能的脚本路径为 `motor/engine_server/cli/main.py`。 | 5 | 在本仓库中，**Engine Server** 指可执行入口 **`engine_server`**（`setup.py` 中 entry point：`engine_server = motor.engine_server.cli.main:main`），实现该功能的脚本路径为 `motor/engine_server/cli/main.py`。 |
| 6 | 6 | ||
| 7 | +> `EngineServer` 是过渡期兼容链路。Node Manager 的直接原生路径由 | ||
| 8 | +> `motor/node_manager/core/services/native_engine/` 管理，并直接拉起 `vllm serve` 或 | ||
| 9 | +> `python3 -m sglang.launch_server`，不会插入 `engine_server` 进程。 | ||
| 10 | + | ||
| 7 | 主要功能如下： | 11 | 主要功能如下： |
| 8 | 12 | ||
| 9 | -- **解析端点配置**：`EndpointConfig.init_endpoint_config()`，经 `ConfigFactory` 得到具体引擎配置（`motor/engine_server/factory/config_factory.py` 中按 `vllm` / `sglang` 等类型选择配置类）。 | 13 | +- **解析端点配置**：`EndpointConfig.init_endpoint_config()`，经 `ConfigFactory` 得到具体引擎配置（`motor/node_manager/core/services/native_engine/config_factory.py` 中按 `vllm` / `sglang` 等类型选择配置类）。 |
| 10 | - **管理面 HTTP（MgmtEndpoint）**：`motor/engine_server/core/mgmt_endpoint.py` 内在 `mgmt_port` 上启动 uvicorn，挂载 Prometheus 相关路由及 **`GET /status`**（路径常量 `STATUS_INTERFACE`，值为 `/status`）。状态字段键为 `STATUS_KEY` 对应常量（与同文件 `NORMAL_STATUS` / `ABNORMAL_STATUS` / `INIT_STATUS` 等配合使用）。可选 **TLS**（`mgmt_tls_config.enable_tls`）。 | 14 | - **管理面 HTTP（MgmtEndpoint）**：`motor/engine_server/core/mgmt_endpoint.py` 内在 `mgmt_port` 上启动 uvicorn，挂载 Prometheus 相关路由及 **`GET /status`**（路径常量 `STATUS_INTERFACE`，值为 `/status`）。状态字段键为 `STATUS_KEY` 对应常量（与同文件 `NORMAL_STATUS` / `ABNORMAL_STATUS` / `INIT_STATUS` 等配合使用）。可选 **TLS**（`mgmt_tls_config.enable_tls`）。 |
| 11 | - **推理面（InferEndpoint）**：由 `EndpointFactory.get_infer_endpoint(config)` 按引擎类型构造（如 `VLLMEndpoint`、`SGLangEndpoint`，见 `motor/engine_server/factory/endpoint_factory.py`），与 `MgmtEndpoint` 并行 `run()`，主线程在 `infer_endpoint.wait()` 阻塞直至退出，再 `shutdown` 两端点。 | 15 | - **推理面（InferEndpoint）**：由 `EndpointFactory.get_infer_endpoint(config)` 按引擎类型构造（如 `VLLMEndpoint`、`SGLangEndpoint`，见 `motor/engine_server/factory/endpoint_factory.py`），与 `MgmtEndpoint` 并行 `run()`，主线程在 `infer_endpoint.wait()` 阻塞直至退出，再 `shutdown` 两端点。 |
| 12 | 16 | ||
| 13 | -Node Manager 侧通过子进程命令 **`engine_server`** 拉起本进程，参数在 `motor/node_manager/core/daemon.py` 的 `pull_engine` 中，包括 `--dp-rank`、`--instance-id`、`--role`、`--host`、`--port`、`--mgmt-port`、`--master-dp-ip`、`--config-path`（值为 `Env.user_config_path`）；单容器模式下还会追加 `--kv-port`、`--dp-rpc-port` 等。 | 17 | +旧部署入口可通过子进程命令 **`engine_server`** 拉起本进程；新 Node Manager 原生路径不再组装该命令。以下 Engine Server 参数说明仅用于兼容链路和本地调试。 |
| 14 | 18 | ||
| 15 | ## 与周边组件的关系 | 19 | ## 与周边组件的关系 |
| 16 | 20 | ||
| @@ -2,12 +2,12 @@ | |||
| 2 | 2 | ||
| 3 | ## 功能介绍 | 3 | ## 功能介绍 |
| 4 | 4 | ||
| 5 | -Node Manager 是部署在推理节点上的管理进程，负责连接 Controller 与本节点的 Engine Server。进程入口为 `motor/node_manager/main.py`，核心逻辑在 `motor/node_manager/node_manager.py`（`NodeManager` 类），继承自 `motor/common/app/application.py`（`Application` 基类）。主要职责如下： | 5 | +Node Manager 是部署在推理节点上的管理进程，负责连接 Controller 与本节点的原生 vLLM/SGLang 引擎。进程入口为 `motor/node_manager/main.py`，核心逻辑在 `motor/node_manager/node_manager.py`（`NodeManager` 类），继承自 `motor/common/app/application.py`（`Application` 基类）。主要职责如下： |
 | |||
| 6 | 6 | ||
| 7 | 1. 加载节点配置，完成端口分配并启动管理面 HTTP 服务。 | 7 | 1. 加载节点配置，完成端口分配并启动管理面 HTTP 服务。 |
| 8 | 2. 向 Controller 注册节点，接收 Controller 下发的实例启动命令。 | 8 | 2. 向 Controller 注册节点，接收 Controller 下发的实例启动命令。 |
| 9 | -3. 按 endpoint 拉起、记录和停止 `engine_server` 子进程。 | 9 | +3. 按 endpoint 直接拉起、监管和停止原生引擎进程组。 |
| 10 | -4. 轮询 Engine Server 状态并向 Controller 上报心跳。 | 10 | +4. 轮询原生业务端口的健康接口并向 Controller 上报心跳。 |
| 11 | 5. 处理优雅暂停、配置热更新、容器快照恢复和软件故障上报。 | 11 | 5. 处理优雅暂停、配置热更新、容器快照恢复和软件故障上报。 |
| 12 | 12 | ||
| 13 | ### 组件结构 | 13 | ### 组件结构 |
| @@ -18,14 +18,15 @@ Node Manager 是部署在推理节点上的管理进程，负责连接 Controlle | |||
| 18 | | `NodeManager` | `motor/node_manager/node_manager.py` | `Application` 子类：组装模块并运行 daemon loop，每 tick 检查自杀标志 | | 18 | | `NodeManager` | `motor/node_manager/node_manager.py` | `Application` 子类：组装模块并运行 daemon loop，每 tick 检查自杀标志 | |
| 19 | | `NodeManagerConfig` | `motor/config/node_manager.py` | 加载、校验和重载节点配置，推导 endpoint 数量与端口 | | 19 | | `NodeManagerConfig` | `motor/config/node_manager.py` | 加载、校验和重载节点配置，推导 endpoint 数量与端口 | |
| 20 | | `NodeManagerAPI` | `motor/node_manager/api_server/node_manager_api.py` | 在后台线程中运行 FastAPI/uvicorn，提供启动、停止和探针接口 | | 20 | | `NodeManagerAPI` | `motor/node_manager/api_server/node_manager_api.py` | 在后台线程中运行 FastAPI/uvicorn，提供启动、停止和探针接口 | |
| 21 | -| `Daemon` | `motor/node_manager/core/daemon.py` | 服务编排器：根据配置发现并实例化 Engine 和 KV-store 服务，维护进程监控器 | | 21 | +| `Daemon` | `motor/node_manager/core/daemon.py` | 服务编排器：根据配置发现并实例化原生引擎和 KV-store 服务，维护进程监控器 | |
| 22 | -| `EngineService` | `motor/node_manager/core/services/engine.py` | Engine 子进程生命周期管理：组装命令、拉起/追踪/停止 `engine_server` 进程 | | 22 | +| `NativeEngineService` | `motor/node_manager/core/services/native_engine/service.py` | 构造不可变 LaunchContext，选择 Native Engine Backend 并管理原生引擎生命周期 | |
| 23 | | `LocalService` | `motor/node_manager/core/services/memcache/lifecycle.py` | memcache 后端生命周期管理：配置准备、子进程拉起（通过 `memcache/worker.py`）、健康检查与重启 | | 23 | | `LocalService` | `motor/node_manager/core/services/memcache/lifecycle.py` | memcache 后端生命周期管理：配置准备、子进程拉起（通过 `memcache/worker.py`）、健康检查与重启 | |
| 24 | +| `ProcessSupervisor` | `motor/node_manager/core/services/native_engine/supervisor.py` | 创建独立进程组，维护运行态、原生健康探测和完整进程树清理 | | ||
| 25 | +| `VllmBackend` / `SGLangBackend` | `motor/node_manager/core/services/native_engine/backends/` | 将统一启动上下文转换为引擎原生命令与 ProbeSpec | | ||
| 24 | | `EngineManager` | `motor/node_manager/core/engine_manager.py` | 注册/重注册、校验启动命令、处理 ranktable、快照元数据和故障上报 | | 26 | | `EngineManager` | `motor/node_manager/core/engine_manager.py` | 注册/重注册、校验启动命令、处理 ranktable、快照元数据和故障上报 | |
| 25 | | `HeartbeatManager` | `motor/node_manager/core/heartbeat_manager.py` | 轮询 endpoint 状态、上报心跳、维护暂停/恢复状态并触发异常自杀 | | 27 | | `HeartbeatManager` | `motor/node_manager/core/heartbeat_manager.py` | 轮询 endpoint 状态、上报心跳、维护暂停/恢复状态并触发异常自杀 | |
| 26 | | `FaultReporter` | `motor/node_manager/core/fault_reporter.py` | 轮询引擎 FT 状态接口并上报软件故障给 Controller | | 28 | | `FaultReporter` | `motor/node_manager/core/fault_reporter.py` | 轮询引擎 FT 状态接口并上报软件故障给 Controller | |
| 27 | | `ControllerApiClient` | `motor/node_manager/api_client/controller_api_client.py` | 调用 Controller 的注册、重注册、心跳和故障上报接口 | | 29 | | `ControllerApiClient` | `motor/node_manager/api_client/controller_api_client.py` | 调用 Controller 的注册、重注册、心跳和故障上报接口 | |
| 28 | -| `EngineServerApiClient` | `motor/node_manager/api_client/engine_server_api_client.py` | 调用 Engine Server 管理面的 `GET /status` | | ||
| 29 | 30 | ||
| 30 | `Daemon`、`EngineManager` 和 `HeartbeatManager` 均为线程安全单例。HTTP 路由和后台线程通过这些单例共享实例、endpoint 和进程状态。`Application` 和 `NodeManager` 不是单例，由 `main.py` 显式创建。 | 31 | `Daemon`、`EngineManager` 和 `HeartbeatManager` 均为线程安全单例。HTTP 路由和后台线程通过这些单例共享实例、endpoint 和进程状态。`Application` 和 `NodeManager` 不是单例，由 `main.py` 显式创建。 |
| 31 | 32 | ||
| @@ -54,18 +55,20 @@ Controller 调用 `POST /node-manager/start` 后，处理流程为： | |||
| 54 | 2. 校验 `job_name`、endpoint 数量以及每个 endpoint 的 IP 是否与本节点配置一致。 | 55 | 2. 校验 `job_name`、endpoint 数量以及每个 endpoint 的 IP 是否与本节点配置一致。 |
| 55 | 3. 保存 `instance_id`、endpoints、`node_rank` 和 D2D peer 信息；如配置了 `RANKTABLE_PATH`，将实例 ranktable 写入该文件。 | 56 | 3. 保存 `instance_id`、endpoints、`node_rank` 和 D2D peer 信息；如配置了 `RANKTABLE_PATH`，将实例 ranktable 写入该文件。 |
| 56 | 4. 准备快照运行目录和元数据。 | 57 | 4. 准备快照运行目录和元数据。 |
| 57 | -5. `Daemon.pull_engine()` 为每个 endpoint 拉起一个 `engine_server` 子进程。 | 58 | +5. `Daemon.pull_engine()` 为每个 endpoint 直接拉起一个原生 vLLM 或 SGLang 进程组。 |
| 58 | 6. 更新 `HeartbeatManager` 中的 endpoint，并启动状态轮询和心跳线程。 | 59 | 6. 更新 `HeartbeatManager` 中的 endpoint，并启动状态轮询和心跳线程。 |
| 59 | 7. 启动 `EngineManager` 中的 `FaultReporter`（仅在故障容忍功能开启时生效）。 | 60 | 7. 启动 `EngineManager` 中的 `FaultReporter`（仅在故障容忍功能开启时生效）。 |
| 60 | 61 | ||
| 61 | -从宿主机侧快照恢复时，第 5 步不会再次拉起 Engine Server，而是更新恢复元数据、endpoint 和恢复状态。 | 62 | +从宿主机侧快照恢复时，第 5 步不会再次拉起引擎，而是更新恢复元数据、endpoint 和恢复状态。 |
| 62 | 63 | ||
| 63 | ### 停止与重调度 | 64 | ### 停止与重调度 |
| 64 | 65 | ||
| 65 | -- 收到 `SIGINT`、`SIGTERM` 或标准输入命令 `stop` 时，`Application._handle_signal()` 设置 `stop_event`，daemon loop 退出后执行 `shutdown()`：按注册逆序调用每个模块的 `stop()`，然后停止配置 watcher。 | 66 | +- 收到 `SIGINT`、`SIGTERM` 或标准输入命令 `stop` 时，`Application._handle_signal()` 设置 `stop_event`；daemon loop 退出后停止配置 watcher，并按初始化的逆序停止模块。 |
| 66 | -- `Daemon.stop()` 遍历所有 service 调用 `stop()`：`EngineService.stop()` 对记录的 Engine Server PID 发送 `SIGKILL`；`LocalService.stop()` 对 memcache worker 子进程发送 `SIGKILL`。 | 67 | +- `Daemon.stop()` 通过 `ProcessSupervisor` 向所有原生引擎进程组发送 `SIGTERM`；宽限期后仍未退出时发送 `SIGKILL` 清理完整进程树。 |
| 67 | -- 任一 endpoint 连续 5 个心跳周期保持 `ABNORMAL` 时，`HeartbeatManager` 设置自杀标志。daemon loop 每 tick 检查该标志，触发后 `stop_event.set()` 并返回 `-1`，用于触发重调度。 | 68 | +- `Daemon.stop()` 同时遍历其他已启用 service；例如 `LocalService.stop()` 会停止 memcache worker 子进程。 |
| 68 | -- `exit_code` 默认返回 `-1`，与旧行为一致（-1 表示 rescheduling）。 | 69 | +- 任一原生引擎进程异常退出时只触发一次 Pod 级恢复，不在 Pod 内重启单个 rank。 |
| 70 | +- 任一 endpoint 连续 5 个心跳周期保持 `ABNORMAL` 时，`HeartbeatManager` 设置自杀标志。主线程执行清理后返回 `-1`，用于触发重调度。 | ||
| 71 | +- 当前 `main()` 正常退出路径同样返回 `-1`；源码注释约定 `-1` 表示 rescheduling、`0` 表示 restart。 | ||
| 69 | 72 | ||
| 70 | ## Node Manager HTTP API | 73 | ## Node Manager HTTP API |
| 71 | 74 | ||
| @@ -73,14 +76,14 @@ Node Manager API 默认监听 `api_config.pod_ip:api_config.node_manager_port` | |||
| 73 | 76 | ||
| 74 | | 方法 | 路径 | 响应 | 说明 | | 77 | | 方法 | 路径 | 响应 | 说明 | |
| 75 | |------|------|----------|------| | 78 | |------|------|----------|------| |
| 76 | -| `POST` | `/node-manager/start` | `200 {}` | 校验启动命令并拉起 Engine Server；快照恢复时执行恢复准备 | | 79 | +| `POST` | `/node-manager/start` | `200 {}` | 校验启动命令并拉起原生引擎；快照恢复时执行恢复准备 | |
| 77 | -| `POST` | `/node-manager/stop` | `200 {"message": "All engine processes stopped successfully."}` | 停止当前 Node Manager 记录的全部 Engine Server 进程 | | 80 | +| `POST` | `/node-manager/stop` | `200 {"message": "All engine processes stopped successfully."}` | 停止当前 Node Manager 监管的全部原生引擎进程组 | |
| 78 | -| `POST` | `/node-manager/pause` | `200 {"status":"ok", ...}` | 将全部 endpoint 标记为 `PAUSED`，并返回 Engine Server 管理地址 | | 81 | +| `POST` | `/node-manager/pause` | `200 {"status":"ok", ...}` | 将全部 endpoint 标记为 `PAUSED`，并返回非 headless 原生引擎 metrics URL | |
| 79 | | `POST` | `/node-manager/resume` | `200 {"status":"ok", ...}` | 仅将 `PAUSED` endpoint 恢复为 `NORMAL` | | 82 | | `POST` | `/node-manager/resume` | `200 {"status":"ok", ...}` | 仅将 `PAUSED` endpoint 恢复为 `NORMAL` | |
| 80 | | `GET` | `/node-manager/status` | `200 {"status": true/false}` | 返回全部 endpoint 是否为 `NORMAL`；无 endpoint 时为 `false` | | 83 | | `GET` | `/node-manager/status` | `200 {"status": true/false}` | 返回全部 endpoint 是否为 `NORMAL`；无 endpoint 时为 `false` | |
| 81 | | `GET` | `/readiness` | `200` 或 `503` | Kubernetes Readiness Probe 接口。实例节点 Pod 默认不配置该探针；仅在容器快照默认应用场景下配置，用于判断执行容器 checkpoint 前的稳态点。未到达稳态点时返回 `503`，到达后返回 `200` | | 84 | | `GET` | `/readiness` | `200` 或 `503` | Kubernetes Readiness Probe 接口。实例节点 Pod 默认不配置该探针；仅在容器快照默认应用场景下配置，用于判断执行容器 checkpoint 前的稳态点。未到达稳态点时返回 `503`，到达后返回 `200` | |
| 82 | 85 | ||
| 83 | -`/node-manager/pause` 用于 PreStop 优雅下线：暂停状态会使 readiness 失败，并通过心跳通知 Controller；状态轮询不会用 Engine Server 返回值覆盖手动设置的 `PAUSED`。如果 PreStop 被取消，可调用 `/node-manager/resume` 恢复调度。 | 86 | +`/node-manager/pause` 用于 PreStop 优雅下线：暂停状态会使 readiness 失败，并通过心跳通知 Controller；原生健康轮询不会覆盖手动设置的 `PAUSED`。响应中的 `engine_metrics_targets` 使用原生业务端口，并排除 headless 成员。如果 PreStop 被取消，可调用 `/node-manager/resume` 恢复调度。 |
| 84 | 87 | ||
| 85 | `/readiness` 仅用于快照默认应用场景，即 MindCluster 实例重调度，不作为 Node Manager 的通用健康检查接口。MindCluster 通过该接口查询实例节点是否到达稳态点。在容器快照的用户自定义应用场景中，可调用 `/node-manager/status` 查询稳态点；接口返回 `200 {"status": true}` 表示已到达稳态点。 | 88 | `/readiness` 仅用于快照默认应用场景，即 MindCluster 实例重调度，不作为 Node Manager 的通用健康检查接口。MindCluster 通过该接口查询实例节点是否到达稳态点。在容器快照的用户自定义应用场景中，可调用 `/node-manager/status` 查询稳态点；接口返回 `200 {"status": true}` 表示已到达稳态点。 |
| 86 | 89 | ||
| @@ -91,7 +94,7 @@ Node Manager API 默认监听 `api_config.pod_ip:api_config.node_manager_port` | |||
| 91 | | `job_name` | string | 是 | 实例任务名，必须与本节点配置一致 | | 94 | | `job_name` | string | 是 | 实例任务名，必须与本节点配置一致 | |
| 92 | | `role` | string | 是 | 实例角色，如 `prefill`、`decode` 或 `union` | | 95 | | `role` | string | 是 | 实例角色，如 `prefill`、`decode` 或 `union` | |
| 93 | | `instance_id` | int | 是 | Controller 分配的实例 ID | | 96 | | `instance_id` | int | 是 | Controller 分配的实例 ID | |
| 94 | -| `endpoints` | array | 是 | 本节点管理的 endpoint；元素包含 `id`、`ip`、`business_port`、`mgmt_port` 等 | | 97 | +| `endpoints` | array | 是 | 本节点管理的 endpoint；元素包含 `id`、`ip`、`business_port`、`mgmt_port`，SGLang PD endpoint 还可包含 `bootstrap_port` | |
| 95 | | `master_dp_ip` | string | 是 | 数据并行主节点 IP | | 98 | | `master_dp_ip` | string | 是 | 数据并行主节点 IP | |
| 96 | | `ranktable` | object/null | 否 | 实例级 ranktable，默认 `null` | | 99 | | `ranktable` | object/null | 否 | 实例级 ranktable，默认 `null` | |
| 97 | | `d2d_peer_ips` | array/null | 否 | D2D 权重传输对端，Controller 使用 `<endpoint_id>:<peer_ip>` 编码，默认 `null` | | 100 | | `d2d_peer_ips` | array/null | 否 | D2D 权重传输对端，Controller 使用 `<endpoint_id>:<peer_ip>` 编码，默认 `null` | |
| @@ -101,7 +104,7 @@ Node Manager API 默认监听 `api_config.pod_ip:api_config.node_manager_port` | |||
| 101 | 104 | ||
| 102 | - 启动命令内部解析异常：`400 Invalid start command payload`。 | 105 | - 启动命令内部解析异常：`400 Invalid start command payload`。 |
| 103 | - `job_name`、endpoint 数量或 endpoint IP 校验失败：`422 Start command validation failed`。 | 106 | - `job_name`、endpoint 数量或 endpoint IP 校验失败：`422 Start command validation failed`。 |
| 104 | -- Engine Server 拉起失败：`500 Failed to start engine server`。 | 107 | +- 原生引擎拉起失败：`500 Failed to start native engine`。 |
| 105 | - 请求 JSON/Pydantic 字段解析异常会被外层异常处理转换为通用 `500`。 | 108 | - 请求 JSON/Pydantic 字段解析异常会被外层异常处理转换为通用 `500`。 |
| 106 | - `/readiness` 在 endpoint 尚未健康或快照恢复后尚未启动时返回 `503`。 | 109 | - `/readiness` 在 endpoint 尚未健康或快照恢复后尚未启动时返回 `503`。 |
| 107 | 110 | ||
| @@ -114,44 +117,36 @@ Node Manager API 默认监听 `api_config.pod_ip:api_config.node_manager_port` | |||
| 114 | | Node Manager → Controller | `POST /controller/register` | 上报角色、模型、端口、并行配置、ranktable、`nnodes` 和快照主节点标记 | | 117 | | Node Manager → Controller | `POST /controller/register` | 上报角色、模型、端口、并行配置、ranktable、`nnodes` 和快照主节点标记 | |
| 115 | | Node Manager → Controller | `POST /controller/reregister` | Controller 重启并对心跳返回 `503` 时，携带实例与 endpoint 信息重新注册 | | 118 | | Node Manager → Controller | `POST /controller/reregister` | Controller 重启并对心跳返回 `503` 时，携带实例与 endpoint 信息重新注册 | |
| 116 | | Node Manager → Controller | `POST /controller/heartbeat` | 按 `heartbeat_interval_seconds` 上报各 endpoint 状态 | | 119 | | Node Manager → Controller | `POST /controller/heartbeat` | 按 `heartbeat_interval_seconds` 上报各 endpoint 状态 | |
| 117 | -| Node Manager → Controller | `POST /controller/report_software_fault` | 转发 Engine Server 软件故障 | | 120 | +| Node Manager → Controller | `POST /controller/report_software_fault` | 转发引擎已有的软件故障信号 | |
| 118 | 121 | ||
| 119 | 心跳使用长连接客户端，单次超时为 5 秒；TCP 请求失败时重建连接，并按 1 秒、2 秒退避重试两次。 | 122 | 心跳使用长连接客户端，单次超时为 5 秒；TCP 请求失败时重建连接，并按 1 秒、2 秒退避重试两次。 |
| 120 | 123 | ||
| 121 | -### Engine Server | 124 | +### 原生引擎运行态 |
| 122 | 125 | ||
| 123 | -`HeartbeatManager` 启动时先等待各 endpoint 管理端口可连接，最长等待 60 秒；之后每秒调用一次 Engine Server 的 `GET /status`，单次请求超时为 5 秒。 | 126 | +`HeartbeatManager` 每秒读取 `ProcessSupervisor` 的运行态。非 headless endpoint 使用原生 `business_port/health`，并沿用 `infer_tls_config`；headless 成员不强造 HTTP frontend，仅检查进程存活并上报 `WAIT2START`。Controller 仅在所有可路由 endpoint 为 `NORMAL`、所有 headless 成员至少已上报 `WAIT2START` 时将实例置为可用。 |
| 124 | 127 | ||
| 125 | 状态轮询具有以下保护逻辑： | 128 | 状态轮询具有以下保护逻辑： |
| 126 | 129 | ||
| 127 | -- Engine Server 启动后的 120 秒宽限期内，探测到 `ABNORMAL` 时保留原状态。 | 130 | +- 进程拉起后进入 `STARTING`；在 `health_check_config.startup_timeout`（默认 1800 秒）内，业务端口尚未监听只表示模型仍在加载。 |
| 131 | +- headless 进程存活时进入 `RUNNING` 并上报 `WAIT2START`，不以进程存活冒充独立服务就绪；进程退出仍上报 `ABNORMAL`。 | ||
| 132 | +- 原生 `/health` 成功后进入 `READY` 并向 Controller 上报 `NORMAL`；单次请求超时仅按 `health_collector_timeout_retry_attempts` 限定次数重试，进程退出或启动窗口外重试耗尽后仍不可用则进入 `UNHEALTHY`。 | ||
| 128 | - 更新 endpoint 时使用 generation 标记，避免旧探测结果覆盖 Controller 新下发的数据。 | 133 | - 更新 endpoint 时使用 generation 标记，避免旧探测结果覆盖 Controller 新下发的数据。 |
| 129 | - 手动设置的 `PAUSED` 状态不会被轮询结果覆盖。 | 134 | - 手动设置的 `PAUSED` 状态不会被轮询结果覆盖。 |
| 130 | -- Engine Server 返回未知状态或无效响应时，按 `ABNORMAL` 处理。 | ||
| 131 | 135 | ||
| 132 | -## Engine Server 拉起参数 | 136 | +## 原生引擎启动 |
| 133 | 137 | ||
| 134 | -`Daemon.pull_engine()` 为每个 endpoint 执行： | 138 | +`NativeEngineService` 构造 `LaunchContext`，Native Engine Backend 加载角色对应配置并直接生成原生命令： |
| 135 | 139 | ||
| 136 | ```text | 140 | ```text |
| 137 | -engine_server \ | 141 | +vllm serve <model> <native vLLM args...> |
| 138 | - --dp-rank <endpoint.id> \ | 142 | + |
| 139 | - --instance-id <instance_id> \ | 143 | +python3 -m sglang.launch_server <native SGLang args...> |
| 140 | - --role <prefill|decode|union> \ | ||
| 141 | - --host <endpoint.ip> \ | ||
| 142 | - --port <endpoint.business_port> \ | ||
| 143 | - --mgmt-port <endpoint.mgmt_port> \ | ||
| 144 | - --master-dp-ip <master_dp_ip> \ | ||
| 145 | - --node-rank <node_rank> \ | ||
| 146 | - --config-path <USER_CONFIG_PATH> | ||
| 147 | ``` | 144 | ``` |
| 148 | 145 | ||
| 149 | -其他参数和环境变量： | 146 | +统一上下文和环境规则： |
| 150 | 147 | ||
| 151 | - 多 endpoint 模式下，按 `local_world_size` 为每个进程计算 `ASCEND_RT_VISIBLE_DEVICES`；设备编号超出末尾时循环分配。 | 148 | - 多 endpoint 模式下，按 `local_world_size` 为每个进程计算 `ASCEND_RT_VISIBLE_DEVICES`；设备编号超出末尾时循环分配。 |
| 152 | -- 单容器模式追加 `--kv-port`、`--dp-rpc-port`，配置存在时再追加 `--lookup-rpc-port`。 | 149 | +- 单容器端口、D2D peer、角色、DP rank 和 node rank 先写入 `LaunchContext`，再由对应 Backend 映射为原生参数。 |
| 153 | -- 开启快照时追加 `--snapshot-metadata`。 | ||
| 154 | -- D2D peer 按 endpoint ID 过滤后，以逗号分隔并通过 `--d2d-peer-ips` 传递。 | ||
| 155 | - 环境中存在 `POD_IP` 且未设置 `VLLM_HOST_IP` 时，自动设置 `VLLM_HOST_IP=POD_IP`。 | 150 | - 环境中存在 `POD_IP` 且未设置 `VLLM_HOST_IP` 时，自动设置 `VLLM_HOST_IP=POD_IP`。 |
| 156 | - `MOONCAKE_ASCEND_IPV6_EXPERIMENT=1` 时，默认设置 `MC_USE_IPV6=1`。 | 151 | - `MOONCAKE_ASCEND_IPV6_EXPERIMENT=1` 时，默认设置 `MC_USE_IPV6=1`。 |
| 157 | 152 | ||
| @@ -159,9 +154,9 @@ endpoint 的业务端口必须处于 `[1024, 65535]`，IP 必须是合法的 IPv | |||
| 159 | 154 | ||
| 160 | ### 跨节点 PCP | 155 | ### 跨节点 PCP |
| 161 | 156 | ||
| 162 | -Node Manager 始终将 Controller 分配的 `node_rank` 作为 `--node-rank` 传给 Engine Server，并将 `master_dp_ip` 作为 `--master-dp-ip` 传递。 | 157 | +Node Manager 始终将 Controller 分配的 `node_rank` 和 `master_dp_ip` 交给 Native Engine Backend。 |
| 163 | 158 | ||
| 164 | -对于 vLLM，Engine Server 在引擎配置包含 `nnodes > 1` 且配置了 `master_port`（兼容 `master-port`）时启用跨节点 PCP： | 159 | +对于 vLLM，引擎配置包含 `nnodes > 1` 且配置了 `master_port`（兼容 `master-port`）时启用跨节点 PCP： |
| 165 | 160 | ||
| 166 | - `master_addr` 使用 `master_dp_ip`。 | 161 | - `master_addr` 使用 `master_dp_ip`。 |
| 167 | - `node_rank == 0` 为主节点。 | 162 | - `node_rank == 0` 为主节点。 |
| @@ -179,19 +174,21 @@ Node Manager 从 `engine_config.nnodes` 推导每节点 `local_world_size`。当 | |||
| 179 | |--------|--------|------| | 174 | |--------|--------|------| |
| 180 | | `api_config.pod_ip` | `Env.pod_ip` 或 `127.0.0.1` | 注册地址和 API 监听地址 | | 175 | | `api_config.pod_ip` | `Env.pod_ip` 或 `127.0.0.1` | 注册地址和 API 监听地址 | |
| 181 | | `api_config.node_manager_port` | `1026` | Node Manager 管理端口 | | 176 | | `api_config.node_manager_port` | `1026` | Node Manager 管理端口 | |
| 182 | -| `endpoint_config.base_port` | `10000` | Engine Server 端口基址；业务/管理端口按偶数/奇数生成 | | 177 | +| `endpoint_config.base_port` | `10000` | 当前控制面端口基址；业务/管理端口仍按偶数/奇数生成，管理端口将在剩余消费者迁移后删除 | |
| 183 | | `basic_config.heartbeat_interval_seconds` | `3` | 向 Controller 上报心跳的周期 | | 178 | | `basic_config.heartbeat_interval_seconds` | `3` | 向 Controller 上报心跳的周期 | |
| 184 | | `basic_config.daemon_loop_interval` | `5.0` | daemon loop 检查间隔（秒），控制自杀标志轮询和 stdin 检查频率，支持热更新 | | 179 | | `basic_config.daemon_loop_interval` | `5.0` | daemon loop 检查间隔（秒），控制自杀标志轮询和 stdin 检查频率，支持热更新 | |
| 185 | | `basic_config.enable_multi_endpoints` | `true` | 是否按 DP 和设备数创建多个 endpoint | | 180 | | `basic_config.enable_multi_endpoints` | `true` | 是否按 DP 和设备数创建多个 endpoint | |
| 186 | | `basic_config.nnodes` | `1` | 从 `engine_config.nnodes` 派生的跨节点数量 | | 181 | | `basic_config.nnodes` | `1` | 从 `engine_config.nnodes` 派生的跨节点数量 | |
| 187 | | `kv_cache_store_config.mode` | `combined` | 部署模式：`combined` 表示 Engine 与 KV-store 在同一 Pod，`separated` 表示 KV-store 独立 Pod（不拉 Engine、不注册、不心跳） | | 182 | | `kv_cache_store_config.mode` | `combined` | 部署模式：`combined` 表示 Engine 与 KV-store 在同一 Pod，`separated` 表示 KV-store 独立 Pod（不拉 Engine、不注册、不心跳） | |
| 188 | -| `mgmt_tls_config.enable_tls` | `false` | Node Manager、Controller 和 Engine Server 管理面通信是否启用 TLS | | 183 | +| `mgmt_tls_config.enable_tls` | `false` | Node Manager 与 Controller 管理面通信是否启用 TLS | |
| 184 | +| `health_check_config.startup_timeout` | `1800` | 原生引擎模型加载启动窗口，窗口内 `/health` 未监听不判死 | | ||
| 185 | +| `health_check_config.health_collector_timeout_retry_attempts` | `3` | 单次原生 `/health` 请求超时后的最大尝试次数；仅超时重试，包含首次请求 | | ||
| 189 | | `fault_tolerance_config.enable_fault_tolerance` | `false` | 显式开启软件故障轮询；引擎 user config 检测到 FT 时自动开启，无需配置 | | 186 | | `fault_tolerance_config.enable_fault_tolerance` | `false` | 显式开启软件故障轮询；引擎 user config 检测到 FT 时自动开启，无需配置 | |
| 190 | | `fault_tolerance_config.poll_interval_sec` | `5.0` | 轮询引擎 FT 状态的时间间隔（秒） | | 187 | | `fault_tolerance_config.poll_interval_sec` | `5.0` | 轮询引擎 FT 状态的时间间隔（秒） | |
| 191 | | `fault_tolerance_config.poll_timeout_sec` | `5.0` | 单次轮询的 HTTP 超时（秒） | | 188 | | `fault_tolerance_config.poll_timeout_sec` | `5.0` | 单次轮询的 HTTP 超时（秒） | |
| 192 | | `fault_tolerance_config.max_poll_failures` | `3` | 连续轮询失败阈值，达到后按 `dead` 上报 | | 189 | | `fault_tolerance_config.max_poll_failures` | `3` | 连续轮询失败阈值，达到后按 `dead` 上报 | |
| 193 | -| `snapshot_config.enable_snapshot` | `false` | 是否启用容器快照流程 | | 190 | +| `snapshot_config.enable_snapshot` | `false` | 原生引擎运行时尚不支持容器快照；设为 `true` 会在配置校验阶段失败 | |
| 194 | -| `snapshot_config.snapshot_metadata_path` | 空 | 自定义快照元数据路径；用户需预先创建并挂载该文件。为空时进入快照默认应用场景，即 MindCluster 实例重调度 | | 191 | +| `snapshot_config.snapshot_metadata_path` | 空 | 预留字段；当前原生引擎运行时不消费该路径 | |
| 195 | | `port_allocator_config.enable` | `true` | 是否在启动时自动检查并调整端口 | | 192 | | `port_allocator_config.enable` | `true` | 是否在启动时自动检查并调整端口 | |
| 196 | 193 | ||
| 197 | `endpoint_num`、`service_ports`、`mgmt_ports`、`device_num`、`parallel_config`、`model_name`、`engine_type` 和 `dispatch_capabilities` 主要由部署配置与引擎配置派生。`dispatch_capabilities` 不接受用户直接覆盖。 | 194 | `endpoint_num`、`service_ports`、`mgmt_ports`、`device_num`、`parallel_config`、`model_name`、`engine_type` 和 `dispatch_capabilities` 主要由部署配置与引擎配置派生。`dispatch_capabilities` 不接受用户直接覆盖。 |
| @@ -220,7 +217,21 @@ services/ | |||
| 220 | __init__.py | 217 | __init__.py |
| 221 | protocols.py ← DaemonService, PreparableService（接口契约） | 218 | protocols.py ← DaemonService, PreparableService（接口契约） |
| 222 | registry.py ← _ServiceRegistry（服务发现与注册中心） | 219 | registry.py ← _ServiceRegistry（服务发现与注册中心） |
| 223 | - engine.py ← EngineService（Engine 子进程生命周期） | 220 | + native_engine/ |
| 221 | + __init__.py | ||
| 222 | + service.py ← NativeEngineService（原生引擎服务编排） | ||
| 223 | + models.py ← LaunchContext、LaunchSpec、ProbeSpec 和 RuntimeState | ||
| 224 | + supervisor.py ← 公共进程组、健康探测和状态管理 | ||
| 225 | + factory.py ← 按 engine_type 选择 Backend | ||
| 226 | + config_factory.py ← 延迟加载引擎配置转换器 | ||
| 227 | + backends/ | ||
| 228 | + base.py ← NativeEngineBackend 和 IConfig 接口 | ||
| 229 | + vllm/ | ||
| 230 | + backend.py ← vLLM 启动策略 | ||
| 231 | + config.py ← vLLM 参数转换与校验 | ||
| 232 | + sglang/ | ||
| 233 | + backend.py ← SGLang 启动策略 | ||
| 234 | + config.py ← SGLang 参数转换与校验 | ||
| 224 | memcache/ | 235 | memcache/ |
| 225 | __init__.py | 236 | __init__.py |
| 226 | worker.py ← memcache worker 子进程入口（DistributedObjectStore） | 237 | worker.py ← memcache worker 子进程入口（DistributedObjectStore） |
| @@ -235,7 +246,7 @@ services/ | |||
| 235 | from motor.node_manager.core.services.registry import register_service | 246 | from motor.node_manager.core.services.registry import register_service |
| 236 | 247 | ||
| 237 | @register_service("engine", backend="engine") | 248 | @register_service("engine", backend="engine") |
| 238 | -class EngineService: | 249 | +class NativeEngineService: |
| 239 | ... | 250 | ... |
| 240 | 251 | ||
| 241 | @register_service("kv_store", backend="memcache", prepare_priority=10) | 252 | @register_service("kv_store", backend="memcache", prepare_priority=10) |
| @@ -255,6 +266,17 @@ class LocalService: | |||
| 255 | 266 | ||
| 256 | `registry.discover(services)` 解析逗号分隔的服务列表，导入对应模块的 `@register_service` 触发注册。 | 267 | `registry.discover(services)` 解析逗号分隔的服务列表，导入对应模块的 `@register_service` 触发注册。 |
| 257 | 268 | ||
| 269 | +### 新增原生引擎 Backend | ||
| 270 | + | ||
| 271 | +vLLM 和 SGLang 是 `NativeEngineService` 的引擎 Backend，不是独立的 Daemon service。新增原生引擎时： | ||
| 272 | + | ||
| 273 | +1. 在 `native_engine/backends/<engine>/` 中实现 `backend.py` 和 `config.py`。 | ||
| 274 | +2. Backend 只负责引擎差异化的配置转换、启动命令、角色校验和探针规格。 | ||
| 275 | +3. 在 `native_engine/factory.py` 注册 Backend，在 `native_engine/config_factory.py` 注册配置转换器。 | ||
| 276 | +4. 进程启动、进程组清理、健康探测和状态机继续复用公共 `ProcessSupervisor`。 | ||
| 277 | + | ||
| 278 | +不要在引擎 Backend 中直接调用 `subprocess.Popen`，也不要为每个引擎复制停止和探测状态机。 | ||
| 279 | + | ||
| 258 | ### 新增后端 | 280 | ### 新增后端 |
| 259 | 281 | ||
| 260 | 新增 KV 后端只需两步，无需修改现有代码： | 282 | 新增 KV 后端只需两步，无需修改现有代码： |
| @@ -299,32 +321,11 @@ GET http://{endpoint.ip}:{endpoint.business_port}/fault_tolerance/status | |||
| 299 | 321 | ||
| 300 | ## 容器快照 | 322 | ## 容器快照 |
| 301 | 323 | ||
| 302 | -开启 `snapshot_config.enable_snapshot` 后： | 324 | +当前 Node Manager 已切换为直接拉起原生 vLLM/SGLang，但原生 `suspend/resume/device-unlock` |
| 325 | +控制接口尚未接入。为避免仅校验 metadata 文件、实际却未执行快照操作，配置 | ||
| 326 | +`snapshot_config.enable_snapshot=true` 会在启动前明确失败。 | ||
| 303 | 327 | ||
| 304 | -- Node Manager 不支持配置热更新, 不启动配置文件 watcher。 | 328 | +待 Native Engine Backend 实现原生快照控制契约后，再恢复 metadata、checkpoint barrier 和恢复编排。 |
| 305 | -- 配置为空时进入快照默认应用场景，即 MindCluster 实例重调度：容器快照镜像由 MindCluster 制作；MindCluster 通过 ConfigMap 挂载快照元数据，NodeManager 将挂载文件复制到默认可写路径 `/snapshot/snapshot_metadata.json` 后交给 Engine Server 使用。 | ||
| 306 | -- 自定义路径场景下，用户需预先创建并挂载快照元数据文件；框架读取或更新该文件，并将其路径传给 Engine Server，不负责该文件的创建和挂载。 | ||
| 307 | -- 快照制作阶段，Engine Server 完成 suspend 后，其管理面状态由 `INIT` 变为 `NORMAL`。当本节点全部 Engine Server 均完成 suspend 时，表示实例节点容器已到达稳态点：快照默认应用场景通过 `/readiness` 返回 `200` 判断；用户自定义应用场景通过 `/node-manager/status` 返回 `200 {"status": true}` 判断。 | ||
| 308 | -- 查询到稳态点后，对实例节点容器执行 checkpoint，并保存容器 Host 快照镜像。 | ||
| 309 | -- 处于容器快照镜像 checkpoint 过程中的实例无法提供服务, 当 NodeManager 状态已正常但 checkpoint 尚未完成时，此时暂停向 Controller 上报心跳。 | ||
| 310 | -- 快照恢复后，先从元数据恢复 `job_name` 和 `namespace`，刷新 Pod IP 与 Controller DNS，再重新注册。 | ||
| 311 | -- Controller 再次调用 `/node-manager/start` 时，Node Manager 只准备快照恢复阶段需要的 `model_load_path` 和 `data_parallel_master_ip` 元数据，但不重新创建 Engine Server 进程。 | ||
| 312 | -- 快照恢复后未收到启动命令前，readiness 始终为未就绪。 | ||
| 313 | - | ||
| 314 | -### 快照元数据字段 | ||
| 315 | - | ||
| 316 | -快照元数据文件必须是 JSON 对象，以下字段的值均为字符串。自定义应用场景下，用户需按字段所处阶段提前准备元数据。 | ||
| 317 | - | ||
| 318 | -| 字段 | 使用阶段 | 准备要求 | 说明 | | ||
| 319 | -|------|----------|----------|------| | ||
| 320 | -| `model_save_path` | 快照制作 | 制作容器快照前必须准备 | Device 快照保存时，容器内运行时权重的落盘路径，必须是宿主机挂载路径 | | ||
| 321 | -| `model_load_path` | 快照恢复 | 从容器快照恢复前必须准备 | Device 快照恢复时，容器内运行时权重的加载路径，必须是宿主机挂载路径 | | ||
| 322 | -| `job_name` | 快照恢复 | 从容器快照恢复前必须准备 | 恢复后注册时用于更新 Node Manager 的任务名 | | ||
| 323 | -| `namespace` | 快照恢复 | Controller 使用集群内 `.svc.cluster.local` DNS 时必须准备 | 恢复后注册时用于将 Controller DNS 更新到快照所属 namespace；非集群 DNS 场景可不配置 | | ||
| 324 | -| `data_parallel_master_ip` | 快照恢复 | 可不预先配置, 由controller下发 | 优先使用文件中的值；未配置时，Node Manager 写入 Controller 下发的 `master_dp_ip` | | ||
| 325 | -| `checkpoint` | 快照制作 | Host 侧 checkpoint 完成后写入 | 用户或 MindCluster 将其更新为 `"done"`，框架据此解锁 Device 并恢复冷启动实例业务 | | ||
| 326 | - | ||
| 327 | -因此，自定义应用场景从容器快照恢复前，至少需要准备 `model_load_path` 和 `job_name`；使用集群内 Controller DNS 时还需准备 `namespace`。元数据中的其他未知字段不会被 Node Manager 使用。 | ||
| 328 | 329 | ||
| 329 | ## 使用样例 | 330 | ## 使用样例 |
| 330 | 331 | ||
| @@ -350,8 +351,8 @@ curl -i http://127.0.0.1:1026/readiness | |||
| 350 | - 日志持续出现 `Registration attempt N failed`：检查 Controller DNS、端口、TLS 配置和网络连通性；Node Manager 会持续重试注册，不会因此退出。 | 351 | - 日志持续出现 `Registration attempt N failed`：检查 Controller DNS、端口、TLS 配置和网络连通性；Node Manager 会持续重试注册，不会因此退出。 |
| 351 | - 日志出现 `Start command validation failed`：检查 Controller 下发的 `job_name`、endpoint 数量和 endpoint IP 是否与 Node Manager 配置一致。 | 352 | - 日志出现 `Start command validation failed`：检查 Controller 下发的 `job_name`、endpoint 数量和 endpoint IP 是否与 Node Manager 配置一致。 |
| 352 | - 日志出现 `Invalid endpoint parameters`：检查 endpoint IP 与业务端口，业务端口必须处于 `[1024, 65535]`。 | 353 | - 日志出现 `Invalid endpoint parameters`：检查 endpoint IP 与业务端口，业务端口必须处于 `[1024, 65535]`。 |
| 353 | -- 日志出现 `Engine process exited immediately`：Engine Server 在 `Popen` 后立即退出，需继续检查 Engine Server 日志、配置路径和启动参数。 | 354 | +- 日志出现 `Engine process exited immediately`：原生引擎在 `Popen` 后立即退出，需继续检查原生引擎日志、配置路径和启动参数。 |
| 354 | - `/readiness` 返回 `503`：无 endpoint、存在非 `NORMAL` endpoint、处于 `PAUSED`，或快照恢复后尚未收到启动命令。 | 355 | - `/readiness` 返回 `503`：无 endpoint、存在非 `NORMAL` endpoint、处于 `PAUSED`，或快照恢复后尚未收到启动命令。 |
| 355 | -- 连续出现 `Consecutive abnormal heartbeat count: 5/5`：Node Manager 将清理 Engine Server 并以 `-1` 退出触发重调度。 | 356 | +- 连续出现 `Consecutive abnormal heartbeat count: 5/5`：Node Manager 将清理原生引擎进程组并以 `-1` 退出触发 Pod 级重调度。 |
| 356 | 357 | ||
| 357 | 相关单元测试位于 `tests/node_manager/`；优雅暂停流程测试位于 `tests/e2e/test_prestop_e2e.py`。 | 358 | 相关单元测试位于 `tests/node_manager/`；优雅暂停流程测试位于 `tests/e2e/test_prestop_e2e.py`。 |
| @@ -280,10 +280,10 @@ IP与端口参见[内部接口的IP/端口](./README.md#内部接口的ip端口) | |||
| 280 | Coordinator 在识别到实例 `engine_type=sglang` 时，会在发往业务口的请求中直接注入原生字段： | 280 | Coordinator 在识别到实例 `engine_type=sglang` 时，会在发往业务口的请求中直接注入原生字段： |
| 281 | 281 | ||
| 282 | - `bootstrap_host`：Prefill 实例 IP | 282 | - `bootstrap_host`：Prefill 实例 IP |
| 283 | -- `bootstrap_port`：环境变量 `DISAGGREGATION_BOOTSTRAP_PORT`（Coordinator 与引擎 Pod 需一致） | 283 | +- `bootstrap_port`：Prefill Endpoint 注册的原生 PD bootstrap 端口 |
| 284 | - `bootstrap_room`：由 `pair_id` + `attempt_seq` 稳定派生 | 284 | - `bootstrap_room`：由 `pair_id` + `attempt_seq` 稳定派生 |
| 285 | 285 | ||
| 286 | -此时**不会**再附带 `_motor_dispatch`。因此 Coordinator 对 SGLang 存在引擎协议耦合；vLLM PD 仍走 `_motor_dispatch` / InferEndpoint（非 native）路径。 | 286 | +此时**不会**再附带 `_motor_dispatch`。SGLang 和 vLLM 的原生 PD 请求差异由 Coordinator 的协议适配器处理。 |
| 287 | 287 | ||
| 288 | ### 取消 / stop | 288 | ### 取消 / stop |
| 289 | 289 | ||
| @@ -219,7 +219,6 @@ DEFAULT_MMC_LOCAL_CONFIG_PATH = "/usr/local/Ascend/pyMotor/conf/mmc-local-inproc | |||
| 219 | ENV_MMC_LOCAL_SERVICE_MODE = "MMC_LOCAL_SERVICE_MODE" | 219 | ENV_MMC_LOCAL_SERVICE_MODE = "MMC_LOCAL_SERVICE_MODE" |
| 220 | MMC_LOCAL_SERVICE_CONFIG_KEY = "local_service_mode" | 220 | MMC_LOCAL_SERVICE_CONFIG_KEY = "local_service_mode" |
| 221 | 221 | ||
| 222 | -ENV_DISAGGREGATION_BOOTSTRAP_PORT = "DISAGGREGATION_BOOTSTRAP_PORT" | ||
| 223 | ENV_ASCEND_MF_STORE_URL = "ASCEND_MF_STORE_URL" | 222 | ENV_ASCEND_MF_STORE_URL = "ASCEND_MF_STORE_URL" |
G 严重程度： 建议 问题： 删除 原因： 本 PR 删除此处常量定义及 coordinator.py/infer_service.py 的注入逻辑，并删除了协调器侧唯一消费方 sglang_native_dispatch.py；但 怎么改： 二选一：a) 若 engine_server 的 SGLang 进程内 dispatch 路径已废弃，删除 sglang_adapter.py 及其测试；b) 若保留，则改为从请求体/上下文（endpoint 注册的 bootstrap_port）读取，删除环境变量依赖，与文档保持一致。    tobking 19 天前 评论： 19 天前 评论： | |||
| 224 | ENV_ASCEND_MF_STORE_PORT = "ASCEND_MF_STORE_PORT" | 223 | ENV_ASCEND_MF_STORE_PORT = "ASCEND_MF_STORE_PORT" |
| 225 | ENV_ASCEND_MF_TRANSFER_PROTOCOL = "ASCEND_MF_TRANSFER_PROTOCOL" | 224 | ENV_ASCEND_MF_TRANSFER_PROTOCOL = "ASCEND_MF_TRANSFER_PROTOCOL" |
| @@ -69,16 +69,6 @@ def modify_coordinator_deployment(deployment_data, user_config): | |||
| 69 | 69 | ||
| 70 | container[C.ENV].extend(k8s_utils.build_kv_store_env_items()) | 70 | container[C.ENV].extend(k8s_utils.build_kv_store_env_items()) |
| 71 | 71 | ||
| 72 | - disaggregation_bootstrap_port = ( | ||
| 73 | - user_config.get(C.MOTOR_ENGINE_PREFILL_CONFIG, {}) | ||
| 74 | - .get(C.ENGINE_CONFIG, {}) | ||
| 75 | - .get("disaggregation_bootstrap_port", "") | ||
| 76 | - ) | ||
| 77 | - if disaggregation_bootstrap_port: | ||
| 78 | - container[C.ENV].append( | ||
| 79 | - {C.NAME: C.ENV_DISAGGREGATION_BOOTSTRAP_PORT, C.VALUE: str(disaggregation_bootstrap_port)} | ||
| 80 | - ) | ||
| 81 | - | ||
| 82 | modify_coordinator_replicas(deployment_data, user_config) | 72 | modify_coordinator_replicas(deployment_data, user_config) |
| 83 | pod_spec = deployment_data[C.SPEC][C.TEMPLATE][C.SPEC] | 73 | pod_spec = deployment_data[C.SPEC][C.TEMPLATE][C.SPEC] |
| 84 | apply_node_selector_override(pod_spec, deploy_config, C.COORDINATOR_NODE_SELECTOR) | 74 | apply_node_selector_override(pod_spec, deploy_config, C.COORDINATOR_NODE_SELECTOR) |
| @@ -169,19 +169,6 @@ def _configure_coordinator_role(infer_doc, user_config): | |||
| 169 | } | 169 | } |
| 170 | ) | 170 | ) |
| 171 | 171 | ||
| 172 | - disaggregation_bootstrap_port = ( | ||
| 173 | - user_config.get(C.MOTOR_ENGINE_PREFILL_CONFIG, {}) | ||
| 174 | - .get(C.ENGINE_CONFIG, {}) | ||
| 175 | - .get("disaggregation_bootstrap_port", "") | ||
| 176 | - ) | ||
| 177 | - if disaggregation_bootstrap_port: | ||
| 178 | - coordinator_env.append( | ||
| 179 | - { | ||
| 180 | - C.NAME: C.ENV_DISAGGREGATION_BOOTSTRAP_PORT, | ||
| 181 | - C.VALUE: str(disaggregation_bootstrap_port), | ||
| 182 | - } | ||
| 183 | - ) | ||
| 184 | - | ||
| 185 | if coordinator_env: | 172 | if coordinator_env: |
| 186 | set_container_env(container, coordinator_env) | 173 | set_container_env(container, coordinator_env) |
| 187 | 174 | ||
| @@ -70,7 +70,7 @@ NodeManager 端口从 `user_config.json` 读取，缺省为 **1026**。 | |||
| 70 | 70 | ||
| 71 | | 接口 | 说明 | | 71 | | 接口 | 说明 | |
| 72 | |------|------| | 72 | |------|------| |
| 73 | -| `POST /node-manager/pause` | PreStop 调用，返回 `engine_mgmt_addrs` | | 73 | +| `POST /node-manager/pause` | PreStop 调用，返回非 headless 原生引擎的 `engine_metrics_targets` | |
| 74 | | `POST /node-manager/resume` | PreStop 取消时恢复（如滚动回滚） | | 74 | | `POST /node-manager/resume` | PreStop 取消时恢复（如滚动回滚） | |
| 75 | 75 | ||
| 76 | ## 日志 | 76 | ## 日志 |
| @@ -14,8 +14,8 @@ Prestop graceful shutdown for Kubernetes PreStop hook. | |||
| 14 | Flow: | 14 | Flow: |
| 15 | 1. Read NodeManager port from user_config.json (api_config.node_manager_port) | 15 | 1. Read NodeManager port from user_config.json (api_config.node_manager_port) |
| 16 | 2. POST /node-manager/pause to local NodeManager → endpoints set to PAUSED | 16 | 2. POST /node-manager/pause to local NodeManager → endpoints set to PAUSED |
| 17 | -3. Read engine_mgmt_addrs from the pause response | 17 | +3. Read native engine_metrics_targets from the pause response |
| 18 | -4. Poll engine /metrics locally (num_requests_waiting / running) until | 18 | +4. Poll native engine /metrics (num_requests_waiting / running) until |
| 19 | all drain to zero, then exit (Pod terminates). | 19 | all drain to zero, then exit (Pod terminates). |
| 20 | 20 | ||
| 21 | Config is read from CONFIG_PATH or CONFIGMAP_PATH env var (same pattern as probe.py). | 21 | Config is read from CONFIG_PATH or CONFIGMAP_PATH env var (same pattern as probe.py). |
| @@ -105,6 +105,31 @@ def get_nm_port(config): | |||
| 105 | return DEFAULT_NM_PORT | 105 | return DEFAULT_NM_PORT |
| 106 | 106 | ||
| 107 | 107 | ||
| 108 | +def get_infer_tls_config(config): | ||
| 109 | + """Return the inference TLS section used by native engine endpoints.""" | ||
| 110 | + tls_config = get_val_by_key_path(config, "motor_deploy_config.tls_config.infer_tls_config") | ||
| 111 | + return tls_config if isinstance(tls_config, dict) else {} | ||
| 112 | + | ||
| 113 | + | ||
| 114 | +def build_curl_tls_args(tls_config): | ||
| 115 | + """Map Motor inference TLS fields to curl arguments.""" | ||
| 116 | + if not tls_config.get("enable_tls", False): | ||
| 117 | + return [] | ||
| 118 | + | ||
| 119 | + args = [] | ||
| 120 | + field_to_option = ( | ||
| 121 | + ("ca_file", "--cacert"), | ||
| 122 | + ("cert_file", "--cert"), | ||
| 123 | + ("key_file", "--key"), | ||
| 124 | + ("crl_file", "--crlfile"), | ||
| 125 | + ) | ||
| 126 | + for field, option in field_to_option: | ||
| 127 | + value = tls_config.get(field) | ||
| 128 | + if value: | ||
| 129 | + args.extend([option, str(value)]) | ||
| 130 | + return args | ||
| 131 | + | ||
| 132 | + | ||
| 108 | def _http_post_json(url, timeout=10): | 133 | def _http_post_json(url, timeout=10): |
| 109 | """HTTP POST via curl. Returns (status_code, body_text) or (None, None).""" | 134 | """HTTP POST via curl. Returns (status_code, body_text) or (None, None).""" |
| 110 | try: | 135 | try: |
| @@ -143,11 +168,14 @@ def _http_post_json(url, timeout=10): | |||
| 143 | return None, None | 168 | return None, None |
| 144 | 169 | ||
| 145 | 170 | ||
| 146 | -def _http_get_text(url, timeout=5): | 171 | +def _http_get_text(url, timeout=5, extra_args=None): |
| 147 | """HTTP GET via curl. Returns response body text or None.""" | 172 | """HTTP GET via curl. Returns response body text or None.""" |
| 148 | try: | 173 | try: |
| 174 | + command = ["curl", "-s", "-X", "GET", url] | ||
| 175 | + command.extend(extra_args or []) | ||
| 176 | + command.extend(["--connect-timeout", str(timeout), "--max-time", str(timeout)]) | ||
| 149 | result = subprocess.run( # nosec B607 | 177 | result = subprocess.run( # nosec B607 |
| 150 | - ["curl", "-s", "-X", "GET", url, "--connect-timeout", str(timeout), "--max-time", str(timeout)], | 178 | + command, |
| 151 | capture_output=True, | 179 | capture_output=True, |
| 152 | text=True, | 180 | text=True, |
| 153 | timeout=timeout + 2, | 181 | timeout=timeout + 2, |
| @@ -180,15 +208,52 @@ def send_pause(node_manager_url, config): | |||
| 180 | return None | 208 | return None |
| 181 | 209 | ||
| 182 | 210 | ||
| 183 | -def get_engine_metrics(engine_mgmt_addr): | 211 | +ENGINE_METRIC_NAMES = { |
| 184 | - """Query engine /metrics locally and sum num_requests_waiting/running.""" | 212 | + "vllm": ("vllm:num_requests_waiting", "vllm:num_requests_running"), |
| 185 | - url = f"http://{engine_mgmt_addr}/metrics" | 213 | + "sglang": ("sglang:num_waiting_reqs", "sglang:num_running_reqs"), |
| 186 | - text = _http_get_text(url, timeout=10) | 214 | +} |
| 215 | + | ||
| 216 | + | ||
| 217 | +def get_engine_type(config): | ||
| 218 | + """Read the active engine type from a role-specific or union section.""" | ||
| 219 | + role = os.environ.get("ROLE", "").strip().lower() | ||
| 220 | + section_names = { | ||
| 221 | + "prefill": ("motor_engine_prefill_config",), | ||
| 222 | + "decode": ("motor_engine_decode_config",), | ||
| 223 | + "union": ("motor_engine_union_config",), | ||
| 224 | + "both": ("motor_engine_union_config",), | ||
| 225 | + }.get(role, ()) | ||
| 226 | + candidates = section_names + ( | ||
| 227 | + "motor_engine_union_config", | ||
| 228 | + "motor_engine_prefill_config", | ||
| 229 | + "motor_engine_decode_config", | ||
| 230 | + ) | ||
| 231 | + for section_name in candidates: | ||
| 232 | + section = config.get(section_name) | ||
| 233 | + if isinstance(section, dict) and section.get("engine_type"): | ||
| 234 | + return str(section["engine_type"]).strip().lower() | ||
| 235 | + return "vllm" | ||
| 236 | + | ||
| 237 | + | ||
| 238 | +def get_engine_metrics(metrics_target, tls_config, engine_type="vllm"): | ||
| 239 | + """Query a native engine metrics URL and sum waiting/running requests.""" | ||
| 240 | + text = _http_get_text( | ||
| 241 | + metrics_target, | ||
| 242 | + timeout=10, | ||
| 243 | + extra_args=build_curl_tls_args(tls_config), | ||
| 244 | + ) | ||
| 187 | if not text: | 245 | if not text: |
| 188 | return None | 246 | return None |
| 189 | 247 | ||
| 248 | + metric_names = ENGINE_METRIC_NAMES.get(str(engine_type).strip().lower()) | ||
| 249 | + if metric_names is None: | ||
| 250 | + logger.error("Unsupported engine type for prestop metrics: %s", engine_type) | ||
| 251 | + return None | ||
| 252 | + waiting_name, running_name = metric_names | ||
| 190 | waiting = 0 | 253 | waiting = 0 |
| 191 | running = 0 | 254 | running = 0 |
| 255 | + found_waiting = False | ||
| 256 | + found_running = False | ||
| 192 | for line in text.split("\n"): | 257 | for line in text.split("\n"): |
| 193 | if line.startswith("#") or not line.strip(): | 258 | if line.startswith("#") or not line.strip(): |
| 194 | continue | 259 | continue |
| @@ -196,14 +261,57 @@ def get_engine_metrics(engine_mgmt_addr): | |||
| 196 | if not m: | 261 | if not m: |
| 197 | continue | 262 | continue |
| 198 | val = int(float(m.group(1))) | 263 | val = int(float(m.group(1))) |
| 199 | - if "num_requests_waiting" in line: | 264 | + metric_name = line.split("{", 1)[0].split(None, 1)[0] |
| 265 | + if metric_name == waiting_name: | ||
| 200 | waiting += val | 266 | waiting += val |
| 201 | - elif "num_requests_running" in line: | 267 | + found_waiting = True |
| 268 | + elif metric_name == running_name: | ||
| 202 | running += val | 269 | running += val |
| 270 | + found_running = True | ||
| 271 | + | ||
| 272 | + if not found_waiting or not found_running: | ||
| 273 | + logger.warning( | ||
| 274 | + "Required %s drain metrics are missing from %s", | ||
| 275 | + engine_type, | ||
| 276 | + metrics_target, | ||
| 277 | + ) | ||
| 278 | + return None | ||
| 203 | 279 | ||
| 204 | return {"waiting": waiting, "running": running} | 280 | return {"waiting": waiting, "running": running} |
| 205 | 281 | ||
| 206 | 282 | ||
| 283 | +def wait_for_engine_drain(engine_metrics_targets, tls_config, engine_type, max_wait, poll_interval): | ||
| 284 | + """Wait until every engine reports zero active requests or the deadline expires.""" | ||
| 285 | + start_time = time.monotonic() | ||
| 286 | + while time.monotonic() - start_time < max_wait: | ||
| 287 | + total_waiting = 0 | ||
| 288 | + total_running = 0 | ||
| 289 | + metrics_available = True | ||
| 290 | + | ||
| 291 | + for target in engine_metrics_targets: | ||
| 292 | + metrics = get_engine_metrics(target, tls_config, engine_type) | ||
| 293 | + if metrics is None: | ||
| 294 | + logger.warning("Engine drain metrics unavailable from %s; retrying", target) | ||
| 295 | + metrics_available = False | ||
| 296 | + break | ||
| 297 | + total_waiting += metrics["waiting"] | ||
| 298 | + total_running += metrics["running"] | ||
| 299 | + | ||
| 300 | + if metrics_available: | ||
| 301 | + total_active = total_waiting + total_running | ||
| 302 | + logger.info("active=%d (waiting=%d, running=%d)", total_active, total_waiting, total_running) | ||
| 303 | + if total_active == 0: | ||
| 304 | + elapsed = time.monotonic() - start_time | ||
| 305 | + logger.info("All requests drained after %.1fs", elapsed) | ||
| 306 | + return True | ||
| 307 | + | ||
| 308 | + time.sleep(poll_interval) | ||
| 309 | + | ||
| 310 | + elapsed = time.monotonic() - start_time | ||
| 311 | + logger.info("Timeout after %.1fs, stopping anyway", elapsed) | ||
| 312 | + return False | ||
| 313 | + | ||
| 314 | + | ||
| 207 | def main(): | 315 | def main(): |
| 208 | """Main prestop function. | 316 | """Main prestop function. |
| 209 | 317 | ||
| @@ -244,47 +352,24 @@ def main(): | |||
| 244 | logger.warning("Terminate request failed, exiting anyway") | 352 | logger.warning("Terminate request failed, exiting anyway") |
| 245 | sys.exit(0) | 353 | sys.exit(0) |
| 246 | 354 | ||
| 247 | - engine_mgmt_addrs = response.get("engine_mgmt_addrs", []) | 355 | + engine_metrics_targets = response.get("engine_metrics_targets", []) |
| 248 | - if not engine_mgmt_addrs: | 356 | + if not engine_metrics_targets: |
| 249 | - logger.error("No engine_mgmt_addrs in pause response, exiting") | 357 | + logger.error("No engine_metrics_targets in pause response, exiting") |
| 250 | sys.exit(0) | 358 | sys.exit(0) |
| 251 | 359 | ||
| 252 | - logger.info("Engine mgmt addresses: %s", engine_mgmt_addrs) | 360 | + logger.info("Native engine metrics targets: %s", engine_metrics_targets) |
| 361 | + infer_tls_config = get_infer_tls_config(config) | ||
| 362 | + engine_type = get_engine_type(config) | ||
| 253 | 363 | ||
| 254 | # Step 2: Poll engine /metrics locally until requests drain | 364 | # Step 2: Poll engine /metrics locally until requests drain |
| 255 | logger.info("Polling engine metrics...") | 365 | logger.info("Polling engine metrics...") |
| 256 | - start_time = time.time() | 366 | + wait_for_engine_drain( |
| 257 | - | 367 | + engine_metrics_targets, |
| 258 | - while time.time() - start_time < args.max_wait: | 368 | + infer_tls_config, |
| 259 | - total_waiting = 0 | 369 | + engine_type, |
| 260 | - total_running = 0 | 370 | + args.max_wait, |
| 261 | - all_ok = True | 371 | + args.poll_interval, |
| 262 | - | 372 | + ) |
| 263 | - for addr in engine_mgmt_addrs: | ||
| 264 | - metrics = get_engine_metrics(addr) | ||
| 265 | - if metrics is None: | ||
| 266 | - logger.debug("Engine %s unreachable", addr) | ||
| 267 | - all_ok = False | ||
| 268 | - break | ||
| 269 | - total_waiting += metrics["waiting"] | ||
| 270 | - total_running += metrics["running"] | ||
| 271 | - | ||
| 272 | - if not all_ok: | ||
| 273 | - logger.warning("Engine unreachable, exiting") | ||
| 274 | - break | ||
| 275 | - | ||
| 276 | - total_active = total_waiting + total_running | ||
| 277 | - logger.info("active=%d (waiting=%d, running=%d)", total_active, total_waiting, total_running) | ||
| 278 | - | ||
| 279 | - if total_active == 0: | ||
| 280 | - elapsed = time.time() - start_time | ||
| 281 | - logger.info("All requests drained after %.1fs", elapsed) | ||
| 282 | - break | ||
| 283 | - | ||
| 284 | - time.sleep(args.poll_interval) | ||
| 285 | - else: | ||
| 286 | - elapsed = time.time() - start_time | ||
| 287 | - logger.info("Timeout after %.1fs, stopping anyway", elapsed) | ||
| 288 | 373 | ||
| 289 | sys.exit(0) | 374 | sys.exit(0) |
| 290 | 375 | ||
| @@ -56,7 +56,7 @@ cp -n .env.example .env # launch.sh 在无 .env 时也会自动从 .env.exampl | |||
| 56 | | 观测能力 | 是否需改 pyMotor 配置 | 需要的配置 | | 56 | | 观测能力 | 是否需改 pyMotor 配置 | 需要的配置 | |
| 57 | |----------|----------------------|-----------| | 57 | |----------|----------------------|-----------| |
| 58 | | Coordinator 基础指标（指标总览 / KV 缓存的请求数、KV、吞吐、延迟等） | **否** | Coordinator 默认在管理端口暴露 `/metrics`、`/instance/metrics`，无需额外配置；只需保证该端口可被观测机或主机端口转发访问 | | 58 | | Coordinator 基础指标（指标总览 / KV 缓存的请求数、KV、吞吐、延迟等） | **否** | Coordinator 默认在管理端口暴露 `/metrics`、`/instance/metrics`，无需额外配置；只需保证该端口可被观测机或主机端口转发访问 | |
| 59 | -| Engine / vLLM 指标 | **否**（默认开启） | Engine 在管理端口（默认 `10001`）暴露 `/metrics`；保证端口可达即可 | | 59 | +| Engine / vLLM 指标 | **否**（默认开启） | 原生引擎在业务端口暴露 `/metrics`；Coordinator 使用 `infer_tls_config` 直接采集，需保证业务端口可达 | |
| 60 | | Tracing（Tempo 链路） | **是** | 见下方「1.4.1 Tracing 接入」 | | 60 | | Tracing（Tempo 链路） | **是** | 见下方「1.4.1 Tracing 接入」 | |
| 61 | | 引擎性能剖析（`vllm_profiling_*`） | **是** | 需安装并开启 `ms_service_metric`，见下方「1.4.2 Profiling 接入」 | | 61 | | 引擎性能剖析（`vllm_profiling_*`） | **是** | 需安装并开启 `ms_service_metric`，见下方「1.4.2 Profiling 接入」 | |
| 62 | 62 | ||
| @@ -329,7 +329,7 @@ MOTOR_NAMESPACE=<namespace> ./launch.sh --native | |||
| 329 | export MOTOR_NAMESPACE=<namespace> # K8s namespace / job_id | 329 | export MOTOR_NAMESPACE=<namespace> # K8s namespace / job_id |
| 330 | export MOTOR_NODE_IP=<node-ip> # NodePort 访问 IP | 330 | export MOTOR_NODE_IP=<node-ip> # NodePort 访问 IP |
| 331 | export MOTOR_USER_CONFIG=/path/user_config.json | 331 | export MOTOR_USER_CONFIG=/path/user_config.json |
| 332 | -export MOTOR_ENGINE_MGMT_PORT=10001 # Engine /metrics 管理端口，默认 10001 | 332 | +export MOTOR_ENGINE_BUSINESS_PORT=10000 # 原生引擎 /metrics 业务端口，默认 10000 |
| 333 | export OBS_HOST=<obs-host> # pyMotor 上报 tracing / OTLP 的观测主机 | 333 | export OBS_HOST=<obs-host> # pyMotor 上报 tracing / OTLP 的观测主机 |
| 334 | export OBS_STACK_MODE=minimal|full # 未传 --minimal/--full 时生效 | 334 | export OBS_STACK_MODE=minimal|full # 未传 --minimal/--full 时生效 |
| 335 | export PROXY_SH=/path/to/pymotor-proxy.env # native runtime 下载二进制（dotenv 格式，可选；见 §2.4） | 335 | export PROXY_SH=/path/to/pymotor-proxy.env # native runtime 下载二进制（dotenv 格式，可选；见 §2.4） |
| @@ -366,7 +366,7 @@ launch.sh | |||
| 366 | - **Namespace**：`--namespace` / `MOTOR_NAMESPACE` → `user_config.json` 的 `motor_deploy_config.job_id` → 扫描含 Coordinator observability NodePort 的 namespace。 | 366 | - **Namespace**：`--namespace` / `MOTOR_NAMESPACE` → `user_config.json` 的 `motor_deploy_config.job_id` → 扫描含 Coordinator observability NodePort 的 namespace。 |
| 367 | - **Node IP**：`--node-ip` / `MOTOR_NODE_IP` → Coordinator Pod `hostIP` → Kubernetes Node `InternalIP`。 | 367 | - **Node IP**：`--node-ip` / `MOTOR_NODE_IP` → Coordinator Pod `hostIP` → Kubernetes Node `InternalIP`。 |
| 368 | - **Coordinator**：自动发现 observability NodePort（默认服务端口 `1027`），生成 `/metrics` 及 `type=instance|role|dp|node` 等指标端点。 | 368 | - **Coordinator**：自动发现 observability NodePort（默认服务端口 `1027`），生成 `/metrics` 及 `type=instance|role|dp|node` 等指标端点。 |
| 369 | -- **Engine**：优先 Engine metrics NodePort；无 NodePort 时回退 PodIP + `MOTOR_ENGINE_MGMT_PORT`（默认 `10001`），识别 `vllm-p0` / `vllm-d0` 等命名并推断 `pd_role` 与 `instance_id`。 | 369 | +- **Engine**：优先原生引擎 metrics NodePort；无 NodePort 时回退 PodIP + `MOTOR_ENGINE_BUSINESS_PORT`（默认 `10000`），识别 `vllm-p0` / `vllm-d0` 等命名并推断 `pd_role` 与 `instance_id`。 |
| 370 | - **Tracing**：写入 `OBS_HOST`、`OTLP_HTTP_ENDPOINT=http://<obs-host>:4318/v1/traces`、`OTLP_GRPC_ENDPOINT=http://<obs-host>:4317`。 | 370 | - **Tracing**：写入 `OBS_HOST`、`OTLP_HTTP_ENDPOINT=http://<obs-host>:4318/v1/traces`、`OTLP_GRPC_ENDPOINT=http://<obs-host>:4317`。 |
| 371 | 371 | ||
| 372 | ### 3.7 默认端口 | 372 | ### 3.7 默认端口 |
| @@ -9,7 +9,7 @@ cd "${SCRIPT_DIR}" | |||
| 9 | NAMESPACE="${MOTOR_NAMESPACE:-}" | 9 | NAMESPACE="${MOTOR_NAMESPACE:-}" |
| 10 | NODE_IP="${MOTOR_NODE_IP:-}" | 10 | NODE_IP="${MOTOR_NODE_IP:-}" |
| 11 | USER_CONFIG="${MOTOR_USER_CONFIG:-}" | 11 | USER_CONFIG="${MOTOR_USER_CONFIG:-}" |
| 12 | -ENGINE_MGMT_PORT="${MOTOR_ENGINE_MGMT_PORT:-10001}" | 12 | +ENGINE_BUSINESS_PORT="${MOTOR_ENGINE_BUSINESS_PORT:-10000}" |
| 13 | OBS_HOST_INPUT="${OBS_HOST:-}" | 13 | OBS_HOST_INPUT="${OBS_HOST:-}" |
| 14 | 14 | ||
| 15 | FORCE_NATIVE=0 | 15 | FORCE_NATIVE=0 |
| @@ -36,7 +36,7 @@ Environment: | |||
| 36 | MOTOR_NAMESPACE | 36 | MOTOR_NAMESPACE |
| 37 | MOTOR_NODE_IP | 37 | MOTOR_NODE_IP |
| 38 | MOTOR_USER_CONFIG | 38 | MOTOR_USER_CONFIG |
| 39 | - MOTOR_ENGINE_MGMT_PORT | 39 | + MOTOR_ENGINE_BUSINESS_PORT |
| 40 | OBS_HOST | 40 | OBS_HOST |
| 41 | PROXY_SH dotenv file for native binary downloads only (see SERVICE_GUIDE.md §2.4) | 41 | PROXY_SH dotenv file for native binary downloads only (see SERVICE_GUIDE.md §2.4) |
| 42 | 42 | ||
| @@ -108,7 +108,7 @@ DISCOVERY_RUNTIME="docker" | |||
| 108 | if [[ "${FORCE_NATIVE}" -eq 1 ]]; then | 108 | if [[ "${FORCE_NATIVE}" -eq 1 ]]; then |
| 109 | DISCOVERY_RUNTIME="native" | 109 | DISCOVERY_RUNTIME="native" |
| 110 | fi | 110 | fi |
| 111 | -DISCOVERY_CMD=(python3 "./scripts/discover-targets.py" "--output-dir" "./generated" "--engine-mgmt-port" "${ENGINE_MGMT_PORT}" "--runtime" "${DISCOVERY_RUNTIME}") | 111 | +DISCOVERY_CMD=(python3 "./scripts/discover-targets.py" "--output-dir" "./generated" "--engine-business-port" "${ENGINE_BUSINESS_PORT}" "--runtime" "${DISCOVERY_RUNTIME}") |
| 112 | [[ -n "${NAMESPACE}" ]] && DISCOVERY_CMD+=("--namespace" "${NAMESPACE}") | 112 | [[ -n "${NAMESPACE}" ]] && DISCOVERY_CMD+=("--namespace" "${NAMESPACE}") |
| 113 | [[ -n "${NODE_IP}" ]] && DISCOVERY_CMD+=("--node-ip" "${NODE_IP}") | 113 | [[ -n "${NODE_IP}" ]] && DISCOVERY_CMD+=("--node-ip" "${NODE_IP}") |
| 114 | [[ -n "${USER_CONFIG}" ]] && DISCOVERY_CMD+=("--user-config" "${USER_CONFIG}") | 114 | [[ -n "${USER_CONFIG}" ]] && DISCOVERY_CMD+=("--user-config" "${USER_CONFIG}") |
| @@ -134,7 +134,7 @@ fi | |||
| 134 | run_native() { | 134 | run_native() { |
| 135 | echo "[launch] starting native runtime..." | 135 | echo "[launch] starting native runtime..." |
| 136 | echo "[launch] refreshing discovery for native runtime..." | 136 | echo "[launch] refreshing discovery for native runtime..." |
| 137 | - NATIVE_DISCOVERY_CMD=(python3 "./scripts/discover-targets.py" "--output-dir" "./generated" "--engine-mgmt-port" "${ENGINE_MGMT_PORT}" "--runtime" "native") | 137 | + NATIVE_DISCOVERY_CMD=(python3 "./scripts/discover-targets.py" "--output-dir" "./generated" "--engine-business-port" "${ENGINE_BUSINESS_PORT}" "--runtime" "native") |
| 138 | [[ -n "${NAMESPACE}" ]] && NATIVE_DISCOVERY_CMD+=("--namespace" "${NAMESPACE}") | 138 | [[ -n "${NAMESPACE}" ]] && NATIVE_DISCOVERY_CMD+=("--namespace" "${NAMESPACE}") |
| 139 | [[ -n "${NODE_IP}" ]] && NATIVE_DISCOVERY_CMD+=("--node-ip" "${NODE_IP}") | 139 | [[ -n "${NODE_IP}" ]] && NATIVE_DISCOVERY_CMD+=("--node-ip" "${NODE_IP}") |
| 140 | [[ -n "${USER_CONFIG}" ]] && NATIVE_DISCOVERY_CMD+=("--user-config" "${USER_CONFIG}") | 140 | [[ -n "${USER_CONFIG}" ]] && NATIVE_DISCOVERY_CMD+=("--user-config" "${USER_CONFIG}") |
| @@ -1,5 +1,12 @@ | |||
| 1 | -#!/usr/bin/env python3 | 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. |
| 2 | -# -*- coding: utf-8 -*- | 2 | +# MindIE is licensed under Mulan PSL v2. |
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 3 | 10 | ||
| 4 | """ | 11 | """ |
| 5 | Auto-discover pyMotor observability endpoints and generate runtime configs. | 12 | Auto-discover pyMotor observability endpoints and generate runtime configs. |
| @@ -16,10 +23,10 @@ import subprocess | |||
| 16 | import sys | 23 | import sys |
| 17 | from dataclasses import dataclass | 24 | from dataclasses import dataclass |
| 18 | from pathlib import Path | 25 | from pathlib import Path |
| 19 | -from typing import Any, Dict, List, Optional, Sequence, Tuple | 26 | +from typing import Any, Sequence |
| 20 | 27 | ||
| 21 | DEFAULT_COORDINATOR_PORT = 1027 | 28 | DEFAULT_COORDINATOR_PORT = 1027 |
| 22 | -DEFAULT_ENGINE_PORT = 10001 | 29 | +DEFAULT_ENGINE_PORT = 10000 |
| 23 | DEFAULT_OBS_HOST = "localhost" | 30 | DEFAULT_OBS_HOST = "localhost" |
| 24 | DEFAULT_PROMETHEUS_PORT = 9090 | 31 | DEFAULT_PROMETHEUS_PORT = 9090 |
| 25 | DEFAULT_GRAFANA_PORT = 3000 | 32 | DEFAULT_GRAFANA_PORT = 3000 |
| @@ -44,14 +51,14 @@ _PROXY_ENV_KEYS = ( | |||
| 44 | ) | 51 | ) |
| 45 | 52 | ||
| 46 | 53 | ||
| 47 | -def _kubectl_env() -> Dict[str, str]: | 54 | +def _kubectl_env() -> dict[str, str]: |
| 48 | env = os.environ.copy() | 55 | env = os.environ.copy() |
| 49 | for key in _PROXY_ENV_KEYS: | 56 | for key in _PROXY_ENV_KEYS: |
| 50 | env.pop(key, None) | 57 | env.pop(key, None) |
| 51 | return env | 58 | return env |
| 52 | 59 | ||
| 53 | 60 | ||
| 54 | -def _kubectl_path() -> Optional[str]: | 61 | +def _kubectl_path() -> str | None: |
| 55 | return shutil.which("kubectl") | 62 | return shutil.which("kubectl") |
| 56 | 63 | ||
| 57 | 64 | ||
| @@ -86,9 +93,9 @@ class DiscoveryResult: | |||
| 86 | coordinator_target: str | 93 | coordinator_target: str |
| 87 | coordinator_pod_name: str | 94 | coordinator_pod_name: str |
| 88 | coordinator_is_primary: bool | 95 | coordinator_is_primary: bool |
| 89 | - engine_targets: List[Dict[str, Any]] | 96 | + engine_targets: list[dict[str, Any]] |
| 90 | - port_forwards: List[PortForwardSpec] | 97 | + port_forwards: list[PortForwardSpec] |
| 91 | - warnings: List[str] | 98 | + warnings: list[str] |
| 92 | 99 | ||
| 93 | 100 | ||
| 94 | def prefill_count(self) -> int: | 101 | def prefill_count(self) -> int: |
| @@ -99,7 +106,7 @@ class DiscoveryResult: | |||
| 99 | return sum(1 for item in self.engine_targets if item["labels"].get("pd_role") == "decode") | 106 | return sum(1 for item in self.engine_targets if item["labels"].get("pd_role") == "decode") |
| 100 | 107 | ||
| 101 | 108 | ||
| 102 | -def _run_kubectl_json(args: Sequence[str]) -> Dict[str, Any]: | 109 | +def _run_kubectl_json(args: Sequence[str]) -> dict[str, Any]: |
| 103 | kubectl = _kubectl_path() | 110 | kubectl = _kubectl_path() |
| 104 | if kubectl is None: | 111 | if kubectl is None: |
| 105 | raise FileNotFoundError("kubectl not found in PATH") | 112 | raise FileNotFoundError("kubectl not found in PATH") |
| @@ -108,7 +115,7 @@ def _run_kubectl_json(args: Sequence[str]) -> Dict[str, Any]: | |||
| 108 | return json.loads(output.stdout) | 115 | return json.loads(output.stdout) |
| 109 | 116 | ||
| 110 | 117 | ||
| 111 | -def _read_user_config_job_id(path: Optional[str]) -> Optional[str]: | 118 | +def _read_user_config_job_id(path: str | None) -> str | None: |
| 112 | if not path: | 119 | if not path: |
| 113 | return None | 120 | return None |
| 114 | cfg = Path(path) | 121 | cfg = Path(path) |
| @@ -150,15 +157,15 @@ def _is_kubectl_ready() -> bool: | |||
| 150 | return False | 157 | return False |
| 151 | 158 | ||
| 152 | 159 | ||
| 153 | -def _service_name(service: Dict[str, Any]) -> str: | 160 | +def _service_name(service: dict[str, Any]) -> str: |
| 154 | return str(service.get("metadata", {}).get("name", "")) | 161 | return str(service.get("metadata", {}).get("name", "")) |
| 155 | 162 | ||
| 156 | 163 | ||
| 157 | -def _service_namespace(service: Dict[str, Any]) -> str: | 164 | +def _service_namespace(service: dict[str, Any]) -> str: |
| 158 | return str(service.get("metadata", {}).get("namespace", "")) | 165 | return str(service.get("metadata", {}).get("namespace", "")) |
| 159 | 166 | ||
| 160 | 167 | ||
| 161 | -def _service_ports(service: Dict[str, Any]) -> List[Dict[str, Any]]: | 168 | +def _service_ports(service: dict[str, Any]) -> list[dict[str, Any]]: |
| 162 | spec = service.get("spec", {}) | 169 | spec = service.get("spec", {}) |
| 163 | ports = spec.get("ports", []) | 170 | ports = spec.get("ports", []) |
| 164 | return ports if isinstance(ports, list) else [] | 171 | return ports if isinstance(ports, list) else [] |
| @@ -175,7 +182,7 @@ def _is_engine_pod(name: str) -> bool: | |||
| 175 | return _has_keyword(name, ("engine", "mindie-motor-engine")) | 182 | return _has_keyword(name, ("engine", "mindie-motor-engine")) |
| 176 | 183 | ||
| 177 | 184 | ||
| 178 | -def _infer_engine_identity_from_pod(name: str, counters: Dict[str, int]) -> Tuple[str, str]: | 185 | +def _infer_engine_identity_from_pod(name: str, counters: dict[str, int]) -> tuple[str, str]: |
| 179 | match = ENGINE_POD_RE.search(name) | 186 | match = ENGINE_POD_RE.search(name) |
| 180 | if match: | 187 | if match: |
| 181 | role_char = (match.group("role") or "p").lower() | 188 | role_char = (match.group("role") or "p").lower() |
| @@ -185,7 +192,7 @@ def _infer_engine_identity_from_pod(name: str, counters: Dict[str, int]) -> Tupl | |||
| 185 | return _infer_engine_identity(name, counters) | 192 | return _infer_engine_identity(name, counters) |
| 186 | 193 | ||
| 187 | 194 | ||
| 188 | -def _pod_ready_counts(pod: Dict[str, Any]) -> Tuple[int, int]: | 195 | +def _pod_ready_counts(pod: dict[str, Any]) -> tuple[int, int]: |
| 189 | """Return (ready, total) container counts, matching kubectl READY column.""" | 196 | """Return (ready, total) container counts, matching kubectl READY column.""" |
| 190 | statuses = pod.get("status", {}).get("containerStatuses") or [] | 197 | statuses = pod.get("status", {}).get("containerStatuses") or [] |
| 191 | if not statuses: | 198 | if not statuses: |
| @@ -195,13 +202,13 @@ def _pod_ready_counts(pod: Dict[str, Any]) -> Tuple[int, int]: | |||
| 195 | return ready, total | 202 | return ready, total |
| 196 | 203 | ||
| 197 | 204 | ||
| 198 | -def _is_running_coordinator_pod(pod: Dict[str, Any]) -> bool: | 205 | +def _is_running_coordinator_pod(pod: dict[str, Any]) -> bool: |
| 199 | phase = str(pod.get("status", {}).get("phase", "")) | 206 | phase = str(pod.get("status", {}).get("phase", "")) |
| 200 | pod_ip = str(pod.get("status", {}).get("podIP", "")) | 207 | pod_ip = str(pod.get("status", {}).get("podIP", "")) |
| 201 | return phase == "Running" and bool(pod_ip) | 208 | return phase == "Running" and bool(pod_ip) |
| 202 | 209 | ||
| 203 | 210 | ||
| 204 | -def _is_primary_coordinator_pod(pod: Dict[str, Any]) -> bool: | 211 | +def _is_primary_coordinator_pod(pod: dict[str, Any]) -> bool: |
| 205 | """Primary coordinator in HA: kubectl READY shows N/N with N > 0 (e.g. 1/1).""" | 212 | """Primary coordinator in HA: kubectl READY shows N/N with N > 0 (e.g. 1/1).""" |
| 206 | if not _is_running_coordinator_pod(pod): | 213 | if not _is_running_coordinator_pod(pod): |
| 207 | return False | 214 | return False |
| @@ -209,9 +216,9 @@ def _is_primary_coordinator_pod(pod: Dict[str, Any]) -> bool: | |||
| 209 | return total > 0 and ready == total | 216 | return total > 0 and ready == total |
| 210 | 217 | ||
| 211 | 218 | ||
| 212 | -def _list_coordinator_pods(namespace: str) -> List[Dict[str, Any]]: | 219 | +def _list_coordinator_pods(namespace: str) -> list[dict[str, Any]]: |
| 213 | pod_json = _run_kubectl_json(["get", "pods", "-n", namespace]) | 220 | pod_json = _run_kubectl_json(["get", "pods", "-n", namespace]) |
| 214 | - pods: List[Dict[str, Any]] = [] | 221 | + pods: list[dict[str, Any]] = [] |
| 215 | for pod in pod_json.get("items", []): | 222 | for pod in pod_json.get("items", []): |
| 216 | pod_name = str(pod.get("metadata", {}).get("name", "")) | 223 | pod_name = str(pod.get("metadata", {}).get("name", "")) |
| 217 | if _has_keyword(pod_name, COORDINATOR_POD_KEYWORDS): | 224 | if _has_keyword(pod_name, COORDINATOR_POD_KEYWORDS): |
| @@ -222,9 +229,9 @@ def _list_coordinator_pods(namespace: str) -> List[Dict[str, Any]]: | |||
| 222 | def _discover_coordinator_pod( | 229 | def _discover_coordinator_pod( |
| 223 | namespace: str, | 230 | namespace: str, |
| 224 | port: int = DEFAULT_COORDINATOR_PORT, | 231 | port: int = DEFAULT_COORDINATOR_PORT, |
| 225 | -) -> Optional[CoordinatorPodMatch]: | 232 | +) -> CoordinatorPodMatch | None: |
| 226 | - primary_pods: List[Dict[str, Any]] = [] | 233 | + primary_pods: list[dict[str, Any]] = [] |
| 227 | - fallback_pods: List[Dict[str, Any]] = [] | 234 | + fallback_pods: list[dict[str, Any]] = [] |
| 228 | for pod in _list_coordinator_pods(namespace): | 235 | for pod in _list_coordinator_pods(namespace): |
| 229 | if _is_primary_coordinator_pod(pod): | 236 | if _is_primary_coordinator_pod(pod): |
| 230 | primary_pods.append(pod) | 237 | primary_pods.append(pod) |
| @@ -246,12 +253,12 @@ def _discover_coordinator_pod( | |||
| 246 | ) | 253 | ) |
| 247 | 254 | ||
| 248 | 255 | ||
| 249 | -def _find_namespace_from_services() -> Optional[str]: | 256 | +def _find_namespace_from_services() -> str | None: |
| 250 | try: | 257 | try: |
| 251 | svc_json = _run_kubectl_json(["get", "svc", "-A"]) | 258 | svc_json = _run_kubectl_json(["get", "svc", "-A"]) |
| 252 | except subprocess.CalledProcessError: | 259 | except subprocess.CalledProcessError: |
| 253 | return None | 260 | return None |
| 254 | - candidates: List[str] = [] | 261 | + candidates: list[str] = [] |
| 255 | for svc in svc_json.get("items", []): | 262 | for svc in svc_json.get("items", []): |
| 256 | name = _service_name(svc) | 263 | name = _service_name(svc) |
| 257 | if not _has_keyword(name, ("coordinator", "mindie-motor-coordinator")): | 264 | if not _has_keyword(name, ("coordinator", "mindie-motor-coordinator")): |
| @@ -295,7 +302,7 @@ def _has_explicit_namespace(args: argparse.Namespace) -> bool: | |||
| 295 | return bool(args.namespace or os.getenv("MOTOR_NAMESPACE")) | 302 | return bool(args.namespace or os.getenv("MOTOR_NAMESPACE")) |
| 296 | 303 | ||
| 297 | 304 | ||
| 298 | -def _resolve_node_ip(namespace: str, args: argparse.Namespace, warnings: List[str]) -> str: | 305 | +def _resolve_node_ip(namespace: str, args: argparse.Namespace, warnings: list[str]) -> str: |
| 299 | if args.node_ip: | 306 | if args.node_ip: |
| 300 | return args.node_ip | 307 | return args.node_ip |
| 301 | env_node_ip = os.getenv("MOTOR_NODE_IP") | 308 | env_node_ip = os.getenv("MOTOR_NODE_IP") |
| @@ -344,10 +351,10 @@ def _resolve_node_ip(namespace: str, args: argparse.Namespace, warnings: List[st | |||
| 344 | 351 | ||
| 345 | 352 | ||
| 346 | def _match_nodeport_for_service( | 353 | def _match_nodeport_for_service( |
| 347 | - service: Dict[str, Any], | 354 | + service: dict[str, Any], |
| 348 | expected_port: int, | 355 | expected_port: int, |
| 349 | allow_keywords: Sequence[str], | 356 | allow_keywords: Sequence[str], |
| 350 | -) -> Optional[int]: | 357 | +) -> int | None: |
| 351 | for port in _service_ports(service): | 358 | for port in _service_ports(service): |
| 352 | node_port = port.get("nodePort") | 359 | node_port = port.get("nodePort") |
| 353 | if not node_port: | 360 | if not node_port: |
| @@ -361,13 +368,13 @@ def _match_nodeport_for_service( | |||
| 361 | 368 | ||
| 362 | 369 | ||
| 363 | def _find_service_nodeport( | 370 | def _find_service_nodeport( |
| 364 | - services: Sequence[Dict[str, Any]], | 371 | + services: Sequence[dict[str, Any]], |
| 365 | expected_port: int, | 372 | expected_port: int, |
| 366 | service_keywords: Sequence[str], | 373 | service_keywords: Sequence[str], |
| 367 | port_keywords: Sequence[str], | 374 | port_keywords: Sequence[str], |
| 368 | -) -> Optional[Tuple[str, int]]: | 375 | +) -> tuple[str, int] | None: |
| 369 | - best_match: Optional[Tuple[str, int]] = None | 376 | + best_match: tuple[str, int] | None = None |
| 370 | - fallback_match: Optional[Tuple[str, int]] = None | 377 | + fallback_match: tuple[str, int] | None = None |
| 371 | for svc in services: | 378 | for svc in services: |
| 372 | name = _service_name(svc) | 379 | name = _service_name(svc) |
| 373 | matched = _match_nodeport_for_service(svc, expected_port, port_keywords) | 380 | matched = _match_nodeport_for_service(svc, expected_port, port_keywords) |
| @@ -381,7 +388,7 @@ def _find_service_nodeport( | |||
| 381 | return best_match or fallback_match | 388 | return best_match or fallback_match |
| 382 | 389 | ||
| 383 | 390 | ||
| 384 | -def _infer_engine_identity(name: str, counters: Dict[str, int]) -> Tuple[str, str]: | 391 | +def _infer_engine_identity(name: str, counters: dict[str, int]) -> tuple[str, str]: |
| 385 | lowered = name.lower() | 392 | lowered = name.lower() |
| 386 | pd_role = "prefill" | 393 | pd_role = "prefill" |
| 387 | if "decode" in lowered: | 394 | if "decode" in lowered: |
| @@ -402,7 +409,7 @@ def _infer_engine_identity(name: str, counters: Dict[str, int]) -> Tuple[str, st | |||
| 402 | return pd_role, instance_id | 409 | return pd_role, instance_id |
| 403 | 410 | ||
| 404 | 411 | ||
| 405 | -def _infer_dp_rank(name: str, labels: Optional[Dict[str, Any]] = None) -> str: | 412 | +def _infer_dp_rank(name: str, labels: dict[str, Any] | None = None) -> str: |
| 406 | labels = labels or {} | 413 | labels = labels or {} |
| 407 | for key in ("dp_rank", "dp-rank", "rank"): | 414 | for key in ("dp_rank", "dp-rank", "rank"): |
| 408 | if labels.get(key) is not None: | 415 | if labels.get(key) is not None: |
| @@ -415,12 +422,12 @@ def _infer_dp_rank(name: str, labels: Optional[Dict[str, Any]] = None) -> str: | |||
| 415 | 422 | ||
| 416 | 423 | ||
| 417 | def _discover_engine_targets_from_services( | 424 | def _discover_engine_targets_from_services( |
| 418 | - services: Sequence[Dict[str, Any]], | 425 | + services: Sequence[dict[str, Any]], |
| 419 | node_ip: str, | 426 | node_ip: str, |
| 420 | engine_port: int, | 427 | engine_port: int, |
| 421 | cluster_label: str, | 428 | cluster_label: str, |
| 422 | -) -> List[Dict[str, Any]]: | 429 | +) -> list[dict[str, Any]]: |
| 423 | - targets: List[Dict[str, Any]] = [] | 430 | + targets: list[dict[str, Any]] = [] |
| 424 | counters = {"prefill": 0, "decode": 0} | 431 | counters = {"prefill": 0, "decode": 0} |
| 425 | for svc in services: | 432 | for svc in services: |
| 426 | name = _service_name(svc) | 433 | name = _service_name(svc) |
| @@ -429,7 +436,7 @@ def _discover_engine_targets_from_services( | |||
| 429 | node_port = _match_nodeport_for_service( | 436 | node_port = _match_nodeport_for_service( |
| 430 | svc, | 437 | svc, |
| 431 | expected_port=engine_port, | 438 | expected_port=engine_port, |
| 432 | - allow_keywords=("metrics", "mgmt", "manage", "prometheus"), | 439 | + allow_keywords=("metrics", "http", "api", "prometheus"), |
| 433 | ) | 440 | ) |
| 434 | if node_port is None: | 441 | if node_port is None: |
| 435 | continue | 442 | continue |
| @@ -454,9 +461,9 @@ def _discover_engine_targets_from_pods( | |||
| 454 | namespace: str, | 461 | namespace: str, |
| 455 | engine_port: int, | 462 | engine_port: int, |
| 456 | cluster_label: str, | 463 | cluster_label: str, |
| 457 | -) -> List[Dict[str, Any]]: | 464 | +) -> list[dict[str, Any]]: |
| 458 | pod_json = _run_kubectl_json(["get", "pods", "-n", namespace]) | 465 | pod_json = _run_kubectl_json(["get", "pods", "-n", namespace]) |
| 459 | - targets: List[Dict[str, Any]] = [] | 466 | + targets: list[dict[str, Any]] = [] |
| 460 | counters = {"prefill": 0, "decode": 0} | 467 | counters = {"prefill": 0, "decode": 0} |
| 461 | for pod in pod_json.get("items", []): | 468 | for pod in pod_json.get("items", []): |
| 462 | metadata = pod.get("metadata", {}) | 469 | metadata = pod.get("metadata", {}) |
| @@ -490,7 +497,7 @@ def _discover_engine_targets_from_pods( | |||
| 490 | return targets | 497 | return targets |
| 491 | 498 | ||
| 492 | 499 | ||
| 493 | -def _split_target(target: str) -> Tuple[str, int]: | 500 | +def _split_target(target: str) -> tuple[str, int]: |
| 494 | host, port = target.rsplit(":", 1) | 501 | host, port = target.rsplit(":", 1) |
| 495 | return host.strip("[]"), int(port) | 502 | return host.strip("[]"), int(port) |
| 496 | 503 | ||
| @@ -502,7 +509,7 @@ def _register_docker_port_forward( | |||
| 502 | remote_port: int, | 509 | remote_port: int, |
| 503 | pod_name: str, | 510 | pod_name: str, |
| 504 | next_port: int, | 511 | next_port: int, |
| 505 | -) -> Tuple[str, int]: | 512 | +) -> tuple[str, int]: |
| 506 | local_port = next_port | 513 | local_port = next_port |
| 507 | result.port_forwards.append( | 514 | result.port_forwards.append( |
| 508 | PortForwardSpec( | 515 | PortForwardSpec( |
| @@ -553,9 +560,9 @@ def _apply_docker_gateway(result: DiscoveryResult, base_port: int) -> None: | |||
| 553 | 560 | ||
| 554 | 561 | ||
| 555 | def _discover(namespace: str, node_ip: str, args: argparse.Namespace) -> DiscoveryResult: | 562 | def _discover(namespace: str, node_ip: str, args: argparse.Namespace) -> DiscoveryResult: |
| 556 | - warnings: List[str] = [] | 563 | + warnings: list[str] = [] |
| 557 | obs_host = args.obs_host or os.getenv("OBS_HOST") or node_ip or DEFAULT_OBS_HOST | 564 | obs_host = args.obs_host or os.getenv("OBS_HOST") or node_ip or DEFAULT_OBS_HOST |
| 558 | - engine_port = args.engine_mgmt_port | 565 | + engine_port = args.engine_business_port |
| 559 | cluster_label = namespace | 566 | cluster_label = namespace |
| 560 | runtime = args.runtime | 567 | runtime = args.runtime |
| 561 | 568 | ||
| @@ -564,8 +571,8 @@ def _discover(namespace: str, node_ip: str, args: argparse.Namespace) -> Discove | |||
| 564 | coordinator_port = DEFAULT_COORDINATOR_PORT | 571 | coordinator_port = DEFAULT_COORDINATOR_PORT |
| 565 | coordinator_pod_name = "" | 572 | coordinator_pod_name = "" |
| 566 | coordinator_is_primary = False | 573 | coordinator_is_primary = False |
| 567 | - engine_targets: List[Dict[str, Any]] = [] | 574 | + engine_targets: list[dict[str, Any]] = [] |
| 568 | - port_forwards: List[PortForwardSpec] = [] | 575 | + port_forwards: list[PortForwardSpec] = [] |
| 569 | 576 | ||
| 570 | if _is_kubectl_ready(): | 577 | if _is_kubectl_ready(): |
| 571 | try: | 578 | try: |
| @@ -710,11 +717,11 @@ def _discover(namespace: str, node_ip: str, args: argparse.Namespace) -> Discove | |||
| 710 | 717 | ||
| 711 | def _render_job( | 718 | def _render_job( |
| 712 | name: str, | 719 | name: str, |
| 713 | - targets: Sequence[Dict[str, Any]], | 720 | + targets: Sequence[dict[str, Any]], |
| 714 | - metrics_path: Optional[str] = None, | 721 | + metrics_path: str | None = None, |
| 715 | honor_labels: bool = False, | 722 | honor_labels: bool = False, |
| 716 | -) -> List[str]: | 723 | +) -> list[str]: |
| 717 | - lines: List[str] = [f" - job_name: {name}"] | 724 | + lines: list[str] = [f" - job_name: {name}"] |
| 718 | if metrics_path: | 725 | if metrics_path: |
| 719 | lines.append(f" metrics_path: {metrics_path}") | 726 | lines.append(f" metrics_path: {metrics_path}") |
| 720 | if honor_labels: | 727 | if honor_labels: |
| @@ -933,10 +940,10 @@ def _parse_args() -> argparse.Namespace: | |||
| 933 | help="First local port for Docker PodIP bridge forwards (default: 19000).", | 940 | help="First local port for Docker PodIP bridge forwards (default: 19000).", |
| 934 | ) | 941 | ) |
| 935 | parser.add_argument( | 942 | parser.add_argument( |
| 936 | - "--engine-mgmt-port", | 943 | + "--engine-business-port", |
| 937 | type=int, | 944 | type=int, |
| 938 | - default=int(os.getenv("MOTOR_ENGINE_MGMT_PORT", str(DEFAULT_ENGINE_PORT))), | 945 | + default=int(os.getenv("MOTOR_ENGINE_BUSINESS_PORT", str(DEFAULT_ENGINE_PORT))), |
| 939 | - help="Engine management metrics port (default: 10001).", | 946 | + help="Native engine business port exposing /metrics (default: 10000).", |
| 940 | ) | 947 | ) |
| 941 | parser.add_argument( | 948 | parser.add_argument( |
| 942 | "--output-dir", | 949 | "--output-dir", |
| @@ -948,7 +955,7 @@ def _parse_args() -> argparse.Namespace: | |||
| 948 | 955 | ||
| 949 | def main() -> int: | 956 | def main() -> int: |
| 950 | args = _parse_args() | 957 | args = _parse_args() |
| 951 | - warnings: List[str] = [] | 958 | + warnings: list[str] = [] |
| 952 | try: | 959 | try: |
| 953 | namespace = _resolve_namespace(args) | 960 | namespace = _resolve_namespace(args) |
| 954 | if _has_explicit_namespace(args) and not _is_kubectl_ready(): | 961 | if _has_explicit_namespace(args) and not _is_kubectl_ready(): |
| @@ -8,6 +8,11 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | +from motor.common.constants import ( | ||
| 12 | + CHAT_COMPLETION_PREFIX as CHAT_COMPLETION_PREFIX, | ||
| 13 | + COMPLETION_PREFIX as COMPLETION_PREFIX, | ||
| 14 | +) | ||
| 15 | + | ||
| 11 | # log dir permission | 16 | # log dir permission |
| 12 | MOTOR_CUSTOM_ZMQ_PRIVILEGE = 0o640 | 17 | MOTOR_CUSTOM_ZMQ_PRIVILEGE = 0o640 |
| 13 | MOTOR_CUSTOM_ZMQ_DIR_PRIVILEGE = 0o750 | 18 | MOTOR_CUSTOM_ZMQ_DIR_PRIVILEGE = 0o750 |
| @@ -9,7 +9,7 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | from enum import Enum | 11 | from enum import Enum |
| 12 | -from typing import Any, Iterable, Literal | 12 | +from typing import Any, Literal |
| 13 | 13 | ||
| 14 | from pydantic import BaseModel, Field, field_validator | 14 | from pydantic import BaseModel, Field, field_validator |
| 15 | 15 | ||
| @@ -100,49 +100,6 @@ def dispatch_capabilities_for_profile(profile: DispatchProfile) -> list[str]: | |||
| 100 | return [] | 100 | return [] |
| 101 | 101 | ||
| 102 | 102 | ||
| 103 | -def dispatch_plans_from_capabilities(values: Iterable[object] | None) -> set[DispatchPlan]: | ||
| 104 | - """Normalize advertised capability values into supported dispatch plans.""" | ||
| 105 | - plans: set[DispatchPlan] = set() | ||
| 106 | - if values is None: | ||
| 107 | - return plans | ||
| 108 | - if isinstance(values, (str, bytes)): | ||
| 109 | - values = (values,) | ||
| 110 | - try: | ||
| 111 | - iterator = iter(values) | ||
| 112 | - except TypeError: | ||
| 113 | - return plans | ||
| 114 | - for value in iterator: | ||
| 115 | - try: | ||
| 116 | - plans.add(DispatchPlan(str(value))) | ||
| 117 | - except ValueError: | ||
| 118 | - continue | ||
| 119 | - return plans | ||
| 120 | - | ||
| 121 | - | ||
| 122 | -def shared_dispatch_plans(prefill: Any, decode: Any) -> set[DispatchPlan]: | ||
| 123 | - """Return dispatch plans advertised by both instances.""" | ||
| 124 | - prefill_plans = dispatch_plans_from_capabilities(getattr(prefill, "dispatch_capabilities", None)) | ||
| 125 | - decode_plans = dispatch_plans_from_capabilities(getattr(decode, "dispatch_capabilities", None)) | ||
| 126 | - return prefill_plans & decode_plans | ||
| 127 | - | ||
| 128 | - | ||
| 129 | -def dispatch_plan_union(instances: Iterable[Any]) -> set[DispatchPlan]: | ||
| 130 | - """Union of dispatch plans advertised across instances; parses each instance once.""" | ||
| 131 | - union: set[DispatchPlan] = set() | ||
| 132 | - for instance in instances: | ||
| 133 | - union |= dispatch_plans_from_capabilities(getattr(instance, "dispatch_capabilities", None)) | ||
| 134 | - return union | ||
| 135 | - | ||
| 136 | - | ||
| 137 | -def has_compatible_dispatch_pair(prefill_instances: Iterable[Any], decode_instances: Iterable[Any]) -> bool: | ||
| 138 | - """Whether at least one P/D instance pair advertises a shared dispatch plan. | ||
| 139 | - | ||
| 140 | - A shared plan exists iff the prefill and decode plan unions intersect, so this | ||
| 141 | - runs in O(P+D) without enumerating instance pairs. | ||
| 142 | - """ | ||
| 143 | - return bool(dispatch_plan_union(prefill_instances) & dispatch_plan_union(decode_instances)) | ||
| 144 | - | ||
| 145 | - | ||
| 146 | def _classify_vllm_kv_transfer_config(kv_transfer_config: Any) -> DispatchProfile: | 103 | def _classify_vllm_kv_transfer_config(kv_transfer_config: Any) -> DispatchProfile: |
| 147 | if not isinstance(kv_transfer_config, dict): | 104 | if not isinstance(kv_transfer_config, dict): |
| 148 | return DispatchProfile.UNKNOWN | 105 | return DispatchProfile.UNKNOWN |
| @@ -213,6 +170,12 @@ class DispatchEndpoint(BaseModel): | |||
| 213 | instance_id: int = Field(..., ge=0, description="Scheduler instance identifier for the target engine") | 170 | instance_id: int = Field(..., ge=0, description="Scheduler instance identifier for the target engine") |
| 214 | endpoint_id: int = Field(..., ge=0, description="Endpoint identifier within the instance") | 171 | endpoint_id: int = Field(..., ge=0, description="Endpoint identifier within the instance") |
| 215 | url: str = Field(..., min_length=1, description="Base HTTP URL for dispatch and stop calls to the engine") | 172 | url: str = Field(..., min_length=1, description="Base HTTP URL for dispatch and stop calls to the engine") |
| 173 | + bootstrap_port: int | None = Field( | ||
| 174 | + default=None, | ||
| 175 | + ge=1, | ||
| 176 | + le=65535, | ||
| 177 | + description="Engine-native PD bootstrap port when it differs from the inference URL port", | ||
| 178 | + ) | ||
| 216 | 179 | ||
| 217 | 180 | ||
| 218 | class DispatchEndpoints(BaseModel): | 181 | class DispatchEndpoints(BaseModel): |
| @@ -67,6 +67,7 @@ class Endpoint(BaseModel): | |||
| 67 | ip: str = Field(..., description="IP address") | 67 | ip: str = Field(..., description="IP address") |
| 68 | business_port: str = Field(..., description="Business port") | 68 | business_port: str = Field(..., description="Business port") |
| 69 | mgmt_port: str = Field(..., description="Management port") | 69 | mgmt_port: str = Field(..., description="Management port") |
| 70 | + bootstrap_port: int | None = Field(default=None, ge=1, le=65535, description="Native PD bootstrap port") | ||
| 70 | status: EndpointStatus = Field(default=EndpointStatus.INITIAL, description="Endpoint status") | 71 | status: EndpointStatus = Field(default=EndpointStatus.INITIAL, description="Endpoint status") |
| 71 | device_infos: list[DeviceInfo] = Field(default_factory=list, description="List of DeviceInfo") | 72 | device_infos: list[DeviceInfo] = Field(default_factory=list, description="List of DeviceInfo") |
| 72 | hb_timestamp: float = Field(default=0, description="Last heartbeat timestamp") | 73 | hb_timestamp: float = Field(default=0, description="Last heartbeat timestamp") |
| @@ -79,6 +80,7 @@ class Endpoint(BaseModel): | |||
| 79 | ip: str, | 80 | ip: str, |
| 80 | business_port: str, | 81 | business_port: str, |
| 81 | mgmt_port: str, | 82 | mgmt_port: str, |
| 83 | + bootstrap_port: int | None = None, | ||
| 82 | status: EndpointStatus | None = None, | 84 | status: EndpointStatus | None = None, |
| 83 | device_infos: list[DeviceInfo] | None = None, | 85 | device_infos: list[DeviceInfo] | None = None, |
| 84 | hb_timestamp: float | None = None, | 86 | hb_timestamp: float | None = None, |
| @@ -90,6 +92,7 @@ class Endpoint(BaseModel): | |||
| 90 | ip=ip, | 92 | ip=ip, |
| 91 | business_port=business_port, | 93 | business_port=business_port, |
| 92 | mgmt_port=mgmt_port, | 94 | mgmt_port=mgmt_port, |
| 95 | + bootstrap_port=bootstrap_port, | ||
| 93 | status=status if status is not None else EndpointStatus.INITIAL, | 96 | status=status if status is not None else EndpointStatus.INITIAL, |
| 94 | device_infos=device_infos if device_infos is not None else [], | 97 | device_infos=device_infos if device_infos is not None else [], |
| 95 | hb_timestamp=hb_timestamp if hb_timestamp is not None else time.time(), | 98 | hb_timestamp=hb_timestamp if hb_timestamp is not None else time.time(), |
| @@ -51,6 +51,7 @@ class RegisterMsg(BaseModel): | |||
| 51 | pod_ip: str = Field(..., description="Pod IP address") | 51 | pod_ip: str = Field(..., description="Pod IP address") |
| 52 | business_port: list[str] = Field(..., description="Business port for all endpoints managed by this nm") | 52 | business_port: list[str] = Field(..., description="Business port for all endpoints managed by this nm") |
| 53 | mgmt_port: list[str] = Field(..., description="Management port for all endpoints managed by this nm") | 53 | mgmt_port: list[str] = Field(..., description="Management port for all endpoints managed by this nm") |
| 54 | + bootstrap_port: int | None = Field(default=None, ge=1, le=65535, description="Native PD bootstrap port") | ||
| 54 | nm_port: str = Field(..., description="Node manager communication port") | 55 | nm_port: str = Field(..., description="Node manager communication port") |
| 55 | parallel_config: ParallelConfig = Field(..., description="Parallel configuration") | 56 | parallel_config: ParallelConfig = Field(..., description="Parallel configuration") |
| 56 | enable_multi_endpoints: bool = Field(default=True, description="Whether to enable multi-endpoints mode") | 57 | enable_multi_endpoints: bool = Field(default=True, description="Whether to enable multi-endpoints mode") |
| @@ -340,12 +340,24 @@ class Instance(BaseModel): | |||
| 340 | return True | 340 | return True |
| 341 | 341 | ||
| 342 | def is_all_endpoints_ready(self) -> bool: | 342 | def is_all_endpoints_ready(self) -> bool: |
| 343 | + """Return whether routable endpoints and their headless workers are ready. | ||
| 344 | + | ||
| 345 | + Headless PCP workers expose no HTTP readiness endpoint. Their | ||
| 346 | + ``WAIT2START`` status confirms process liveness without advertising a | ||
| 347 | + routable service; the non-headless master remains the readiness gate. | ||
| 348 | + """ | ||
| 343 | with self._lock: | 349 | with self._lock: |
| 350 | + has_routable_endpoint = False | ||
| 344 | for pod_endpoints in self.endpoints.values(): | 351 | for pod_endpoints in self.endpoints.values(): |
| 345 | for endpoint in pod_endpoints.values(): | 352 | for endpoint in pod_endpoints.values(): |
| 353 | + if endpoint.headless: | ||
| 354 | + if endpoint.status not in (EndpointStatus.NORMAL, EndpointStatus.WAIT2START): | ||
| 355 | + return False | ||
| 356 | + continue | ||
| 357 | + has_routable_endpoint = True | ||
| 346 | if endpoint.status != EndpointStatus.NORMAL: | 358 | if endpoint.status != EndpointStatus.NORMAL: |
| 347 | return False | 359 | return False |
| 348 | - return True | 360 | + return has_routable_endpoint |
| 349 | 361 | ||
| 350 | def is_have_one_endpoint_abnormal(self) -> bool: | 362 | def is_have_one_endpoint_abnormal(self) -> bool: |
| 351 | abnormal_endpoints: dict[str, list[int]] = {} # pod_ip -> [endpoint_id] | 363 | abnormal_endpoints: dict[str, list[int]] = {} # pod_ip -> [endpoint_id] |
| @@ -12,8 +12,16 @@ import os | |||
| 12 | from typing import Callable, Any | 12 | from typing import Callable, Any |
| 13 | from typing import List, Tuple | 13 | from typing import List, Tuple |
| 14 | 14 | ||
| 15 | -from motor.engine_server.constants.constants import (MAX_SIZE, MIN_SIZE, MIN_RANK_SIZE, MAX_RANK_SIZE, | 15 | +from motor.common.engine_constants import ( |
| 16 | - MAX_FILE_NUMS, MIN_DEVICE_NUM, MAX_DEVICE_NUM, MUSK_PRIVILEGE) | 16 | + MAX_DEVICE_NUM, |
| 17 | + MAX_FILE_NUMS, | ||
| 18 | + MAX_RANK_SIZE, | ||
| 19 | + MAX_SIZE, | ||
| 20 | + MIN_DEVICE_NUM, | ||
| 21 | + MIN_RANK_SIZE, | ||
| 22 | + MIN_SIZE, | ||
| 23 | + MUSK_PRIVILEGE, | ||
| 24 | +) | ||
| 17 | 25 | ||
| 18 | 26 | ||
| 19 | class Validator: | 27 | class Validator: |
| @@ -50,7 +58,7 @@ class Validator: | |||
| 50 | if self.is_valid_state is None: | 58 | if self.is_valid_state is None: |
| 51 | try: | 59 | try: |
| 52 | self.check() | 60 | self.check() |
| 53 | - except Exception as e: | 61 | + except Exception: |
| 54 | self.is_valid_state = False | 62 | self.is_valid_state = False |
| 55 | return self.is_valid_state | 63 | return self.is_valid_state |
| 56 | 64 | ||
| @@ -933,7 +933,7 @@ class CoordinatorConfig: | |||
| 933 | return reload_dataclass_config_from_json( | 933 | return reload_dataclass_config_from_json( |
| 934 | self, | 934 | self, |
| 935 | self.from_json, | 935 | self.from_json, |
| 936 | - skip=frozenset({"worker_index", "worker_metaserver_port"}), | 936 | + skip=frozenset({"worker_index"}), |
| 937 | skip_private=True, | 937 | skip_private=True, |
| 938 | ) | 938 | ) |
| 939 | 939 | ||
| @@ -20,9 +20,9 @@ from motor.common.resources.dispatch import DISPATCH_PROFILE_KEY | |||
| 20 | from motor.config.config_utils import _update_engine_server_tls_config | 20 | from motor.config.config_utils import _update_engine_server_tls_config |
| 21 | from motor.config.resolver import ConfigResolver, normalize_keys | 21 | from motor.config.resolver import ConfigResolver, normalize_keys |
| 22 | from motor.config.tls_config import TLSConfig | 22 | from motor.config.tls_config import TLSConfig |
| 23 | -from motor.engine_server.constants import constants | 23 | +from motor.common import engine_constants as constants |
| 24 | -from motor.engine_server.utils.ip import ip_valid_check, port_valid_check | 24 | +from motor.common.utils.ip import ip_valid_check, port_valid_check |
| 25 | -from motor.engine_server.utils.validators import FileValidator | 25 | +from motor.common.utils.validators import FileValidator |
| 26 | 26 | ||
| 27 | logger = get_logger(__name__) | 27 | logger = get_logger(__name__) |
| 28 | 28 | ||
| @@ -111,6 +111,9 @@ class HealthCheckConfig: | |||
| 111 | health_collector_timeout: int = 5 | 111 | health_collector_timeout: int = 5 |
| 112 | # Max attempts for /health probe when request times out (timeout-only retry). | 112 | # Max attempts for /health probe when request times out (timeout-only retry). |
| 113 | health_collector_timeout_retry_attempts: int = 3 | 113 | health_collector_timeout_retry_attempts: int = 3 |
| 114 | + # Native engines can spend a long time loading model weights before their | ||
| 115 | + # HTTP endpoint starts accepting connections. | ||
| 116 | + startup_timeout: int = 1800 | ||
| 114 | npu_usage_threshold: int = 3 | 117 | npu_usage_threshold: int = 3 |
| 115 | enable_virtual_inference: bool = False | 118 | enable_virtual_inference: bool = False |
| 116 | max_failure_count: int = 6 | 119 | max_failure_count: int = 6 |
| @@ -129,6 +132,7 @@ class HealthCheckConfig: | |||
| 129 | self.health_collector_timeout_retry_attempts = self._as_positive_int( | 132 | self.health_collector_timeout_retry_attempts = self._as_positive_int( |
| 130 | "health_collector_timeout_retry_attempts", self.health_collector_timeout_retry_attempts | 133 | "health_collector_timeout_retry_attempts", self.health_collector_timeout_retry_attempts |
| 131 | ) | 134 | ) |
| 135 | + self.startup_timeout = self._as_positive_int("startup_timeout", self.startup_timeout) | ||
| 132 | 136 | ||
| 133 | 137 | ||
| 134 | def from_dict(cls, data: dict[str, Any]) -> "HealthCheckConfig": | 138 | def from_dict(cls, data: dict[str, Any]) -> "HealthCheckConfig": |
| @@ -62,6 +62,8 @@ MULTICONNECTOR = "MultiConnector" | |||
| 62 | KV_CONNECTOR_EXTRA_CONFIG_KEY = "kv_connector_extra_config" | 62 | KV_CONNECTOR_EXTRA_CONFIG_KEY = "kv_connector_extra_config" |
| 63 | CONNECTORS_KEY = "connectors" | 63 | CONNECTORS_KEY = "connectors" |
| 64 | KV_PORT_KEY = "kv_port" | 64 | KV_PORT_KEY = "kv_port" |
| 65 | +DISAGGREGATION_BOOTSTRAP_PORT_KEY = "disaggregation_bootstrap_port" | ||
| 66 | +DISAGGREGATION_BOOTSTRAP_PORT_CLI_KEY = "disaggregation-bootstrap-port" | ||
| 65 | LOOPUP_RPC_PORT_KEY = "lookup_rpc_port" | 67 | LOOPUP_RPC_PORT_KEY = "lookup_rpc_port" |
| 66 | UCM_CONNECTOR = "UCMConnector" | 68 | UCM_CONNECTOR = "UCMConnector" |
| 67 | SERVER_LIST = "server_list" | 69 | SERVER_LIST = "server_list" |
| @@ -154,6 +156,7 @@ class EndpointConfig: | |||
| 154 | base_port: int = 10000 | 156 | base_port: int = 10000 |
| 155 | mgmt_ports: list[str] = field(default_factory=list) | 157 | mgmt_ports: list[str] = field(default_factory=list) |
| 156 | service_ports: list[str] = field(default_factory=list) | 158 | service_ports: list[str] = field(default_factory=list) |
| 159 | + bootstrap_port: int | None = None | ||
| 157 | 160 | ||
| 158 | 161 | ||
| 159 | 162 | ||
| @@ -381,6 +384,7 @@ class NodeManagerConfig: | |||
| 381 | config_data = raw | 384 | config_data = raw |
| 382 | 385 | ||
| 383 | cls._update_from_config_data(config, config_data) | 386 | cls._update_from_config_data(config, config_data) |
| 387 | + cls._set_native_bootstrap_port(config, raw) | ||
| 384 | else: | 388 | else: |
| 385 | logger.warning("Config file does not exist, using default configuration: %s", config_path_obj) | 389 | logger.warning("Config file does not exist, using default configuration: %s", config_path_obj) |
| 386 | 390 | ||
| @@ -506,6 +510,46 @@ class NodeManagerConfig: | |||
| 506 | 510 | ||
| 507 | return config_data | 511 | return config_data |
| 508 | 512 | ||
| 513 | + | ||
| 514 | + def _set_native_bootstrap_port(cls, config: "NodeManagerConfig", user_cfg: dict[str, Any]) -> None: | ||
| 515 | + """Resolve SGLang bootstrap metadata from the selected engine section.""" | ||
| 516 | + engine_key = { | ||
| 517 | + "encode": MOTOR_ENGINE_ENCODE_CONFIG_KEY, | ||
| 518 | + "prefill": MOTOR_ENGINE_PREFILL_CONFIG_KEY, | ||
| 519 | + "decode": MOTOR_ENGINE_DECODE_CONFIG_KEY, | ||
| 520 | + "union": MOTOR_ENGINE_UNION_CONFIG_KEY, | ||
| 521 | + "both": MOTOR_ENGINE_UNION_CONFIG_KEY, | ||
| 522 | + }.get(Env.role) | ||
| 523 | + engine_section = user_cfg.get(engine_key, {}) if engine_key else {} | ||
| 524 | + if not isinstance(engine_section, dict): | ||
| 525 | + return | ||
| 526 | + engine_config = engine_section.get(ENGINE_CONFIG_KEY, {}) | ||
| 527 | + if not isinstance(engine_config, dict): | ||
| 528 | + return | ||
| 529 | + | ||
| 530 | + if str(engine_section.get(ENGINE_TYPE_KEY, "")).strip().lower() != ENGINE_TYPE_SGLANG: | ||
| 531 | + return | ||
| 532 | + port = engine_config.get( | ||
| 533 | + DISAGGREGATION_BOOTSTRAP_PORT_KEY, | ||
| 534 | + engine_config.get(DISAGGREGATION_BOOTSTRAP_PORT_CLI_KEY), | ||
| 535 | + ) | ||
| 536 | + config.endpoint_config.bootstrap_port = cls._parse_native_port( | ||
| 537 | + DISAGGREGATION_BOOTSTRAP_PORT_KEY, | ||
| 538 | + port, | ||
| 539 | + ) | ||
| 540 | + | ||
| 541 | + | ||
| 542 | + def _parse_native_port(field_name: str, value: Any) -> int | None: | ||
| 543 | + if value in (None, ""): | ||
| 544 | + return None | ||
| 545 | + try: | ||
| 546 | + port = int(value) | ||
| 547 | + except (TypeError, ValueError) as exc: | ||
| 548 | + raise ValueError(f"{field_name} must be an integer port") from exc | ||
| 549 | + if not 1 <= port <= 65535: | ||
| 550 | + raise ValueError(f"{field_name} must be in range 1-65535") | ||
| 551 | + return port | ||
| 552 | + | ||
| 509 | 553 | ||
| 510 | def _infer_dispatch_capabilities(cls, engine_config: dict[str, Any]) -> list[str]: | 554 | def _infer_dispatch_capabilities(cls, engine_config: dict[str, Any]) -> list[str]: |
| 511 | """Infer Motor dispatch capabilities from engine-native config.""" | 555 | """Infer Motor dispatch capabilities from engine-native config.""" |
| @@ -521,7 +565,9 @@ class NodeManagerConfig: | |||
| 521 | native_engine_config, | 565 | native_engine_config, |
| 522 | explicit_profile=engine_config.get(DISPATCH_PROFILE_KEY), | 566 | explicit_profile=engine_config.get(DISPATCH_PROFILE_KEY), |
| 523 | ) | 567 | ) |
| 524 | - capabilities = dispatch_capabilities_for_profile(profile) | 568 | + # The native vLLM runtime implements only the explicit handoff contract. |
| 569 | + # Do not advertise trigger/concurrent support that runtime validation rejects. | ||
| 570 | + capabilities = dispatch_capabilities_for_profile(profile) if profile == DispatchProfile.HANDOFF else [] | ||
| 525 | if not capabilities and profile == DispatchProfile.UNKNOWN: | 571 | if not capabilities and profile == DispatchProfile.UNKNOWN: |
| 526 | logger.warning( | 572 | logger.warning( |
| 527 | "Unable to infer vLLM dispatch capability from kv_transfer_config. " | 573 | "Unable to infer vLLM dispatch capability from kv_transfer_config. " |
| @@ -747,6 +793,9 @@ class NodeManagerConfig: | |||
| 747 | if self.basic_config.heartbeat_interval_seconds <= 0: | 793 | if self.basic_config.heartbeat_interval_seconds <= 0: |
| 748 | errors.append("heartbeat_interval_seconds must be greater than 0") | 794 | errors.append("heartbeat_interval_seconds must be greater than 0") |
| 749 | 795 | ||
| 796 | + if self.snapshot_config.enable_snapshot: | ||
G 严重程度： 提示 问题： snapshot 能力被移除（配置期硬校验），但快照相关代码路径未同步清理，成不可达死代码。 原因： 此处新增无条件校验 怎么改： 在同一 PR 内清理死代码（删除 engine_manager 的快照方法及其调用点、node_manager.py:77 的 watcher snapshot 分支、engine_server 的 snapshot_monitor/snapshot_sentinel），或至少将校验错误信息改为明确提示“该能力已在原生引擎运行时移除”并附迁移说明，同时更新 node_manager.md 中 /readiness 段落（仍写着“用于快照默认应用场景”）。 tobking 19 天前 评论： 19 天前 评论： | |||
| 797 | + errors.append("Native engine runtime does not support snapshot yet; enable_snapshot must be false") | ||
| 798 | + | ||
| 750 | # Validate logging configuration | 799 | # Validate logging configuration |
| 751 | valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR"] | 800 | valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR"] |
| 752 | if self.logging_config.log_level.upper() not in valid_log_levels: | 801 | if self.logging_config.log_level.upper() not in valid_log_levels: |
| @@ -489,6 +489,7 @@ class InstanceAssembler(ThreadSafeSingleton): | |||
| 489 | ip=msg.pod_ip, | 489 | ip=msg.pod_ip, |
| 490 | business_port=msg.business_port[0], | 490 | business_port=msg.business_port[0], |
| 491 | mgmt_port=msg.mgmt_port[0], | 491 | mgmt_port=msg.mgmt_port[0], |
| 492 | + bootstrap_port=msg.bootstrap_port, | ||
| 492 | device_infos=device_infos, | 493 | device_infos=device_infos, |
| 493 | ) | 494 | ) |
| 494 | } | 495 | } |
| @@ -533,6 +534,7 @@ class InstanceAssembler(ThreadSafeSingleton): | |||
| 533 | ip=msg.pod_ip, | 534 | ip=msg.pod_ip, |
| 534 | business_port=port, | 535 | business_port=port, |
| 535 | mgmt_port=msg.mgmt_port[i], | 536 | mgmt_port=msg.mgmt_port[i], |
| 537 | + bootstrap_port=msg.bootstrap_port, | ||
| 536 | device_infos=device_infos, | 538 | device_infos=device_infos, |
| 537 | ) | 539 | ) |
| 538 | logger.debug("Built %d endpoints for pod %s", len(pod_endpoints), msg.pod_ip) | 540 | logger.debug("Built %d endpoints for pod %s", len(pod_endpoints), msg.pod_ip) |
Rmotor/coordinator/api_client/engine_server_api_client.py→motor/coordinator/api_client/native_engine_api_client.py+17-35
| @@ -1,5 +1,4 @@ | |||
| 1 | -# -*- coding: utf-8 -*- | 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. |
| 2 | -# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 3 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 4 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| 5 | # You may obtain a copy of Mulan PSL v2 at: | 4 | # You may obtain a copy of Mulan PSL v2 at: |
| @@ -11,49 +10,32 @@ | |||
| 11 | 10 | ||
| 12 | from motor.common.http.http_client import SafeHTTPSClient | 11 | from motor.common.http.http_client import SafeHTTPSClient |
| 13 | from motor.common.logger import get_logger | 12 | from motor.common.logger import get_logger |
| 14 | -from motor.common.logger.rate_limited_logger import RateLimitedLogger | 13 | +from motor.config.tls_config import TLSConfig |
| 15 | -from motor.config.coordinator import CoordinatorConfig | ||
| 16 | 14 | ||
| 17 | logger = get_logger(__name__) | 15 | logger = get_logger(__name__) |
| 18 | -_rl = RateLimitedLogger(logger) | ||
| 19 | 16 | ||
| 20 | 17 | ||
| 21 | -class EngineServerApiClient: | 18 | +class NativeEngineApiClient: |
| 22 | - tls_config = CoordinatorConfig.from_json().mgmt_tls_config | 19 | + """Read native engine operational endpoints over the inference channel.""" |
| 23 | 20 | ||
| 24 | 21 | ||
| 25 | - def query_metrics(address: str): | 22 | + def query_metrics(address: str, tls_config: TLSConfig | None) -> str: |
G 严重程度： 提示 问题： 原因： 基线 engine_server_api_client 的 query_metrics 明确检查 怎么改：
恢复状态码检查： tobking 19 天前 评论： 19 天前 评论： | |||
| 26 | - client_args = EngineServerApiClient._generate_client_args(address) | ||
| 27 | try: | 23 | try: |
| 28 | - client = SafeHTTPSClient(timeout=2, **client_args) | 24 | + with SafeHTTPSClient(address=address, tls_config=tls_config, timeout=2) as client: |
| 29 | - response = client.do_get("/metrics") | 25 | + response = client.do_get("/metrics") |
| 30 | - if response.status_code == 200: | 26 | + if response.status_code != 200: |
| 31 | - data = response.text | ||
| 32 | - return data | ||
| 33 | - else: | ||
| 34 | logger.warning( | 27 | logger.warning( |
| 35 | - "Coordinator->EngineServer query_metrics non-2xx. " | 28 | + "Coordinator native metrics request returned status=%s. address=%s", |
| 36 | - "address=%s, status_code=%s. " | ||
| 37 | - "Possible causes: 1) engine_server not ready 2) wrong endpoint 3) auth failure.", | ||
| 38 | - address, | ||
| 39 | response.status_code, | 29 | response.status_code, |
| 30 | + address, | ||
| 40 | ) | 31 | ) |
| 41 | - except Exception as e: | 32 | + return "" |
| 33 | + return response.text | ||
| 34 | + except Exception as err: | ||
| 42 | logger.warning( | 35 | logger.warning( |
| 43 | - "Coordinator->EngineServer query_metrics failed. address=%s, error=%s. " | 36 | + "Coordinator native metrics request failed. address=%s, error=%s. " |
| 44 | - "Possible causes: 1) engine_server down 2) network unreachable 3) tls mismatch. " | 37 | + "Check native engine readiness, network reachability, and inference TLS configuration.", |
| 45 | - "Check: ping %s, engine_server process status.", | ||
| 46 | - address, | ||
| 47 | - e, | ||
| 48 | address, | 38 | address, |
| 39 | + err, | ||
| 49 | ) | 40 | ) |
| 50 | - | 41 | + return "" |
| 51 | - return "" | ||
| 52 | - | ||
| 53 | - | ||
| 54 | - def _generate_client_args(cls, address) -> dict[str, str]: | ||
| 55 | - client_ars = { | ||
| 56 | - "address": f"{address}", | ||
| 57 | - "tls_config": cls.tls_config, | ||
| 58 | - } | ||
| 59 | - return client_ars | ||
| @@ -501,6 +501,7 @@ class InferenceServer(BaseCoordinatorServer): | |||
| 501 | self.coordinator_config, | 501 | self.coordinator_config, |
| 502 | scheduler=self._get_scheduler_client(), | 502 | scheduler=self._get_scheduler_client(), |
| 503 | request_manager=request_manager, | 503 | request_manager=request_manager, |
| 504 | + request_json=body_json, | ||
| 504 | ) | 505 | ) |
| 505 | except HTTPException: | 506 | except HTTPException: |
| 506 | raise | 507 | raise |
| @@ -531,6 +532,7 @@ class InferenceServer(BaseCoordinatorServer): | |||
| 531 | self.coordinator_config, | 532 | self.coordinator_config, |
| 532 | scheduler=self._get_scheduler_client(), | 533 | scheduler=self._get_scheduler_client(), |
| 533 | request_manager=request_manager, | 534 | request_manager=request_manager, |
| 535 | + request_json=body_json, | ||
| 534 | ) | 536 | ) |
| 535 | except HTTPException: | 537 | except HTTPException: |
| 536 | raise | 538 | raise |
| @@ -13,7 +13,6 @@ from typing import Mapping | |||
| 13 | 13 | ||
| 14 | from motor.common.logger import get_logger | 14 | from motor.common.logger import get_logger |
| 15 | from motor.common.logger.rate_limited_logger import RateLimitedLogger | 15 | from motor.common.logger.rate_limited_logger import RateLimitedLogger |
| 16 | -from motor.common.resources.dispatch import has_compatible_dispatch_pair | ||
| 17 | from motor.common.resources.instance import Instance, PDRole, Workload, Endpoint | 16 | from motor.common.resources.instance import Instance, PDRole, Workload, Endpoint |
| 18 | from motor.common.resources.http_msg_spec import EventType | 17 | from motor.common.resources.http_msg_spec import EventType |
| 19 | from motor.config.coordinator import CoordinatorConfig | 18 | from motor.config.coordinator import CoordinatorConfig |
| @@ -361,49 +360,8 @@ class InstanceManager: | |||
| 361 | len(self._decode_pool), | 360 | len(self._decode_pool), |
| 362 | len(self._hybrid_pool), | 361 | len(self._hybrid_pool), |
| 363 | ) | 362 | ) |
| 364 | - self._log_dispatch_capabilities() | ||
| 365 | return result | 363 | return result |
| 366 | 364 | ||
| 367 | - def _log_dispatch_capabilities(self) -> None: | ||
| 368 | - """Log per-instance dispatch_capabilities and flag an incompatible P/D pool. | ||
| 369 | - | ||
| 370 | - Readiness gates on a shared P/D dispatch capability; when P and D are both present | ||
| 371 | - but advertise no common capability the service stays not-ready with a vague | ||
| 372 | - instances_status=unknown. This makes the actual capabilities (empty vs mismatched) visible. | ||
| 373 | - """ | ||
| 374 | - | ||
| 375 | - def _role(inst: Instance) -> str: | ||
| 376 | - role = getattr(inst, "role", None) | ||
| 377 | - return role.value if hasattr(role, "value") else str(role) | ||
| 378 | - | ||
| 379 | - summary = ", ".join( | ||
| 380 | - f"{inst.id}({_role(inst)})={list(getattr(inst, 'dispatch_capabilities', []) or [])}" | ||
| 381 | - for pool in (self._prefill_pool, self._decode_pool, self._encode_pool, self._hybrid_pool) | ||
| 382 | - for inst in pool.values() | ||
| 383 | - ) | ||
| 384 | - logger.info("Instance dispatch_capabilities: %s", summary or "(none)") | ||
| 385 | - | ||
| 386 | - if ( | ||
| 387 | - self._prefill_pool | ||
| 388 | - and self._decode_pool | ||
| 389 | - and not has_compatible_dispatch_pair(self._prefill_pool.values(), self._decode_pool.values()) | ||
| 390 | - ): | ||
| 391 | - if self._hybrid_pool: | ||
| 392 | - logger.warning( | ||
| 393 | - "P/D instances are online but advertise no shared dispatch capability; " | ||
| 394 | - "requests will fall back to PDHybridRouter via union instances. " | ||
| 395 | - "Check the engine kv_connector is recognized or set dispatch_profile explicitly. " | ||
| 396 | - "capabilities: %s", | ||
| 397 | - summary or "(none)", | ||
| 398 | - ) | ||
| 399 | - else: | ||
| 400 | - logger.warning( | ||
| 401 | - "P/D instances are online but advertise no shared dispatch capability " | ||
| 402 | - "(readiness will report instances_status=unknown). Check the engine kv_connector " | ||
| 403 | - "is recognized or set dispatch_profile explicitly. capabilities: %s", | ||
| 404 | - summary or "(none)", | ||
| 405 | - ) | ||
| 406 | - | ||
| 407 | def _find_available_pool(self, instance_id: int) -> dict[int, Instance] | None: | 365 | def _find_available_pool(self, instance_id: int) -> dict[int, Instance] | None: |
| 408 | # This is a private method that should only be called within locked contexts | 366 | # This is a private method that should only be called within locked contexts |
| 409 | instance = self._available_pool.get(instance_id) | 367 | instance = self._available_pool.get(instance_id) |
| @@ -19,7 +19,6 @@ from typing import Iterable, Protocol | |||
| 19 | 19 | ||
| 20 | from pydantic import BaseModel | 20 | from pydantic import BaseModel |
| 21 | 21 | ||
| 22 | -from motor.common.resources.dispatch import has_compatible_dispatch_pair | ||
| 23 | from motor.common.resources.endpoint import Endpoint, Workload, WorkloadAction | 22 | from motor.common.resources.endpoint import Endpoint, Workload, WorkloadAction |
| 24 | from motor.common.resources.instance import Instance, PDRole | 23 | from motor.common.resources.instance import Instance, PDRole |
| 25 | from motor.coordinator.models.request import RequestInfo | 24 | from motor.coordinator.models.request import RequestInfo |
| @@ -53,7 +52,7 @@ class InstanceReadiness(str, Enum): | |||
| 53 | 52 | ||
| 54 | 53 | ||
| 55 | def readiness_from_instances(instances: Iterable[Instance]) -> InstanceReadiness: | 54 | def readiness_from_instances(instances: Iterable[Instance]) -> InstanceReadiness: |
| 56 | - """Infer readiness from available roles and compatible P/D dispatch capabilities.""" | 55 | + """Infer readiness from the available native-engine roles.""" |
| 57 | encode_instances = [] | 56 | encode_instances = [] |
| 58 | prefill_instances = [] | 57 | prefill_instances = [] |
| 59 | decode_instances = [] | 58 | decode_instances = [] |
| @@ -70,13 +69,10 @@ def readiness_from_instances(instances: Iterable[Instance]) -> InstanceReadiness | |||
| 70 | elif role_value in (PDRole.ROLE_U.value, "both", "hybrid"): | 69 | elif role_value in (PDRole.ROLE_U.value, "both", "hybrid"): |
| 71 | union_instances.append(instance) | 70 | union_instances.append(instance) |
| 72 | 71 | ||
| 73 | - has_compatible_pd = has_compatible_dispatch_pair(prefill_instances, decode_instances) | 72 | + if prefill_instances and decode_instances: |
| 74 | - if has_compatible_pd: | ||
| 75 | return InstanceReadiness.REQUIRED_MET_EPD if encode_instances else InstanceReadiness.REQUIRED_MET | 73 | return InstanceReadiness.REQUIRED_MET_EPD if encode_instances else InstanceReadiness.REQUIRED_MET |
| 76 | if union_instances: | 74 | if union_instances: |
| 77 | return InstanceReadiness.REQUIRED_MET | 75 | return InstanceReadiness.REQUIRED_MET |
| 78 | - if prefill_instances and decode_instances: | ||
| 79 | - return InstanceReadiness.UNKNOWN | ||
| 80 | if encode_instances and prefill_instances: | 76 | if encode_instances and prefill_instances: |
| 81 | return InstanceReadiness.ENCODE_PREFILL | 77 | return InstanceReadiness.ENCODE_PREFILL |
| 82 | if prefill_instances: | 78 | if prefill_instances: |
| @@ -126,10 +122,12 @@ class SchedulingFacade(Protocol): | |||
| 126 | req_info: RequestInfo, | 122 | req_info: RequestInfo, |
| 127 | *, | 123 | *, |
| 128 | target_instance_id: int | None = None, | 124 | target_instance_id: int | None = None, |
| 125 | + required_engine_type: str | None = None, | ||
| 129 | ) -> tuple[Instance, Endpoint, Workload] | None: | 126 | ) -> tuple[Instance, Endpoint, Workload] | None: |
| 130 | """ | 127 | """ |
| 131 | Atomic: select instance + one workload allocation (ALLOCATION). | 128 | Atomic: select instance + one workload allocation (ALLOCATION). |
| 132 | When target_instance_id is set, pin to that instance (skip policy selection). | 129 | When target_instance_id is set, pin to that instance (skip policy selection). |
| 130 | + When required_engine_type is set, only matching engine instances are eligible. | ||
| 133 | Returns (instance, endpoint, allocation_workload). Caller records allocation_workload for release. | 131 | Returns (instance, endpoint, allocation_workload). Caller records allocation_workload for release. |
| 134 | """ | 132 | """ |
| 135 | ... | 133 | ... |
| @@ -153,10 +151,6 @@ class SchedulingFacade(Protocol): | |||
| 153 | """Return roles currently present in the scheduler's local instance view.""" | 151 | """Return roles currently present in the scheduler's local instance view.""" |
| 154 | ... | 152 | ... |
| 155 | 153 | ||
| 156 | - async def has_compatible_pd_pair(self) -> bool: | ||
| 157 | - """Return whether the local scheduler view contains a compatible P/D pair.""" | ||
| 158 | - ... | ||
| 159 | - | ||
| 160 | async def report_cb_event(self, instance_id: int, event: str) -> None: | 154 | async def report_cb_event(self, instance_id: int, event: str) -> None: |
| 161 | """Report a circuit-breaker event ("failure" | "success") for an instance. | 155 | """Report a circuit-breaker event ("failure" | "success") for an instance. |
| 162 | 156 | ||
| @@ -27,7 +27,7 @@ if TYPE_CHECKING: | |||
| 27 | logger = get_logger(__name__) | 27 | logger = get_logger(__name__) |
| 28 | 28 | ||
| 29 | # (p_instance_id or None, d_instance_id) | 29 | # (p_instance_id or None, d_instance_id) |
| 30 | -# p_instance_id 为 None 表示 CDP/PD_SEPARATE 模式下 coordinator 侧不感知 P 实例 | 30 | +# p_instance_id 为 None 表示请求由单个 Hybrid/Union 实例完成。 |
| 31 | PDGroupKey = tuple[int | None, int] | 31 | PDGroupKey = tuple[int | None, int] |
| 32 | 32 | ||
| 33 | 33 | ||
| @@ -20,9 +20,6 @@ Unknown metrics fall back to type-based defaults (gauge→sum, counter→sum, hi | |||
| 20 | from dataclasses import dataclass | 20 | from dataclasses import dataclass |
| 21 | from enum import Enum | 21 | from enum import Enum |
| 22 | 22 | ||
| 23 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 24 | - | ||
| 25 | - | ||
| 26 | # --------------------------------------------------------------------------- | 23 | # --------------------------------------------------------------------------- |
| 27 | # Semantic taxonomy | 24 | # Semantic taxonomy |
| 28 | # --------------------------------------------------------------------------- | 25 | # --------------------------------------------------------------------------- |
| @@ -257,24 +254,11 @@ class MetricRegistry: | |||
| 257 | def get_effective_role_scope( | 254 | def get_effective_role_scope( |
| 258 | cls, | 255 | cls, |
| 259 | metric_name: str, | 256 | metric_name: str, |
| 260 | - dispatch_capabilities: set[str] | None = None, | ||
| 261 | ) -> str | None: | 257 | ) -> str | None: |
| 262 | - """Get the effective role scope for a metric, considering connector capabilities. | 258 | + """Get the configured role scope for a metric.""" |
| 263 | - | ||
| 264 | - Handoff connectors expose meaningful TTFT on both P and D instances. | ||
| 265 | - Concurrent connectors use D's TTFT as the authoritative service value. | ||
| 266 | - """ | ||
| 267 | config = cls.get_semantic(metric_name) | 259 | config = cls.get_semantic(metric_name) |
| 268 | if config is None or config.role_scope is None: | 260 | if config is None or config.role_scope is None: |
| 269 | return None | 261 | return None |
| 270 | - | ||
| 271 | - if ( | ||
| 272 | - metric_name == "vllm:time_to_first_token_seconds" | ||
| 273 | - and dispatch_capabilities | ||
| 274 | - and DispatchPlan.PREFILL_HANDOFF_DECODE.value in dispatch_capabilities | ||
| 275 | - ): | ||
| 276 | - return None | ||
| 277 | - | ||
| 278 | return config.role_scope | 262 | return config.role_scope |
| 279 | 263 | ||
| 280 | 264 | ||
| @@ -35,7 +35,7 @@ class AggregationScope(Enum): | |||
| 35 | class AggregationContext: | 35 | class AggregationContext: |
| 36 | scope: AggregationScope = AggregationScope.INSTANCE | 36 | scope: AggregationScope = AggregationScope.INSTANCE |
| 37 | instance_roles: dict[int, str] | None = None | 37 | instance_roles: dict[int, str] | None = None |
| 38 | - instance_dispatch_capabilities: dict[int, set[str]] | None = None | 38 | + instance_engine_types: dict[int, str] | None = None |
| 39 | ins_ids: list[int] | None = None | 39 | ins_ids: list[int] | None = None |
| 40 | 40 | ||
| 41 | 41 | ||
| @@ -17,14 +17,13 @@ from collections.abc import Callable | |||
| 17 | from typing import Any | 17 | from typing import Any |
| 18 | import requests | 18 | import requests |
| 19 | 19 | ||
| 20 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 21 | from motor.common.resources.instance import Instance | 20 | from motor.common.resources.instance import Instance |
| 22 | from motor.common.logger import get_logger | 21 | from motor.common.logger import get_logger |
| 23 | from motor.common.logger.rate_limited_logger import RateLimitedLogger | 22 | from motor.common.logger.rate_limited_logger import RateLimitedLogger |
| 24 | from motor.common.utils.net import format_address | 23 | from motor.common.utils.net import format_address |
| 25 | from motor.common.utils.singleton import ThreadSafeSingleton | 24 | from motor.common.utils.singleton import ThreadSafeSingleton |
| 26 | from motor.config.coordinator import CoordinatorConfig | 25 | from motor.config.coordinator import CoordinatorConfig |
| 27 | -from motor.coordinator.api_client.engine_server_api_client import EngineServerApiClient | 26 | +from motor.coordinator.api_client.native_engine_api_client import NativeEngineApiClient |
| 28 | from motor.coordinator.metrics.metric_types import ( | 27 | from motor.coordinator.metrics.metric_types import ( |
| 29 | AggregationContext, | 28 | AggregationContext, |
| 30 | AggregationScope, | 29 | AggregationScope, |
| @@ -243,6 +242,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 243 | config = CoordinatorConfig() | 242 | config = CoordinatorConfig() |
| 244 | self._prometheus_metrics_config = config.prometheus_metrics_config | 243 | self._prometheus_metrics_config = config.prometheus_metrics_config |
| 245 | self._deploy_config = config.deploy_config | 244 | self._deploy_config = config.deploy_config |
| 245 | + self._infer_tls_config = config.infer_tls_config | ||
| 246 | 246 | ||
| 247 | # Initial metrics state | 247 | # Initial metrics state |
| 248 | self._inactive_instance_metrics_aggregate: dict[str, list[Metric]] = {} | 248 | self._inactive_instance_metrics_aggregate: dict[str, list[Metric]] = {} |
| @@ -301,6 +301,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 301 | with self._config_lock: | 301 | with self._config_lock: |
| 302 | self._prometheus_metrics_config = config.prometheus_metrics_config | 302 | self._prometheus_metrics_config = config.prometheus_metrics_config |
| 303 | self._deploy_config = config.deploy_config | 303 | self._deploy_config = config.deploy_config |
| 304 | + self._infer_tls_config = config.infer_tls_config | ||
| 304 | logger.info("MetricsCollector configuration updated") | 305 | logger.info("MetricsCollector configuration updated") |
| 305 | 306 | ||
| 306 | def get_metrics( | 307 | def get_metrics( |
| @@ -845,7 +846,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 845 | collect["role"] = ins_info.role | 846 | collect["role"] = ins_info.role |
| 846 | collect["job_name"] = ins_info.job_name | 847 | collect["job_name"] = ins_info.job_name |
| 847 | collect["model_name"] = ins_info.model_name | 848 | collect["model_name"] = ins_info.model_name |
| 848 | - collect["dispatch_capabilities"] = list(ins_info.dispatch_capabilities or []) | 849 | + collect["engine_type"] = ins_info.engine_type |
| 849 | collects[ins_info.id] = collect | 850 | collects[ins_info.id] = collect |
| 850 | 851 | ||
| 851 | return collects | 852 | return collects |
| @@ -868,7 +869,10 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 868 | collect = {"endpoints": {}} | 869 | collect = {"endpoints": {}} |
| 869 | 870 | ||
| 870 | for en_info in ins_info.get_all_endpoints(): | 871 | for en_info in ins_info.get_all_endpoints(): |
| 871 | - metrics_str = EngineServerApiClient.query_metrics(f"{en_info.ip}:{en_info.mgmt_port}") | 872 | + metrics_str = NativeEngineApiClient.query_metrics( |
| 873 | + format_address(en_info.ip, en_info.business_port), | ||
| 874 | + self._infer_tls_config, | ||
| 875 | + ) | ||
| 872 | if not metrics_str: | 876 | if not metrics_str: |
| 873 | return {} | 877 | return {} |
| 874 | collect["endpoints"][en_info.id] = { | 878 | collect["endpoints"][en_info.id] = { |
| @@ -882,7 +886,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 882 | self, | 886 | self, |
| 883 | collects: dict[int, dict[str, Any]], | 887 | collects: dict[int, dict[str, Any]], |
| 884 | instance_roles: dict[int, str], | 888 | instance_roles: dict[int, str], |
| 885 | - instance_dispatch_capabilities: dict[int, set[str]], | 889 | + instance_engine_types: dict[int, str], |
| 886 | ) -> list[Metric]: | 890 | ) -> list[Metric]: |
| 887 | """Aggreagte metrics of all instances.""" | 891 | """Aggreagte metrics of all instances.""" |
| 888 | 892 | ||
| @@ -909,7 +913,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 909 | scope=AggregationScope.SERVICE, | 913 | scope=AggregationScope.SERVICE, |
| 910 | ins_ids=ins_ids, | 914 | ins_ids=ins_ids, |
| 911 | instance_roles=instance_roles, | 915 | instance_roles=instance_roles, |
| 912 | - instance_dispatch_capabilities=instance_dispatch_capabilities, | 916 | + instance_engine_types=instance_engine_types, |
| 913 | ) | 917 | ) |
| 914 | return self._aggregate_metrics(aggr_input, ctx=ctx) | 918 | return self._aggregate_metrics(aggr_input, ctx=ctx) |
| 915 | 919 | ||
| @@ -943,8 +947,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 943 | result: list[Metric] = [] | 947 | result: list[Metric] = [] |
| 944 | for name, entries in aggr_input.items(): | 948 | for name, entries in aggr_input.items(): |
| 945 | if ctx is not None and ctx.scope == AggregationScope.SERVICE and ctx.instance_roles is not None: | 949 | if ctx is not None and ctx.scope == AggregationScope.SERVICE and ctx.instance_roles is not None: |
| 946 | - if name == "vllm:time_to_first_token_seconds" and ctx.instance_dispatch_capabilities is not None: | 950 | + if name == "vllm:time_to_first_token_seconds" and ctx.instance_engine_types is not None: |
| 947 | - handoff = DispatchPlan.PREFILL_HANDOFF_DECODE.value | ||
| 948 | entries = [ | 951 | entries = [ |
| 949 | (ins_id, m) | 952 | (ins_id, m) |
| 950 | for ins_id, m in entries | 953 | for ins_id, m in entries |
| @@ -953,7 +956,7 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 953 | or ctx.instance_roles.get(ins_id) in {"decode", "union", "both", "hybrid"} | 956 | or ctx.instance_roles.get(ins_id) in {"decode", "union", "both", "hybrid"} |
| 954 | or ( | 957 | or ( |
| 955 | ctx.instance_roles.get(ins_id) == "prefill" | 958 | ctx.instance_roles.get(ins_id) == "prefill" |
| 956 | - and handoff in ctx.instance_dispatch_capabilities.get(ins_id, set()) | 959 | + and ctx.instance_engine_types.get(ins_id, "").strip().lower() == "vllm" |
| 957 | ) | 960 | ) |
| 958 | ) | 961 | ) |
| 959 | ] | 962 | ] |
| @@ -993,13 +996,11 @@ class MetricsCollector(ThreadSafeSingleton): | |||
| 993 | for ins_id, metrics_list in instance_metrics.items(): | 996 | for ins_id, metrics_list in instance_metrics.items(): |
| 994 | self._instance_metrics_cached[ins_id] = {self.METRICS_KEY: metrics_list} | 997 | self._instance_metrics_cached[ins_id] = {self.METRICS_KEY: metrics_list} |
| 995 | instance_roles = {ins_id: data.get("role", "") for ins_id, data in collects.items() if isinstance(data, dict)} | 998 | instance_roles = {ins_id: data.get("role", "") for ins_id, data in collects.items() if isinstance(data, dict)} |
| 996 | - instance_dispatch_capabilities = { | 999 | + instance_engine_types = { |
| 997 | - ins_id: set(data.get("dispatch_capabilities", [])) | 1000 | + ins_id: str(data.get("engine_type", "")) for ins_id, data in collects.items() if isinstance(data, dict) |
| 998 | - for ins_id, data in collects.items() | ||
| 999 | - if isinstance(data, dict) | ||
| 1000 | } | 1001 | } |
| 1001 | aggregate = self._aggregation_engine.post_process( | 1002 | aggregate = self._aggregation_engine.post_process( |
| 1002 | - self._aggregate_metrics_all_instance(collects, instance_roles, instance_dispatch_capabilities) | 1003 | + self._aggregate_metrics_all_instance(collects, instance_roles, instance_engine_types) |
| 1003 | ) | 1004 | ) |
| 1004 | with self._config_lock: | 1005 | with self._config_lock: |
| 1005 | deploy_config = self._deploy_config | 1006 | deploy_config = self._deploy_config |
| @@ -92,9 +92,6 @@ class RequestInfo(BaseModel): | |||
| 92 | prompt_tokens_details: dict = Field(default={}, description="prefill prompt_tokens_details") | 92 | prompt_tokens_details: dict = Field(default={}, description="prefill prompt_tokens_details") |
| 93 | prompt_token_ids: list = Field(default=[], description="prefill prompt_token_ids") | 93 | prompt_token_ids: list = Field(default=[], description="prefill prompt_token_ids") |
| 94 | cached_token_ids: list = Field(default=[], description="Cached token_ids") | 94 | cached_token_ids: list = Field(default=[], description="Cached token_ids") |
| 95 | - p_instance_id: int | None = Field( | ||
| 96 | - default=None, description="P instance ID set by metaserver handler in CDP-like modes" | ||
| 97 | - ) | ||
| 98 | scheduling_constraint: SchedulingConstraint | None = Field( | 95 | scheduling_constraint: SchedulingConstraint | None = Field( |
| 99 | default=None, | 96 | default=None, |
| 100 | description="Internal pin-to-instance constraint (e.g. precision probe); not from client API", | 97 | description="Internal pin-to-instance constraint (e.g. precision probe); not from client API", |
| @@ -1,9 +1,20 @@ | |||
| 1 | -# -*- coding: utf-8 -*- | ||
| 2 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 3 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 4 | 3 | ||
| 5 | """Coordinator HTTP routing layer: dispatch entry and strategy implementations.""" | 4 | """Coordinator HTTP routing layer: dispatch entry and strategy implementations.""" |
| 6 | 5 | ||
| 6 | +from typing import TYPE_CHECKING | ||
| 7 | + | ||
| 8 | +if TYPE_CHECKING: | ||
| 9 | + from motor.coordinator.router.dispatch import handle_request | ||
| 10 | + | ||
| 7 | __all__ = ["handle_request"] | 11 | __all__ = ["handle_request"] |
| 8 | 12 | ||
| 9 | -from motor.coordinator.router.dispatch import handle_request | 13 | + |
| 14 | +def __getattr__(name: str): | ||
| 15 | + """Load the HTTP entry lazily so adapter imports do not initialize the full domain graph.""" | ||
| 16 | + if name == "handle_request": | ||
| 17 | + from motor.coordinator.router.dispatch import handle_request | ||
| 18 | + | ||
| 19 | + return handle_request | ||
| 20 | + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") | ||
| @@ -1,9 +1,23 @@ | |||
| 1 | -# -*- coding: utf-8 -*- | ||
| 2 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 3 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 4 | 3 | ||
| 5 | """Response format adapters (OpenAI Completion <-> Chat).""" | 4 | """Response format adapters (OpenAI Completion <-> Chat).""" |
| 6 | 5 | ||
| 6 | +from typing import TYPE_CHECKING | ||
| 7 | + | ||
| 8 | +if TYPE_CHECKING: | ||
| 9 | + from motor.coordinator.router.adapters.completion_to_chat import ( | ||
| 10 | + adapt_completion_nonstream_to_chat, | ||
| 11 | + adapt_completion_stream_chunk_to_chat, | ||
| 12 | + is_completion_like_stream_chunk, | ||
| 13 | + ) | ||
| 14 | + from motor.coordinator.router.adapters.stream import ( | ||
| 15 | + encode_stream_chunk_bytes, | ||
| 16 | + parse_stream_chunk_json, | ||
| 17 | + strip_nonstream_response_body_for_client, | ||
| 18 | + strip_stream_chunk_bytes_for_client, | ||
| 19 | + ) | ||
| 20 | + | ||
| 7 | __all__ = [ | 21 | __all__ = [ |
| 8 | "adapt_completion_nonstream_to_chat", | 22 | "adapt_completion_nonstream_to_chat", |
| 9 | "adapt_completion_stream_chunk_to_chat", | 23 | "adapt_completion_stream_chunk_to_chat", |
| @@ -14,15 +28,28 @@ __all__ = [ | |||
| 14 | "strip_stream_chunk_bytes_for_client", | 28 | "strip_stream_chunk_bytes_for_client", |
| 15 | ] | 29 | ] |
| 16 | 30 | ||
| 17 | -from motor.coordinator.router.adapters.completion_to_chat import ( | ||
| 18 | - adapt_completion_nonstream_to_chat, | ||
| 19 | - adapt_completion_stream_chunk_to_chat, | ||
| 20 | - is_completion_like_stream_chunk, | ||
| 21 | -) | ||
| 22 | 31 | ||
| 23 | -from motor.coordinator.router.adapters.stream import ( | 32 | +_COMPLETION_EXPORTS = { |
| 24 | - encode_stream_chunk_bytes, | 33 | + "adapt_completion_nonstream_to_chat", |
| 25 | - parse_stream_chunk_json, | 34 | + "adapt_completion_stream_chunk_to_chat", |
| 26 | - strip_nonstream_response_body_for_client, | 35 | + "is_completion_like_stream_chunk", |
| 27 | - strip_stream_chunk_bytes_for_client, | 36 | +} |
| 28 | -) | 37 | +_STREAM_EXPORTS = { |
| 38 | + "encode_stream_chunk_bytes", | ||
| 39 | + "parse_stream_chunk_json", | ||
| 40 | + "strip_nonstream_response_body_for_client", | ||
| 41 | + "strip_stream_chunk_bytes_for_client", | ||
| 42 | +} | ||
| 43 | + | ||
| 44 | + | ||
| 45 | +def __getattr__(name: str): | ||
| 46 | + """Resolve adapter exports lazily to keep protocol-only imports dependency-free.""" | ||
| 47 | + if name in _COMPLETION_EXPORTS: | ||
| 48 | + from motor.coordinator.router.adapters import completion_to_chat | ||
| 49 | + | ||
| 50 | + return getattr(completion_to_chat, name) | ||
| 51 | + if name in _STREAM_EXPORTS: | ||
| 52 | + from motor.coordinator.router.adapters import stream | ||
| 53 | + | ||
| 54 | + return getattr(stream, name) | ||
| 55 | + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") | ||
| @@ -1,4 +1,3 @@ | |||
| 1 | -# -*- coding: utf-8 -*- | ||
| 2 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 3 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 4 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| @@ -30,6 +29,14 @@ def _chat_completion_id(req_id: str) -> str: | |||
| 30 | return f"chatcmpl-{base}" | 29 | return f"chatcmpl-{base}" |
| 31 | 30 | ||
| 32 | 31 | ||
| 32 | +def is_completion_like_body(body: dict[str, Any]) -> bool: | ||
| 33 | + """Return whether a non-stream response has the OpenAI Completion shape.""" | ||
| 34 | + if body.get("object") == "text_completion": | ||
| 35 | + return True | ||
| 36 | + choices = body.get(OpenAIField.CHOICES) or [] | ||
| 37 | + return bool(choices and isinstance(choices[0], dict) and OpenAIField.TEXT in choices[0]) | ||
| 38 | + | ||
| 39 | + | ||
| 33 | def is_completion_like_stream_chunk(chunk_json: dict[str, Any]) -> bool: | 40 | def is_completion_like_stream_chunk(chunk_json: dict[str, Any]) -> bool: |
| 34 | """True if chunk looks like a text_completion stream object (not chat chunk).""" | 41 | """True if chunk looks like a text_completion stream object (not chat chunk).""" |
| 35 | if chunk_json.get("object") == "text_completion": | 42 | if chunk_json.get("object") == "text_completion": |
| @@ -0,0 +1,261 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""Stateless native-engine protocol mapping for PD requests.""" | ||
| 12 | + | ||
| 13 | +import hashlib | ||
| 14 | +from collections.abc import Mapping | ||
| 15 | +from copy import deepcopy | ||
| 16 | +from dataclasses import dataclass | ||
| 17 | +from enum import Enum | ||
| 18 | +from types import MappingProxyType | ||
| 19 | +from typing import Any, Protocol | ||
| 20 | + | ||
| 21 | + | ||
| 22 | +class CoordinationMode(str, Enum): | ||
| 23 | + HANDOFF = "handoff" | ||
| 24 | + BOOTSTRAP = "bootstrap" | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +class EngineEndpointMetadata: | ||
| 29 | + host: str | ||
| 30 | + bootstrap_port: int | None = None | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +class LegContext: | ||
| 35 | + engine_request_id: str | ||
| 36 | + pair_id: str | ||
| 37 | + attempt_seq: int | ||
| 38 | + api: str | ||
| 39 | + endpoint: EngineEndpointMetadata | ||
| 40 | + peer_endpoint: EngineEndpointMetadata | None = None | ||
| 41 | + | ||
| 42 | + | ||
| 43 | + | ||
| 44 | +class EngineRequest: | ||
| 45 | + api: str | ||
| 46 | + body: dict[str, Any] | ||
| 47 | + | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +class PrefillMetadata: | ||
| 51 | + handoff_ticket: dict[str, Any] | None = None | ||
| 52 | + usage: dict[str, Any] | None = None | ||
| 53 | + | ||
| 54 | + | ||
| 55 | +class EngineProtocolError(RuntimeError): | ||
| 56 | + """A native engine response or request cannot satisfy the PD contract.""" | ||
| 57 | + | ||
| 58 | + def __init__(self, *, engine_type: str, phase: str, message: str) -> None: | ||
| 59 | + self.engine_type = engine_type | ||
| 60 | + self.phase = phase | ||
| 61 | + self.message = message | ||
| 62 | + super().__init__(f"{engine_type} {phase} protocol error: {message}") | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +class PDProtocolAdapter(Protocol): | ||
| 66 | + engine_type: str | ||
| 67 | + coordination_mode: CoordinationMode | ||
| 68 | + internal_response_fields: frozenset[str] | ||
| 69 | + | ||
| 70 | + def build_prefill_request( | ||
| 71 | + self, | ||
| 72 | + request: Mapping[str, Any], | ||
| 73 | + context: LegContext, | ||
| 74 | + ) -> EngineRequest: ... | ||
| 75 | + | ||
| 76 | + def parse_prefill_response( | ||
| 77 | + self, | ||
| 78 | + response: dict[str, Any], | ||
| 79 | + ) -> PrefillMetadata: ... | ||
| 80 | + | ||
| 81 | + def build_decode_request( | ||
| 82 | + self, | ||
| 83 | + request: Mapping[str, Any], | ||
| 84 | + context: LegContext, | ||
| 85 | + prefill: PrefillMetadata | None, | ||
| 86 | + ) -> EngineRequest: ... | ||
| 87 | + | ||
| 88 | + def inject_request_id(self, body: dict[str, Any], request_id: str) -> None: ... | ||
| 89 | + | ||
| 90 | + def build_abort_request(self, context: LegContext) -> EngineRequest | None: ... | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +class VllmProtocolAdapter: | ||
| 94 | + engine_type = "vllm" | ||
| 95 | + coordination_mode = CoordinationMode.HANDOFF | ||
| 96 | + internal_response_fields = frozenset({"kv_transfer_params"}) | ||
| 97 | + | ||
| 98 | + def build_prefill_request( | ||
| 99 | + self, | ||
| 100 | + request: Mapping[str, Any], | ||
| 101 | + context: LegContext, | ||
| 102 | + ) -> EngineRequest: | ||
| 103 | + body = deepcopy(dict(request)) | ||
| 104 | + self.inject_request_id(body, context.engine_request_id) | ||
| 105 | + body["stream"] = False | ||
| 106 | + body["max_tokens"] = 1 | ||
| 107 | + body["min_tokens"] = 1 | ||
| 108 | + body.pop("stream_options", None) | ||
| 109 | + if "max_completion_tokens" in body: | ||
| 110 | + body["max_completion_tokens"] = 1 | ||
| 111 | + body["kv_transfer_params"] = { | ||
| 112 | + "do_remote_decode": True, | ||
| 113 | + "do_remote_prefill": False, | ||
| 114 | + "remote_engine_id": None, | ||
| 115 | + "remote_block_ids": None, | ||
| 116 | + "remote_host": None, | ||
| 117 | + "remote_port": None, | ||
| 118 | + } | ||
| 119 | + return EngineRequest(api=context.api, body=body) | ||
| 120 | + | ||
| 121 | + def parse_prefill_response( | ||
| 122 | + self, | ||
| 123 | + response: dict[str, Any], | ||
| 124 | + ) -> PrefillMetadata: | ||
| 125 | + kv_params = response.get("kv_transfer_params") | ||
| 126 | + if not isinstance(kv_params, dict) or not kv_params: | ||
| 127 | + raise EngineProtocolError( | ||
| 128 | + engine_type=self.engine_type, | ||
| 129 | + phase="prefill", | ||
| 130 | + message="Missing kv_transfer_params", | ||
| 131 | + ) | ||
| 132 | + if kv_params.get("do_remote_prefill") is not True: | ||
| 133 | + raise EngineProtocolError( | ||
| 134 | + engine_type=self.engine_type, | ||
| 135 | + phase="prefill", | ||
| 136 | + message="do_remote_prefill must be true", | ||
| 137 | + ) | ||
| 138 | + usage = response.get("usage") | ||
| 139 | + return PrefillMetadata( | ||
| 140 | + handoff_ticket=deepcopy(kv_params), | ||
| 141 | + usage=deepcopy(usage) if isinstance(usage, dict) else None, | ||
| 142 | + ) | ||
| 143 | + | ||
| 144 | + def build_decode_request( | ||
| 145 | + self, | ||
| 146 | + request: Mapping[str, Any], | ||
| 147 | + context: LegContext, | ||
| 148 | + prefill: PrefillMetadata | None, | ||
| 149 | + ) -> EngineRequest: | ||
| 150 | + if prefill is None or not prefill.handoff_ticket: | ||
| 151 | + raise EngineProtocolError( | ||
| 152 | + engine_type=self.engine_type, | ||
| 153 | + phase="decode", | ||
| 154 | + message="Missing handoff ticket", | ||
| 155 | + ) | ||
| 156 | + body = deepcopy(dict(request)) | ||
| 157 | + self.inject_request_id(body, context.engine_request_id) | ||
| 158 | + body["kv_transfer_params"] = deepcopy(prefill.handoff_ticket) | ||
| 159 | + return EngineRequest(api=context.api, body=body) | ||
| 160 | + | ||
| 161 | + def inject_request_id(self, body: dict[str, Any], request_id: str) -> None: | ||
| 162 | + body.pop("rid", None) | ||
| 163 | + body["request_id"] = request_id | ||
| 164 | + | ||
| 165 | + def build_abort_request(self, context: LegContext) -> EngineRequest | None: | ||
| 166 | + del context | ||
| 167 | + | ||
| 168 | + | ||
| 169 | +class SglangProtocolAdapter: | ||
| 170 | + engine_type = "sglang" | ||
| 171 | + coordination_mode = CoordinationMode.BOOTSTRAP | ||
| 172 | + internal_response_fields = frozenset({"bootstrap_host", "bootstrap_port", "bootstrap_room"}) | ||
| 173 | + | ||
| 174 | + def build_prefill_request( | ||
G 严重程度： 建议 问题： sglang P 腿不再强制 stream=False， 原因： 此构建方法不重写请求的 stream 字段（客户端 stream=True 时 P 腿也是流式请求），基线 unified_pd.py 强制 怎么改： 用 SSE 标准解析替代 splitlines：按 \n\n 切分事件块，对块内多行 data 用 "\n".join 合并后再 parse_stream_chunk_json；或 P 腿统一改写为 stream=False（保留基线语义，绕开 SSE 缓冲路径）。建议在 test_pd_protocol_adapter.py 补一条用例固定 P 腿流式场景的 usage 帧解析行为。 tobking 19 天前 评论： 19 天前 评论： | |||
| 175 | + self, | ||
| 176 | + request: Mapping[str, Any], | ||
| 177 | + context: LegContext, | ||
| 178 | + ) -> EngineRequest: | ||
| 179 | + engine_request = self._build_request(request, context, context.endpoint, phase="prefill") | ||
| 180 | + engine_request.body["stream"] = False | ||
| 181 | + engine_request.body.pop("stream_options", None) | ||
| 182 | + return engine_request | ||
| 183 | + | ||
| 184 | + def parse_prefill_response( | ||
| 185 | + self, | ||
| 186 | + response: dict[str, Any], | ||
| 187 | + ) -> PrefillMetadata: | ||
| 188 | + usage = response.get("usage") | ||
| 189 | + return PrefillMetadata(usage=deepcopy(usage) if isinstance(usage, dict) else None) | ||
| 190 | + | ||
| 191 | + def build_decode_request( | ||
| 192 | + self, | ||
| 193 | + request: Mapping[str, Any], | ||
| 194 | + context: LegContext, | ||
| 195 | + prefill: PrefillMetadata | None, | ||
| 196 | + ) -> EngineRequest: | ||
| 197 | + del prefill | ||
| 198 | + if context.peer_endpoint is None: | ||
| 199 | + raise EngineProtocolError( | ||
| 200 | + engine_type=self.engine_type, | ||
| 201 | + phase="decode", | ||
| 202 | + message="Missing prefill endpoint metadata", | ||
| 203 | + ) | ||
| 204 | + return self._build_request(request, context, context.peer_endpoint, phase="decode") | ||
| 205 | + | ||
| 206 | + def inject_request_id(self, body: dict[str, Any], request_id: str) -> None: | ||
| 207 | + body.pop("request_id", None) | ||
| 208 | + body["rid"] = request_id | ||
| 209 | + | ||
| 210 | + def build_abort_request(self, context: LegContext) -> EngineRequest | None: | ||
| 211 | + return EngineRequest(api="abort_request", body={"rid": context.engine_request_id}) | ||
| 212 | + | ||
| 213 | + def _build_request( | ||
| 214 | + self, | ||
| 215 | + request: Mapping[str, Any], | ||
| 216 | + context: LegContext, | ||
| 217 | + prefill_endpoint: EngineEndpointMetadata, | ||
| 218 | + *, | ||
| 219 | + phase: str, | ||
| 220 | + ) -> EngineRequest: | ||
| 221 | + self._validate_prefill_endpoint(prefill_endpoint, phase=phase) | ||
| 222 | + body = deepcopy(dict(request)) | ||
| 223 | + self.inject_request_id(body, context.engine_request_id) | ||
| 224 | + body.update( | ||
| 225 | + { | ||
| 226 | + "bootstrap_host": prefill_endpoint.host, | ||
| 227 | + "bootstrap_port": prefill_endpoint.bootstrap_port, | ||
| 228 | + "bootstrap_room": self._stable_bootstrap_room(context.pair_id, context.attempt_seq), | ||
| 229 | + } | ||
| 230 | + ) | ||
| 231 | + return EngineRequest(api=context.api, body=body) | ||
| 232 | + | ||
| 233 | + | ||
| 234 | + def _validate_prefill_endpoint(cls, endpoint: EngineEndpointMetadata, *, phase: str) -> None: | ||
| 235 | + if not endpoint.host: | ||
| 236 | + raise EngineProtocolError( | ||
| 237 | + engine_type=cls.engine_type, | ||
| 238 | + phase=phase, | ||
| 239 | + message="Missing prefill bootstrap host", | ||
| 240 | + ) | ||
| 241 | + port = endpoint.bootstrap_port | ||
| 242 | + if not isinstance(port, int) or isinstance(port, bool) or not 1 <= port <= 65535: | ||
| 243 | + raise EngineProtocolError( | ||
| 244 | + engine_type=cls.engine_type, | ||
| 245 | + phase=phase, | ||
| 246 | + message="Missing or invalid prefill bootstrap port", | ||
| 247 | + ) | ||
| 248 | + | ||
| 249 | + | ||
| 250 | + def _stable_bootstrap_room(pair_id: str, attempt_seq: int) -> int: | ||
| 251 | + raw = f"{pair_id}:{attempt_seq}".encode("utf-8") | ||
| 252 | + digest = hashlib.blake2b(raw, digest_size=8).digest() | ||
| 253 | + return int.from_bytes(digest, "big") & ((1 << 63) - 1) | ||
| 254 | + | ||
| 255 | + | ||
| 256 | +ADAPTERS: Mapping[str, PDProtocolAdapter] = MappingProxyType( | ||
| 257 | + { | ||
| 258 | + VllmProtocolAdapter.engine_type: VllmProtocolAdapter(), | ||
| 259 | + SglangProtocolAdapter.engine_type: SglangProtocolAdapter(), | ||
| 260 | + } | ||
| 261 | +) | ||
| @@ -179,29 +179,23 @@ async def select_router_class( | |||
| 179 | ) -> type["BaseRouter"]: | 179 | ) -> type["BaseRouter"]: |
| 180 | """Select the router implementation from the live instance topology. | 180 | """Select the router implementation from the live instance topology. |
| 181 | 181 | ||
| 182 | - Routing is derived from the roles currently present plus whether a P/D pair shares a | 182 | + Routing is derived from the live roles and circuit-breaker state. Native protocol |
| 183 | - dispatch capability — no deploy_mode. Shared by user traffic (handle_request) and the | 183 | + selection happens inside UnifiedPDRouter from the selected instance engine_type. |
| 184 | - internal precision probe so both route identically. | 184 | + Shared by user traffic and the internal precision probe so both route identically. |
| 185 | 185 | ||
| 186 | Raises HTTPException(503) when no routable topology is available. | 186 | Raises HTTPException(503) when no routable topology is available. |
| 187 | """ | 187 | """ |
| 188 | roles = await scheduler.get_available_instance_roles() | 188 | roles = await scheduler.get_available_instance_roles() |
| 189 | has_pd_roles = PDRole.ROLE_P in roles and PDRole.ROLE_D in roles | 189 | has_pd_roles = PDRole.ROLE_P in roles and PDRole.ROLE_D in roles |
| 190 | - has_compatible_pair = False | 190 | + has_routable_pd_pair = has_pd_roles |
| 191 | if has_pd_roles: | 191 | if has_pd_roles: |
| 192 | - compatibility_check = getattr(scheduler, "has_compatible_pd_pair", None) | 192 | + get_unblocked = getattr(scheduler, "get_unblocked_instances", None) |
| 193 | - has_compatible_pair = await compatibility_check() if compatibility_check is not None else True | 193 | + if get_unblocked is not None: |
| 194 | - if has_compatible_pair: | 194 | + unblocked_p = await get_unblocked(PDRole.ROLE_P) |
| 195 | - # Check circuit breaker: only treat as compatible if there are | 195 | + unblocked_d = await get_unblocked(PDRole.ROLE_D) |
| 196 | - # non-blocked instances for BOTH roles | 196 | + has_routable_pd_pair = bool(unblocked_p and unblocked_d) |
| 197 | - get_unblocked = getattr(scheduler, "get_unblocked_instances", None) | ||
| 198 | - if get_unblocked is not None: | ||
| 199 | - unblocked_p = await get_unblocked(PDRole.ROLE_P) | ||
| 200 | - unblocked_d = await get_unblocked(PDRole.ROLE_D) | ||
| 201 | - if not unblocked_p or not unblocked_d: | ||
| 202 | - has_compatible_pair = False | ||
| 203 | 197 | ||
| 204 | - if has_compatible_pair: | 198 | + if has_routable_pd_pair: |
| 205 | return UnifiedPDRouter | 199 | return UnifiedPDRouter |
| 206 | 200 | ||
| 207 | # Degrade to hybrid mode if any unblocked instance is available | 201 | # Degrade to hybrid mode if any unblocked instance is available |
| @@ -225,7 +219,7 @@ async def select_router_class( | |||
| 225 | is_hybrid_deploy = _is_pd_hybrid_deploy(config) | 219 | is_hybrid_deploy = _is_pd_hybrid_deploy(config) |
| 226 | if not fallback_enabled and not is_hybrid_deploy: | 220 | if not fallback_enabled and not is_hybrid_deploy: |
| 227 | if has_pd_roles: | 221 | if has_pd_roles: |
| 228 | - message = "PD separate service has no compatible P/D pair and fallback to hybrid is disabled" | 222 | + message = "PD separate service has no circuit-breaker-available P/D pair and fallback is disabled" |
| 229 | else: | 223 | else: |
| 230 | message = "PD separate service is unavailable and fallback to hybrid is disabled" | 224 | message = "PD separate service is unavailable and fallback to hybrid is disabled" |
| 231 | if req_info is not None: | 225 | if req_info is not None: |
| @@ -234,17 +228,13 @@ async def select_router_class( | |||
| 234 | raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=message) | 228 | raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=message) |
| 235 | 229 | ||
| 236 | if PDRole.ROLE_U in roles or PDRole.ROLE_P in roles: | 230 | if PDRole.ROLE_U in roles or PDRole.ROLE_P in roles: |
| 237 | - if has_pd_roles and not has_compatible_pair and PDRole.ROLE_U in roles: | 231 | + if has_pd_roles and not has_routable_pd_pair and PDRole.ROLE_U in roles: |
| 238 | - message = ( | 232 | + message = "P/D instances are unavailable; falling back to PDHybridRouter via union instances" |
| 239 | - "P/D instances are online but advertise no shared dispatch capability; " | ||
| 240 | - "falling back to PDHybridRouter via union instances. " | ||
| 241 | - "Check the engine kv_connector is recognized or set dispatch_profile explicitly." | ||
| 242 | - ) | ||
| 243 | if req_info is not None: | 233 | if req_info is not None: |
| 244 | req_info.trace_obj.set_trace_error_message(message) | 234 | req_info.trace_obj.set_trace_error_message(message) |
| 245 | logger.warning(message) | 235 | logger.warning(message) |
| 246 | - elif has_pd_roles and not has_compatible_pair and req_info is not None: | 236 | + elif has_pd_roles and not has_routable_pd_pair and req_info is not None: |
| 247 | - error_message = "PD separate service degraded to hybrid: P or D instances circuit-broken or incompatible" | 237 | + error_message = "PD separate service degraded to hybrid: P or D instances are circuit-broken" |
| 248 | req_info.trace_obj.set_trace_error_message(error_message) | 238 | req_info.trace_obj.set_trace_error_message(error_message) |
| 249 | logger.warning(error_message) | 239 | logger.warning(error_message) |
| 250 | elif req_info is not None and PDRole.ROLE_U not in roles: | 240 | elif req_info is not None and PDRole.ROLE_U not in roles: |
| @@ -265,12 +255,14 @@ async def handle_request( | |||
| 265 | scheduler=None, | 255 | scheduler=None, |
| 266 | *, | 256 | *, |
| 267 | request_manager: RequestManager, | 257 | request_manager: RequestManager, |
| 258 | + request_json: dict | None = None, | ||
| 268 | ) -> Response: | 259 | ) -> Response: |
| 269 | """Handle incoming requests and route them to appropriate router implementation | 260 | """Handle incoming requests and route them to appropriate router implementation |
| 270 | 261 | ||
| 271 | Args: | 262 | Args: |
| 272 | raw_request: The incoming FastAPI request object | 263 | raw_request: The incoming FastAPI request object |
| 273 | request_manager: RequestManager instance (required, injected by InferenceServer) | 264 | request_manager: RequestManager instance (required, injected by InferenceServer) |
| 265 | + request_json: Body already parsed by the API ingress, when available | ||
| 274 | 266 | ||
| 275 | Returns: | 267 | Returns: |
| 276 | Response: The response from the selected router implementation (stream, non-stream, or error) | 268 | Response: The response from the selected router implementation (stream, non-stream, or error) |
| @@ -279,7 +271,7 @@ async def handle_request( | |||
| 279 | HTTPException: If request body is empty or request fail | 271 | HTTPException: If request body is empty or request fail |
| 280 | """ | 272 | """ |
| 281 | 273 | ||
| 282 | - req_info = await __create_request_info(raw_request, request_manager) | 274 | + req_info = await __create_request_info(raw_request, request_manager, request_json=request_json) |
| 283 | 275 | ||
| 284 | if TracerManager().contains_trace_headers(raw_request.headers): | 276 | if TracerManager().contains_trace_headers(raw_request.headers): |
| 285 | req_info.trace_obj.parent_context = TracerManager().extract_trace_context(raw_request.headers) | 277 | req_info.trace_obj.parent_context = TracerManager().extract_trace_context(raw_request.headers) |
| @@ -344,16 +336,18 @@ async def handle_request( | |||
| 344 | async def __create_request_info( | 336 | async def __create_request_info( |
| 345 | raw_request: Request, | 337 | raw_request: Request, |
| 346 | request_manager: RequestManager, | 338 | request_manager: RequestManager, |
| 339 | + request_json: dict | None = None, | ||
| 347 | ) -> RequestInfo: | 340 | ) -> RequestInfo: |
| 348 | request_body = await raw_request.body() | 341 | request_body = await raw_request.body() |
| 349 | if not request_body: | 342 | if not request_body: |
| 350 | raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Empty request body") | 343 | raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Empty request body") |
| 351 | 344 | ||
| 352 | - try: | 345 | + if request_json is None: |
| 353 | - request_json = await raw_request.json() | 346 | + try: |
| 354 | - except Exception as e: | 347 | + request_json = await raw_request.json() |
| 355 | - logger.warning("JSON parse failed: %s", e) | 348 | + except Exception as e: |
| 356 | - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid JSON format") from e | 349 | + logger.warning("JSON parse failed: %s", e) |
| 350 | + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid JSON format") from e | ||
| 357 | 351 | ||
| 358 | if not request_json: | 352 | if not request_json: |
| 359 | raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Empty request json") | 353 | raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Empty request json") |
| @@ -1,53 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -from motor.common.resources.dispatch import DispatchPlan, dispatch_plans_from_capabilities, shared_dispatch_plans | ||
| 12 | -from motor.coordinator.domain import ScheduledResource | ||
| 13 | - | ||
| 14 | - | ||
| 15 | -class DispatchPlanNotSupported(RuntimeError): | ||
| 16 | - pass | ||
| 17 | - | ||
| 18 | - | ||
| 19 | -def select_dispatch_plan_for_pair( | ||
| 20 | - *, | ||
| 21 | - prefill: ScheduledResource | None, | ||
| 22 | - decode: ScheduledResource | None, | ||
| 23 | -) -> DispatchPlan: | ||
| 24 | - explicit_plan = _select_explicit_plan(prefill, decode) | ||
| 25 | - if explicit_plan is not None: | ||
| 26 | - return explicit_plan | ||
| 27 | - raise DispatchPlanNotSupported( | ||
| 28 | - "Selected P/D instances do not advertise a shared dispatch capability; " | ||
| 29 | - "configure a supported engine connector or dispatch_profile" | ||
| 30 | - ) | ||
| 31 | - | ||
| 32 | - | ||
| 33 | -def _select_explicit_plan( | ||
| 34 | - prefill: ScheduledResource | None, | ||
| 35 | - decode: ScheduledResource | None, | ||
| 36 | -) -> DispatchPlan | None: | ||
| 37 | - prefill_instance = prefill.instance if prefill is not None else None | ||
| 38 | - decode_instance = decode.instance if decode is not None else None | ||
| 39 | - prefill_plans = dispatch_plans_from_capabilities(getattr(prefill_instance, "dispatch_capabilities", None)) | ||
| 40 | - decode_plans = dispatch_plans_from_capabilities(getattr(decode_instance, "dispatch_capabilities", None)) | ||
| 41 | - if not prefill_plans or not decode_plans: | ||
| 42 | - return None | ||
| 43 | - | ||
| 44 | - supported = shared_dispatch_plans(prefill_instance, decode_instance) | ||
| 45 | - | ||
| 46 | - preferred = [ | ||
| 47 | - DispatchPlan.CONCURRENT_ENGINE_SYNC, | ||
| 48 | - DispatchPlan.PREFILL_HANDOFF_DECODE, | ||
| 49 | - ] | ||
| 50 | - for plan in preferred: | ||
| 51 | - if plan in supported: | ||
| 52 | - return plan | ||
| 53 | - raise DispatchPlanNotSupported("Selected P/D instances have no shared dispatch capability") | ||
| @@ -9,20 +9,11 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | import asyncio | 11 | import asyncio |
| 12 | -import time | ||
| 13 | import uuid | 12 | import uuid |
| 14 | from dataclasses import dataclass, field | 13 | from dataclasses import dataclass, field |
| 15 | from enum import Enum | 14 | from enum import Enum |
| 16 | 15 | ||
| 17 | from motor.common.http import HTTPClientPool | 16 | from motor.common.http import HTTPClientPool |
| 18 | -from motor.common.resources.dispatch import ( | ||
| 19 | - DispatchEndpoint, | ||
| 20 | - DispatchEndpoints, | ||
| 21 | - MotorDispatch, | ||
| 22 | - PrefillContextBudget, | ||
| 23 | -) | ||
| 24 | -from motor.common.resources.instance import PDRole | ||
| 25 | -from motor.common.utils.net import format_address | ||
| 26 | from motor.config.coordinator import CoordinatorConfig | 17 | from motor.config.coordinator import CoordinatorConfig |
| 27 | from motor.coordinator.domain import ScheduledResource | 18 | from motor.coordinator.domain import ScheduledResource |
| 28 | 19 | ||
| @@ -41,6 +32,12 @@ class AttemptState(str, Enum): | |||
| 41 | STOPPED = "stopped" | 32 | STOPPED = "stopped" |
| 42 | 33 | ||
| 43 | 34 | ||
| 35 | +class AttemptStopReason(str, Enum): | ||
| 36 | + CLIENT_DISCONNECT = "client_disconnect" | ||
| 37 | + PEER_FAILED = "peer_failed" | ||
| 38 | + OTHER = "other" | ||
| 39 | + | ||
| 40 | + | ||
| 44 | 41 | ||
| 45 | class AttemptReleaseFlags: | 42 | class AttemptReleaseFlags: |
| 46 | prefill_tokens: bool = False | 43 | prefill_tokens: bool = False |
| @@ -52,20 +49,18 @@ class AttemptContext: | |||
| 52 | root_request_id: str | 49 | root_request_id: str |
| 53 | attempt_seq: int | 50 | attempt_seq: int |
| 54 | pair_id: str | 51 | pair_id: str |
| 55 | - prefill_context_budget: PrefillContextBudget | None = None | ||
| 56 | prefill_resource: ScheduledResource | None = None | 52 | prefill_resource: ScheduledResource | None = None |
| 57 | decode_resource: ScheduledResource | None = None | 53 | decode_resource: ScheduledResource | None = None |
| 58 | state: AttemptState = AttemptState.CREATED | 54 | state: AttemptState = AttemptState.CREATED |
| 59 | - first_visible_sent: bool = False | ||
| 60 | - stop_sent: bool = False | ||
| 61 | release_flags: AttemptReleaseFlags = field(default_factory=AttemptReleaseFlags) | 55 | release_flags: AttemptReleaseFlags = field(default_factory=AttemptReleaseFlags) |
| 62 | - created_at: float = field(default_factory=time.time) | ||
| 63 | - updated_at: float = field(default_factory=time.time) | ||
| 64 | prefill_task: asyncio.Task | None = None | 56 | prefill_task: asyncio.Task | None = None |
| 65 | decode_task: asyncio.Task | None = None | 57 | decode_task: asyncio.Task | None = None |
| 58 | + prefill_dispatched: bool = False | ||
| 59 | + prefill_completed: bool = False | ||
| 60 | + decode_dispatched: bool = False | ||
| 61 | + decode_completed: bool = False | ||
| 66 | stop_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) | 62 | stop_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) |
| 67 | config: CoordinatorConfig | None = None | 63 | config: CoordinatorConfig | None = None |
| 68 | - fail_reason: str | None = None | ||
| 69 | 64 | ||
| 70 | def transition(self, state: AttemptState) -> bool: | 65 | def transition(self, state: AttemptState) -> bool: |
| 71 | if self.state == AttemptState.STOPPED: | 66 | if self.state == AttemptState.STOPPED: |
| @@ -73,37 +68,16 @@ class AttemptContext: | |||
| 73 | if self.state == AttemptState.STOPPING: | 68 | if self.state == AttemptState.STOPPING: |
| 74 | if state == AttemptState.STOPPED: | 69 | if state == AttemptState.STOPPED: |
| 75 | self.state = state | 70 | self.state = state |
| 76 | - self.updated_at = time.time() | ||
| 77 | return True | 71 | return True |
| 78 | return state == AttemptState.STOPPING | 72 | return state == AttemptState.STOPPING |
| 79 | if self.state == AttemptState.DONE: | 73 | if self.state == AttemptState.DONE: |
| 80 | return state == AttemptState.DONE | 74 | return state == AttemptState.DONE |
| 81 | self.state = state | 75 | self.state = state |
| 82 | - self.updated_at = time.time() | ||
| 83 | - if state == AttemptState.FIRST_VISIBLE: | ||
| 84 | - self.first_visible_sent = True | ||
| 85 | return True | 76 | return True |
| 86 | 77 | ||
| 87 | def stop(self) -> None: | 78 | def stop(self) -> None: |
| 88 | if self.state not in (AttemptState.DONE, AttemptState.STOPPED): | 79 | if self.state not in (AttemptState.DONE, AttemptState.STOPPED): |
| 89 | self.state = AttemptState.STOPPING | 80 | self.state = AttemptState.STOPPING |
| 90 | - self.stop_sent = True | ||
| 91 | - self.updated_at = time.time() | ||
| 92 | - | ||
| 93 | - def dispatch_for(self, role: PDRole, dispatch_mode: str) -> MotorDispatch: | ||
| 94 | - return MotorDispatch( | ||
| 95 | - root_request_id=self.root_request_id, | ||
| 96 | - engine_request_id=f"{self.root_request_id}#a{self.attempt_seq}", | ||
| 97 | - pair_id=self.pair_id, | ||
| 98 | - attempt_seq=self.attempt_seq, | ||
| 99 | - role="prefill" if role == PDRole.ROLE_P else "decode", | ||
| 100 | - dispatch_mode=dispatch_mode, | ||
| 101 | - prefill_context_budget=self.prefill_context_budget, | ||
| 102 | - endpoints=DispatchEndpoints( | ||
| 103 | - prefill=_dispatch_endpoint(self.prefill_resource), | ||
| 104 | - decode=_dispatch_endpoint(self.decode_resource), | ||
| 105 | - ), | ||
| 106 | - ) | ||
| 107 | 81 | ||
| 108 | def register_prefill_task(self, task: asyncio.Task) -> asyncio.Task: | 82 | def register_prefill_task(self, task: asyncio.Task) -> asyncio.Task: |
| 109 | self.prefill_task = task | 83 | self.prefill_task = task |
| @@ -113,25 +87,45 @@ class AttemptContext: | |||
| 113 | self.decode_task = task | 87 | self.decode_task = task |
| 114 | return task | 88 | return task |
| 115 | 89 | ||
| 90 | + def mark_dispatched(self, role: str) -> None: | ||
| 91 | + if role == "prefill": | ||
| 92 | + self.prefill_dispatched = True | ||
| 93 | + else: | ||
| 94 | + self.decode_dispatched = True | ||
| 95 | + | ||
| 96 | + def mark_completed(self, role: str) -> None: | ||
| 97 | + if role == "prefill": | ||
| 98 | + self.prefill_completed = True | ||
| 99 | + else: | ||
| 100 | + self.decode_completed = True | ||
| 101 | + | ||
| 102 | + def needs_abort(self, role: str) -> bool: | ||
| 103 | + if role == "prefill": | ||
| 104 | + return self.prefill_dispatched and not self.prefill_completed | ||
| 105 | + return self.decode_dispatched and not self.decode_completed | ||
| 106 | + | ||
| 116 | async def cancel(self, reason: str = ""): | 107 | async def cancel(self, reason: str = ""): |
| 117 | - self.fail_reason = reason | 108 | + tasks = [] |
| 118 | - task = [] | ||
| 119 | if self.prefill_task and not self.prefill_task.done() and not self.prefill_task.cancelled(): | 109 | if self.prefill_task and not self.prefill_task.done() and not self.prefill_task.cancelled(): |
| 120 | logger.info( | 110 | logger.info( |
| 121 | - f"Cancelling prefill task: {self.prefill_resource.endpoint.ip} {self.prefill_resource.instance.job_name}" | 111 | + "Cancelling prefill task: %s %s because %s", |
| 122 | - f" because {reason}" | 112 | + self.prefill_resource.endpoint.ip, |
| 113 | + self.prefill_resource.instance.job_name, | ||
| 114 | + reason, | ||
| 123 | ) | 115 | ) |
| 124 | self.prefill_task.cancel(msg=reason) | 116 | self.prefill_task.cancel(msg=reason) |
| 125 | - task.append(self.prefill_task) | 117 | + tasks.append(self.prefill_task) |
| 126 | if self.decode_task and not self.decode_task.done() and not self.decode_task.cancelled(): | 118 | if self.decode_task and not self.decode_task.done() and not self.decode_task.cancelled(): |
| 127 | logger.info( | 119 | logger.info( |
| 128 | - f"Cancelling decode task: {self.decode_resource.endpoint.ip} {self.decode_resource.instance.job_name}" | 120 | + "Cancelling decode task: %s %s because %s", |
| 129 | - f" because {reason}" | 121 | + self.decode_resource.endpoint.ip, |
| 122 | + self.decode_resource.instance.job_name, | ||
| 123 | + reason, | ||
| 130 | ) | 124 | ) |
| 131 | self.decode_task.cancel(msg=reason) | 125 | self.decode_task.cancel(msg=reason) |
| 132 | - task.append(self.decode_task) | 126 | + tasks.append(self.decode_task) |
| 133 | - if task: | 127 | + if tasks: |
| 134 | - await asyncio.gather(*task, return_exceptions=True) | 128 | + await asyncio.gather(*tasks, return_exceptions=True) |
| 135 | 129 | ||
| 136 | def register_canceller(self): | 130 | def register_canceller(self): |
| 137 | pool = HTTPClientPool() | 131 | pool = HTTPClientPool() |
| @@ -168,8 +162,7 @@ class AttemptContext: | |||
| 168 | ) | 162 | ) |
| 169 | pool.unregister_canceller(d_key, self.pair_id) | 163 | pool.unregister_canceller(d_key, self.pair_id) |
| 170 | except Exception as e: | 164 | except Exception as e: |
| 171 | - logger.error(f"Unregister error {e=}") | 165 | + logger.error("Unregister error: %s", e) |
| 172 | - pass | ||
| 173 | 166 | ||
| 174 | def unregister_prefill_canceller(self): | 167 | def unregister_prefill_canceller(self): |
| 175 | try: | 168 | try: |
| @@ -182,8 +175,7 @@ class AttemptContext: | |||
| 182 | ) | 175 | ) |
| 183 | pool.unregister_canceller(p_key, self.pair_id) | 176 | pool.unregister_canceller(p_key, self.pair_id) |
| 184 | except Exception as e: | 177 | except Exception as e: |
| 185 | - logger.error(f"Unregister error {e=}") | 178 | + logger.error("Unregister error: %s", e) |
| 186 | - pass | ||
| 187 | 179 | ||
| 188 | def register_decode_canceller(self): | 180 | def register_decode_canceller(self): |
| 189 | if not self.decode_resource: | 181 | if not self.decode_resource: |
| @@ -198,47 +190,23 @@ class AttemptContext: | |||
| 198 | 190 | ||
| 199 | 191 | ||
| 200 | class PDDispatchSession: | 192 | class PDDispatchSession: |
| 201 | - def __init__( | 193 | + def __init__(self, root_request_id: str) -> None: |
| 202 | - self, | ||
| 203 | - root_request_id: str, | ||
| 204 | - prefill_context_budget: PrefillContextBudget | None = None, | ||
| 205 | - ) -> None: | ||
| 206 | self.root_request_id = root_request_id | 194 | self.root_request_id = root_request_id |
| 207 | - self.prefill_context_budget = prefill_context_budget | ||
| 208 | self._attempt_seq = 0 | 195 | self._attempt_seq = 0 |
| 209 | - self.attempts: dict[int, AttemptContext] = {} | ||
| 210 | 196 | ||
| 211 | def new_attempt( | 197 | def new_attempt( |
| 212 | self, | 198 | self, |
| 213 | prefill_resource: ScheduledResource | None, | 199 | prefill_resource: ScheduledResource | None, |
| 214 | decode_resource: ScheduledResource | None, | 200 | decode_resource: ScheduledResource | None, |
| 215 | config: CoordinatorConfig, | 201 | config: CoordinatorConfig, |
| 216 | - *, | ||
| 217 | - consumed_output_tokens: int = 0, | ||
| 218 | ) -> AttemptContext: | 202 | ) -> AttemptContext: |
| 219 | self._attempt_seq += 1 | 203 | self._attempt_seq += 1 |
| 220 | - budget = self.prefill_context_budget | ||
| 221 | - if budget is not None: | ||
| 222 | - budget = budget.after_output_tokens(consumed_output_tokens) | ||
| 223 | attempt = AttemptContext( | 204 | attempt = AttemptContext( |
| 224 | root_request_id=self.root_request_id, | 205 | root_request_id=self.root_request_id, |
| 225 | attempt_seq=self._attempt_seq, | 206 | attempt_seq=self._attempt_seq, |
| 226 | pair_id=uuid.uuid4().hex, | 207 | pair_id=uuid.uuid4().hex, |
| 227 | - prefill_context_budget=budget, | ||
| 228 | prefill_resource=prefill_resource, | 208 | prefill_resource=prefill_resource, |
| 229 | decode_resource=decode_resource, | 209 | decode_resource=decode_resource, |
| 230 | config=config, | 210 | config=config, |
| 231 | ) | 211 | ) |
| 232 | - self.attempts[attempt.attempt_seq] = attempt | ||
| 233 | return attempt | 212 | return attempt |
| 234 | - | ||
| 235 | - | ||
| 236 | -def _dispatch_endpoint(resource: ScheduledResource | None) -> DispatchEndpoint | None: | ||
| 237 | - if not resource or not resource.instance or not resource.endpoint: | ||
| 238 | - return None | ||
| 239 | - endpoint = resource.endpoint | ||
| 240 | - return DispatchEndpoint( | ||
| 241 | - instance_id=int(resource.instance.id), | ||
| 242 | - endpoint_id=int(endpoint.id), | ||
| 243 | - url=f"http://{format_address(endpoint.ip, endpoint.business_port)}", | ||
| 244 | - ) | ||
| @@ -37,7 +37,7 @@ def build_decode_sample( | |||
| 37 | """Construct a DecodeSample from request_info and context. | 37 | """Construct a DecodeSample from request_info and context. |
| 38 | 38 | ||
| 39 | Args: | 39 | Args: |
| 40 | - p_instance_id: P instance id (None if unknown in CDP mode). | 40 | + p_instance_id: P instance id (None for a Hybrid/Union request). |
| 41 | d_instance_id: D instance id. | 41 | d_instance_id: D instance id. |
| 42 | request_info: Per-request mutable dict populated by | 42 | request_info: Per-request mutable dict populated by |
| 43 | update_token_id_cache / update_logprob_cache. | 43 | update_token_id_cache / update_logprob_cache. |
| @@ -8,7 +8,8 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -"""Shared helpers for PD/CDP rescheduler (token-cache retry and request rewrite). | 11 | +"""Shared helpers for P/D rescheduling (token-cache retry and request rewrite). |
| 12 | + | ||
| 12 | 13 | ||
| 13 | Performance note: per-chunk JSON re-serialization is CPU-bound; ``async`` does not | 14 | Performance note: per-chunk JSON re-serialization is CPU-bound; ``async`` does not |
| 14 | make it faster. Mitigations here use :mod:`msgspec` (already a dependency) for | 15 | make it faster. Mitigations here use :mod:`msgspec` (already a dependency) for |
| @@ -1,77 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -from __future__ import annotations | ||
| 12 | - | ||
| 13 | -import hashlib | ||
| 14 | -import os | ||
| 15 | -from typing import Any | ||
| 16 | - | ||
| 17 | -from motor.coordinator.domain import ScheduledResource | ||
| 18 | -from motor.coordinator.router.dispatch_session import AttemptContext | ||
| 19 | - | ||
| 20 | -_SGLANG_ENGINE_TYPE = "sglang" | ||
| 21 | - | ||
| 22 | - | ||
| 23 | -def is_sglang_resource(resource: ScheduledResource | None) -> bool: | ||
| 24 | - """Return True if the scheduled resource reports engine_type=sglang.""" | ||
| 25 | - if resource is None or resource.instance is None: | ||
| 26 | - return False | ||
| 27 | - engine_type = str(getattr(resource.instance, "engine_type", "") or "").strip().lower() | ||
| 28 | - return engine_type == _SGLANG_ENGINE_TYPE | ||
| 29 | - | ||
| 30 | - | ||
| 31 | -def ensure_sglang_pd_pair(attempt: AttemptContext) -> None: | ||
| 32 | - """Require both P and D legs to be SGLang for native bootstrap PD.""" | ||
| 33 | - prefill = attempt.prefill_resource | ||
| 34 | - decode = attempt.decode_resource | ||
| 35 | - if not is_sglang_resource(prefill) or not is_sglang_resource(decode): | ||
| 36 | - prefill_engine = prefill.instance.engine_type if prefill and prefill.instance else None | ||
| 37 | - decode_engine = decode.instance.engine_type if decode and decode.instance else None | ||
| 38 | - raise RuntimeError( | ||
| 39 | - "SGLang native PD requires both prefill and decode engine_type=sglang, " | ||
| 40 | - f"got prefill={prefill_engine!r} decode={decode_engine!r}." | ||
| 41 | - ) | ||
| 42 | - | ||
| 43 | - | ||
| 44 | -def _prefill_bootstrap_host(attempt: AttemptContext) -> str: | ||
| 45 | - """Return Prefill endpoint IP.""" | ||
| 46 | - resource = attempt.prefill_resource | ||
| 47 | - if resource is None or resource.endpoint is None: | ||
| 48 | - raise RuntimeError("SGLang PD requires a scheduled prefill endpoint for bootstrap_host") | ||
| 49 | - return resource.endpoint.ip | ||
| 50 | - | ||
| 51 | - | ||
| 52 | -def _bootstrap_port() -> str: | ||
| 53 | - """Read and validate DISAGGREGATION_BOOTSTRAP_PORT.""" | ||
| 54 | - raw = os.getenv("DISAGGREGATION_BOOTSTRAP_PORT", "").strip() | ||
| 55 | - try: | ||
| 56 | - port = int(raw) | ||
| 57 | - except ValueError as e: | ||
| 58 | - raise RuntimeError(f"DISAGGREGATION_BOOTSTRAP_PORT must be an integer (e.g. 8998), got {raw!r}.") from e | ||
| 59 | - if not 1 <= port <= 65535: | ||
| 60 | - raise RuntimeError(f"DISAGGREGATION_BOOTSTRAP_PORT must be in range 1-65535, got {port}.") | ||
| 61 | - return str(port) | ||
| 62 | - | ||
| 63 | - | ||
| 64 | -def _stable_bootstrap_room(pair_id: str, attempt_seq: int) -> int: | ||
| 65 | - """Derive a stable positive int63 bootstrap_room for a P/D attempt.""" | ||
| 66 | - raw = f"{pair_id}:{attempt_seq}".encode("utf-8") | ||
| 67 | - digest = hashlib.blake2b(raw, digest_size=8).digest() | ||
| 68 | - return int.from_bytes(digest, "big") & ((1 << 63) - 1) | ||
| 69 | - | ||
| 70 | - | ||
| 71 | -def inject_sglang_pd_fields(req: dict[str, Any], attempt: AttemptContext) -> None: | ||
| 72 | - """Mutate ``req`` in place with stock SGLang PD fields (no ``_motor_dispatch``).""" | ||
| 73 | - ensure_sglang_pd_pair(attempt) | ||
| 74 | - req["bootstrap_host"] = _prefill_bootstrap_host(attempt) | ||
| 75 | - req["bootstrap_port"] = _bootstrap_port() | ||
| 76 | - req["bootstrap_room"] = _stable_bootstrap_room(attempt.pair_id, attempt.attempt_seq) | ||
| 77 | - req.setdefault("request_id", f"{attempt.root_request_id}#a{attempt.attempt_seq}") | ||
| @@ -1,155 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -import time | ||
| 12 | - | ||
| 13 | -import httpx | ||
| 14 | - | ||
| 15 | -from motor.common.http.http_client import HTTPClientPool | ||
| 16 | -from motor.common.logger import get_logger | ||
| 17 | -from motor.common.resources.dispatch import ( | ||
| 18 | - DispatchStopReason, | ||
| 19 | - DispatchStopRequest, | ||
| 20 | - DispatchStopResponse, | ||
| 21 | - DispatchStopState, | ||
| 22 | -) | ||
| 23 | -from motor.config.coordinator import CoordinatorConfig | ||
| 24 | -from motor.coordinator.domain import ScheduledResource | ||
| 25 | -from motor.coordinator.router.dispatch_session import AttemptContext | ||
| 26 | -from motor.coordinator.router.sglang_native_dispatch import is_sglang_resource | ||
| 27 | - | ||
| 28 | -logger = get_logger(__name__) | ||
| 29 | - | ||
| 30 | - | ||
| 31 | -def _engine_request_id(attempt: AttemptContext) -> str: | ||
| 32 | - return f"{attempt.root_request_id}#a{attempt.attempt_seq}" | ||
| 33 | - | ||
| 34 | - | ||
| 35 | -class DispatchStopClient: | ||
| 36 | - def __init__(self, config: CoordinatorConfig) -> None: | ||
| 37 | - self._config = config | ||
| 38 | - | ||
| 39 | - async def stop( | ||
| 40 | - self, | ||
| 41 | - resource: ScheduledResource, | ||
| 42 | - attempt: AttemptContext, | ||
| 43 | - reason: DispatchStopReason, | ||
| 44 | - timeout: float = 1.0, | ||
| 45 | - ) -> DispatchStopResponse | None: | ||
| 46 | - if not resource or not resource.endpoint: | ||
| 47 | - return None | ||
| 48 | - | ||
| 49 | - if is_sglang_resource(resource): | ||
| 50 | - return await self._stop_sglang_native(resource, attempt, reason, timeout) | ||
| 51 | - return await self._stop_motor_dispatch(resource, attempt, reason, timeout) | ||
| 52 | - | ||
| 53 | - async def _stop_sglang_native( | ||
| 54 | - self, | ||
| 55 | - resource: ScheduledResource, | ||
| 56 | - attempt: AttemptContext, | ||
| 57 | - reason: DispatchStopReason, | ||
| 58 | - timeout: float, | ||
| 59 | - ) -> DispatchStopResponse | None: | ||
| 60 | - """Abort via stock SGLang ``POST /abort_request`` (no InferEndpoint stop API).""" | ||
| 61 | - endpoint = resource.endpoint | ||
| 62 | - engine_request_id = _engine_request_id(attempt) | ||
| 63 | - try: | ||
| 64 | - client = await HTTPClientPool().get_client( | ||
| 65 | - ip=endpoint.ip, | ||
| 66 | - port=endpoint.business_port, | ||
| 67 | - tls_config=self._config.infer_tls_config, | ||
| 68 | - ) | ||
| 69 | - response = await client.post( | ||
| 70 | - "/abort_request", | ||
| 71 | - json={"rid": engine_request_id}, | ||
| 72 | - timeout=timeout, | ||
| 73 | - ) | ||
| 74 | - response.raise_for_status() | ||
| 75 | - logger.info( | ||
| 76 | - "SGLang abort_request accepted root_request_id=%s attempt_seq=%s endpoint=%s:%s rid=%s reason=%s", | ||
| 77 | - attempt.root_request_id, | ||
| 78 | - attempt.attempt_seq, | ||
| 79 | - endpoint.ip, | ||
| 80 | - endpoint.business_port, | ||
| 81 | - engine_request_id, | ||
| 82 | - reason.value, | ||
| 83 | - ) | ||
| 84 | - return DispatchStopResponse( | ||
| 85 | - root_request_id=attempt.root_request_id, | ||
| 86 | - attempt_seq=attempt.attempt_seq, | ||
| 87 | - accepted=True, | ||
| 88 | - state=DispatchStopState.STOPPED, | ||
| 89 | - message="sglang:/abort_request", | ||
| 90 | - ) | ||
| 91 | - except httpx.HTTPError as e: | ||
| 92 | - logger.warning( | ||
| 93 | - "SGLang abort_request failed root_request_id=%s attempt_seq=%s endpoint=%s:%s rid=%s error=%s", | ||
| 94 | - attempt.root_request_id, | ||
| 95 | - attempt.attempt_seq, | ||
| 96 | - endpoint.ip, | ||
| 97 | - endpoint.business_port, | ||
| 98 | - engine_request_id, | ||
| 99 | - e, | ||
| 100 | - ) | ||
| 101 | - except Exception as e: | ||
| 102 | - logger.warning( | ||
| 103 | - "SGLang abort_request unexpected error root_request_id=%s attempt_seq=%s error=%s", | ||
| 104 | - attempt.root_request_id, | ||
| 105 | - attempt.attempt_seq, | ||
| 106 | - e, | ||
| 107 | - ) | ||
| 108 | - return None | ||
| 109 | - | ||
| 110 | - async def _stop_motor_dispatch( | ||
| 111 | - self, | ||
| 112 | - resource: ScheduledResource, | ||
| 113 | - attempt: AttemptContext, | ||
| 114 | - reason: DispatchStopReason, | ||
| 115 | - timeout: float, | ||
| 116 | - ) -> DispatchStopResponse | None: | ||
| 117 | - endpoint = resource.endpoint | ||
| 118 | - request = DispatchStopRequest( | ||
| 119 | - root_request_id=attempt.root_request_id, | ||
| 120 | - engine_request_id=_engine_request_id(attempt), | ||
| 121 | - attempt_seq=attempt.attempt_seq, | ||
| 122 | - pair_id=attempt.pair_id, | ||
| 123 | - reason=reason.value, | ||
| 124 | - sent_at_ms=int(time.time() * 1000), | ||
| 125 | - ) | ||
| 126 | - try: | ||
| 127 | - client = await HTTPClientPool().get_client( | ||
| 128 | - ip=endpoint.ip, | ||
| 129 | - port=endpoint.business_port, | ||
| 130 | - tls_config=self._config.infer_tls_config, | ||
| 131 | - ) | ||
| 132 | - response = await client.post( | ||
| 133 | - "/v1/dispatch/stop", | ||
| 134 | - json=request.model_dump(mode="json"), | ||
| 135 | - timeout=timeout, | ||
| 136 | - ) | ||
| 137 | - response.raise_for_status() | ||
| 138 | - return DispatchStopResponse.model_validate(response.json()) | ||
| 139 | - except httpx.HTTPError as e: | ||
| 140 | - logger.warning( | ||
| 141 | - "Dispatch stop failed root_request_id=%s attempt_seq=%s endpoint=%s:%s error=%s", | ||
| 142 | - attempt.root_request_id, | ||
| 143 | - attempt.attempt_seq, | ||
| 144 | - endpoint.ip, | ||
| 145 | - endpoint.business_port, | ||
| 146 | - e, | ||
| 147 | - ) | ||
| 148 | - except Exception as e: | ||
| 149 | - logger.warning( | ||
| 150 | - "Dispatch stop response invalid root_request_id=%s attempt_seq=%s error=%s", | ||
| 151 | - attempt.root_request_id, | ||
| 152 | - attempt.attempt_seq, | ||
| 153 | - e, | ||
| 154 | - ) | ||
| 155 | - return None | ||
| @@ -35,7 +35,6 @@ from motor.coordinator.models.constants import ( | |||
| 35 | OpenAIField, | 35 | OpenAIField, |
| 36 | REQUEST_ID_KEY, | 36 | REQUEST_ID_KEY, |
| 37 | ) | 37 | ) |
| 38 | -from motor.common.resources.dispatch import MOTOR_DISPATCH_KEY | ||
| 39 | from motor.coordinator.models.response import ErrorResponse | 38 | from motor.coordinator.models.response import ErrorResponse |
| 40 | from motor.coordinator.domain import ScheduledResource | 39 | from motor.coordinator.domain import ScheduledResource |
| 41 | from motor.coordinator.models.request import RequestInfo, ReqState | 40 | from motor.coordinator.models.request import RequestInfo, ReqState |
| @@ -106,7 +105,7 @@ class RequestLoggerAdapter(logging.LoggerAdapter): | |||
| 106 | 105 | ||
| 107 | 106 | ||
| 108 | class RecomputeState: | 107 | class RecomputeState: |
| 109 | - """Per-request recompute counters and flags (PD/CDP routers).""" | 108 | + """Per-request recompute counters and flags for P/D routers.""" |
| 110 | 109 | ||
| 111 | retry_count: int = 0 | 110 | retry_count: int = 0 |
| 112 | wants_retry: bool = False | 111 | wants_retry: bool = False |
| @@ -180,9 +179,8 @@ class BaseRouter(ABC): | |||
| 180 | return p_req | 179 | return p_req |
| 181 | 180 | ||
| 182 | def _forward_request_id(self, req_data: dict) -> str: | 181 | def _forward_request_id(self, req_data: dict) -> str: |
| 183 | - dispatch_data = req_data.get(MOTOR_DISPATCH_KEY) | 182 | + for field in ("request_id", "rid"): |
| 184 | - if isinstance(dispatch_data, dict): | 183 | + engine_request_id = req_data.get(field) |
| 185 | - engine_request_id = dispatch_data.get("engine_request_id") | ||
| 186 | if isinstance(engine_request_id, str) and engine_request_id: | 184 | if isinstance(engine_request_id, str) and engine_request_id: |
| 187 | return engine_request_id | 185 | return engine_request_id |
| 188 | return self.req_info.req_id | 186 | return self.req_info.req_id |
| @@ -642,7 +640,7 @@ class BaseRouter(ABC): | |||
| 642 | trace_obj = self.req_info.trace_obj | 640 | trace_obj = self.req_info.trace_obj |
| 643 | headers = trace_obj.get_trace_headers_dict(self.is_meta) | 641 | headers = trace_obj.get_trace_headers_dict(self.is_meta) |
| 644 | trace_context = TracerManager().extract_trace_context(headers) | 642 | trace_context = TracerManager().extract_trace_context(headers) |
| 645 | - with TracerManager().tracer.start_as_current_span("CDP_Encode", context=trace_context) as span: | 643 | + with TracerManager().tracer.start_as_current_span("PD_Encode", context=trace_context) as span: |
| 646 | self.is_meta = True | 644 | self.is_meta = True |
| 647 | trace_obj.meta_span = span | 645 | trace_obj.meta_span = span |
| 648 | trace_obj.meta_trace_headers = TracerManager().inject_trace_context() | 646 | trace_obj.meta_trace_headers = TracerManager().inject_trace_context() |
| @@ -20,16 +20,6 @@ from fastapi.responses import JSONResponse, Response | |||
| 20 | 20 | ||
| 21 | import motor.common.utils.error as cancel_error | 21 | import motor.common.utils.error as cancel_error |
| 22 | from motor.common.utils.error import RequestCancelledError | 22 | from motor.common.utils.error import RequestCancelledError |
| 23 | -from motor.common.resources.dispatch import ( | ||
| 24 | - DispatchPlan, | ||
| 25 | - DispatchStopReason, | ||
| 26 | - dispatch_plans_from_capabilities, | ||
| 27 | - MOTOR_DISPATCH_KEY, | ||
| 28 | - MOTOR_PREFILL_RESULT_KEY, | ||
| 29 | - PrefillResult, | ||
| 30 | - PrefillResultStatus, | ||
| 31 | - PrefillContextBudget, | ||
| 32 | -) | ||
| 33 | from motor.common.resources.endpoint import WorkloadAction | 23 | from motor.common.resources.endpoint import WorkloadAction |
| 34 | from motor.common.resources.instance import PDRole | 24 | from motor.common.resources.instance import PDRole |
| 35 | from motor.config.coordinator import CoordinatorConfig | 25 | from motor.config.coordinator import CoordinatorConfig |
| @@ -39,19 +29,13 @@ from motor.coordinator.domain import ( | |||
| 39 | UpdateWorkloadParams, | 29 | UpdateWorkloadParams, |
| 40 | ) | 30 | ) |
| 41 | from motor.coordinator.domain.request_manager import RequestManager | 31 | from motor.coordinator.domain.request_manager import RequestManager |
| 42 | -from motor.coordinator.models.constants import OpenAIField | ||
| 43 | from motor.coordinator.models.request import RequestInfo, ReqState | 32 | from motor.coordinator.models.request import RequestInfo, ReqState |
| 44 | from motor.coordinator.router.dispatch_session import ( | 33 | from motor.coordinator.router.dispatch_session import ( |
| 45 | AttemptContext, | 34 | AttemptContext, |
| 46 | AttemptState, | 35 | AttemptState, |
| 36 | + AttemptStopReason, | ||
| 47 | PDDispatchSession, | 37 | PDDispatchSession, |
| 48 | ) | 38 | ) |
| 49 | -from motor.coordinator.router.dispatch_capability import select_dispatch_plan_for_pair | ||
| 50 | -from motor.coordinator.router.sglang_native_dispatch import ( | ||
| 51 | - inject_sglang_pd_fields, | ||
| 52 | - is_sglang_resource, | ||
| 53 | -) | ||
| 54 | -from motor.coordinator.router.stop_client import DispatchStopClient | ||
| 55 | from motor.coordinator.router.strategies.base import BaseRouter, check_cancel_error | 39 | from motor.coordinator.router.strategies.base import BaseRouter, check_cancel_error |
| 56 | from motor.coordinator.router.strategies.pd_hybrid import PDHybridRouter | 40 | from motor.coordinator.router.strategies.pd_hybrid import PDHybridRouter |
| 57 | from motor.coordinator.router.rescheduler.rescheduler import ( | 41 | from motor.coordinator.router.rescheduler.rescheduler import ( |
| @@ -60,9 +44,25 @@ from motor.coordinator.router.rescheduler.rescheduler import ( | |||
| 60 | ) | 44 | ) |
| 61 | from motor.coordinator.router.workload import WorkloadActionHandler | 45 | from motor.coordinator.router.workload import WorkloadActionHandler |
| 62 | from motor.coordinator.router.precision_sample.request import inject_logprobs | 46 | from motor.coordinator.router.precision_sample.request import inject_logprobs |
| 47 | +from motor.coordinator.router.adapters.completion_to_chat import ( | ||
| 48 | + adapt_completion_nonstream_to_chat, | ||
| 49 | + is_completion_like_body, | ||
| 50 | +) | ||
| 51 | +from motor.coordinator.router.adapters.pd_protocol import ( | ||
| 52 | + ADAPTERS, | ||
| 53 | + CoordinationMode, | ||
| 54 | + EngineEndpointMetadata, | ||
| 55 | + EngineProtocolError, | ||
| 56 | + EngineRequest, | ||
| 57 | + LegContext, | ||
| 58 | + PDProtocolAdapter, | ||
| 59 | + PrefillMetadata, | ||
| 60 | +) | ||
| 63 | from motor.coordinator.router.adapters.stream import ( | 61 | from motor.coordinator.router.adapters.stream import ( |
| 64 | parse_stream_chunk_json, | 62 | parse_stream_chunk_json, |
| 65 | encode_stream_chunk_bytes, | 63 | encode_stream_chunk_bytes, |
| 64 | + strip_nonstream_response_body_for_client, | ||
| 65 | + strip_stream_chunk_bytes_for_client, | ||
| 66 | ) | 66 | ) |
| 67 | from motor.coordinator.router.stream_response import ( | 67 | from motor.coordinator.router.stream_response import ( |
| 68 | CommitAwareStreamingResponse, | 68 | CommitAwareStreamingResponse, |
| @@ -74,6 +74,8 @@ from motor.coordinator.router.upstream_error import ( | |||
| 74 | is_retryable_upstream_error, | 74 | is_retryable_upstream_error, |
| 75 | ) | 75 | ) |
| 76 | 76 | ||
| 77 | +_SGLANG_PROTOCOL_ADAPTER = ADAPTERS["sglang"] | ||
| 78 | + | ||
| 77 | 79 | ||
| 78 | 80 | ||
| 79 | class ReleaseWorkItem: | 81 | class ReleaseWorkItem: |
| @@ -114,14 +116,12 @@ class _ReleaseTaskRecord: | |||
| 114 | 116 | ||
| 115 | 117 | ||
| 116 | class UnifiedPDRouter(BaseRouter): | 118 | class UnifiedPDRouter(BaseRouter): |
| 117 | - """Unified P/D pair router behind feature flag. | 119 | + """Native-engine P/D router. |
| 118 | 120 | ||
| 119 | - Coordinator owns lifecycle orchestration; EngineServer dispatch adapters own | 121 | + Coordinator owns lifecycle orchestration and uses stateless protocol adapters |
| 120 | - engine-specific request and response normalization. | 122 | + for native-engine request mapping. |
| 121 | """ | 123 | """ |
| 122 | 124 | ||
| 123 | - _DISPATCH_MODE = "pd_pair" | ||
| 124 | - | ||
| 125 | # Background workload-release RPC retry policy. | 125 | # Background workload-release RPC retry policy. |
| 126 | _RELEASE_RPC_ATTEMPTS = 3 | 126 | _RELEASE_RPC_ATTEMPTS = 3 |
| 127 | _RELEASE_RPC_BACKOFF_BASE_S = 0.05 | 127 | _RELEASE_RPC_BACKOFF_BASE_S = 0.05 |
| @@ -177,6 +177,30 @@ class UnifiedPDRouter(BaseRouter): | |||
| 177 | self._capture_prompt_tokens_details(body) | 177 | self._capture_prompt_tokens_details(body) |
| 178 | self.req_info.update_state(ReqState.PREFILL_END) | 178 | self.req_info.update_state(ReqState.PREFILL_END) |
| 179 | 179 | ||
| 180 | + def _record_bootstrap_prefill_response( | ||
| 181 | + self, | ||
| 182 | + attempt: AttemptContext, | ||
| 183 | + response: httpx.Response, | ||
| 184 | + ) -> None: | ||
| 185 | + if not self._is_native_sglang_attempt(attempt): | ||
| 186 | + self._record_prefill_complete(response.json()) | ||
| 187 | + return | ||
| 188 | + | ||
| 189 | + try: | ||
| 190 | + body = response.json() | ||
| 191 | + except ValueError: | ||
| 192 | + # SGLang's native gateway forwards the client's stream flag to both | ||
| 193 | + # legs. A streaming P response is internal: consume its SSE frames, | ||
| 194 | + # retain usage metadata when present, and never forward P chunks. | ||
| 195 | + for line in response.content.splitlines(): | ||
| 196 | + frame_body = parse_stream_chunk_json(line, self.logger) | ||
| 197 | + if frame_body is not None: | ||
| 198 | + self._capture_prompt_tokens_details(frame_body) | ||
| 199 | + self.req_info.update_state(ReqState.PREFILL_END) | ||
| 200 | + return | ||
| 201 | + | ||
| 202 | + self._record_prefill_complete(body) | ||
| 203 | + | ||
| 180 | def _merge_prompt_tokens_details(self, body: dict[str, Any]) -> bool: | 204 | def _merge_prompt_tokens_details(self, body: dict[str, Any]) -> bool: |
| 181 | details = self.req_info.prompt_tokens_details | 205 | details = self.req_info.prompt_tokens_details |
| 182 | usage = body.get("usage") | 206 | usage = body.get("usage") |
| @@ -199,6 +223,60 @@ class UnifiedPDRouter(BaseRouter): | |||
| 199 | return chunk | 223 | return chunk |
| 200 | return encode_stream_chunk_bytes(chunk, chunk_json) | 224 | return encode_stream_chunk_bytes(chunk, chunk_json) |
| 201 | 225 | ||
| 226 | + | ||
| 227 | + def _is_native_sglang_attempt(attempt: AttemptContext | None) -> bool: | ||
| 228 | + return UnifiedPDRouter._adapter_for_attempt(attempt) is _SGLANG_PROTOCOL_ADAPTER | ||
| 229 | + | ||
| 230 | + | ||
| 231 | + def _adapter_for_attempt(attempt: AttemptContext | None) -> PDProtocolAdapter | None: | ||
| 232 | + if attempt is None or attempt.prefill_resource is None: | ||
| 233 | + return None | ||
| 234 | + engine_type = getattr(attempt.prefill_resource.instance, "engine_type", None) | ||
| 235 | + if not isinstance(engine_type, str): | ||
| 236 | + return None | ||
| 237 | + normalized = engine_type.strip().lower() | ||
| 238 | + return ADAPTERS.get(normalized) | ||
| 239 | + | ||
| 240 | + | ||
| 241 | + def _require_adapter_for_attempt(attempt: AttemptContext) -> PDProtocolAdapter: | ||
| 242 | + adapter = UnifiedPDRouter._adapter_for_attempt(attempt) | ||
| 243 | + if adapter is not None: | ||
| 244 | + return adapter | ||
| 245 | + engine_type = None | ||
| 246 | + if attempt.prefill_resource is not None: | ||
| 247 | + engine_type = getattr(attempt.prefill_resource.instance, "engine_type", None) | ||
| 248 | + raise RuntimeError(f"Unsupported native engine type for P/D coordination: {engine_type!r}") | ||
| 249 | + | ||
| 250 | + | ||
| 251 | + def _strip_native_internal_fields_from_body(body: dict[str, Any], attempt: AttemptContext) -> None: | ||
| 252 | + adapter = UnifiedPDRouter._adapter_for_attempt(attempt) | ||
| 253 | + if adapter is None: | ||
| 254 | + return | ||
| 255 | + for field in adapter.internal_response_fields: | ||
| 256 | + body.pop(field, None) | ||
| 257 | + | ||
| 258 | + def _strip_native_internal_fields_from_stream_chunk( | ||
| 259 | + self, | ||
| 260 | + chunk: bytes, | ||
| 261 | + attempt: AttemptContext, | ||
| 262 | + ) -> bytes: | ||
| 263 | + adapter = self._adapter_for_attempt(attempt) | ||
| 264 | + if adapter is None: | ||
| 265 | + return chunk | ||
| 266 | + fields = adapter.internal_response_fields | ||
| 267 | + if not any(f'"{field}"'.encode() in chunk for field in fields): | ||
| 268 | + return chunk | ||
| 269 | + chunk_json = parse_stream_chunk_json(chunk, self.logger) | ||
| 270 | + if chunk_json is None: | ||
| 271 | + self.logger.warning("Dropping invalid native stream chunk containing internal protocol fields") | ||
| 272 | + return b"" | ||
| 273 | + mutated = False | ||
| 274 | + for field in fields: | ||
| 275 | + if field in chunk_json: | ||
| 276 | + chunk_json.pop(field) | ||
| 277 | + mutated = True | ||
| 278 | + return encode_stream_chunk_bytes(chunk, chunk_json) if mutated else chunk | ||
| 279 | + | ||
| 202 | async def handle_request(self) -> Response: | 280 | async def handle_request(self) -> Response: |
| 203 | await self.do_encode() | 281 | await self.do_encode() |
| 204 | self.is_meta = False | 282 | self.is_meta = False |
| @@ -235,6 +313,8 @@ class UnifiedPDRouter(BaseRouter): | |||
| 235 | Shared by the non-stream / stream-restart / stream-resume fallback gates so the | 313 | Shared by the non-stream / stream-restart / stream-resume fallback gates so the |
| 236 | ``get_unblocked_instances`` availability probe stays in one place. | 314 | ``get_unblocked_instances`` availability probe stays in one place. |
| 237 | """ | 315 | """ |
| 316 | + if not self.config.scheduler_config.enable_pd_separation_fallback_to_hybrid: | ||
| 317 | + return False | ||
| 238 | get_unblocked = getattr(self._scheduler, "get_unblocked_instances", None) | 318 | get_unblocked = getattr(self._scheduler, "get_unblocked_instances", None) |
| 239 | if get_unblocked is None: | 319 | if get_unblocked is None: |
| 240 | return False | 320 | return False |
| @@ -244,12 +324,12 @@ class UnifiedPDRouter(BaseRouter): | |||
| 244 | unblocked_p = await get_unblocked(PDRole.ROLE_P) | 324 | unblocked_p = await get_unblocked(PDRole.ROLE_P) |
| 245 | return bool(unblocked_u or unblocked_p) | 325 | return bool(unblocked_u or unblocked_p) |
| 246 | 326 | ||
| 247 | - async def _should_trigger_hybrid_fallback_nonstream(self) -> bool: | 327 | + async def _should_use_hybrid_fallback_nonstream(self) -> bool: |
| 248 | if self.req_info.req_data.get("stream", False): | 328 | if self.req_info.req_data.get("stream", False): |
| 249 | return False | 329 | return False |
| 250 | return await self._hybrid_fallback_feasible() | 330 | return await self._hybrid_fallback_feasible() |
| 251 | 331 | ||
| 252 | - async def _should_trigger_hybrid_fallback_stream_restart(self) -> bool: | 332 | + async def _should_use_hybrid_fallback_stream_restart(self) -> bool: |
| 253 | if not self.req_info.req_data.get("stream", False): | 333 | if not self.req_info.req_data.get("stream", False): |
| 254 | return False | 334 | return False |
| 255 | if self._hybrid_stream_fallback_attempted: | 335 | if self._hybrid_stream_fallback_attempted: |
| @@ -260,7 +340,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 260 | return False | 340 | return False |
| 261 | return await self._hybrid_fallback_feasible() | 341 | return await self._hybrid_fallback_feasible() |
| 262 | 342 | ||
| 263 | - async def _should_trigger_hybrid_fallback_stream_resume(self) -> bool: | 343 | + async def _should_use_hybrid_fallback_stream_resume(self) -> bool: |
| 264 | """Post-commit continuation gate: visible output already sent, decode pool gone. | 344 | """Post-commit continuation gate: visible output already sent, decode pool gone. |
| 265 | 345 | ||
| 266 | Requires a replayable token cache (prompt + generated ids) so a single hybrid | 346 | Requires a replayable token cache (prompt + generated ids) so a single hybrid |
| @@ -367,15 +447,12 @@ class UnifiedPDRouter(BaseRouter): | |||
| 367 | trace_obj = self.req_info.trace_obj | 447 | trace_obj = self.req_info.trace_obj |
| 368 | with self._trace_span("UnifiedPD_Stream", True): | 448 | with self._trace_span("UnifiedPD_Stream", True): |
| 369 | max_retry = max(self.config.exception_config.transport_retry_limit, 1) | 449 | max_retry = max(self.config.exception_config.transport_retry_limit, 1) |
| 370 | - session = PDDispatchSession( | 450 | + session = PDDispatchSession(self.req_info.req_id) |
| 371 | - self.req_info.req_id, | ||
| 372 | - prefill_context_budget=self._prefill_context_budget(), | ||
| 373 | - ) | ||
| 374 | 451 | ||
| 375 | async with self._manage_request_context(): | 452 | async with self._manage_request_context(): |
| 376 | for attempt_index in range(max_retry): | 453 | for attempt_index in range(max_retry): |
| 377 | attempt: AttemptContext | None = None | 454 | attempt: AttemptContext | None = None |
| 378 | - cleanup_reason = DispatchStopReason.OTHER | 455 | + cleanup_reason = AttemptStopReason.OTHER |
| 379 | try: | 456 | try: |
| 380 | self._active_retry_plan = None | 457 | self._active_retry_plan = None |
| 381 | if attempt_index > 0: | 458 | if attempt_index > 0: |
| @@ -387,7 +464,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 387 | if not self._stream_commit_controller.commit_sealed: | 464 | if not self._stream_commit_controller.commit_sealed: |
| 388 | self._stream_commit_controller.begin_attempt(attempt.attempt_seq) | 465 | self._stream_commit_controller.begin_attempt(attempt.attempt_seq) |
| 389 | attempt.register_canceller() | 466 | attempt.register_canceller() |
| 390 | - dispatch_plan = self._select_dispatch_plan(attempt) | 467 | + coordination_mode = self._select_coordination_mode(attempt) |
| 391 | if attempt_index > 0: | 468 | if attempt_index > 0: |
| 392 | self.logger.warning( | 469 | self.logger.warning( |
| 393 | "Rescheduling %d/%d: P=[%s], D=[%s]", | 470 | "Rescheduling %d/%d: P=[%s], D=[%s]", |
| @@ -397,7 +474,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 397 | self._resource_label(attempt.decode_resource), | 474 | self._resource_label(attempt.decode_resource), |
| 398 | ) | 475 | ) |
| 399 | attempt.transition(AttemptState.DISPATCHING) | 476 | attempt.transition(AttemptState.DISPATCHING) |
| 400 | - async with aclosing(self._run_stream_attempt(attempt, dispatch_plan)) as attempt_stream: | 477 | + async with aclosing(self._run_stream_attempt(attempt, coordination_mode)) as attempt_stream: |
| 401 | async for chunk in attempt_stream: | 478 | async for chunk in attempt_stream: |
| 402 | attempt.transition(AttemptState.FIRST_VISIBLE) | 479 | attempt.transition(AttemptState.FIRST_VISIBLE) |
| 403 | yield chunk | 480 | yield chunk |
| @@ -414,7 +491,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 414 | self.logger.info(trace_obj.set_end_and_ttft_tpot()) | 491 | self.logger.info(trace_obj.set_end_and_ttft_tpot()) |
| 415 | return | 492 | return |
| 416 | except GeneratorExit: | 493 | except GeneratorExit: |
| 417 | - cleanup_reason = DispatchStopReason.CLIENT_DISCONNECT | 494 | + cleanup_reason = AttemptStopReason.CLIENT_DISCONNECT |
| 418 | raise | 495 | raise |
| 419 | except (asyncio.CancelledError, Exception) as e: | 496 | except (asyncio.CancelledError, Exception) as e: |
| 420 | error, retry = await self._process_response_error( | 497 | error, retry = await self._process_response_error( |
| @@ -423,14 +500,14 @@ class UnifiedPDRouter(BaseRouter): | |||
| 423 | e, | 500 | e, |
| 424 | allow_retry=self._stream_retry_allowed(), | 501 | allow_retry=self._stream_retry_allowed(), |
| 425 | ) | 502 | ) |
| 426 | - if await self._should_trigger_hybrid_fallback_stream_restart(): | 503 | + if await self._should_use_hybrid_fallback_stream_restart(): |
| 427 | async with aclosing( | 504 | async with aclosing( |
| 428 | self._run_hybrid_fallback_stream_restart(attempt, attempt_index) | 505 | self._run_hybrid_fallback_stream_restart(attempt, attempt_index) |
| 429 | ) as fallback_stream: | 506 | ) as fallback_stream: |
| 430 | async for chunk in fallback_stream: | 507 | async for chunk in fallback_stream: |
| 431 | yield chunk | 508 | yield chunk |
| 432 | return | 509 | return |
| 433 | - if await self._should_trigger_hybrid_fallback_stream_resume(): | 510 | + if await self._should_use_hybrid_fallback_stream_resume(): |
| 434 | async with aclosing( | 511 | async with aclosing( |
| 435 | self._run_hybrid_fallback_stream_resume(attempt, attempt_index) | 512 | self._run_hybrid_fallback_stream_resume(attempt, attempt_index) |
| 436 | ) as fallback_stream: | 513 | ) as fallback_stream: |
| @@ -453,10 +530,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 453 | trace_obj = self.req_info.trace_obj | 530 | trace_obj = self.req_info.trace_obj |
| 454 | with self._trace_span("UnifiedPD", False): | 531 | with self._trace_span("UnifiedPD", False): |
| 455 | max_retry = max(self.config.exception_config.transport_retry_limit, 1) | 532 | max_retry = max(self.config.exception_config.transport_retry_limit, 1) |
| 456 | - session = PDDispatchSession( | 533 | + session = PDDispatchSession(self.req_info.req_id) |
| 457 | - self.req_info.req_id, | ||
| 458 | - prefill_context_budget=self._prefill_context_budget(), | ||
| 459 | - ) | ||
| 460 | 534 | ||
| 461 | async with self._manage_request_context(): | 535 | async with self._manage_request_context(): |
| 462 | for attempt_index in range(max_retry): | 536 | for attempt_index in range(max_retry): |
| @@ -464,7 +538,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 464 | try: | 538 | try: |
| 465 | attempt = await self._create_attempt(session) | 539 | attempt = await self._create_attempt(session) |
| 466 | attempt.register_canceller() | 540 | attempt.register_canceller() |
| 467 | - dispatch_plan = self._select_dispatch_plan(attempt) | 541 | + coordination_mode = self._select_coordination_mode(attempt) |
| 468 | if attempt_index > 0: | 542 | if attempt_index > 0: |
| 469 | self.rescheduler.retry_count = attempt_index | 543 | self.rescheduler.retry_count = attempt_index |
| 470 | self.logger.warning( | 544 | self.logger.warning( |
| @@ -475,14 +549,21 @@ class UnifiedPDRouter(BaseRouter): | |||
| 475 | self._resource_label(attempt.decode_resource), | 549 | self._resource_label(attempt.decode_resource), |
| 476 | ) | 550 | ) |
| 477 | attempt.transition(AttemptState.DISPATCHING) | 551 | attempt.transition(AttemptState.DISPATCHING) |
| 478 | - body = await self._run_nonstream_attempt(attempt, dispatch_plan) | 552 | + body = await self._run_nonstream_attempt(attempt, coordination_mode) |
| 479 | attempt.unregister_canceller() | 553 | attempt.unregister_canceller() |
| 480 | attempt.transition(AttemptState.DONE) | 554 | attempt.transition(AttemptState.DONE) |
| 481 | self._merge_prompt_tokens_details(body) | 555 | self._merge_prompt_tokens_details(body) |
| 556 | + if "chat" in self.req_info.effective_entry_api() and is_completion_like_body(body): | ||
| 557 | + adapt_completion_nonstream_to_chat(body, req_id=self.req_info.req_id) | ||
| 558 | + self._strip_native_internal_fields_from_body(body, attempt) | ||
| 559 | + strip_nonstream_response_body_for_client( | ||
| 560 | + body, | ||
| 561 | + client_return_token_ids=self.req_info.client_expects_token_ids, | ||
| 562 | + ) | ||
| 482 | return JSONResponse(content=body) | 563 | return JSONResponse(content=body) |
| 483 | except (asyncio.CancelledError, Exception) as e: | 564 | except (asyncio.CancelledError, Exception) as e: |
| 484 | error, retry = await self._process_response_error(attempt, attempt_index, e) | 565 | error, retry = await self._process_response_error(attempt, attempt_index, e) |
| 485 | - if await self._should_trigger_hybrid_fallback_nonstream(): | 566 | + if await self._should_use_hybrid_fallback_nonstream(): |
| 486 | return await self._run_hybrid_fallback_nonstream() | 567 | return await self._run_hybrid_fallback_nonstream() |
| 487 | if not retry: | 568 | if not retry: |
| 488 | raise error | 569 | raise error |
| @@ -512,7 +593,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 512 | error = RequestCancelledError(reason_str) | 593 | error = RequestCancelledError(reason_str) |
| 513 | label = f"Unified PD cancelled {attempt_index}/{max_retry}" | 594 | label = f"Unified PD cancelled {attempt_index}/{max_retry}" |
| 514 | else: | 595 | else: |
| 515 | - reason = DispatchStopReason.PEER_FAILED | 596 | + reason = AttemptStopReason.PEER_FAILED |
| 516 | reason_str = str(error) | 597 | reason_str = str(error) |
| 517 | retry = allow_retry and attempt_index < max_retry - 1 | 598 | retry = allow_retry and attempt_index < max_retry - 1 |
| 518 | if isinstance(error, HTTPException): | 599 | if isinstance(error, HTTPException): |
| @@ -542,12 +623,12 @@ class UnifiedPDRouter(BaseRouter): | |||
| 542 | return error, retry | 623 | return error, retry |
| 543 | 624 | ||
| 544 | 625 | ||
| 545 | - def _cancel_stop_reason(reason: str) -> DispatchStopReason: | 626 | + def _cancel_stop_reason(reason: str) -> AttemptStopReason: |
| 546 | if reason.startswith(cancel_error.NODE_FAULT): | 627 | if reason.startswith(cancel_error.NODE_FAULT): |
| 547 | - return DispatchStopReason.PEER_FAILED | 628 | + return AttemptStopReason.PEER_FAILED |
| 548 | if reason == cancel_error.CLIENT_DISCONNECT: | 629 | if reason == cancel_error.CLIENT_DISCONNECT: |
| 549 | - return DispatchStopReason.CLIENT_DISCONNECT | 630 | + return AttemptStopReason.CLIENT_DISCONNECT |
| 550 | - return DispatchStopReason.OTHER | 631 | + return AttemptStopReason.OTHER |
| 551 | 632 | ||
| 552 | 633 | ||
| 553 | def _resource_label(resource: ScheduledResource | None) -> str: | 634 | def _resource_label(resource: ScheduledResource | None) -> str: |
| @@ -557,24 +638,20 @@ class UnifiedPDRouter(BaseRouter): | |||
| 557 | 638 | ||
| 558 | async def _create_attempt(self, session: PDDispatchSession) -> AttemptContext: | 639 | async def _create_attempt(self, session: PDDispatchSession) -> AttemptContext: |
| 559 | attempt_seq = session._attempt_seq + 1 | 640 | attempt_seq = session._attempt_seq + 1 |
| 560 | - consumed_output_tokens = ( | ||
| 561 | - self._active_retry_plan.cached_output_tokens if self._active_retry_plan is not None else 0 | ||
| 562 | - ) | ||
| 563 | p_resource = await self._prepare_attempt_resource(PDRole.ROLE_P, attempt_seq) | 641 | p_resource = await self._prepare_attempt_resource(PDRole.ROLE_P, attempt_seq) |
| 564 | 642 | ||
| 565 | # Handoff connectors (CPCD-style) do not need a concrete decode endpoint while prefill runs. | 643 | # Handoff connectors (CPCD-style) do not need a concrete decode endpoint while prefill runs. |
| 566 | # Allocate D after prefill completes so long prompts do not reserve stale decode workload for | 644 | # Allocate D after prefill completes so long prompts do not reserve stale decode workload for |
| 567 | # the entire prefill window. Concurrent connectors still allocate both legs up front. | 645 | # the entire prefill window. Concurrent connectors still allocate both legs up front. |
| 568 | if self._should_defer_decode_allocation(p_resource): | 646 | if self._should_defer_decode_allocation(p_resource): |
| 569 | - return session.new_attempt( | 647 | + return session.new_attempt(p_resource, None, self.config) |
| 570 | - p_resource, | ||
| 571 | - None, | ||
| 572 | - self.config, | ||
| 573 | - consumed_output_tokens=consumed_output_tokens, | ||
| 574 | - ) | ||
| 575 | 648 | ||
| 576 | try: | 649 | try: |
| 577 | - d_resource = await self._prepare_attempt_resource(PDRole.ROLE_D, attempt_seq) | 650 | + d_resource = await self._prepare_attempt_resource( |
| 651 | + PDRole.ROLE_D, | ||
| 652 | + attempt_seq, | ||
| 653 | + required_engine_type=str(p_resource.instance.engine_type), | ||
| 654 | + ) | ||
| 578 | except Exception as e: | 655 | except Exception as e: |
| 579 | error_message = ( | 656 | error_message = ( |
| 580 | f"Unified PD D allocation failed after P allocated " | 657 | f"Unified PD D allocation failed after P allocated " |
| @@ -584,12 +661,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 584 | self.logger.warning(error_message) | 661 | self.logger.warning(error_message) |
| 585 | await self._release_attempt_resource(p_resource, attempt_seq, WorkloadAction.RELEASE_TOKENS) | 662 | await self._release_attempt_resource(p_resource, attempt_seq, WorkloadAction.RELEASE_TOKENS) |
| 586 | raise | 663 | raise |
| 587 | - return session.new_attempt( | 664 | + return session.new_attempt(p_resource, d_resource, self.config) |
| 588 | - p_resource, | ||
| 589 | - d_resource, | ||
| 590 | - self.config, | ||
| 591 | - consumed_output_tokens=consumed_output_tokens, | ||
| 592 | - ) | ||
| 593 | 665 | ||
| 594 | 666 | ||
| 595 | def _should_defer_decode_allocation( | 667 | def _should_defer_decode_allocation( |
| @@ -597,22 +669,25 @@ class UnifiedPDRouter(BaseRouter): | |||
| 597 | ) -> bool: | 669 | ) -> bool: |
| 598 | if prefill_resource is None: | 670 | if prefill_resource is None: |
| 599 | return False | 671 | return False |
| 600 | - plans = dispatch_plans_from_capabilities(getattr(prefill_resource.instance, "dispatch_capabilities", None)) | 672 | + engine_type = getattr(prefill_resource.instance, "engine_type", None) |
| 601 | - return DispatchPlan.PREFILL_HANDOFF_DECODE in plans and DispatchPlan.CONCURRENT_ENGINE_SYNC not in plans | 673 | + if not isinstance(engine_type, str): |
| 674 | + return False | ||
| 675 | + adapter = ADAPTERS.get(engine_type.strip().lower()) | ||
| 676 | + return adapter is not None and adapter.coordination_mode == CoordinationMode.HANDOFF | ||
| 602 | 677 | ||
| 603 | async def _run_stream_attempt( | 678 | async def _run_stream_attempt( |
| 604 | - self, attempt: AttemptContext, dispatch_plan: DispatchPlan | 679 | + self, attempt: AttemptContext, coordination_mode: CoordinationMode |
| 605 | ) -> AsyncGenerator[str, None]: | 680 | ) -> AsyncGenerator[str, None]: |
| 606 | - if dispatch_plan == DispatchPlan.PREFILL_HANDOFF_DECODE: | 681 | + if coordination_mode == CoordinationMode.HANDOFF: |
| 607 | run_func = self._run_handoff_stream_attempt | 682 | run_func = self._run_handoff_stream_attempt |
| 608 | else: | 683 | else: |
| 609 | - run_func = self._run_concurrent_stream_attempt | 684 | + run_func = self._run_bootstrap_stream_attempt |
| 610 | async with aclosing(run_func(attempt)) as attempt_stream: | 685 | async with aclosing(run_func(attempt)) as attempt_stream: |
| 611 | async for chunk in attempt_stream: | 686 | async for chunk in attempt_stream: |
| 612 | yield chunk | 687 | yield chunk |
| 613 | return | 688 | return |
| 614 | 689 | ||
| 615 | - async def _run_concurrent_stream_attempt(self, attempt: AttemptContext) -> AsyncGenerator[str, None]: | 690 | + async def _run_bootstrap_stream_attempt(self, attempt: AttemptContext) -> AsyncGenerator[str, None]: |
| 616 | attempt.transition(AttemptState.ACTIVE) | 691 | attempt.transition(AttemptState.ACTIVE) |
| 617 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) | 692 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) |
| 618 | d_req, d_api = self._request_for_attempt(attempt, PDRole.ROLE_D) | 693 | d_req, d_api = self._request_for_attempt(attempt, PDRole.ROLE_D) |
| @@ -626,15 +701,19 @@ class UnifiedPDRouter(BaseRouter): | |||
| 626 | 701 | ||
| 627 | async def prefill_task(): | 702 | async def prefill_task(): |
| 628 | try: | 703 | try: |
| 704 | + attempt.mark_dispatched(PDRole.ROLE_P.value) | ||
| 629 | response = await self.forward_request( | 705 | response = await self.forward_request( |
| 630 | p_api, p_req, p_client, self.config.exception_config.first_token_timeout | 706 | p_api, p_req, p_client, self.config.exception_config.first_token_timeout |
| 631 | ) | 707 | ) |
| 632 | - self._record_prefill_complete(response.json()) | 708 | + attempt.mark_completed(PDRole.ROLE_P.value) |
| 709 | + self._record_bootstrap_prefill_response(attempt, response) | ||
| 633 | self._stream_commit_controller.mark_ready("prefill", attempt.attempt_seq) | 710 | self._stream_commit_controller.mark_ready("prefill", attempt.attempt_seq) |
| 634 | await self._scheduler.report_cb_event(p_instance_id, "success") | 711 | await self._scheduler.report_cb_event(p_instance_id, "success") |
| 635 | except asyncio.CancelledError: # pylint: disable=try-except-raise | 712 | except asyncio.CancelledError: # pylint: disable=try-except-raise |
| 636 | raise | 713 | raise |
| 637 | except Exception as e: | 714 | except Exception as e: |
| 715 | + if isinstance(e, UpstreamHTTPError): | ||
| 716 | + attempt.mark_completed(PDRole.ROLE_P.value) | ||
| 638 | if is_cb_reportable_failure(e): | 717 | if is_cb_reportable_failure(e): |
| 639 | await self._scheduler.report_cb_event(p_instance_id, "failure") | 718 | await self._scheduler.report_cb_event(p_instance_id, "failure") |
| 640 | raise | 719 | raise |
| @@ -704,6 +783,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 704 | 783 | ||
| 705 | async def decode_task() -> None: | 784 | async def decode_task() -> None: |
| 706 | try: | 785 | try: |
| 786 | + attempt.mark_dispatched(PDRole.ROLE_D.value) | ||
| 707 | async for chunk in self.forward_stream_request( | 787 | async for chunk in self.forward_stream_request( |
| 708 | d_api, | 788 | d_api, |
| 709 | d_req, | 789 | d_req, |
| @@ -713,6 +793,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 713 | ): | 793 | ): |
| 714 | if chunk: | 794 | if chunk: |
| 715 | await queue.put(chunk) | 795 | await queue.put(chunk) |
| 796 | + attempt.mark_completed(PDRole.ROLE_D.value) | ||
| 716 | if not terminal.done(): | 797 | if not terminal.done(): |
| 717 | terminal.set_result(("done", None)) | 798 | terminal.set_result(("done", None)) |
| 718 | await self._scheduler.report_cb_event(d_instance_id, "success") | 799 | await self._scheduler.report_cb_event(d_instance_id, "success") |
| @@ -720,6 +801,8 @@ class UnifiedPDRouter(BaseRouter): | |||
| 720 | if not terminal.done(): | 801 | if not terminal.done(): |
| 721 | terminal.set_result(("cancel", e)) | 802 | terminal.set_result(("cancel", e)) |
| 722 | except Exception as e: | 803 | except Exception as e: |
| 804 | + if isinstance(e, UpstreamHTTPError): | ||
| 805 | + attempt.mark_completed(PDRole.ROLE_D.value) | ||
| 723 | if is_cb_reportable_failure(e): | 806 | if is_cb_reportable_failure(e): |
| 724 | await self._scheduler.report_cb_event(d_instance_id, "failure") | 807 | await self._scheduler.report_cb_event(d_instance_id, "failure") |
| 725 | if not terminal.done(): | 808 | if not terminal.done(): |
| @@ -810,6 +893,11 @@ class UnifiedPDRouter(BaseRouter): | |||
| 810 | ) | 893 | ) |
| 811 | else: | 894 | else: |
| 812 | value = self._merge_prompt_tokens_details_into_stream_chunk(value) | 895 | value = self._merge_prompt_tokens_details_into_stream_chunk(value) |
| 896 | + value = strip_stream_chunk_bytes_for_client( | ||
| 897 | + value, | ||
| 898 | + client_return_token_ids=self.req_info.client_expects_token_ids, | ||
| 899 | + ) | ||
| 900 | + value = self._strip_native_internal_fields_from_stream_chunk(value, attempt) | ||
| 813 | yield value | 901 | yield value |
| 814 | elif key == "done": | 902 | elif key == "done": |
| 815 | if prefill_task is not None: | 903 | if prefill_task is not None: |
| @@ -837,13 +925,15 @@ class UnifiedPDRouter(BaseRouter): | |||
| 837 | if not terminal.done(): | 925 | if not terminal.done(): |
| 838 | terminal.cancel() | 926 | terminal.cancel() |
| 839 | 927 | ||
| 840 | - async def _run_nonstream_attempt(self, attempt: AttemptContext, dispatch_plan: DispatchPlan) -> dict[str, Any]: | 928 | + async def _run_nonstream_attempt( |
| 841 | - if dispatch_plan == DispatchPlan.PREFILL_HANDOFF_DECODE: | 929 | + self, attempt: AttemptContext, coordination_mode: CoordinationMode |
| 930 | + ) -> dict[str, Any]: | ||
| 931 | + if coordination_mode == CoordinationMode.HANDOFF: | ||
| 842 | return await self._run_handoff_nonstream_attempt(attempt) | 932 | return await self._run_handoff_nonstream_attempt(attempt) |
| 843 | else: | 933 | else: |
| 844 | - return await self._run_concurrent_nonstream_attempt(attempt) | 934 | + return await self._run_bootstrap_nonstream_attempt(attempt) |
| 845 | 935 | ||
| 846 | - async def _run_concurrent_nonstream_attempt(self, attempt: AttemptContext) -> dict[str, Any]: | 936 | + async def _run_bootstrap_nonstream_attempt(self, attempt: AttemptContext) -> dict[str, Any]: |
| 847 | attempt.transition(AttemptState.ACTIVE) | 937 | attempt.transition(AttemptState.ACTIVE) |
| 848 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) | 938 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) |
| 849 | d_req, d_api = self._request_for_attempt(attempt, PDRole.ROLE_D) | 939 | d_req, d_api = self._request_for_attempt(attempt, PDRole.ROLE_D) |
| @@ -856,14 +946,18 @@ class UnifiedPDRouter(BaseRouter): | |||
| 856 | 946 | ||
| 857 | async def prefill_task(): | 947 | async def prefill_task(): |
| 858 | try: | 948 | try: |
| 949 | + attempt.mark_dispatched(PDRole.ROLE_P.value) | ||
| 859 | response = await self.forward_request( | 950 | response = await self.forward_request( |
| 860 | p_api, p_req, p_client, self.config.exception_config.first_token_timeout | 951 | p_api, p_req, p_client, self.config.exception_config.first_token_timeout |
| 861 | ) | 952 | ) |
| 862 | - self._record_prefill_complete(response.json()) | 953 | + attempt.mark_completed(PDRole.ROLE_P.value) |
| 954 | + self._record_bootstrap_prefill_response(attempt, response) | ||
| 863 | await self._scheduler.report_cb_event(p_instance_id, "success") | 955 | await self._scheduler.report_cb_event(p_instance_id, "success") |
| 864 | except asyncio.CancelledError: # pylint: disable=try-except-raise | 956 | except asyncio.CancelledError: # pylint: disable=try-except-raise |
| 865 | raise | 957 | raise |
| 866 | except Exception as e: | 958 | except Exception as e: |
| 959 | + if isinstance(e, UpstreamHTTPError): | ||
| 960 | + attempt.mark_completed(PDRole.ROLE_P.value) | ||
| 867 | if is_cb_reportable_failure(e): | 961 | if is_cb_reportable_failure(e): |
| 868 | await self._scheduler.report_cb_event(p_instance_id, "failure") | 962 | await self._scheduler.report_cb_event(p_instance_id, "failure") |
| 869 | raise | 963 | raise |
| @@ -892,14 +986,18 @@ class UnifiedPDRouter(BaseRouter): | |||
| 892 | 986 | ||
| 893 | async def decode_task() -> tuple[Any, Any]: | 987 | async def decode_task() -> tuple[Any, Any]: |
| 894 | try: | 988 | try: |
| 989 | + attempt.mark_dispatched(PDRole.ROLE_D.value) | ||
| 895 | response = await self.forward_request( | 990 | response = await self.forward_request( |
| 896 | d_api, d_req, d_client, self.config.exception_config.infer_timeout | 991 | d_api, d_req, d_client, self.config.exception_config.infer_timeout |
| 897 | ) | 992 | ) |
| 993 | + attempt.mark_completed(PDRole.ROLE_D.value) | ||
| 898 | await self._scheduler.report_cb_event(d_instance_id, "success") | 994 | await self._scheduler.report_cb_event(d_instance_id, "success") |
| 899 | return response.json(), None | 995 | return response.json(), None |
| 900 | except asyncio.CancelledError: # pylint: disable=try-except-raise | 996 | except asyncio.CancelledError: # pylint: disable=try-except-raise |
| 901 | raise | 997 | raise |
| 902 | except Exception as e: | 998 | except Exception as e: |
| 999 | + if isinstance(e, UpstreamHTTPError): | ||
| 1000 | + attempt.mark_completed(PDRole.ROLE_D.value) | ||
| 903 | if is_cb_reportable_failure(e): | 1001 | if is_cb_reportable_failure(e): |
| 904 | await self._scheduler.report_cb_event(d_instance_id, "failure") | 1002 | await self._scheduler.report_cb_event(d_instance_id, "failure") |
| 905 | return None, e | 1003 | return None, e |
| @@ -1011,7 +1109,12 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1011 | if attempt.decode_resource is not None: | 1109 | if attempt.decode_resource is not None: |
| 1012 | return False | 1110 | return False |
| 1013 | start = time.perf_counter() | 1111 | start = time.perf_counter() |
| 1014 | - d_resource = await self._prepare_attempt_resource(PDRole.ROLE_D, attempt.attempt_seq) | 1112 | + adapter = self._require_adapter_for_attempt(attempt) |
| 1113 | + d_resource = await self._prepare_attempt_resource( | ||
| 1114 | + PDRole.ROLE_D, | ||
| 1115 | + attempt.attempt_seq, | ||
| 1116 | + required_engine_type=adapter.engine_type, | ||
| 1117 | + ) | ||
| 1015 | elapsed_ms = (time.perf_counter() - start) * 1000 | 1118 | elapsed_ms = (time.perf_counter() - start) * 1000 |
| 1016 | self.logger.info( | 1119 | self.logger.info( |
| 1017 | "Scheduling latency stage=late_select_d elapsed_ms=%.2f instance_id=%s endpoint_id=%s req_id=%s", | 1120 | "Scheduling latency stage=late_select_d elapsed_ms=%.2f instance_id=%s endpoint_id=%s req_id=%s", |
| @@ -1022,12 +1125,10 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1022 | ) | 1125 | ) |
| 1023 | attempt.decode_resource = d_resource | 1126 | attempt.decode_resource = d_resource |
| 1024 | try: | 1127 | try: |
| 1025 | - dispatch_plan = select_dispatch_plan_for_pair( | 1128 | + adapter = self._require_adapter_for_attempt(attempt) |
| 1026 | - prefill=attempt.prefill_resource, | 1129 | + decode_engine_type = getattr(attempt.decode_resource.instance, "engine_type", None) |
| 1027 | - decode=attempt.decode_resource, | 1130 | + if not isinstance(decode_engine_type, str) or decode_engine_type.strip().lower() != adapter.engine_type: |
| 1028 | - ) | 1131 | + raise RuntimeError(f"{adapter.engine_type} handoff requires a matching decode instance") |
| 1029 | - if dispatch_plan != DispatchPlan.PREFILL_HANDOFF_DECODE: | ||
| 1030 | - raise RuntimeError(f"Late decode allocation selected unsupported plan: {dispatch_plan}") | ||
| 1031 | except Exception: | 1132 | except Exception: |
| 1032 | await self._release_attempt_resource( | 1133 | await self._release_attempt_resource( |
| 1033 | d_resource, | 1134 | d_resource, |
| @@ -1039,7 +1140,11 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1039 | raise | 1140 | raise |
| 1040 | return True | 1141 | return True |
| 1041 | 1142 | ||
| 1042 | - async def _await_handoff_prefill(self, attempt: AttemptContext, p_client) -> PrefillResult: | 1143 | + async def _await_handoff_prefill( |
| 1144 | + self, | ||
| 1145 | + attempt: AttemptContext, | ||
| 1146 | + p_client, | ||
| 1147 | + ) -> PrefillMetadata: | ||
| 1043 | p_instance_id = attempt.prefill_resource.instance.id | 1148 | p_instance_id = attempt.prefill_resource.instance.id |
| 1044 | 1149 | ||
| 1045 | async def prefill_task(): | 1150 | async def prefill_task(): |
| @@ -1063,15 +1168,23 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1063 | attempt: AttemptContext, | 1168 | attempt: AttemptContext, |
| 1064 | role: PDRole, | 1169 | role: PDRole, |
| 1065 | *, | 1170 | *, |
| 1066 | - prefill_result: PrefillResult | None = None, | 1171 | + prefill_result: PrefillMetadata | None = None, |
| 1067 | ) -> (dict[str, Any], str): | 1172 | ) -> (dict[str, Any], str): |
| 1173 | + req, api = self._base_request_for_attempt(role) | ||
| 1174 | + adapter = self._require_adapter_for_attempt(attempt) | ||
| 1175 | + return self._native_request_for_attempt( | ||
| 1176 | + attempt, | ||
| 1177 | + role, | ||
| 1178 | + adapter=adapter, | ||
| 1179 | + req=req, | ||
| 1180 | + api=api, | ||
| 1181 | + prefill_metadata=prefill_result, | ||
| 1182 | + ) | ||
| 1183 | + | ||
| 1184 | + def _base_request_for_attempt(self, role: PDRole) -> tuple[dict[str, Any], str]: | ||
| 1068 | api = self.req_info.entry_api | 1185 | api = self.req_info.entry_api |
| 1069 | req = self.req_info.req_data.copy() | 1186 | req = self.req_info.req_data.copy() |
| 1070 | stream = self.req_info.req_data.get("stream", False) | 1187 | stream = self.req_info.req_data.get("stream", False) |
| 1071 | - req["request_id"] = f"{attempt.root_request_id}#a{attempt.attempt_seq}" | ||
| 1072 | - if role == PDRole.ROLE_P: | ||
| 1073 | - req["stream"] = False | ||
| 1074 | - req = self._apply_prefill_params(req, set_min_tokens=False) | ||
| 1075 | if stream and self.config.exception_config.reschedule_enabled: | 1188 | if stream and self.config.exception_config.reschedule_enabled: |
| 1076 | req["return_token_ids"] = True | 1189 | req["return_token_ids"] = True |
| 1077 | if self._active_retry_plan is not None: | 1190 | if self._active_retry_plan is not None: |
| @@ -1086,82 +1199,121 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1086 | and self._sampling_manager is not None | 1199 | and self._sampling_manager is not None |
| 1087 | ): | 1200 | ): |
| 1088 | inject_logprobs(req, self.config.precision_detection_config, req_id=self.req_info.req_id) | 1201 | inject_logprobs(req, self.config.precision_detection_config, req_id=self.req_info.req_id) |
| 1089 | - # SGLang pure-native: stock launch_server needs bootstrap_* on the business | 1202 | + return req, api |
| 1090 | - # port. Do not attach _motor_dispatch (no InferEndpoint adapter in native mode). | ||
| 1091 | - # Prefill endpoint must still be held here so Decode-leg inject can read | ||
| 1092 | - # bootstrap_host from attempt.prefill_resource. | ||
| 1093 | - target = attempt.prefill_resource if role == PDRole.ROLE_P else attempt.decode_resource | ||
| 1094 | - if is_sglang_resource(target): | ||
| 1095 | - inject_sglang_pd_fields(req, attempt) | ||
| 1096 | - return (req, api) | ||
| 1097 | - req[MOTOR_DISPATCH_KEY] = attempt.dispatch_for(role, self._DISPATCH_MODE).model_dump(mode="json") | ||
| 1098 | - if prefill_result is not None: | ||
| 1099 | - req[MOTOR_PREFILL_RESULT_KEY] = prefill_result.model_dump(mode="json") | ||
| 1100 | - return (req, api) | ||
| 1101 | 1203 | ||
| 1102 | - def _prefill_context_budget(self) -> PrefillContextBudget | None: | 1204 | + def _native_request_for_attempt( |
| 1103 | - """Return the client budget before the prefill leg is rewritten to one token.""" | 1205 | + self, |
| 1104 | - for field in (OpenAIField.MAX_COMPLETION_TOKENS, OpenAIField.MAX_TOKENS): | 1206 | + attempt: AttemptContext, |
| 1105 | - value = self.req_info.req_data.get(field) | 1207 | + role: PDRole, |
| 1106 | - if isinstance(value, int) and not isinstance(value, bool) and value >= 0: | 1208 | + *, |
| 1107 | - return PrefillContextBudget( | 1209 | + adapter: PDProtocolAdapter, |
| 1108 | - max_output_tokens=value, | 1210 | + req: dict[str, Any], |
| 1109 | - parameter=field.value, | 1211 | + api: str, |
| 1110 | - ) | 1212 | + prefill_metadata: PrefillMetadata | None = None, |
| 1111 | - return None | 1213 | + ) -> tuple[dict[str, Any], str]: |
| 1214 | + context = self._native_leg_context(attempt, role, api) | ||
| 1215 | + if role == PDRole.ROLE_P: | ||
| 1216 | + engine_request = adapter.build_prefill_request(req, context) | ||
| 1217 | + else: | ||
| 1218 | + engine_request = adapter.build_decode_request(req, context, prefill_metadata) | ||
| 1219 | + return engine_request.body, engine_request.api | ||
| 1112 | 1220 | ||
| 1113 | - async def _request_prefill_result(self, attempt: AttemptContext, p_client) -> PrefillResult: | 1221 | + @staticmethod |
| 1222 | + def _native_leg_context(attempt: AttemptContext, role: PDRole, api: str) -> LegContext: | ||
| 1223 | + resource = attempt.prefill_resource if role == PDRole.ROLE_P else attempt.decode_resource | ||
| 1224 | + if resource is None: | ||
| 1225 | + raise RuntimeError(f"Missing {role.value} resource for native engine request") | ||
| 1226 | + peer_resource = attempt.decode_resource if role == PDRole.ROLE_P else attempt.prefill_resource | ||
| 1227 | + return LegContext( | ||
| 1228 | + engine_request_id=f"{attempt.root_request_id}#a{attempt.attempt_seq}", | ||
| 1229 | + pair_id=attempt.pair_id, | ||
| 1230 | + attempt_seq=attempt.attempt_seq, | ||
| 1231 | + api=api, | ||
| 1232 | + endpoint=UnifiedPDRouter._native_endpoint_metadata(resource), | ||
| 1233 | + peer_endpoint=( | ||
| 1234 | + UnifiedPDRouter._native_endpoint_metadata(peer_resource) if peer_resource is not None else None | ||
| 1235 | + ), | ||
| 1236 | + ) | ||
| 1237 | + | ||
| 1238 | + | ||
| 1239 | + def _native_endpoint_metadata(resource: ScheduledResource) -> EngineEndpointMetadata: | ||
| 1240 | + endpoint = resource.endpoint | ||
| 1241 | + return EngineEndpointMetadata( | ||
| 1242 | + host=endpoint.ip, | ||
| 1243 | + bootstrap_port=endpoint.bootstrap_port, | ||
| 1244 | + ) | ||
| 1245 | + | ||
| 1246 | + async def _request_prefill_result( | ||
| 1247 | + self, | ||
| 1248 | + attempt: AttemptContext, | ||
| 1249 | + p_client, | ||
| 1250 | + ) -> PrefillMetadata: | ||
| 1114 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) | 1251 | p_req, p_api = self._request_for_attempt(attempt, PDRole.ROLE_P) |
| 1115 | - response = await self.forward_request(p_api, p_req, p_client, self.config.exception_config.first_token_timeout) | 1252 | + attempt.mark_dispatched(PDRole.ROLE_P.value) |
| 1253 | + try: | ||
| 1254 | + response = await self.forward_request( | ||
| 1255 | + p_api, | ||
| 1256 | + p_req, | ||
| 1257 | + p_client, | ||
| 1258 | + self.config.exception_config.first_token_timeout, | ||
| 1259 | + ) | ||
| 1260 | + except UpstreamHTTPError: | ||
| 1261 | + attempt.mark_completed(PDRole.ROLE_P.value) | ||
| 1262 | + raise | ||
| 1263 | + attempt.mark_completed(PDRole.ROLE_P.value) | ||
| 1116 | response_body = response.json() | 1264 | response_body = response.json() |
| 1117 | - self._capture_prompt_tokens_details(response_body) | 1265 | + adapter = self._require_adapter_for_attempt(attempt) |
| 1118 | - prefill_result = PrefillResult.model_validate(response_body) | 1266 | + try: |
| 1119 | - self._validate_prefill_result(attempt, prefill_result, expected_status=PrefillResultStatus.COMPLETED) | 1267 | + prefill_metadata = adapter.parse_prefill_response(response_body) |
| 1268 | + except EngineProtocolError as error: | ||
| 1269 | + raise UpstreamHTTPError( | ||
| 1270 | + status_code=502, | ||
| 1271 | + body=str(error).encode("utf-8"), | ||
| 1272 | + headers={"content-type": "text/plain; charset=utf-8"}, | ||
| 1273 | + phase="prefill", | ||
| 1274 | + ) from error | ||
| 1275 | + if prefill_metadata.usage is not None: | ||
| 1276 | + self._capture_prompt_tokens_details({"usage": prefill_metadata.usage}) | ||
| 1120 | self.req_info.update_state(ReqState.PREFILL_END) | 1277 | self.req_info.update_state(ReqState.PREFILL_END) |
| 1121 | if self._stream_commit_controller is not None: | 1278 | if self._stream_commit_controller is not None: |
| 1122 | self._stream_commit_controller.mark_ready("prefill", attempt.attempt_seq) | 1279 | self._stream_commit_controller.mark_ready("prefill", attempt.attempt_seq) |
| 1123 | - return prefill_result | 1280 | + return prefill_metadata |
| 1124 | 1281 | ||
| 1125 | - @staticmethod | 1282 | + def _select_coordination_mode(self, attempt: AttemptContext) -> CoordinationMode: |
G 严重程度： 提示 问题： 混合 engine_type 的 P/D 池在实例分配后才报裸 RuntimeError（HTTP 500 且无诊断信息）。 原因： 怎么改： 把 RuntimeError 换成统一的 5xx 响应并带上实例标识（engine_type 与实例地址），或在 select_router_class 阶段对 P/D 池做一次 engine_type 一致性预检，不匹配时直接返回 503 并打 warning 日志。 tobking 19 天前 评论： 19 天前 评论： | |||
| 1126 | - def _validate_prefill_result( | 1283 | + adapter = self._require_adapter_for_attempt(attempt) |
| 1127 | - attempt: AttemptContext, | 1284 | + if adapter.coordination_mode == CoordinationMode.BOOTSTRAP and attempt.decode_resource is None: |
| 1128 | - prefill_result: PrefillResult, | 1285 | + raise RuntimeError(f"{adapter.engine_type} bootstrap requires a decode instance") |
| 1286 | + if attempt.decode_resource is not None: | ||
| 1287 | + decode_engine_type = getattr(attempt.decode_resource.instance, "engine_type", None) | ||
| 1288 | + if not isinstance(decode_engine_type, str) or decode_engine_type.strip().lower() != adapter.engine_type: | ||
| 1289 | + raise RuntimeError( | ||
| 1290 | + f"P/D engine types must match: prefill={adapter.engine_type}, decode={decode_engine_type!r}" | ||
| 1291 | + ) | ||
| 1292 | + return adapter.coordination_mode | ||
| 1293 | + | ||
| 1294 | + async def _prepare_attempt_resource( | ||
| 1295 | + self, | ||
| 1296 | + role: PDRole, | ||
| 1297 | + attempt_seq: int, | ||
| 1129 | *, | 1298 | *, |
| 1130 | - expected_status: PrefillResultStatus, | 1299 | + required_engine_type: str | None = None, |
| 1131 | - ) -> None: | 1300 | + ) -> ScheduledResource: |
| 1132 | - if ( | ||
| 1133 | - prefill_result.root_request_id != attempt.root_request_id | ||
| 1134 | - or prefill_result.pair_id != attempt.pair_id | ||
| 1135 | - or prefill_result.attempt_seq != attempt.attempt_seq | ||
| 1136 | - ): | ||
| 1137 | - raise RuntimeError("PrefillResult does not match current dispatch attempt") | ||
| 1138 | - if prefill_result.status != expected_status.value: | ||
| 1139 | - raise RuntimeError(f"Unexpected PrefillResult status: {prefill_result.status}") | ||
| 1140 | - | ||
| 1141 | - def _select_dispatch_plan(self, attempt: AttemptContext) -> DispatchPlan: | ||
| 1142 | - if attempt.decode_resource is None: | ||
| 1143 | - if self._should_defer_decode_allocation(attempt.prefill_resource): | ||
| 1144 | - return DispatchPlan.PREFILL_HANDOFF_DECODE | ||
| 1145 | - raise RuntimeError("Decode resource is required before selecting a concurrent P/D dispatch plan") | ||
| 1146 | - return select_dispatch_plan_for_pair( | ||
| 1147 | - prefill=attempt.prefill_resource, | ||
| 1148 | - decode=attempt.decode_resource, | ||
| 1149 | - ) | ||
| 1150 | - | ||
| 1151 | - async def _prepare_attempt_resource(self, role: PDRole, attempt_seq: int) -> ScheduledResource: | ||
| 1152 | self.req_info.update_state(ReqState.P_SCHEDULING if role == PDRole.ROLE_P else ReqState.D_SCHEDULING) | 1301 | self.req_info.update_state(ReqState.P_SCHEDULING if role == PDRole.ROLE_P else ReqState.D_SCHEDULING) |
| 1153 | target_instance_id = None | 1302 | target_instance_id = None |
| 1154 | constraint = self.req_info.scheduling_constraint | 1303 | constraint = self.req_info.scheduling_constraint |
| 1155 | if constraint is not None: | 1304 | if constraint is not None: |
| 1156 | target_instance_id = constraint.target_for_role(role) | 1305 | target_instance_id = constraint.target_for_role(role) |
| 1157 | - result = await self._scheduler.select_and_allocate( | 1306 | + scheduler_kwargs = {"target_instance_id": target_instance_id} |
| 1158 | - role, | 1307 | + if required_engine_type is not None: |
| 1159 | - self.req_info, | 1308 | + scheduler_kwargs["required_engine_type"] = required_engine_type |
| 1160 | - target_instance_id=target_instance_id, | 1309 | + result = await self._scheduler.select_and_allocate(role, self.req_info, **scheduler_kwargs) |
| 1161 | - ) | ||
| 1162 | if result is None: | 1310 | if result is None: |
| 1163 | error_message = f"No instance available for role {role}" | 1311 | error_message = f"No instance available for role {role}" |
| 1312 | + if required_engine_type is not None: | ||
| 1313 | + error_message += f" with engine_type={required_engine_type}" | ||
| 1164 | self.req_info.trace_obj.set_trace_error_message(error_message) | 1314 | self.req_info.trace_obj.set_trace_error_message(error_message) |
| 1315 | + if required_engine_type is not None: | ||
| 1316 | + raise HTTPException(status_code=503, detail=error_message) | ||
| 1165 | raise RuntimeError(error_message) | 1317 | raise RuntimeError(error_message) |
| 1166 | ins, endpoint, workload = result | 1318 | ins, endpoint, workload = result |
| 1167 | await self._record_attempt_workload(attempt_seq, role, workload) | 1319 | await self._record_attempt_workload(attempt_seq, role, workload) |
| @@ -1587,7 +1739,7 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1587 | return "release_d_tokens" | 1739 | return "release_d_tokens" |
| 1588 | return f"release_{role.value}_{action.value.lower()}" | 1740 | return f"release_{role.value}_{action.value.lower()}" |
| 1589 | 1741 | ||
| 1590 | - async def _stop_attempt(self, attempt: AttemptContext | None, reason: DispatchStopReason) -> None: | 1742 | + async def _stop_attempt(self, attempt: AttemptContext | None, reason: AttemptStopReason) -> None: |
G 严重程度： 建议 问题： 原因： 新版只做本地 怎么改：
恢复对 sglang 腿的 abort 调用，rid 与 若不做，至少在 PR 描述中明确标注该能力回退。 tobking 19 天前 评论： 19 天前 评论： | |||
| 1591 | if attempt is None: | 1743 | if attempt is None: |
| 1592 | return | 1744 | return |
| 1593 | async with attempt.stop_lock: | 1745 | async with attempt.stop_lock: |
| @@ -1596,15 +1748,8 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1596 | return | 1748 | return |
| 1597 | 1749 | ||
| 1598 | attempt.stop() | 1750 | attempt.stop() |
| 1751 | + await self._abort_native_attempt(attempt) | ||
| 1599 | await attempt.cancel(reason.value) | 1752 | await attempt.cancel(reason.value) |
| 1600 | - client = DispatchStopClient(self.config) | ||
| 1601 | - tasks = [] | ||
| 1602 | - if attempt.prefill_resource: | ||
| 1603 | - tasks.append(client.stop(attempt.prefill_resource, attempt, reason)) | ||
| 1604 | - if attempt.decode_resource: | ||
| 1605 | - tasks.append(client.stop(attempt.decode_resource, attempt, reason)) | ||
| 1606 | - if tasks: | ||
| 1607 | - await asyncio.gather(*tasks, return_exceptions=True) | ||
| 1608 | await self._release_attempt(attempt, wait=False) | 1753 | await self._release_attempt(attempt, wait=False) |
| 1609 | try: | 1754 | try: |
| 1610 | await self._drain_release_tasks() | 1755 | await self._drain_release_tasks() |
| @@ -1617,6 +1762,55 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1617 | ) | 1762 | ) |
| 1618 | attempt.transition(AttemptState.STOPPED) | 1763 | attempt.transition(AttemptState.STOPPED) |
| 1619 | 1764 | ||
| 1765 | + async def _abort_native_attempt(self, attempt: AttemptContext) -> None: | ||
| 1766 | + """Best-effort native cancellation; failures never mask the original stop reason.""" | ||
| 1767 | + try: | ||
| 1768 | + adapter = self._require_adapter_for_attempt(attempt) | ||
| 1769 | + except Exception: | ||
| 1770 | + return | ||
| 1771 | + aborts = [] | ||
| 1772 | + for role, resource in ( | ||
| 1773 | + (PDRole.ROLE_P, attempt.prefill_resource), | ||
| 1774 | + (PDRole.ROLE_D, attempt.decode_resource), | ||
| 1775 | + ): | ||
| 1776 | + if resource is None or not attempt.needs_abort(role.value): | ||
| 1777 | + continue | ||
| 1778 | + context = self._native_leg_context(attempt, role, self.req_info.entry_api) | ||
| 1779 | + abort_request = adapter.build_abort_request(context) | ||
| 1780 | + if abort_request is not None: | ||
| 1781 | + aborts.append(self._send_native_abort(resource, abort_request)) | ||
| 1782 | + if aborts: | ||
| 1783 | + await asyncio.gather(*aborts) | ||
| 1784 | + | ||
| 1785 | + async def _send_native_abort(self, resource: ScheduledResource, request: EngineRequest) -> None: | ||
| 1786 | + try: | ||
| 1787 | + async with self._client_for(resource) as client: | ||
| 1788 | + timeout = min(float(self.config.exception_config.first_token_timeout), 1.0) | ||
| 1789 | + response = await asyncio.wait_for( | ||
| 1790 | + client.post( | ||
| 1791 | + f"/{request.api}", | ||
| 1792 | + json=request.body, | ||
| 1793 | + timeout=timeout, | ||
| 1794 | + ), | ||
| 1795 | + timeout=timeout, | ||
| 1796 | + ) | ||
| 1797 | + if response.status_code >= 400: | ||
| 1798 | + self.logger.warning( | ||
| 1799 | + "Native abort rejected engine=%s instance_id=%s endpoint_id=%s status=%s", | ||
| 1800 | + resource.instance.engine_type, | ||
| 1801 | + resource.instance.id, | ||
| 1802 | + resource.endpoint.id, | ||
| 1803 | + response.status_code, | ||
| 1804 | + ) | ||
| 1805 | + except Exception as err: | ||
| 1806 | + self.logger.warning( | ||
| 1807 | + "Native abort failed engine=%s instance_id=%s endpoint_id=%s error=%s", | ||
| 1808 | + resource.instance.engine_type, | ||
| 1809 | + resource.instance.id, | ||
| 1810 | + resource.endpoint.id, | ||
| 1811 | + err, | ||
| 1812 | + ) | ||
| 1813 | + | ||
| 1620 | def _client_for(self, resource: ScheduledResource): | 1814 | def _client_for(self, resource: ScheduledResource): |
| 1621 | if resource is None: | 1815 | if resource is None: |
| 1622 | raise RuntimeError("Scheduled resource is missing") | 1816 | raise RuntimeError("Scheduled resource is missing") |
| @@ -1661,49 +1855,6 @@ class UnifiedPDRouter(BaseRouter): | |||
| 1661 | if await self._sampling_manager.confirm_sample((p_id, d_id), time.time()): | 1855 | if await self._sampling_manager.confirm_sample((p_id, d_id), time.time()): |
| 1662 | await self._submit_token_sample(p_id, d_id, info, attempt.decode_resource) | 1856 | await self._submit_token_sample(p_id, d_id, info, attempt.decode_resource) |
| 1663 | 1857 | ||
| 1664 | - # ------------------------------------------------------------------ | ||
| 1665 | - # Metaserver forward entry point (CDP mode: D-side prefill → P instance) | ||
| 1666 | - # ------------------------------------------------------------------ | ||
| 1667 | - | ||
| 1668 | - async def handle_metaserver_request(self) -> dict[str, Any]: | ||
| 1669 | - self.is_meta = True | ||
| 1670 | - schedule_resource: ScheduledResource = None | ||
| 1671 | - try: | ||
| 1672 | - schedule_resource = await self.prepare_resource(PDRole.ROLE_P) | ||
| 1673 | - req_data = self.req_info.req_data.copy() | ||
| 1674 | - req_data["stream"] = False | ||
| 1675 | - async with self._client_for(schedule_resource) as client: | ||
| 1676 | - response = await self.forward_request( | ||
| 1677 | - self.req_info.api, | ||
| 1678 | - req_data, | ||
| 1679 | - client, | ||
| 1680 | - self.config.exception_config.first_token_timeout, | ||
| 1681 | - ) | ||
| 1682 | - resp_json = response.json() | ||
| 1683 | - self.logger.debug("Prefill response received") | ||
| 1684 | - self._capture_prompt_tokens_details(resp_json) | ||
| 1685 | - self.req_info.update_state(ReqState.PREFILL_END) | ||
| 1686 | - if hasattr(self.req_info, "p_instance_id"): | ||
| 1687 | - self.req_info.p_instance_id = schedule_resource.instance.id | ||
| 1688 | - return resp_json | ||
| 1689 | - except asyncio.CancelledError: | ||
| 1690 | - self.req_info.trace_obj.set_trace_prompt(self.req_info.req_data) | ||
| 1691 | - self.logger.info("Metaserver request was cancelled") | ||
| 1692 | - self.req_info.cancel_scope() | ||
| 1693 | - raise | ||
| 1694 | - except Exception: | ||
| 1695 | - self.req_info.trace_obj.set_trace_prompt(self.req_info.req_data) | ||
| 1696 | - self.req_info.cancel_scope() | ||
| 1697 | - self.req_info.update_state(ReqState.EXCEPTION) | ||
| 1698 | - raise | ||
| 1699 | - finally: | ||
| 1700 | - if schedule_resource and self.req_info.state != ReqState.PREFILL_END: | ||
| 1701 | - if not await self.release_all(schedule_resource): | ||
| 1702 | - self.logger.debug( | ||
| 1703 | - "release_all(prefill) returned False instance_id=%s", | ||
| 1704 | - schedule_resource.instance.id, | ||
| 1705 | - ) | ||
| 1706 | - | ||
| 1707 | 1858 | ||
| 1708 | async def _cancel_task_quietly(task: asyncio.Task | None) -> None: | 1859 | async def _cancel_task_quietly(task: asyncio.Task | None) -> None: |
| 1709 | if task is None or task.done(): | 1860 | if task is None or task.done(): |
| @@ -80,9 +80,8 @@ def is_cb_reportable_failure(error: BaseException) -> bool: | |||
| 80 | penalize the instance. HTTP 5xx errors indicate instance health issues. | 80 | penalize the instance. HTTP 5xx errors indicate instance health issues. |
| 81 | Connection-level errors always indicate instance issues. | 81 | Connection-level errors always indicate instance issues. |
| 82 | 82 | ||
| 83 | - The engine_server (serving_error.py) is the single source of truth for classifying exceptions | 83 | + Native engines classify request-validation failures as 4xx and engine faults as 5xx. |
| 84 | - as client errors vs engine faults: it maps all known request-validation errors to 4xx before | 84 | + The coordinator therefore only needs to inspect the upstream status code. |
| 85 | - they reach the coordinator. The coordinator therefore only needs to inspect the status code. | ||
| 86 | """ | 85 | """ |
| 87 | if isinstance(error, UpstreamHTTPError): | 86 | if isinstance(error, UpstreamHTTPError): |
| 88 | return error.status_code >= 500 | 87 | return error.status_code >= 500 |
| @@ -20,9 +20,6 @@ from typing import Any, Awaitable, Callable | |||
| 20 | import msgspec | 20 | import msgspec |
| 21 | import zmq | 21 | import zmq |
| 22 | 22 | ||
| 23 | -from motor.common.resources.dispatch import ( | ||
| 24 | - has_compatible_dispatch_pair, | ||
| 25 | -) | ||
| 26 | from motor.common.resources.instance import Instance, PDRole | 23 | from motor.common.resources.instance import Instance, PDRole |
G 严重程度： 提示 问题： 删除 原因： 本 PR 删除此处的 import 与方法（连同 scheduler.py 的对应调用、dispatch_capability.py），但 怎么改： 若测试仍有价值，保留并在 dispatch.py 代码注释标注仅测试使用；否则连测试一起删除。 tobking 19 天前 评论： 19 天前 评论： | |||
| 27 | from motor.common.resources.endpoint import Endpoint, Workload | 24 | from motor.common.resources.endpoint import Endpoint, Workload |
| 28 | from motor.coordinator.domain import ( | 25 | from motor.coordinator.domain import ( |
| @@ -740,10 +737,13 @@ class AsyncSchedulerClient: | |||
| 740 | req_info: RequestInfo, | 737 | req_info: RequestInfo, |
| 741 | role: PDRole | None = None, | 738 | role: PDRole | None = None, |
| 742 | top_k: int = 1, | 739 | top_k: int = 1, |
| 740 | + required_engine_type: str | None = None, | ||
| 743 | ) -> tuple[list[tuple[Instance, Endpoint, float]], str]: | 741 | ) -> tuple[list[tuple[Instance, Endpoint, float]], str]: |
| 744 | """Select endpoint candidates from cache or fresh instances.""" | 742 | """Select endpoint candidates from cache or fresh instances.""" |
| 745 | cache_role = role if role is not None else PDRole.ROLE_U | 743 | cache_role = role if role is not None else PDRole.ROLE_U |
| 746 | - cached_instances = self._cache.get_instances(cache_role) | 744 | + cached_instances = self._filter_instances_by_engine_type( |
| 745 | + self._cache.get_instances(cache_role), required_engine_type | ||
| 746 | + ) | ||
| 747 | if cached_instances: | 747 | if cached_instances: |
| 748 | # Cache stores instances sorted by id (see replace_all call sites); use as-is for RR | 748 | # Cache stores instances sorted by id (see replace_all call sites); use as-is for RR |
| 749 | candidates, candidate_policy = self._select_endpoint_candidates_from_list_with_policy( | 749 | candidates, candidate_policy = self._select_endpoint_candidates_from_list_with_policy( |
| @@ -762,7 +762,9 @@ class AsyncSchedulerClient: | |||
| 762 | return [], self._scheduler_type or CANDIDATE_POLICY_ROUND_ROBIN | 762 | return [], self._scheduler_type or CANDIDATE_POLICY_ROUND_ROBIN |
| 763 | 763 | ||
| 764 | # get_available_instances already wrote sorted list to cache; build sorted list once for this path | 764 | # get_available_instances already wrote sorted list to cache; build sorted list once for this path |
| 765 | - instance_list = sorted(instances.values(), key=lambda i: i.id) | 765 | + instance_list = self._filter_instances_by_engine_type( |
| 766 | + sorted(instances.values(), key=lambda i: i.id), required_engine_type | ||
| 767 | + ) | ||
| 766 | candidates, candidate_policy = self._select_endpoint_candidates_from_list_with_policy( | 768 | candidates, candidate_policy = self._select_endpoint_candidates_from_list_with_policy( |
| 767 | instance_list, cache_role, req_info, top_k=top_k | 769 | instance_list, cache_role, req_info, top_k=top_k |
| 768 | ) | 770 | ) |
| @@ -775,6 +777,17 @@ class AsyncSchedulerClient: | |||
| 775 | ) | 777 | ) |
| 776 | return candidates, candidate_policy | 778 | return candidates, candidate_policy |
| 777 | 779 | ||
| 780 | + | ||
| 781 | + def _filter_instances_by_engine_type(instances: list[Instance], required_engine_type: str | None) -> list[Instance]: | ||
| 782 | + normalized = str(required_engine_type or "").strip().lower() | ||
| 783 | + if not normalized: | ||
| 784 | + return instances | ||
| 785 | + return [ | ||
| 786 | + instance | ||
| 787 | + for instance in instances | ||
| 788 | + if str(getattr(instance, "engine_type", "")).strip().lower() == normalized | ||
| 789 | + ] | ||
| 790 | + | ||
| 778 | async def _refresh_cache_from_workload_reader(self, role: PDRole | None = None) -> None: | 791 | async def _refresh_cache_from_workload_reader(self, role: PDRole | None = None) -> None: |
| 779 | """Patch live workload into the local cache and pull a fresh instance list on | 792 | """Patch live workload into the local cache and pull a fresh instance list on |
| 780 | heartbeat-stale or instance-version change. | 793 | heartbeat-stale or instance-version change. |
| @@ -824,6 +837,7 @@ class AsyncSchedulerClient: | |||
| 824 | req_info: RequestInfo, | 837 | req_info: RequestInfo, |
| 825 | *, | 838 | *, |
| 826 | target_instance_id: int | None = None, | 839 | target_instance_id: int | None = None, |
| 840 | + required_engine_type: str | None = None, | ||
| 827 | ) -> tuple[Instance, Endpoint, Workload] | None: | 841 | ) -> tuple[Instance, Endpoint, Workload] | None: |
| 828 | """Select instance locally + ALLOCATE_ONLY RPC. Allocation workload is decided here (RR=zero, LB=demand).""" | 842 | """Select instance locally + ALLOCATE_ONLY RPC. Allocation workload is decided here (RR=zero, LB=demand).""" |
| 829 | role_str = role.value if role is not None else (getattr(PDRole.ROLE_U, "value", "union")) | 843 | role_str = role.value if role is not None else (getattr(PDRole.ROLE_U, "value", "union")) |
| @@ -834,8 +848,16 @@ class AsyncSchedulerClient: | |||
| 834 | # affinity-discounted prefill cost so the scheduler re-ranks globally by fresh load. | 848 | # affinity-discounted prefill cost so the scheduler re-ranks globally by fresh load. |
| 835 | global_affinity = False | 849 | global_affinity = False |
| 836 | 850 | ||
| 851 | + normalized_engine_type = str(required_engine_type or "").strip().lower() | ||
| 852 | + | ||
| 837 | if target_instance_id is not None: | 853 | if target_instance_id is not None: |
| 838 | instances = await self.get_available_instances(role) | 854 | instances = await self.get_available_instances(role) |
| 855 | + if normalized_engine_type: | ||
| 856 | + instances = { | ||
| 857 | + instance_id: candidate | ||
| 858 | + for instance_id, candidate in instances.items() | ||
| 859 | + if str(getattr(candidate, "engine_type", "")).strip().lower() == normalized_engine_type | ||
| 860 | + } | ||
| 839 | instance = resolve_pinned_instance(instances, target_instance_id) | 861 | instance = resolve_pinned_instance(instances, target_instance_id) |
| 840 | if instance is None: | 862 | if instance is None: |
| 841 | logger.warning( | 863 | logger.warning( |
| @@ -877,7 +899,12 @@ class AsyncSchedulerClient: | |||
| 877 | ( | 899 | ( |
| 878 | candidates, | 900 | candidates, |
| 879 | candidate_policy, | 901 | candidate_policy, |
| 880 | - ) = await self._select_endpoint_candidates_with_policy(req_info, role, top_k=request_top_k) | 902 | + ) = await self._select_endpoint_candidates_with_policy( |
| 903 | + req_info, | ||
| 904 | + role, | ||
| 905 | + top_k=request_top_k, | ||
| 906 | + required_engine_type=normalized_engine_type or None, | ||
| 907 | + ) | ||
| 881 | if not candidates: | 908 | if not candidates: |
| 882 | return None | 909 | return None |
| 883 | instance, endpoint, _ = candidates[0] | 910 | instance, endpoint, _ = candidates[0] |
| @@ -893,6 +920,13 @@ class AsyncSchedulerClient: | |||
| 893 | and any(rec[2] is not None for rec in affinity_debug.values()) | 920 | and any(rec[2] is not None for rec in affinity_debug.values()) |
| 894 | ) | 921 | ) |
| 895 | if global_affinity: | 922 | if global_affinity: |
| 923 | + allowed_instance_ids = { | ||
| 924 | + candidate.id | ||
| 925 | + for candidate in self._filter_instances_by_engine_type( | ||
| 926 | + self._cache.get_instances(role), | ||
| 927 | + normalized_engine_type or None, | ||
| 928 | + ) | ||
| 929 | + } | ||
| 896 | candidate_endpoints = [ | 930 | candidate_endpoints = [ |
| 897 | { | 931 | { |
| 898 | "instance_id": ins_id, | 932 | "instance_id": ins_id, |
| @@ -901,7 +935,7 @@ class AsyncSchedulerClient: | |||
| 901 | "prefill_cost": rec[2], | 935 | "prefill_cost": rec[2], |
| 902 | } | 936 | } |
| 903 | for (ins_id, ep_id), rec in affinity_debug.items() | 937 | for (ins_id, ep_id), rec in affinity_debug.items() |
| 904 | - if rec[2] is not None | 938 | + if rec[2] is not None and ins_id in allowed_instance_ids |
| 905 | ] | 939 | ] |
| 906 | elif candidate_policy == CANDIDATE_POLICY_KV_CACHE_AFFINITY and isinstance(affinity_debug, dict): | 940 | elif candidate_policy == CANDIDATE_POLICY_KV_CACHE_AFFINITY and isinstance(affinity_debug, dict): |
| 907 | # load_gated (and other affinity modes without prefill_cost): still forward | 941 | # load_gated (and other affinity modes without prefill_cost): still forward |
| @@ -949,6 +983,7 @@ class AsyncSchedulerClient: | |||
| 949 | # Non-affinity demand as a raw float (RR sends 0.0). Affinity overrides via isl/matched. | 983 | # Non-affinity demand as a raw float (RR sends 0.0). Affinity overrides via isl/matched. |
| 950 | "workload_active_tokens": workload.active_tokens, | 984 | "workload_active_tokens": workload.active_tokens, |
| 951 | "candidate_policy": candidate_policy, | 985 | "candidate_policy": candidate_policy, |
| 986 | + "required_engine_type": normalized_engine_type or None, | ||
| 952 | } | 987 | } |
| 953 | if isl > 0: | 988 | if isl > 0: |
| 954 | req_data["isl"] = isl | 989 | req_data["isl"] = isl |
| @@ -1337,23 +1372,6 @@ class AsyncSchedulerClient: | |||
| 1337 | roles = self._roles_from_cache() | 1372 | roles = self._roles_from_cache() |
| 1338 | return roles | 1373 | return roles |
| 1339 | 1374 | ||
| 1340 | - async def has_compatible_pd_pair(self) -> bool: | ||
| 1341 | - """Return whether cached P/D pools contain a compatible pair. | ||
| 1342 | - | ||
| 1343 | - Assumes the cache was warmed by a preceding get_available_instance_roles in the same | ||
| 1344 | - routing decision; falls back to a warm-up fetch if both pools look empty. | ||
| 1345 | - """ | ||
| 1346 | - prefill = self._cache.get_instances(PDRole.ROLE_P) | ||
| 1347 | - decode = self._cache.get_instances(PDRole.ROLE_D) | ||
| 1348 | - if not prefill and not decode: | ||
| 1349 | - try: | ||
| 1350 | - await self.get_available_instances(None) | ||
| 1351 | - except Exception as e: | ||
| 1352 | - logger.debug("has_compatible_pd_pair: warm-up fetch failed: %s", e) | ||
| 1353 | - prefill = self._cache.get_instances(PDRole.ROLE_P) | ||
| 1354 | - decode = self._cache.get_instances(PDRole.ROLE_D) | ||
| 1355 | - return has_compatible_dispatch_pair(prefill, decode) | ||
| 1356 | - | ||
| 1357 | async def get_unblocked_instances(self, role: PDRole) -> list[int]: | 1375 | async def get_unblocked_instances(self, role: PDRole) -> list[int]: |
| 1358 | """Return instance IDs of the given role that are NOT blocked by circuit breaker.""" | 1376 | """Return instance IDs of the given role that are NOT blocked by circuit breaker.""" |
| 1359 | cached = self._cache.get_instances(role) | 1377 | cached = self._cache.get_instances(role) |
| @@ -133,6 +133,7 @@ _KEY_MATCHED_TOKENS = "matched_tokens" | |||
| 133 | _KEY_PREFILL_COST = "prefill_cost" | 133 | _KEY_PREFILL_COST = "prefill_cost" |
| 134 | _KEY_LOAD_WEIGHT = "load_weight" | 134 | _KEY_LOAD_WEIGHT = "load_weight" |
| 135 | _KEY_PREFILL_LOAD_SCALE = "prefill_load_scale" | 135 | _KEY_PREFILL_LOAD_SCALE = "prefill_load_scale" |
| 136 | +_KEY_REQUIRED_ENGINE_TYPE = "required_engine_type" | ||
| 136 | 137 | ||
| 137 | 138 | ||
| 138 | def _should_log_scheduling_sample(sample_key: str) -> bool: | 139 | def _should_log_scheduling_sample(sample_key: str) -> bool: |
| @@ -169,7 +170,6 @@ def _serialize_instance_minimal(instance: Instance | None) -> dict: | |||
| 169 | "job_name": instance.job_name, | 170 | "job_name": instance.job_name, |
| 170 | "model_name": instance.model_name, | 171 | "model_name": instance.model_name, |
| 171 | "engine_type": instance.engine_type, | 172 | "engine_type": instance.engine_type, |
| 172 | - "dispatch_capabilities": list(instance.dispatch_capabilities or []), | ||
| 173 | } | 173 | } |
| 174 | 174 | ||
| 175 | 175 | ||
| @@ -182,6 +182,7 @@ def _serialize_endpoint_minimal(endpoint: Endpoint | None) -> dict: | |||
| 182 | "ip": endpoint.ip, | 182 | "ip": endpoint.ip, |
| 183 | "business_port": endpoint.business_port, | 183 | "business_port": endpoint.business_port, |
| 184 | "mgmt_port": getattr(endpoint, "mgmt_port", "") or "", | 184 | "mgmt_port": getattr(endpoint, "mgmt_port", "") or "", |
| 185 | + "bootstrap_port": endpoint.bootstrap_port, | ||
| 185 | } | 186 | } |
| 186 | if hasattr(endpoint, "status") and endpoint.status is not None: | 187 | if hasattr(endpoint, "status") and endpoint.status is not None: |
| 187 | out["status"] = endpoint.status.value if hasattr(endpoint.status, "value") else str(endpoint.status) | 188 | out["status"] = endpoint.status.value if hasattr(endpoint.status, "value") else str(endpoint.status) |
| @@ -660,6 +661,7 @@ class _SchedulerRequestDispatcher: | |||
| 660 | worker_load_weight = self._parse_optional_float(request.data.get(_KEY_LOAD_WEIGHT)) | 661 | worker_load_weight = self._parse_optional_float(request.data.get(_KEY_LOAD_WEIGHT)) |
| 661 | worker_prefill_load_scale = self._parse_optional_float(request.data.get(_KEY_PREFILL_LOAD_SCALE)) | 662 | worker_prefill_load_scale = self._parse_optional_float(request.data.get(_KEY_PREFILL_LOAD_SCALE)) |
| 662 | isl = self._parse_optional_float(request.data.get(_KEY_ISL)) | 663 | isl = self._parse_optional_float(request.data.get(_KEY_ISL)) |
| 664 | + required_engine_type = str(request.data.get(_KEY_REQUIRED_ENGINE_TYPE) or "").strip().lower() or None | ||
| 663 | 665 | ||
| 664 | if instance_id is None or endpoint_id is None: | 666 | if instance_id is None or endpoint_id is None: |
| 665 | return SchedulerResponse( | 667 | return SchedulerResponse( |
| @@ -714,7 +716,7 @@ class _SchedulerRequestDispatcher: | |||
| 714 | role, | 716 | role, |
| 715 | ) | 717 | ) |
| 716 | selected = ( | 718 | selected = ( |
| 717 | - self._select_valid_candidate(selected_candidate, role) | 719 | + self._select_valid_candidate(selected_candidate, role, required_engine_type) |
| 718 | if fast_path | 720 | if fast_path |
| 719 | else self._select_authoritative_allocate_candidate( | 721 | else self._select_authoritative_allocate_candidate( |
| 720 | selected_candidate, | 722 | selected_candidate, |
| @@ -724,6 +726,7 @@ class _SchedulerRequestDispatcher: | |||
| 724 | affinity_candidates, | 726 | affinity_candidates, |
| 725 | worker_prefill_load_scale, | 727 | worker_prefill_load_scale, |
| 726 | worker_load_weight, | 728 | worker_load_weight, |
| 729 | + required_engine_type, | ||
| 727 | ) | 730 | ) |
| 728 | ) | 731 | ) |
| 729 | if fast_path and selected is None: | 732 | if fast_path and selected is None: |
| @@ -735,6 +738,7 @@ class _SchedulerRequestDispatcher: | |||
| 735 | affinity_candidates, | 738 | affinity_candidates, |
| 736 | worker_prefill_load_scale, | 739 | worker_prefill_load_scale, |
| 737 | worker_load_weight, | 740 | worker_load_weight, |
| 741 | + required_engine_type, | ||
| 738 | ) | 742 | ) |
| 739 | fast_path = False | 743 | fast_path = False |
| 740 | if selected is None: | 744 | if selected is None: |
| @@ -909,9 +913,10 @@ class _SchedulerRequestDispatcher: | |||
| 909 | self, | 913 | self, |
| 910 | candidate: tuple[int, int], | 914 | candidate: tuple[int, int], |
| 911 | role: PDRole, | 915 | role: PDRole, |
| 916 | + required_engine_type: str | None = None, | ||
| 912 | ) -> tuple[Instance, Endpoint, float] | None: | 917 | ) -> tuple[Instance, Endpoint, float] | None: |
| 913 | """Select the best candidate using SchedulerServer's current workload ledger.""" | 918 | """Select the best candidate using SchedulerServer's current workload ledger.""" |
| 914 | - return self._select_valid_candidate(candidate, role) | 919 | + return self._select_valid_candidate(candidate, role, required_engine_type) |
| 915 | 920 | ||
| 916 | def _select_authoritative_allocate_candidate( | 921 | def _select_authoritative_allocate_candidate( |
| 917 | self, | 922 | self, |
| @@ -922,6 +927,7 @@ class _SchedulerRequestDispatcher: | |||
| 922 | affinity_candidates: list[tuple[int, int, float]] | None = None, | 927 | affinity_candidates: list[tuple[int, int, float]] | None = None, |
| 923 | prefill_load_scale: float | None = None, | 928 | prefill_load_scale: float | None = None, |
| 924 | load_weight: float | None = None, | 929 | load_weight: float | None = None, |
| 930 | + required_engine_type: str | None = None, | ||
| 925 | ) -> tuple[Instance, Endpoint, float] | None: | 931 | ) -> tuple[Instance, Endpoint, float] | None: |
| 926 | """ | 932 | """ |
| 927 | Select allocation target using SchedulerServer's authoritative workload view. | 933 | Select allocation target using SchedulerServer's authoritative workload view. |
| @@ -934,19 +940,21 @@ class _SchedulerRequestDispatcher: | |||
| 934 | worker's ranked alternates". Other policies keep the worker-proposed endpoint. | 940 | worker's ranked alternates". Other policies keep the worker-proposed endpoint. |
| 935 | """ | 941 | """ |
| 936 | if self._should_scan_global_load_balance(candidate_policy): | 942 | if self._should_scan_global_load_balance(candidate_policy): |
| 937 | - selected = self._select_global_load_balance_candidate(role) | 943 | + selected = self._select_global_load_balance_candidate(role, required_engine_type) |
| 938 | if selected is not None: | 944 | if selected is not None: |
| 939 | return selected | 945 | return selected |
| 940 | if candidate_policy == CANDIDATE_POLICY_KV_CACHE_AFFINITY: | 946 | if candidate_policy == CANDIDATE_POLICY_KV_CACHE_AFFINITY: |
| 941 | if affinity_candidates: | 947 | if affinity_candidates: |
| 942 | - selected = self._select_affinity_global(affinity_candidates, role, prefill_load_scale, load_weight) | 948 | + selected = self._select_affinity_global( |
| 949 | + affinity_candidates, role, prefill_load_scale, load_weight, required_engine_type | ||
| 950 | + ) | ||
| 943 | if selected is not None: | 951 | if selected is not None: |
| 944 | return selected | 952 | return selected |
| 945 | elif len(candidates) > 1: | 953 | elif len(candidates) > 1: |
| 946 | - selected = self._select_lowest_load_among_candidates(candidates, role) | 954 | + selected = self._select_lowest_load_among_candidates(candidates, role, required_engine_type) |
| 947 | if selected is not None: | 955 | if selected is not None: |
| 948 | return selected | 956 | return selected |
| 949 | - return self._select_authoritative_candidate(candidate, role) | 957 | + return self._select_authoritative_candidate(candidate, role, required_engine_type) |
| 950 | 958 | ||
| 951 | def _select_affinity_global( | 959 | def _select_affinity_global( |
| 952 | self, | 960 | self, |
| @@ -954,6 +962,7 @@ class _SchedulerRequestDispatcher: | |||
| 954 | role: PDRole, | 962 | role: PDRole, |
| 955 | prefill_load_scale: float | None, | 963 | prefill_load_scale: float | None, |
| 956 | load_weight: float | None, | 964 | load_weight: float | None, |
| 965 | + required_engine_type: str | None = None, | ||
| 957 | ) -> tuple[Instance, Endpoint, float] | None: | 966 | ) -> tuple[Instance, Endpoint, float] | None: |
| 958 | """ | 967 | """ |
| 959 | Global kv_cache_affinity unified selection over EVERY worker-reported endpoint. | 968 | Global kv_cache_affinity unified selection over EVERY worker-reported endpoint. |
| @@ -974,6 +983,8 @@ class _SchedulerRequestDispatcher: | |||
| 974 | if found is None: | 983 | if found is None: |
| 975 | continue | 984 | continue |
| 976 | instance, endpoint = found | 985 | instance, endpoint = found |
| 986 | + if not self._matches_engine_type(instance, required_engine_type): | ||
| 987 | + continue | ||
| 977 | try: | 988 | try: |
| 978 | instance_role = PDRole(instance.role) | 989 | instance_role = PDRole(instance.role) |
| 979 | except ValueError: | 990 | except ValueError: |
| @@ -1008,6 +1019,7 @@ class _SchedulerRequestDispatcher: | |||
| 1008 | self, | 1019 | self, |
| 1009 | candidates: list[tuple[int, int]], | 1020 | candidates: list[tuple[int, int]], |
| 1010 | role: PDRole, | 1021 | role: PDRole, |
| 1022 | + required_engine_type: str | None = None, | ||
| 1011 | ) -> tuple[Instance, Endpoint, float] | None: | 1023 | ) -> tuple[Instance, Endpoint, float] | None: |
| 1012 | """ | 1024 | """ |
| 1013 | Among the worker's affinity-ranked candidates, pick the lowest current endpoint score from | 1025 | Among the worker's affinity-ranked candidates, pick the lowest current endpoint score from |
| @@ -1022,6 +1034,8 @@ class _SchedulerRequestDispatcher: | |||
| 1022 | if found is None: | 1034 | if found is None: |
| 1023 | continue | 1035 | continue |
| 1024 | instance, endpoint = found | 1036 | instance, endpoint = found |
| 1037 | + if not self._matches_engine_type(instance, required_engine_type): | ||
| 1038 | + continue | ||
| 1025 | try: | 1039 | try: |
| 1026 | instance_role = PDRole(instance.role) | 1040 | instance_role = PDRole(instance.role) |
| 1027 | except ValueError: | 1041 | except ValueError: |
| @@ -1207,13 +1221,18 @@ class _SchedulerRequestDispatcher: | |||
| 1207 | def _select_global_load_balance_candidate( | 1221 | def _select_global_load_balance_candidate( |
| 1208 | self, | 1222 | self, |
| 1209 | role: PDRole, | 1223 | role: PDRole, |
| 1224 | + required_engine_type: str | None = None, | ||
| 1210 | ) -> tuple[Instance, Endpoint, float] | None: | 1225 | ) -> tuple[Instance, Endpoint, float] | None: |
| 1211 | """Select the globally lowest-score endpoint for role from SchedulerServer's local pool. | 1226 | """Select the globally lowest-score endpoint for role from SchedulerServer's local pool. |
| 1212 | 1227 | ||
| 1213 | Circuit-broken endpoints are filtered so the authoritative re-scan never picks one | 1228 | Circuit-broken endpoints are filtered so the authoritative re-scan never picks one |
| 1214 | that the local PUB cache may not yet know about. | 1229 | that the local PUB cache may not yet know about. |
| 1215 | """ | 1230 | """ |
| 1216 | - instances = self._instance_manager.get_available_instances(role).values() | 1231 | + instances = [ |
| 1232 | + instance | ||
| 1233 | + for instance in self._instance_manager.get_available_instances(role).values() | ||
| 1234 | + if self._matches_engine_type(instance, required_engine_type) | ||
| 1235 | + ] | ||
| 1217 | candidates = LoadBalancePolicy.select_endpoint_candidates_from_list( | 1236 | candidates = LoadBalancePolicy.select_endpoint_candidates_from_list( |
| 1218 | instances, | 1237 | instances, |
| 1219 | role=role, | 1238 | role=role, |
| @@ -1258,6 +1277,7 @@ class _SchedulerRequestDispatcher: | |||
| 1258 | self, | 1277 | self, |
| 1259 | candidate: tuple[int, int], | 1278 | candidate: tuple[int, int], |
| 1260 | role: PDRole, | 1279 | role: PDRole, |
| 1280 | + required_engine_type: str | None = None, | ||
| 1261 | ) -> tuple[Instance, Endpoint, float] | None: | 1281 | ) -> tuple[Instance, Endpoint, float] | None: |
| 1262 | """ | 1282 | """ |
| 1263 | Validate one worker-selected candidate and calculate its current score for observability. | 1283 | Validate one worker-selected candidate and calculate its current score for observability. |
| @@ -1272,6 +1292,8 @@ class _SchedulerRequestDispatcher: | |||
| 1272 | if found is None: | 1292 | if found is None: |
| 1273 | return None | 1293 | return None |
| 1274 | instance, endpoint = found | 1294 | instance, endpoint = found |
| 1295 | + if not self._matches_engine_type(instance, required_engine_type): | ||
| 1296 | + return None | ||
| 1275 | try: | 1297 | try: |
| 1276 | instance_role = PDRole(instance.role) | 1298 | instance_role = PDRole(instance.role) |
| 1277 | except ValueError: | 1299 | except ValueError: |
| @@ -1295,6 +1317,12 @@ class _SchedulerRequestDispatcher: | |||
| 1295 | return None | 1317 | return None |
| 1296 | return (instance, endpoint, score) | 1318 | return (instance, endpoint, score) |
| 1297 | 1319 | ||
| 1320 | + | ||
| 1321 | + def _matches_engine_type(instance: Instance, required_engine_type: str | None) -> bool: | ||
| 1322 | + if not required_engine_type: | ||
| 1323 | + return True | ||
| 1324 | + return str(getattr(instance, "engine_type", "")).strip().lower() == required_engine_type | ||
| 1325 | + | ||
| 1298 | def _find_available_instance_endpoint( | 1326 | def _find_available_instance_endpoint( |
| 1299 | self, | 1327 | self, |
| 1300 | instance_id: int, | 1328 | instance_id: int, |
| @@ -11,9 +11,6 @@ | |||
| 11 | import asyncio | 11 | import asyncio |
| 12 | import uuid | 12 | import uuid |
| 13 | 13 | ||
| 14 | -from motor.common.resources.dispatch import ( | ||
| 15 | - has_compatible_dispatch_pair, | ||
| 16 | -) | ||
| 17 | from motor.common.resources.instance import Instance, PDRole | 14 | from motor.common.resources.instance import Instance, PDRole |
| 18 | from motor.common.resources.endpoint import WorkloadAction, Workload | 15 | from motor.common.resources.endpoint import WorkloadAction, Workload |
| 19 | from motor.coordinator.domain import ( | 16 | from motor.coordinator.domain import ( |
| @@ -116,6 +113,7 @@ class Scheduler: | |||
| 116 | req_info: RequestInfo, | 113 | req_info: RequestInfo, |
| 117 | *, | 114 | *, |
| 118 | target_instance_id: int | None = None, | 115 | target_instance_id: int | None = None, |
| 116 | + required_engine_type: str | None = None, | ||
| 119 | ): | 117 | ): |
| 120 | """ | 118 | """ |
| 121 | Atomic: select instance + one workload allocation (ALLOCATION). | 119 | Atomic: select instance + one workload allocation (ALLOCATION). |
| @@ -125,8 +123,15 @@ class Scheduler: | |||
| 125 | (Instance, Endpoint, Workload) tuple or None (no instance or update_workload failed). | 123 | (Instance, Endpoint, Workload) tuple or None (no instance or update_workload failed). |
| 126 | The returned Workload is what was allocated; caller records it for release. | 124 | The returned Workload is what was allocated; caller records it for release. |
| 127 | """ | 125 | """ |
| 126 | + pool = self._instance_provider.get_available_instances(role) | ||
| 127 | + if required_engine_type is not None: | ||
| 128 | + normalized_engine_type = required_engine_type.strip().lower() | ||
| 129 | + pool = { | ||
| 130 | + instance_id: instance | ||
| 131 | + for instance_id, instance in pool.items() | ||
| 132 | + if str(getattr(instance, "engine_type", "")).strip().lower() == normalized_engine_type | ||
| 133 | + } | ||
| 128 | if target_instance_id is not None: | 134 | if target_instance_id is not None: |
| 129 | - pool = self._instance_provider.get_available_instances(role) | ||
| 130 | instance = resolve_pinned_instance(pool, target_instance_id) | 135 | instance = resolve_pinned_instance(pool, target_instance_id) |
| 131 | if instance is None: | 136 | if instance is None: |
| 132 | logger.warning( | 137 | logger.warning( |
| @@ -147,7 +152,11 @@ class Scheduler: | |||
| 147 | ) | 152 | ) |
| 148 | return None | 153 | return None |
| 149 | else: | 154 | else: |
| 150 | - r = self._scheduling_policy.select_instance_and_endpoint(role) | 155 | + r = self._scheduling_policy.select_instance_and_endpoint_from_list( |
| 156 | + list(pool.values()), | ||
| 157 | + role, | ||
| 158 | + req_info, | ||
| 159 | + ) | ||
| 151 | result = (await r) if asyncio.iscoroutine(r) else r | 160 | result = (await r) if asyncio.iscoroutine(r) else r |
| 152 | if result is None: | 161 | if result is None: |
| 153 | return None | 162 | return None |
| @@ -240,12 +249,6 @@ class Scheduler: | |||
| 240 | roles.add(aliases[normalized]) | 249 | roles.add(aliases[normalized]) |
| 241 | return roles | 250 | return roles |
| 242 | 251 | ||
| 243 | - async def has_compatible_pd_pair(self) -> bool: | ||
| 244 | - """Return whether the in-process instance view has a compatible P/D pair.""" | ||
| 245 | - prefill = self._instance_provider.get_available_instances(PDRole.ROLE_P).values() | ||
| 246 | - decode = self._instance_provider.get_available_instances(PDRole.ROLE_D).values() | ||
| 247 | - return has_compatible_dispatch_pair(prefill, decode) | ||
| 248 | - | ||
| 249 | async def get_unblocked_instances(self, role: PDRole) -> list[int]: | 252 | async def get_unblocked_instances(self, role: PDRole) -> list[int]: |
| 250 | """Return all instance IDs for the role (in-process scheduler has no CB).""" | 253 | """Return all instance IDs for the role (in-process scheduler has no CB).""" |
| 251 | return [inst.id for inst in self._instance_provider.get_available_instances(role).values()] | 254 | return [inst.id for inst in self._instance_provider.get_available_instances(role).values()] |
| @@ -44,7 +44,7 @@ set_process_title(f"EngineServer-DP{_dp_rank_from_argv()}") | |||
| 44 | # ruff: noqa: E402 | 44 | # ruff: noqa: E402 |
| 45 | from motor.common.logger import get_logger | 45 | from motor.common.logger import get_logger |
| 46 | from motor.config.endpoint import EndpointConfig | 46 | from motor.config.endpoint import EndpointConfig |
| 47 | -from motor.engine_server.factory.config_factory import ConfigFactory | 47 | +from motor.node_manager.core.services.native_engine.config_factory import ConfigFactory |
| 48 | from motor.engine_server.factory.endpoint_factory import EndpointFactory | 48 | from motor.engine_server.factory.endpoint_factory import EndpointFactory |
| 49 | from motor.engine_server.utils.prometheus import setup_multiprocess_prometheus | 49 | from motor.engine_server.utils.prometheus import setup_multiprocess_prometheus |
| 50 | 50 | ||
| @@ -31,7 +31,7 @@ from motor.common.resources.dispatch import ( | |||
| 31 | MotorDispatch, | 31 | MotorDispatch, |
| 32 | PrefillResult, | 32 | PrefillResult, |
| 33 | ) | 33 | ) |
| 34 | -from motor.engine_server.core.config import IConfig | 34 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 35 | from motor.engine_server.core.vllm.prefill_context_validation import PrefillContextCheck | 35 | from motor.engine_server.core.vllm.prefill_context_validation import PrefillContextCheck |
| 36 | from motor.engine_server.core.errors.sanitizer import sanitize_error_message | 36 | from motor.engine_server.core.errors.sanitizer import sanitize_error_message |
| 37 | 37 | ||
| @@ -8,7 +8,7 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -from motor.engine_server.core.config import IConfig | 11 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 12 | from motor.engine_server.core.dispatch_adapter.base import DispatchAdapter | 12 | from motor.engine_server.core.dispatch_adapter.base import DispatchAdapter |
| 13 | from motor.engine_server.core.dispatch_adapter.sglang_adapter import SGLangDispatchAdapter | 13 | from motor.engine_server.core.dispatch_adapter.sglang_adapter import SGLangDispatchAdapter |
| 14 | from motor.engine_server.core.dispatch_adapter.vllm_adapter import VLLMDispatchAdapter | 14 | from motor.engine_server.core.dispatch_adapter.vllm_adapter import VLLMDispatchAdapter |
| @@ -9,7 +9,6 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | import hashlib | 11 | import hashlib |
| 12 | -import os | ||
| 13 | from typing import Any | 12 | from typing import Any |
| 14 | from urllib.parse import urlparse | 13 | from urllib.parse import urlparse |
| 15 | 14 | ||
| @@ -27,11 +26,11 @@ class SGLangDispatchAdapter(DispatchAdapter): | |||
| 27 | return body | 26 | return body |
| 28 | 27 | ||
| 29 | parsed = urlparse(prefill.url) | 28 | parsed = urlparse(prefill.url) |
| 30 | - bootstrap_port = os.getenv("DISAGGREGATION_BOOTSTRAP_PORT", "").strip() | 29 | + bootstrap_port = prefill.bootstrap_port |
| 31 | - if not bootstrap_port: | 30 | + if bootstrap_port is None: |
| 32 | raise HTTPException( | 31 | raise HTTPException( |
| 33 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, | 32 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 34 | - detail=("DISAGGREGATION_BOOTSTRAP_PORT must be set for SGLang dispatch requests."), | 33 | + detail=("Prefill endpoint bootstrap_port is required for SGLang dispatch requests."), |
| 35 | ) | 34 | ) |
| 36 | body.update( | 35 | body.update( |
| 37 | { | 36 | { |
| @@ -14,7 +14,7 @@ from motor.common.http.http_client import AsyncSafeHTTPSClient | |||
| 14 | from motor.common.logger import get_logger | 14 | from motor.common.logger import get_logger |
| 15 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, get_pod_ip | 15 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, get_pod_ip |
| 16 | from motor.config.endpoint import EndpointConfig | 16 | from motor.config.endpoint import EndpointConfig |
| 17 | -from motor.engine_server.utils.ip import build_endpoint | 17 | +from motor.common.utils.ip import build_endpoint |
| 18 | 18 | ||
| 19 | logger = get_logger(__name__) | 19 | logger = get_logger(__name__) |
| 20 | 20 | ||
| @@ -24,7 +24,7 @@ from motor.common.http.cert_util import CertUtil | |||
| 24 | from motor.common.logger import get_logger | 24 | from motor.common.logger import get_logger |
| 25 | from motor.common.resources.dispatch import DispatchStopState, MotorDispatch | 25 | from motor.common.resources.dispatch import DispatchStopState, MotorDispatch |
| 26 | from motor.engine_server.core.dispatch_adapter import create_dispatch_adapter | 26 | from motor.engine_server.core.dispatch_adapter import create_dispatch_adapter |
| 27 | -from motor.engine_server.core.config import IConfig | 27 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 28 | from motor.engine_server.core.dispatch_adapter.base import DispatchResponseContext | 28 | from motor.engine_server.core.dispatch_adapter.base import DispatchResponseContext |
| 29 | from motor.common.utils.net import format_address | 29 | from motor.common.utils.net import format_address |
| 30 | from motor.engine_server.core.endpoint import Endpoint | 30 | from motor.engine_server.core.endpoint import Endpoint |
| @@ -24,13 +24,13 @@ from motor.common.http.cert_util import CertUtil | |||
| 24 | from motor.common.logger import get_logger | 24 | from motor.common.logger import get_logger |
| 25 | from motor.common.utils.net import format_address | 25 | from motor.common.utils.net import format_address |
| 26 | from motor.config.endpoint import EndpointConfig | 26 | from motor.config.endpoint import EndpointConfig |
| 27 | -from motor.engine_server.core.config import IConfig | 27 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 28 | from motor.engine_server.core.endpoint import Endpoint | 28 | from motor.engine_server.core.endpoint import Endpoint |
| 29 | from motor.engine_server.core.health_collector import HealthCollector | 29 | from motor.engine_server.core.health_collector import HealthCollector |
| 30 | from motor.engine_server.core.snapshot_monitor import SnapshotMonitor | 30 | from motor.engine_server.core.snapshot_monitor import SnapshotMonitor |
| 31 | from motor.engine_server.core.sim_inference import SimInference | 31 | from motor.engine_server.core.sim_inference import SimInference |
| 32 | -from motor.engine_server.constants import constants | 32 | +from motor.common import engine_constants as constants |
| 33 | -from motor.engine_server.constants.constants import ( | 33 | +from motor.common.engine_constants import ( |
| 34 | METRICS_INTERFACE, | 34 | METRICS_INTERFACE, |
| 35 | STATUS_INTERFACE, | 35 | STATUS_INTERFACE, |
| 36 | STATUS_KEY, | 36 | STATUS_KEY, |
| @@ -22,7 +22,7 @@ | |||
| 22 | from typing import Any | 22 | from typing import Any |
| 23 | 23 | ||
| 24 | from motor.common.logger import get_logger | 24 | from motor.common.logger import get_logger |
| 25 | -from motor.engine_server.core.config import IConfig | 25 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 26 | from motor.engine_server.core.engine import Engine | 26 | from motor.engine_server.core.engine import Engine |
| 27 | 27 | ||
| 28 | logger = get_logger(__name__) | 28 | logger = get_logger(__name__) |
| @@ -33,13 +33,13 @@ def _kill_engine_children() -> None: | |||
| 33 | try: | 33 | try: |
| 34 | from sglang.srt.utils import kill_process_tree | 34 | from sglang.srt.utils import kill_process_tree |
| 35 | import os | 35 | import os |
| 36 | + | ||
| 36 | kill_process_tree(os.getpid(), include_parent=False) | 37 | kill_process_tree(os.getpid(), include_parent=False) |
| 37 | except Exception as e: | 38 | except Exception as e: |
| 38 | logger.exception("Error killing SGLang engine child processes: %s", e) | 39 | logger.exception("Error killing SGLang engine child processes: %s", e) |
| 39 | 40 | ||
| 40 | 41 | ||
| 41 | class SGLangEngine(Engine): | 42 | class SGLangEngine(Engine): |
| 42 | - | ||
| 43 | def __init__(self, config: IConfig): | 43 | def __init__(self, config: IConfig): |
| 44 | self._template_manager: Any | None = None | 44 | self._template_manager: Any | None = None |
| 45 | self._tokenizer_manager: Any | None = None | 45 | self._tokenizer_manager: Any | None = None |
| @@ -47,7 +47,7 @@ class SGLangEngine(Engine): | |||
| 47 | self.config = config | 47 | self.config = config |
| 48 | 48 | ||
| 49 | def launch(self) -> Any: | 49 | def launch(self) -> Any: |
| 50 | - from sglang.srt.entrypoints.engine import (_launch_subprocesses) | 50 | + from sglang.srt.entrypoints.engine import _launch_subprocesses |
| 51 | 51 | ||
| 52 | server_args = self.config.get_args() | 52 | server_args = self.config.get_args() |
| 53 | if server_args is None: | 53 | if server_args is None: |
| @@ -74,16 +74,14 @@ class SGLangEngine(Engine): | |||
| 74 | 74 | ||
| 75 | if server_args.tokenizer_worker_num > 1: | 75 | if server_args.tokenizer_worker_num > 1: |
| 76 | from sglang.srt.managers.multi_tokenizer_mixin import write_data_for_multi_tokenizer | 76 | from sglang.srt.managers.multi_tokenizer_mixin import write_data_for_multi_tokenizer |
| 77 | - self._multi_tokenizer_args_shm = write_data_for_multi_tokenizer( | 77 | + |
| 78 | - port_args, server_args, scheduler_infos[0] | 78 | + self._multi_tokenizer_args_shm = write_data_for_multi_tokenizer(port_args, server_args, scheduler_infos[0]) |
| 79 | - ) | ||
| 80 | 79 | ||
| 81 | return self._tokenizer_manager, self._template_manager | 80 | return self._tokenizer_manager, self._template_manager |
| 82 | 81 | ||
| 83 | def shutdown(self) -> None: | 82 | def shutdown(self) -> None: |
| 84 | if self._multi_tokenizer_args_shm is not None: | 83 | if self._multi_tokenizer_args_shm is not None: |
| 85 | self._multi_tokenizer_args_shm.close() | 84 | self._multi_tokenizer_args_shm.close() |
| 86 | - self._multi_tokenizer_args_shm = None | 85 | + self._multi_tokenizer_args_shm = None |
| 87 | - self._tokenizer_manager.socket_mapping.clear_all_sockets() | 86 | + self._tokenizer_manager.socket_mapping.clear_all_sockets() |
| 88 | _kill_engine_children() | 87 | _kill_engine_children() |
| 89 | - | ||
| @@ -18,10 +18,10 @@ from motor.common.resources.dispatch import DispatchProfile, infer_vllm_dispatch | |||
| 18 | from motor.common.http.http_client import AsyncSafeHTTPSClient | 18 | from motor.common.http.http_client import AsyncSafeHTTPSClient |
| 19 | from motor.common.logger import get_logger | 19 | from motor.common.logger import get_logger |
| 20 | from motor.common.utils.net import format_address | 20 | from motor.common.utils.net import format_address |
| 21 | -from motor.engine_server.core.config import IConfig | 21 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 22 | from motor.engine_server.utils.ai_cube import get_ai_cube_usage, is_ai_cube_usage_watch_supported | 22 | from motor.engine_server.utils.ai_cube import get_ai_cube_usage, is_ai_cube_usage_watch_supported |
| 23 | -from motor.engine_server.constants import constants | 23 | +from motor.common import engine_constants as constants |
| 24 | -from motor.engine_server.utils.ip import build_endpoint | 24 | +from motor.common.utils.ip import build_endpoint |
| 25 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, get_pod_ip | 25 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, get_pod_ip |
| 26 | 26 | ||
| 27 | logger = get_logger(__name__) | 27 | logger = get_logger(__name__) |
| @@ -21,7 +21,7 @@ from motor.common.utils.snapshot_utils import ( | |||
| 21 | get_pod_ip, | 21 | get_pod_ip, |
| 22 | RETRY_LOG_FREQUENCY, | 22 | RETRY_LOG_FREQUENCY, |
| 23 | ) | 23 | ) |
| 24 | -from motor.engine_server.utils.ip import build_endpoint | 24 | +from motor.common.utils.ip import build_endpoint |
| 25 | 25 | ||
| 26 | logger = get_logger(__name__) | 26 | logger = get_logger(__name__) |
| 27 | 27 | ||
| @@ -18,7 +18,7 @@ from vllm.v1.executor.multiproc_executor import MultiprocExecutor | |||
| 18 | 18 | ||
| 19 | from motor.common.logger import get_logger | 19 | from motor.common.logger import get_logger |
| 20 | from motor.common.logger import attach_to_vllm_logger | 20 | from motor.common.logger import attach_to_vllm_logger |
| 21 | -from motor.engine_server.core.config import IConfig | 21 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 22 | from motor.engine_server.core.engine import Engine | 22 | from motor.engine_server.core.engine import Engine |
| 23 | from motor.engine_server.core.vllm.vllm_openai_compat import cli_env_setup | 23 | from motor.engine_server.core.vllm.vllm_openai_compat import cli_env_setup |
| 24 | 24 | ||
| @@ -21,7 +21,7 @@ import importlib | |||
| 21 | from typing import Any | 21 | from typing import Any |
| 22 | 22 | ||
| 23 | from motor.common.logger import get_logger | 23 | from motor.common.logger import get_logger |
| 24 | -from motor.engine_server.core.config import IConfig | 24 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 25 | 25 | ||
| 26 | logger = get_logger(__name__) | 26 | logger = get_logger(__name__) |
| 27 | 27 | ||
| @@ -37,10 +37,7 @@ class EndpointFactory: | |||
| 37 | target = self._CREATOR_MAP.get(engine_type) | 37 | target = self._CREATOR_MAP.get(engine_type) |
| 38 | if not target: | 38 | if not target: |
| 39 | supported_types = list(self._CREATOR_MAP.keys()) | 39 | supported_types = list(self._CREATOR_MAP.keys()) |
| 40 | - raise ValueError( | 40 | + raise ValueError(f"Unsupported engine type: {engine_type}. Supported types are: {supported_types}.") |
| 41 | - f"Unsupported engine type: {engine_type}. " | ||
| 42 | - f"Supported types are: {supported_types}." | ||
| 43 | - ) | ||
| 44 | 41 | ||
| 45 | try: | 42 | try: |
| 46 | module_path, class_name = target.rsplit(".", 1) | 43 | module_path, class_name = target.rsplit(".", 1) |
| @@ -1,40 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -from motor.common.http.http_client import SafeHTTPSClient | ||
| 12 | -from motor.common.logger import get_logger | ||
| 13 | -from motor.common.logger.rate_limited_logger import RateLimitedLogger | ||
| 14 | -from motor.config.node_manager import NodeManagerConfig | ||
| 15 | - | ||
| 16 | -logger = get_logger(__name__) | ||
| 17 | -_rl = RateLimitedLogger(logger) | ||
| 18 | - | ||
| 19 | - | ||
| 20 | -class EngineServerApiClient: | ||
| 21 | - tls_config = NodeManagerConfig.from_json().mgmt_tls_config | ||
| 22 | - | ||
| 23 | - | ||
| 24 | - def query_status(address: str): | ||
| 25 | - client_args = EngineServerApiClient._generate_client_args(address) | ||
| 26 | - client = SafeHTTPSClient(**client_args, timeout=5) | ||
| 27 | - response = client.get("/status") | ||
| 28 | - _rl.record_success(f"node_manager.engine_server.query_status.{address}") | ||
| 29 | - _rl.emit_periodic( | ||
| 30 | - f"node_manager.engine_server.query_status.{address}", | ||
| 31 | - "NodeManager->EngineServer query_status periodic summary: succeeded {count} times in last 60s", | ||
| 32 | - level="DEBUG", | ||
| 33 | - ) | ||
| 34 | - logger.debug(f"Query engine server status success, response: {response}, address: {client_args['address']}") | ||
| 35 | - return response | ||
| 36 | - | ||
| 37 | - | ||
| 38 | - def _generate_client_args(cls, address: str) -> dict[str, str]: | ||
| 39 | - client_ars = {"address": f"{address}", "tls_config": cls.tls_config} | ||
| 40 | - return client_ars | ||
| @@ -99,7 +99,7 @@ async def start_instance(request: Request): | |||
| 99 | except Exception as pull_err: | 99 | except Exception as pull_err: |
| 100 | logger.error("Failed to pull engine: %s", pull_err) | 100 | logger.error("Failed to pull engine: %s", pull_err) |
| 101 | raise HTTPException( | 101 | raise HTTPException( |
| 102 | - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to start engine server" | 102 | + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to start native engine" |
| 103 | ) from pull_err | 103 | ) from pull_err |
| 104 | 104 | ||
| 105 | # Start KV store service (if backend is configured) | 105 | # Start KV store service (if backend is configured) |
| @@ -154,11 +154,11 @@ async def pause_instance(request: Request): | |||
| 154 | try: | 154 | try: |
| 155 | await asyncio.to_thread(HeartbeatManager().pause_all_endpoints) | 155 | await asyncio.to_thread(HeartbeatManager().pause_all_endpoints) |
| 156 | hm = HeartbeatManager() | 156 | hm = HeartbeatManager() |
| 157 | - engine_mgmt_addrs = hm.get_engine_mgmt_addrs() | 157 | + engine_metrics_targets = hm.get_engine_metrics_targets() |
| 158 | content = { | 158 | content = { |
| 159 | "status": "ok", | 159 | "status": "ok", |
| 160 | "message": "Endpoints set to PAUSED", | 160 | "message": "Endpoints set to PAUSED", |
| 161 | - "engine_mgmt_addrs": engine_mgmt_addrs, | 161 | + "engine_metrics_targets": engine_metrics_targets, |
| 162 | } | 162 | } |
| 163 | return Response(status_code=status.HTTP_200_OK, content=json.dumps(content)) | 163 | return Response(status_code=status.HTTP_200_OK, content=json.dumps(content)) |
| 164 | except Exception as err: | 164 | except Exception as err: |
| @@ -12,6 +12,7 @@ import threading | |||
| 12 | 12 | ||
| 13 | from motor.common.resources.instance import PDRole | 13 | from motor.common.resources.instance import PDRole |
| 14 | from motor.common.resources.endpoint import Endpoint | 14 | from motor.common.resources.endpoint import Endpoint |
| 15 | +from motor.node_manager.core.services.native_engine.models import RuntimeState | ||
| 15 | from motor.common.utils.singleton import ThreadSafeSingleton | 16 | from motor.common.utils.singleton import ThreadSafeSingleton |
| 16 | from motor.common.logger import get_logger | 17 | from motor.common.logger import get_logger |
| 17 | from motor.config.node_manager import NodeManagerConfig | 18 | from motor.config.node_manager import NodeManagerConfig |
| @@ -125,6 +126,19 @@ class Daemon(ThreadSafeSingleton): | |||
| 125 | """True when the engine service is active (i.e. this pod runs inference).""" | 126 | """True when the engine service is active (i.e. this pod runs inference).""" |
| 126 | return SERVICE_ENGINE in self._services | 127 | return SERVICE_ENGINE in self._services |
| 127 | 128 | ||
| 129 | + def get_engine_runtime_state(self, endpoint: Endpoint) -> RuntimeState: | ||
| 130 | + """Return the native runtime state for one locally managed endpoint.""" | ||
| 131 | + engine = self._services.get(SERVICE_ENGINE) | ||
| 132 | + if engine is None: | ||
| 133 | + return RuntimeState.STOPPED | ||
| 134 | + return engine.runtime_state(endpoint) # type: ignore[attr-defined] | ||
| 135 | + | ||
| 136 | + def get_engine_metrics_target(self, endpoint: Endpoint) -> str | None: | ||
| 137 | + engine = self._services.get(SERVICE_ENGINE) | ||
| 138 | + if engine is None: | ||
| 139 | + return None | ||
| 140 | + return engine.metrics_target(endpoint) # type: ignore[attr-defined] | ||
| 141 | + | ||
| 128 | def stop(self) -> None: | 142 | def stop(self) -> None: |
| 129 | self._monitor_stop.set() | 143 | self._monitor_stop.set() |
| 130 | if self._monitor_thread is not None and self._monitor_thread.is_alive(): | 144 | if self._monitor_thread is not None and self._monitor_thread.is_alive(): |
| @@ -364,6 +364,7 @@ class EngineManager(ThreadSafeSingleton): | |||
| 364 | pod_ip = self._config.api_config.pod_ip | 364 | pod_ip = self._config.api_config.pod_ip |
| 365 | business_port = self._config.endpoint_config.service_ports | 365 | business_port = self._config.endpoint_config.service_ports |
| 366 | mgmt_port = self._config.endpoint_config.mgmt_ports | 366 | mgmt_port = self._config.endpoint_config.mgmt_ports |
| 367 | + bootstrap_port = self._config.endpoint_config.bootstrap_port | ||
| 367 | node_manager_port = self._config.api_config.node_manager_port | 368 | node_manager_port = self._config.api_config.node_manager_port |
| 368 | parallel_config = self._config.basic_config.parallel_config | 369 | parallel_config = self._config.basic_config.parallel_config |
| 369 | enable_multi_endpoints = self._config.basic_config.enable_multi_endpoints | 370 | enable_multi_endpoints = self._config.basic_config.enable_multi_endpoints |
| @@ -379,6 +380,7 @@ class EngineManager(ThreadSafeSingleton): | |||
| 379 | pod_ip=pod_ip, | 380 | pod_ip=pod_ip, |
| 380 | business_port=business_port, | 381 | business_port=business_port, |
| 381 | mgmt_port=mgmt_port, | 382 | mgmt_port=mgmt_port, |
| 383 | + bootstrap_port=bootstrap_port, | ||
| 382 | nm_port=str(node_manager_port), | 384 | nm_port=str(node_manager_port), |
| 383 | parallel_config=parallel_config, | 385 | parallel_config=parallel_config, |
| 384 | enable_multi_endpoints=enable_multi_endpoints, | 386 | enable_multi_endpoints=enable_multi_endpoints, |
| @@ -8,10 +8,8 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -import os | ||
| 12 | import threading | 11 | import threading |
| 13 | import time | 12 | import time |
| 14 | -import socket | ||
| 15 | 13 | ||
| 16 | from motor.common.resources.endpoint import Endpoint, EndpointStatus | 14 | from motor.common.resources.endpoint import Endpoint, EndpointStatus |
| 17 | from motor.common.resources.http_msg_spec import StartCmdMsg, HeartbeatMsg | 15 | from motor.common.resources.http_msg_spec import StartCmdMsg, HeartbeatMsg |
| @@ -22,9 +20,9 @@ from motor.common.utils.singleton import ThreadSafeSingleton | |||
| 22 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, RETRY_LOG_FREQUENCY | 20 | from motor.common.utils.snapshot_utils import is_restored_from_host_side_snapshot, RETRY_LOG_FREQUENCY |
| 23 | from motor.config.node_manager import NodeManagerConfig | 21 | from motor.config.node_manager import NodeManagerConfig |
| 24 | from motor.node_manager.api_client.controller_api_client import ControllerApiClient | 22 | from motor.node_manager.api_client.controller_api_client import ControllerApiClient |
| 25 | -from motor.node_manager.api_client.engine_server_api_client import EngineServerApiClient | ||
| 26 | from motor.node_manager.core.engine_manager import EngineManager | 23 | from motor.node_manager.core.engine_manager import EngineManager |
| 27 | from motor.node_manager.core.daemon import Daemon | 24 | from motor.node_manager.core.daemon import Daemon |
| 25 | +from motor.node_manager.core.services.native_engine.models import RuntimeState | ||
| 28 | 26 | ||
| 29 | 27 | ||
| 30 | logger = get_logger(__name__) | 28 | logger = get_logger(__name__) |
| @@ -55,14 +53,12 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 55 | daemon=True, | 53 | daemon=True, |
| 56 | name="heartbeat_report", | 54 | name="heartbeat_report", |
| 57 | ) | 55 | ) |
| 58 | - self._engine_server_status_thread = threading.Thread( | 56 | + self._engine_status_thread = threading.Thread( |
| 59 | target=self._refresh_endpoints_status_loop, | 57 | target=self._refresh_endpoints_status_loop, |
| 60 | daemon=True, | 58 | daemon=True, |
| 61 | name="endpoint_status_fetch", | 59 | name="endpoint_status_fetch", |
| 62 | ) | 60 | ) |
| 63 | self._thread_started = False | 61 | self._thread_started = False |
| 64 | - self._engine_status_thread_start_time = None | ||
| 65 | - self._is_within_grace_period = True | ||
| 66 | self._consecutive_abnormal_count = 0 | 62 | self._consecutive_abnormal_count = 0 |
| 67 | self._abnormal_count_lock = threading.Lock() | 63 | self._abnormal_count_lock = threading.Lock() |
| 68 | self._should_suicide = False | 64 | self._should_suicide = False |
| @@ -81,8 +77,7 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 81 | def start(self): | 77 | def start(self): |
| 82 | if self._thread_started is False: | 78 | if self._thread_started is False: |
| 83 | self._heartbeat_report_thread.start() | 79 | self._heartbeat_report_thread.start() |
| 84 | - self._engine_server_status_thread.start() | 80 | + self._engine_status_thread.start() |
| 85 | - self._engine_status_thread_start_time = time.time() | ||
| 86 | self._thread_started = True | 81 | self._thread_started = True |
| 87 | else: | 82 | else: |
| 88 | logger.info("Heartbeat thread has been started...") | 83 | logger.info("Heartbeat thread has been started...") |
| @@ -109,9 +104,6 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 109 | # Reset suicide flag when endpoints are updated | 104 | # Reset suicide flag when endpoints are updated |
| 110 | with self._suicide_lock: | 105 | with self._suicide_lock: |
| 111 | self._should_suicide = False | 106 | self._should_suicide = False |
| 112 | - self._is_within_grace_period = True | ||
| 113 | - if self._thread_started: | ||
| 114 | - self._engine_status_thread_start_time = time.time() | ||
| 115 | 107 | ||
| 116 | def should_suicide(self) -> bool: | 108 | def should_suicide(self) -> bool: |
| 117 | """ | 109 | """ |
| @@ -125,8 +117,8 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 125 | self.stop_event.set() | 117 | self.stop_event.set() |
| 126 | if self._heartbeat_report_thread.is_alive(): | 118 | if self._heartbeat_report_thread.is_alive(): |
| 127 | self._heartbeat_report_thread.join(timeout=2.0) | 119 | self._heartbeat_report_thread.join(timeout=2.0) |
| 128 | - if self._engine_server_status_thread.is_alive(): | 120 | + if self._engine_status_thread.is_alive(): |
| 129 | - self._engine_server_status_thread.join(timeout=2.0) | 121 | + self._engine_status_thread.join(timeout=2.0) |
| 130 | logger.info("HeartBeatManager stopped.") | 122 | logger.info("HeartBeatManager stopped.") |
| 131 | 123 | ||
| 132 | def check_all_endpoints_normal(self) -> bool: | 124 | def check_all_endpoints_normal(self) -> bool: |
| @@ -166,10 +158,12 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 166 | endpoint.status = EndpointStatus.PAUSED | 158 | endpoint.status = EndpointStatus.PAUSED |
| 167 | logger.info("All endpoints set to PAUSED for graceful shutdown") | 159 | logger.info("All endpoints set to PAUSED for graceful shutdown") |
| 168 | 160 | ||
| 169 | - def get_engine_mgmt_addrs(self) -> list[str]: | 161 | + def get_engine_metrics_targets(self) -> list[str]: |
| 170 | - """Return engine management addresses for local metrics polling.""" | 162 | + """Return native metrics URLs for routable local endpoints.""" |
| 171 | with self._endpoint_lock: | 163 | with self._endpoint_lock: |
| 172 | - return [format_address(ep.ip, ep.mgmt_port) for ep in self._endpoints] | 164 | + endpoints = [endpoint for endpoint in self._endpoints if not endpoint.headless] |
| 165 | + daemon = Daemon() | ||
| 166 | + return [target for endpoint in endpoints if (target := daemon.get_engine_metrics_target(endpoint)) is not None] | ||
| 173 | 167 | ||
| 174 | def resume_all_endpoints(self) -> None: | 168 | def resume_all_endpoints(self) -> None: |
| 175 | """Resume all endpoints from PAUSED back to NORMAL status. | 169 | """Resume all endpoints from PAUSED back to NORMAL status. |
| @@ -192,35 +186,11 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 192 | self._is_started_after_restore = is_started | 186 | self._is_started_after_restore = is_started |
| 193 | 187 | ||
| 194 | def _refresh_endpoints_status_loop(self) -> None: | 188 | def _refresh_endpoints_status_loop(self) -> None: |
| 195 | - # Poll each engine server's mgmt port until it responds (max 60s) | ||
| 196 | - self._wait_for_engine_servers_ready(timeout=60) | ||
| 197 | while not self.stop_event.is_set(): | 189 | while not self.stop_event.is_set(): |
| 198 | - self._get_engine_server_status() | 190 | + self._refresh_native_engine_status() |
| 199 | - time.sleep(1) | 191 | + self.stop_event.wait(1) |
| 200 | 192 | ||
| 201 | - def _wait_for_engine_servers_ready(self, timeout: float = 60) -> None: | 193 | + def _refresh_native_engine_status(self) -> None: |
| 202 | - """Poll each endpoint's mgmt port until it accepts connections or timeout.""" | ||
| 203 | - with self._endpoint_lock: | ||
| 204 | - endpoints = list(self._endpoints) | ||
| 205 | - | ||
| 206 | - deadline = time.time() + timeout | ||
| 207 | - daemon = Daemon() | ||
| 208 | - for endpoint in endpoints: | ||
| 209 | - address = f"{endpoint.ip}:{endpoint.mgmt_port}" | ||
| 210 | - logger.info("Waiting for engine server at %s to become ready...", address) | ||
| 211 | - while not self.stop_event.is_set() and time.time() < deadline: | ||
| 212 | - try: | ||
| 213 | - with socket.create_connection((endpoint.ip, int(endpoint.mgmt_port)), timeout=2): | ||
| 214 | - logger.info("Engine server at %s is ready.", address) | ||
| 215 | - break | ||
| 216 | - except (OSError, ConnectionRefusedError, TimeoutError): | ||
| 217 | - # Check if engine process is still alive | ||
| 218 | - if not any(os.path.isdir(f"/proc/{pid}") for pid in daemon.engine_pids): | ||
| 219 | - logger.error("Engine process for %s is no longer running, aborting wait.", address) | ||
| 220 | - break | ||
| 221 | - time.sleep(1) | ||
| 222 | - | ||
| 223 | - def _get_engine_server_status(self) -> None: | ||
| 224 | with self._endpoint_lock: | 194 | with self._endpoint_lock: |
| 225 | endpoints_snapshot = list(self._endpoints) | 195 | endpoints_snapshot = list(self._endpoints) |
| 226 | generation_at_start = self._endpoints_generation | 196 | generation_at_start = self._endpoints_generation |
| @@ -228,50 +198,26 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 228 | if not endpoints_snapshot: | 198 | if not endpoints_snapshot: |
| 229 | return | 199 | return |
| 230 | 200 | ||
| 231 | - # Check if within one minute after startup | ||
| 232 | - if self._is_within_grace_period and self._engine_status_thread_start_time is not None: | ||
| 233 | - elapsed_time = time.time() - self._engine_status_thread_start_time | ||
| 234 | - self._is_within_grace_period = elapsed_time < 120 | ||
| 235 | - | ||
| 236 | updated_endpoints = [] | 201 | updated_endpoints = [] |
| 237 | - client = None | 202 | + daemon = Daemon() |
| 238 | for item in endpoints_snapshot: | 203 | for item in endpoints_snapshot: |
| 239 | original_status = item.status | 204 | original_status = item.status |
| 240 | - client = None | ||
| 241 | - detected_status = None | ||
| 242 | - engine_server_base_url = format_address(item.ip, item.mgmt_port) | ||
| 243 | try: | 205 | try: |
| 244 | - response = EngineServerApiClient.query_status(engine_server_base_url) | 206 | + runtime_state = daemon.get_engine_runtime_state(item) |
| 245 | - if isinstance(response, dict) and "status" in response: | ||
| 246 | - status_value = response.get("status") | ||
| 247 | - try: | ||
| 248 | - detected_status = EndpointStatus(status_value) | ||
| 249 | - except ValueError: | ||
| 250 | - logger.error( | ||
| 251 | - "Invalid status value '%s' from Engine Server %d: %s", | ||
| 252 | - status_value, | ||
| 253 | - item.id, | ||
| 254 | - engine_server_base_url, | ||
| 255 | - ) | ||
| 256 | - detected_status = EndpointStatus.ABNORMAL | ||
| 257 | - else: | ||
| 258 | - logger.error( | ||
| 259 | - "Invalid response format from Engine Server%d: %s: %s", | ||
| 260 | - item.id, | ||
| 261 | - engine_server_base_url, | ||
| 262 | - response, | ||
| 263 | - ) | ||
| 264 | - detected_status = EndpointStatus.ABNORMAL | ||
| 265 | except Exception as e: | 207 | except Exception as e: |
| 266 | - if not self._is_within_grace_period: | 208 | + logger.error( |
| 267 | - logger.error("Failed to get engine server status from %s: %s", engine_server_base_url, e) | 209 | + "Failed to probe native engine at %s: %s", |
| 268 | - detected_status = EndpointStatus.ABNORMAL | 210 | + format_address(item.ip, item.business_port), |
| 269 | - finally: | 211 | + e, |
| 270 | - if client is not None: | 212 | + ) |
| 271 | - try: | 213 | + runtime_state = RuntimeState.UNHEALTHY |
| 272 | - client.close() | 214 | + |
| 273 | - except Exception as e: | 215 | + detected_status = { |
| 274 | - logger.error("Failed to close client: %s", e) | 216 | + RuntimeState.RUNNING: EndpointStatus.WAIT2START, |
| 217 | + RuntimeState.READY: EndpointStatus.NORMAL, | ||
| 218 | + RuntimeState.UNHEALTHY: EndpointStatus.ABNORMAL, | ||
| 219 | + RuntimeState.STOPPED: EndpointStatus.ABNORMAL, | ||
| 220 | + }.get(runtime_state) | ||
| 275 | 221 | ||
| 276 | if is_restored_from_host_side_snapshot() and item.ip != self._config.api_config.pod_ip: | 222 | if is_restored_from_host_side_snapshot() and item.ip != self._config.api_config.pod_ip: |
| 277 | # If restored from host side snapshot and not started after restore(pod_ip do not refresh yet), keep original status | 223 | # If restored from host side snapshot and not started after restore(pod_ip do not refresh yet), keep original status |
| @@ -281,11 +227,13 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 281 | original_status, | 227 | original_status, |
| 282 | ) | 228 | ) |
| 283 | item.status = original_status | 229 | item.status = original_status |
| 284 | - elif self._is_within_grace_period and detected_status == EndpointStatus.ABNORMAL: | 230 | + elif runtime_state in (RuntimeState.STARTING, RuntimeState.STOPPING): |
G 严重程度： 提示 问题： 原因： 怎么改： 要么让 state() 在 STOPPING 时短路返回 STOPPING（探测前先查 state），要么删掉心跳侧的死分支并补充注释说明“停止中按 ABNORMAL 处理”。 tobking 19 天前 评论： 19 天前 评论： | |||
| 285 | - # If within grace period and abnormal status detected, do not update status | 231 | + # Loading is not a failure. Keep INITIAL (or the last reported |
| 232 | + # status) until the native readiness endpoint succeeds. | ||
| 286 | logger.debug( | 233 | logger.debug( |
| 287 | - "Engine server %s status is abnormal within grace period, keeping original status: %s", | 234 | + "Native engine %s is %s, keeping status %s", |
| 288 | - engine_server_base_url, | 235 | + format_address(item.ip, item.business_port), |
| 236 | + runtime_state.value, | ||
| 289 | original_status, | 237 | original_status, |
| 290 | ) | 238 | ) |
| 291 | item.status = original_status | 239 | item.status = original_status |
| @@ -297,7 +245,7 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 297 | 245 | ||
| 298 | if item.status != original_status: | 246 | if item.status != original_status: |
| 299 | logger.info( | 247 | logger.info( |
| 300 | - "Engine Server rank %d, status change from %s to %s ", | 248 | + "Native engine rank %d, status change from %s to %s ", |
| 301 | item.id, | 249 | item.id, |
| 302 | original_status, | 250 | original_status, |
| 303 | item.status, | 251 | item.status, |
| @@ -316,8 +264,7 @@ class HeartbeatManager(ThreadSafeSingleton): | |||
| 316 | is_normal = True | 264 | is_normal = True |
| 317 | try: | 265 | try: |
| 318 | with self._endpoint_lock: | 266 | with self._endpoint_lock: |
| 319 | - # Check if any endpoint has abnormal status (only after grace period) | 267 | + # Check actual endpoint status, not the reported status. |
| 320 | - # Check actual endpoint status, not the reported status | ||
| 321 | has_abnormal = any(item.status == EndpointStatus.ABNORMAL for item in self._endpoints) | 268 | has_abnormal = any(item.status == EndpointStatus.ABNORMAL for item in self._endpoints) |
| 322 | is_normal = all(item.status == EndpointStatus.NORMAL for item in self._endpoints) | 269 | is_normal = all(item.status == EndpointStatus.NORMAL for item in self._endpoints) |
| 323 | 270 | ||
| @@ -0,0 +1,11 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""NodeManager service for launching and supervising native inference engines.""" | ||
Rmotor/engine_server/constants/__init__.py→motor/node_manager/core/services/native_engine/backends/__init__.py+4-4
| @@ -1,6 +1,4 @@ | |||
| 1 | -#!/usr/bin/env python3 | 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. |
| 2 | -# -*- coding: utf-8 -*- | ||
| 3 | -# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 4 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 5 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| 6 | # You may obtain a copy of Mulan PSL v2 at: | 4 | # You may obtain a copy of Mulan PSL v2 at: |
| @@ -8,4 +6,6 @@ | |||
| 8 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, |
| 9 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | 7 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, |
| 10 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 11 | -# See the Mulan PSL v2 for more details. | 9 | +# See the Mulan PSL v2 for more details. |
| 10 | + | ||
| 11 | +"""Engine-specific native launch backends.""" | ||
| @@ -0,0 +1,132 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +import argparse | ||
| 12 | +from abc import ABC, abstractmethod | ||
| 13 | +from typing import Protocol | ||
| 14 | + | ||
| 15 | +from motor.common.logger import get_logger | ||
| 16 | +from motor.config.endpoint import EndpointConfig | ||
| 17 | +from motor.node_manager.core.services.native_engine.models import ( | ||
| 18 | + CommandSpec, | ||
| 19 | + LaunchContext, | ||
| 20 | + LaunchSpec, | ||
| 21 | + ProbeSpec, | ||
| 22 | +) | ||
| 23 | + | ||
| 24 | +logger = get_logger(__name__) | ||
| 25 | + | ||
| 26 | +supported_engine = ["vllm", "sglang"] | ||
| 27 | +supported_role = ["prefill", "decode", "union"] | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +class IConfig(ABC): | ||
| 31 | + | ||
| 32 | + def initialize(self): | ||
| 33 | + pass | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + def validate(self): | ||
| 37 | + pass | ||
| 38 | + | ||
| 39 | + | ||
| 40 | + def convert(self): | ||
| 41 | + pass | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + def get_args(self) -> argparse.Namespace | None: | ||
| 45 | + pass | ||
| 46 | + | ||
| 47 | + | ||
| 48 | + def get_endpoint_config(self) -> EndpointConfig | None: | ||
| 49 | + pass | ||
| 50 | + | ||
| 51 | + | ||
| 52 | + def get_cli_args(self) -> list[str]: | ||
| 53 | + """Return the CLI argument list suitable for native engine launch (vllm serve / sglang.launch_server).""" | ||
| 54 | + pass | ||
| 55 | + | ||
| 56 | + | ||
| 57 | +class NativeEngineBackend(Protocol): | ||
| 58 | + """Stateless conversion from launch context to a native engine launch specification.""" | ||
| 59 | + | ||
| 60 | + engine_type: str | ||
| 61 | + | ||
| 62 | + def prepare(self, context: LaunchContext) -> LaunchSpec: | ||
| 63 | + """Validate one endpoint and build its immutable launch specification.""" | ||
| 64 | + ... | ||
| 65 | + | ||
| 66 | + | ||
| 67 | +class BaseNativeEngineBackend: | ||
| 68 | + """Shared native launch-spec construction for engine-specific backends.""" | ||
| 69 | + | ||
| 70 | + engine_type: str | ||
| 71 | + command_prefix: tuple[str, ...] | ||
| 72 | + | ||
| 73 | + def prepare(self, context: LaunchContext) -> LaunchSpec: | ||
| 74 | + self._validate_context(context) | ||
| 75 | + endpoint_config = build_endpoint_config(context, self.engine_type) | ||
| 76 | + if endpoint_config.engine_type != self.engine_type: | ||
| 77 | + raise ValueError( | ||
| 78 | + f"Configured engine type {endpoint_config.engine_type} does not match " | ||
| 79 | + f"Node Manager engine type {self.engine_type}" | ||
| 80 | + ) | ||
| 81 | + self._validate_endpoint_config(endpoint_config) | ||
| 82 | + | ||
| 83 | + # Import lazily so ConfigFactory can type against IConfig without a module cycle. | ||
| 84 | + from motor.node_manager.core.services.native_engine.config_factory import ConfigFactory | ||
| 85 | + | ||
| 86 | + config = ConfigFactory(endpoint_config=endpoint_config).build_cli_config() | ||
G 严重程度： 建议 问题： 原生路径不执行 convert()/validate()：非法 CLI 参数只有在 vllm/sglang 进程退出时才以笼统错误暴露。 原因： 此处只调用 怎么改：
在 prepare()（或 start() 前）以“仅校验不启动”方式跑一遍参数校验：对 失败时抛带具体参数名的 ValueError，让 pull 的报错包含真实原因。 tobking 19 天前 评论： 19 天前 评论： | |||
| 87 | + config.convert() | ||
| 88 | + config.validate() | ||
| 89 | + health_config = endpoint_config.deploy_config.health_check_config | ||
| 90 | + return LaunchSpec( | ||
| 91 | + command=CommandSpec( | ||
| 92 | + argv=self.command_prefix + tuple(config.get_cli_args()), | ||
| 93 | + env=context.environment, | ||
| 94 | + ), | ||
| 95 | + probe=ProbeSpec( | ||
| 96 | + path="/health", | ||
| 97 | + timeout_seconds=float(health_config.health_collector_timeout), | ||
G 严重程度： 提示 问题： 原因： 此处只取 health_collector_timeout；supervisor.py:86-95 每次 state() 新建 SafeHTTPSClient（新 requests.Session，TLS 时新 SSL context）做单次 do_get。基线 HealthCollector 对超时按 health_collector_timeout_retry_attempts（默认 3）重试。HealthCheckConfig 仍保留该字段，但新探针完全忽略，配置字段失效；多个 endpoint 在模型加载期（默认 startup_timeout 1800s）每秒各建一个 Session+SSL context，属不必要的开销。 怎么改： supervisor 内缓存按 (host,port,tls) 复用的客户端，或在 state() 中对超时类异常按 attempts 重试一次；同时决定该配置字段的去留（删除或实现）。 tobking 19 天前 评论： 19 天前 评论： | |||
| 98 | + startup_timeout_seconds=float(health_config.startup_timeout), | ||
| 99 | + max_attempts=health_config.health_collector_timeout_retry_attempts, | ||
| 100 | + tls_config=endpoint_config.deploy_config.infer_tls_config, | ||
| 101 | + process_only=context.headless, | ||
| 102 | + ), | ||
| 103 | + ) | ||
| 104 | + | ||
| 105 | + def _validate_context(self, context: LaunchContext) -> None: | ||
| 106 | + pass | ||
| 107 | + | ||
| 108 | + def _validate_endpoint_config(self, endpoint_config: EndpointConfig) -> None: | ||
| 109 | + pass | ||
| 110 | + | ||
| 111 | + | ||
| 112 | +def build_endpoint_config(context: LaunchContext, engine_type: str) -> EndpointConfig: | ||
| 113 | + """Rebuild and validate one role-specific endpoint configuration.""" | ||
| 114 | + endpoint_config = EndpointConfig( | ||
| 115 | + engine_type=engine_type, | ||
| 116 | + host=context.host, | ||
| 117 | + role=context.role.value, | ||
| 118 | + kv_port=context.kv_port, | ||
| 119 | + lookup_rpc_port=context.lookup_rpc_port, | ||
| 120 | + master_dp_ip=context.master_dp_ip, | ||
| 121 | + dp_rpc_port=context.dp_rpc_port, | ||
| 122 | + port=context.business_port, | ||
| 123 | + mgmt_port=context.mgmt_port, | ||
| 124 | + instance_id=context.instance_id, | ||
| 125 | + dp_rank=context.dp_rank, | ||
| 126 | + node_rank=context.node_rank, | ||
| 127 | + config_path=context.config_path, | ||
| 128 | + d2d_peer_ips=",".join(context.d2d_peer_ips) if context.d2d_peer_ips else None, | ||
| 129 | + ) | ||
| 130 | + endpoint_config.validate() | ||
| 131 | + endpoint_config.load_deploy_config() | ||
| 132 | + return endpoint_config | ||
| @@ -0,0 +1,11 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""Native SGLang backend.""" | ||
| @@ -0,0 +1,24 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from motor.common.resources.instance import PDRole | ||
| 12 | +from motor.node_manager.core.services.native_engine.backends.base import BaseNativeEngineBackend | ||
| 13 | +from motor.node_manager.core.services.native_engine.models import LaunchContext | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +class SGLangBackend(BaseNativeEngineBackend): | ||
| 17 | + """Build and validate native SGLang launch specifications.""" | ||
| 18 | + | ||
| 19 | + engine_type = "sglang" | ||
| 20 | + command_prefix = ("python3", "-m", "sglang.launch_server") | ||
| 21 | + | ||
| 22 | + def _validate_context(self, context: LaunchContext) -> None: | ||
| 23 | + if context.role == PDRole.ROLE_E: | ||
| 24 | + raise ValueError("SGLang encode role is not supported") | ||
Rmotor/engine_server/core/sglang/sglang_config.py→motor/node_manager/core/services/native_engine/backends/sglang/config.py+20-6
| @@ -10,21 +10,22 @@ | |||
| 10 | 10 | ||
| 11 | import argparse | 11 | import argparse |
| 12 | import json | 12 | import json |
| 13 | -import sys | ||
| 14 | from dataclasses import dataclass | 13 | from dataclasses import dataclass |
| 15 | from typing import Any | 14 | from typing import Any |
| 16 | 15 | ||
| 17 | from motor.common.logger import get_logger | 16 | from motor.common.logger import get_logger |
| 18 | from motor.common.utils.net import format_address | 17 | from motor.common.utils.net import format_address |
| 19 | from motor.config.endpoint import EndpointConfig | 18 | from motor.config.endpoint import EndpointConfig |
| 20 | -from motor.engine_server.constants import constants | 19 | +from motor.common import engine_constants as constants |
| 21 | -from motor.engine_server.core.config import IConfig | 20 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 22 | 21 | ||
| 23 | logger = get_logger(__name__) | 22 | logger = get_logger(__name__) |
| 24 | 23 | ||
| 25 | 24 | ||
| 26 | def _add_argument_to_list(arg_list: list, key: str, value: Any): | 25 | def _add_argument_to_list(arg_list: list, key: str, value: Any): |
| 27 | """Append key-value to arg_list as CLI args (e.g. --key value).""" | 26 | """Append key-value to arg_list as CLI args (e.g. --key value).""" |
| 27 | + if value is None: | ||
| 28 | + return | ||
| 28 | if isinstance(value, bool): | 29 | if isinstance(value, bool): |
| 29 | if value: | 30 | if value: |
| 30 | arg_list.append(f"--{key}") | 31 | arg_list.append(f"--{key}") |
| @@ -58,12 +59,11 @@ class SGLangConfig(IConfig): | |||
| 58 | arg_list = self._get_param_list() | 59 | arg_list = self._get_param_list() |
| 59 | logger.info("engine server sglang arg_list: %s", arg_list) | 60 | logger.info("engine server sglang arg_list: %s", arg_list) |
| 60 | 61 | ||
| 61 | - sys.argv = ["serve"] + arg_list | ||
| 62 | from sglang.srt.server_args import ServerArgs | 62 | from sglang.srt.server_args import ServerArgs |
| 63 | 63 | ||
| 64 | parser = argparse.ArgumentParser() | 64 | parser = argparse.ArgumentParser() |
| 65 | ServerArgs.add_cli_args(parser) | 65 | ServerArgs.add_cli_args(parser) |
| 66 | - raw_args = parser.parse_args() | 66 | + raw_args = parser.parse_args(arg_list) |
| 67 | self.args = ServerArgs.from_cli_args(raw_args) | 67 | self.args = ServerArgs.from_cli_args(raw_args) |
| 68 | 68 | ||
| 69 | def get_args(self) -> argparse.Namespace: | 69 | def get_args(self) -> argparse.Namespace: |
| @@ -87,11 +87,25 @@ class SGLangConfig(IConfig): | |||
| 87 | role = self.endpoint_config.role | 87 | role = self.endpoint_config.role |
| 88 | 88 | ||
| 89 | flattened.update(deploy_config.engine_config.configs) | 89 | flattened.update(deploy_config.engine_config.configs) |
| 90 | + # Coordinator metrics aggregation and Kubernetes drain depend on this endpoint. | ||
| 91 | + flattened["enable-metrics"] = True | ||
| 90 | 92 | ||
| 91 | flattened["host"] = self.endpoint_config.host | 93 | flattened["host"] = self.endpoint_config.host |
| 92 | flattened["port"] = self.endpoint_config.port | 94 | flattened["port"] = self.endpoint_config.port |
| 93 | 95 | ||
| 94 | - if flattened.get("nnodes", 1) > 1: | 96 | + raw_nnodes = flattened.get("nnodes", 1) |
| 97 | + try: | ||
| 98 | + nnodes = int(raw_nnodes) | ||
| 99 | + except (TypeError, ValueError) as exc: | ||
| 100 | + raise ValueError(f"nnodes must be an integer, got {raw_nnodes!r}") from exc | ||
| 101 | + if nnodes < 1: | ||
| 102 | + raise ValueError(f"nnodes must be greater than 0, got {nnodes}") | ||
| 103 | + if "nnodes" in flattened: | ||
| 104 | + flattened["nnodes"] = nnodes | ||
| 105 | + | ||
| 106 | + if nnodes > 1: | ||
| 107 | + if not self.endpoint_config.master_dp_ip: | ||
| 108 | + raise ValueError("master_dp_ip is required when nnodes > 1") | ||
| 95 | parallel_config = deploy_config.get_parallel_config(role) | 109 | parallel_config = deploy_config.get_parallel_config(role) |
| 96 | flattened["dist-init-addr"] = format_address(self.endpoint_config.master_dp_ip, parallel_config.dp_rpc_port) | 110 | flattened["dist-init-addr"] = format_address(self.endpoint_config.master_dp_ip, parallel_config.dp_rpc_port) |
| 97 | flattened["node-rank"] = self.endpoint_config.node_rank | 111 | flattened["node-rank"] = self.endpoint_config.node_rank |
| @@ -0,0 +1,11 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""Native vLLM backend.""" | ||
| @@ -0,0 +1,34 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from motor.common.resources.dispatch import DispatchProfile, classify_vllm_dispatch_profile | ||
| 12 | +from motor.common.resources.instance import PDRole | ||
| 13 | +from motor.config.endpoint import EndpointConfig | ||
| 14 | +from motor.node_manager.core.services.native_engine.backends.base import BaseNativeEngineBackend | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +class VllmBackend(BaseNativeEngineBackend): | ||
| 18 | + """Build and validate native ``vllm serve`` launch specifications.""" | ||
| 19 | + | ||
| 20 | + engine_type = "vllm" | ||
| 21 | + command_prefix = ("vllm", "serve") | ||
| 22 | + | ||
| 23 | + def _validate_endpoint_config(self, endpoint_config: EndpointConfig) -> None: | ||
| 24 | + if endpoint_config.role not in (PDRole.ROLE_P.value, PDRole.ROLE_D.value): | ||
| 25 | + return | ||
| 26 | + deploy_config = endpoint_config.deploy_config | ||
| 27 | + profile = classify_vllm_dispatch_profile( | ||
| 28 | + deploy_config.engine_config, | ||
| 29 | + explicit_profile=deploy_config.dispatch_profile, | ||
| 30 | + ) | ||
| 31 | + if profile != DispatchProfile.HANDOFF: | ||
| 32 | + raise ValueError( | ||
| 33 | + f"Native vLLM P/D launch only supports handoff connectors; resolved dispatch profile is {profile.value}" | ||
| 34 | + ) | ||
Rmotor/engine_server/core/vllm/vllm_config.py→motor/node_manager/core/services/native_engine/backends/vllm/config.py+9-8
| @@ -9,23 +9,22 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | import argparse | 11 | import argparse |
| 12 | -import sys | ||
| 13 | import json | 12 | import json |
| 14 | from typing import Any | 13 | from typing import Any |
| 15 | from dataclasses import dataclass, field | 14 | from dataclasses import dataclass, field |
| 16 | 15 | ||
| 17 | -from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args | ||
| 18 | - | ||
| 19 | from motor.config.endpoint import EndpointConfig | 16 | from motor.config.endpoint import EndpointConfig |
| 20 | -from motor.engine_server.core.config import IConfig | 17 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 21 | from motor.common.logger import get_logger | 18 | from motor.common.logger import get_logger |
| 22 | from motor.common.utils.net import format_address | 19 | from motor.common.utils.net import format_address |
| 23 | -from motor.engine_server.constants import constants | 20 | +from motor.common import engine_constants as constants |
| 24 | 21 | ||
| 25 | logger = get_logger(__name__) | 22 | logger = get_logger(__name__) |
| 26 | 23 | ||
| 27 | 24 | ||
| 28 | def _add_argument_to_list(arg_list: list, key: str, value: Any): | 25 | def _add_argument_to_list(arg_list: list, key: str, value: Any): |
| 26 | + if value is None: | ||
| 27 | + return | ||
| 29 | if isinstance(value, bool): | 28 | if isinstance(value, bool): |
| 30 | if value: | 29 | if value: |
| 31 | arg_list.append(f"--{key}") | 30 | arg_list.append(f"--{key}") |
| @@ -78,21 +77,23 @@ class VLLMConfig(IConfig): | |||
| 78 | 77 | ||
| 79 | def validate(self): | 78 | def validate(self): |
| 80 | if self.args is not None: | 79 | if self.args is not None: |
| 80 | + from vllm.entrypoints.openai.cli_args import validate_parsed_serve_args | ||
| 81 | + | ||
| 81 | validate_parsed_serve_args(self.args) | 82 | validate_parsed_serve_args(self.args) |
| 82 | 83 | ||
| 83 | def convert(self): | 84 | def convert(self): |
| 84 | arg_list = self._get_param_list() | 85 | arg_list = self._get_param_list() |
| 85 | logger.info(f'engine server parsed arg_list: {arg_list}') | 86 | logger.info(f'engine server parsed arg_list: {arg_list}') |
| 86 | 87 | ||
| 87 | - sys.argv = ["serve"] + arg_list | ||
| 88 | - | ||
| 89 | try: | 88 | try: |
| 90 | from vllm.utils import FlexibleArgumentParser | 89 | from vllm.utils import FlexibleArgumentParser |
| 91 | except ImportError: | 90 | except ImportError: |
| 92 | from vllm.utils.argparse_utils import FlexibleArgumentParser | 91 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 92 | + from vllm.entrypoints.openai.cli_args import make_arg_parser | ||
| 93 | + | ||
| 93 | parser = FlexibleArgumentParser(description="vLLM parser") | 94 | parser = FlexibleArgumentParser(description="vLLM parser") |
| 94 | parser = make_arg_parser(parser) | 95 | parser = make_arg_parser(parser) |
| 95 | - self.args = parser.parse_args() | 96 | + self.args = parser.parse_args(arg_list) |
| 96 | 97 | ||
| 97 | def get_args(self) -> argparse.Namespace: | 98 | def get_args(self) -> argparse.Namespace: |
| 98 | return self.args | 99 | return self.args |
Rmotor/engine_server/factory/config_factory.py→motor/node_manager/core/services/native_engine/config_factory.py+15-10
| @@ -11,22 +11,22 @@ | |||
| 11 | import importlib | 11 | import importlib |
| 12 | 12 | ||
| 13 | from motor.config.endpoint import EndpointConfig | 13 | from motor.config.endpoint import EndpointConfig |
| 14 | -from motor.engine_server.core.config import IConfig | 14 | +from motor.node_manager.core.services.native_engine.backends.base import IConfig |
| 15 | -from motor.common.logger import get_logger | ||
| 16 | - | ||
| 17 | -logger = get_logger(__name__) | ||
| 18 | 15 | ||
| 19 | 16 | ||
| 20 | class ConfigFactory: | 17 | class ConfigFactory: |
| 18 | + """Lazily construct engine-specific native CLI configuration adapters.""" | ||
| 19 | + | ||
| 21 | _ENGINE_CONFIG_MAP: dict[str, str] = { | 20 | _ENGINE_CONFIG_MAP: dict[str, str] = { |
| 22 | - "vllm": "motor.engine_server.core.vllm.vllm_config.VLLMConfig", | 21 | + "vllm": "motor.node_manager.core.services.native_engine.backends.vllm.config.VLLMConfig", |
| 23 | - "sglang": "motor.engine_server.core.sglang.sglang_config.SGLangConfig", | 22 | + "sglang": "motor.node_manager.core.services.native_engine.backends.sglang.config.SGLangConfig", |
| 24 | } | 23 | } |
| 25 | 24 | ||
| 26 | def __init__(self, endpoint_config: EndpointConfig): | 25 | def __init__(self, endpoint_config: EndpointConfig): |
| 27 | self.endpoint_config = endpoint_config | 26 | self.endpoint_config = endpoint_config |
| 28 | 27 | ||
| 29 | - def parse(self) -> IConfig: | 28 | + def build_cli_config(self) -> IConfig: |
| 29 | + """Build native CLI configuration without importing engine parsers.""" | ||
| 30 | engine_type = self.endpoint_config.engine_type | 30 | engine_type = self.endpoint_config.engine_type |
| 31 | config_class_path = self._ENGINE_CONFIG_MAP.get(engine_type) | 31 | config_class_path = self._ENGINE_CONFIG_MAP.get(engine_type) |
| 32 | 32 | ||
| @@ -38,14 +38,19 @@ class ConfigFactory: | |||
| 38 | ) | 38 | ) |
| 39 | 39 | ||
| 40 | try: | 40 | try: |
| 41 | - module_path, class_name = config_class_path.rsplit('.', 1) | 41 | + module_path, class_name = config_class_path.rsplit(".", 1) |
| 42 | module = importlib.import_module(module_path) | 42 | module = importlib.import_module(module_path) |
| 43 | config_class = getattr(module, class_name) | 43 | config_class = getattr(module, class_name) |
| 44 | 44 | ||
| 45 | config_instance = config_class(endpoint_config=self.endpoint_config) | 45 | config_instance = config_class(endpoint_config=self.endpoint_config) |
| 46 | config_instance.initialize() | 46 | config_instance.initialize() |
| 47 | - config_instance.convert() | ||
| 48 | - config_instance.validate() | ||
| 49 | return config_instance | 47 | return config_instance |
| 50 | except (ImportError, AttributeError) as e: | 48 | except (ImportError, AttributeError) as e: |
| 51 | raise ValueError(f"Failed to load config class for {engine_type}") from e | 49 | raise ValueError(f"Failed to load config class for {engine_type}") from e |
| 50 | + | ||
| 51 | + def parse(self) -> IConfig: | ||
| 52 | + """Build and validate the transitional in-process EngineServer configuration.""" | ||
| 53 | + config_instance = self.build_cli_config() | ||
| 54 | + config_instance.convert() | ||
| 55 | + config_instance.validate() | ||
| 56 | + return config_instance | ||
| @@ -8,40 +8,21 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -import argparse | 11 | +from motor.node_manager.core.services.native_engine.backends.base import NativeEngineBackend |
| 12 | -from abc import ABC, abstractmethod | 12 | +from motor.node_manager.core.services.native_engine.backends.sglang.backend import SGLangBackend |
| 13 | - | 13 | +from motor.node_manager.core.services.native_engine.backends.vllm.backend import VllmBackend |
| 14 | -from motor.common.logger import get_logger | ||
| 15 | -from motor.config.endpoint import EndpointConfig | ||
| 16 | - | ||
| 17 | -logger = get_logger(__name__) | ||
| 18 | - | ||
| 19 | -supported_engine = ["vllm", "sglang"] | ||
| 20 | -supported_role = ["prefill", "decode", "union"] | ||
| 21 | 14 | ||
| 22 | 15 | ||
| 23 | -class IConfig(ABC): | 16 | +_BACKENDS: dict[str, NativeEngineBackend] = { |
| 24 | - @abstractmethod | 17 | + VllmBackend.engine_type: VllmBackend(), |
| 25 | - def initialize(self): | 18 | + SGLangBackend.engine_type: SGLangBackend(), |
| 26 | - pass | 19 | +} |
| 27 | 20 | ||
| 28 | - | ||
| 29 | - def validate(self): | ||
| 30 | - pass | ||
| 31 | 21 | ||
| 32 | - @abstractmethod | 22 | +def get_backend(engine_type: str | None) -> NativeEngineBackend: |
| 33 | - def convert(self): | 23 | + """Return the stateless backend for the configured native engine type.""" |
| 34 | - pass | 24 | + normalized = str(engine_type or "").strip().lower() |
| 35 | - | 25 | + backend = _BACKENDS.get(normalized) |
| 36 | - @abstractmethod | 26 | + if backend is None: |
| 37 | - def get_args(self) -> argparse.Namespace | None: | 27 | + raise ValueError(f"Unsupported engine type: {engine_type}") |
| 38 | - pass | 28 | + return backend |
| 39 | - | ||
| 40 | - | ||
| 41 | - def get_endpoint_config(self) -> EndpointConfig | None: | ||
| 42 | - pass | ||
| 43 | - | ||
| 44 | - | ||
| 45 | - def get_cli_args(self) -> list[str]: | ||
| 46 | - """Return the CLI argument list suitable for native engine launch (vllm serve / sglang.launch_server).""" | ||
| 47 | - pass | ||
| @@ -0,0 +1,104 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from collections.abc import Mapping | ||
| 12 | +from dataclasses import dataclass | ||
| 13 | +from enum import Enum | ||
| 14 | +from types import MappingProxyType | ||
| 15 | + | ||
| 16 | +from motor.common.resources.instance import PDRole | ||
| 17 | +from motor.config.tls_config import TLSConfig | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +class LaunchContext: | ||
| 22 | + """Complete, immutable input required to build one native engine command.""" | ||
| 23 | + | ||
| 24 | + role: PDRole | ||
| 25 | + instance_id: int | ||
| 26 | + dp_rank: int | ||
| 27 | + node_rank: int | ||
| 28 | + host: str | ||
| 29 | + business_port: int | ||
| 30 | + mgmt_port: int | ||
| 31 | + config_path: str | ||
| 32 | + master_dp_ip: str | None | ||
| 33 | + kv_port: int | None | ||
| 34 | + lookup_rpc_port: int | None | ||
| 35 | + dp_rpc_port: int | None | ||
| 36 | + d2d_peer_ips: tuple[str, ...] | ||
| 37 | + environment: Mapping[str, str] | ||
| 38 | + headless: bool = False | ||
| 39 | + | ||
| 40 | + def __post_init__(self) -> None: | ||
| 41 | + object.__setattr__(self, "d2d_peer_ips", tuple(self.d2d_peer_ips)) | ||
| 42 | + object.__setattr__(self, "environment", MappingProxyType(dict(self.environment))) | ||
| 43 | + | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +class CommandSpec: | ||
| 47 | + """Native process invocation produced by a runtime adapter.""" | ||
| 48 | + | ||
| 49 | + argv: tuple[str, ...] | ||
| 50 | + env: Mapping[str, str] | ||
| 51 | + cwd: str | None = None | ||
| 52 | + | ||
| 53 | + def __post_init__(self) -> None: | ||
| 54 | + object.__setattr__(self, "argv", tuple(self.argv)) | ||
| 55 | + object.__setattr__(self, "env", MappingProxyType(dict(self.env))) | ||
| 56 | + | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +class ProbeSpec: | ||
| 60 | + """Native readiness probe associated with one engine process.""" | ||
| 61 | + | ||
| 62 | + path: str | ||
| 63 | + timeout_seconds: float | ||
| 64 | + startup_timeout_seconds: float | ||
| 65 | + max_attempts: int = 1 | ||
| 66 | + tls_config: TLSConfig | None = None | ||
| 67 | + process_only: bool = False | ||
| 68 | + | ||
| 69 | + def __post_init__(self) -> None: | ||
| 70 | + if not isinstance(self.max_attempts, int) or isinstance(self.max_attempts, bool) or self.max_attempts < 1: | ||
| 71 | + raise ValueError("max_attempts must be a positive integer") | ||
| 72 | + | ||
| 73 | + | ||
| 74 | + | ||
| 75 | +class LaunchSpec: | ||
| 76 | + """Command and readiness probe built from one validated engine config.""" | ||
| 77 | + | ||
| 78 | + command: CommandSpec | ||
| 79 | + probe: ProbeSpec | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +class RuntimeState(str, Enum): | ||
| 83 | + """Node-local lifecycle state; it is not part of the control-plane API.""" | ||
| 84 | + | ||
| 85 | + STARTING = "starting" | ||
| 86 | + RUNNING = "running" | ||
| 87 | + READY = "ready" | ||
| 88 | + UNHEALTHY = "unhealthy" | ||
| 89 | + STOPPING = "stopping" | ||
| 90 | + STOPPED = "stopped" | ||
| 91 | + | ||
| 92 | + | ||
| 93 | + | ||
| 94 | +class RuntimeProcess: | ||
| 95 | + """Mutable process record owned exclusively by ProcessSupervisor.""" | ||
| 96 | + | ||
| 97 | + endpoint_id: int | ||
| 98 | + process: object | ||
| 99 | + command: CommandSpec | ||
| 100 | + probe: ProbeSpec | ||
| 101 | + started_at: float | ||
| 102 | + process_group_id: int | None = None | ||
| 103 | + state: RuntimeState = RuntimeState.STARTING | ||
| 104 | + ready_at: float | None = None | ||
Rmotor/node_manager/core/services/engine.py→motor/node_manager/core/services/native_engine/service.py+111-155
| @@ -8,37 +8,32 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -import ipaddress | ||
| 12 | import os | 11 | import os |
| 13 | import signal | 12 | import signal |
| 14 | -import subprocess | ||
| 15 | import threading | 13 | import threading |
| 16 | 14 | ||
| 17 | -from motor.common.resources.instance import PDRole | ||
| 18 | from motor.common.resources.endpoint import Endpoint | 15 | from motor.common.resources.endpoint import Endpoint |
| 16 | +from motor.common.resources.instance import PDRole | ||
| 19 | from motor.common.utils.env import Env | 17 | from motor.common.utils.env import Env |
| 20 | from motor.common.logger import get_logger | 18 | from motor.common.logger import get_logger |
| 21 | -from motor.common.utils.snapshot_utils import MOTOR_SNAPSHOT_METADATA_PATH | 19 | +from motor.common.utils.net import format_address |
| 22 | -from motor.node_manager.core.services.registry import register_service, SERVICE_ENGINE | 20 | +from motor.node_manager.core.services.native_engine.factory import get_backend |
| 21 | +from motor.node_manager.core.services.native_engine.models import LaunchContext, RuntimeState | ||
| 22 | +from motor.node_manager.core.services.native_engine.supervisor import ProcessSupervisor | ||
| 23 | +from motor.node_manager.core.services.registry import SERVICE_ENGINE, register_service | ||
| 23 | 24 | ||
| 24 | logger = get_logger(__name__) | 25 | logger = get_logger(__name__) |
| 25 | -MAX_PORT = 65535 | ||
| 26 | -MIN_PORT = 1024 | ||
| 27 | 26 | ||
| 28 | 27 | ||
| 29 | -def _create_engine(hardware_type: str, config): | 28 | +def _create_native_engine(hardware_type: str, config): |
| 30 | - """Factory for EngineService — keeps constructor details out of the daemon.""" | 29 | + """Factory for NativeEngineService — keeps constructor details out of the daemon.""" |
| 31 | - return EngineService( | 30 | + del hardware_type |
| 32 | - hardware_type=hardware_type, | 31 | + return NativeEngineService( |
| 32 | + engine_type=config.basic_config.engine_type, | ||
| 33 | + config_path=config.config_path, | ||
| 33 | device_num=config.basic_config.device_num, | 34 | device_num=config.basic_config.device_num, |
| 34 | parallel_config=config.basic_config.parallel_config, | 35 | parallel_config=config.basic_config.parallel_config, |
| 35 | enable_multi_endpoints=config.basic_config.enable_multi_endpoints, | 36 | enable_multi_endpoints=config.basic_config.enable_multi_endpoints, |
| 36 | - enable_snapshot=config.snapshot_config.enable_snapshot, | ||
| 37 | - snapshot_metadata_path=( | ||
| 38 | - config.snapshot_config.snapshot_metadata_path | ||
| 39 | - if config.snapshot_config.snapshot_metadata_path != "" | ||
| 40 | - else MOTOR_SNAPSHOT_METADATA_PATH | ||
| 41 | - ), | ||
| 42 | single_container_flag=config.single_container_config.single_container_flag, | 37 | single_container_flag=config.single_container_config.single_container_flag, |
| 43 | device_offset=config.single_container_config.device_offset, | 38 | device_offset=config.single_container_config.device_offset, |
| 44 | kv_port=config.single_container_config.kv_port, | 39 | kv_port=config.single_container_config.kv_port, |
| @@ -47,60 +42,39 @@ def _create_engine(hardware_type: str, config): | |||
| 47 | ) | 42 | ) |
| 48 | 43 | ||
| 49 | 44 | ||
| 50 | -@register_service(SERVICE_ENGINE, backend="engine", factory=_create_engine) | 45 | +@register_service(SERVICE_ENGINE, backend="engine", factory=_create_native_engine) |
| 51 | -class EngineService: | 46 | +class NativeEngineService: |
| 52 | """Manage engine subprocess lifecycle: start, track PIDs, stop.""" | 47 | """Manage engine subprocess lifecycle: start, track PIDs, stop.""" |
| 53 | 48 | ||
| 54 | def __init__( | 49 | def __init__( |
| 55 | self, | 50 | self, |
| 56 | - hardware_type: str, | 51 | + engine_type: str, |
| 52 | + config_path: str, | ||
| 57 | device_num: int, | 53 | device_num: int, |
| 58 | parallel_config, | 54 | parallel_config, |
| 59 | enable_multi_endpoints: bool, | 55 | enable_multi_endpoints: bool, |
| 60 | - enable_snapshot: bool, | ||
| 61 | - snapshot_metadata_path: str, | ||
| 62 | single_container_flag: bool = False, | 56 | single_container_flag: bool = False, |
| 63 | device_offset: int = 0, | 57 | device_offset: int = 0, |
| 64 | kv_port: int | None = None, | 58 | kv_port: int | None = None, |
| 65 | lookup_rpc_port: int | None = None, | 59 | lookup_rpc_port: int | None = None, |
| 66 | dp_rpc_port: int | None = None, | 60 | dp_rpc_port: int | None = None, |
| 67 | ): | 61 | ): |
| 68 | - self.hardware_type = hardware_type | 62 | + self.engine_type = str(engine_type).strip().lower() |
| 63 | + self.config_path = config_path | ||
| 69 | self.device_num = device_num | 64 | self.device_num = device_num |
| 70 | self.parallel_config = parallel_config | 65 | self.parallel_config = parallel_config |
| 71 | self.enable_multi_endpoints = enable_multi_endpoints | 66 | self.enable_multi_endpoints = enable_multi_endpoints |
| 72 | - self.enable_snapshot = enable_snapshot | ||
| 73 | - self.snapshot_metadata_path = snapshot_metadata_path | ||
| 74 | self.single_container_flag = single_container_flag | 67 | self.single_container_flag = single_container_flag |
| 75 | self.device_offset = device_offset | 68 | self.device_offset = device_offset |
| 76 | self.kv_port = kv_port | 69 | self.kv_port = kv_port |
| 77 | self.lookup_rpc_port = lookup_rpc_port | 70 | self.lookup_rpc_port = lookup_rpc_port |
| 78 | self.dp_rpc_port = dp_rpc_port | 71 | self.dp_rpc_port = dp_rpc_port |
| 72 | + self.backend = get_backend(self.engine_type) | ||
| 73 | + self.supervisor = ProcessSupervisor() | ||
| 74 | + self._pull_lock = threading.Lock() | ||
| 79 | 75 | ||
| 80 | self.restart_on_failure = Env.motor_restart_engine | 76 | self.restart_on_failure = Env.motor_restart_engine |
| 81 | - | 77 | + self._recovery_requested = False |
| 82 | - self.engine_pids: list[int] = [] | ||
| 83 | - self._pids_lock = threading.Lock() | ||
| 84 | - | ||
| 85 | - | ||
| 86 | - def _check_params(params: Endpoint) -> bool: | ||
| 87 | - try: | ||
| 88 | - port = int(params.business_port) | ||
| 89 | - if not (MIN_PORT <= port <= MAX_PORT): | ||
| 90 | - logger.error("Port %s is out of valid range", port) | ||
| 91 | - return False | ||
| 92 | - except ValueError: | ||
| 93 | - logger.error("Invalid port value: %s", params.business_port) | ||
| 94 | - return False | ||
| 95 | - try: | ||
| 96 | - ipaddress.ip_address(params.ip) | ||
| 97 | - except ValueError: | ||
| 98 | - logger.error("Invalid IP address: %s", params.ip) | ||
| 99 | - return False | ||
| 100 | - except Exception as e: | ||
| 101 | - logger.error("Error validating IP address %s: %s", params.ip, e) | ||
| 102 | - return False | ||
| 103 | - return True | ||
| 104 | 78 | ||
| 105 | def pull( | 79 | def pull( |
| 106 | self, | 80 | self, |
| @@ -111,133 +85,103 @@ class EngineService: | |||
| 111 | d2d_peer_ips: list[str] | None = None, | 85 | d2d_peer_ips: list[str] | None = None, |
| 112 | node_rank: int = 0, | 86 | node_rank: int = 0, |
| 113 | ): | 87 | ): |
| 114 | - """Launch engine_server subprocesses for every endpoint on this node.""" | 88 | + """Launch native engine subprocesses for every endpoint on this node.""" |
| 89 | + with self._pull_lock: | ||
| 90 | + self._pull(pd_role_info, endpoints_info, instance_id, master_dp_ip, d2d_peer_ips, node_rank) | ||
| 91 | + | ||
| 92 | + def _pull( | ||
| 93 | + self, | ||
| 94 | + pd_role_info: PDRole, | ||
| 95 | + endpoints_info: list[Endpoint], | ||
| 96 | + instance_id: int, | ||
| 97 | + master_dp_ip: str, | ||
| 98 | + d2d_peer_ips: list[str] | None, | ||
| 99 | + node_rank: int, | ||
| 100 | + ) -> None: | ||
| 101 | + started_endpoint_ids: list[int] = [] | ||
| 115 | try: | 102 | try: |
| 116 | - env = os.environ.copy() | 103 | + base_env = os.environ.copy() |
| 117 | - pod_ip = env.get("POD_IP") | 104 | + pod_ip = base_env.get("POD_IP") |
| 118 | - if pod_ip and not env.get("VLLM_HOST_IP"): | 105 | + if pod_ip and not base_env.get("VLLM_HOST_IP"): |
| 119 | - env["VLLM_HOST_IP"] = pod_ip | 106 | + base_env["VLLM_HOST_IP"] = pod_ip |
| 120 | - if env.get("MOONCAKE_ASCEND_IPV6_EXPERIMENT") == "1": | 107 | + if base_env.get("MOONCAKE_ASCEND_IPV6_EXPERIMENT") == "1": |
| 121 | - env["MC_USE_IPV6"] = env.get("MC_USE_IPV6", "1") | 108 | + base_env["MC_USE_IPV6"] = base_env.get("MC_USE_IPV6", "1") |
| 122 | device_size = self.device_num | 109 | device_size = self.device_num |
| 123 | for i, endpoint in enumerate(endpoints_info): | 110 | for i, endpoint in enumerate(endpoints_info): |
| 124 | - if not self._check_params(endpoint): | 111 | + env = base_env.copy() |
| 125 | - raise ValueError("Invalid endpoint parameters") | ||
| 126 | - | ||
| 127 | if self.enable_multi_endpoints: | 112 | if self.enable_multi_endpoints: |
| 128 | device_ids_str = self._calc_visible_device_ids(i, device_size) | 113 | device_ids_str = self._calc_visible_device_ids(i, device_size) |
| 129 | logger.info("Device IDs: %s", device_ids_str) | 114 | logger.info("Device IDs: %s", device_ids_str) |
| 130 | env["ASCEND_RT_VISIBLE_DEVICES"] = device_ids_str | 115 | env["ASCEND_RT_VISIBLE_DEVICES"] = device_ids_str |
| 131 | 116 | ||
| 132 | - cmd = [ | 117 | + peer_ips = self._get_d2d_peer_ips(endpoint.id, d2d_peer_ips) |
| 133 | - "engine_server", | ||
| 134 | - "--dp-rank", | ||
| 135 | - str(endpoint.id), | ||
| 136 | - "--instance-id", | ||
| 137 | - str(instance_id), | ||
| 138 | - "--role", | ||
| 139 | - str(pd_role_info.value), | ||
| 140 | - "--host", | ||
| 141 | - str(endpoint.ip), | ||
| 142 | - "--port", | ||
| 143 | - str(int(endpoint.business_port)), | ||
| 144 | - "--mgmt-port", | ||
| 145 | - str(int(endpoint.mgmt_port)), | ||
| 146 | - "--master-dp-ip", | ||
| 147 | - master_dp_ip, | ||
| 148 | - "--node-rank", | ||
| 149 | - str(node_rank), | ||
| 150 | - "--config-path", | ||
| 151 | - str(Env.user_config_path), | ||
| 152 | - ] | ||
| 153 | - if self.enable_snapshot: | ||
| 154 | - cmd.extend(["--snapshot-metadata", self.snapshot_metadata_path]) | ||
| 155 | - if self.single_container_flag: | ||
| 156 | - if self.kv_port is not None: | ||
| 157 | - cmd.extend(["--kv-port", str(self.kv_port)]) | ||
| 158 | - if self.dp_rpc_port is not None: | ||
| 159 | - cmd.extend(["--dp-rpc-port", str(self.dp_rpc_port)]) | ||
| 160 | - if self.lookup_rpc_port is not None: | ||
| 161 | - cmd.extend(["--lookup-rpc-port", str(self.lookup_rpc_port)]) | ||
| 162 | if d2d_peer_ips: | 118 | if d2d_peer_ips: |
| 163 | - ep_id = str(endpoint.id) | 119 | + logger.info("D2D peer IPs for ep_id %s: %s", endpoint.id, list(peer_ips)) |
| 164 | - peer_ips = [] | 120 | + |
| 165 | - for entry in d2d_peer_ips: | 121 | + context = LaunchContext( |
| 166 | - encoded_ep_id, ip = entry.split(":", 1) | 122 | + role=pd_role_info, |
| 167 | - if encoded_ep_id == ep_id: | 123 | + instance_id=instance_id, |
| 168 | - peer_ips.append(ip) | 124 | + dp_rank=endpoint.id, |
| 169 | - if peer_ips: | 125 | + node_rank=node_rank, |
| 170 | - cmd.extend(["--d2d-peer-ips", ",".join(peer_ips)]) | 126 | + host=endpoint.ip, |
| 171 | - logger.info("D2D peer IPs for ep_id %s: %s", endpoint.id, peer_ips) | 127 | + business_port=int(endpoint.business_port), |
| 128 | + mgmt_port=int(endpoint.mgmt_port), | ||
| 129 | + config_path=self.config_path, | ||
| 130 | + master_dp_ip=master_dp_ip, | ||
| 131 | + kv_port=self.kv_port if self.single_container_flag else None, | ||
| 132 | + lookup_rpc_port=self.lookup_rpc_port if self.single_container_flag else None, | ||
| 133 | + dp_rpc_port=self.dp_rpc_port if self.single_container_flag else None, | ||
| 134 | + d2d_peer_ips=peer_ips, | ||
| 135 | + environment=env, | ||
| 136 | + headless=endpoint.headless, | ||
| 137 | + ) | ||
| 138 | + launch_spec = self.backend.prepare(context) | ||
| 139 | + cmd = list(launch_spec.command.argv) | ||
| 172 | logger.info(" ".join(cmd)) | 140 | logger.info(" ".join(cmd)) |
| 173 | - process = subprocess.Popen(cmd, shell=False, env=env) # pylint: disable=consider-using-with | 141 | + if self.supervisor.start(endpoint.id, launch_spec.command, launch_spec.probe): |
| 174 | - if process.poll() is not None: | 142 | + started_endpoint_ids.append(endpoint.id) |
| 175 | - raise RuntimeError("Engine process exited immediately with code %s" % process.returncode) | 143 | + |
| 176 | - with self._pids_lock: | 144 | + self._recovery_requested = False |
| 177 | - self.engine_pids.append(process.pid) | ||
| 178 | 145 | ||
| 179 | except Exception as e: | 146 | except Exception as e: |
| 180 | - self.stop() | 147 | + for endpoint_id in reversed(started_endpoint_ids): |
| 148 | + self.supervisor.stop(endpoint_id) | ||
| 181 | raise RuntimeError("Failed to pull engine: %s" % e) from e | 149 | raise RuntimeError("Failed to pull engine: %s" % e) from e |
| 182 | 150 | ||
| 183 | def stop(self) -> list[int]: | 151 | def stop(self) -> list[int]: |
| 184 | - """Kill all engine processes and return the list of PIDs that were killed.""" | 152 | + """Gracefully stop all native process groups, then force-kill on timeout.""" |
| 185 | - with self._pids_lock: | 153 | + return self.supervisor.stop_all() |
| 186 | - pids = list(self.engine_pids) | ||
| 187 | - self.engine_pids.clear() | ||
| 188 | - for pid in pids: | ||
| 189 | - try: | ||
| 190 | - os.kill(pid, signal.SIGKILL) | ||
| 191 | - logger.info("Killed engine process with PID: %s", pid) | ||
| 192 | - except ProcessLookupError: | ||
| 193 | - logger.info("Process %s already terminated", pid) | ||
| 194 | - except PermissionError: | ||
| 195 | - logger.error("No permission to kill process %s", pid) | ||
| 196 | - except Exception as e: | ||
| 197 | - logger.error("Failed to kill process %s: %s", pid, e) | ||
| 198 | - return pids | ||
| 199 | 154 | ||
| 200 | def pid_list(self) -> list[int]: | 155 | def pid_list(self) -> list[int]: |
| 201 | - with self._pids_lock: | 156 | + return self.supervisor.pid_list() |
| 202 | - return list(self.engine_pids) | ||
| 203 | 157 | ||
| 204 | - def remove_pid(self, pid: int) -> None: | 158 | + def runtime_state(self, endpoint: Endpoint) -> RuntimeState: |
| 205 | - with self._pids_lock: | 159 | + return self.supervisor.state(endpoint.id, endpoint.ip, int(endpoint.business_port)) |
| 206 | - if pid in self.engine_pids: | 160 | + |
| 207 | - self.engine_pids.remove(pid) | 161 | + def metrics_target(self, endpoint: Endpoint) -> str | None: |
| 162 | + """Return the native metrics URL for a routable local endpoint.""" | ||
| 163 | + if endpoint.headless: | ||
| 164 | + return None | ||
| 165 | + probe = self.supervisor.probe_spec(endpoint.id) | ||
| 166 | + if probe is None: | ||
| 167 | + return None | ||
| 168 | + scheme = "https" if probe.tls_config and probe.tls_config.enable_tls else "http" | ||
| 169 | + return f"{scheme}://{format_address(endpoint.ip, endpoint.business_port)}/metrics" | ||
| 208 | 170 | ||
| 209 | def health_check(self) -> None: | 171 | def health_check(self) -> None: |
| 210 | - """Check engine PIDs; trigger pod restart on failure (DaemonService protocol). | 172 | + """Trigger Pod-level recovery when any native engine process exits.""" |
| 211 | - | 173 | + dead_pids = self.supervisor.dead_pids() |
| 212 | - Two-layer suicide strategy: | 174 | + if not dead_pids: |
| 213 | - 1. Fast path (here): individual engine subprocess death → immediate SIGTERM | 175 | + return |
| 214 | - to own process, controlled by ``Env.motor_restart_engine``. This gives | 176 | + logger.warning( |
| 215 | - ~5 s detection latency via the Daemon process monitor loop. | 177 | + "Engine PIDs %s died (restart_on_failure=%s)", |
| 216 | - | 178 | + dead_pids, |
| 217 | - 2. Slow path (HeartbeatManager): after 5 consecutive abnormal heartbeats, | 179 | + self.restart_on_failure, |
| 218 | - HeartbeatManager sets its ``should_suicide`` flag; the main loop calls | 180 | + ) |
| 219 | - ``suicide_procedure()`` which returns exit code -1 (pod rescheduling). | 181 | + if self.restart_on_failure and not self._recovery_requested: |
G 严重程度： 建议 问题： 原因： 怎么改：
在 pull()/_pull() 成功入口将 tobking 19 天前 评论： 19 天前 评论： | |||
| 220 | - This catches controller-level node unhealthiness that may not involve | 182 | + self._recovery_requested = True |
| 221 | - an individual engine PID dying. | 183 | + logger.info("Engine failure requires Pod-level recovery") |
| 222 | - | 184 | + os.kill(os.getpid(), signal.SIGTERM) |
| 223 | - These two paths are independent: the fast path handles local engine crashes | ||
| 224 | - immediately; the slow path handles orchestration-level failure detection. | ||
| 225 | - """ | ||
| 226 | - for pid in self.pid_list(): | ||
| 227 | - try: | ||
| 228 | - os.kill(pid, 0) | ||
| 229 | - except ProcessLookupError: | ||
| 230 | - logger.warning( | ||
| 231 | - "Engine PID %s died (restart_on_failure=%s)", | ||
| 232 | - pid, | ||
| 233 | - self.restart_on_failure, | ||
| 234 | - ) | ||
| 235 | - self.remove_pid(pid) | ||
| 236 | - if self.restart_on_failure: | ||
| 237 | - logger.info("Engine restart requested — triggering suicide for k8s pod restart") | ||
| 238 | - os.kill(os.getpid(), signal.SIGTERM) | ||
| 239 | - except PermissionError: | ||
| 240 | - pass | ||
| 241 | 185 | ||
| 242 | def _calc_visible_device_ids(self, index: int, device_size: int) -> str: | 186 | def _calc_visible_device_ids(self, index: int, device_size: int) -> str: |
| 243 | local_world_size = self.parallel_config.local_world_size | 187 | local_world_size = self.parallel_config.local_world_size |
| @@ -250,3 +194,15 @@ class EngineService: | |||
| 250 | if self.single_container_flag: | 194 | if self.single_container_flag: |
| 251 | device_ids = [x + self.device_offset for x in device_ids] | 195 | device_ids = [x + self.device_offset for x in device_ids] |
| 252 | return ",".join(map(str, device_ids)) | 196 | return ",".join(map(str, device_ids)) |
| 197 | + | ||
| 198 | + | ||
| 199 | + def _get_d2d_peer_ips(endpoint_id: int, d2d_peer_ips: list[str] | None) -> tuple[str, ...]: | ||
| 200 | + if not d2d_peer_ips: | ||
| 201 | + return () | ||
| 202 | + endpoint_id_text = str(endpoint_id) | ||
| 203 | + peers = [] | ||
| 204 | + for entry in d2d_peer_ips: | ||
| 205 | + encoded_endpoint_id, ip = entry.split(":", 1) | ||
| 206 | + if encoded_endpoint_id == endpoint_id_text: | ||
| 207 | + peers.append(ip) | ||
| 208 | + return tuple(peers) | ||
| @@ -0,0 +1,261 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +import os | ||
| 12 | +import signal | ||
| 13 | +import subprocess | ||
| 14 | +import threading | ||
| 15 | +import time | ||
| 16 | + | ||
| 17 | +import requests | ||
| 18 | + | ||
| 19 | +from motor.common.http.http_client import SafeHTTPSClient | ||
| 20 | +from motor.common.logger import get_logger | ||
| 21 | +from motor.common.utils.net import format_address | ||
| 22 | +from motor.node_manager.core.services.native_engine.models import CommandSpec, ProbeSpec, RuntimeProcess, RuntimeState | ||
| 23 | + | ||
| 24 | +logger = get_logger(__name__) | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +class ProcessSupervisor: | ||
| 28 | + """Own native engine process groups and their node-local runtime state.""" | ||
| 29 | + | ||
| 30 | + def __init__(self, stop_grace_seconds: float = 10.0) -> None: | ||
| 31 | + self._processes: dict[int, RuntimeProcess] = {} | ||
| 32 | + self._lock = threading.Lock() | ||
| 33 | + self._stop_grace_seconds = stop_grace_seconds | ||
| 34 | + | ||
| 35 | + def start(self, endpoint_id: int, command: CommandSpec, probe: ProbeSpec) -> bool: | ||
| 36 | + """Start one endpoint atomically; return False when it is already running.""" | ||
| 37 | + with self._lock: | ||
| 38 | + existing = self._processes.get(endpoint_id) | ||
| 39 | + if existing is not None and existing.process.poll() is None: | ||
| 40 | + if existing.command == command and existing.probe == probe: | ||
| 41 | + return False | ||
| 42 | + raise RuntimeError(f"Engine endpoint {endpoint_id} is already running with a different launch spec") | ||
| 43 | + self._processes.pop(endpoint_id, None) | ||
| 44 | + | ||
| 45 | + process = subprocess.Popen( # pylint: disable=consider-using-with | ||
| 46 | + list(command.argv), | ||
| 47 | + shell=False, | ||
| 48 | + env=dict(command.env), | ||
| 49 | + cwd=command.cwd, | ||
| 50 | + start_new_session=True, | ||
| 51 | + ) | ||
| 52 | + if process.poll() is not None: | ||
| 53 | + raise RuntimeError(f"Engine process exited immediately with code {process.returncode}") | ||
| 54 | + | ||
| 55 | + self._processes[endpoint_id] = RuntimeProcess( | ||
| 56 | + endpoint_id=endpoint_id, | ||
| 57 | + process=process, | ||
| 58 | + command=command, | ||
| 59 | + probe=probe, | ||
| 60 | + started_at=time.monotonic(), | ||
| 61 | + process_group_id=process.pid if os.name == "posix" else None, | ||
| 62 | + ) | ||
| 63 | + return True | ||
| 64 | + | ||
| 65 | + def pid_list(self) -> list[int]: | ||
| 66 | + with self._lock: | ||
| 67 | + return [runtime.process.pid for runtime in self._processes.values()] | ||
| 68 | + | ||
| 69 | + def probe_spec(self, endpoint_id: int) -> ProbeSpec | None: | ||
| 70 | + with self._lock: | ||
| 71 | + runtime = self._processes.get(endpoint_id) | ||
| 72 | + return runtime.probe if runtime is not None else None | ||
| 73 | + | ||
| 74 | + def state(self, endpoint_id: int, host: str, port: int) -> RuntimeState: | ||
G 严重程度： 建议 问题： 原因： 怎么改： 把 state/ready_at 的读写收敛到锁内，或给 RuntimeProcess 加专用小锁： state() 的 HTTP 探测可在锁外进行，但回写统一走 set_state。 tobking 19 天前 评论： 19 天前 评论： | |||
| 75 | + with self._lock: | ||
| 76 | + runtime = self._processes.get(endpoint_id) | ||
| 77 | + if runtime is None: | ||
| 78 | + return RuntimeState.STOPPED | ||
| 79 | + if runtime.state == RuntimeState.STOPPING: | ||
| 80 | + return RuntimeState.STOPPING | ||
| 81 | + | ||
| 82 | + process = runtime.process | ||
| 83 | + if process.poll() is not None: | ||
| 84 | + return self._commit_state(endpoint_id, runtime, RuntimeState.STOPPED) | ||
| 85 | + if runtime.probe.process_only: | ||
G 严重程度： 严重 问题： headless 端点（跨节点 PCP slave）失去“引擎加载完成”就绪门控，进程一 spawn 即上报 READY→NORMAL。 原因： 怎么改：
与作者确认 vllm tobking 19 天前 评论： 19 天前 评论： | |||
| 86 | + return self._commit_state(endpoint_id, runtime, RuntimeState.RUNNING) | ||
| 87 | + try: | ||
| 88 | + address = format_address(host, port) | ||
| 89 | + with SafeHTTPSClient( | ||
| 90 | + address=address, | ||
| 91 | + tls_config=runtime.probe.tls_config, | ||
| 92 | + timeout=runtime.probe.timeout_seconds, | ||
| 93 | + ) as client: | ||
| 94 | + for attempt in range(1, runtime.probe.max_attempts + 1): | ||
| 95 | + try: | ||
| 96 | + client.do_get(runtime.probe.path) | ||
| 97 | + break | ||
| 98 | + except Exception as err: | ||
| 99 | + if not self._is_timeout_error(err) or attempt == runtime.probe.max_attempts: | ||
| 100 | + raise | ||
| 101 | + logger.debug( | ||
| 102 | + "Native health probe timed out for endpoint %s (attempt %s/%s), retrying", | ||
| 103 | + endpoint_id, | ||
| 104 | + attempt, | ||
| 105 | + runtime.probe.max_attempts, | ||
| 106 | + ) | ||
| 107 | + return self._commit_state(endpoint_id, runtime, RuntimeState.READY, ready=True) | ||
| 108 | + except Exception as err: | ||
| 109 | + elapsed = time.monotonic() - runtime.started_at | ||
| 110 | + if runtime.state == RuntimeState.STARTING and elapsed < runtime.probe.startup_timeout_seconds: | ||
| 111 | + logger.debug( | ||
| 112 | + "Engine endpoint %s is still starting after %.1fs: %s", | ||
| 113 | + endpoint_id, | ||
| 114 | + elapsed, | ||
| 115 | + err, | ||
| 116 | + ) | ||
| 117 | + return self._current_state(endpoint_id, runtime) | ||
| 118 | + logger.warning("Native health probe failed for endpoint %s: %s", endpoint_id, err) | ||
| 119 | + return self._commit_state(endpoint_id, runtime, RuntimeState.UNHEALTHY) | ||
| 120 | + | ||
| 121 | + def dead_pids(self) -> list[int]: | ||
G 严重程度： 建议 问题： 死进程记录永不从 原因： 怎么改： 在 dead_pids() 中把已确认死亡的记录移出字典： 注意保持 tobking 19 天前 评论： 19 天前 评论： | |||
| 122 | + with self._lock: | ||
| 123 | + dead = [] | ||
| 124 | + dead_runtimes = [] | ||
| 125 | + for endpoint_id, runtime in list(self._processes.items()): | ||
| 126 | + if runtime.state == RuntimeState.STOPPING or runtime.process.poll() is None: | ||
| 127 | + continue | ||
| 128 | + runtime.state = RuntimeState.STOPPED | ||
| 129 | + dead.append(runtime.process.pid) | ||
| 130 | + dead_runtimes.append(runtime) | ||
| 131 | + self._processes.pop(endpoint_id, None) | ||
| 132 | + | ||
| 133 | + # The launcher may exit before its workers. Clean the cached process group | ||
| 134 | + # after removing the record so the same death is reported only once. | ||
| 135 | + for runtime in dead_runtimes: | ||
| 136 | + self._kill_group(runtime) | ||
| 137 | + return dead | ||
| 138 | + | ||
| 139 | + def stop_all(self) -> list[int]: | ||
| 140 | + with self._lock: | ||
| 141 | + runtimes = list(self._processes.values()) | ||
| 142 | + for runtime in runtimes: | ||
| 143 | + runtime.state = RuntimeState.STOPPING | ||
| 144 | + | ||
| 145 | + stopped = self._stop_runtimes(runtimes) | ||
| 146 | + self._remove_stopped(runtimes) | ||
| 147 | + return stopped | ||
| 148 | + | ||
| 149 | + def stop(self, endpoint_id: int) -> int | None: | ||
| 150 | + """Stop one endpoint without affecting unrelated engine processes.""" | ||
| 151 | + with self._lock: | ||
| 152 | + runtime = self._processes.get(endpoint_id) | ||
| 153 | + if runtime is not None: | ||
| 154 | + runtime.state = RuntimeState.STOPPING | ||
| 155 | + if runtime is None: | ||
| 156 | + return None | ||
| 157 | + stopped = self._stop_runtimes([runtime])[0] | ||
| 158 | + self._remove_stopped([runtime]) | ||
| 159 | + return stopped | ||
| 160 | + | ||
| 161 | + def _stop_runtimes(self, runtimes: list[RuntimeProcess]) -> list[int]: | ||
G 严重程度： 建议 问题： 进程组清理依赖 leader 存活：leader 先退而 worker 存活时，既收不到 SIGTERM 也收不到 SIGKILL，引擎 worker 泄漏。 原因： 怎么改：
start() 时缓存 pgid（Popen 成功后立即 | |||
| 162 | + for runtime in runtimes: | ||
| 163 | + runtime.state = RuntimeState.STOPPING | ||
| 164 | + self._terminate_group(runtime) | ||
| 165 | + | ||
| 166 | + deadline = time.monotonic() + self._stop_grace_seconds | ||
| 167 | + for runtime in runtimes: | ||
| 168 | + remaining = max(0.0, deadline - time.monotonic()) | ||
| 169 | + try: | ||
| 170 | + runtime.process.wait(timeout=remaining) | ||
| 171 | + except subprocess.TimeoutExpired: | ||
| 172 | + self._kill_group(runtime) | ||
| 173 | + else: | ||
| 174 | + # The launcher can exit while child workers still hold the same PGID. | ||
| 175 | + # Reap the remaining group instead of treating a leader exit as cleanup. | ||
| 176 | + if self._group_exists(runtime): | ||
| 177 | + self._kill_group(runtime) | ||
| 178 | + return [runtime.process.pid for runtime in runtimes] | ||
| 179 | + | ||
| 180 | + def _current_state(self, endpoint_id: int, runtime: RuntimeProcess) -> RuntimeState: | ||
| 181 | + with self._lock: | ||
| 182 | + current = self._processes.get(endpoint_id) | ||
| 183 | + return runtime.state if current is runtime else RuntimeState.STOPPED | ||
| 184 | + | ||
| 185 | + def _commit_state( | ||
| 186 | + self, | ||
| 187 | + endpoint_id: int, | ||
| 188 | + runtime: RuntimeProcess, | ||
| 189 | + state: RuntimeState, | ||
| 190 | + *, | ||
| 191 | + ready: bool = False, | ||
| 192 | + ) -> RuntimeState: | ||
| 193 | + with self._lock: | ||
| 194 | + current = self._processes.get(endpoint_id) | ||
| 195 | + if current is not runtime: | ||
| 196 | + return RuntimeState.STOPPED | ||
| 197 | + if runtime.state == RuntimeState.STOPPING: | ||
| 198 | + return RuntimeState.STOPPING | ||
| 199 | + runtime.state = state | ||
| 200 | + if ready: | ||
| 201 | + runtime.ready_at = runtime.ready_at or time.monotonic() | ||
| 202 | + return runtime.state | ||
| 203 | + | ||
| 204 | + def _remove_stopped(self, runtimes: list[RuntimeProcess]) -> None: | ||
| 205 | + with self._lock: | ||
| 206 | + for runtime in runtimes: | ||
| 207 | + if self._processes.get(runtime.endpoint_id) is runtime: | ||
| 208 | + runtime.state = RuntimeState.STOPPED | ||
| 209 | + self._processes.pop(runtime.endpoint_id, None) | ||
| 210 | + | ||
| 211 | + | ||
| 212 | + def _is_timeout_error(error: Exception) -> bool: | ||
| 213 | + """Return whether SafeHTTPSClient surfaced a requests timeout.""" | ||
| 214 | + current: BaseException | None = error | ||
| 215 | + while current is not None: | ||
| 216 | + if isinstance(current, requests.exceptions.Timeout): | ||
| 217 | + return True | ||
| 218 | + current = current.__cause__ | ||
| 219 | + return False | ||
| 220 | + | ||
| 221 | + | ||
| 222 | + def _terminate_group(runtime: RuntimeProcess) -> None: | ||
| 223 | + process = runtime.process | ||
| 224 | + try: | ||
| 225 | + if os.name == "posix" and runtime.process_group_id is not None: | ||
| 226 | + getattr(os, "killpg")(runtime.process_group_id, signal.SIGTERM) | ||
| 227 | + else: | ||
| 228 | + process.terminate() | ||
| 229 | + except ProcessLookupError: | ||
| 230 | + pass | ||
| 231 | + except (OSError, PermissionError) as err: | ||
| 232 | + logger.warning("Failed to terminate engine process group %s: %s", process.pid, err) | ||
| 233 | + | ||
| 234 | + | ||
| 235 | + def _kill_group(runtime: RuntimeProcess) -> None: | ||
| 236 | + process = runtime.process | ||
| 237 | + try: | ||
| 238 | + if os.name == "posix" and runtime.process_group_id is not None: | ||
| 239 | + getattr(os, "killpg")(runtime.process_group_id, getattr(signal, "SIGKILL")) | ||
| 240 | + else: | ||
| 241 | + process.kill() | ||
| 242 | + except ProcessLookupError: | ||
| 243 | + pass | ||
| 244 | + except (OSError, PermissionError) as err: | ||
| 245 | + logger.error("Failed to kill engine process group %s: %s", process.pid, err) | ||
| 246 | + | ||
| 247 | + | ||
| 248 | + def _group_exists(runtime: RuntimeProcess) -> bool: | ||
| 249 | + """Return whether a cached POSIX process group still has a member.""" | ||
| 250 | + if os.name != "posix" or runtime.process_group_id is None: | ||
| 251 | + return False | ||
| 252 | + try: | ||
| 253 | + getattr(os, "killpg")(runtime.process_group_id, 0) | ||
| 254 | + except ProcessLookupError: | ||
| 255 | + return False | ||
| 256 | + except PermissionError: | ||
| 257 | + return True | ||
| 258 | + except OSError as err: | ||
| 259 | + logger.warning("Failed to inspect engine process group %s: %s", runtime.process_group_id, err) | ||
| 260 | + return True | ||
| 261 | + return True | ||
| @@ -41,7 +41,7 @@ class PreparableService(DaemonService, Protocol): | |||
| 41 | """A service that needs pre-flight preparation before the engine starts.""" | 41 | """A service that needs pre-flight preparation before the engine starts.""" |
| 42 | 42 | ||
| 43 | def prepare(self, **kwargs) -> None: | 43 | def prepare(self, **kwargs) -> None: |
| 44 | - """Run before ``EngineService.pull()``. | 44 | + """Run before ``NativeEngineService.pull()``. |
| 45 | 45 | ||
| 46 | *kwargs* include ``endpoints_count`` (int) so the service can | 46 | *kwargs* include ``endpoints_count`` (int) so the service can |
| 47 | divide per-node DRAM across DP ranks. | 47 | divide per-node DRAM across DP ranks. |
| @@ -15,7 +15,7 @@ Usage:: | |||
| 15 | from motor.node_manager.core.services.registry import register_service | 15 | from motor.node_manager.core.services.registry import register_service |
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | - class EngineService: | 18 | + class NativeEngineService: |
| 19 | ... | 19 | ... |
| 20 | 20 | ||
| 21 | 21 | ||
| @@ -49,7 +49,7 @@ SERVICE_KV_STORE: str = "kv_store" | |||
| 49 | # that backend appears in the ``Env.kv_store_backend`` list (comma-separated). | 49 | # that backend appears in the ``Env.kv_store_backend`` list (comma-separated). |
| 50 | # ``None`` key = always-active (imported unconditionally). | 50 | # ``None`` key = always-active (imported unconditionally). |
| 51 | _DEFAULT_MODULE_MAP: dict[str | None, list[str]] = { | 51 | _DEFAULT_MODULE_MAP: dict[str | None, list[str]] = { |
| 52 | - "engine": ["motor.node_manager.core.services.engine"], | 52 | + "engine": ["motor.node_manager.core.services.native_engine.service"], |
| 53 | "memcache": ["motor.node_manager.core.services.memcache.lifecycle"], | 53 | "memcache": ["motor.node_manager.core.services.memcache.lifecycle"], |
| 54 | } | 54 | } |
| 55 | 55 | ||
| @@ -8,62 +8,15 @@ | |||
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | -"""Unit tests for dispatch helpers in motor.common.resources.dispatch.""" | 11 | +"""Unit tests for dispatch-profile inference.""" |
| 12 | 12 | ||
| 13 | -from dataclasses import dataclass, field | ||
| 14 | from types import SimpleNamespace | 13 | from types import SimpleNamespace |
| 15 | 14 | ||
| 16 | from motor.common.resources.dispatch import ( | 15 | from motor.common.resources.dispatch import ( |
| 17 | - DispatchPlan, | ||
| 18 | DispatchProfile, | 16 | DispatchProfile, |
| 19 | - dispatch_plan_union, | ||
| 20 | - has_compatible_dispatch_pair, | ||
| 21 | infer_vllm_dispatch_profile_from_config, | 17 | infer_vllm_dispatch_profile_from_config, |
| 22 | - shared_dispatch_plans, | ||
| 23 | ) | 18 | ) |
| 24 | 19 | ||
| 25 | -CONCURRENT = DispatchPlan.CONCURRENT_ENGINE_SYNC.value | ||
| 26 | -HANDOFF = DispatchPlan.PREFILL_HANDOFF_DECODE.value | ||
| 27 | - | ||
| 28 | - | ||
| 29 | - | ||
| 30 | -class _Inst: | ||
| 31 | - """Minimal stand-in carrying only the dispatch_capabilities attribute the helpers read.""" | ||
| 32 | - | ||
| 33 | - id: int = 0 | ||
| 34 | - dispatch_capabilities: list = field(default_factory=list) | ||
| 35 | - | ||
| 36 | - | ||
| 37 | -def _pairwise_compatible(prefill_instances, decode_instances): | ||
| 38 | - """Reference O(P*D) definition the optimized helpers must stay equivalent to.""" | ||
| 39 | - decode_list = list(decode_instances) | ||
| 40 | - return any(shared_dispatch_plans(p, d) for p in prefill_instances for d in decode_list) | ||
| 41 | - | ||
| 42 | - | ||
| 43 | -def test_dispatch_plan_union_aggregates_and_ignores_unknown_values(): | ||
| 44 | - instances = [_Inst(dispatch_capabilities=[CONCURRENT]), _Inst(dispatch_capabilities=[HANDOFF, "bogus"])] | ||
| 45 | - assert dispatch_plan_union(instances) == { | ||
| 46 | - DispatchPlan.CONCURRENT_ENGINE_SYNC, | ||
| 47 | - DispatchPlan.PREFILL_HANDOFF_DECODE, | ||
| 48 | - } | ||
| 49 | - assert dispatch_plan_union([]) == set() | ||
| 50 | - assert dispatch_plan_union([_Inst(dispatch_capabilities=[])]) == set() | ||
| 51 | - | ||
| 52 | - | ||
| 53 | -def test_has_compatible_dispatch_pair_matches_pairwise_definition(): | ||
| 54 | - cases = [ | ||
| 55 | - ([_Inst(dispatch_capabilities=[CONCURRENT])], [_Inst(dispatch_capabilities=[CONCURRENT])]), | ||
| 56 | - ([_Inst(dispatch_capabilities=[CONCURRENT])], [_Inst(dispatch_capabilities=[HANDOFF])]), | ||
| 57 | - ( | ||
| 58 | - [_Inst(dispatch_capabilities=[CONCURRENT]), _Inst(dispatch_capabilities=[HANDOFF])], | ||
| 59 | - [_Inst(dispatch_capabilities=[HANDOFF])], | ||
| 60 | - ), | ||
| 61 | - ([_Inst(dispatch_capabilities=[])], [_Inst(dispatch_capabilities=[CONCURRENT])]), | ||
| 62 | - ([], [_Inst(dispatch_capabilities=[CONCURRENT])]), | ||
| 63 | - ] | ||
| 64 | - for prefill, decode in cases: | ||
| 65 | - assert has_compatible_dispatch_pair(prefill, decode) == _pairwise_compatible(prefill, decode) | ||
| 66 | - | ||
| 67 | 20 | ||
| 68 | class _EngineConfig: | 21 | class _EngineConfig: |
| 69 | def __init__(self, configs): | 22 | def __init__(self, configs): |
| @@ -1,5 +1,3 @@ | |||
| 1 | -#!/usr/bin/env python3 | ||
| 2 | -# -*- coding: utf-8 -*- | ||
| 3 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 4 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 5 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| @@ -321,6 +319,68 @@ def test_endpoint_headless_defaults_to_false() -> None: | |||
| 321 | assert endpoint.headless is False | 319 | assert endpoint.headless is False |
| 322 | 320 | ||
| 323 | 321 | ||
| 322 | +def test_instance_readiness_requires_live_headless_workers_and_routable_master() -> None: | ||
| 323 | + instance = Instance( | ||
| 324 | + job_name="test_headless_readiness", | ||
| 325 | + model_name="test_model", | ||
| 326 | + id=1, | ||
| 327 | + role="prefill", | ||
| 328 | + parallel_config=ParallelConfig(dp_size=1, tp_size=4), | ||
| 329 | + enable_multi_endpoints=True, | ||
| 330 | + ) | ||
| 331 | + master = Endpoint( | ||
| 332 | + id=0, | ||
| 333 | + ip="10.0.0.1", | ||
| 334 | + business_port="8000", | ||
| 335 | + mgmt_port="9000", | ||
| 336 | + status=EndpointStatus.NORMAL, | ||
| 337 | + ) | ||
| 338 | + worker = Endpoint( | ||
| 339 | + id=1, | ||
| 340 | + ip="10.0.0.2", | ||
| 341 | + business_port="8000", | ||
| 342 | + mgmt_port="9000", | ||
| 343 | + status=EndpointStatus.INITIAL, | ||
| 344 | + headless=True, | ||
| 345 | + ) | ||
| 346 | + instance.add_endpoints(master.ip, {master.id: master}) | ||
| 347 | + instance.add_endpoints(worker.ip, {worker.id: worker}) | ||
| 348 | + | ||
| 349 | + assert instance.is_all_endpoints_ready() is False | ||
| 350 | + | ||
| 351 | + worker.status = EndpointStatus.WAIT2START | ||
| 352 | + assert instance.is_all_endpoints_ready() is True | ||
| 353 | + | ||
| 354 | + master.status = EndpointStatus.INITIAL | ||
| 355 | + assert instance.is_all_endpoints_ready() is False | ||
| 356 | + | ||
| 357 | + master.status = EndpointStatus.NORMAL | ||
| 358 | + worker.status = EndpointStatus.ABNORMAL | ||
| 359 | + assert instance.is_all_endpoints_ready() is False | ||
| 360 | + | ||
| 361 | + | ||
| 362 | +def test_instance_with_only_headless_workers_is_not_ready() -> None: | ||
| 363 | + instance = Instance( | ||
| 364 | + job_name="test_headless_only_readiness", | ||
| 365 | + model_name="test_model", | ||
| 366 | + id=1, | ||
| 367 | + role="prefill", | ||
| 368 | + parallel_config=ParallelConfig(dp_size=1, tp_size=4), | ||
| 369 | + enable_multi_endpoints=True, | ||
| 370 | + ) | ||
| 371 | + worker = Endpoint( | ||
| 372 | + id=0, | ||
| 373 | + ip="10.0.0.2", | ||
| 374 | + business_port="8000", | ||
| 375 | + mgmt_port="9000", | ||
| 376 | + status=EndpointStatus.WAIT2START, | ||
| 377 | + headless=True, | ||
| 378 | + ) | ||
| 379 | + instance.add_endpoints(worker.ip, {worker.id: worker}) | ||
| 380 | + | ||
| 381 | + assert instance.is_all_endpoints_ready() is False | ||
| 382 | + | ||
| 383 | + | ||
| 324 | def test_is_endpoints_enough_counts_all_endpoints() -> None: | 384 | def test_is_endpoints_enough_counts_all_endpoints() -> None: |
| 325 | """is_endpoints_enough counts all endpoints regardless of headless flag. | 385 | """is_endpoints_enough counts all endpoints regardless of headless flag. |
| 326 | This is safe because headless is only set when nnodes>1, and in that | 386 | This is safe because headless is only set when nnodes>1, and in that |
| @@ -167,6 +167,21 @@ def create_reregister_msg(job_name: str, pod_ip: str, instance_id: int, config: | |||
| 167 | ) | 167 | ) |
| 168 | 168 | ||
| 169 | 169 | ||
| 170 | +def test_build_endpoints_preserves_bootstrap_port(instance_assembler, test_config): | ||
| 171 | + msg = create_register_msg( | ||
| 172 | + "native-transfer", | ||
| 173 | + test_config["pod_ip1"], | ||
| 174 | + test_config, | ||
| 175 | + bootstrap_port=9100, | ||
| 176 | + ) | ||
| 177 | + | ||
| 178 | + endpoints = instance_assembler._build_multi_endpoints(msg, 0) | ||
| 179 | + | ||
| 180 | + assert endpoints | ||
| 181 | + for endpoint in endpoints.values(): | ||
| 182 | + assert endpoint.bootstrap_port == 9100 | ||
| 183 | + | ||
| 184 | + | ||
| 170 | def register_instance_with_pods(assembler: InstanceAssembler, job_name: str, config: dict, pod_count: int = 2) -> bool: | 185 | def register_instance_with_pods(assembler: InstanceAssembler, job_name: str, config: dict, pod_count: int = 2) -> bool: |
| 171 | """Register pods for an instance and return whether assembly is complete""" | 186 | """Register pods for an instance and return whether assembly is complete""" |
| 172 | pod_ips = [f"127.0.0.{i + 1}" for i in range(pod_count)] | 187 | pod_ips = [f"127.0.0.{i + 1}" for i in range(pod_count)] |
| @@ -558,6 +558,22 @@ def test_state_transitions(instance_manager, test_config): | |||
| 558 | assert instance.status == InsStatus.INACTIVE | 558 | assert instance.status == InsStatus.INACTIVE |
| 559 | 559 | ||
| 560 | 560 | ||
| 561 | +def test_initial_instance_activates_with_ready_master_and_live_headless_worker(): | ||
| 562 | + """A headless worker's process liveness is sufficient only when the routable master is ready.""" | ||
| 563 | + manager = create_instance_manager_with_config() | ||
| 564 | + instance = create_test_instance(205, "test_headless_worker", ["192.168.1.1", "192.168.1.2"]) | ||
| 565 | + worker = instance.endpoints["192.168.1.2"][0] | ||
| 566 | + worker.headless = True | ||
| 567 | + manager.add_instance(instance) | ||
| 568 | + instance.endpoints["192.168.1.1"][0].status = EndpointStatus.NORMAL | ||
| 569 | + worker.status = EndpointStatus.WAIT2START | ||
| 570 | + | ||
| 571 | + result = manager._handle_state_transition(instance) | ||
| 572 | + | ||
| 573 | + assert result is True | ||
| 574 | + assert instance.status == InsStatus.ACTIVE | ||
| 575 | + | ||
| 576 | + | ||
| 561 | def test_inactive_with_mixed_paused_normal_goes_to_paused(instance_manager): | 577 | def test_inactive_with_mixed_paused_normal_goes_to_paused(instance_manager): |
| 562 | """INACTIVE + mixed PAUSED/NORMAL (no ABNORMAL) → PAUSED, NOT INITIAL""" | 578 | """INACTIVE + mixed PAUSED/NORMAL (no ABNORMAL) → PAUSED, NOT INITIAL""" |
| 563 | manager = create_instance_manager_with_config() | 579 | manager = create_instance_manager_with_config() |
| @@ -0,0 +1,46 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from unittest.mock import MagicMock, patch | ||
| 12 | + | ||
| 13 | +from motor.config.tls_config import TLSConfig | ||
| 14 | +from motor.coordinator.api_client.native_engine_api_client import NativeEngineApiClient | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +def test_query_metrics_uses_inference_tls(mock_client): | ||
| 19 | + tls_config = TLSConfig(enable_tls=True) | ||
| 20 | + response = MagicMock(text="# TYPE ready gauge\nready 1", status_code=200) | ||
| 21 | + mock_client.return_value.__enter__.return_value.do_get.return_value = response | ||
| 22 | + | ||
| 23 | + result = NativeEngineApiClient.query_metrics("10.0.0.8:8000", tls_config) | ||
| 24 | + | ||
| 25 | + assert result == response.text | ||
| 26 | + mock_client.assert_called_once_with( | ||
| 27 | + address="10.0.0.8:8000", | ||
| 28 | + tls_config=tls_config, | ||
| 29 | + timeout=2, | ||
| 30 | + ) | ||
| 31 | + mock_client.return_value.__enter__.return_value.do_get.assert_called_once_with("/metrics") | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +def test_query_metrics_failure_returns_empty_text(mock_client): | ||
| 36 | + mock_client.side_effect = RuntimeError("connection refused") | ||
| 37 | + | ||
| 38 | + assert NativeEngineApiClient.query_metrics("10.0.0.8:8000", None) == "" | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +def test_query_metrics_non_success_status_returns_empty_text(mock_client): | ||
| 43 | + response = MagicMock(text="internal error", status_code=503) | ||
| 44 | + mock_client.return_value.__enter__.return_value.do_get.return_value = response | ||
| 45 | + | ||
| 46 | + assert NativeEngineApiClient.query_metrics("10.0.0.8:8000", None) == "" | ||
| @@ -13,7 +13,6 @@ import threading | |||
| 13 | 13 | ||
| 14 | import pytest | 14 | import pytest |
| 15 | 15 | ||
| 16 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 17 | from motor.common.resources import Instance, PDRole, Workload, Endpoint, EventType | 16 | from motor.common.resources import Instance, PDRole, Workload, Endpoint, EventType |
| 18 | from motor.config.coordinator import CoordinatorConfig | 17 | from motor.config.coordinator import CoordinatorConfig |
| 19 | from motor.coordinator.domain.instance_manager import InstanceManager, UpdateInstanceMode | 18 | from motor.coordinator.domain.instance_manager import InstanceManager, UpdateInstanceMode |
| @@ -34,7 +33,6 @@ class TestInstanceManager: | |||
| 34 | model_name="test-model", | 33 | model_name="test-model", |
| 35 | id=1, | 34 | id=1, |
| 36 | role=PDRole.ROLE_P, | 35 | role=PDRole.ROLE_P, |
| 37 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 38 | endpoints={}, | 36 | endpoints={}, |
| 39 | ) | 37 | ) |
| 40 | 38 | ||
| @@ -43,7 +41,6 @@ class TestInstanceManager: | |||
| 43 | model_name="test-model", | 41 | model_name="test-model", |
| 44 | id=2, | 42 | id=2, |
| 45 | role=PDRole.ROLE_D, | 43 | role=PDRole.ROLE_D, |
| 46 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 47 | endpoints={}, | 44 | endpoints={}, |
| 48 | ) | 45 | ) |
| 49 | 46 | ||
| @@ -74,13 +71,6 @@ class TestInstanceManager: | |||
| 74 | self.instance_manager._add_instance_to_available_pool(self.decode_instance) | 71 | self.instance_manager._add_instance_to_available_pool(self.decode_instance) |
| 75 | assert self.instance_manager.has_required_instances() is True | 72 | assert self.instance_manager.has_required_instances() is True |
| 76 | 73 | ||
| 77 | - def test_has_required_instances_rejects_incompatible_pd_pair(self): | ||
| 78 | - self.decode_instance.dispatch_capabilities = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | ||
| 79 | - self.instance_manager._add_instance_to_available_pool(self.prefill_instance) | ||
| 80 | - self.instance_manager._add_instance_to_available_pool(self.decode_instance) | ||
| 81 | - | ||
| 82 | - assert self.instance_manager.has_required_instances() is False | ||
| 83 | - | ||
| 84 | def test_has_required_instances_hybrid(self): | 74 | def test_has_required_instances_hybrid(self): |
| 85 | result = self.instance_manager._add_instance_to_available_pool(self.hybrid_instance) | 75 | result = self.instance_manager._add_instance_to_available_pool(self.hybrid_instance) |
| 86 | assert result is True | 76 | assert result is True |
| @@ -962,7 +952,6 @@ class TestInstanceManagerThreadSafety: | |||
| 962 | model_name="test-model", | 952 | model_name="test-model", |
| 963 | id=1, | 953 | id=1, |
| 964 | role=PDRole.ROLE_P, | 954 | role=PDRole.ROLE_P, |
| 965 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 966 | endpoints={}, | 955 | endpoints={}, |
| 967 | ) | 956 | ) |
| 968 | 957 | ||
| @@ -971,7 +960,6 @@ class TestInstanceManagerThreadSafety: | |||
| 971 | model_name="test-model", | 960 | model_name="test-model", |
| 972 | id=2, | 961 | id=2, |
| 973 | role=PDRole.ROLE_D, | 962 | role=PDRole.ROLE_D, |
| 974 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 975 | endpoints={}, | 963 | endpoints={}, |
| 976 | ) | 964 | ) |
| 977 | 965 | ||
| @@ -29,6 +29,7 @@ from motor.coordinator.metrics.metrics_collector import ( | |||
| 29 | Metric, | 29 | Metric, |
| 30 | _filter_kvstore_metrics, | 30 | _filter_kvstore_metrics, |
| 31 | ) | 31 | ) |
| 32 | +from motor.coordinator.api_client.native_engine_api_client import NativeEngineApiClient | ||
| 32 | from motor.config.coordinator import CoordinatorConfig | 33 | from motor.config.coordinator import CoordinatorConfig |
| 33 | from motor.common.utils.singleton import ThreadSafeSingleton | 34 | from motor.common.utils.singleton import ThreadSafeSingleton |
| 34 | from motor.coordinator.metrics.metric_computer import ( | 35 | from motor.coordinator.metrics.metric_computer import ( |
| @@ -873,6 +874,29 @@ http_request_duration_seconds_created{handler="/v1/chat/completions",method="POS | |||
| 873 | for port in [8004, 8005]: | 874 | for port in [8004, 8005]: |
| 874 | assert requests.get(f"http://localhost:{port}/metrics").status_code == 404 | 875 | assert requests.get(f"http://localhost:{port}/metrics").status_code == 404 |
| 875 | 876 | ||
| 877 | + | ||
| 878 | + | ||
| 879 | + def test_fetch_endpoint_metrics_uses_native_business_port_and_infer_tls(self, query_metrics): | ||
| 880 | + endpoint = Endpoint( | ||
| 881 | + id=7, | ||
| 882 | + ip="2001:db8::7", | ||
| 883 | + business_port="8007", | ||
| 884 | + mgmt_port="9007", | ||
| 885 | + ) | ||
| 886 | + instance = Instance( | ||
| 887 | + job_name="native-metrics", | ||
| 888 | + model_name="test-model", | ||
| 889 | + id=7, | ||
| 890 | + role=PDRole.ROLE_U, | ||
| 891 | + endpoints={"2001:db8::7": {7: endpoint}}, | ||
| 892 | + ) | ||
| 893 | + collector = MetricsCollector(self.config) | ||
| 894 | + | ||
| 895 | + result = collector._fetch_endpoint_metrics(instance) | ||
| 896 | + | ||
| 897 | + query_metrics.assert_called_once_with("[2001:db8::7]:8007", self.config.infer_tls_config) | ||
| 898 | + assert result["endpoints"][7]["metrics_str"] == "# TYPE ready gauge\nready 1" | ||
| 899 | + | ||
| 876 | def test_prometheus_metrics_handler(self, mock_metrics_collector): # pylint: disable=redefined-outer-name | 900 | def test_prometheus_metrics_handler(self, mock_metrics_collector): # pylint: disable=redefined-outer-name |
| 877 | mock_metrics_collector._last_metrics = None | 901 | mock_metrics_collector._last_metrics = None |
| 878 | mock_metrics_collector.get_metrics.return_value = "" | 902 | mock_metrics_collector.get_metrics.return_value = "" |
| @@ -16,7 +16,6 @@ Unit tests for semantic metrics modules: | |||
| 16 | import math | 16 | import math |
| 17 | from unittest.mock import patch, MagicMock | 17 | from unittest.mock import patch, MagicMock |
| 18 | 18 | ||
| 19 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 20 | from motor.coordinator.metrics.metric_types import ( | 19 | from motor.coordinator.metrics.metric_types import ( |
| 21 | AggregationContext, | 20 | AggregationContext, |
| 22 | AggregationScope, | 21 | AggregationScope, |
| @@ -582,8 +581,7 @@ def _apply_scope_filter( | |||
| 582 | ) -> list[tuple[int, Metric]]: | 581 | ) -> list[tuple[int, Metric]]: |
| 583 | """Mirror MetricsCollector._aggregate_metrics SERVICE-scope filtering.""" | 582 | """Mirror MetricsCollector._aggregate_metrics SERVICE-scope filtering.""" |
| 584 | if ctx is not None and ctx.scope == AggregationScope.SERVICE and ctx.instance_roles is not None: | 583 | if ctx is not None and ctx.scope == AggregationScope.SERVICE and ctx.instance_roles is not None: |
| 585 | - if name == "vllm:time_to_first_token_seconds" and ctx.instance_dispatch_capabilities is not None: | 584 | + if name == "vllm:time_to_first_token_seconds" and ctx.instance_engine_types is not None: |
| 586 | - handoff = DispatchPlan.PREFILL_HANDOFF_DECODE.value | ||
| 587 | entries = [ | 585 | entries = [ |
| 588 | (ins_id, metric) | 586 | (ins_id, metric) |
| 589 | for ins_id, metric in entries | 587 | for ins_id, metric in entries |
| @@ -592,7 +590,7 @@ def _apply_scope_filter( | |||
| 592 | or ctx.instance_roles.get(ins_id) in {"decode", "union", "both", "hybrid"} | 590 | or ctx.instance_roles.get(ins_id) in {"decode", "union", "both", "hybrid"} |
| 593 | or ( | 591 | or ( |
| 594 | ctx.instance_roles.get(ins_id) == "prefill" | 592 | ctx.instance_roles.get(ins_id) == "prefill" |
| 595 | - and handoff in ctx.instance_dispatch_capabilities.get(ins_id, set()) | 593 | + and ctx.instance_engine_types.get(ins_id, "").strip().lower() == "vllm" |
| 596 | ) | 594 | ) |
| 597 | ) | 595 | ) |
| 598 | ] | 596 | ] |
| @@ -619,30 +617,26 @@ class TestAggregationScope: | |||
| 619 | ctx = AggregationContext( | 617 | ctx = AggregationContext( |
| 620 | scope=AggregationScope.SERVICE, | 618 | scope=AggregationScope.SERVICE, |
| 621 | instance_roles={1: "prefill", 2: "decode"}, | 619 | instance_roles={1: "prefill", 2: "decode"}, |
| 622 | - instance_dispatch_capabilities={}, | 620 | + instance_engine_types={}, |
| 623 | ) | 621 | ) |
| 624 | entries = _apply_scope_filter("vllm:time_to_first_token_seconds", [(1, p_hist), (2, d_hist)], ctx) | 622 | entries = _apply_scope_filter("vllm:time_to_first_token_seconds", [(1, p_hist), (2, d_hist)], ctx) |
| 625 | assert len(entries) == 1 | 623 | assert len(entries) == 1 |
| 626 | ttft = self.engine.aggregate("vllm:time_to_first_token_seconds", [m for _, m in entries]) | 624 | ttft = self.engine.aggregate("vllm:time_to_first_token_seconds", [m for _, m in entries]) |
| 627 | assert ttft.value[-1] == 8.0 | 625 | assert ttft.value[-1] == 8.0 |
| 628 | 626 | ||
| 629 | - def test_service_scope_includes_only_handoff_prefill_ttft(self): | 627 | + def test_service_scope_includes_only_vllm_prefill_ttft(self): |
| 630 | - p_concurrent = _make_ttft_histogram(10.0, 5.0) | 628 | + p_sglang = _make_ttft_histogram(10.0, 5.0) |
| 631 | - p_handoff = _make_ttft_histogram(20.0, 15.0) | 629 | + p_vllm = _make_ttft_histogram(20.0, 15.0) |
| 632 | d_hist = _make_ttft_histogram(30.0, 9.0) | 630 | d_hist = _make_ttft_histogram(30.0, 9.0) |
| 633 | ctx = AggregationContext( | 631 | ctx = AggregationContext( |
| 634 | scope=AggregationScope.SERVICE, | 632 | scope=AggregationScope.SERVICE, |
| 635 | instance_roles={1: "prefill", 2: "prefill", 3: "decode"}, | 633 | instance_roles={1: "prefill", 2: "prefill", 3: "decode"}, |
| 636 | - instance_dispatch_capabilities={ | 634 | + instance_engine_types={1: "sglang", 2: "vllm", 3: "vllm"}, |
| 637 | - 1: {DispatchPlan.CONCURRENT_ENGINE_SYNC.value}, | ||
| 638 | - 2: {DispatchPlan.PREFILL_HANDOFF_DECODE.value}, | ||
| 639 | - 3: {DispatchPlan.PREFILL_HANDOFF_DECODE.value}, | ||
| 640 | - }, | ||
| 641 | ) | 635 | ) |
| 642 | 636 | ||
| 643 | entries = _apply_scope_filter( | 637 | entries = _apply_scope_filter( |
| 644 | "vllm:time_to_first_token_seconds", | 638 | "vllm:time_to_first_token_seconds", |
| 645 | - [(1, p_concurrent), (2, p_handoff), (3, d_hist)], | 639 | + [(1, p_sglang), (2, p_vllm), (3, d_hist)], |
| 646 | ctx, | 640 | ctx, |
| 647 | ) | 641 | ) |
| 648 | 642 | ||
| @@ -695,13 +689,6 @@ class TestRoleScope: | |||
| 695 | scope = MetricRegistry.get_effective_role_scope("vllm:time_to_first_token_seconds") | 689 | scope = MetricRegistry.get_effective_role_scope("vllm:time_to_first_token_seconds") |
| 696 | assert scope == "decode" | 690 | assert scope == "decode" |
| 697 | 691 | ||
| 698 | - def test_ttft_effective_role_scope_handoff_connector(self): | ||
| 699 | - scope = MetricRegistry.get_effective_role_scope( | ||
| 700 | - "vllm:time_to_first_token_seconds", | ||
| 701 | - {DispatchPlan.PREFILL_HANDOFF_DECODE.value}, | ||
| 702 | - ) | ||
| 703 | - assert scope is None # No filtering | ||
| 704 | - | ||
| 705 | def test_effective_role_scope_unknown_metric(self): | 692 | def test_effective_role_scope_unknown_metric(self): |
| 706 | scope = MetricRegistry.get_effective_role_scope("some_unknown_metric") | 693 | scope = MetricRegistry.get_effective_role_scope("some_unknown_metric") |
| 707 | assert scope is None | 694 | assert scope is None |
| @@ -1,101 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -from unittest.mock import AsyncMock, MagicMock, patch | ||
| 12 | - | ||
| 13 | -import pytest | ||
| 14 | - | ||
| 15 | -from motor.common.resources.dispatch import DispatchStopReason, DispatchStopState | ||
| 16 | -from motor.common.resources.endpoint import Endpoint, EndpointStatus | ||
| 17 | -from motor.common.resources.instance import Instance, InsStatus, PDRole, ParallelConfig | ||
| 18 | -from motor.coordinator.domain import ScheduledResource | ||
| 19 | -from motor.coordinator.router.dispatch_session import AttemptContext | ||
| 20 | -from motor.coordinator.router.stop_client import DispatchStopClient | ||
| 21 | - | ||
| 22 | - | ||
| 23 | -def _config(): | ||
| 24 | - cfg = MagicMock() | ||
| 25 | - cfg.infer_tls_config = None | ||
| 26 | - return cfg | ||
| 27 | - | ||
| 28 | - | ||
| 29 | -def _resource(engine_type: str) -> ScheduledResource: | ||
| 30 | - endpoint = Endpoint( | ||
| 31 | - id=1, | ||
| 32 | - ip="127.0.0.1", | ||
| 33 | - business_port="8000", | ||
| 34 | - mgmt_port="1026", | ||
| 35 | - status=EndpointStatus.NORMAL, | ||
| 36 | - ) | ||
| 37 | - instance = Instance( | ||
| 38 | - job_name="job-1", | ||
| 39 | - model_name="m", | ||
| 40 | - engine_type=engine_type, | ||
| 41 | - id=1, | ||
| 42 | - role=PDRole.ROLE_D, | ||
| 43 | - status=InsStatus.ACTIVE, | ||
| 44 | - parallel_config=ParallelConfig(dp_size=1), | ||
| 45 | - endpoints={endpoint.ip: {endpoint.id: endpoint}}, | ||
| 46 | - ) | ||
| 47 | - return ScheduledResource(instance=instance, endpoint=endpoint) | ||
| 48 | - | ||
| 49 | - | ||
| 50 | - | ||
| 51 | -async def test_dispatch_stop_uses_sglang_abort_request(): | ||
| 52 | - client = DispatchStopClient(_config()) | ||
| 53 | - attempt = AttemptContext(root_request_id="r1", attempt_seq=1, pair_id="p1") | ||
| 54 | - mock_http = AsyncMock() | ||
| 55 | - mock_response = MagicMock() | ||
| 56 | - mock_response.raise_for_status = MagicMock() | ||
| 57 | - mock_http.post = AsyncMock(return_value=mock_response) | ||
| 58 | - | ||
| 59 | - with patch("motor.coordinator.router.stop_client.HTTPClientPool") as pool_cls: | ||
| 60 | - pool_cls.return_value.get_client = AsyncMock(return_value=mock_http) | ||
| 61 | - result = await client.stop( | ||
| 62 | - _resource("sglang"), | ||
| 63 | - attempt, | ||
| 64 | - DispatchStopReason.PEER_FAILED, | ||
| 65 | - ) | ||
| 66 | - | ||
| 67 | - assert result is not None | ||
| 68 | - assert result.accepted is True | ||
| 69 | - assert result.state == DispatchStopState.STOPPED | ||
| 70 | - mock_http.post.assert_awaited_once() | ||
| 71 | - assert mock_http.post.await_args.args[0] == "/abort_request" | ||
| 72 | - assert mock_http.post.await_args.kwargs["json"] == {"rid": "r1#a1"} | ||
| 73 | - | ||
| 74 | - | ||
| 75 | - | ||
| 76 | -async def test_dispatch_stop_posts_for_non_sglang(): | ||
| 77 | - client = DispatchStopClient(_config()) | ||
| 78 | - attempt = AttemptContext(root_request_id="r1", attempt_seq=1, pair_id="p1") | ||
| 79 | - mock_http = AsyncMock() | ||
| 80 | - mock_response = MagicMock() | ||
| 81 | - mock_response.raise_for_status = MagicMock() | ||
| 82 | - mock_response.json.return_value = { | ||
| 83 | - "root_request_id": "r1", | ||
| 84 | - "attempt_seq": 1, | ||
| 85 | - "accepted": True, | ||
| 86 | - "state": DispatchStopState.STOPPED.value, | ||
| 87 | - "message": "", | ||
| 88 | - } | ||
| 89 | - mock_http.post = AsyncMock(return_value=mock_response) | ||
| 90 | - | ||
| 91 | - with patch("motor.coordinator.router.stop_client.HTTPClientPool") as pool_cls: | ||
| 92 | - pool_cls.return_value.get_client = AsyncMock(return_value=mock_http) | ||
| 93 | - result = await client.stop( | ||
| 94 | - _resource("vllm"), | ||
| 95 | - attempt, | ||
| 96 | - DispatchStopReason.PEER_FAILED, | ||
| 97 | - ) | ||
| 98 | - | ||
| 99 | - assert result is not None | ||
| 100 | - mock_http.post.assert_awaited_once() | ||
| 101 | - assert mock_http.post.await_args.args[0] == "/v1/dispatch/stop" | ||
| @@ -9,6 +9,7 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | import asyncio | 11 | import asyncio |
| 12 | +import json | ||
| 12 | 13 | ||
| 13 | import pytest | 14 | import pytest |
| 14 | from fastapi import FastAPI, HTTPException, Request | 15 | from fastapi import FastAPI, HTTPException, Request |
| @@ -17,7 +18,6 @@ from fastapi.testclient import TestClient | |||
| 17 | 18 | ||
| 18 | import motor.common.utils.error as cancel_error | 19 | import motor.common.utils.error as cancel_error |
| 19 | from motor.common.logger.logger import _resolve_logger_name | 20 | from motor.common.logger.logger import _resolve_logger_name |
| 20 | -from motor.common.resources.dispatch import DispatchPlan, has_compatible_dispatch_pair | ||
| 21 | from motor.config.coordinator import CoordinatorConfig | 21 | from motor.config.coordinator import CoordinatorConfig |
| 22 | from motor.common.resources.instance import Instance, PDRole | 22 | from motor.common.resources.instance import Instance, PDRole |
| 23 | from motor.coordinator.domain import InstanceReadiness | 23 | from motor.coordinator.domain import InstanceReadiness |
| @@ -44,13 +44,6 @@ class _Scheduler: | |||
| 44 | return InstanceReadiness.ONLY_PREFILL | 44 | return InstanceReadiness.ONLY_PREFILL |
| 45 | return InstanceReadiness.REQUIRED_MET | 45 | return InstanceReadiness.REQUIRED_MET |
| 46 | 46 | ||
| 47 | - async def has_compatible_pd_pair(self): | ||
| 48 | - if self._instances is None: | ||
| 49 | - return False | ||
| 50 | - prefill = [instance for instance in self._instances.values() if instance.role == PDRole.ROLE_P.value] | ||
| 51 | - decode = [instance for instance in self._instances.values() if instance.role == PDRole.ROLE_D.value] | ||
| 52 | - return has_compatible_dispatch_pair(prefill, decode) | ||
| 53 | - | ||
| 54 | 47 | ||
| 55 | def _app(config: CoordinatorConfig, scheduler: _Scheduler) -> FastAPI: | 48 | def _app(config: CoordinatorConfig, scheduler: _Scheduler) -> FastAPI: |
| 56 | app = FastAPI() | 49 | app = FastAPI() |
| @@ -90,14 +83,12 @@ def test_dispatch_uses_unified_router_by_default(monkeypatch): | |||
| 90 | model_name="m", | 83 | model_name="m", |
| 91 | id=1, | 84 | id=1, |
| 92 | role=PDRole.ROLE_P.value, | 85 | role=PDRole.ROLE_P.value, |
| 93 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 94 | ), | 86 | ), |
| 95 | 2: Instance( | 87 | 2: Instance( |
| 96 | job_name="d", | 88 | job_name="d", |
| 97 | model_name="m", | 89 | model_name="m", |
| 98 | id=2, | 90 | id=2, |
| 99 | role=PDRole.ROLE_D.value, | 91 | role=PDRole.ROLE_D.value, |
| 100 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 101 | ), | 92 | ), |
| 102 | } | 93 | } |
| 103 | client = TestClient(_app(_config(), _Scheduler(instances))) | 94 | client = TestClient(_app(_config(), _Scheduler(instances))) |
| @@ -165,103 +156,43 @@ def test_dispatch_rejects_only_prefill_when_hybrid_fallback_disabled(monkeypatch | |||
| 165 | assert response.json()["detail"] == "PD separate service is unavailable and fallback to hybrid is disabled" | 156 | assert response.json()["detail"] == "PD separate service is unavailable and fallback to hybrid is disabled" |
| 166 | 157 | ||
| 167 | 158 | ||
| 168 | -def test_dispatch_rejects_incompatible_pd_topology(): | 159 | +def test_dispatch_reuses_request_json_parsed_at_ingress(monkeypatch): |
| 169 | - instances = { | 160 | + """A body parsed by the API layer is not parsed again while building RequestInfo.""" |
| 170 | - 1: Instance( | 161 | + calls = [] |
| 171 | - job_name="p", | ||
| 172 | - model_name="m", | ||
| 173 | - id=1, | ||
| 174 | - role=PDRole.ROLE_P.value, | ||
| 175 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 176 | - ), | ||
| 177 | - 2: Instance( | ||
| 178 | - job_name="d", | ||
| 179 | - model_name="m", | ||
| 180 | - id=2, | ||
| 181 | - role=PDRole.ROLE_D.value, | ||
| 182 | - dispatch_capabilities=[DispatchPlan.PREFILL_HANDOFF_DECODE.value], | ||
| 183 | - ), | ||
| 184 | - } | ||
| 185 | 162 | ||
| 186 | - response = TestClient(_app(_config(), _Scheduler(instances))).post( | ||
| 187 | - "/v1/completions", | ||
| 188 | - json={"model": "m", "prompt": "hi"}, | ||
| 189 | - ) | ||
| 190 | - | ||
| 191 | - assert response.status_code == 503 | ||
| 192 | - | ||
| 193 | - | ||
| 194 | -def test_dispatch_falls_back_to_union_for_incompatible_pd(monkeypatch): | ||
| 195 | class _FakeHybridRouter: | 163 | class _FakeHybridRouter: |
| 196 | def __init__(self, req_info, config, scheduler=None, request_manager=None, sampling_manager=None): | 164 | def __init__(self, req_info, config, scheduler=None, request_manager=None, sampling_manager=None): |
| 197 | - pass | 165 | + calls.append(req_info.req_data) |
| 198 | 166 | ||
| 199 | async def handle_request(self): | 167 | async def handle_request(self): |
| 200 | return JSONResponse({"router": "hybrid"}) | 168 | return JSONResponse({"router": "hybrid"}) |
| 201 | 169 | ||
| 202 | monkeypatch.setattr(dispatch, "PDHybridRouter", _FakeHybridRouter) | 170 | monkeypatch.setattr(dispatch, "PDHybridRouter", _FakeHybridRouter) |
| 203 | - instances = { | 171 | + app = FastAPI() |
| 204 | - 1: Instance( | 172 | + config = CoordinatorConfig() |
| 205 | - job_name="p", | 173 | + request_manager = RequestManager(config) |
| 206 | - model_name="m", | 174 | + scheduler = _Scheduler({1: Instance(job_name="p", model_name="m", id=1, role=PDRole.ROLE_P.value)}) |
| 207 | - id=1, | ||
| 208 | - role=PDRole.ROLE_P.value, | ||
| 209 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 210 | - ), | ||
| 211 | - 2: Instance( | ||
| 212 | - job_name="d", | ||
| 213 | - model_name="m", | ||
| 214 | - id=2, | ||
| 215 | - role=PDRole.ROLE_D.value, | ||
| 216 | - dispatch_capabilities=[DispatchPlan.PREFILL_HANDOFF_DECODE.value], | ||
| 217 | - ), | ||
| 218 | - 3: Instance(job_name="u", model_name="m", id=3, role=PDRole.ROLE_U.value), | ||
| 219 | - } | ||
| 220 | 175 | ||
| 221 | - response = TestClient(_app(_config(), _Scheduler(instances))).post( | 176 | + @app.post("/v1/completions") |
| 222 | - "/v1/completions", | 177 | + async def completions(request: Request): |
| 223 | - json={"model": "m", "prompt": "hi"}, | 178 | + request_json = json.loads((await request.body()).decode("utf-8")) |
| 224 | - ) | 179 | + |
| 180 | + async def unexpected_second_parse(): | ||
| 181 | + raise AssertionError("request.json() must not be called after ingress parsing") | ||
| 182 | + | ||
| 183 | + monkeypatch.setattr(request, "json", unexpected_second_parse) | ||
| 184 | + return await dispatch.handle_request( | ||
| 185 | + request, | ||
| 186 | + config, | ||
| 187 | + scheduler=scheduler, | ||
| 188 | + request_manager=request_manager, | ||
| 189 | + request_json=request_json, | ||
| 190 | + ) | ||
| 191 | + | ||
| 192 | + response = TestClient(app).post("/v1/completions", json={"model": "m", "prompt": "hi"}) | ||
| 225 | 193 | ||
| 226 | assert response.status_code == 200 | 194 | assert response.status_code == 200 |
| 227 | - assert response.json() == {"router": "hybrid"} | 195 | + assert calls == [{"model": "m", "prompt": "hi"}] |
| 228 | - | ||
| 229 | - | ||
| 230 | -def test_dispatch_rejects_union_fallback_when_hybrid_fallback_disabled(monkeypatch): | ||
| 231 | - class _FakeHybridRouter: | ||
| 232 | - def __init__(self, req_info, config, scheduler=None, request_manager=None, sampling_manager=None): | ||
| 233 | - raise AssertionError("PDHybridRouter should not be used when fallback is disabled") | ||
| 234 | - | ||
| 235 | - monkeypatch.setattr(dispatch, "PDHybridRouter", _FakeHybridRouter) | ||
| 236 | - instances = { | ||
| 237 | - 1: Instance( | ||
| 238 | - job_name="p", | ||
| 239 | - model_name="m", | ||
| 240 | - id=1, | ||
| 241 | - role=PDRole.ROLE_P.value, | ||
| 242 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 243 | - ), | ||
| 244 | - 2: Instance( | ||
| 245 | - job_name="d", | ||
| 246 | - model_name="m", | ||
| 247 | - id=2, | ||
| 248 | - role=PDRole.ROLE_D.value, | ||
| 249 | - dispatch_capabilities=[DispatchPlan.PREFILL_HANDOFF_DECODE.value], | ||
| 250 | - ), | ||
| 251 | - 3: Instance(job_name="u", model_name="m", id=3, role=PDRole.ROLE_U.value), | ||
| 252 | - } | ||
| 253 | - config = CoordinatorConfig() | ||
| 254 | - config.scheduler_config.enable_pd_separation_fallback_to_hybrid = False | ||
| 255 | - | ||
| 256 | - response = TestClient(_app(config, _Scheduler(instances))).post( | ||
| 257 | - "/v1/completions", | ||
| 258 | - json={"model": "m", "prompt": "hi"}, | ||
| 259 | - ) | ||
| 260 | - | ||
| 261 | - assert response.status_code == 503 | ||
| 262 | - assert ( | ||
| 263 | - response.json()["detail"] == "PD separate service has no compatible P/D pair and fallback to hybrid is disabled" | ||
| 264 | - ) | ||
| 265 | 196 | ||
| 266 | 197 | ||
| 267 | def test_dispatch_preserves_upstream_http_error(monkeypatch): | 198 | def test_dispatch_preserves_upstream_http_error(monkeypatch): |
| @@ -286,14 +217,12 @@ def test_dispatch_preserves_upstream_http_error(monkeypatch): | |||
| 286 | model_name="m", | 217 | model_name="m", |
| 287 | id=1, | 218 | id=1, |
| 288 | role=PDRole.ROLE_P.value, | 219 | role=PDRole.ROLE_P.value, |
| 289 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 290 | ), | 220 | ), |
| 291 | 2: Instance( | 221 | 2: Instance( |
| 292 | job_name="d", | 222 | job_name="d", |
| 293 | model_name="m", | 223 | model_name="m", |
| 294 | id=2, | 224 | id=2, |
| 295 | role=PDRole.ROLE_D.value, | 225 | role=PDRole.ROLE_D.value, |
| 296 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 297 | ), | 226 | ), |
| 298 | } | 227 | } |
| 299 | 228 | ||
| @@ -329,14 +258,12 @@ def test_dispatch_request_cancelled_does_not_log_error(monkeypatch, caplog): | |||
| 329 | model_name="m", | 258 | model_name="m", |
| 330 | id=1, | 259 | id=1, |
| 331 | role=PDRole.ROLE_P.value, | 260 | role=PDRole.ROLE_P.value, |
| 332 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 333 | ), | 261 | ), |
| 334 | 2: Instance( | 262 | 2: Instance( |
| 335 | job_name="d", | 263 | job_name="d", |
| 336 | model_name="m", | 264 | model_name="m", |
| 337 | id=2, | 265 | id=2, |
| 338 | role=PDRole.ROLE_D.value, | 266 | role=PDRole.ROLE_D.value, |
| 339 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 340 | ), | 267 | ), |
| 341 | } | 268 | } |
| 342 | 269 | ||
| @@ -373,14 +300,12 @@ def test_dispatch_dispatch_abort_still_returns_500(monkeypatch, caplog): | |||
| 373 | model_name="m", | 300 | model_name="m", |
| 374 | id=1, | 301 | id=1, |
| 375 | role=PDRole.ROLE_P.value, | 302 | role=PDRole.ROLE_P.value, |
| 376 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 377 | ), | 303 | ), |
| 378 | 2: Instance( | 304 | 2: Instance( |
| 379 | job_name="d", | 305 | job_name="d", |
| 380 | model_name="m", | 306 | model_name="m", |
| 381 | id=2, | 307 | id=2, |
| 382 | role=PDRole.ROLE_D.value, | 308 | role=PDRole.ROLE_D.value, |
| 383 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 384 | ), | 309 | ), |
| 385 | } | 310 | } |
| 386 | 311 | ||
| @@ -451,14 +376,12 @@ def test_dispatch_unexpected_exception_still_logs_error(monkeypatch, caplog): | |||
| 451 | model_name="m", | 376 | model_name="m", |
| 452 | id=1, | 377 | id=1, |
| 453 | role=PDRole.ROLE_P.value, | 378 | role=PDRole.ROLE_P.value, |
| 454 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 455 | ), | 379 | ), |
| 456 | 2: Instance( | 380 | 2: Instance( |
| 457 | job_name="d", | 381 | job_name="d", |
| 458 | model_name="m", | 382 | model_name="m", |
| 459 | id=2, | 383 | id=2, |
| 460 | role=PDRole.ROLE_D.value, | 384 | role=PDRole.ROLE_D.value, |
| 461 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 462 | ), | 385 | ), |
| 463 | } | 386 | } |
| 464 | 387 | ||
| @@ -0,0 +1,501 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""HTTP-contract tests for Coordinator-to-native-engine P/D traffic. | ||
| 12 | + | ||
| 13 | +These tests use real ``httpx`` request serialization and ASGI applications as | ||
| 14 | +fake native engines. Only endpoint-to-transport binding is replaced; the | ||
| 15 | +Coordinator request, response, SSE, error, and protocol-adapter paths run | ||
| 16 | +unchanged. | ||
| 17 | +""" | ||
| 18 | + | ||
| 19 | +from contextlib import asynccontextmanager | ||
| 20 | +import json | ||
| 21 | +from unittest.mock import AsyncMock | ||
| 22 | + | ||
| 23 | +from fastapi import FastAPI, Request | ||
| 24 | +from fastapi.responses import JSONResponse, StreamingResponse | ||
| 25 | +import httpx | ||
| 26 | +import pytest | ||
| 27 | + | ||
| 28 | +from motor.common.http import HTTPClientPool | ||
| 29 | +from motor.common.resources.endpoint import Endpoint, EndpointStatus, Workload | ||
| 30 | +from motor.common.resources.instance import Instance, InsStatus, ParallelConfig, PDRole | ||
| 31 | +from motor.config.coordinator import CoordinatorConfig, ExceptionConfig | ||
| 32 | +from motor.config.tls_config import TLSConfig | ||
| 33 | +from motor.coordinator.domain import ScheduledResource | ||
| 34 | +from motor.coordinator.domain.request_manager import RequestManager | ||
| 35 | +from motor.coordinator.models.request import RequestInfo, ReqState | ||
| 36 | +from motor.coordinator.router.strategies.unified_pd import UnifiedPDRouter | ||
| 37 | +from motor.coordinator.router.upstream_error import UpstreamHTTPError | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +def _instance( | ||
| 41 | + instance_id: int, | ||
| 42 | + role: PDRole, | ||
| 43 | + engine_type: str, | ||
| 44 | + *, | ||
| 45 | + bootstrap_port: int | None = None, | ||
| 46 | +) -> Instance: | ||
| 47 | + endpoint = Endpoint( | ||
| 48 | + id=instance_id, | ||
| 49 | + ip="127.0.0.1", | ||
| 50 | + business_port=str(8300 + instance_id), | ||
| 51 | + mgmt_port=str(9300 + instance_id), | ||
| 52 | + bootstrap_port=bootstrap_port, | ||
| 53 | + status=EndpointStatus.NORMAL, | ||
| 54 | + ) | ||
| 55 | + return Instance( | ||
| 56 | + job_name=f"native-{engine_type}-{role.value}-{instance_id}", | ||
| 57 | + model_name="test-model", | ||
| 58 | + engine_type=engine_type, | ||
| 59 | + id=instance_id, | ||
| 60 | + role=role, | ||
| 61 | + status=InsStatus.ACTIVE, | ||
| 62 | + parallel_config=ParallelConfig(dp_size=1, tp_size=1), | ||
| 63 | + endpoints={endpoint.ip: {endpoint.id: endpoint}}, | ||
| 64 | + ) | ||
| 65 | + | ||
| 66 | + | ||
| 67 | +class _ContractScheduler: | ||
| 68 | + def __init__(self, engine_type: str): | ||
| 69 | + bootstrap_port = 21001 if engine_type == "sglang" else None | ||
| 70 | + self.p = _instance(1, PDRole.ROLE_P, engine_type, bootstrap_port=bootstrap_port) | ||
| 71 | + self.d = _instance(2, PDRole.ROLE_D, engine_type) | ||
| 72 | + self.cb_events: list[tuple[int, str]] = [] | ||
| 73 | + self.workload_updates = [] | ||
| 74 | + | ||
| 75 | + async def select_and_allocate(self, role, _req_info, **_kwargs): | ||
| 76 | + instance = self.p if role == PDRole.ROLE_P else self.d | ||
| 77 | + endpoint = next(iter(next(iter(instance.endpoints.values())).values())) | ||
| 78 | + return instance, endpoint, Workload(active_tokens=1) | ||
| 79 | + | ||
| 80 | + async def update_workload(self, params): | ||
| 81 | + self.workload_updates.append(params) | ||
| 82 | + return True | ||
| 83 | + | ||
| 84 | + async def report_cb_event(self, instance_id: int, event: str) -> None: | ||
| 85 | + self.cb_events.append((instance_id, event)) | ||
| 86 | + | ||
| 87 | + async def get_unblocked_instances(self, role) -> list[int]: | ||
| 88 | + if role == PDRole.ROLE_P: | ||
| 89 | + return [self.p.id] | ||
| 90 | + if role == PDRole.ROLE_D: | ||
| 91 | + return [self.d.id] | ||
| 92 | + return [] | ||
| 93 | + | ||
| 94 | + | ||
| 95 | +def _config() -> CoordinatorConfig: | ||
| 96 | + config = CoordinatorConfig() | ||
| 97 | + config.exception_config = ExceptionConfig(max_retry=1, retry_delay=0) | ||
| 98 | + return config | ||
| 99 | + | ||
| 100 | + | ||
| 101 | +def _router(req_info: RequestInfo, scheduler: _ContractScheduler) -> UnifiedPDRouter: | ||
| 102 | + config = _config() | ||
| 103 | + return UnifiedPDRouter( | ||
| 104 | + req_info, | ||
| 105 | + config, | ||
| 106 | + scheduler=scheduler, | ||
| 107 | + request_manager=RequestManager(config), | ||
| 108 | + ) | ||
| 109 | + | ||
| 110 | + | ||
| 111 | + | ||
| 112 | +async def _bind_native_apps(monkeypatch, router: UnifiedPDRouter, p_app: FastAPI, d_app: FastAPI): | ||
| 113 | + async with ( | ||
| 114 | + httpx.AsyncClient( | ||
| 115 | + transport=httpx.ASGITransport(app=p_app), | ||
| 116 | + base_url="http://native-prefill", | ||
| 117 | + ) as p_client, | ||
| 118 | + httpx.AsyncClient( | ||
| 119 | + transport=httpx.ASGITransport(app=d_app), | ||
| 120 | + base_url="http://native-decode", | ||
| 121 | + ) as d_client, | ||
| 122 | + ): | ||
| 123 | + | ||
| 124 | + | ||
| 125 | + async def _client_for(resource: ScheduledResource): | ||
| 126 | + yield p_client if resource.instance.role == PDRole.ROLE_P else d_client | ||
| 127 | + | ||
| 128 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 129 | + try: | ||
| 130 | + yield | ||
| 131 | + finally: | ||
| 132 | + await HTTPClientPool().close_all() | ||
| 133 | + | ||
| 134 | + | ||
| 135 | +async def _collect_stream_body(response) -> bytes: | ||
| 136 | + chunks = [] | ||
| 137 | + async for chunk in response.body_iterator: | ||
| 138 | + chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode()) | ||
| 139 | + return b"".join(chunks) | ||
| 140 | + | ||
| 141 | + | ||
| 142 | +def _sse(payload: dict) -> bytes: | ||
| 143 | + return b"data: " + json.dumps(payload, separators=(",", ":")).encode() + b"\n\n" | ||
| 144 | + | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +async def test_vllm_handoff_preserves_chat_tools_logprobs_and_usage_over_http(monkeypatch): | ||
| 148 | + p_requests = [] | ||
| 149 | + d_requests = [] | ||
| 150 | + p_headers = [] | ||
| 151 | + d_headers = [] | ||
| 152 | + p_app = FastAPI() | ||
| 153 | + d_app = FastAPI() | ||
| 154 | + | ||
| 155 | + | ||
| 156 | + async def prefill(request: Request): | ||
| 157 | + body = await request.json() | ||
| 158 | + p_requests.append(body) | ||
| 159 | + p_headers.append(dict(request.headers)) | ||
| 160 | + return { | ||
| 161 | + "kv_transfer_params": { | ||
| 162 | + "do_remote_prefill": True, | ||
| 163 | + "remote_request_id": body["request_id"], | ||
| 164 | + "remote_host": "10.0.0.8", | ||
| 165 | + "remote_port": 25000, | ||
| 166 | + "connector_private": {"ticket": "opaque"}, | ||
| 167 | + }, | ||
| 168 | + "usage": { | ||
| 169 | + "prompt_tokens": 9, | ||
| 170 | + "prompt_tokens_details": {"cached_tokens": 4}, | ||
| 171 | + }, | ||
| 172 | + } | ||
| 173 | + | ||
| 174 | + | ||
| 175 | + async def decode(request: Request): | ||
| 176 | + body = await request.json() | ||
| 177 | + d_requests.append(body) | ||
| 178 | + d_headers.append(dict(request.headers)) | ||
| 179 | + return { | ||
| 180 | + "id": "chatcmpl-native", | ||
| 181 | + "object": "chat.completion", | ||
| 182 | + "choices": [ | ||
| 183 | + { | ||
| 184 | + "index": 0, | ||
| 185 | + "message": { | ||
| 186 | + "role": "assistant", | ||
| 187 | + "content": None, | ||
| 188 | + "tool_calls": [ | ||
| 189 | + { | ||
| 190 | + "id": "call_weather", | ||
| 191 | + "type": "function", | ||
| 192 | + "function": { | ||
| 193 | + "name": "weather", | ||
| 194 | + "arguments": '{"city":"Shenzhen"}', | ||
| 195 | + }, | ||
| 196 | + } | ||
| 197 | + ], | ||
| 198 | + }, | ||
| 199 | + "finish_reason": "tool_calls", | ||
| 200 | + "logprobs": { | ||
| 201 | + "content": [ | ||
| 202 | + { | ||
| 203 | + "token": "weather", | ||
| 204 | + "logprob": -0.1, | ||
| 205 | + "top_logprobs": [], | ||
| 206 | + } | ||
| 207 | + ] | ||
| 208 | + }, | ||
| 209 | + "token_ids": [42], | ||
| 210 | + } | ||
| 211 | + ], | ||
| 212 | + "usage": { | ||
| 213 | + "prompt_tokens": 9, | ||
| 214 | + "completion_tokens": 1, | ||
| 215 | + "total_tokens": 10, | ||
| 216 | + }, | ||
| 217 | + "kv_transfer_params": {"must_not_leak": True}, | ||
| 218 | + } | ||
| 219 | + | ||
| 220 | + tools = [ | ||
| 221 | + { | ||
| 222 | + "type": "function", | ||
| 223 | + "function": { | ||
| 224 | + "name": "weather", | ||
| 225 | + "description": "Get weather", | ||
| 226 | + "parameters": { | ||
| 227 | + "type": "object", | ||
| 228 | + "properties": {"city": {"type": "string"}}, | ||
| 229 | + }, | ||
| 230 | + }, | ||
| 231 | + } | ||
| 232 | + ] | ||
| 233 | + req_info = RequestInfo( | ||
| 234 | + req_id="contract-vllm-chat", | ||
| 235 | + req_data={ | ||
| 236 | + "model": "test-model", | ||
| 237 | + "messages": [{"role": "user", "content": "Weather?"}], | ||
| 238 | + "tools": tools, | ||
| 239 | + "tool_choice": "auto", | ||
| 240 | + "logprobs": True, | ||
| 241 | + "top_logprobs": 2, | ||
| 242 | + "max_completion_tokens": 6, | ||
| 243 | + "stream": False, | ||
| 244 | + }, | ||
| 245 | + api="v1/chat/completions", | ||
| 246 | + entry_api="v1/chat/completions", | ||
| 247 | + req_len=9, | ||
| 248 | + ) | ||
| 249 | + scheduler = _ContractScheduler("vllm") | ||
| 250 | + router = _router(req_info, scheduler) | ||
| 251 | + | ||
| 252 | + async with _bind_native_apps(monkeypatch, router, p_app, d_app): | ||
| 253 | + response = await router.handle_request() | ||
| 254 | + | ||
| 255 | + body = json.loads(response.body) | ||
| 256 | + assert p_requests[0]["stream"] is False | ||
| 257 | + assert p_requests[0]["max_tokens"] == 1 | ||
| 258 | + assert p_requests[0]["max_completion_tokens"] == 1 | ||
| 259 | + assert p_requests[0]["tools"] == tools | ||
| 260 | + assert d_requests[0]["max_completion_tokens"] == 6 | ||
| 261 | + assert d_requests[0]["tools"] == tools | ||
| 262 | + assert d_requests[0]["kv_transfer_params"]["connector_private"] == {"ticket": "opaque"} | ||
| 263 | + assert p_headers[0]["x-request-id"] == d_headers[0]["x-request-id"] == "contract-vllm-chat#a1" | ||
| 264 | + assert "authorization" not in p_headers[0] | ||
| 265 | + assert "authorization" not in d_headers[0] | ||
| 266 | + assert body["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "weather" | ||
| 267 | + assert body["choices"][0]["logprobs"]["content"][0]["logprob"] == -0.1 | ||
| 268 | + assert body["usage"]["prompt_tokens_details"] == {"cached_tokens": 4} | ||
| 269 | + assert "token_ids" not in body["choices"][0] | ||
| 270 | + assert "kv_transfer_params" not in body | ||
| 271 | + assert req_info.state == ReqState.DECODE_END | ||
| 272 | + assert scheduler.cb_events == [(1, "success"), (2, "success")] | ||
| 273 | + | ||
| 274 | + | ||
| 275 | + | ||
| 276 | +async def test_sglang_bootstrap_preserves_stream_sse_tools_logprobs_and_usage_over_http(monkeypatch): | ||
| 277 | + p_requests = [] | ||
| 278 | + d_requests = [] | ||
| 279 | + p_app = FastAPI() | ||
| 280 | + d_app = FastAPI() | ||
| 281 | + | ||
| 282 | + | ||
| 283 | + async def prefill(request: Request): | ||
| 284 | + p_requests.append(await request.json()) | ||
| 285 | + | ||
| 286 | + async def frames(): | ||
| 287 | + yield _sse( | ||
| 288 | + { | ||
| 289 | + "choices": [], | ||
| 290 | + "usage": { | ||
| 291 | + "prompt_tokens": 7, | ||
| 292 | + "prompt_tokens_details": {"cached_tokens": 3}, | ||
| 293 | + }, | ||
| 294 | + } | ||
| 295 | + ) | ||
| 296 | + yield b"data: [DONE]\n\n" | ||
| 297 | + | ||
| 298 | + return StreamingResponse(frames(), media_type="text/event-stream") | ||
| 299 | + | ||
| 300 | + | ||
| 301 | + async def decode(request: Request): | ||
| 302 | + d_requests.append(await request.json()) | ||
| 303 | + | ||
| 304 | + async def frames(): | ||
| 305 | + yield _sse( | ||
| 306 | + { | ||
| 307 | + "id": "chatcmpl-sglang", | ||
| 308 | + "choices": [ | ||
| 309 | + { | ||
| 310 | + "index": 0, | ||
| 311 | + "delta": { | ||
| 312 | + "tool_calls": [ | ||
| 313 | + { | ||
| 314 | + "index": 0, | ||
| 315 | + "id": "call_time", | ||
| 316 | + "type": "function", | ||
| 317 | + "function": { | ||
| 318 | + "name": "local_time", | ||
| 319 | + "arguments": '{"city":"Shanghai"}', | ||
| 320 | + }, | ||
| 321 | + } | ||
| 322 | + ] | ||
| 323 | + }, | ||
| 324 | + "logprobs": { | ||
| 325 | + "content": [ | ||
| 326 | + { | ||
| 327 | + "token": "time", | ||
| 328 | + "logprob": -0.2, | ||
| 329 | + "top_logprobs": [], | ||
| 330 | + } | ||
| 331 | + ] | ||
| 332 | + }, | ||
| 333 | + "token_ids": [51], | ||
| 334 | + "finish_reason": "tool_calls", | ||
| 335 | + } | ||
| 336 | + ], | ||
| 337 | + "bootstrap_host": "127.0.0.1", | ||
| 338 | + "bootstrap_port": 21001, | ||
| 339 | + "bootstrap_room": d_requests[0]["bootstrap_room"], | ||
| 340 | + } | ||
| 341 | + ) | ||
| 342 | + yield _sse( | ||
| 343 | + { | ||
| 344 | + "choices": [], | ||
| 345 | + "usage": { | ||
| 346 | + "prompt_tokens": 7, | ||
| 347 | + "completion_tokens": 1, | ||
| 348 | + "total_tokens": 8, | ||
| 349 | + }, | ||
| 350 | + "bootstrap_room": d_requests[0]["bootstrap_room"], | ||
| 351 | + } | ||
| 352 | + ) | ||
| 353 | + yield b"data: [DONE]\n\n" | ||
| 354 | + | ||
| 355 | + return StreamingResponse(frames(), media_type="text/event-stream") | ||
| 356 | + | ||
| 357 | + tools = [ | ||
| 358 | + { | ||
| 359 | + "type": "function", | ||
| 360 | + "function": { | ||
| 361 | + "name": "local_time", | ||
| 362 | + "parameters": { | ||
| 363 | + "type": "object", | ||
| 364 | + "properties": {"city": {"type": "string"}}, | ||
| 365 | + }, | ||
| 366 | + }, | ||
| 367 | + } | ||
| 368 | + ] | ||
| 369 | + req_info = RequestInfo( | ||
| 370 | + req_id="contract-sglang-stream", | ||
| 371 | + req_data={ | ||
| 372 | + "model": "test-model", | ||
| 373 | + "messages": [{"role": "user", "content": "Time?"}], | ||
| 374 | + "tools": tools, | ||
| 375 | + "logprobs": True, | ||
| 376 | + "top_logprobs": 2, | ||
| 377 | + "max_tokens": 5, | ||
| 378 | + "stream": True, | ||
| 379 | + "stream_options": {"include_usage": True}, | ||
| 380 | + }, | ||
| 381 | + api="v1/chat/completions", | ||
| 382 | + entry_api="v1/chat/completions", | ||
| 383 | + req_len=7, | ||
| 384 | + ) | ||
| 385 | + scheduler = _ContractScheduler("sglang") | ||
| 386 | + router = _router(req_info, scheduler) | ||
| 387 | + | ||
| 388 | + async with _bind_native_apps(monkeypatch, router, p_app, d_app): | ||
| 389 | + response = await router.handle_request() | ||
| 390 | + body = await _collect_stream_body(response) | ||
| 391 | + | ||
| 392 | + assert p_requests[0]["rid"] == d_requests[0]["rid"] == "contract-sglang-stream#a1" | ||
| 393 | + assert p_requests[0]["bootstrap_room"] == d_requests[0]["bootstrap_room"] | ||
| 394 | + assert p_requests[0]["bootstrap_host"] == d_requests[0]["bootstrap_host"] == "127.0.0.1" | ||
| 395 | + assert p_requests[0]["bootstrap_port"] == d_requests[0]["bootstrap_port"] == 21001 | ||
| 396 | + assert p_requests[0]["tools"] == d_requests[0]["tools"] == tools | ||
| 397 | + assert b'"tool_calls"' in body | ||
| 398 | + assert b'"logprobs"' in body | ||
| 399 | + assert b'"prompt_tokens_details":{"cached_tokens":3}' in body | ||
| 400 | + assert b"token_ids" not in body | ||
| 401 | + assert b"bootstrap_host" not in body | ||
| 402 | + assert b"bootstrap_port" not in body | ||
| 403 | + assert b"bootstrap_room" not in body | ||
| 404 | + assert body.count(b"data: [DONE]") == 1 | ||
| 405 | + assert req_info.state == ReqState.DECODE_END | ||
| 406 | + assert scheduler.cb_events == [(1, "success"), (2, "success")] | ||
| 407 | + | ||
| 408 | + | ||
| 409 | + | ||
| 410 | + ("status_code", "expected_cb_event"), | ||
| 411 | + [ | ||
| 412 | + (400, None), | ||
| 413 | + (503, (1, "failure")), | ||
| 414 | + ], | ||
| 415 | +) | ||
| 416 | + | ||
| 417 | +async def test_vllm_prefill_http_error_preserves_status_and_circuit_breaker_semantics( | ||
| 418 | + monkeypatch, | ||
| 419 | + status_code, | ||
| 420 | + expected_cb_event, | ||
| 421 | +): | ||
| 422 | + p_app = FastAPI() | ||
| 423 | + d_app = FastAPI() | ||
| 424 | + d_requests = [] | ||
| 425 | + | ||
| 426 | + | ||
| 427 | + async def prefill(_request: Request): | ||
| 428 | + return JSONResponse( | ||
| 429 | + status_code=status_code, | ||
| 430 | + content={"error": {"message": f"native prefill {status_code}"}}, | ||
| 431 | + headers={"retry-after": "2"}, | ||
| 432 | + ) | ||
| 433 | + | ||
| 434 | + | ||
| 435 | + async def decode(request: Request): | ||
| 436 | + d_requests.append(await request.json()) | ||
| 437 | + return {"choices": [{"text": "must not run"}]} | ||
| 438 | + | ||
| 439 | + req_info = RequestInfo( | ||
| 440 | + req_id=f"contract-vllm-error-{status_code}", | ||
| 441 | + req_data={ | ||
| 442 | + "model": "test-model", | ||
| 443 | + "prompt": "hello", | ||
| 444 | + "max_tokens": 4, | ||
| 445 | + "stream": False, | ||
| 446 | + }, | ||
| 447 | + api="v1/completions", | ||
| 448 | + entry_api="v1/completions", | ||
| 449 | + req_len=3, | ||
| 450 | + ) | ||
| 451 | + scheduler = _ContractScheduler("vllm") | ||
| 452 | + router = _router(req_info, scheduler) | ||
| 453 | + | ||
| 454 | + async with _bind_native_apps(monkeypatch, router, p_app, d_app): | ||
| 455 | + with pytest.raises(UpstreamHTTPError) as exc_info: | ||
| 456 | + await router.handle_request() | ||
| 457 | + | ||
| 458 | + assert exc_info.value.status_code == status_code | ||
| 459 | + assert exc_info.value.headers["retry-after"] == "2" | ||
| 460 | + assert d_requests == [] | ||
| 461 | + assert scheduler.cb_events == ([] if expected_cb_event is None else [expected_cb_event]) | ||
| 462 | + | ||
| 463 | + | ||
| 464 | + | ||
| 465 | +async def test_native_pd_client_uses_coordinator_inference_tls_config(monkeypatch): | ||
| 466 | + req_info = RequestInfo( | ||
| 467 | + req_id="contract-native-tls", | ||
| 468 | + req_data={"model": "test-model", "prompt": "hello", "max_tokens": 1}, | ||
| 469 | + api="v1/completions", | ||
| 470 | + entry_api="v1/completions", | ||
| 471 | + req_len=1, | ||
| 472 | + ) | ||
| 473 | + scheduler = _ContractScheduler("vllm") | ||
| 474 | + config = _config() | ||
| 475 | + config.infer_tls_config = TLSConfig( | ||
| 476 | + enable_tls=True, | ||
| 477 | + ca_file="/certs/ca.pem", | ||
| 478 | + cert_file="/certs/client.pem", | ||
| 479 | + key_file="/certs/client.key", | ||
| 480 | + ) | ||
| 481 | + router = UnifiedPDRouter( | ||
| 482 | + req_info, | ||
| 483 | + config, | ||
| 484 | + scheduler=scheduler, | ||
| 485 | + request_manager=RequestManager(config), | ||
| 486 | + ) | ||
| 487 | + endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) | ||
| 488 | + resource = ScheduledResource(instance=scheduler.p, endpoint=endpoint) | ||
| 489 | + fake_client = object() | ||
| 490 | + pool = HTTPClientPool() | ||
| 491 | + get_client = AsyncMock(return_value=fake_client) | ||
| 492 | + monkeypatch.setattr(pool, "get_client", get_client) | ||
| 493 | + | ||
| 494 | + async with router._client_for(resource) as client: | ||
| 495 | + assert client is fake_client | ||
| 496 | + | ||
| 497 | + get_client.assert_awaited_once_with( | ||
| 498 | + ip=endpoint.ip, | ||
| 499 | + port=endpoint.business_port, | ||
| 500 | + tls_config=config.infer_tls_config, | ||
| 501 | + ) | ||
| @@ -0,0 +1,271 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from copy import deepcopy | ||
| 12 | + | ||
| 13 | +import pytest | ||
| 14 | + | ||
| 15 | +from motor.coordinator.router.adapters.pd_protocol import ( | ||
| 16 | + ADAPTERS, | ||
| 17 | + CoordinationMode, | ||
| 18 | + EngineEndpointMetadata, | ||
| 19 | + EngineProtocolError, | ||
| 20 | + LegContext, | ||
| 21 | + PrefillMetadata, | ||
| 22 | + SglangProtocolAdapter, | ||
| 23 | + VllmProtocolAdapter, | ||
| 24 | +) | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +def _endpoint(host: str, port: int | None = None) -> EngineEndpointMetadata: | ||
| 28 | + return EngineEndpointMetadata( | ||
| 29 | + host=host, | ||
| 30 | + bootstrap_port=port, | ||
| 31 | + ) | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def _context( | ||
| 35 | + *, | ||
| 36 | + endpoint: EngineEndpointMetadata | None = None, | ||
| 37 | + peer_endpoint: EngineEndpointMetadata | None = None, | ||
| 38 | + attempt_seq: int = 2, | ||
| 39 | +) -> LegContext: | ||
| 40 | + return LegContext( | ||
| 41 | + engine_request_id="engine-1", | ||
| 42 | + pair_id="pair-1", | ||
| 43 | + attempt_seq=attempt_seq, | ||
| 44 | + api="/v1/chat/completions", | ||
| 45 | + endpoint=endpoint or _endpoint("prefill.local", 8998), | ||
| 46 | + peer_endpoint=peer_endpoint, | ||
| 47 | + ) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +def test_adapter_registry_is_static(): | ||
| 51 | + assert ADAPTERS["vllm"].coordination_mode is CoordinationMode.HANDOFF | ||
| 52 | + assert ADAPTERS["sglang"].coordination_mode is CoordinationMode.BOOTSTRAP | ||
| 53 | + | ||
| 54 | + with pytest.raises(TypeError): | ||
| 55 | + ADAPTERS["other"] = VllmProtocolAdapter() | ||
| 56 | + | ||
| 57 | + | ||
| 58 | +def test_vllm_prefill_request_matches_native_handoff_contract_without_mutation(): | ||
| 59 | + adapter = VllmProtocolAdapter() | ||
| 60 | + request = { | ||
| 61 | + "model": "glm", | ||
| 62 | + "messages": [{"role": "user", "content": "hello"}], | ||
| 63 | + "stream": True, | ||
| 64 | + "stream_options": {"include_usage": True}, | ||
| 65 | + "max_tokens": 128, | ||
| 66 | + "max_completion_tokens": 96, | ||
| 67 | + "request_id": "client-request-id", | ||
| 68 | + "rid": "wrong-engine-id", | ||
| 69 | + "extra": {"nested": [1, 2]}, | ||
| 70 | + } | ||
| 71 | + original = deepcopy(request) | ||
| 72 | + | ||
| 73 | + engine_request = adapter.build_prefill_request(request, _context()) | ||
| 74 | + | ||
| 75 | + assert request == original | ||
| 76 | + assert engine_request.api == "/v1/chat/completions" | ||
| 77 | + assert engine_request.body["request_id"] == "engine-1" | ||
| 78 | + assert "rid" not in engine_request.body | ||
| 79 | + assert engine_request.body["stream"] is False | ||
| 80 | + assert engine_request.body["max_tokens"] == 1 | ||
| 81 | + assert engine_request.body["max_completion_tokens"] == 1 | ||
| 82 | + assert engine_request.body["min_tokens"] == 1 | ||
| 83 | + assert "stream_options" not in engine_request.body | ||
| 84 | + assert engine_request.body["kv_transfer_params"] == { | ||
| 85 | + "do_remote_decode": True, | ||
| 86 | + "do_remote_prefill": False, | ||
| 87 | + "remote_engine_id": None, | ||
| 88 | + "remote_block_ids": None, | ||
| 89 | + "remote_host": None, | ||
| 90 | + "remote_port": None, | ||
| 91 | + } | ||
| 92 | + engine_request.body["extra"]["nested"].append(3) | ||
| 93 | + assert request == original | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +def test_vllm_prefill_does_not_add_max_completion_tokens(): | ||
| 97 | + request = {"max_tokens": 32} | ||
| 98 | + | ||
| 99 | + body = VllmProtocolAdapter().build_prefill_request(request, _context()).body | ||
| 100 | + | ||
| 101 | + assert "max_completion_tokens" not in body | ||
| 102 | + | ||
| 103 | + | ||
| 104 | +def test_vllm_prefill_response_returns_copied_ticket_and_usage(): | ||
| 105 | + response = { | ||
| 106 | + "kv_transfer_params": { | ||
| 107 | + "do_remote_prefill": True, | ||
| 108 | + "remote_host": "10.0.0.1", | ||
| 109 | + "connector_private": {"blocks": [1, 2]}, | ||
| 110 | + }, | ||
| 111 | + "usage": {"prompt_tokens": 8, "details": {"cached_tokens": 4}}, | ||
| 112 | + } | ||
| 113 | + | ||
| 114 | + metadata = VllmProtocolAdapter().parse_prefill_response(response) | ||
| 115 | + | ||
| 116 | + assert metadata.handoff_ticket == response["kv_transfer_params"] | ||
| 117 | + assert metadata.usage == response["usage"] | ||
| 118 | + response["kv_transfer_params"]["connector_private"]["blocks"].append(3) | ||
| 119 | + response["usage"]["details"]["cached_tokens"] = 0 | ||
| 120 | + assert metadata.handoff_ticket["connector_private"]["blocks"] == [1, 2] | ||
| 121 | + assert metadata.usage["details"]["cached_tokens"] == 4 | ||
| 122 | + | ||
| 123 | + | ||
| 124 | + | ||
| 125 | + ("response", "message"), | ||
| 126 | + [ | ||
| 127 | + ({}, "Missing kv_transfer_params"), | ||
| 128 | + ({"kv_transfer_params": {}}, "Missing kv_transfer_params"), | ||
| 129 | + ({"kv_transfer_params": {"do_remote_prefill": False}}, "do_remote_prefill must be true"), | ||
| 130 | + ], | ||
| 131 | +) | ||
| 132 | +def test_vllm_prefill_response_rejects_invalid_handoff(response, message): | ||
| 133 | + with pytest.raises(EngineProtocolError, match=message) as exc_info: | ||
| 134 | + VllmProtocolAdapter().parse_prefill_response(response) | ||
| 135 | + | ||
| 136 | + assert exc_info.value.engine_type == "vllm" | ||
| 137 | + assert exc_info.value.phase == "prefill" | ||
| 138 | + | ||
| 139 | + | ||
| 140 | +def test_vllm_decode_restores_original_generation_budget_and_copies_ticket(): | ||
| 141 | + adapter = VllmProtocolAdapter() | ||
| 142 | + request = { | ||
| 143 | + "stream": True, | ||
| 144 | + "max_tokens": 128, | ||
| 145 | + "max_completion_tokens": 96, | ||
| 146 | + "extra": {"nested": [1]}, | ||
| 147 | + } | ||
| 148 | + original = deepcopy(request) | ||
| 149 | + ticket = {"do_remote_prefill": True, "remote_host": "10.0.0.1", "private": {"blocks": [7]}} | ||
| 150 | + | ||
| 151 | + engine_request = adapter.build_decode_request( | ||
| 152 | + request, | ||
| 153 | + _context(), | ||
| 154 | + PrefillMetadata(handoff_ticket=ticket), | ||
| 155 | + ) | ||
| 156 | + | ||
| 157 | + assert request == original | ||
| 158 | + assert engine_request.body["request_id"] == "engine-1" | ||
| 159 | + assert engine_request.body["stream"] is True | ||
| 160 | + assert engine_request.body["max_tokens"] == 128 | ||
| 161 | + assert engine_request.body["max_completion_tokens"] == 96 | ||
| 162 | + assert engine_request.body["kv_transfer_params"] == ticket | ||
| 163 | + engine_request.body["kv_transfer_params"]["private"]["blocks"].append(8) | ||
| 164 | + assert ticket["private"]["blocks"] == [7] | ||
| 165 | + | ||
| 166 | + | ||
| 167 | + | ||
| 168 | +def test_vllm_decode_requires_handoff_ticket(metadata): | ||
| 169 | + with pytest.raises(EngineProtocolError, match="Missing handoff ticket") as exc_info: | ||
| 170 | + VllmProtocolAdapter().build_decode_request({}, _context(), metadata) | ||
| 171 | + | ||
| 172 | + assert exc_info.value.phase == "decode" | ||
| 173 | + | ||
| 174 | + | ||
| 175 | +def test_vllm_declares_internal_response_field(): | ||
| 176 | + assert VllmProtocolAdapter.internal_response_fields == frozenset({"kv_transfer_params"}) | ||
| 177 | + | ||
| 178 | + | ||
| 179 | +def test_sglang_prefill_and_decode_share_prefill_bootstrap_address_and_room(): | ||
| 180 | + adapter = SglangProtocolAdapter() | ||
| 181 | + prefill_endpoint = _endpoint("prefill.local", 8998) | ||
| 182 | + decode_endpoint = _endpoint("decode.local", 8999) | ||
| 183 | + request = { | ||
| 184 | + "model": "glm", | ||
| 185 | + "messages": [{"role": "user", "content": "hello"}], | ||
| 186 | + "request_id": "wrong-engine-id", | ||
| 187 | + "rid": "client-request-id", | ||
| 188 | + } | ||
| 189 | + original = deepcopy(request) | ||
| 190 | + prefill_context = _context(endpoint=prefill_endpoint, peer_endpoint=decode_endpoint) | ||
| 191 | + decode_context = _context(endpoint=decode_endpoint, peer_endpoint=prefill_endpoint) | ||
| 192 | + | ||
| 193 | + prefill = adapter.build_prefill_request(request, prefill_context) | ||
| 194 | + decode = adapter.build_decode_request(request, decode_context, None) | ||
| 195 | + | ||
| 196 | + assert request == original | ||
| 197 | + assert prefill.body["rid"] == "engine-1" | ||
| 198 | + assert decode.body["rid"] == "engine-1" | ||
| 199 | + assert "request_id" not in prefill.body | ||
| 200 | + assert "request_id" not in decode.body | ||
| 201 | + assert prefill.body["stream"] is False | ||
| 202 | + for body in (prefill.body, decode.body): | ||
| 203 | + assert body["bootstrap_host"] == "prefill.local" | ||
| 204 | + assert body["bootstrap_port"] == 8998 | ||
| 205 | + assert prefill.body["bootstrap_room"] == decode.body["bootstrap_room"] | ||
| 206 | + prefill.body["messages"][0]["content"] = "changed" | ||
| 207 | + assert request == original | ||
| 208 | + | ||
| 209 | + | ||
| 210 | +def test_sglang_bootstrap_room_is_stable_per_attempt(): | ||
| 211 | + adapter = SglangProtocolAdapter() | ||
| 212 | + | ||
| 213 | + first = adapter.build_prefill_request({}, _context(attempt_seq=2)).body["bootstrap_room"] | ||
| 214 | + repeated = adapter.build_prefill_request({}, _context(attempt_seq=2)).body["bootstrap_room"] | ||
| 215 | + retried = adapter.build_prefill_request({}, _context(attempt_seq=3)).body["bootstrap_room"] | ||
| 216 | + | ||
| 217 | + assert first == repeated | ||
| 218 | + assert first != retried | ||
| 219 | + assert 0 <= first < 1 << 63 | ||
| 220 | + | ||
| 221 | + | ||
| 222 | + | ||
| 223 | + "endpoint", | ||
| 224 | + [ | ||
| 225 | + EngineEndpointMetadata(host="", bootstrap_port=8998), | ||
| 226 | + _endpoint("prefill.local"), | ||
| 227 | + _endpoint("prefill.local", 0), | ||
| 228 | + _endpoint("prefill.local", 65536), | ||
| 229 | + ], | ||
| 230 | +) | ||
| 231 | +def test_sglang_prefill_rejects_invalid_bootstrap_endpoint(endpoint): | ||
| 232 | + with pytest.raises(EngineProtocolError) as exc_info: | ||
| 233 | + SglangProtocolAdapter().build_prefill_request({}, _context(endpoint=endpoint)) | ||
| 234 | + | ||
| 235 | + assert exc_info.value.engine_type == "sglang" | ||
| 236 | + assert exc_info.value.phase == "prefill" | ||
| 237 | + | ||
| 238 | + | ||
| 239 | +def test_sglang_decode_requires_prefill_endpoint(): | ||
| 240 | + with pytest.raises(EngineProtocolError, match="Missing prefill endpoint metadata") as exc_info: | ||
| 241 | + SglangProtocolAdapter().build_decode_request({}, _context(), None) | ||
| 242 | + | ||
| 243 | + assert exc_info.value.phase == "decode" | ||
| 244 | + | ||
| 245 | + | ||
| 246 | +def test_sglang_prefill_response_copies_usage_without_handoff_ticket(): | ||
| 247 | + response = {"usage": {"prompt_tokens": 8, "details": {"cached_tokens": 4}}} | ||
| 248 | + | ||
| 249 | + metadata = SglangProtocolAdapter().parse_prefill_response(response) | ||
| 250 | + | ||
| 251 | + assert metadata.handoff_ticket is None | ||
| 252 | + assert metadata.usage == response["usage"] | ||
| 253 | + response["usage"]["details"]["cached_tokens"] = 0 | ||
| 254 | + assert metadata.usage["details"]["cached_tokens"] == 4 | ||
| 255 | + | ||
| 256 | + | ||
| 257 | +def test_sglang_declares_internal_response_fields(): | ||
| 258 | + assert SglangProtocolAdapter.internal_response_fields == frozenset( | ||
| 259 | + {"bootstrap_host", "bootstrap_port", "bootstrap_room"} | ||
| 260 | + ) | ||
| 261 | + | ||
| 262 | + | ||
| 263 | +def test_sglang_abort_uses_native_request_id(): | ||
| 264 | + request = SglangProtocolAdapter().build_abort_request(_context()) | ||
| 265 | + | ||
| 266 | + assert request.api == "abort_request" | ||
| 267 | + assert request.body == {"rid": "engine-1"} | ||
| 268 | + | ||
| 269 | + | ||
| 270 | +def test_vllm_does_not_claim_an_unverified_abort_endpoint(): | ||
| 271 | + assert VllmProtocolAdapter().build_abort_request(_context()) is None | ||
Rtests/coordinator/router/test_router_cdp_separation.py→tests/coordinator/router/test_router_native_handoff.py+957-1119
| @@ -1,1119 +1,957 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 2 | -# MindIE is licensed under Mulan PSL v2. | 2 | +# MindIE is licensed under Mulan PSL v2. |
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. |
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | 4 | +# You may obtain a copy of Mulan PSL v2 at: |
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | 5 | +# http://license.coscl.org.cn/MulanPSL2 |
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, |
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, |
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | -# See the Mulan PSL v2 for more details. | 9 | +# See the Mulan PSL v2 for more details. |
| 10 | - | 10 | + |
| 11 | -from pytest import MonkeyPatch | 11 | +from pytest import MonkeyPatch |
| 12 | -from fastapi import FastAPI, status, Request | 12 | +from fastapi import FastAPI, status, Request |
| 13 | -from fastapi.responses import JSONResponse | 13 | +from unittest.mock import patch, MagicMock, AsyncMock |
| 14 | -from unittest.mock import patch, MagicMock, AsyncMock | 14 | +from fastapi.testclient import TestClient |
| 15 | -from fastapi.testclient import TestClient | 15 | +import asyncio |
| 16 | -import asyncio | 16 | +from contextlib import asynccontextmanager |
| 17 | -from contextlib import asynccontextmanager | 17 | +import httpx |
| 18 | -import httpx | 18 | +import json |
| 19 | -import json | 19 | +import pytest |
| 20 | -import pytest | 20 | + |
| 21 | - | 21 | +from motor.common.resources.endpoint import ( |
| 22 | -from motor.common.resources.dispatch import MOTOR_DISPATCH_KEY | 22 | + Endpoint, |
| 23 | -from motor.common.resources.endpoint import ( | 23 | + EndpointStatus, |
| 24 | - Endpoint, | 24 | + Workload, |
| 25 | - EndpointStatus, | 25 | + WorkloadAction, |
| 26 | - Workload, | 26 | +) |
| 27 | - WorkloadAction, | 27 | +from motor.common.resources.instance import PDRole, Instance, InsStatus, ParallelConfig |
| 28 | -) | 28 | +from motor.config.coordinator import CoordinatorConfig, ExceptionConfig, SchedulerType |
| 29 | -from motor.common.resources.dispatch import DispatchPlan | 29 | +from motor.coordinator.domain.instance_manager import InstanceManager |
| 30 | -from motor.common.resources.instance import PDRole, Instance, InsStatus, ParallelConfig | 30 | +from motor.coordinator.domain import InstanceReadiness, ScheduledResource |
| 31 | -from motor.config.coordinator import CoordinatorConfig, ExceptionConfig, SchedulerType | 31 | +from motor.coordinator.models.request import ReqState, RequestInfo |
| 32 | -from motor.coordinator.domain.instance_manager import InstanceManager | 32 | +from motor.coordinator.router.strategies.unified_pd import UnifiedPDRouter |
| 33 | -from motor.coordinator.domain import InstanceReadiness, ScheduledResource | 33 | +from motor.coordinator.tracer.tracing import TracerManager |
| 34 | -from motor.coordinator.models.request import ReqState, RequestInfo | 34 | +from motor.coordinator.scheduler.scheduler import Scheduler |
| 35 | -from motor.coordinator.router.strategies.unified_pd import ( | 35 | +from motor.coordinator.domain.request_manager import RequestManager |
| 36 | - UnifiedPDRouter as SeparateCDPRouter, | 36 | +from tests.coordinator.router.mock_openai_request import create_mock_request_info |
| 37 | -) | 37 | +import motor.coordinator.router.dispatch as router |
| 38 | -from motor.coordinator.tracer.tracing import TracerManager | 38 | + |
| 39 | -from motor.coordinator.scheduler.scheduler import Scheduler | 39 | +TracerManager() |
| 40 | -from motor.coordinator.domain.request_manager import RequestManager | 40 | + |
| 41 | -from tests.coordinator.router.mock_openai_request import ( | 41 | + |
| 42 | - MockStreamResponse, | 42 | +def _native_prefill_body(request_body, *, usage=None): |
| 43 | - create_mock_request_info, | 43 | + body = { |
| 44 | -) | 44 | + "kv_transfer_params": { |
| 45 | -import motor.coordinator.router.dispatch as router | 45 | + "do_remote_prefill": True, |
| 46 | - | 46 | + "remote_request_id": request_body["request_id"], |
| 47 | -TracerManager() | 47 | + "remote_host": "127.0.0.1", |
| 48 | - | 48 | + "remote_port": 9000, |
| 49 | - | 49 | + } |
| 50 | -class _UnifiedPDPrefillClient: | 50 | + } |
| 51 | - def __init__(self, *, exc: Exception | None = None, post_fail_times: int = 0): | 51 | + if usage is not None: |
| 52 | - self.exc = exc | 52 | + body["usage"] = usage |
| 53 | - self.post_fail_times = post_fail_times | 53 | + return body |
| 54 | - self.post_fail_count = 0 | 54 | + |
| 55 | - self.requests = [] | 55 | + |
| 56 | - self.base_url = "http://prefill" | 56 | +class _UnifiedPDPrefillClient: |
| 57 | - self.timeout = 1 | 57 | + def __init__(self, *, exc: Exception | None = None, post_fail_times: int = 0): |
| 58 | - | 58 | + self.exc = exc |
| 59 | - async def post(self, path, json=None, headers=None, timeout=None): | 59 | + self.post_fail_times = post_fail_times |
| 60 | - self.requests.append(json) | 60 | + self.post_fail_count = 0 |
| 61 | - if self.exc is not None and (self.post_fail_times == 0 or self.post_fail_count < self.post_fail_times): | 61 | + self.requests = [] |
| 62 | - self.post_fail_count += 1 | 62 | + self.base_url = "http://prefill" |
| 63 | - raise self.exc | 63 | + self.timeout = 1 |
| 64 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | 64 | + |
| 65 | - return httpx.Response(status_code=200, json={"status": "cached"}, request=request) | 65 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 66 | - | 66 | + self.requests.append(json) |
| 67 | - | 67 | + if self.exc is not None and (self.post_fail_times == 0 or self.post_fail_count < self.post_fail_times): |
| 68 | -class _UnifiedPDStreamResponse: | 68 | + self.post_fail_count += 1 |
| 69 | - def __init__(self, chunks, exc: Exception | None = None): | 69 | + raise self.exc |
| 70 | - self.chunks = list(chunks) | 70 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 71 | - self.exc = exc | 71 | + return httpx.Response( |
| 72 | - if isinstance(exc, httpx.HTTPStatusError): | 72 | + status_code=200, |
| 73 | - self.status_code = exc.response.status_code | 73 | + json=_native_prefill_body(json), |
| 74 | - self.is_success = False | 74 | + request=request, |
| 75 | - self.text = exc.response.text | 75 | + ) |
| 76 | - self.headers = exc.response.headers | 76 | + |
| 77 | - self._error_body = exc.response.content or str(exc).encode() | 77 | + |
| 78 | - else: | 78 | +class _UnifiedPDStreamResponse: |
| 79 | - self.status_code = 200 | 79 | + def __init__(self, chunks, exc: Exception | None = None): |
| 80 | - self.is_success = True | 80 | + self.chunks = list(chunks) |
| 81 | - self.text = "" | 81 | + self.exc = exc |
| 82 | - self.headers = {} | 82 | + if isinstance(exc, httpx.HTTPStatusError): |
| 83 | - self._error_body = b"" | 83 | + self.status_code = exc.response.status_code |
| 84 | - | 84 | + self.is_success = False |
| 85 | - async def __aenter__(self): | 85 | + self.text = exc.response.text |
| 86 | - if isinstance(self.exc, httpx.RequestError): | 86 | + self.headers = exc.response.headers |
| 87 | - raise self.exc | 87 | + self._error_body = exc.response.content or str(exc).encode() |
| 88 | - return self | 88 | + else: |
| 89 | - | 89 | + self.status_code = 200 |
| 90 | - async def __aexit__(self, exc_type, exc_val, exc_tb): | 90 | + self.is_success = True |
| 91 | - return False | 91 | + self.text = "" |
| 92 | - | 92 | + self.headers = {} |
| 93 | - async def aread(self): | 93 | + self._error_body = b"" |
| 94 | - return self._error_body | 94 | + |
| 95 | - | 95 | + async def __aenter__(self): |
| 96 | - def raise_for_status(self): | 96 | + if isinstance(self.exc, httpx.RequestError): |
| 97 | - if self.exc is not None and not self.chunks: | 97 | + raise self.exc |
| 98 | - raise self.exc | 98 | + return self |
| 99 | - | 99 | + |
| 100 | - async def aiter_bytes(self): | 100 | + async def __aexit__(self, exc_type, exc_val, exc_tb): |
| 101 | - if not self.is_success: | 101 | + return False |
| 102 | - yield self._error_body | 102 | + |
| 103 | - return | 103 | + async def aread(self): |
| 104 | - for chunk in self.chunks: | 104 | + return self._error_body |
| 105 | - yield chunk | 105 | + |
| 106 | - if self.exc is not None and not isinstance(self.exc, (httpx.HTTPStatusError, httpx.RequestError)): | 106 | + def raise_for_status(self): |
| 107 | - raise self.exc | 107 | + if self.exc is not None and not self.chunks: |
| 108 | - | 108 | + raise self.exc |
| 109 | - | 109 | + |
| 110 | -class _UnifiedPDDecodeClient: | 110 | + async def aiter_bytes(self): |
| 111 | - def __init__( | 111 | + if not self.is_success: |
| 112 | - self, | 112 | + yield self._error_body |
| 113 | - *, | 113 | + return |
| 114 | - stream_chunks=None, | 114 | + for chunk in self.chunks: |
| 115 | - stream_exc: Exception | None = None, | 115 | + yield chunk |
| 116 | - stream_fail_times: int = 0, | 116 | + if self.exc is not None and not isinstance(self.exc, (httpx.HTTPStatusError, httpx.RequestError)): |
| 117 | - post_exc: Exception | None = None, | 117 | + raise self.exc |
| 118 | - post_fail_times: int = 0, | 118 | + |
| 119 | - ): | 119 | + |
| 120 | - self.stream_chunks = stream_chunks or [ | 120 | +class _UnifiedPDDecodeClient: |
| 121 | - b'data: {"choices":[{"delta":{"content":"decoded chunk"},"index":0,"finish_reason":null}]}\n\n', | 121 | + def __init__( |
| 122 | - ] | 122 | + self, |
| 123 | - self.stream_exc = stream_exc | 123 | + *, |
| 124 | - self.stream_fail_times = stream_fail_times | 124 | + stream_chunks=None, |
| 125 | - self.post_exc = post_exc | 125 | + stream_exc: Exception | None = None, |
| 126 | - self.post_fail_times = post_fail_times | 126 | + stream_fail_times: int = 0, |
| 127 | - self.requests = [] | 127 | + post_exc: Exception | None = None, |
| 128 | - self.stream_count = 0 | 128 | + post_fail_times: int = 0, |
| 129 | - self.stream_fail_count = 0 | 129 | + ): |
| 130 | - self.post_count = 0 | 130 | + self.stream_chunks = stream_chunks or [ |
| 131 | - self.post_fail_count = 0 | 131 | + b'data: {"choices":[{"delta":{"content":"decoded chunk"},"index":0,"finish_reason":null}]}\n\n', |
| 132 | - self.base_url = "http://decode" | 132 | + ] |
| 133 | - self.timeout = 1 | 133 | + self.stream_exc = stream_exc |
| 134 | - | 134 | + self.stream_fail_times = stream_fail_times |
| 135 | - def stream(self, method, url, json=None, headers=None, timeout=None): | 135 | + self.post_exc = post_exc |
| 136 | - self.stream_count += 1 | 136 | + self.post_fail_times = post_fail_times |
| 137 | - if json: | 137 | + self.requests = [] |
| 138 | - self.requests.append(json) | 138 | + self.stream_count = 0 |
| 139 | - if self.stream_exc is not None and ( | 139 | + self.stream_fail_count = 0 |
| 140 | - self.stream_fail_times == 0 or self.stream_fail_count < self.stream_fail_times | 140 | + self.post_count = 0 |
| 141 | - ): | 141 | + self.post_fail_count = 0 |
| 142 | - self.stream_fail_count += 1 | 142 | + self.base_url = "http://decode" |
| 143 | - return _UnifiedPDStreamResponse([], exc=self.stream_exc) | 143 | + self.timeout = 1 |
| 144 | - return _UnifiedPDStreamResponse(self.stream_chunks) | 144 | + |
| 145 | - | 145 | + def stream(self, method, url, json=None, headers=None, timeout=None): |
| 146 | - async def post(self, path, json=None, headers=None, timeout=None): | 146 | + self.stream_count += 1 |
| 147 | - self.post_count += 1 | 147 | + if json: |
| 148 | - if json: | 148 | + self.requests.append(json) |
| 149 | - self.requests.append(json) | 149 | + if self.stream_exc is not None and ( |
| 150 | - if self.post_exc is not None and (self.post_fail_times == 0 or self.post_fail_count < self.post_fail_times): | 150 | + self.stream_fail_times == 0 or self.stream_fail_count < self.stream_fail_times |
| 151 | - self.post_fail_count += 1 | 151 | + ): |
| 152 | - raise self.post_exc | 152 | + self.stream_fail_count += 1 |
| 153 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | 153 | + return _UnifiedPDStreamResponse([], exc=self.stream_exc) |
| 154 | - return httpx.Response( | 154 | + return _UnifiedPDStreamResponse(self.stream_chunks) |
| 155 | - status_code=200, | 155 | + |
| 156 | - json={"choices": [{"message": {"content": "test response"}}]}, | 156 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 157 | - request=request, | 157 | + self.post_count += 1 |
| 158 | - ) | 158 | + if json: |
| 159 | - | 159 | + self.requests.append(json) |
| 160 | - | 160 | + if self.post_exc is not None and (self.post_fail_times == 0 or self.post_fail_count < self.post_fail_times): |
| 161 | -def _patch_unified_pd_clients(monkeypatch, router_obj, p_client, d_client): | 161 | + self.post_fail_count += 1 |
| 162 | - @asynccontextmanager | 162 | + raise self.post_exc |
| 163 | - async def _client_for(resource: ScheduledResource): | 163 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 164 | - if resource.instance.role == PDRole.ROLE_P: | 164 | + return httpx.Response( |
| 165 | - yield p_client | 165 | + status_code=200, |
| 166 | - else: | 166 | + json={"choices": [{"message": {"content": "test response"}}]}, |
| 167 | - yield d_client | 167 | + request=request, |
| 168 | - | 168 | + ) |
| 169 | - monkeypatch.setattr(router_obj, "_client_for", _client_for) | 169 | + |
| 170 | - | 170 | + |
| 171 | - | 171 | +def _patch_unified_pd_clients(monkeypatch, router_obj, p_client, d_client): |
| 172 | -def _patch_unified_pd_router_clients(monkeypatch, p_client, d_client): | 172 | + @asynccontextmanager |
| 173 | - """Patch UnifiedPDRouter._client_for at class level (for app-level integration tests).""" | 173 | + async def _client_for(resource: ScheduledResource): |
| 174 | - | 174 | + if resource.instance.role == PDRole.ROLE_P: |
| 175 | - def _client_for(self, resource: ScheduledResource): | 175 | + yield p_client |
| 176 | - @asynccontextmanager | 176 | + else: |
| 177 | - async def _cm(): | 177 | + yield d_client |
| 178 | - if resource.instance.role == PDRole.ROLE_P: | 178 | + |
| 179 | - yield p_client | 179 | + monkeypatch.setattr(router_obj, "_client_for", _client_for) |
| 180 | - else: | 180 | + |
| 181 | - yield d_client | 181 | + |
| 182 | - | 182 | +def _patch_unified_pd_router_clients(monkeypatch, p_client, d_client): |
| 183 | - return _cm() | 183 | + """Patch UnifiedPDRouter._client_for at class level (for app-level integration tests).""" |
| 184 | - | 184 | + |
| 185 | - monkeypatch.setattr(SeparateCDPRouter, "_client_for", _client_for) | 185 | + def _client_for(self, resource: ScheduledResource): |
| 186 | - | 186 | + @asynccontextmanager |
| 187 | - | 187 | + async def _cm(): |
| 188 | -async def _collect_stream_chunks(response) -> str: | 188 | + if resource.instance.role == PDRole.ROLE_P: |
| 189 | - chunks = [] | 189 | + yield p_client |
| 190 | - async for chunk in response.body_iterator: | 190 | + else: |
| 191 | - if isinstance(chunk, bytes): | 191 | + yield d_client |
| 192 | - chunks.append(chunk.decode("utf-8", errors="replace")) | 192 | + |
| 193 | - else: | 193 | + return _cm() |
| 194 | - chunks.append(chunk) | 194 | + |
| 195 | - return "".join(chunks) | 195 | + monkeypatch.setattr(UnifiedPDRouter, "_client_for", _client_for) |
| 196 | - | 196 | + |
| 197 | - | 197 | + |
| 198 | -def _parse_stream_error_payload(chunk_str: str) -> dict: | 198 | +async def _collect_stream_chunks(response) -> str: |
| 199 | - data_lines = [ | 199 | + chunks = [] |
| 200 | - line.removeprefix("data: ").strip() | 200 | + async for chunk in response.body_iterator: |
| 201 | - for line in chunk_str.splitlines() | 201 | + if isinstance(chunk, bytes): |
| 202 | - if line.startswith("data: ") and line.strip() != "data: [DONE]" | 202 | + chunks.append(chunk.decode("utf-8", errors="replace")) |
| 203 | - ] | 203 | + else: |
| 204 | - assert data_lines, f"expected streaming error chunk, got: {chunk_str!r}" | 204 | + chunks.append(chunk) |
| 205 | - return json.loads(data_lines[-1]) | 205 | + return "".join(chunks) |
| 206 | - | 206 | + |
| 207 | - | 207 | + |
| 208 | -def _assert_stream_error_chunk( | 208 | +def _parse_stream_error_payload(chunk_str: str) -> dict: |
| 209 | - chunk_str: str, | 209 | + data_lines = [ |
| 210 | - *, | 210 | + line.removeprefix("data: ").strip() |
| 211 | - error_message: str, | 211 | + for line in chunk_str.splitlines() |
| 212 | - error_type: str | None = None, | 212 | + if line.startswith("data: ") and line.strip() != "data: [DONE]" |
| 213 | -) -> None: | 213 | + ] |
| 214 | - """Assert the SSE error payload propagates the expected message (and optional type).""" | 214 | + assert data_lines, f"expected streaming error chunk, got: {chunk_str!r}" |
| 215 | - payload = _parse_stream_error_payload(chunk_str) | 215 | + return json.loads(data_lines[-1]) |
| 216 | - # Coordinator-synthesized stream errors use the {"error": {...}} envelope (matching the | 216 | + |
| 217 | - # pre-commit / non-stream shape); unwrap it. Engine-verbatim bodies may already be flat. | 217 | + |
| 218 | - if isinstance(payload.get("error"), dict): | 218 | +def _assert_stream_error_chunk( |
| 219 | - payload = payload["error"] | 219 | + chunk_str: str, |
| 220 | - assert error_message in payload["message"], ( | 220 | + *, |
| 221 | - f"expected {error_message!r} in error message, got {payload['message']!r}" | 221 | + error_message: str, |
| 222 | - ) | 222 | + error_type: str | None = None, |
| 223 | - if error_type is not None: | 223 | +) -> None: |
| 224 | - assert payload["type"] == error_type | 224 | + """Assert the SSE error payload propagates the expected message (and optional type).""" |
| 225 | - | 225 | + payload = _parse_stream_error_payload(chunk_str) |
| 226 | - | 226 | + # Coordinator-synthesized stream errors use the {"error": {...}} envelope (matching the |
| 227 | -app = FastAPI() | 227 | + # pre-commit / non-stream shape); unwrap it. Engine-verbatim bodies may already be flat. |
| 228 | -_config = CoordinatorConfig() | 228 | + if isinstance(payload.get("error"), dict): |
| 229 | -# CDP separate mode requires worker metaserver; set so app-based tests have a valid config | 229 | + payload = payload["error"] |
| 230 | -_config.worker_metaserver_port = getattr(_config, "worker_metaserver_port", None) or 12000 | 230 | + assert error_message in payload["message"], ( |
| 231 | -_scheduler = Scheduler(instance_provider=InstanceManager(_config), config=_config) | 231 | + f"expected {error_message!r} in error message, got {payload['message']!r}" |
| 232 | -_request_manager = RequestManager(_config) | 232 | + ) |
| 233 | - | 233 | + if error_type is not None: |
| 234 | - | 234 | + assert payload["type"] == error_type |
| 235 | -@app.post("/v1/chat/completions") | 235 | + |
| 236 | -async def handle_completions(request: Request): | 236 | + |
| 237 | - return await router.handle_request(request, _config, scheduler=_scheduler, request_manager=_request_manager) | 237 | +app = FastAPI() |
| 238 | - | 238 | +_config = CoordinatorConfig() |
| 239 | - | 239 | +_scheduler = Scheduler(instance_provider=InstanceManager(_config), config=_config) |
| 240 | -@app.post("/v1/metaserver") | 240 | +_request_manager = RequestManager(_config) |
| 241 | -async def handle_metaserver(request: Request): | 241 | + |
| 242 | - """Legacy metaserver stub kept for unused MockAsyncClient helpers.""" | 242 | + |
| 243 | - await request.json() | 243 | +@app.post("/v1/chat/completions") |
| 244 | - return JSONResponse(content={"status": "ok"}) | 244 | +async def handle_completions(request: Request): |
| 245 | - | 245 | + return await router.handle_request(request, _config, scheduler=_scheduler, request_manager=_request_manager) |
| 246 | - | 246 | + |
| 247 | -class MockAsyncClient: | 247 | + |
| 248 | - def __init__( | 248 | +class TestRouterNativeHandoff: |
| 249 | - self, | 249 | + @pytest.fixture(autouse=True) |
| 250 | - post_exc: Exception = None, | 250 | + def fast_retry(self, monkeypatch: MonkeyPatch): |
| 251 | - stream_exc: Exception = None, | 251 | + """Skip real backoff in transport retry tests.""" |
| 252 | - post_fail_times: int = 1, | 252 | + |
| 253 | - stream_fail_times: int = 1, | 253 | + async def _instant_sleep(*_args, **_kwargs): |
| 254 | - ): | 254 | + return None |
| 255 | - self.post_exc = post_exc | 255 | + |
| 256 | - self.post_fail_times = post_fail_times | 256 | + monkeypatch.setattr(asyncio, "sleep", _instant_sleep) |
| 257 | - self.post_count = 0 | 257 | + |
| 258 | - self.post_fail_count = 0 | 258 | + @pytest.fixture |
| 259 | - | 259 | + def client(self): |
| 260 | - self.stream_exc = stream_exc | 260 | + return TestClient(app) |
| 261 | - self.stream_fail_times = stream_fail_times | 261 | + |
| 262 | - self.stream_count = 0 | 262 | + @classmethod |
| 263 | - self.stream_fail_count = 0 | 263 | + def create_mock_instance(cls, instance_id, role): |
| 264 | - | 264 | + """Create a proper mock Instance object""" |
| 265 | - self.req_data_from_metaserver = {} | 265 | + mock_instance = Instance( |
| 266 | - self.req_data_d_request = {} # D request (with metaserver URL), not overwritten by inner post() | 266 | + job_name=f"test-job-{instance_id}", |
| 267 | - self.req_headers_from_router = {} | 267 | + model_name=f"test-model-{instance_id}", |
| 268 | - | 268 | + engine_type="vllm", |
| 269 | - self.base_url = "test-base-url" | 269 | + id=instance_id, |
| 270 | - self.timeout = 1 | 270 | + role=role, |
| 271 | - self.is_closed = True | 271 | + status=InsStatus.ACTIVE, |
| 272 | - | 272 | + parallel_config=ParallelConfig(dp_size=1, tp_size=1), |
| 273 | - async def __aenter__(self): | 273 | + endpoints={}, |
| 274 | - return self | 274 | + ) |
| 275 | - | 275 | + return mock_instance |
| 276 | - async def __aexit__(self, exc_type, exc_val, exc_tb): | 276 | + |
| 277 | - pass | 277 | + def _make_router(self, req_info, monkeypatch, p_client, d_client): |
| 278 | - | 278 | + router_obj = UnifiedPDRouter( |
| 279 | - async def aclose(self): | 279 | + req_info, |
| 280 | - pass | 280 | + CoordinatorConfig(), |
| 281 | - | 281 | + scheduler=Scheduler( |
| 282 | - async def post(self, url, json=None, headers=None, **kwargs): | 282 | + instance_provider=InstanceManager(CoordinatorConfig()), |
| 283 | - self.post_count += 1 | 283 | + config=CoordinatorConfig(), |
| 284 | - if self.post_exc and self.post_fail_count < self.post_fail_times: | 284 | + ), |
| 285 | - self.post_fail_count += 1 | 285 | + request_manager=_request_manager, |
| 286 | - mock_response_fail = MagicMock() | 286 | + ) |
| 287 | - mock_response_fail.raise_for_status = MagicMock(side_effect=self.post_exc) | 287 | + _patch_unified_pd_clients(monkeypatch, router_obj, p_client, d_client) |
| 288 | - return mock_response_fail | 288 | + return router_obj |
| 289 | - | 289 | + |
| 290 | - self.req_data_from_metaserver = json | 290 | + @pytest.fixture |
| 291 | - request = httpx.Request("POST", url, headers=headers or {}, json=json) | 291 | + def setup_native_handoff(self, monkeypatch: MonkeyPatch): |
| 292 | - | 292 | + host = "127.0.0.1" |
| 293 | - return httpx.Response( | 293 | + # Create proper instances for separate P/D flow |
| 294 | - status_code=status.HTTP_200_OK, | 294 | + mock_instance_p = self.create_mock_instance(0, PDRole.ROLE_P) |
| 295 | - json={ | 295 | + mock_endpoint_p = Endpoint( |
| 296 | - "choices": [ | 296 | + id=0, |
| 297 | - { | 297 | + ip=host, |
| 298 | - "delta": {"content": "decoded chunk"}, | 298 | + business_port="8000", |
| 299 | - "index": 0, | 299 | + mgmt_port="8000", |
| 300 | - "finish_reason": None, | 300 | + status=EndpointStatus.NORMAL, |
| 301 | - } | 301 | + ) |
| 302 | - ], | 302 | + mock_instance_p.endpoints = {host: {0: mock_endpoint_p}} |
| 303 | - "id": "chatcmpl-123", | 303 | + |
| 304 | - }, | 304 | + mock_instance_d = self.create_mock_instance(1, PDRole.ROLE_D) |
| 305 | - request=request, | 305 | + mock_endpoint_d = Endpoint( |
| 306 | - ) | 306 | + id=1, |
| 307 | - | 307 | + ip=host, |
| 308 | - def stream(self, method, url, json=None, headers=None, **kwargs): | 308 | + business_port="8001", |
| 309 | - self.stream_count += 1 | 309 | + mgmt_port="8001", |
| 310 | - if json: | 310 | + status=EndpointStatus.NORMAL, |
| 311 | - self.req_data_from_metaserver = json | 311 | + ) |
| 312 | - self.req_data_d_request = json # keep D request; post() may overwrite req_data_from_metaserver | 312 | + mock_instance_d.endpoints = {host: {1: mock_endpoint_d}} |
| 313 | - # logger.info(f"----------req_data_from_coordinator:{json}") | 313 | + |
| 314 | - if self.stream_exc and self.stream_fail_count < self.stream_fail_times: | 314 | + # Mock functions (Scheduler uses get_required_instances_status for readiness) |
| 315 | - self.stream_fail_count += 1 | 315 | + def mock_get_required_instances_status(self): |
| 316 | - return MockStreamResponse(json or {}, recomputed=False, exc=self.stream_exc) | 316 | + return InstanceReadiness.REQUIRED_MET |
| 317 | - | 317 | + |
| 318 | - from urllib.parse import urlparse | 318 | + def mock_has_required_instances(self): |
| 319 | - | 319 | + return True |
| 320 | - client = TestClient(app) | 320 | + |
| 321 | - self.req_headers_from_router = headers | 321 | + def mock_get_available_instances(self, role=None): |
| 322 | - | 322 | + if role is None: |
| 323 | - url = json["kv_transfer_params"]["metaserver"] | 323 | + return { |
| 324 | - parsed_url = urlparse(url) | 324 | + mock_instance_p.id: mock_instance_p, |
| 325 | - | 325 | + mock_instance_d.id: mock_instance_d, |
| 326 | - # Forward request to metaserver | 326 | + } |
| 327 | - response = None | 327 | + if role == PDRole.ROLE_U: # PD hybrid role |
| 328 | - try: | 328 | + return {} # No PD hybrid instances, will use separate P/D |
| 329 | - response = client.post( | 329 | + if role == PDRole.ROLE_P: |
| 330 | - parsed_url.path, | 330 | + return {mock_instance_p.id: mock_instance_p} |
| 331 | - json={ | 331 | + if role == PDRole.ROLE_D: |
| 332 | - "request_id": headers.get("X-Request-Id"), | 332 | + return {mock_instance_d.id: mock_instance_d} |
| 333 | - "do_remote_decode": False, | 333 | + return {} |
| 334 | - "do_remote_prefill": True, | 334 | + |
| 335 | - "remote_engine_id": "test-engine", | 335 | + async def mock_select_instance_and_endpoint(self, role): |
| 336 | - "remote_host": parsed_url.hostname, | 336 | + if role == PDRole.ROLE_P: |
| 337 | - "remote_port": str(parsed_url.port), | 337 | + return mock_instance_p, mock_endpoint_p |
| 338 | - }, | 338 | + elif role == PDRole.ROLE_D: |
| 339 | - ) | 339 | + return mock_instance_d, mock_endpoint_d |
| 340 | - response.raise_for_status() | 340 | + return None, None |
| 341 | - except Exception as e: | 341 | + |
| 342 | - err_text = getattr(response, "text", str(e)) if response is not None else str(e) | 342 | + async def mock_select_and_allocate( |
| 343 | - err_status = getattr(response, "status_code", 500) if response is not None else 500 | 343 | + self, |
| 344 | - return MockStreamResponse( | 344 | + role, |
| 345 | - json or {}, | 345 | + req_info, |
| 346 | - recomputed=False, | 346 | + *, |
| 347 | - exc=httpx.HTTPStatusError( | 347 | + target_instance_id=None, |
| 348 | - message=err_text, | 348 | + required_engine_type=None, |
| 349 | - request=MagicMock(), | 349 | + ): |
| 350 | - response=httpx.Response(status_code=err_status, text=err_text), | 350 | + del required_engine_type |
| 351 | - ), | 351 | + if role == PDRole.ROLE_P: |
| 352 | - ) | 352 | + return ( |
| 353 | - | 353 | + mock_instance_p, |
| 354 | - # Return an async context manager | 354 | + mock_endpoint_p, |
| 355 | - return MockStreamResponse(json or {}, recomputed=False, exc=None) | 355 | + Workload(active_tokens=1), |
| 356 | - | 356 | + ) |
| 357 | - | 357 | + if role == PDRole.ROLE_D: |
| 358 | -class MockAsyncClientFirstStreamRecompute(MockAsyncClient): | 358 | + return ( |
| 359 | - """First decode stream simulates recompute after partial output; second completes.""" | 359 | + mock_instance_d, |
| 360 | - | 360 | + mock_endpoint_d, |
| 361 | - def stream(self, method, url, json=None, headers=None, **kwargs): | 361 | + Workload(active_tokens=1), |
| 362 | - self.stream_count += 1 | 362 | + ) |
| 363 | - if json: | 363 | + return None |
| 364 | - self.req_data_from_metaserver = json | 364 | + |
| 365 | - self.req_data_d_request = json | 365 | + async def mock_update_workload(self, params): |
| 366 | - if self.stream_exc and self.stream_fail_count < self.stream_fail_times: | 366 | + return True |
| 367 | - self.stream_fail_count += 1 | 367 | + |
| 368 | - return MockStreamResponse(json or {}, recomputed=False, exc=self.stream_exc) | 368 | + monkeypatch.setattr( |
| 369 | - | 369 | + InstanceManager, |
| 370 | - from urllib.parse import urlparse | 370 | + "get_required_instances_status", |
| 371 | - | 371 | + mock_get_required_instances_status, |
| 372 | - client = TestClient(app) | 372 | + ) |
| 373 | - self.req_headers_from_router = headers or {} | 373 | + monkeypatch.setattr(InstanceManager, "has_required_instances", mock_has_required_instances) |
| 374 | - | 374 | + monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances) |
| 375 | - url_ms = json["kv_transfer_params"]["metaserver"] | 375 | + monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint) |
| 376 | - parsed_url = urlparse(url_ms) | 376 | + monkeypatch.setattr(Scheduler, "select_and_allocate", mock_select_and_allocate) |
| 377 | - | 377 | + monkeypatch.setattr(Scheduler, "update_workload", mock_update_workload) |
| 378 | - response = None | 378 | + |
| 379 | - try: | 379 | + mock_scheduler_config = MagicMock() |
| 380 | - response = client.post( | 380 | + mock_scheduler_config.scheduler_type = SchedulerType.LOAD_BALANCE |
| 381 | - parsed_url.path, | 381 | + # Real ExceptionConfig so transport_retry_limit and rescheduling settings work; |
| 382 | - json={ | 382 | + # MagicMock lacks @property implementation and breaks decode transport loops (range / last-attempt check). |
| 383 | - "request_id": headers.get("X-Request-Id"), | 383 | + mock_exception_config = ExceptionConfig(max_retry=5, retry_delay=0.0001) |
| 384 | - "do_remote_decode": False, | 384 | + mock_api_config = MagicMock() |
| 385 | - "do_remote_prefill": True, | 385 | + mock_api_config.coordinator_api_host = "127.0.0.1" |
| 386 | - "remote_engine_id": "test-engine", | 386 | + mock_tls_config = MagicMock() |
| 387 | - "remote_host": parsed_url.hostname, | 387 | + mock_tls_config.enable_tls = False |
| 388 | - "remote_port": str(parsed_url.port), | 388 | + |
| 389 | - }, | 389 | + mock_config = MagicMock() |
| 390 | - ) | 390 | + mock_config.scheduler_config = mock_scheduler_config |
| 391 | - response.raise_for_status() | 391 | + mock_config.exception_config = mock_exception_config |
| 392 | - except Exception as e: | 392 | + mock_config.api_config = mock_api_config |
| 393 | - err_text = getattr(response, "text", str(e)) if response is not None else str(e) | 393 | + mock_config.infer_tls_config = mock_tls_config |
| 394 | - err_status = getattr(response, "status_code", 500) if response is not None else 500 | 394 | + mock_config.mgmt_tls_config = mock_tls_config |
| 395 | - return MockStreamResponse( | 395 | + monkeypatch.setattr(CoordinatorConfig, "__new__", lambda cls: mock_config) |
| 396 | - json or {}, | 396 | + _config.exception_config = mock_exception_config |
| 397 | - recomputed=False, | 397 | + |
| 398 | - exc=httpx.HTTPStatusError( | 398 | + @pytest.fixture |
| 399 | - message=err_text, | 399 | + def mock_raw_request(self): |
| 400 | - request=MagicMock(), | 400 | + # Mock Request |
| 401 | - response=httpx.Response(status_code=err_status, text=err_text), | 401 | + mock_req = MagicMock(spec=Request) |
| 402 | - ), | 402 | + mock_req.body = AsyncMock(return_value=b'{"model": "test"}') |
| 403 | - ) | 403 | + mock_req.json = AsyncMock(return_value={"model": "test"}) |
| 404 | - | 404 | + mock_req.headers = {} |
| 405 | - recomputed = self.stream_count == 1 | 405 | + mock_req.url.path = "/v1/chat/completions" |
| 406 | - return MockStreamResponse(json or {}, recomputed=recomputed, exc=None) | 406 | + |
| 407 | - | 407 | + # Must be awaitable so listen_for_disconnect() does not raise; never completes so handler wins. |
| 408 | - | 408 | + async def _never_receive(): |
| 409 | -class TestRouterCDPSeparation: | 409 | + await asyncio.Event().wait() |
| 410 | - @pytest.fixture(autouse=True) | 410 | + |
| 411 | - def fast_retry(self, monkeypatch: MonkeyPatch): | 411 | + mock_req.receive = AsyncMock(side_effect=_never_receive) |
| 412 | - """Skip real backoff and dispatch-stop HTTP in transport retry tests.""" | 412 | + return mock_req |
| 413 | - | 413 | + |
| 414 | - async def _instant_sleep(*_args, **_kwargs): | 414 | + @pytest.mark.asyncio |
| 415 | - return None | 415 | + async def test_successful_native_handoff_request(self, client, monkeypatch: MonkeyPatch, setup_native_handoff): |
| 416 | - | 416 | + """Test case: native handoff request succeeds. |
| 417 | - async def _noop_dispatch_stop(self, resource, attempt, reason, timeout=1.0): | 417 | + Expected behavior: |
| 418 | - return None | 418 | + 1) Check request status is DecodeEnd |
| 419 | - | 419 | + 2) Return normal response |
| 420 | - monkeypatch.setattr(asyncio, "sleep", _instant_sleep) | 420 | + """ |
| 421 | - monkeypatch.setattr( | 421 | + p_client = _UnifiedPDPrefillClient() |
| 422 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | 422 | + d_client = _UnifiedPDDecodeClient() |
| 423 | - _noop_dispatch_stop, | 423 | + |
| 424 | - ) | 424 | + req_info = await create_mock_request_info() |
| 425 | - | 425 | + origin_req_id = req_info.req_id |
| 426 | - @pytest.fixture | 426 | + origin_req_len = req_info.req_len |
| 427 | - def client(self): | 427 | + origin_req_data = req_info.req_data |
| 428 | - return TestClient(app) | 428 | + |
| 429 | - | 429 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 430 | - @classmethod | 430 | + response = await handoff_router.handle_request() |
| 431 | - def create_mock_instance(cls, instance_id, role): | 431 | + chunks = [] |
| 432 | - """Create a proper mock Instance object""" | 432 | + async for chunk in response.body_iterator: |
| 433 | - mock_instance = Instance( | 433 | + chunks.append(chunk) |
| 434 | - job_name=f"test-job-{instance_id}", | 434 | + |
| 435 | - model_name=f"test-model-{instance_id}", | 435 | + assert response.status_code == status.HTTP_200_OK |
| 436 | - engine_type="vllm", | 436 | + assert "text/event-stream" in response.headers.get("content-type") |
| 437 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | 437 | + |
| 438 | - id=instance_id, | 438 | + assert len(p_client.requests) == 1 |
| 439 | - role=role, | 439 | + assert len(d_client.requests) == 1 |
| 440 | - status=InsStatus.ACTIVE, | 440 | + req_data_p = p_client.requests[0] |
| 441 | - parallel_config=ParallelConfig(dp_size=1, tp_size=1), | 441 | + assert req_data_p["stream"] is False |
| 442 | - endpoints={}, | 442 | + assert req_data_p["max_tokens"] == 1 |
| 443 | - ) | 443 | + assert "_motor_dispatch" not in req_data_p |
| 444 | - return mock_instance | 444 | + assert "_motor_dispatch" not in d_client.requests[0] |
| 445 | - | 445 | + assert d_client.requests[0]["request_id"] == req_data_p["request_id"] |
| 446 | - def _make_router(self, req_info, monkeypatch, p_client, d_client): | 446 | + assert d_client.requests[0]["kv_transfer_params"] == { |
| 447 | - router_obj = SeparateCDPRouter( | 447 | + "do_remote_prefill": True, |
| 448 | - req_info, | 448 | + "remote_request_id": req_data_p["request_id"], |
| 449 | - CoordinatorConfig(), | 449 | + "remote_host": "127.0.0.1", |
| 450 | - scheduler=Scheduler( | 450 | + "remote_port": 9000, |
| 451 | - instance_provider=InstanceManager(CoordinatorConfig()), | 451 | + } |
| 452 | - config=CoordinatorConfig(), | 452 | + |
| 453 | - ), | 453 | + assert req_info.req_id == origin_req_id |
| 454 | - request_manager=_request_manager, | 454 | + assert req_info.req_len == origin_req_len |
| 455 | - ) | 455 | + assert req_info.req_data == origin_req_data |
| 456 | - _patch_unified_pd_clients(monkeypatch, router_obj, p_client, d_client) | 456 | + |
| 457 | - return router_obj | 457 | + assert req_info.state == ReqState.DECODE_END |
| 458 | - | 458 | + assert req_info.status[ReqState.D_ALLOCATED] >= req_info.status[ReqState.ARRIVE] |
| 459 | - @pytest.fixture | 459 | + assert req_info.status[ReqState.P_ALLOCATED] >= req_info.status[ReqState.ARRIVE] |
| 460 | - def setup_cdp_separation(self, monkeypatch: MonkeyPatch): | 460 | + allocated_at = min( |
| 461 | - host = "127.0.0.1" | 461 | + req_info.status[ReqState.P_ALLOCATED], |
| 462 | - # Create proper instances for separate P/D flow | 462 | + req_info.status[ReqState.D_ALLOCATED], |
| 463 | - mock_instance_p = self.create_mock_instance(0, PDRole.ROLE_P) | 463 | + ) |
| 464 | - mock_endpoint_p = Endpoint( | 464 | + assert req_info.status[ReqState.DECODE_END] >= allocated_at |
| 465 | - id=0, | 465 | + |
| 466 | - ip=host, | 466 | + @pytest.mark.asyncio |
| 467 | - business_port="8000", | 467 | + async def test_native_handoff_stream_recompute_after_partial_output_continues( |
| 468 | - mgmt_port="8000", | 468 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 469 | - status=EndpointStatus.NORMAL, | 469 | + ): |
| 470 | - ) | 470 | + """Decode stream with recomputed stop_reason still completes successfully.""" |
| 471 | - mock_instance_p.endpoints = {host: {0: mock_endpoint_p}} | 471 | + d_client = _UnifiedPDDecodeClient( |
| 472 | - | 472 | + stream_chunks=[ |
| 473 | - mock_instance_d = self.create_mock_instance(1, PDRole.ROLE_D) | 473 | + b'data: {"choices":[{"delta":{"content":"partial "},"index":0,"stop_reason":"recomputed","token_ids":[1,2]}],"prompt_token_ids":[10,11]}\n\n', |
| 474 | - mock_endpoint_d = Endpoint( | 474 | + b'data: {"choices":[{"delta":{"content":"continuation"},"index":0,"finish_reason":"stop"}]}\n\n', |
| 475 | - id=1, | 475 | + ] |
| 476 | - ip=host, | 476 | + ) |
| 477 | - business_port="8001", | 477 | + p_client = _UnifiedPDPrefillClient() |
| 478 | - mgmt_port="8001", | 478 | + req_info = await create_mock_request_info() |
| 479 | - status=EndpointStatus.NORMAL, | 479 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 480 | - ) | 480 | + |
| 481 | - mock_instance_d.endpoints = {host: {1: mock_endpoint_d}} | 481 | + response = await handoff_router.handle_request() |
| 482 | - | 482 | + chunk_str = await _collect_stream_chunks(response) |
| 483 | - # Mock functions (Scheduler uses get_required_instances_status for readiness) | 483 | + |
| 484 | - def mock_get_required_instances_status(self): | 484 | + assert req_info.state == ReqState.DECODE_END, chunk_str |
| 485 | - return InstanceReadiness.REQUIRED_MET | 485 | + assert d_client.stream_count == 1 |
| 486 | - | 486 | + assert "partial" in chunk_str |
| 487 | - def mock_has_required_instances(self): | 487 | + assert "recompute after first chunk" not in chunk_str |
| 488 | - return True | 488 | + |
| 489 | - | 489 | + @pytest.mark.asyncio |
| 490 | - def mock_get_available_instances(self, role=None): | 490 | + async def test_native_engine_decode_4xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_native_handoff): |
| 491 | - if role is None: | 491 | + """Test case: native decode engine returns a 4XX status code. |
| 492 | - return { | 492 | + Expected behavior: |
| 493 | - mock_instance_p.id: mock_instance_p, | 493 | + 1) No request retry triggered |
| 494 | - mock_instance_d.id: mock_instance_d, | 494 | + 2) Directly return error message |
| 495 | - } | 495 | + """ |
| 496 | - if role == PDRole.ROLE_U: # PD hybrid role | 496 | + error_message = "Test Bad Request" |
| 497 | - return {} # No PD hybrid instances, will use separate P/D | 497 | + d_client = _UnifiedPDDecodeClient( |
| 498 | - if role == PDRole.ROLE_P: | 498 | + stream_exc=httpx.HTTPStatusError( |
| 499 | - return {mock_instance_p.id: mock_instance_p} | 499 | + message=error_message, |
| 500 | - if role == PDRole.ROLE_D: | 500 | + request=MagicMock(), |
| 501 | - return {mock_instance_d.id: mock_instance_d} | 501 | + response=httpx.Response(status_code=status.HTTP_400_BAD_REQUEST, text=error_message), |
| 502 | - return {} | 502 | + ) |
| 503 | - | 503 | + ) |
| 504 | - async def mock_select_instance_and_endpoint(self, role): | 504 | + p_client = _UnifiedPDPrefillClient() |
| 505 | - if role == PDRole.ROLE_P: | 505 | + req_info = await create_mock_request_info() |
| 506 | - return mock_instance_p, mock_endpoint_p | 506 | + |
| 507 | - elif role == PDRole.ROLE_D: | 507 | + release_p_tokens = 0 |
| 508 | - return mock_instance_d, mock_endpoint_d | 508 | + release_d_tokens = 0 |
| 509 | - return None, None | 509 | + original_release = UnifiedPDRouter._release_attempt_resource |
| 510 | - | 510 | + |
| 511 | - async def mock_select_and_allocate(self, role, req_info, *, target_instance_id=None): | 511 | + async def mock_release_attempt_resource(self, resource, attempt_seq, action, attempt=None, **kwargs): |
| 512 | - if role == PDRole.ROLE_P: | 512 | + nonlocal release_p_tokens, release_d_tokens |
| 513 | - return ( | 513 | + if resource.instance.role == PDRole.ROLE_P: |
| 514 | - mock_instance_p, | 514 | + if action == WorkloadAction.RELEASE_TOKENS: |
| 515 | - mock_endpoint_p, | 515 | + release_p_tokens += 1 |
| 516 | - Workload(active_tokens=1), | 516 | + elif resource.instance.role == PDRole.ROLE_D: |
| 517 | - ) | 517 | + if action == WorkloadAction.RELEASE_TOKENS: |
| 518 | - if role == PDRole.ROLE_D: | 518 | + release_d_tokens += 1 |
| 519 | - return ( | 519 | + await original_release(self, resource, attempt_seq, action, attempt, **kwargs) |
| 520 | - mock_instance_d, | 520 | + |
| 521 | - mock_endpoint_d, | 521 | + monkeypatch.setattr( |
| 522 | - Workload(active_tokens=1), | 522 | + UnifiedPDRouter, |
| 523 | - ) | 523 | + "_release_attempt_resource", |
| 524 | - return None | 524 | + mock_release_attempt_resource, |
| 525 | - | 525 | + ) |
| 526 | - async def mock_update_workload(self, params): | 526 | + |
| 527 | - return True | 527 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 528 | - | 528 | + response = await handoff_router.handle_request() |
| 529 | - monkeypatch.setattr( | 529 | + chunk_str = await _collect_stream_chunks(response) |
| 530 | - InstanceManager, | 530 | + |
| 531 | - "get_required_instances_status", | 531 | + assert req_info.state == ReqState.EXCEPTION |
| 532 | - mock_get_required_instances_status, | 532 | + _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="UpstreamHTTPError") |
| 533 | - ) | 533 | + assert str(status.HTTP_400_BAD_REQUEST) in chunk_str |
| 534 | - monkeypatch.setattr(InstanceManager, "has_required_instances", mock_has_required_instances) | 534 | + assert d_client.stream_count == 1 |
| 535 | - monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances) | 535 | + assert release_d_tokens >= 1 |
| 536 | - monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint) | 536 | + assert release_p_tokens >= 1 |
| 537 | - monkeypatch.setattr(Scheduler, "select_and_allocate", mock_select_and_allocate) | 537 | + |
| 538 | - monkeypatch.setattr(Scheduler, "update_workload", mock_update_workload) | 538 | + @pytest.mark.asyncio |
| 539 | - | 539 | + async def test_native_engine_decode_continuous_5xx_status_code( |
| 540 | - mock_scheduler_config = MagicMock() | 540 | + self, |
| 541 | - mock_scheduler_config.scheduler_type = SchedulerType.LOAD_BALANCE | 541 | + client, |
| 542 | - # Real ExceptionConfig so transport_retry_limit and rescheduling settings work; | 542 | + monkeypatch: MonkeyPatch, |
| 543 | - # MagicMock lacks @property implementation and breaks decode transport loops (range / last-attempt check). | 543 | + setup_native_handoff, |
| 544 | - mock_exception_config = ExceptionConfig(max_retry=5, retry_delay=0.0001) | 544 | + caplog: pytest.LogCaptureFixture, |
| 545 | - mock_api_config = MagicMock() | 545 | + ): |
| 546 | - mock_api_config.coordinator_api_host = "127.0.0.1" | 546 | + """Decode keeps getting 5XX with the same message: retries exhaust, error chunk returned; |
| 547 | - mock_tls_config = MagicMock() | 547 | + identical-error logs: one ERROR + (max_retry-1) WARNING dedup lines. |
| 548 | - mock_tls_config.enable_tls = False | 548 | + """ |
| 549 | - | 549 | + error_message = "Test Internal Server Error" |
| 550 | - mock_config = MagicMock() | 550 | + max_retry = CoordinatorConfig().exception_config.transport_retry_limit |
| 551 | - mock_config.scheduler_config = mock_scheduler_config | 551 | + d_client = _UnifiedPDDecodeClient( |
| 552 | - mock_config.exception_config = mock_exception_config | 552 | + stream_exc=httpx.HTTPStatusError( |
| 553 | - mock_config.api_config = mock_api_config | 553 | + message=error_message, |
| 554 | - mock_config.infer_tls_config = mock_tls_config | 554 | + request=MagicMock(), |
| 555 | - mock_config.mgmt_tls_config = mock_tls_config | 555 | + response=httpx.Response( |
| 556 | - # CDP separate requires worker metaserver; use a fixed port for test | 556 | + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 557 | - mock_config.worker_metaserver_port = 12000 | 557 | + text=error_message, |
| 558 | - | 558 | + ), |
| 559 | - monkeypatch.setattr(CoordinatorConfig, "__new__", lambda cls: mock_config) | 559 | + ) |
| 560 | - _config.exception_config = mock_exception_config | 560 | + ) |
| 561 | - | 561 | + p_client = _UnifiedPDPrefillClient() |
| 562 | - @pytest.fixture | 562 | + req_info = await create_mock_request_info() |
| 563 | - def mock_raw_request(self): | 563 | + |
| 564 | - # Mock Request | 564 | + exec_release = 0 |
| 565 | - mock_req = MagicMock(spec=Request) | 565 | + original_release = UnifiedPDRouter._release_attempt_resource |
| 566 | - mock_req.body = AsyncMock(return_value=b'{"model": "test"}') | 566 | + |
| 567 | - mock_req.json = AsyncMock(return_value={"model": "test"}) | 567 | + async def mock_release_attempt_resource(self, resource, attempt_seq, action, attempt=None, **kwargs): |
| 568 | - mock_req.headers = {} | 568 | + nonlocal exec_release |
| 569 | - mock_req.url.path = "/v1/chat/completions" | 569 | + exec_release += 1 |
| 570 | - | 570 | + await original_release(self, resource, attempt_seq, action, attempt, **kwargs) |
| 571 | - # Must be awaitable so listen_for_disconnect() does not raise; never completes so handler wins. | 571 | + |
| 572 | - async def _never_receive(): | 572 | + monkeypatch.setattr( |
| 573 | - await asyncio.Event().wait() | 573 | + UnifiedPDRouter, |
| 574 | - | 574 | + "_release_attempt_resource", |
| 575 | - mock_req.receive = AsyncMock(side_effect=_never_receive) | 575 | + mock_release_attempt_resource, |
| 576 | - return mock_req | 576 | + ) |
| 577 | - | 577 | + |
| 578 | - @pytest.mark.asyncio | 578 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 579 | - async def test_successful_request_with_separate_cdp(self, client, monkeypatch: MonkeyPatch, setup_cdp_separation): | 579 | + response = await handoff_router.handle_request() |
| 580 | - """Test case: CDP separation mode request success | 580 | + chunk_str = await _collect_stream_chunks(response) |
| 581 | - Expected behavior: | 581 | + |
| 582 | - 1) Check request status is DecodeEnd | 582 | + assert req_info.state == ReqState.EXCEPTION |
| 583 | - 2) Return normal response | 583 | + _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="UpstreamHTTPError") |
| 584 | - """ | 584 | + assert str(status.HTTP_500_INTERNAL_SERVER_ERROR) in chunk_str |
| 585 | - p_client = _UnifiedPDPrefillClient() | 585 | + assert d_client.stream_count == max_retry |
| 586 | - d_client = _UnifiedPDDecodeClient() | 586 | + assert exec_release >= 1 |
| 587 | - | 587 | + |
| 588 | - req_info = await create_mock_request_info() | 588 | + @pytest.mark.asyncio |
| 589 | - origin_req_id = req_info.req_id | 589 | + async def test_native_engine_decode_once_5xx_status_code( |
| 590 | - origin_req_len = req_info.req_len | 590 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 591 | - origin_req_data = req_info.req_data | 591 | + ): |
| 592 | - | 592 | + """Test case: native decode request first returns 5XX, then 200. |
| 593 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 593 | + Expected behavior: |
| 594 | - response = await cdp_router.handle_request() | 594 | + 1) Check request status is DecodeEnd |
| 595 | - chunks = [] | 595 | + 2) Trigger request retry |
| 596 | - async for chunk in response.body_iterator: | 596 | + 3) Request retry succeeds |
| 597 | - chunks.append(chunk) | 597 | + """ |
| 598 | - | 598 | + error_message = "Test Internal Server Error" |
| 599 | - assert response.status_code == status.HTTP_200_OK | 599 | + d_client = _UnifiedPDDecodeClient( |
| 600 | - assert "text/event-stream" in response.headers.get("content-type") | 600 | + stream_exc=httpx.HTTPStatusError( |
| 601 | - | 601 | + message=error_message, |
| 602 | - assert len(p_client.requests) == 1 | 602 | + request=MagicMock(), |
| 603 | - assert len(d_client.requests) == 1 | 603 | + response=httpx.Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR), |
| 604 | - req_data_p = p_client.requests[0] | 604 | + ), |
| 605 | - assert req_data_p["stream"] is False | 605 | + stream_fail_times=1, |
| 606 | - assert req_data_p["max_tokens"] == 1 | 606 | + ) |
| 607 | - assert req_data_p[MOTOR_DISPATCH_KEY]["role"] == "prefill" | 607 | + p_client = _UnifiedPDPrefillClient() |
| 608 | - assert d_client.requests[0][MOTOR_DISPATCH_KEY]["role"] == "decode" | 608 | + req_info = await create_mock_request_info() |
| 609 | - assert req_data_p[MOTOR_DISPATCH_KEY]["pair_id"] == d_client.requests[0][MOTOR_DISPATCH_KEY]["pair_id"] | 609 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 610 | - | 610 | + |
| 611 | - assert req_info.req_id == origin_req_id | 611 | + response = await handoff_router.handle_request() |
| 612 | - assert req_info.req_len == origin_req_len | 612 | + chunk_str = await _collect_stream_chunks(response) |
| 613 | - assert req_info.req_data == origin_req_data | 613 | + |
| 614 | - | 614 | + assert response.status_code == status.HTTP_200_OK |
| 615 | - assert req_info.state == ReqState.DECODE_END | 615 | + assert d_client.stream_fail_count == 1 |
| 616 | - assert req_info.status[ReqState.D_ALLOCATED] >= req_info.status[ReqState.ARRIVE] | 616 | + assert d_client.stream_count >= 2 |
| 617 | - assert req_info.status[ReqState.P_ALLOCATED] >= req_info.status[ReqState.ARRIVE] | 617 | + assert req_info.state == ReqState.DECODE_END |
| 618 | - allocated_at = min( | 618 | + assert "decoded chunk" in chunk_str |
| 619 | - req_info.status[ReqState.P_ALLOCATED], | 619 | + |
| 620 | - req_info.status[ReqState.D_ALLOCATED], | 620 | + @pytest.mark.asyncio |
| 621 | - ) | 621 | + async def test_native_engine_decode_network_exception(self, client, monkeypatch: MonkeyPatch, setup_native_handoff): |
| 622 | - assert req_info.status[ReqState.DECODE_END] >= allocated_at | 622 | + """Test case: native decode engine raises a network exception. |
| 623 | - | 623 | + Expected behavior: |
| 624 | - @pytest.mark.asyncio | 624 | + 1) Check request status is Exception |
| 625 | - async def test_cdp_stream_recompute_after_partial_output_continues( | 625 | + 2) Retries exhaust transport_retry_limit |
| 626 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | 626 | + 3) Directly return error message |
| 627 | - ): | 627 | + """ |
| 628 | - """Decode stream with recomputed stop_reason still completes successfully.""" | 628 | + error_message = "Connection error" |
| 629 | - d_client = _UnifiedPDDecodeClient( | 629 | + max_retry = CoordinatorConfig().exception_config.transport_retry_limit |
| 630 | - stream_chunks=[ | 630 | + d_client = _UnifiedPDDecodeClient( |
| 631 | - b'data: {"choices":[{"delta":{"content":"partial "},"index":0,"stop_reason":"recomputed","token_ids":[1,2]}],"prompt_token_ids":[10,11]}\n\n', | 631 | + stream_exc=httpx.ConnectError(error_message, request=MagicMock()), |
| 632 | - b'data: {"choices":[{"delta":{"content":"continuation"},"index":0,"finish_reason":"stop"}]}\n\n', | 632 | + stream_fail_times=max_retry, |
| 633 | - ] | 633 | + ) |
| 634 | - ) | 634 | + p_client = _UnifiedPDPrefillClient() |
| 635 | - p_client = _UnifiedPDPrefillClient() | 635 | + req_info = await create_mock_request_info() |
| 636 | - req_info = await create_mock_request_info() | 636 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 637 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 637 | + |
| 638 | - | 638 | + response = await handoff_router.handle_request() |
| 639 | - response = await cdp_router.handle_request() | 639 | + chunk_str = await _collect_stream_chunks(response) |
| 640 | - chunk_str = await _collect_stream_chunks(response) | 640 | + |
| 641 | - | 641 | + _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="ConnectError") |
| 642 | - assert req_info.state == ReqState.DECODE_END, chunk_str | 642 | + assert d_client.stream_count == max_retry |
| 643 | - assert d_client.stream_count == 1 | 643 | + assert d_client.stream_fail_count == max_retry |
| 644 | - assert "partial" in chunk_str | 644 | + assert req_info.state == ReqState.EXCEPTION |
| 645 | - assert "recompute after first chunk" not in chunk_str | 645 | + |
| 646 | - | 646 | + @pytest.mark.asyncio |
| 647 | - @pytest.mark.asyncio | 647 | + async def test_native_handoff_decode_non_stream_retry_exhausts_transport_limit( |
| 648 | - async def test_cdp_requires_worker_metaserver_port(self, setup_cdp_separation, monkeypatch: MonkeyPatch): | 648 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 649 | - """UnifiedPD no longer requires worker_metaserver_port on the coordinator.""" | 649 | + ): |
| 650 | - req_info = await create_mock_request_info() | 650 | + """Non-stream decode failures retry whole transport attempts until limit.""" |
| 651 | - p_client = _UnifiedPDPrefillClient() | 651 | + error_message = "Same post Decode error every retry" |
| 652 | - d_client = _UnifiedPDDecodeClient() | 652 | + max_retry = CoordinatorConfig().exception_config.transport_retry_limit |
| 653 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 653 | + d_client = _UnifiedPDDecodeClient( |
| 654 | - response = await cdp_router.handle_request() | 654 | + post_exc=httpx.HTTPStatusError( |
| 655 | - assert response.status_code == status.HTTP_200_OK | 655 | + message=error_message, |
| 656 | - await _collect_stream_chunks(response) | 656 | + request=MagicMock(), |
| 657 | - assert req_info.state == ReqState.DECODE_END | 657 | + response=httpx.Response(status_code=status.HTTP_502_BAD_GATEWAY, text=error_message), |
| 658 | - | 658 | + ), |
| 659 | - @pytest.mark.asyncio | 659 | + post_fail_times=max_retry, |
| 660 | - async def test_engine_server_decode_4xx_status_code(self, client, monkeypatch: MonkeyPatch, setup_cdp_separation): | 660 | + ) |
| 661 | - """Test case: Decode EngineServer returns 4XX status code | 661 | + p_client = _UnifiedPDPrefillClient() |
| 662 | - Expected behavior: | 662 | + req_info = await create_mock_request_info(stream=False) |
| 663 | - 1) No request retry triggered | 663 | + |
| 664 | - 2) Directly return error message | 664 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 665 | - """ | 665 | + |
| 666 | - error_message = "Test Bad Request" | 666 | + with pytest.raises(httpx.HTTPStatusError): |
| 667 | - d_client = _UnifiedPDDecodeClient( | 667 | + await handoff_router.handle_request() |
| 668 | - stream_exc=httpx.HTTPStatusError( | 668 | + |
| 669 | - message=error_message, | 669 | + assert d_client.post_count == max_retry |
| 670 | - request=MagicMock(), | 670 | + assert d_client.post_fail_count == max_retry |
| 671 | - response=httpx.Response(status_code=status.HTTP_400_BAD_REQUEST, text=error_message), | 671 | + assert req_info.state == ReqState.EXCEPTION |
| 672 | - ) | 672 | + |
| 673 | - ) | 673 | + @pytest.mark.asyncio |
| 674 | - p_client = _UnifiedPDPrefillClient() | 674 | + async def test_native_engine_prefill_network_exception( |
| 675 | - req_info = await create_mock_request_info() | 675 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 676 | - | 676 | + ): |
| 677 | - release_p_tokens = 0 | 677 | + """Test case: native prefill engine raises a network exception. |
| 678 | - release_d_tokens = 0 | 678 | + Expected behavior: |
| 679 | - original_release = SeparateCDPRouter._release_attempt_resource | 679 | + 1) Check request status is Exception |
| 680 | - | 680 | + 2) Retries exhaust transport_retry_limit |
| 681 | - async def mock_release_attempt_resource(self, resource, attempt_seq, action, attempt=None, **kwargs): | 681 | + 3) Directly return error message |
| 682 | - nonlocal release_p_tokens, release_d_tokens | 682 | + """ |
| 683 | - if resource.instance.role == PDRole.ROLE_P: | 683 | + error_message = "Connection error" |
| 684 | - if action == WorkloadAction.RELEASE_TOKENS: | 684 | + retry_times = CoordinatorConfig().exception_config.transport_retry_limit |
| 685 | - release_p_tokens += 1 | 685 | + p_client = _UnifiedPDPrefillClient( |
| 686 | - elif resource.instance.role == PDRole.ROLE_D: | 686 | + exc=httpx.ConnectError(message=error_message, request=MagicMock()), |
| 687 | - if action == WorkloadAction.RELEASE_TOKENS: | 687 | + post_fail_times=retry_times, |
| 688 | - release_d_tokens += 1 | 688 | + ) |
| 689 | - await original_release(self, resource, attempt_seq, action, attempt, **kwargs) | 689 | + d_client = _UnifiedPDDecodeClient() |
| 690 | - | 690 | + _patch_unified_pd_router_clients(monkeypatch, p_client, d_client) |
| 691 | - monkeypatch.setattr( | 691 | + |
| 692 | - SeparateCDPRouter, | 692 | + state: ReqState = None |
| 693 | - "_release_attempt_resource", | 693 | + |
| 694 | - mock_release_attempt_resource, | 694 | + def mock_update_state(self, new_state: ReqState): |
| 695 | - ) | 695 | + nonlocal state |
| 696 | - | 696 | + state = new_state |
| 697 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 697 | + |
| 698 | - response = await cdp_router.handle_request() | 698 | + monkeypatch.setattr(RequestInfo, "update_state", mock_update_state) |
| 699 | - chunk_str = await _collect_stream_chunks(response) | 699 | + |
| 700 | - | 700 | + with client.stream( |
| 701 | - assert req_info.state == ReqState.EXCEPTION | 701 | + "POST", |
| 702 | - _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="UpstreamHTTPError") | 702 | + "/v1/chat/completions", |
| 703 | - assert str(status.HTTP_400_BAD_REQUEST) in chunk_str | 703 | + json={ |
| 704 | - assert d_client.stream_count == 1 | 704 | + "model": "test-model", |
| 705 | - assert release_d_tokens >= 1 | 705 | + "messages": [{"role": "user", "content": "Hello"}], |
| 706 | - assert release_p_tokens >= 1 | 706 | + }, |
| 707 | - | 707 | + ) as response: |
| 708 | - @pytest.mark.asyncio | 708 | + chunks = [] |
| 709 | - async def test_engine_server_decode_continuous_5xx_status_code( | 709 | + for chunk in response.iter_lines(): |
| 710 | - self, | 710 | + chunks.append(chunk) |
| 711 | - client, | 711 | + chunk_str = "".join(chunks) |
| 712 | - monkeypatch: MonkeyPatch, | 712 | + |
| 713 | - setup_cdp_separation, | 713 | + assert error_message in chunk_str |
| 714 | - caplog: pytest.LogCaptureFixture, | 714 | + assert p_client.post_fail_count == retry_times |
| 715 | - ): | 715 | + assert state == ReqState.EXCEPTION |
| 716 | - """Decode keeps getting 5XX with the same message: retries exhaust, error chunk returned; | 716 | + |
| 717 | - identical-error logs: one ERROR + (max_retry-1) WARNING dedup lines. | 717 | + @pytest.mark.asyncio |
| 718 | - """ | 718 | + async def test_degradation_to_single_node( |
| 719 | - error_message = "Test Internal Server Error" | 719 | + self, monkeypatch: MonkeyPatch, setup_native_handoff, mock_raw_request, client |
| 720 | - max_retry = CoordinatorConfig().exception_config.transport_retry_limit | 720 | + ): |
| 721 | - d_client = _UnifiedPDDecodeClient( | 721 | + """ |
| 722 | - stream_exc=httpx.HTTPStatusError( | 722 | + Test that when no ROLE_D instances are available, the router degrades to SINGLE_NODE mode |
| 723 | - message=error_message, | 723 | + and uses PDHybridRouter. |
| 724 | - request=MagicMock(), | 724 | + """ |
| 725 | - response=httpx.Response( | 725 | + # Let listen_for_disconnect() exit immediately so the task does not hang (avoids WSL Terminated |
| 726 | - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, | 726 | + # when handler and disconnect run concurrently and disconnect awaits a never-completing receive). |
| 727 | - text=error_message, | 727 | + disconnect_msg = {"type": "http.disconnect"} |
| 728 | - ), | 728 | + mock_raw_request.receive = AsyncMock(return_value=disconnect_msg) |
| 729 | - ) | 729 | + |
| 730 | - ) | 730 | + # Mock InstanceManager.get_available_instances |
| 731 | - p_client = _UnifiedPDPrefillClient() | 731 | + host = "127.0.0.1" |
| 732 | - req_info = await create_mock_request_info() | 732 | + mock_instance_p = self.create_mock_instance(0, PDRole.ROLE_P) |
| 733 | - | 733 | + mock_endpoint_p = Endpoint(id=0, ip=host, business_port="8000", mgmt_port="8000") |
| 734 | - exec_release = 0 | 734 | + mock_instance_p.endpoints = {host: {0: mock_endpoint_p}} |
| 735 | - original_release = SeparateCDPRouter._release_attempt_resource | 735 | + |
| 736 | - | 736 | + def mock_get_available_instances(self, role=None): |
| 737 | - async def mock_release_attempt_resource(self, resource, attempt_seq, action, attempt=None, **kwargs): | 737 | + if role is None: |
| 738 | - nonlocal exec_release | 738 | + return {mock_instance_p.id: mock_instance_p} |
| 739 | - exec_release += 1 | 739 | + if role == PDRole.ROLE_U: # PD hybrid role |
| 740 | - await original_release(self, resource, attempt_seq, action, attempt, **kwargs) | 740 | + return {} |
| 741 | - | 741 | + elif role == PDRole.ROLE_P: |
| 742 | - monkeypatch.setattr( | 742 | + return {mock_instance_p.id: mock_instance_p} |
| 743 | - SeparateCDPRouter, | 743 | + elif role == PDRole.ROLE_D: |
| 744 | - "_release_attempt_resource", | 744 | + return {} |
| 745 | - mock_release_attempt_resource, | 745 | + return {} |
| 746 | - ) | 746 | + |
| 747 | - | 747 | + monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances) |
| 748 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 748 | + |
| 749 | - response = await cdp_router.handle_request() | 749 | + # So router chooses SINGLE_NODE (PDHybridRouter) before creating the router |
| 750 | - chunk_str = await _collect_stream_chunks(response) | 750 | + def mock_get_required_instances_status(self): |
| 751 | - | 751 | + return InstanceReadiness.ONLY_PREFILL # not ready -> fallback to SINGLE_NODE |
| 752 | - assert req_info.state == ReqState.EXCEPTION | 752 | + |
| 753 | - _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="UpstreamHTTPError") | 753 | + monkeypatch.setattr( |
| 754 | - assert str(status.HTTP_500_INTERNAL_SERVER_ERROR) in chunk_str | 754 | + InstanceManager, |
| 755 | - assert d_client.stream_count == max_retry | 755 | + "get_required_instances_status", |
| 756 | - assert exec_release >= 1 | 756 | + mock_get_required_instances_status, |
| 757 | - | 757 | + ) |
| 758 | - @pytest.mark.asyncio | 758 | + |
| 759 | - async def test_engine_server_decode_once_5xx_status_code( | 759 | + def mock_has_required_instances(self): |
| 760 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | 760 | + return False |
| 761 | - ): | 761 | + |
| 762 | - """Test case: EngineServer Decode request first returns 5XX, then 200. | 762 | + monkeypatch.setattr(InstanceManager, "has_required_instances", mock_has_required_instances) |
| 763 | - Expected behavior: | 763 | + |
| 764 | - 1) Check request status is DecodeEnd | 764 | + def mock_select_instance_and_endpoint(self, role): |
| 765 | - 2) Trigger request retry | 765 | + if role == PDRole.ROLE_P: |
| 766 | - 3) Request retry succeeds | 766 | + return mock_instance_p, mock_endpoint_p |
| 767 | - """ | 767 | + elif role == PDRole.ROLE_D: |
| 768 | - error_message = "Test Internal Server Error" | 768 | + return None, None |
| 769 | - d_client = _UnifiedPDDecodeClient( | 769 | + return None, None |
| 770 | - stream_exc=httpx.HTTPStatusError( | 770 | + |
| 771 | - message=error_message, | 771 | + monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint) |
| 772 | - request=MagicMock(), | 772 | + |
| 773 | - response=httpx.Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR), | 773 | + # Mock PDHybridRouter response |
| 774 | - ), | 774 | + mock_response = "mock_message" |
| 775 | - stream_fail_times=1, | 775 | + with patch( |
| 776 | - ) | 776 | + "motor.coordinator.router.dispatch.PDHybridRouter.handle_request", |
| 777 | - p_client = _UnifiedPDPrefillClient() | 777 | + new_callable=AsyncMock, |
| 778 | - req_info = await create_mock_request_info() | 778 | + return_value=mock_response, |
| 779 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 779 | + ) as mock_handle_request: |
| 780 | - | 780 | + response = await router.handle_request( |
| 781 | - response = await cdp_router.handle_request() | 781 | + mock_raw_request, |
| 782 | - chunk_str = await _collect_stream_chunks(response) | 782 | + CoordinatorConfig(), |
| 783 | - | 783 | + scheduler=_scheduler, |
| 784 | - assert response.status_code == status.HTTP_200_OK | 784 | + request_manager=_request_manager, |
| 785 | - assert d_client.stream_fail_count == 1 | 785 | + ) |
| 786 | - assert d_client.stream_count >= 2 | 786 | + # Verify PDHybridRouter.handle_request was called |
| 787 | - assert req_info.state == ReqState.DECODE_END | 787 | + mock_handle_request.assert_called_once() |
| 788 | - assert "decoded chunk" in chunk_str | 788 | + # Verify response |
| 789 | - | 789 | + assert response == mock_response |
| 790 | - @pytest.mark.asyncio | 790 | + |
| 791 | - async def test_engine_server_decode_network_exception(self, client, monkeypatch: MonkeyPatch, setup_cdp_separation): | 791 | + @pytest.mark.asyncio |
| 792 | - """Test case: EngineServer Decode network exception | 792 | + async def test_no_degradation_when_d_instances_exist(self, monkeypatch, setup_native_handoff, mock_raw_request): |
| 793 | - Expected behavior: | 793 | + """ |
| 794 | - 1) Check request status is Exception | 794 | + Test that when ROLE_D instances are available, the router uses the configured mode (PD_SEPARATE). |
| 795 | - 2) Retries exhaust transport_retry_limit | 795 | + """ |
| 796 | - 3) Directly return error message | 796 | + # Let listen_for_disconnect() exit immediately (same as test_degradation_to_single_node). |
| 797 | - """ | 797 | + disconnect_msg = {"type": "http.disconnect"} |
| 798 | - error_message = "Connection error" | 798 | + mock_raw_request.receive = AsyncMock(return_value=disconnect_msg) |
| 799 | - max_retry = CoordinatorConfig().exception_config.transport_retry_limit | 799 | + |
| 800 | - d_client = _UnifiedPDDecodeClient( | 800 | + # Mock UnifiedPDRouter response |
| 801 | - stream_exc=httpx.ConnectError(error_message, request=MagicMock()), | 801 | + mock_response = "mock_message" |
| 802 | - stream_fail_times=max_retry, | 802 | + with patch( |
| 803 | - ) | 803 | + "motor.coordinator.router.dispatch.UnifiedPDRouter.handle_request", |
| 804 | - p_client = _UnifiedPDPrefillClient() | 804 | + new_callable=AsyncMock, |
| 805 | - req_info = await create_mock_request_info() | 805 | + return_value=mock_response, |
| 806 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 806 | + ) as mock_handle_request: |
| 807 | - | 807 | + response = await router.handle_request( |
| 808 | - response = await cdp_router.handle_request() | 808 | + mock_raw_request, |
| 809 | - chunk_str = await _collect_stream_chunks(response) | 809 | + CoordinatorConfig(), |
| 810 | - | 810 | + scheduler=_scheduler, |
| 811 | - _assert_stream_error_chunk(chunk_str, error_message=error_message, error_type="ConnectError") | 811 | + request_manager=_request_manager, |
| 812 | - assert d_client.stream_count == max_retry | 812 | + ) |
| 813 | - assert d_client.stream_fail_count == max_retry | 813 | + |
| 814 | - assert req_info.state == ReqState.EXCEPTION | 814 | + # Verify UnifiedPDRouter.handle_request was called |
| 815 | - | 815 | + mock_handle_request.assert_called_once() |
| 816 | - @pytest.mark.asyncio | 816 | + # Verify response |
| 817 | - async def test_cdp_decode_non_stream_retry_exhausts_transport_limit( | 817 | + assert response == mock_response |
| 818 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | 818 | + |
| 819 | - ): | 819 | + @pytest.mark.parametrize( |
| 820 | - """Non-stream decode failures retry whole transport attempts until limit.""" | 820 | + ("prefill_details", "expected_details"), |
| 821 | - error_message = "Same post Decode error every retry" | 821 | + [ |
| 822 | - max_retry = CoordinatorConfig().exception_config.transport_retry_limit | 822 | + ({"cached_tokens": 10}, {"cached_tokens": 10}), |
| 823 | - d_client = _UnifiedPDDecodeClient( | 823 | + (None, {"cached_tokens": 0}), |
| 824 | - post_exc=httpx.HTTPStatusError( | 824 | + ], |
| 825 | - message=error_message, | 825 | + ) |
| 826 | - request=MagicMock(), | 826 | + @pytest.mark.asyncio |
| 827 | - response=httpx.Response(status_code=status.HTTP_502_BAD_GATEWAY, text=error_message), | 827 | + async def test_prompt_tokens_details_propagation( |
| 828 | - ), | 828 | + self, |
| 829 | - post_fail_times=max_retry, | 829 | + client, |
| 830 | - ) | 830 | + monkeypatch: MonkeyPatch, |
| 831 | - p_client = _UnifiedPDPrefillClient() | 831 | + setup_native_handoff, |
| 832 | - req_info = await create_mock_request_info(stream=False) | 832 | + prefill_details, |
| 833 | - | 833 | + expected_details, |
| 834 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | 834 | + ): |
| 835 | - | 835 | + """UnifiedPD returns prompt cache details collected from the prefill response.""" |
| 836 | - with pytest.raises(httpx.HTTPStatusError): | 836 | + req_info = await create_mock_request_info(stream=False) |
| 837 | - await cdp_router.handle_request() | 837 | + |
| 838 | - | 838 | + class _PrefillClient(_UnifiedPDPrefillClient): |
| 839 | - assert d_client.post_count == max_retry | 839 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 840 | - assert d_client.post_fail_count == max_retry | 840 | + self.requests.append(json) |
| 841 | - assert req_info.state == ReqState.EXCEPTION | 841 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 842 | - | 842 | + return httpx.Response( |
| 843 | - @pytest.mark.asyncio | 843 | + status_code=200, |
| 844 | - async def test_engine_server_prefill_network_exception( | 844 | + json=_native_prefill_body( |
| 845 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | 845 | + json, |
| 846 | - ): | 846 | + usage={"prompt_tokens_details": prefill_details}, |
| 847 | - """Test case: EngineServer prefill network exception | 847 | + ), |
| 848 | - Expected behavior: | 848 | + request=request, |
| 849 | - 1) Check request status is Exception | 849 | + ) |
| 850 | - 2) Retries exhaust transport_retry_limit | 850 | + |
| 851 | - 3) Directly return error message | 851 | + class _DecodeClient(_UnifiedPDDecodeClient): |
| 852 | - """ | 852 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 853 | - error_message = "Connection error" | 853 | + self.post_count += 1 |
| 854 | - retry_times = CoordinatorConfig().exception_config.transport_retry_limit | 854 | + if json: |
| 855 | - p_client = _UnifiedPDPrefillClient( | 855 | + self.requests.append(json) |
| 856 | - exc=httpx.ConnectError(message=error_message, request=MagicMock()), | 856 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 857 | - post_fail_times=retry_times, | 857 | + return httpx.Response( |
| 858 | - ) | 858 | + status_code=200, |
| 859 | - d_client = _UnifiedPDDecodeClient() | 859 | + json={ |
| 860 | - _patch_unified_pd_router_clients(monkeypatch, p_client, d_client) | 860 | + "choices": [{"message": {"content": "test response"}}], |
| 861 | - | 861 | + "usage": { |
| 862 | - state: ReqState = None | 862 | + "prompt_tokens": 15, |
| 863 | - | 863 | + "completion_tokens": 1, |
| 864 | - def mock_update_state(self, new_state: ReqState): | 864 | + "total_tokens": 16, |
| 865 | - nonlocal state | 865 | + }, |
| 866 | - state = new_state | 866 | + }, |
| 867 | - | 867 | + request=request, |
| 868 | - monkeypatch.setattr(RequestInfo, "update_state", mock_update_state) | 868 | + ) |
| 869 | - | 869 | + |
| 870 | - with client.stream( | 870 | + p_client = _PrefillClient() |
| 871 | - "POST", | 871 | + d_client = _DecodeClient() |
| 872 | - "/v1/chat/completions", | 872 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 873 | - json={ | 873 | + |
| 874 | - "model": "test-model", | 874 | + response = await handoff_router.handle_request() |
| 875 | - "messages": [{"role": "user", "content": "Hello"}], | 875 | + response_json = response.body.decode() if hasattr(response.body, "decode") else response.body |
| 876 | - }, | 876 | + response_data = json.loads(response_json) |
| 877 | - ) as response: | 877 | + |
| 878 | - chunks = [] | 878 | + assert req_info.prompt_tokens_details == expected_details |
| 879 | - for chunk in response.iter_lines(): | 879 | + assert response_data["choices"][0]["message"]["content"] == "test response" |
| 880 | - chunks.append(chunk) | 880 | + assert response_data["usage"]["prompt_tokens_details"] == expected_details |
| 881 | - chunk_str = "".join(chunks) | 881 | + |
| 882 | - | 882 | + @pytest.mark.asyncio |
| 883 | - assert error_message in chunk_str | 883 | + async def test_stream_prompt_tokens_details_when_recompute_disabled( |
| 884 | - assert p_client.post_fail_count == retry_times | 884 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 885 | - assert state == ReqState.EXCEPTION | 885 | + ): |
| 886 | - | 886 | + """Prompt cache details are independent of the recompute feature switch.""" |
| 887 | - @pytest.mark.asyncio | 887 | + prompt_tokens_details = {"cached_tokens": 10} |
| 888 | - async def test_degradation_to_single_node( | 888 | + req_info = await create_mock_request_info(stream=True) |
| 889 | - self, monkeypatch: MonkeyPatch, setup_cdp_separation, mock_raw_request, client | 889 | + |
| 890 | - ): | 890 | + usage_chunk = b'data: {"choices":[],"usage":{"prompt_tokens":15,"completion_tokens":1,"total_tokens":16}}\n\n' |
| 891 | - """ | 891 | + |
| 892 | - Test that when no ROLE_D instances are available, the router degrades to SINGLE_NODE mode | 892 | + class _DelayedPrefillClient(_UnifiedPDPrefillClient): |
| 893 | - and uses PDHybridRouter. | 893 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 894 | - """ | 894 | + self.requests.append(json) |
| 895 | - # Let listen_for_disconnect() exit immediately so the task does not hang (avoids WSL Terminated | 895 | + await asyncio.sleep(0.01) |
| 896 | - # when handler and disconnect run concurrently and disconnect awaits a never-completing receive). | 896 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 897 | - disconnect_msg = {"type": "http.disconnect"} | 897 | + return httpx.Response( |
| 898 | - mock_raw_request.receive = AsyncMock(return_value=disconnect_msg) | 898 | + status_code=200, |
| 899 | - | 899 | + json=_native_prefill_body( |
| 900 | - # Mock InstanceManager.get_available_instances | 900 | + json, |
| 901 | - host = "127.0.0.1" | 901 | + usage={"prompt_tokens_details": prompt_tokens_details}, |
| 902 | - mock_instance_p = self.create_mock_instance(0, PDRole.ROLE_P) | 902 | + ), |
| 903 | - mock_endpoint_p = Endpoint(id=0, ip=host, business_port="8000", mgmt_port="8000") | 903 | + request=request, |
| 904 | - mock_instance_p.endpoints = {host: {0: mock_endpoint_p}} | 904 | + ) |
| 905 | - | 905 | + |
| 906 | - def mock_get_available_instances(self, role=None): | 906 | + p_client = _DelayedPrefillClient() |
| 907 | - if role is None: | 907 | + d_client = _UnifiedPDDecodeClient(stream_chunks=[usage_chunk]) |
| 908 | - return {mock_instance_p.id: mock_instance_p} | 908 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 909 | - if role == PDRole.ROLE_U: # PD hybrid role | 909 | + handoff_router.config.exception_config.reschedule_enabled = False |
| 910 | - return {} | 910 | + |
| 911 | - elif role == PDRole.ROLE_P: | 911 | + response = await handoff_router.handle_request() |
| 912 | - return {mock_instance_p.id: mock_instance_p} | 912 | + chunks = [chunk async for chunk in response.body_iterator] |
| 913 | - elif role == PDRole.ROLE_D: | 913 | + response_data = json.loads(chunks[0].removeprefix(b"data: ").strip()) |
| 914 | - return {} | 914 | + |
| 915 | - return {} | 915 | + assert response_data["usage"]["prompt_tokens_details"] == prompt_tokens_details |
| 916 | - | 916 | + |
| 917 | - monkeypatch.setattr(InstanceManager, "get_available_instances", mock_get_available_instances) | 917 | + @pytest.mark.asyncio |
| 918 | - | 918 | + async def test_native_handoff_nonstream_strips_token_ids_at_coordinator_exit( |
| 919 | - # So router chooses SINGLE_NODE (PDHybridRouter) before creating the router | 919 | + self, client, monkeypatch: MonkeyPatch, setup_native_handoff |
| 920 | - def mock_get_required_instances_status(self): | 920 | + ): |
| 921 | - return InstanceReadiness.ONLY_PREFILL # not ready -> fallback to SINGLE_NODE | 921 | + """UnifiedPD non-stream strips token ids and normalizes stop_reason before returning to client.""" |
| 922 | - | 922 | + req_info = await create_mock_request_info(stream=False) |
| 923 | - monkeypatch.setattr( | 923 | + req_info.entry_api = req_info.api |
| 924 | - InstanceManager, | 924 | + |
| 925 | - "get_required_instances_status", | 925 | + recomputed_body = { |
| 926 | - mock_get_required_instances_status, | 926 | + "prompt_token_ids": [1, 2], |
| 927 | - ) | 927 | + "choices": [ |
| 928 | - | 928 | + { |
| 929 | - def mock_has_required_instances(self): | 929 | + "message": {"role": "assistant", "content": "partial "}, |
| 930 | - return False | 930 | + "stop_reason": "recomputed", |
| 931 | - | 931 | + "token_ids": [3, 4], |
| 932 | - monkeypatch.setattr(InstanceManager, "has_required_instances", mock_has_required_instances) | 932 | + } |
| 933 | - | 933 | + ], |
| 934 | - def mock_select_instance_and_endpoint(self, role): | 934 | + "usage": {"completion_tokens": 2}, |
| 935 | - if role == PDRole.ROLE_P: | 935 | + } |
| 936 | - return mock_instance_p, mock_endpoint_p | 936 | + |
| 937 | - elif role == PDRole.ROLE_D: | 937 | + class _RecomputedDecodeClient(_UnifiedPDDecodeClient): |
| 938 | - return None, None | 938 | + async def post(self, path, json=None, headers=None, timeout=None): |
| 939 | - return None, None | 939 | + self.post_count += 1 |
| 940 | - | 940 | + if json: |
| 941 | - monkeypatch.setattr(Scheduler, "select_instance_and_endpoint", mock_select_instance_and_endpoint) | 941 | + self.requests.append(json) |
| 942 | - | 942 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 943 | - # Mock PDHybridRouter response | 943 | + return httpx.Response(status_code=200, json=recomputed_body, request=request) |
| 944 | - mock_response = "mock_message" | 944 | + |
| 945 | - with patch( | 945 | + p_client = _UnifiedPDPrefillClient() |
| 946 | - "motor.coordinator.router.dispatch.PDHybridRouter.handle_request", | 946 | + d_client = _RecomputedDecodeClient() |
| 947 | - new_callable=AsyncMock, | 947 | + handoff_router = self._make_router(req_info, monkeypatch, p_client, d_client) |
| 948 | - return_value=mock_response, | 948 | + |
| 949 | - ) as mock_handle_request: | 949 | + response = await handoff_router.handle_request() |
| 950 | - response = await router.handle_request( | 950 | + response_json = response.body.decode() if hasattr(response.body, "decode") else response.body |
| 951 | - mock_raw_request, | 951 | + response_data = json.loads(response_json) |
| 952 | - CoordinatorConfig(), | 952 | + |
| 953 | - scheduler=_scheduler, | 953 | + assert response_data["choices"][0]["message"]["content"] == "partial " |
| 954 | - request_manager=_request_manager, | 954 | + assert response_data["choices"][0]["stop_reason"] == "stop" |
| 955 | - ) | 955 | + assert "token_ids" not in response_data["choices"][0] |
| 956 | - # Verify PDHybridRouter.handle_request was called | 956 | + assert "prompt_token_ids" not in response_data |
| 957 | - mock_handle_request.assert_called_once() | 957 | + assert req_info.state == ReqState.DECODE_END |
| 958 | - # Verify response | ||
| 959 | - assert response == mock_response | ||
| 960 | - | ||
| 961 | - | ||
| 962 | - async def test_no_degradation_when_d_instances_exist(self, monkeypatch, setup_cdp_separation, mock_raw_request): | ||
| 963 | - """ | ||
| 964 | - Test that when ROLE_D instances are available, the router uses the configured mode (PD_SEPARATE). | ||
| 965 | - """ | ||
| 966 | - # Let listen_for_disconnect() exit immediately (same as test_degradation_to_single_node). | ||
| 967 | - disconnect_msg = {"type": "http.disconnect"} | ||
| 968 | - mock_raw_request.receive = AsyncMock(return_value=disconnect_msg) | ||
| 969 | - | ||
| 970 | - # Mock UnifiedPDRouter response | ||
| 971 | - mock_response = "mock_message" | ||
| 972 | - with patch( | ||
| 973 | - "motor.coordinator.router.dispatch.UnifiedPDRouter.handle_request", | ||
| 974 | - new_callable=AsyncMock, | ||
| 975 | - return_value=mock_response, | ||
| 976 | - ) as mock_handle_request: | ||
| 977 | - response = await router.handle_request( | ||
| 978 | - mock_raw_request, | ||
| 979 | - CoordinatorConfig(), | ||
| 980 | - scheduler=_scheduler, | ||
| 981 | - request_manager=_request_manager, | ||
| 982 | - ) | ||
| 983 | - | ||
| 984 | - # Verify UnifiedPDRouter.handle_request was called | ||
| 985 | - mock_handle_request.assert_called_once() | ||
| 986 | - # Verify response | ||
| 987 | - assert response == mock_response | ||
| 988 | - | ||
| 989 | - | ||
| 990 | - ("prefill_details", "expected_details"), | ||
| 991 | - [ | ||
| 992 | - ({"cached_tokens": 10}, {"cached_tokens": 10}), | ||
| 993 | - (None, {"cached_tokens": 0}), | ||
| 994 | - ], | ||
| 995 | - ) | ||
| 996 | - | ||
| 997 | - async def test_prompt_tokens_details_propagation( | ||
| 998 | - self, | ||
| 999 | - client, | ||
| 1000 | - monkeypatch: MonkeyPatch, | ||
| 1001 | - setup_cdp_separation, | ||
| 1002 | - prefill_details, | ||
| 1003 | - expected_details, | ||
| 1004 | - ): | ||
| 1005 | - """UnifiedPD returns prompt cache details collected from the prefill response.""" | ||
| 1006 | - req_info = await create_mock_request_info(stream=False) | ||
| 1007 | - | ||
| 1008 | - class _PrefillClient(_UnifiedPDPrefillClient): | ||
| 1009 | - async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1010 | - self.requests.append(json) | ||
| 1011 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 1012 | - return httpx.Response( | ||
| 1013 | - status_code=200, | ||
| 1014 | - json={"usage": {"prompt_tokens_details": prefill_details}}, | ||
| 1015 | - request=request, | ||
| 1016 | - ) | ||
| 1017 | - | ||
| 1018 | - class _DecodeClient(_UnifiedPDDecodeClient): | ||
| 1019 | - async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1020 | - self.post_count += 1 | ||
| 1021 | - if json: | ||
| 1022 | - self.requests.append(json) | ||
| 1023 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 1024 | - return httpx.Response( | ||
| 1025 | - status_code=200, | ||
| 1026 | - json={ | ||
| 1027 | - "choices": [{"message": {"content": "test response"}}], | ||
| 1028 | - "usage": { | ||
| 1029 | - "prompt_tokens": 15, | ||
| 1030 | - "completion_tokens": 1, | ||
| 1031 | - "total_tokens": 16, | ||
| 1032 | - }, | ||
| 1033 | - }, | ||
| 1034 | - request=request, | ||
| 1035 | - ) | ||
| 1036 | - | ||
| 1037 | - p_client = _PrefillClient() | ||
| 1038 | - d_client = _DecodeClient() | ||
| 1039 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | ||
| 1040 | - | ||
| 1041 | - response = await cdp_router.handle_request() | ||
| 1042 | - response_json = response.body.decode() if hasattr(response.body, "decode") else response.body | ||
| 1043 | - response_data = json.loads(response_json) | ||
| 1044 | - | ||
| 1045 | - assert req_info.prompt_tokens_details == expected_details | ||
| 1046 | - assert response_data["choices"][0]["message"]["content"] == "test response" | ||
| 1047 | - assert response_data["usage"]["prompt_tokens_details"] == expected_details | ||
| 1048 | - | ||
| 1049 | - | ||
| 1050 | - async def test_stream_prompt_tokens_details_when_recompute_disabled( | ||
| 1051 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | ||
| 1052 | - ): | ||
| 1053 | - """Prompt cache details are independent of the recompute feature switch.""" | ||
| 1054 | - prompt_tokens_details = {"cached_tokens": 10} | ||
| 1055 | - req_info = await create_mock_request_info(stream=True) | ||
| 1056 | - | ||
| 1057 | - usage_chunk = b'data: {"choices":[],"usage":{"prompt_tokens":15,"completion_tokens":1,"total_tokens":16}}\n\n' | ||
| 1058 | - | ||
| 1059 | - class _DelayedPrefillClient(_UnifiedPDPrefillClient): | ||
| 1060 | - async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1061 | - self.requests.append(json) | ||
| 1062 | - await asyncio.sleep(0.01) | ||
| 1063 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 1064 | - return httpx.Response( | ||
| 1065 | - status_code=200, | ||
| 1066 | - json={"usage": {"prompt_tokens_details": prompt_tokens_details}}, | ||
| 1067 | - request=request, | ||
| 1068 | - ) | ||
| 1069 | - | ||
| 1070 | - p_client = _DelayedPrefillClient() | ||
| 1071 | - d_client = _UnifiedPDDecodeClient(stream_chunks=[usage_chunk]) | ||
| 1072 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | ||
| 1073 | - cdp_router.config.exception_config.reschedule_enabled = False | ||
| 1074 | - | ||
| 1075 | - response = await cdp_router.handle_request() | ||
| 1076 | - chunks = [chunk async for chunk in response.body_iterator] | ||
| 1077 | - response_data = json.loads(chunks[0].removeprefix(b"data: ").strip()) | ||
| 1078 | - | ||
| 1079 | - assert response_data["usage"]["prompt_tokens_details"] == prompt_tokens_details | ||
| 1080 | - | ||
| 1081 | - | ||
| 1082 | - async def test_cdp_nonstream_recompute_returns_decode_body_as_is( | ||
| 1083 | - self, client, monkeypatch: MonkeyPatch, setup_cdp_separation | ||
| 1084 | - ): | ||
| 1085 | - """UnifiedPD non-stream returns the decode engine body without coordinator-side recompute merge.""" | ||
| 1086 | - req_info = await create_mock_request_info(stream=False) | ||
| 1087 | - req_info.entry_api = req_info.api | ||
| 1088 | - | ||
| 1089 | - recomputed_body = { | ||
| 1090 | - "prompt_token_ids": [1, 2], | ||
| 1091 | - "choices": [ | ||
| 1092 | - { | ||
| 1093 | - "message": {"role": "assistant", "content": "partial "}, | ||
| 1094 | - "stop_reason": "recomputed", | ||
| 1095 | - "token_ids": [3, 4], | ||
| 1096 | - } | ||
| 1097 | - ], | ||
| 1098 | - "usage": {"completion_tokens": 2}, | ||
| 1099 | - } | ||
| 1100 | - | ||
| 1101 | - class _RecomputedDecodeClient(_UnifiedPDDecodeClient): | ||
| 1102 | - async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1103 | - self.post_count += 1 | ||
| 1104 | - if json: | ||
| 1105 | - self.requests.append(json) | ||
| 1106 | - request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 1107 | - return httpx.Response(status_code=200, json=recomputed_body, request=request) | ||
| 1108 | - | ||
| 1109 | - p_client = _UnifiedPDPrefillClient() | ||
| 1110 | - d_client = _RecomputedDecodeClient() | ||
| 1111 | - cdp_router = self._make_router(req_info, monkeypatch, p_client, d_client) | ||
| 1112 | - | ||
| 1113 | - response = await cdp_router.handle_request() | ||
| 1114 | - response_json = response.body.decode() if hasattr(response.body, "decode") else response.body | ||
| 1115 | - response_data = json.loads(response_json) | ||
| 1116 | - | ||
| 1117 | - assert response_data["choices"][0]["message"]["content"] == "partial " | ||
| 1118 | - assert response_data["choices"][0]["stop_reason"] == "recomputed" | ||
| 1119 | - assert req_info.state == ReqState.DECODE_END | ||
| @@ -1,135 +0,0 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | -# MindIE is licensed under Mulan PSL v2. | ||
| 3 | -# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | -# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | -# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | -# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | -# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | -# See the Mulan PSL v2 for more details. | ||
| 10 | - | ||
| 11 | -import pytest | ||
| 12 | - | ||
| 13 | -from motor.common.resources.dispatch import MOTOR_DISPATCH_KEY | ||
| 14 | -from motor.common.resources.endpoint import Endpoint, EndpointStatus | ||
| 15 | -from motor.common.resources.instance import Instance, InsStatus, PDRole, ParallelConfig | ||
| 16 | -from motor.coordinator.domain import ScheduledResource | ||
| 17 | -from motor.coordinator.router.dispatch_session import AttemptContext | ||
| 18 | -from motor.coordinator.router.sglang_native_dispatch import ( | ||
| 19 | - _bootstrap_port, | ||
| 20 | - _stable_bootstrap_room, | ||
| 21 | - ensure_sglang_pd_pair, | ||
| 22 | - inject_sglang_pd_fields, | ||
| 23 | - is_sglang_resource, | ||
| 24 | -) | ||
| 25 | - | ||
| 26 | - | ||
| 27 | -def _resource(engine_type: str | None, *, instance_id: int = 1) -> ScheduledResource: | ||
| 28 | - endpoint = Endpoint( | ||
| 29 | - id=instance_id, | ||
| 30 | - ip="10.0.0.8", | ||
| 31 | - business_port="8000", | ||
| 32 | - mgmt_port="1026", | ||
| 33 | - status=EndpointStatus.NORMAL, | ||
| 34 | - ) | ||
| 35 | - instance = Instance( | ||
| 36 | - job_name=f"job-{instance_id}", | ||
| 37 | - model_name="m", | ||
| 38 | - engine_type=engine_type, | ||
| 39 | - id=instance_id, | ||
| 40 | - role=PDRole.ROLE_P if instance_id == 1 else PDRole.ROLE_D, | ||
| 41 | - status=InsStatus.ACTIVE, | ||
| 42 | - parallel_config=ParallelConfig(dp_size=1), | ||
| 43 | - endpoints={endpoint.ip: {endpoint.id: endpoint}}, | ||
| 44 | - ) | ||
| 45 | - return ScheduledResource(instance=instance, endpoint=endpoint) | ||
| 46 | - | ||
| 47 | - | ||
| 48 | -def test_stable_bootstrap_room_is_stable_and_positive(): | ||
| 49 | - room_a = _stable_bootstrap_room("pair-1", 1) | ||
| 50 | - room_b = _stable_bootstrap_room("pair-1", 1) | ||
| 51 | - room_c = _stable_bootstrap_room("pair-1", 2) | ||
| 52 | - assert room_a == room_b | ||
| 53 | - assert room_a != room_c | ||
| 54 | - assert room_a >= 0 | ||
| 55 | - assert room_a < (1 << 63) | ||
| 56 | - | ||
| 57 | - | ||
| 58 | -def test_bootstrap_port(monkeypatch): | ||
| 59 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "8998") | ||
| 60 | - assert _bootstrap_port() == "8998" | ||
| 61 | - | ||
| 62 | - monkeypatch.delenv("DISAGGREGATION_BOOTSTRAP_PORT", raising=False) | ||
| 63 | - with pytest.raises(RuntimeError, match="must be an integer"): | ||
| 64 | - _bootstrap_port() | ||
| 65 | - | ||
| 66 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "abc") | ||
| 67 | - with pytest.raises(RuntimeError, match="must be an integer"): | ||
| 68 | - _bootstrap_port() | ||
| 69 | - | ||
| 70 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "70000") | ||
| 71 | - with pytest.raises(RuntimeError, match="1-65535"): | ||
| 72 | - _bootstrap_port() | ||
| 73 | - | ||
| 74 | - | ||
| 75 | -def test_is_sglang_resource(): | ||
| 76 | - sglang_p = _resource("sglang", instance_id=1) | ||
| 77 | - vllm_d = _resource("vllm", instance_id=2) | ||
| 78 | - assert is_sglang_resource(sglang_p) is True | ||
| 79 | - assert is_sglang_resource(vllm_d) is False | ||
| 80 | - assert is_sglang_resource(None) is False | ||
| 81 | - | ||
| 82 | - | ||
| 83 | -def test_ensure_sglang_pd_pair_rejects_mixed_engines(): | ||
| 84 | - attempt = AttemptContext( | ||
| 85 | - root_request_id="r1", | ||
| 86 | - attempt_seq=1, | ||
| 87 | - pair_id="pair", | ||
| 88 | - prefill_resource=_resource("sglang", instance_id=1), | ||
| 89 | - decode_resource=_resource("vllm", instance_id=2), | ||
| 90 | - ) | ||
| 91 | - with pytest.raises(RuntimeError, match="both prefill and decode"): | ||
| 92 | - ensure_sglang_pd_pair(attempt) | ||
| 93 | - | ||
| 94 | - | ||
| 95 | -def test_inject_sglang_pd_fields(monkeypatch): | ||
| 96 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "9100") | ||
| 97 | - attempt = AttemptContext( | ||
| 98 | - root_request_id="r1", | ||
| 99 | - attempt_seq=3, | ||
| 100 | - pair_id="pair-xyz", | ||
| 101 | - prefill_resource=_resource("sglang", instance_id=1), | ||
| 102 | - decode_resource=_resource("sglang", instance_id=2), | ||
| 103 | - ) | ||
| 104 | - req = {"model": "m", "prompt": "hi", "request_id": "r1#a3"} | ||
| 105 | - assert inject_sglang_pd_fields(req, attempt) is None | ||
| 106 | - assert req["bootstrap_host"] == "10.0.0.8" | ||
| 107 | - assert req["bootstrap_port"] == "9100" | ||
| 108 | - assert req["bootstrap_room"] == _stable_bootstrap_room("pair-xyz", 3) | ||
| 109 | - assert MOTOR_DISPATCH_KEY not in req | ||
| 110 | - | ||
| 111 | - | ||
| 112 | -def test_inject_sglang_pd_fields_requires_bootstrap_port(monkeypatch): | ||
| 113 | - monkeypatch.delenv("DISAGGREGATION_BOOTSTRAP_PORT", raising=False) | ||
| 114 | - attempt = AttemptContext( | ||
| 115 | - root_request_id="r1", | ||
| 116 | - attempt_seq=1, | ||
| 117 | - pair_id="pair", | ||
| 118 | - prefill_resource=_resource("sglang", instance_id=1), | ||
| 119 | - decode_resource=_resource("sglang", instance_id=2), | ||
| 120 | - ) | ||
| 121 | - with pytest.raises(RuntimeError, match="DISAGGREGATION_BOOTSTRAP_PORT"): | ||
| 122 | - inject_sglang_pd_fields({"model": "m"}, attempt) | ||
| 123 | - | ||
| 124 | - | ||
| 125 | -def test_inject_sglang_pd_fields_rejects_mixed_engines(monkeypatch): | ||
| 126 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "9100") | ||
| 127 | - attempt = AttemptContext( | ||
| 128 | - root_request_id="r1", | ||
| 129 | - attempt_seq=1, | ||
| 130 | - pair_id="pair", | ||
| 131 | - prefill_resource=_resource("sglang", instance_id=1), | ||
| 132 | - decode_resource=_resource("vllm", instance_id=2), | ||
| 133 | - ) | ||
| 134 | - with pytest.raises(RuntimeError, match="both prefill and decode"): | ||
| 135 | - inject_sglang_pd_fields({"model": "m"}, attempt) | ||
| @@ -12,23 +12,17 @@ import asyncio | |||
| 12 | import json | 12 | import json |
| 13 | import logging | 13 | import logging |
| 14 | from contextlib import asynccontextmanager | 14 | from contextlib import asynccontextmanager |
| 15 | -from types import SimpleNamespace | ||
| 16 | from unittest.mock import AsyncMock, MagicMock | 15 | from unittest.mock import AsyncMock, MagicMock |
| 17 | 16 | ||
| 18 | import httpx | 17 | import httpx |
| 19 | import pytest | 18 | import pytest |
| 19 | +from fastapi import HTTPException | ||
| 20 | +from fastapi.responses import JSONResponse | ||
| 20 | from starlette.requests import ClientDisconnect | 21 | from starlette.requests import ClientDisconnect |
| 21 | 22 | ||
| 22 | import motor.common.utils.error as cancel_error | 23 | import motor.common.utils.error as cancel_error |
| 23 | from motor.common.http import HTTPClientPool | 24 | from motor.common.http import HTTPClientPool |
| 24 | from motor.common.logger.logger import _resolve_logger_name | 25 | from motor.common.logger.logger import _resolve_logger_name |
| 25 | -from motor.common.resources.dispatch import ( | ||
| 26 | - DispatchPlan, | ||
| 27 | - DispatchStopReason, | ||
| 28 | - MOTOR_DISPATCH_KEY, | ||
| 29 | - MOTOR_PREFILL_RESULT_KEY, | ||
| 30 | - PrefillContextBudget, | ||
| 31 | -) | ||
| 32 | from motor.common.resources.endpoint import ( | 26 | from motor.common.resources.endpoint import ( |
| 33 | Endpoint, | 27 | Endpoint, |
| 34 | EndpointStatus, | 28 | EndpointStatus, |
| @@ -40,15 +34,13 @@ from motor.config.coordinator import CoordinatorConfig, ExceptionConfig, Schedul | |||
| 40 | from motor.coordinator.domain import ScheduledResource | 34 | from motor.coordinator.domain import ScheduledResource |
| 41 | from motor.coordinator.domain.request_manager import RequestManager | 35 | from motor.coordinator.domain.request_manager import RequestManager |
| 42 | from motor.coordinator.models.request import RequestInfo, ReqState | 36 | from motor.coordinator.models.request import RequestInfo, ReqState |
| 37 | +from motor.coordinator.router.adapters.pd_protocol import EngineProtocolError | ||
| 43 | from motor.coordinator.router.dispatch_session import ( | 38 | from motor.coordinator.router.dispatch_session import ( |
| 44 | AttemptContext, | 39 | AttemptContext, |
| 45 | AttemptState, | 40 | AttemptState, |
| 41 | + AttemptStopReason, | ||
| 46 | PDDispatchSession, | 42 | PDDispatchSession, |
| 47 | ) | 43 | ) |
| 48 | -from motor.coordinator.router.dispatch_capability import ( | ||
| 49 | - DispatchPlanNotSupported, | ||
| 50 | - select_dispatch_plan_for_pair, | ||
| 51 | -) | ||
| 52 | from motor.common.utils.error import RequestCancelledError | 44 | from motor.common.utils.error import RequestCancelledError |
| 53 | from motor.coordinator.router.rescheduler.rescheduler import Rescheduler, RetryRequestPlan | 45 | from motor.coordinator.router.rescheduler.rescheduler import Rescheduler, RetryRequestPlan |
| 54 | from motor.coordinator.router.strategies.unified_pd import UnifiedPDRouter | 46 | from motor.coordinator.router.strategies.unified_pd import UnifiedPDRouter |
| @@ -61,21 +53,21 @@ def _instance( | |||
| 61 | instance_id: int, | 53 | instance_id: int, |
| 62 | role: PDRole, | 54 | role: PDRole, |
| 63 | *, | 55 | *, |
| 64 | - engine_type: str | None = None, | 56 | + engine_type: str = "sglang", |
| 65 | - dispatch_capabilities: list[str] | None = None, | 57 | + bootstrap_port: int | None = None, |
| 66 | ) -> Instance: | 58 | ) -> Instance: |
| 67 | endpoint = Endpoint( | 59 | endpoint = Endpoint( |
| 68 | id=instance_id, | 60 | id=instance_id, |
| 69 | ip="127.0.0.1", | 61 | ip="127.0.0.1", |
| 70 | business_port=str(8100 + instance_id), | 62 | business_port=str(8100 + instance_id), |
| 71 | mgmt_port=str(9100 + instance_id), | 63 | mgmt_port=str(9100 + instance_id), |
| 64 | + bootstrap_port=bootstrap_port, | ||
| 72 | status=EndpointStatus.NORMAL, | 65 | status=EndpointStatus.NORMAL, |
| 73 | ) | 66 | ) |
| 74 | return Instance( | 67 | return Instance( |
| 75 | job_name=f"job-{instance_id}", | 68 | job_name=f"job-{instance_id}", |
| 76 | - model_name=engine_type or "model", | 69 | + model_name=engine_type, |
| 77 | engine_type=engine_type, | 70 | engine_type=engine_type, |
| 78 | - dispatch_capabilities=dispatch_capabilities or [], | ||
| 79 | id=instance_id, | 71 | id=instance_id, |
| 80 | role=role, | 72 | role=role, |
| 81 | status=InsStatus.ACTIVE, | 73 | status=InsStatus.ACTIVE, |
| @@ -88,31 +80,28 @@ class _Scheduler: | |||
| 88 | def __init__( | 80 | def __init__( |
| 89 | self, | 81 | self, |
| 90 | *, | 82 | *, |
| 91 | - prefill_engine_type: str | None = None, | 83 | + prefill_engine_type: str = "sglang", |
| 92 | - decode_engine_type: str | None = None, | 84 | + decode_engine_type: str = "sglang", |
| 93 | - prefill_capabilities: list[str] | None = None, | 85 | + prefill_bootstrap_port: int | None = 20001, |
| 94 | - decode_capabilities: list[str] | None = None, | ||
| 95 | ): | 86 | ): |
| 96 | - if prefill_capabilities is None: | ||
| 97 | - prefill_capabilities = [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | ||
| 98 | - if decode_capabilities is None: | ||
| 99 | - decode_capabilities = [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | ||
| 100 | self.p = _instance( | 87 | self.p = _instance( |
| 101 | 1, | 88 | 1, |
| 102 | PDRole.ROLE_P, | 89 | PDRole.ROLE_P, |
| 103 | engine_type=prefill_engine_type, | 90 | engine_type=prefill_engine_type, |
| 104 | - dispatch_capabilities=prefill_capabilities, | 91 | + bootstrap_port=prefill_bootstrap_port, |
| 105 | ) | 92 | ) |
| 106 | self.d = _instance( | 93 | self.d = _instance( |
| 107 | 2, | 94 | 2, |
| 108 | PDRole.ROLE_D, | 95 | PDRole.ROLE_D, |
| 109 | engine_type=decode_engine_type, | 96 | engine_type=decode_engine_type, |
| 110 | - dispatch_capabilities=decode_capabilities, | ||
| 111 | ) | 97 | ) |
| 112 | self.update_workload = AsyncMock(return_value=True) | 98 | self.update_workload = AsyncMock(return_value=True) |
| 113 | 99 | ||
| 114 | async def select_and_allocate(self, role, req_info, **_kwargs): | 100 | async def select_and_allocate(self, role, req_info, **_kwargs): |
| 115 | instance = self.p if role == PDRole.ROLE_P else self.d | 101 | instance = self.p if role == PDRole.ROLE_P else self.d |
| 102 | + required_engine_type = _kwargs.get("required_engine_type") | ||
| 103 | + if required_engine_type and instance.engine_type != required_engine_type: | ||
| 104 | + return None | ||
| 116 | endpoint = next(iter(next(iter(instance.endpoints.values())).values())) | 105 | endpoint = next(iter(next(iter(instance.endpoints.values())).values())) |
| 117 | return instance, endpoint, Workload(active_tokens=1) | 106 | return instance, endpoint, Workload(active_tokens=1) |
| 118 | 107 | ||
| @@ -129,11 +118,16 @@ class _Client: | |||
| 129 | self.name = name | 118 | self.name = name |
| 130 | self.exc = exc | 119 | self.exc = exc |
| 131 | self.requests = [] | 120 | self.requests = [] |
| 121 | + self.abort_requests = [] | ||
| 132 | self.headers = [] | 122 | self.headers = [] |
| 133 | self.base_url = f"http://{name}" | 123 | self.base_url = f"http://{name}" |
| 134 | self.timeout = 1 | 124 | self.timeout = 1 |
| 135 | 125 | ||
| 136 | async def post(self, path, json=None, headers=None, timeout=None): | 126 | async def post(self, path, json=None, headers=None, timeout=None): |
| 127 | + if str(path).rstrip("/").endswith("abort_request"): | ||
| 128 | + self.abort_requests.append(json) | ||
| 129 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 130 | + return httpx.Response(status_code=200, json={}, request=request) | ||
| 137 | self.requests.append(json) | 131 | self.requests.append(json) |
| 138 | self.headers.append(headers or {}) | 132 | self.headers.append(headers or {}) |
| 139 | if self.exc is not None: | 133 | if self.exc is not None: |
| @@ -142,7 +136,7 @@ class _Client: | |||
| 142 | if self.name == "prefill": | 136 | if self.name == "prefill": |
| 143 | return httpx.Response( | 137 | return httpx.Response( |
| 144 | status_code=200, | 138 | status_code=200, |
| 145 | - json={"status": "cached", "id": json["request_id"]}, | 139 | + json={"status": "cached", "id": json.get("request_id") or json.get("rid")}, |
| 146 | request=request, | 140 | request=request, |
| 147 | ) | 141 | ) |
| 148 | return httpx.Response( | 142 | return httpx.Response( |
| @@ -179,24 +173,54 @@ class _DelayedHTTPErrorClient(_HTTPErrorClient): | |||
| 179 | return await super().post(path, json=json, headers=headers, timeout=timeout) | 173 | return await super().post(path, json=json, headers=headers, timeout=timeout) |
| 180 | 174 | ||
| 181 | 175 | ||
| 182 | -class _PrefillResultClient(_Client): | 176 | +class _NativeHandoffPrefillClient(_Client): |
| 183 | async def post(self, path, json=None, headers=None, timeout=None): | 177 | async def post(self, path, json=None, headers=None, timeout=None): |
| 184 | self.requests.append(json) | 178 | self.requests.append(json) |
| 185 | self.headers.append(headers or {}) | 179 | self.headers.append(headers or {}) |
| 186 | request = httpx.Request("POST", path, headers=headers or {}, json=json) | 180 | request = httpx.Request("POST", path, headers=headers or {}, json=json) |
| 187 | - dispatch = json[MOTOR_DISPATCH_KEY] | ||
| 188 | return httpx.Response( | 181 | return httpx.Response( |
| 189 | status_code=200, | 182 | status_code=200, |
| 190 | json={ | 183 | json={ |
| 191 | - "object": "motor.prefill_result", | 184 | + "kv_transfer_params": { |
| 192 | - "schema_version": "1.0", | 185 | + "do_remote_prefill": True, |
| 193 | - "root_request_id": dispatch["root_request_id"], | 186 | + "remote_request_id": json["request_id"], |
| 194 | - "engine_request_id": dispatch["engine_request_id"], | 187 | + "remote_host": "10.0.0.1", |
| 195 | - "pair_id": dispatch["pair_id"], | 188 | + "remote_port": 9000, |
| 196 | - "attempt_seq": dispatch["attempt_seq"], | 189 | + "connector_private": {"opaque": "kv"}, |
| 197 | - "status": "completed", | 190 | + }, |
| 198 | - "handoff_mode": "handoff", | 191 | + "usage": { |
| 199 | - "payload": {"opaque": "kv"}, | 192 | + "prompt_tokens": 5, |
| 193 | + "prompt_tokens_details": {"cached_tokens": 2}, | ||
| 194 | + }, | ||
| 195 | + }, | ||
| 196 | + request=request, | ||
| 197 | + ) | ||
| 198 | + | ||
| 199 | + | ||
| 200 | +class _NativeSglangPrefillClient(_Client): | ||
| 201 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 202 | + self.requests.append(json) | ||
| 203 | + self.headers.append(headers or {}) | ||
| 204 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 205 | + if json.get("stream"): | ||
| 206 | + return httpx.Response( | ||
| 207 | + status_code=200, | ||
| 208 | + content=( | ||
| 209 | + b'data: {"choices":[],"usage":{"prompt_tokens":5,' | ||
| 210 | + b'"prompt_tokens_details":{"cached_tokens":2}}}\n\n' | ||
| 211 | + b"data: [DONE]\n\n" | ||
| 212 | + ), | ||
| 213 | + headers={"content-type": "text/event-stream"}, | ||
| 214 | + request=request, | ||
| 215 | + ) | ||
| 216 | + return httpx.Response( | ||
| 217 | + status_code=200, | ||
| 218 | + json={ | ||
| 219 | + "id": json["rid"], | ||
| 220 | + "usage": { | ||
| 221 | + "prompt_tokens": 5, | ||
| 222 | + "prompt_tokens_details": {"cached_tokens": 2}, | ||
| 223 | + }, | ||
| 200 | }, | 224 | }, |
| 201 | request=request, | 225 | request=request, |
| 202 | ) | 226 | ) |
| @@ -403,11 +427,11 @@ async def _invoke_asgi_response(response) -> list[dict]: | |||
| 403 | [ | 427 | [ |
| 404 | ( | 428 | ( |
| 405 | f"{cancel_error.NODE_FAULT}: http://127.0.0.1:8102", | 429 | f"{cancel_error.NODE_FAULT}: http://127.0.0.1:8102", |
| 406 | - DispatchStopReason.PEER_FAILED, | 430 | + AttemptStopReason.PEER_FAILED, |
| 407 | ), | 431 | ), |
| 408 | - (cancel_error.CLIENT_DISCONNECT, DispatchStopReason.CLIENT_DISCONNECT), | 432 | + (cancel_error.CLIENT_DISCONNECT, AttemptStopReason.CLIENT_DISCONNECT), |
| 409 | - (cancel_error.DISPATCH_ABORT, DispatchStopReason.OTHER), | 433 | + (cancel_error.DISPATCH_ABORT, AttemptStopReason.OTHER), |
| 410 | - (cancel_error.SCOPE_ABORT, DispatchStopReason.OTHER), | 434 | + (cancel_error.SCOPE_ABORT, AttemptStopReason.OTHER), |
| 411 | ], | 435 | ], |
| 412 | ) | 436 | ) |
| 413 | def test_unified_pd_cancel_stop_reason_mapping(reason, expected): | 437 | def test_unified_pd_cancel_stop_reason_mapping(reason, expected): |
| @@ -454,236 +478,8 @@ async def test_unified_pd_process_response_error_wraps_cancelled_as_request_canc | |||
| 454 | assert cancel_error.CLIENT_DISCONNECT in caplog.text | 478 | assert cancel_error.CLIENT_DISCONNECT in caplog.text |
| 455 | 479 | ||
| 456 | 480 | ||
| 457 | -def test_dispatch_plan_prefers_explicit_capability_over_engine_fallback(): | ||
| 458 | - scheduler = _Scheduler( | ||
| 459 | - prefill_engine_type="vllm", | ||
| 460 | - decode_engine_type="sglang", | ||
| 461 | - prefill_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 462 | - decode_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 463 | - ) | ||
| 464 | - p_endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) | ||
| 465 | - d_endpoint = next(iter(next(iter(scheduler.d.endpoints.values())).values())) | ||
| 466 | - | ||
| 467 | - plan = select_dispatch_plan_for_pair( | ||
| 468 | - prefill=ScheduledResource(instance=scheduler.p, endpoint=p_endpoint), | ||
| 469 | - decode=ScheduledResource(instance=scheduler.d, endpoint=d_endpoint), | ||
| 470 | - ) | ||
| 471 | - | ||
| 472 | - assert plan == DispatchPlan.CONCURRENT_ENGINE_SYNC | ||
| 473 | - | ||
| 474 | - | ||
| 475 | -def test_dispatch_plan_requires_connector_capability(): | ||
| 476 | - scheduler = _Scheduler(prefill_capabilities=[], decode_capabilities=[]) | ||
| 477 | - p_endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) | ||
| 478 | - d_endpoint = next(iter(next(iter(scheduler.d.endpoints.values())).values())) | ||
| 479 | - | ||
| 480 | - with pytest.raises(DispatchPlanNotSupported, match="do not advertise"): | ||
| 481 | - select_dispatch_plan_for_pair( | ||
| 482 | - prefill=ScheduledResource(instance=scheduler.p, endpoint=p_endpoint), | ||
| 483 | - decode=ScheduledResource(instance=scheduler.d, endpoint=d_endpoint), | ||
| 484 | - ) | ||
| 485 | - | ||
| 486 | - | ||
| 487 | -def test_dispatch_plan_requires_capability_from_both_instances(): | ||
| 488 | - scheduler = _Scheduler( | ||
| 489 | - prefill_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 490 | - decode_capabilities=[], | ||
| 491 | - ) | ||
| 492 | - p_endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) | ||
| 493 | - d_endpoint = next(iter(next(iter(scheduler.d.endpoints.values())).values())) | ||
| 494 | - | ||
| 495 | - with pytest.raises(DispatchPlanNotSupported, match="do not advertise"): | ||
| 496 | - select_dispatch_plan_for_pair( | ||
| 497 | - prefill=ScheduledResource(instance=scheduler.p, endpoint=p_endpoint), | ||
| 498 | - decode=ScheduledResource(instance=scheduler.d, endpoint=d_endpoint), | ||
| 499 | - ) | ||
| 500 | - | ||
| 501 | - | ||
| 502 | 481 | ||
| 503 | -async def test_unified_pd_nonstream_dispatches_prefill_and_decode_with_same_attempt( | 482 | +async def test_unified_pd_vllm_uses_native_handoff_before_selecting_decode(monkeypatch): |
| 504 | - monkeypatch, | ||
| 505 | -): | ||
| 506 | - req_info = RequestInfo( | ||
| 507 | - req_id="root-1", | ||
| 508 | - req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 509 | - api="v1/completions", | ||
| 510 | - entry_api="v1/completions", | ||
| 511 | - req_len=10, | ||
| 512 | - ) | ||
| 513 | - scheduler = _Scheduler() | ||
| 514 | - router = UnifiedPDRouter( | ||
| 515 | - req_info, | ||
| 516 | - _config(), | ||
| 517 | - scheduler=scheduler, | ||
| 518 | - request_manager=RequestManager(_config()), | ||
| 519 | - ) | ||
| 520 | - p_client = _Client("prefill") | ||
| 521 | - d_client = _Client("decode") | ||
| 522 | - | ||
| 523 | - | ||
| 524 | - async def _client_for(resource: ScheduledResource): | ||
| 525 | - if resource.instance.role == PDRole.ROLE_P: | ||
| 526 | - yield p_client | ||
| 527 | - else: | ||
| 528 | - yield d_client | ||
| 529 | - | ||
| 530 | - monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 531 | - | ||
| 532 | - response = await router.handle_request() | ||
| 533 | - | ||
| 534 | - assert response.body == b'{"choices":[{"message":{"role":"assistant","content":"ok"}}]}' | ||
| 535 | - assert len(p_client.requests) == 1 | ||
| 536 | - assert len(d_client.requests) == 1 | ||
| 537 | - | ||
| 538 | - p_dispatch = p_client.requests[0][MOTOR_DISPATCH_KEY] | ||
| 539 | - d_dispatch = d_client.requests[0][MOTOR_DISPATCH_KEY] | ||
| 540 | - assert p_dispatch["role"] == "prefill" | ||
| 541 | - assert d_dispatch["role"] == "decode" | ||
| 542 | - assert p_dispatch["root_request_id"] == "root-1" | ||
| 543 | - assert d_dispatch["root_request_id"] == "root-1" | ||
| 544 | - assert p_dispatch["attempt_seq"] == d_dispatch["attempt_seq"] == 1 | ||
| 545 | - assert p_dispatch["pair_id"] == d_dispatch["pair_id"] | ||
| 546 | - assert p_client.requests[0]["request_id"] == "root-1#a1" | ||
| 547 | - assert d_client.requests[0]["request_id"] == "root-1#a1" | ||
| 548 | - assert p_client.headers[0]["X-Request-Id"] == "root-1#a1" | ||
| 549 | - assert d_client.headers[0]["X-Request-Id"] == "root-1#a1" | ||
| 550 | - assert scheduler.update_workload.await_count == 2 | ||
| 551 | - assert ReqState.PREFILL_END in req_info.status | ||
| 552 | - assert req_info.status[ReqState.P_ALLOCATED] <= req_info.status[ReqState.PREFILL_END] | ||
| 553 | - assert req_info.status[ReqState.PREFILL_END] <= req_info.status[ReqState.DECODE_END] | ||
| 554 | - | ||
| 555 | - | ||
| 556 | - | ||
| 557 | -async def test_unified_pd_decode_failure_stops_both_legs(monkeypatch): | ||
| 558 | - req_info = RequestInfo( | ||
| 559 | - req_id="root-stop", | ||
| 560 | - req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 561 | - api="v1/completions", | ||
| 562 | - entry_api="v1/completions", | ||
| 563 | - req_len=10, | ||
| 564 | - ) | ||
| 565 | - req_info.trace_obj.set_trace_prompt = MagicMock() | ||
| 566 | - scheduler = _Scheduler() | ||
| 567 | - router = UnifiedPDRouter( | ||
| 568 | - req_info, | ||
| 569 | - _config(), | ||
| 570 | - scheduler=scheduler, | ||
| 571 | - request_manager=RequestManager(_config()), | ||
| 572 | - ) | ||
| 573 | - p_client = _Client("prefill") | ||
| 574 | - d_client = _Client("decode", exc=httpx.ConnectError("decode down")) | ||
| 575 | - stop_calls = [] | ||
| 576 | - | ||
| 577 | - | ||
| 578 | - async def _client_for(resource: ScheduledResource): | ||
| 579 | - if resource.instance.role == PDRole.ROLE_P: | ||
| 580 | - yield p_client | ||
| 581 | - else: | ||
| 582 | - yield d_client | ||
| 583 | - | ||
| 584 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 585 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason.value)) | ||
| 586 | - return None | ||
| 587 | - | ||
| 588 | - monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 589 | - monkeypatch.setattr( | ||
| 590 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 591 | - _stop, | ||
| 592 | - ) | ||
| 593 | - | ||
| 594 | - with pytest.raises(httpx.ConnectError): | ||
| 595 | - await router.handle_request() | ||
| 596 | - | ||
| 597 | - assert len(stop_calls) == 2 | ||
| 598 | - req_info.trace_obj.set_trace_prompt.assert_called_with(req_info.req_data) | ||
| 599 | - assert {call[0] for call in stop_calls} == {PDRole.ROLE_P, PDRole.ROLE_D} | ||
| 600 | - assert all(call[1] == 1 for call in stop_calls) | ||
| 601 | - assert scheduler.update_workload.await_count == 2 | ||
| 602 | - | ||
| 603 | - | ||
| 604 | - | ||
| 605 | -async def test_unified_pd_dual_dispatch_uses_dispatch_context_not_bootstrap_fields( | ||
| 606 | - monkeypatch, | ||
| 607 | -): | ||
| 608 | - req_info = RequestInfo( | ||
| 609 | - req_id="root-vllm-concurrent", | ||
| 610 | - req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 611 | - api="v1/completions", | ||
| 612 | - entry_api="v1/completions", | ||
| 613 | - req_len=10, | ||
| 614 | - ) | ||
| 615 | - scheduler = _Scheduler(prefill_engine_type="vllm", decode_engine_type="vllm") | ||
| 616 | - router = UnifiedPDRouter( | ||
| 617 | - req_info, | ||
| 618 | - _config(), | ||
| 619 | - scheduler=scheduler, | ||
| 620 | - request_manager=RequestManager(_config()), | ||
| 621 | - ) | ||
| 622 | - p_client = _Client("prefill") | ||
| 623 | - d_client = _Client("decode") | ||
| 624 | - | ||
| 625 | - | ||
| 626 | - async def _client_for(resource: ScheduledResource): | ||
| 627 | - if resource.instance.role == PDRole.ROLE_P: | ||
| 628 | - yield p_client | ||
| 629 | - else: | ||
| 630 | - yield d_client | ||
| 631 | - | ||
| 632 | - monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 633 | - | ||
| 634 | - await router.handle_request() | ||
| 635 | - | ||
| 636 | - for request_body in (p_client.requests[0], d_client.requests[0]): | ||
| 637 | - assert "bootstrap_host" not in request_body | ||
| 638 | - assert "bootstrap_port" not in request_body | ||
| 639 | - assert "bootstrap_room" not in request_body | ||
| 640 | - assert request_body[MOTOR_DISPATCH_KEY]["dispatch_mode"] == "pd_pair" | ||
| 641 | - | ||
| 642 | - | ||
| 643 | - | ||
| 644 | -async def test_unified_pd_sglang_injects_bootstrap_fields_not_motor_dispatch( | ||
| 645 | - monkeypatch, | ||
| 646 | -): | ||
| 647 | - """SGLang pure-native PD: Coordinator injects stock bootstrap_* on both legs.""" | ||
| 648 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "9100") | ||
| 649 | - req_info = RequestInfo( | ||
| 650 | - req_id="root-sglang", | ||
| 651 | - req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 652 | - api="v1/completions", | ||
| 653 | - entry_api="v1/completions", | ||
| 654 | - req_len=10, | ||
| 655 | - ) | ||
| 656 | - scheduler = _Scheduler(prefill_engine_type="sglang", decode_engine_type="sglang") | ||
| 657 | - router = UnifiedPDRouter( | ||
| 658 | - req_info, | ||
| 659 | - _config(), | ||
| 660 | - scheduler=scheduler, | ||
| 661 | - request_manager=RequestManager(_config()), | ||
| 662 | - ) | ||
| 663 | - p_client = _Client("prefill") | ||
| 664 | - d_client = _Client("decode") | ||
| 665 | - | ||
| 666 | - | ||
| 667 | - async def _client_for(resource: ScheduledResource): | ||
| 668 | - if resource.instance.role == PDRole.ROLE_P: | ||
| 669 | - yield p_client | ||
| 670 | - else: | ||
| 671 | - yield d_client | ||
| 672 | - | ||
| 673 | - monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 674 | - | ||
| 675 | - await router.handle_request() | ||
| 676 | - | ||
| 677 | - for request_body in (p_client.requests[0], d_client.requests[0]): | ||
| 678 | - assert MOTOR_DISPATCH_KEY not in request_body | ||
| 679 | - assert request_body["bootstrap_host"] == "127.0.0.1" | ||
| 680 | - assert request_body["bootstrap_port"] == "9100" | ||
| 681 | - assert isinstance(request_body["bootstrap_room"], int) | ||
| 682 | - assert p_client.requests[0]["bootstrap_room"] == d_client.requests[0]["bootstrap_room"] | ||
| 683 | - | ||
| 684 | - | ||
| 685 | - | ||
| 686 | -async def test_unified_pd_cpcd_waits_for_prefill_result_before_decode(monkeypatch): | ||
| 687 | req_info = RequestInfo( | 483 | req_info = RequestInfo( |
| 688 | req_id="root-cpcd", | 484 | req_id="root-cpcd", |
| 689 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | 485 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, |
| @@ -691,8 +487,10 @@ async def test_unified_pd_cpcd_waits_for_prefill_result_before_decode(monkeypatc | |||
| 691 | entry_api="v1/completions", | 487 | entry_api="v1/completions", |
| 692 | req_len=10, | 488 | req_len=10, |
| 693 | ) | 489 | ) |
| 694 | - handoff = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | 490 | + scheduler = _Scheduler( |
| 695 | - scheduler = _Scheduler(prefill_capabilities=handoff, decode_capabilities=handoff) | 491 | + prefill_engine_type="vllm", |
| 492 | + decode_engine_type="vllm", | ||
| 493 | + ) | ||
| 696 | events = [] | 494 | events = [] |
| 697 | select_and_allocate = scheduler.select_and_allocate | 495 | select_and_allocate = scheduler.select_and_allocate |
| 698 | 496 | ||
| @@ -713,10 +511,10 @@ async def test_unified_pd_cpcd_waits_for_prefill_result_before_decode(monkeypatc | |||
| 713 | request_manager=RequestManager(_config()), | 511 | request_manager=RequestManager(_config()), |
| 714 | ) | 512 | ) |
| 715 | 513 | ||
| 716 | - class _RecordingPrefillClient(_PrefillResultClient): | 514 | + class _RecordingPrefillClient(_NativeHandoffPrefillClient): |
| 717 | async def post(self, path, json=None, headers=None, timeout=None): | 515 | async def post(self, path, json=None, headers=None, timeout=None): |
| 718 | response = await super().post(path, json=json, headers=headers, timeout=timeout) | 516 | response = await super().post(path, json=json, headers=headers, timeout=timeout) |
| 719 | - events.append(("prefill_result",)) | 517 | + events.append(("native_prefill_result",)) |
| 720 | return response | 518 | return response |
| 721 | 519 | ||
| 722 | p_client = _RecordingPrefillClient("prefill") | 520 | p_client = _RecordingPrefillClient("prefill") |
| @@ -735,26 +533,254 @@ async def test_unified_pd_cpcd_waits_for_prefill_result_before_decode(monkeypatc | |||
| 735 | 533 | ||
| 736 | assert len(p_client.requests) == 1 | 534 | assert len(p_client.requests) == 1 |
| 737 | assert len(d_client.requests) == 1 | 535 | assert len(d_client.requests) == 1 |
| 738 | - p_dispatch = p_client.requests[0][MOTOR_DISPATCH_KEY] | 536 | + p_request = p_client.requests[0] |
| 739 | - d_dispatch = d_client.requests[0][MOTOR_DISPATCH_KEY] | 537 | + d_request = d_client.requests[0] |
| 740 | - assert p_dispatch["attempt_seq"] == d_dispatch["attempt_seq"] == 1 | 538 | + assert "_motor_dispatch" not in p_request |
| 741 | - assert p_dispatch["pair_id"] == d_dispatch["pair_id"] | 539 | + assert "_motor_dispatch" not in d_request |
| 742 | - prefill_result = d_client.requests[0][MOTOR_PREFILL_RESULT_KEY] | 540 | + assert "_motor_prefill_result" not in d_request |
| 743 | - assert prefill_result["status"] == "completed" | 541 | + assert p_request["request_id"] == d_request["request_id"] == "root-cpcd#a1" |
| 744 | - assert prefill_result["handoff_mode"] == "handoff" | 542 | + assert p_request["stream"] is False |
| 745 | - assert prefill_result["payload"] == {"opaque": "kv"} | 543 | + assert p_request["max_tokens"] == 1 |
| 544 | + assert p_request["min_tokens"] == 1 | ||
| 545 | + assert d_request["max_tokens"] == 8 | ||
| 546 | + assert d_request["kv_transfer_params"] == { | ||
| 547 | + "do_remote_prefill": True, | ||
| 548 | + "remote_request_id": "root-cpcd#a1", | ||
| 549 | + "remote_host": "10.0.0.1", | ||
| 550 | + "remote_port": 9000, | ||
| 551 | + "connector_private": {"opaque": "kv"}, | ||
| 552 | + } | ||
| 553 | + assert p_client.headers[0]["X-Request-Id"] == "root-cpcd#a1" | ||
| 554 | + assert d_client.headers[0]["X-Request-Id"] == "root-cpcd#a1" | ||
| 746 | assert scheduler.update_workload.await_count == 2 | 555 | assert scheduler.update_workload.await_count == 2 |
| 747 | assert [event for event in events if event[0] == "select"] == [ | 556 | assert [event for event in events if event[0] == "select"] == [ |
| 748 | ("select", PDRole.ROLE_P), | 557 | ("select", PDRole.ROLE_P), |
| 749 | ("select", PDRole.ROLE_D), | 558 | ("select", PDRole.ROLE_D), |
| 750 | ] | 559 | ] |
| 751 | - assert events.index(("prefill_result",)) < events.index(("select", PDRole.ROLE_D)) | 560 | + assert events.index(("native_prefill_result",)) < events.index(("select", PDRole.ROLE_D)) |
| 752 | assert ("release", PDRole.ROLE_P, WorkloadAction.RELEASE_TOKENS) in events | 561 | assert ("release", PDRole.ROLE_P, WorkloadAction.RELEASE_TOKENS) in events |
| 562 | + assert req_info.prompt_tokens_details == {"cached_tokens": 2} | ||
| 753 | assert ReqState.PREFILL_END in req_info.status | 563 | assert ReqState.PREFILL_END in req_info.status |
| 754 | assert req_info.status[ReqState.P_ALLOCATED] <= req_info.status[ReqState.PREFILL_END] | 564 | assert req_info.status[ReqState.P_ALLOCATED] <= req_info.status[ReqState.PREFILL_END] |
| 755 | assert req_info.status[ReqState.PREFILL_END] <= req_info.status[ReqState.DECODE_END] | 565 | assert req_info.status[ReqState.PREFILL_END] <= req_info.status[ReqState.DECODE_END] |
| 756 | 566 | ||
| 757 | 567 | ||
| 568 | + | ||
| 569 | +async def test_unified_pd_vllm_rejects_missing_native_ticket_without_decode_or_peer_stop( | ||
| 570 | + monkeypatch, | ||
| 571 | +): | ||
| 572 | + req_info = RequestInfo( | ||
| 573 | + req_id="root-missing-ticket", | ||
| 574 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 575 | + api="v1/completions", | ||
| 576 | + entry_api="v1/completions", | ||
| 577 | + req_len=10, | ||
| 578 | + ) | ||
| 579 | + scheduler = _Scheduler( | ||
| 580 | + prefill_engine_type="vllm", | ||
| 581 | + decode_engine_type="vllm", | ||
| 582 | + ) | ||
| 583 | + select_and_allocate = AsyncMock(side_effect=scheduler.select_and_allocate) | ||
| 584 | + scheduler.select_and_allocate = select_and_allocate | ||
| 585 | + router = UnifiedPDRouter( | ||
| 586 | + req_info, | ||
| 587 | + _config(), | ||
| 588 | + scheduler=scheduler, | ||
| 589 | + request_manager=RequestManager(_config()), | ||
| 590 | + ) | ||
| 591 | + p_client = _Client("prefill") | ||
| 592 | + d_client = _Client("decode") | ||
| 593 | + | ||
| 594 | + | ||
| 595 | + async def _client_for(resource: ScheduledResource): | ||
| 596 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 597 | + yield p_client | ||
| 598 | + else: | ||
| 599 | + yield d_client | ||
| 600 | + | ||
| 601 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 602 | + | ||
| 603 | + with pytest.raises(UpstreamHTTPError) as exc_info: | ||
| 604 | + await router.handle_request() | ||
| 605 | + | ||
| 606 | + assert exc_info.value.status_code == 502 | ||
| 607 | + assert b"Missing kv_transfer_params" in exc_info.value.body | ||
| 608 | + assert select_and_allocate.await_count == 1 | ||
| 609 | + assert len(p_client.requests) == 1 | ||
| 610 | + assert d_client.requests == [] | ||
| 611 | + assert scheduler.update_workload.await_count == 1 | ||
| 612 | + release = scheduler.update_workload.await_args.args[0] | ||
| 613 | + assert release.role == PDRole.ROLE_P | ||
| 614 | + assert release.workload_action == WorkloadAction.RELEASE_TOKENS | ||
| 615 | + | ||
| 616 | + | ||
| 617 | + | ||
| 618 | +async def test_unified_pd_vllm_strips_native_internal_fields_from_nonstream_response( | ||
| 619 | + monkeypatch, | ||
| 620 | +): | ||
| 621 | + req_info = RequestInfo( | ||
| 622 | + req_id="root-native-strip", | ||
| 623 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 624 | + api="v1/completions", | ||
| 625 | + entry_api="v1/completions", | ||
| 626 | + req_len=10, | ||
| 627 | + ) | ||
| 628 | + scheduler = _Scheduler( | ||
| 629 | + prefill_engine_type="vllm", | ||
| 630 | + decode_engine_type="vllm", | ||
| 631 | + ) | ||
| 632 | + router = UnifiedPDRouter( | ||
| 633 | + req_info, | ||
| 634 | + _config(), | ||
| 635 | + scheduler=scheduler, | ||
| 636 | + request_manager=RequestManager(_config()), | ||
| 637 | + ) | ||
| 638 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 639 | + | ||
| 640 | + class _DecodeClient(_Client): | ||
| 641 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 642 | + self.requests.append(json) | ||
| 643 | + self.headers.append(headers or {}) | ||
| 644 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 645 | + return httpx.Response( | ||
| 646 | + status_code=200, | ||
| 647 | + json={ | ||
| 648 | + "choices": [{"text": "ok", "index": 0}], | ||
| 649 | + "kv_transfer_params": {"remote_host": "10.0.0.1"}, | ||
| 650 | + }, | ||
| 651 | + request=request, | ||
| 652 | + ) | ||
| 653 | + | ||
| 654 | + d_client = _DecodeClient("decode") | ||
| 655 | + | ||
| 656 | + | ||
| 657 | + async def _client_for(resource: ScheduledResource): | ||
| 658 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 659 | + yield p_client | ||
| 660 | + else: | ||
| 661 | + yield d_client | ||
| 662 | + | ||
| 663 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 664 | + | ||
| 665 | + response = await router.handle_request() | ||
| 666 | + body = json.loads(response.body) | ||
| 667 | + | ||
| 668 | + assert body["choices"][0]["text"] == "ok" | ||
| 669 | + assert "kv_transfer_params" not in body | ||
| 670 | + assert d_client.requests[0]["kv_transfer_params"]["remote_host"] == "10.0.0.1" | ||
| 671 | + | ||
| 672 | + | ||
| 673 | + | ||
| 674 | +async def test_unified_pd_vllm_native_handoff_preserves_chat_output_budget_and_shape( | ||
| 675 | + monkeypatch, | ||
| 676 | +): | ||
| 677 | + req_info = RequestInfo( | ||
| 678 | + req_id="root-native-chat", | ||
| 679 | + req_data={ | ||
| 680 | + "model": "m", | ||
| 681 | + "messages": [{"role": "user", "content": "hello"}], | ||
| 682 | + "stream": False, | ||
| 683 | + "max_completion_tokens": 6, | ||
| 684 | + }, | ||
| 685 | + api="v1/chat/completions", | ||
| 686 | + entry_api="v1/chat/completions", | ||
| 687 | + req_len=10, | ||
| 688 | + ) | ||
| 689 | + scheduler = _Scheduler( | ||
| 690 | + prefill_engine_type="vllm", | ||
| 691 | + decode_engine_type="vllm", | ||
| 692 | + ) | ||
| 693 | + router = UnifiedPDRouter( | ||
| 694 | + req_info, | ||
| 695 | + _config(), | ||
| 696 | + scheduler=scheduler, | ||
| 697 | + request_manager=RequestManager(_config()), | ||
| 698 | + ) | ||
| 699 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 700 | + | ||
| 701 | + class _CompletionDecodeClient(_Client): | ||
| 702 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 703 | + self.requests.append(json) | ||
| 704 | + self.headers.append(headers or {}) | ||
| 705 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 706 | + return httpx.Response( | ||
| 707 | + status_code=200, | ||
| 708 | + json={ | ||
| 709 | + "object": "text_completion", | ||
| 710 | + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop"}], | ||
| 711 | + }, | ||
| 712 | + request=request, | ||
| 713 | + ) | ||
| 714 | + | ||
| 715 | + d_client = _CompletionDecodeClient("decode") | ||
| 716 | + | ||
| 717 | + | ||
| 718 | + async def _client_for(resource: ScheduledResource): | ||
| 719 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 720 | + yield p_client | ||
| 721 | + else: | ||
| 722 | + yield d_client | ||
| 723 | + | ||
| 724 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 725 | + | ||
| 726 | + response = await router.handle_request() | ||
| 727 | + body = json.loads(response.body) | ||
| 728 | + | ||
| 729 | + assert p_client.requests[0]["max_tokens"] == 1 | ||
| 730 | + assert p_client.requests[0]["max_completion_tokens"] == 1 | ||
| 731 | + assert d_client.requests[0]["max_completion_tokens"] == 6 | ||
| 732 | + assert body["object"] == "chat.completion" | ||
| 733 | + assert body["choices"][0]["message"] == {"role": "assistant", "content": "ok"} | ||
| 734 | + | ||
| 735 | + | ||
| 736 | + | ||
| 737 | +async def test_unified_pd_vllm_strips_native_internal_fields_from_stream_response( | ||
| 738 | + monkeypatch, | ||
| 739 | +): | ||
| 740 | + req_info = RequestInfo( | ||
| 741 | + req_id="root-native-stream-strip", | ||
| 742 | + req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | ||
| 743 | + api="v1/completions", | ||
| 744 | + entry_api="v1/completions", | ||
| 745 | + req_len=10, | ||
| 746 | + ) | ||
| 747 | + scheduler = _Scheduler( | ||
| 748 | + prefill_engine_type="vllm", | ||
| 749 | + decode_engine_type="vllm", | ||
| 750 | + ) | ||
| 751 | + router = UnifiedPDRouter( | ||
| 752 | + req_info, | ||
| 753 | + _config(), | ||
| 754 | + scheduler=scheduler, | ||
| 755 | + request_manager=RequestManager(_config()), | ||
| 756 | + ) | ||
| 757 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 758 | + d_client = _SequenceStreamClient( | ||
| 759 | + "decode", | ||
| 760 | + [ | ||
| 761 | + _StreamResponse( | ||
| 762 | + [(b'data: {"choices":[{"text":"A","index":0}],"kv_transfer_params":{"remote_host":"10.0.0.1"}}\n\n')] | ||
| 763 | + ) | ||
| 764 | + ], | ||
| 765 | + ) | ||
| 766 | + | ||
| 767 | + | ||
| 768 | + async def _client_for(resource: ScheduledResource): | ||
| 769 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 770 | + yield p_client | ||
| 771 | + else: | ||
| 772 | + yield d_client | ||
| 773 | + | ||
| 774 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 775 | + | ||
| 776 | + response = await router.handle_request() | ||
| 777 | + chunks = [chunk async for chunk in response.body_iterator] | ||
| 778 | + | ||
| 779 | + assert len(chunks) == 1 | ||
| 780 | + assert b'"text":"A"' in chunks[0] | ||
| 781 | + assert b"kv_transfer_params" not in chunks[0] | ||
| 782 | + | ||
| 783 | + | ||
| 758 | 784 | ||
| 759 | async def test_unified_pd_handoff_registers_decode_canceller_after_client_open( | 785 | async def test_unified_pd_handoff_registers_decode_canceller_after_client_open( |
| 760 | monkeypatch, | 786 | monkeypatch, |
| @@ -769,8 +795,10 @@ async def test_unified_pd_handoff_registers_decode_canceller_after_client_open( | |||
| 769 | entry_api="v1/completions", | 795 | entry_api="v1/completions", |
| 770 | req_len=10, | 796 | req_len=10, |
| 771 | ) | 797 | ) |
| 772 | - handoff = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | 798 | + scheduler = _Scheduler( |
| 773 | - scheduler = _Scheduler(prefill_capabilities=handoff, decode_capabilities=handoff) | 799 | + prefill_engine_type="vllm", |
| 800 | + decode_engine_type="vllm", | ||
| 801 | + ) | ||
| 774 | router = UnifiedPDRouter( | 802 | router = UnifiedPDRouter( |
| 775 | req_info, | 803 | req_info, |
| 776 | _config(), | 804 | _config(), |
| @@ -787,7 +815,7 @@ async def test_unified_pd_handoff_registers_decode_canceller_after_client_open( | |||
| 787 | 815 | ||
| 788 | monkeypatch.setattr(AttemptContext, "register_decode_canceller", _register_decode_canceller) | 816 | monkeypatch.setattr(AttemptContext, "register_decode_canceller", _register_decode_canceller) |
| 789 | 817 | ||
| 790 | - p_client = _PrefillResultClient("prefill") | 818 | + p_client = _NativeHandoffPrefillClient("prefill") |
| 791 | d_client = _Client("decode") | 819 | d_client = _Client("decode") |
| 792 | 820 | ||
| 793 | 821 | ||
| @@ -838,14 +866,7 @@ async def test_unified_pd_nonretryable_upstream_error_is_not_retried(monkeypatch | |||
| 838 | else: | 866 | else: |
| 839 | yield d_client | 867 | yield d_client |
| 840 | 868 | ||
| 841 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 842 | - return None | ||
| 843 | - | ||
| 844 | monkeypatch.setattr(router, "_client_for", _client_for) | 869 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 845 | - monkeypatch.setattr( | ||
| 846 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 847 | - _stop, | ||
| 848 | - ) | ||
| 849 | 870 | ||
| 850 | with pytest.raises(UpstreamHTTPError) as exc_info: | 871 | with pytest.raises(UpstreamHTTPError) as exc_info: |
| 851 | await router.handle_request() | 872 | await router.handle_request() |
| @@ -888,14 +909,7 @@ async def test_unified_pd_stream_prefill_rejection_is_returned_before_first_deco | |||
| 888 | else: | 909 | else: |
| 889 | yield d_client | 910 | yield d_client |
| 890 | 911 | ||
| 891 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 892 | - return None | ||
| 893 | - | ||
| 894 | monkeypatch.setattr(router, "_client_for", _client_for) | 912 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 895 | - monkeypatch.setattr( | ||
| 896 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 897 | - _stop, | ||
| 898 | - ) | ||
| 899 | monkeypatch.setattr( | 913 | monkeypatch.setattr( |
| 900 | router.rescheduler, | 914 | router.rescheduler, |
| 901 | "process_stream_chunk", | 915 | "process_stream_chunk", |
| @@ -913,8 +927,7 @@ async def test_unified_pd_stream_prefill_rejection_is_returned_before_first_deco | |||
| 913 | 927 | ||
| 914 | 928 | ||
| 915 | 929 | ||
| 916 | -async def test_unified_pd_cpcd_sglang_uses_concurrent_plan(monkeypatch): | 930 | +async def test_unified_pd_sglang_uses_native_bootstrap_concurrently(monkeypatch): |
| 917 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "9100") | ||
| 918 | req_info = RequestInfo( | 931 | req_info = RequestInfo( |
| 919 | req_id="root-cpcd-sglang", | 932 | req_id="root-cpcd-sglang", |
| 920 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | 933 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, |
| @@ -922,12 +935,10 @@ async def test_unified_pd_cpcd_sglang_uses_concurrent_plan(monkeypatch): | |||
| 922 | entry_api="v1/completions", | 935 | entry_api="v1/completions", |
| 923 | req_len=10, | 936 | req_len=10, |
| 924 | ) | 937 | ) |
| 925 | - concurrent = [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | ||
| 926 | scheduler = _Scheduler( | 938 | scheduler = _Scheduler( |
| 927 | prefill_engine_type="sglang", | 939 | prefill_engine_type="sglang", |
| 928 | decode_engine_type="sglang", | 940 | decode_engine_type="sglang", |
| 929 | - prefill_capabilities=concurrent, | 941 | + prefill_bootstrap_port=8998, |
| 930 | - decode_capabilities=concurrent, | ||
| 931 | ) | 942 | ) |
| 932 | router = UnifiedPDRouter( | 943 | router = UnifiedPDRouter( |
| 933 | req_info, | 944 | req_info, |
| @@ -935,7 +946,7 @@ async def test_unified_pd_cpcd_sglang_uses_concurrent_plan(monkeypatch): | |||
| 935 | scheduler=scheduler, | 946 | scheduler=scheduler, |
| 936 | request_manager=RequestManager(_config()), | 947 | request_manager=RequestManager(_config()), |
| 937 | ) | 948 | ) |
| 938 | - p_client = _Client("prefill") | 949 | + p_client = _NativeSglangPrefillClient("prefill") |
| 939 | d_client = _Client("decode") | 950 | d_client = _Client("decode") |
| 940 | 951 | ||
| 941 | 952 | ||
| @@ -951,24 +962,78 @@ async def test_unified_pd_cpcd_sglang_uses_concurrent_plan(monkeypatch): | |||
| 951 | 962 | ||
| 952 | assert len(p_client.requests) == 1 | 963 | assert len(p_client.requests) == 1 |
| 953 | assert len(d_client.requests) == 1 | 964 | assert len(d_client.requests) == 1 |
| 954 | - assert MOTOR_PREFILL_RESULT_KEY not in d_client.requests[0] | 965 | + p_request = p_client.requests[0] |
| 955 | - assert MOTOR_DISPATCH_KEY not in d_client.requests[0] | 966 | + d_request = d_client.requests[0] |
| 956 | - assert d_client.requests[0]["bootstrap_port"] == "9100" | 967 | + assert "_motor_dispatch" not in p_request |
| 968 | + assert "_motor_dispatch" not in d_request | ||
| 969 | + assert "_motor_prefill_result" not in d_request | ||
| 970 | + assert p_request["rid"] == d_request["rid"] == "root-cpcd-sglang#a1" | ||
| 971 | + for request in (p_request, d_request): | ||
| 972 | + assert request["bootstrap_host"] == "127.0.0.1" | ||
| 973 | + assert request["bootstrap_port"] == 8998 | ||
| 974 | + assert p_request["bootstrap_room"] == d_request["bootstrap_room"] | ||
| 975 | + assert p_client.headers[0]["X-Request-Id"] == "root-cpcd-sglang#a1" | ||
| 976 | + assert d_client.headers[0]["X-Request-Id"] == "root-cpcd-sglang#a1" | ||
| 957 | assert scheduler.update_workload.await_count == 2 | 977 | assert scheduler.update_workload.await_count == 2 |
| 958 | 978 | ||
| 959 | 979 | ||
| 960 | 980 | ||
| 961 | -async def test_unified_pd_rejects_pair_without_shared_connector_capability(monkeypatch): | 981 | +async def test_unified_pd_sglang_rejects_missing_bootstrap_port_before_dispatch( |
| 982 | + monkeypatch, | ||
| 983 | +): | ||
| 962 | req_info = RequestInfo( | 984 | req_info = RequestInfo( |
| 963 | - req_id="root-mixed", | 985 | + req_id="root-sglang-no-port", |
| 964 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | 986 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, |
| 965 | api="v1/completions", | 987 | api="v1/completions", |
| 966 | entry_api="v1/completions", | 988 | entry_api="v1/completions", |
| 967 | req_len=10, | 989 | req_len=10, |
| 968 | ) | 990 | ) |
| 969 | scheduler = _Scheduler( | 991 | scheduler = _Scheduler( |
| 970 | - prefill_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | 992 | + prefill_engine_type="sglang", |
| 971 | - decode_capabilities=[DispatchPlan.PREFILL_HANDOFF_DECODE.value], | 993 | + decode_engine_type="sglang", |
| 994 | + ) | ||
| 995 | + next(iter(next(iter(scheduler.p.endpoints.values())).values())).bootstrap_port = None | ||
| 996 | + router = UnifiedPDRouter( | ||
| 997 | + req_info, | ||
| 998 | + _config(), | ||
| 999 | + scheduler=scheduler, | ||
| 1000 | + request_manager=RequestManager(_config()), | ||
| 1001 | + ) | ||
| 1002 | + p_client = _NativeSglangPrefillClient("prefill") | ||
| 1003 | + d_client = _Client("decode") | ||
| 1004 | + | ||
| 1005 | + | ||
| 1006 | + async def _client_for(resource: ScheduledResource): | ||
| 1007 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 1008 | + yield p_client | ||
| 1009 | + else: | ||
| 1010 | + yield d_client | ||
| 1011 | + | ||
| 1012 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 1013 | + | ||
| 1014 | + with pytest.raises(EngineProtocolError, match="bootstrap port"): | ||
| 1015 | + await router.handle_request() | ||
| 1016 | + | ||
| 1017 | + assert p_client.requests == [] | ||
| 1018 | + assert d_client.requests == [] | ||
| 1019 | + assert scheduler.update_workload.await_count == 2 | ||
| 1020 | + | ||
| 1021 | + | ||
| 1022 | + | ||
| 1023 | +async def test_unified_pd_sglang_rejects_mixed_engine_pair_without_dispatch( | ||
| 1024 | + monkeypatch, | ||
| 1025 | +): | ||
| 1026 | + req_info = RequestInfo( | ||
| 1027 | + req_id="root-sglang-mixed", | ||
| 1028 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 1029 | + api="v1/completions", | ||
| 1030 | + entry_api="v1/completions", | ||
| 1031 | + req_len=10, | ||
| 1032 | + ) | ||
| 1033 | + scheduler = _Scheduler( | ||
| 1034 | + prefill_engine_type="sglang", | ||
| 1035 | + decode_engine_type="vllm", | ||
| 1036 | + prefill_bootstrap_port=8998, | ||
| 972 | ) | 1037 | ) |
| 973 | router = UnifiedPDRouter( | 1038 | router = UnifiedPDRouter( |
| 974 | req_info, | 1039 | req_info, |
| @@ -976,21 +1041,211 @@ async def test_unified_pd_rejects_pair_without_shared_connector_capability(monke | |||
| 976 | scheduler=scheduler, | 1041 | scheduler=scheduler, |
| 977 | request_manager=RequestManager(_config()), | 1042 | request_manager=RequestManager(_config()), |
| 978 | ) | 1043 | ) |
| 979 | - stop_calls = [] | 1044 | + with pytest.raises(HTTPException) as exc_info: |
| 980 | - | ||
| 981 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 982 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason.value)) | ||
| 983 | - return None | ||
| 984 | - | ||
| 985 | - monkeypatch.setattr( | ||
| 986 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 987 | - _stop, | ||
| 988 | - ) | ||
| 989 | - | ||
| 990 | - with pytest.raises(RuntimeError, match="no shared dispatch capability"): | ||
| 991 | await router.handle_request() | 1045 | await router.handle_request() |
| 992 | 1046 | ||
| 993 | - assert {call[0] for call in stop_calls} == {PDRole.ROLE_P, PDRole.ROLE_D} | 1047 | + assert exc_info.value.status_code == 503 |
| 1048 | + assert "engine_type=sglang" in exc_info.value.detail | ||
| 1049 | + assert scheduler.update_workload.await_count == 1 | ||
| 1050 | + | ||
| 1051 | + | ||
| 1052 | + | ||
| 1053 | +async def test_unified_pd_sglang_decode_failure_cancels_prefill_without_peer_stop( | ||
| 1054 | + monkeypatch, | ||
| 1055 | +): | ||
| 1056 | + req_info = RequestInfo( | ||
| 1057 | + req_id="root-sglang-cancel", | ||
| 1058 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 1059 | + api="v1/completions", | ||
| 1060 | + entry_api="v1/completions", | ||
| 1061 | + req_len=10, | ||
| 1062 | + ) | ||
| 1063 | + scheduler = _Scheduler( | ||
| 1064 | + prefill_engine_type="sglang", | ||
| 1065 | + decode_engine_type="sglang", | ||
| 1066 | + prefill_bootstrap_port=8998, | ||
| 1067 | + ) | ||
| 1068 | + router = UnifiedPDRouter( | ||
| 1069 | + req_info, | ||
| 1070 | + _config(), | ||
| 1071 | + scheduler=scheduler, | ||
| 1072 | + request_manager=RequestManager(_config()), | ||
| 1073 | + ) | ||
| 1074 | + prefill_started = asyncio.Event() | ||
| 1075 | + prefill_cancelled = asyncio.Event() | ||
| 1076 | + | ||
| 1077 | + class _BlockingPrefillClient(_Client): | ||
| 1078 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1079 | + if str(path).rstrip("/").endswith("abort_request"): | ||
| 1080 | + return await super().post(path, json=json, headers=headers, timeout=timeout) | ||
| 1081 | + self.requests.append(json) | ||
| 1082 | + self.headers.append(headers or {}) | ||
| 1083 | + prefill_started.set() | ||
| 1084 | + try: | ||
| 1085 | + await asyncio.Event().wait() | ||
| 1086 | + finally: | ||
| 1087 | + prefill_cancelled.set() | ||
| 1088 | + | ||
| 1089 | + class _FailingDecodeClient(_Client): | ||
| 1090 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1091 | + if str(path).rstrip("/").endswith("abort_request"): | ||
| 1092 | + return await super().post(path, json=json, headers=headers, timeout=timeout) | ||
| 1093 | + self.requests.append(json) | ||
| 1094 | + self.headers.append(headers or {}) | ||
| 1095 | + await prefill_started.wait() | ||
| 1096 | + raise httpx.ConnectError("decode down") | ||
| 1097 | + | ||
| 1098 | + p_client = _BlockingPrefillClient("prefill") | ||
| 1099 | + d_client = _FailingDecodeClient("decode") | ||
| 1100 | + | ||
| 1101 | + | ||
| 1102 | + async def _client_for(resource: ScheduledResource): | ||
| 1103 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 1104 | + yield p_client | ||
| 1105 | + else: | ||
| 1106 | + yield d_client | ||
| 1107 | + | ||
| 1108 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 1109 | + | ||
| 1110 | + with pytest.raises(httpx.ConnectError, match="decode down"): | ||
| 1111 | + await router.handle_request() | ||
| 1112 | + | ||
| 1113 | + await asyncio.wait_for(prefill_cancelled.wait(), timeout=1) | ||
| 1114 | + assert p_client.requests[0]["bootstrap_room"] == d_client.requests[0]["bootstrap_room"] | ||
| 1115 | + assert p_client.abort_requests == [{"rid": "root-sglang-cancel#a1"}] | ||
| 1116 | + assert d_client.abort_requests == [{"rid": "root-sglang-cancel#a1"}] | ||
| 1117 | + assert scheduler.update_workload.await_count == 2 | ||
| 1118 | + | ||
| 1119 | + | ||
| 1120 | + | ||
| 1121 | +async def test_unified_pd_sglang_retry_uses_new_matched_bootstrap_room( | ||
| 1122 | + monkeypatch, | ||
| 1123 | +): | ||
| 1124 | + req_info = RequestInfo( | ||
| 1125 | + req_id="root-sglang-retry", | ||
| 1126 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 1127 | + api="v1/completions", | ||
| 1128 | + entry_api="v1/completions", | ||
| 1129 | + req_len=10, | ||
| 1130 | + ) | ||
| 1131 | + scheduler = _Scheduler( | ||
| 1132 | + prefill_engine_type="sglang", | ||
| 1133 | + decode_engine_type="sglang", | ||
| 1134 | + prefill_bootstrap_port=8998, | ||
| 1135 | + ) | ||
| 1136 | + config = _config() | ||
| 1137 | + config.exception_config.transport_max_retry = 2 | ||
| 1138 | + router = UnifiedPDRouter( | ||
| 1139 | + req_info, | ||
| 1140 | + config, | ||
| 1141 | + scheduler=scheduler, | ||
| 1142 | + request_manager=RequestManager(config), | ||
| 1143 | + ) | ||
| 1144 | + p_client = _NativeSglangPrefillClient("prefill") | ||
| 1145 | + | ||
| 1146 | + class _FailOnceDecodeClient(_Client): | ||
| 1147 | + async def post(self, path, json=None, headers=None, timeout=None): | ||
| 1148 | + if str(path).rstrip("/").endswith("abort_request"): | ||
| 1149 | + return await super().post(path, json=json, headers=headers, timeout=timeout) | ||
| 1150 | + self.requests.append(json) | ||
| 1151 | + self.headers.append(headers or {}) | ||
| 1152 | + if len(self.requests) == 1: | ||
| 1153 | + raise httpx.ConnectError("retry decode") | ||
| 1154 | + request = httpx.Request("POST", path, headers=headers or {}, json=json) | ||
| 1155 | + return httpx.Response( | ||
| 1156 | + status_code=200, | ||
| 1157 | + json={"choices": [{"text": "ok", "index": 0}]}, | ||
| 1158 | + request=request, | ||
| 1159 | + ) | ||
| 1160 | + | ||
| 1161 | + d_client = _FailOnceDecodeClient("decode") | ||
| 1162 | + | ||
| 1163 | + | ||
| 1164 | + async def _client_for(resource: ScheduledResource): | ||
| 1165 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 1166 | + yield p_client | ||
| 1167 | + else: | ||
| 1168 | + yield d_client | ||
| 1169 | + | ||
| 1170 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 1171 | + | ||
| 1172 | + response = await router.handle_request() | ||
| 1173 | + | ||
| 1174 | + assert json.loads(response.body)["choices"][0]["text"] == "ok" | ||
| 1175 | + assert len(p_client.requests) == len(d_client.requests) == 2 | ||
| 1176 | + for attempt_seq, (p_request, d_request) in enumerate( | ||
| 1177 | + zip(p_client.requests, d_client.requests, strict=True), | ||
| 1178 | + start=1, | ||
| 1179 | + ): | ||
| 1180 | + assert p_request["rid"] == d_request["rid"] == f"root-sglang-retry#a{attempt_seq}" | ||
| 1181 | + assert p_request["bootstrap_room"] == d_request["bootstrap_room"] | ||
| 1182 | + assert p_client.requests[0]["bootstrap_room"] != p_client.requests[1]["bootstrap_room"] | ||
| 1183 | + | ||
| 1184 | + | ||
| 1185 | + | ||
| 1186 | +async def test_unified_pd_sglang_chat_stream_uses_native_fields_and_strips_them( | ||
| 1187 | + monkeypatch, | ||
| 1188 | +): | ||
| 1189 | + req_info = RequestInfo( | ||
| 1190 | + req_id="root-sglang-chat-stream", | ||
| 1191 | + req_data={ | ||
| 1192 | + "model": "m", | ||
| 1193 | + "messages": [{"role": "user", "content": "hello"}], | ||
| 1194 | + "stream": True, | ||
| 1195 | + "max_completion_tokens": 6, | ||
| 1196 | + }, | ||
| 1197 | + api="v1/chat/completions", | ||
| 1198 | + entry_api="v1/chat/completions", | ||
| 1199 | + req_len=10, | ||
| 1200 | + ) | ||
| 1201 | + scheduler = _Scheduler( | ||
| 1202 | + prefill_engine_type="sglang", | ||
| 1203 | + decode_engine_type="sglang", | ||
| 1204 | + prefill_bootstrap_port=8998, | ||
| 1205 | + ) | ||
| 1206 | + router = UnifiedPDRouter( | ||
| 1207 | + req_info, | ||
| 1208 | + _config(), | ||
| 1209 | + scheduler=scheduler, | ||
| 1210 | + request_manager=RequestManager(_config()), | ||
| 1211 | + ) | ||
| 1212 | + p_client = _NativeSglangPrefillClient("prefill") | ||
| 1213 | + d_client = _SequenceStreamClient( | ||
| 1214 | + "decode", | ||
| 1215 | + [ | ||
| 1216 | + _StreamResponse( | ||
| 1217 | + [ | ||
| 1218 | + ( | ||
| 1219 | + b'data: {"choices":[{"delta":{"content":"A"},"index":0}],' | ||
| 1220 | + b'"bootstrap_host":"127.0.0.1","bootstrap_port":8998,' | ||
| 1221 | + b'"bootstrap_room":123}\n\n' | ||
| 1222 | + ) | ||
| 1223 | + ] | ||
| 1224 | + ) | ||
| 1225 | + ], | ||
| 1226 | + ) | ||
| 1227 | + | ||
| 1228 | + | ||
| 1229 | + async def _client_for(resource: ScheduledResource): | ||
| 1230 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 1231 | + yield p_client | ||
| 1232 | + else: | ||
| 1233 | + yield d_client | ||
| 1234 | + | ||
| 1235 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 1236 | + | ||
| 1237 | + response = await router.handle_request() | ||
| 1238 | + chunks = [chunk async for chunk in response.body_iterator] | ||
| 1239 | + | ||
| 1240 | + assert len(chunks) == 1 | ||
| 1241 | + assert b'"content":"A"' in chunks[0] | ||
| 1242 | + assert b"bootstrap_" not in chunks[0] | ||
| 1243 | + assert p_client.requests[0]["stream"] is False | ||
| 1244 | + assert d_client.requests[0]["stream"] is True | ||
| 1245 | + for request in (p_client.requests[0], d_client.requests[0]): | ||
| 1246 | + assert request["max_completion_tokens"] == 6 | ||
| 1247 | + assert "_motor_dispatch" not in request | ||
| 1248 | + assert p_client.requests[0]["bootstrap_room"] == d_client.requests[0]["bootstrap_room"] | ||
| 994 | 1249 | ||
| 995 | 1250 | ||
| 996 | 1251 | ||
| @@ -1270,7 +1525,7 @@ async def test_unified_pd_background_release_uses_single_tracked_task(): | |||
| 1270 | 1525 | ||
| 1271 | 1526 | ||
| 1272 | 1527 | ||
| 1273 | -async def test_unified_pd_release_failure_retains_local_ledger_and_is_drained(caplog): | 1528 | +async def test_unified_pd_release_failure_keeps_local_release_and_is_drained(caplog): |
| 1274 | req_info = RequestInfo( | 1529 | req_info = RequestInfo( |
| 1275 | req_id="root-release-failure", | 1530 | req_id="root-release-failure", |
| 1276 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | 1531 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, |
| @@ -1298,6 +1553,7 @@ async def test_unified_pd_release_failure_retains_local_ledger_and_is_drained(ca | |||
| 1298 | PDRole.ROLE_P, | 1553 | PDRole.ROLE_P, |
| 1299 | ) | 1554 | ) |
| 1300 | assert original is not None | 1555 | assert original is not None |
| 1556 | + original_active_tokens = original.active_tokens | ||
| 1301 | 1557 | ||
| 1302 | with caplog.at_level(logging.DEBUG, logger=_ROUTER_LOGGER): | 1558 | with caplog.at_level(logging.DEBUG, logger=_ROUTER_LOGGER): |
| 1303 | submitted = await router._release_attempt_resource( | 1559 | submitted = await router._release_attempt_resource( |
| @@ -1315,11 +1571,11 @@ async def test_unified_pd_release_failure_retains_local_ledger_and_is_drained(ca | |||
| 1315 | attempt.attempt_seq, | 1571 | attempt.attempt_seq, |
| 1316 | PDRole.ROLE_P, | 1572 | PDRole.ROLE_P, |
| 1317 | ) | 1573 | ) |
| 1318 | - # Delete-after-ACK: failed RPC keeps the local ledger so a later re-enqueue can | 1574 | + assert current is not None |
| 1319 | - # recompute the full negative delta; release_flags stay unmarked either way. | 1575 | + assert current.active_tokens == original_active_tokens |
| 1320 | - assert current == original | ||
| 1321 | assert not attempt.release_flags.prefill_tokens | 1576 | assert not attempt.release_flags.prefill_tokens |
| 1322 | - assert scheduler.update_workload.await_count == UnifiedPDRouter._RELEASE_RPC_ATTEMPTS | 1577 | + assert scheduler.update_workload.await_count == 3 |
| 1578 | + assert "Release workload background task failed" in caplog.text | ||
| 1323 | assert "Release workload rolled back locally" not in caplog.text | 1579 | assert "Release workload rolled back locally" not in caplog.text |
| 1324 | finally: | 1580 | finally: |
| 1325 | await request_manager.del_req_info(req_info.req_id) | 1581 | await request_manager.del_req_info(req_info.req_id) |
| @@ -1486,7 +1742,7 @@ async def test_unified_pd_drain_double_cancel_keeps_release_cleanup(): | |||
| 1486 | 1742 | ||
| 1487 | 1743 | ||
| 1488 | 1744 | ||
| 1489 | -async def test_unified_pd_concurrent_stream_tail_release_survives_iterator_cancellation( | 1745 | +async def test_unified_pd_bootstrap_stream_tail_release_survives_iterator_cancellation( |
| 1490 | monkeypatch, | 1746 | monkeypatch, |
| 1491 | ): | 1747 | ): |
| 1492 | req_info = RequestInfo( | 1748 | req_info = RequestInfo( |
| @@ -1565,8 +1821,10 @@ async def test_unified_pd_handoff_stream_yields_before_prefill_token_release_fin | |||
| 1565 | entry_api="v1/completions", | 1821 | entry_api="v1/completions", |
| 1566 | req_len=10, | 1822 | req_len=10, |
| 1567 | ) | 1823 | ) |
| 1568 | - handoff = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | 1824 | + scheduler = _Scheduler( |
| 1569 | - scheduler = _Scheduler(prefill_capabilities=handoff, decode_capabilities=handoff) | 1825 | + prefill_engine_type="vllm", |
| 1826 | + decode_engine_type="vllm", | ||
| 1827 | + ) | ||
| 1570 | token_release_compute_started = asyncio.Event() | 1828 | token_release_compute_started = asyncio.Event() |
| 1571 | allow_token_release_compute = asyncio.Event() | 1829 | allow_token_release_compute = asyncio.Event() |
| 1572 | token_release_started = asyncio.Event() | 1830 | token_release_started = asyncio.Event() |
| @@ -1605,7 +1863,7 @@ async def test_unified_pd_handoff_stream_yields_before_prefill_token_release_fin | |||
| 1605 | return await compute_and_update(resource, req_id, action, req_info_arg, **kwargs) | 1863 | return await compute_and_update(resource, req_id, action, req_info_arg, **kwargs) |
| 1606 | 1864 | ||
| 1607 | monkeypatch.setattr(router._workload_action_handler, "compute_and_update", _compute_and_update) | 1865 | monkeypatch.setattr(router._workload_action_handler, "compute_and_update", _compute_and_update) |
| 1608 | - p_client = _PrefillResultClient("prefill") | 1866 | + p_client = _NativeHandoffPrefillClient("prefill") |
| 1609 | d_client = _StreamClient("decode") | 1867 | d_client = _StreamClient("decode") |
| 1610 | 1868 | ||
| 1611 | 1869 | ||
| @@ -1671,18 +1929,18 @@ async def test_unified_pd_stream_dispatches_context_and_yields_visible_chunk( | |||
| 1671 | assert chunks == [b'data: {"choices":[{"delta":{"content":"A"},"index":0}]}\n\n'] | 1929 | assert chunks == [b'data: {"choices":[{"delta":{"content":"A"},"index":0}]}\n\n'] |
| 1672 | assert len(p_client.requests) == 1 | 1930 | assert len(p_client.requests) == 1 |
| 1673 | assert len(d_client.requests) == 1 | 1931 | assert len(d_client.requests) == 1 |
| 1674 | - assert d_client.requests[0][MOTOR_DISPATCH_KEY]["role"] == "decode" | 1932 | + assert d_client.requests[0]["rid"] == p_client.requests[0]["rid"] |
| 1675 | - assert p_client.requests[0][MOTOR_DISPATCH_KEY]["pair_id"] == d_client.requests[0][MOTOR_DISPATCH_KEY]["pair_id"] | 1933 | + assert p_client.requests[0]["bootstrap_room"] == d_client.requests[0]["bootstrap_room"] |
| 1676 | assert scheduler.update_workload.await_count == 2 | 1934 | assert scheduler.update_workload.await_count == 2 |
| 1677 | assert ReqState.PREFILL_END in req_info.status | 1935 | assert ReqState.PREFILL_END in req_info.status |
| 1678 | 1936 | ||
| 1679 | 1937 | ||
| 1680 | 1938 | ||
| 1681 | -async def test_unified_pd_concurrent_stream_prefill_release_does_not_block_first_chunk( | 1939 | +async def test_unified_pd_bootstrap_stream_prefill_release_does_not_block_first_chunk( |
| 1682 | monkeypatch, | 1940 | monkeypatch, |
| 1683 | ): | 1941 | ): |
| 1684 | req_info = RequestInfo( | 1942 | req_info = RequestInfo( |
| 1685 | - req_id="root-concurrent-prefill-release-background", | 1943 | + req_id="root-bootstrap-prefill-release-background", |
| 1686 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | 1944 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, |
| 1687 | api="v1/completions", | 1945 | api="v1/completions", |
| 1688 | entry_api="v1/completions", | 1946 | entry_api="v1/completions", |
| @@ -1739,12 +1997,11 @@ async def test_unified_pd_concurrent_stream_prefill_release_does_not_block_first | |||
| 1739 | 1997 | ||
| 1740 | 1998 | ||
| 1741 | 1999 | ||
| 1742 | -async def test_unified_pd_concurrent_nonstream_holds_prefill_until_decode_returns( | 2000 | +async def test_unified_pd_bootstrap_nonstream_holds_prefill_tokens_until_decode_finishes( |
| 1743 | monkeypatch, | 2001 | monkeypatch, |
| 1744 | ): | 2002 | ): |
| 1745 | - """Trigger/concurrent: P prepared ACK must not release load before decode completes.""" | ||
| 1746 | req_info = RequestInfo( | 2003 | req_info = RequestInfo( |
| 1747 | - req_id="root-concurrent-nonstream-prefill-release", | 2004 | + req_id="root-bootstrap-nonstream-prefill-release", |
| 1748 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | 2005 | req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, |
| 1749 | api="v1/chat/completions", | 2006 | api="v1/chat/completions", |
| 1750 | entry_api="v1/chat/completions", | 2007 | entry_api="v1/chat/completions", |
| @@ -1797,13 +2054,13 @@ async def test_unified_pd_concurrent_nonstream_holds_prefill_until_decode_return | |||
| 1797 | 2054 | ||
| 1798 | response_task = asyncio.create_task(router.handle_request()) | 2055 | response_task = asyncio.create_task(router.handle_request()) |
| 1799 | await asyncio.wait_for(decode_started.wait(), timeout=1) | 2056 | await asyncio.wait_for(decode_started.wait(), timeout=1) |
| 1800 | - # Prefill HTTP has returned (prepared), but P load must still be held. | ||
| 1801 | - await asyncio.sleep(0.05) | ||
| 1802 | assert not p_token_compute_started.is_set() | 2057 | assert not p_token_compute_started.is_set() |
| 1803 | assert not response_task.done() | 2058 | assert not response_task.done() |
| 1804 | 2059 | ||
| 1805 | allow_decode.set() | 2060 | allow_decode.set() |
| 1806 | await asyncio.wait_for(p_token_compute_started.wait(), timeout=1) | 2061 | await asyncio.wait_for(p_token_compute_started.wait(), timeout=1) |
| 2062 | + assert not response_task.done() | ||
| 2063 | + | ||
| 1807 | allow_p_release_compute.set() | 2064 | allow_p_release_compute.set() |
| 1808 | response = await asyncio.wait_for(response_task, timeout=1) | 2065 | response = await asyncio.wait_for(response_task, timeout=1) |
| 1809 | 2066 | ||
| @@ -1811,7 +2068,7 @@ async def test_unified_pd_concurrent_nonstream_holds_prefill_until_decode_return | |||
| 1811 | 2068 | ||
| 1812 | 2069 | ||
| 1813 | 2070 | ||
| 1814 | -async def test_unified_pd_client_disconnect_cancels_tasks_and_stops_engine(monkeypatch): | 2071 | +async def test_unified_pd_client_disconnect_cancels_tasks_and_releases_resources(monkeypatch): |
| 1815 | req_info = RequestInfo( | 2072 | req_info = RequestInfo( |
| 1816 | req_id="root-stream-disconnect", | 2073 | req_id="root-stream-disconnect", |
| 1817 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | 2074 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, |
| @@ -1832,7 +2089,6 @@ async def test_unified_pd_client_disconnect_cancels_tasks_and_stops_engine(monke | |||
| 1832 | d_response = _BlockingStreamResponse() | 2089 | d_response = _BlockingStreamResponse() |
| 1833 | d_client = _BlockingStreamClient("decode", d_response) | 2090 | d_client = _BlockingStreamClient("decode", d_response) |
| 1834 | attempts = [] | 2091 | attempts = [] |
| 1835 | - stop_calls = [] | ||
| 1836 | original_create_attempt = router._create_attempt | 2092 | original_create_attempt = router._create_attempt |
| 1837 | 2093 | ||
| 1838 | 2094 | ||
| @@ -1847,10 +2103,6 @@ async def test_unified_pd_client_disconnect_cancels_tasks_and_stops_engine(monke | |||
| 1847 | attempts.append(attempt) | 2103 | attempts.append(attempt) |
| 1848 | return attempt | 2104 | return attempt |
| 1849 | 2105 | ||
| 1850 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 1851 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason)) | ||
| 1852 | - return None | ||
| 1853 | - | ||
| 1854 | async def receive(): | 2106 | async def receive(): |
| 1855 | await d_response.started.wait() | 2107 | await d_response.started.wait() |
| 1856 | return {"type": "http.disconnect"} | 2108 | return {"type": "http.disconnect"} |
| @@ -1860,10 +2112,6 @@ async def test_unified_pd_client_disconnect_cancels_tasks_and_stops_engine(monke | |||
| 1860 | 2112 | ||
| 1861 | monkeypatch.setattr(router, "_client_for", _client_for) | 2113 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 1862 | monkeypatch.setattr(router, "_create_attempt", _create_attempt) | 2114 | monkeypatch.setattr(router, "_create_attempt", _create_attempt) |
| 1863 | - monkeypatch.setattr( | ||
| 1864 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 1865 | - _stop, | ||
| 1866 | - ) | ||
| 1867 | 2115 | ||
| 1868 | response = await router.handle_request() | 2116 | response = await router.handle_request() |
| 1869 | await asyncio.wait_for(response(_asgi_scope(), receive, send), timeout=1) | 2117 | await asyncio.wait_for(response(_asgi_scope(), receive, send), timeout=1) |
| @@ -1872,13 +2120,8 @@ async def test_unified_pd_client_disconnect_cancels_tasks_and_stops_engine(monke | |||
| 1872 | attempt = attempts[0] | 2120 | attempt = attempts[0] |
| 1873 | await asyncio.wait_for(d_response.closed.wait(), timeout=1) | 2121 | await asyncio.wait_for(d_response.closed.wait(), timeout=1) |
| 1874 | assert attempt.state == AttemptState.STOPPED | 2122 | assert attempt.state == AttemptState.STOPPED |
| 1875 | - assert attempt.stop_sent is True | ||
| 1876 | assert attempt.prefill_task.done() | 2123 | assert attempt.prefill_task.done() |
| 1877 | assert attempt.decode_task.done() | 2124 | assert attempt.decode_task.done() |
| 1878 | - assert set(stop_calls) == { | ||
| 1879 | - (PDRole.ROLE_P, 1, DispatchStopReason.CLIENT_DISCONNECT), | ||
| 1880 | - (PDRole.ROLE_D, 1, DispatchStopReason.CLIENT_DISCONNECT), | ||
| 1881 | - } | ||
| 1882 | assert scheduler.update_workload.await_count == 2 | 2125 | assert scheduler.update_workload.await_count == 2 |
| 1883 | assert not any( | 2126 | assert not any( |
| 1884 | task.get_name() == "unified-pd-queue-root-stream-disconnect-a1" and not task.done() | 2127 | task.get_name() == "unified-pd-queue-root-stream-disconnect-a1" and not task.done() |
| @@ -1906,20 +2149,12 @@ async def test_unified_pd_stop_attempt_drains_release_failures(monkeypatch, capl | |||
| 1906 | request_manager=request_manager, | 2149 | request_manager=request_manager, |
| 1907 | ) | 2150 | ) |
| 1908 | 2151 | ||
| 1909 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 1910 | - return None | ||
| 1911 | - | ||
| 1912 | - monkeypatch.setattr( | ||
| 1913 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 1914 | - _stop, | ||
| 1915 | - ) | ||
| 1916 | - | ||
| 1917 | await request_manager.add_req_info(req_info) | 2152 | await request_manager.add_req_info(req_info) |
| 1918 | try: | 2153 | try: |
| 1919 | attempt = await router._create_attempt(PDDispatchSession(req_info.req_id)) | 2154 | attempt = await router._create_attempt(PDDispatchSession(req_info.req_id)) |
| 1920 | 2155 | ||
| 1921 | with caplog.at_level(logging.DEBUG, logger=_ROUTER_LOGGER): | 2156 | with caplog.at_level(logging.DEBUG, logger=_ROUTER_LOGGER): |
| 1922 | - await router._stop_attempt(attempt, DispatchStopReason.CLIENT_DISCONNECT) | 2157 | + await router._stop_attempt(attempt, AttemptStopReason.CLIENT_DISCONNECT) |
| 1923 | 2158 | ||
| 1924 | assert attempt.state == AttemptState.STOPPED | 2159 | assert attempt.state == AttemptState.STOPPED |
| 1925 | assert scheduler.update_workload.await_count == 6 | 2160 | assert scheduler.update_workload.await_count == 6 |
| @@ -1932,7 +2167,7 @@ async def test_unified_pd_stop_attempt_drains_release_failures(monkeypatch, capl | |||
| 1932 | 2167 | ||
| 1933 | 2168 | ||
| 1934 | 2169 | ||
| 1935 | -async def test_unified_pd_send_failure_closes_decode_and_stops_engine(monkeypatch): | 2170 | +async def test_unified_pd_send_failure_closes_decode_and_releases_resources(monkeypatch): |
| 1936 | req_info = RequestInfo( | 2171 | req_info = RequestInfo( |
| 1937 | req_id="root-stream-send-failure", | 2172 | req_id="root-stream-send-failure", |
| 1938 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | 2173 | req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, |
| @@ -1954,7 +2189,6 @@ async def test_unified_pd_send_failure_closes_decode_and_stops_engine(monkeypatc | |||
| 1954 | ) | 2189 | ) |
| 1955 | d_client = _BlockingStreamClient("decode", d_response) | 2190 | d_client = _BlockingStreamClient("decode", d_response) |
| 1956 | attempts = [] | 2191 | attempts = [] |
| 1957 | - stop_calls = [] | ||
| 1958 | original_create_attempt = router._create_attempt | 2192 | original_create_attempt = router._create_attempt |
| 1959 | 2193 | ||
| 1960 | 2194 | ||
| @@ -1969,20 +2203,12 @@ async def test_unified_pd_send_failure_closes_decode_and_stops_engine(monkeypatc | |||
| 1969 | attempts.append(attempt) | 2203 | attempts.append(attempt) |
| 1970 | return attempt | 2204 | return attempt |
| 1971 | 2205 | ||
| 1972 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 1973 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason)) | ||
| 1974 | - return None | ||
| 1975 | - | ||
| 1976 | async def send(message): | 2206 | async def send(message): |
| 1977 | if message["type"] == "http.response.body" and message.get("body"): | 2207 | if message["type"] == "http.response.body" and message.get("body"): |
| 1978 | raise OSError("client socket closed") | 2208 | raise OSError("client socket closed") |
| 1979 | 2209 | ||
| 1980 | monkeypatch.setattr(router, "_client_for", _client_for) | 2210 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 1981 | monkeypatch.setattr(router, "_create_attempt", _create_attempt) | 2211 | monkeypatch.setattr(router, "_create_attempt", _create_attempt) |
| 1982 | - monkeypatch.setattr( | ||
| 1983 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 1984 | - _stop, | ||
| 1985 | - ) | ||
| 1986 | 2212 | ||
| 1987 | response = await router.handle_request() | 2213 | response = await router.handle_request() |
| 1988 | with pytest.raises(ClientDisconnect): | 2214 | with pytest.raises(ClientDisconnect): |
| @@ -1992,13 +2218,8 @@ async def test_unified_pd_send_failure_closes_decode_and_stops_engine(monkeypatc | |||
| 1992 | attempt = attempts[0] | 2218 | attempt = attempts[0] |
| 1993 | await asyncio.wait_for(d_response.closed.wait(), timeout=1) | 2219 | await asyncio.wait_for(d_response.closed.wait(), timeout=1) |
| 1994 | assert attempt.state == AttemptState.STOPPED | 2220 | assert attempt.state == AttemptState.STOPPED |
| 1995 | - assert attempt.stop_sent is True | ||
| 1996 | assert attempt.prefill_task.done() | 2221 | assert attempt.prefill_task.done() |
| 1997 | assert attempt.decode_task.done() | 2222 | assert attempt.decode_task.done() |
| 1998 | - assert set(stop_calls) == { | ||
| 1999 | - (PDRole.ROLE_P, 1, DispatchStopReason.CLIENT_DISCONNECT), | ||
| 2000 | - (PDRole.ROLE_D, 1, DispatchStopReason.CLIENT_DISCONNECT), | ||
| 2001 | - } | ||
| 2002 | assert scheduler.update_workload.await_count == 2 | 2223 | assert scheduler.update_workload.await_count == 2 |
| 2003 | 2224 | ||
| 2004 | 2225 | ||
| @@ -2056,14 +2277,7 @@ async def test_unified_pd_stream_error_after_visible_chunk_reschedules_with_toke | |||
| 2056 | else: | 2277 | else: |
| 2057 | yield d_client | 2278 | yield d_client |
| 2058 | 2279 | ||
| 2059 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2060 | - return None | ||
| 2061 | - | ||
| 2062 | monkeypatch.setattr(router, "_client_for", _client_for) | 2280 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 2063 | - monkeypatch.setattr( | ||
| 2064 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2065 | - _stop, | ||
| 2066 | - ) | ||
| 2067 | 2281 | ||
| 2068 | response = await router.handle_request() | 2282 | response = await router.handle_request() |
| 2069 | messages = await _invoke_asgi_response(response) | 2283 | messages = await _invoke_asgi_response(response) |
| @@ -2075,7 +2289,7 @@ async def test_unified_pd_stream_error_after_visible_chunk_reschedules_with_toke | |||
| 2075 | assert len(d_client.requests) == 2 | 2289 | assert len(d_client.requests) == 2 |
| 2076 | assert d_client.requests[0]["return_token_ids"] is True | 2290 | assert d_client.requests[0]["return_token_ids"] is True |
| 2077 | assert p_client.requests[1]["prompt"] == [1, 2, 10] | 2291 | assert p_client.requests[1]["prompt"] == [1, 2, 10] |
| 2078 | - assert p_client.requests[1]["max_tokens"] == 1 | 2292 | + assert p_client.requests[1]["max_tokens"] == 8 |
| 2079 | assert p_client.requests[1]["stream"] is False | 2293 | assert p_client.requests[1]["stream"] is False |
| 2080 | assert d_client.requests[1]["prompt"] == [1, 2, 10] | 2294 | assert d_client.requests[1]["prompt"] == [1, 2, 10] |
| 2081 | assert d_client.requests[1]["max_tokens"] == 7 | 2295 | assert d_client.requests[1]["max_tokens"] == 7 |
| @@ -2114,7 +2328,6 @@ async def test_unified_pd_retry_plan_validation_fails_before_new_attempt_allocat | |||
| 2114 | scheduler=scheduler, | 2328 | scheduler=scheduler, |
| 2115 | request_manager=RequestManager(config), | 2329 | request_manager=RequestManager(config), |
| 2116 | ) | 2330 | ) |
| 2117 | - stop_calls = [] | ||
| 2118 | build_plan_calls = [] | 2331 | build_plan_calls = [] |
| 2119 | original_build_retry_plan = router.rescheduler.build_retry_plan | 2332 | original_build_retry_plan = router.rescheduler.build_retry_plan |
| 2120 | 2333 | ||
| @@ -2124,19 +2337,11 @@ async def test_unified_pd_retry_plan_validation_fails_before_new_attempt_allocat | |||
| 2124 | raise httpx.ReadError("after token cache") | 2337 | raise httpx.ReadError("after token cache") |
| 2125 | yield b"" # pylint: disable=unreachable | 2338 | yield b"" # pylint: disable=unreachable |
| 2126 | 2339 | ||
| 2127 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2128 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason)) | ||
| 2129 | - return None | ||
| 2130 | - | ||
| 2131 | def _build_retry_plan(req_data): | 2340 | def _build_retry_plan(req_data): |
| 2132 | build_plan_calls.append(req_data) | 2341 | build_plan_calls.append(req_data) |
| 2133 | return original_build_retry_plan(req_data) | 2342 | return original_build_retry_plan(req_data) |
| 2134 | 2343 | ||
| 2135 | monkeypatch.setattr(router, "_run_stream_attempt", _run_stream_attempt) | 2344 | monkeypatch.setattr(router, "_run_stream_attempt", _run_stream_attempt) |
| 2136 | - monkeypatch.setattr( | ||
| 2137 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2138 | - _stop, | ||
| 2139 | - ) | ||
| 2140 | monkeypatch.setattr(router.rescheduler, "build_retry_plan", _build_retry_plan) | 2345 | monkeypatch.setattr(router.rescheduler, "build_retry_plan", _build_retry_plan) |
| 2141 | 2346 | ||
| 2142 | response = await router.handle_request() | 2347 | response = await router.handle_request() |
| @@ -2147,10 +2352,6 @@ async def test_unified_pd_retry_plan_validation_fails_before_new_attempt_allocat | |||
| 2147 | assert "parallel sampling" in body["detail"] | 2352 | assert "parallel sampling" in body["detail"] |
| 2148 | assert len(build_plan_calls) == 1 | 2353 | assert len(build_plan_calls) == 1 |
| 2149 | assert scheduler.select_and_allocate.await_count == 2 | 2354 | assert scheduler.select_and_allocate.await_count == 2 |
| 2150 | - assert set(stop_calls) == { | ||
| 2151 | - (PDRole.ROLE_P, 1, DispatchStopReason.PEER_FAILED), | ||
| 2152 | - (PDRole.ROLE_D, 1, DispatchStopReason.PEER_FAILED), | ||
| 2153 | - } | ||
| 2154 | 2355 | ||
| 2155 | 2356 | ||
| 2156 | 2357 | ||
| @@ -2164,8 +2365,10 @@ async def test_unified_pd_handoff_stream_retry_replays_same_prompt_through_prefi | |||
| 2164 | entry_api="v1/completions", | 2365 | entry_api="v1/completions", |
| 2165 | req_len=10, | 2366 | req_len=10, |
| 2166 | ) | 2367 | ) |
| 2167 | - handoff = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | 2368 | + scheduler = _Scheduler( |
| 2168 | - scheduler = _Scheduler(prefill_capabilities=handoff, decode_capabilities=handoff) | 2369 | + prefill_engine_type="vllm", |
| 2370 | + decode_engine_type="vllm", | ||
| 2371 | + ) | ||
| 2169 | config = _config() | 2372 | config = _config() |
| 2170 | config.exception_config.transport_max_retry = 2 | 2373 | config.exception_config.transport_max_retry = 2 |
| 2171 | config.exception_config.reschedule_enabled = True | 2374 | config.exception_config.reschedule_enabled = True |
| @@ -2175,7 +2378,7 @@ async def test_unified_pd_handoff_stream_retry_replays_same_prompt_through_prefi | |||
| 2175 | scheduler=scheduler, | 2378 | scheduler=scheduler, |
| 2176 | request_manager=RequestManager(config), | 2379 | request_manager=RequestManager(config), |
| 2177 | ) | 2380 | ) |
| 2178 | - p_client = _PrefillResultClient("prefill") | 2381 | + p_client = _NativeHandoffPrefillClient("prefill") |
| 2179 | d_client = _SequenceStreamClient( | 2382 | d_client = _SequenceStreamClient( |
| 2180 | "decode", | 2383 | "decode", |
| 2181 | [ | 2384 | [ |
| @@ -2200,14 +2403,7 @@ async def test_unified_pd_handoff_stream_retry_replays_same_prompt_through_prefi | |||
| 2200 | else: | 2403 | else: |
| 2201 | yield d_client | 2404 | yield d_client |
| 2202 | 2405 | ||
| 2203 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2204 | - return None | ||
| 2205 | - | ||
| 2206 | monkeypatch.setattr(router, "_client_for", _client_for) | 2406 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 2207 | - monkeypatch.setattr( | ||
| 2208 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2209 | - _stop, | ||
| 2210 | - ) | ||
| 2211 | 2407 | ||
| 2212 | response = await router.handle_request() | 2408 | response = await router.handle_request() |
| 2213 | messages = await _invoke_asgi_response(response) | 2409 | messages = await _invoke_asgi_response(response) |
| @@ -2218,26 +2414,21 @@ async def test_unified_pd_handoff_stream_retry_replays_same_prompt_through_prefi | |||
| 2218 | assert len(d_client.requests) == 2 | 2414 | assert len(d_client.requests) == 2 |
| 2219 | assert p_client.requests[1]["prompt"] == [1, 2, 10] | 2415 | assert p_client.requests[1]["prompt"] == [1, 2, 10] |
| 2220 | assert p_client.requests[1]["max_tokens"] == 1 | 2416 | assert p_client.requests[1]["max_tokens"] == 1 |
| 2221 | - assert p_client.requests[0][MOTOR_DISPATCH_KEY]["prefill_context_budget"] == { | 2417 | + assert all("_motor_dispatch" not in request for request in p_client.requests) |
| 2222 | - "max_output_tokens": 8, | 2418 | + assert all("_motor_dispatch" not in request for request in d_client.requests) |
| 2223 | - "parameter": "max_tokens", | 2419 | + assert all("_motor_prefill_result" not in request for request in d_client.requests) |
| 2224 | - } | 2420 | + assert p_client.requests[0]["request_id"] == d_client.requests[0]["request_id"] == "root-handoff-reschedule#a1" |
| 2225 | - assert p_client.requests[1][MOTOR_DISPATCH_KEY]["prefill_context_budget"] == { | 2421 | + assert p_client.requests[1]["request_id"] == d_client.requests[1]["request_id"] == "root-handoff-reschedule#a2" |
| 2226 | - "max_output_tokens": 7, | ||
| 2227 | - "parameter": "max_tokens", | ||
| 2228 | - } | ||
| 2229 | assert d_client.requests[1]["prompt"] == [1, 2, 10] | 2422 | assert d_client.requests[1]["prompt"] == [1, 2, 10] |
| 2230 | assert d_client.requests[1]["max_tokens"] == 7 | 2423 | assert d_client.requests[1]["max_tokens"] == 7 |
| 2231 | - retry_prefill_result = d_client.requests[1][MOTOR_PREFILL_RESULT_KEY] | 2424 | + assert d_client.requests[1]["kv_transfer_params"]["remote_request_id"] == "root-handoff-reschedule#a2" |
| 2232 | - assert retry_prefill_result["attempt_seq"] == 2 | ||
| 2233 | - assert retry_prefill_result["pair_id"] == d_client.requests[1][MOTOR_DISPATCH_KEY]["pair_id"] | ||
| 2234 | assert b'"text":"A"' in body | 2425 | assert b'"text":"A"' in body |
| 2235 | assert b'"text":"B"' in body | 2426 | assert b'"text":"B"' in body |
| 2236 | assert b"ReadError" not in body | 2427 | assert b"ReadError" not in body |
| 2237 | 2428 | ||
| 2238 | 2429 | ||
| 2239 | 2430 | ||
| 2240 | -async def test_unified_pd_pool_node_fault_reschedules_and_stops_as_peer_failure( | 2431 | +async def test_unified_pd_pool_node_fault_reschedules_without_peer_stop( |
| 2241 | monkeypatch, | 2432 | monkeypatch, |
| 2242 | ): | 2433 | ): |
| 2243 | req_info = RequestInfo( | 2434 | req_info = RequestInfo( |
| @@ -2260,14 +2451,13 @@ async def test_unified_pd_pool_node_fault_reschedules_and_stops_as_peer_failure( | |||
| 2260 | decode_started = asyncio.Event() | 2451 | decode_started = asyncio.Event() |
| 2261 | decode_calls = [] | 2452 | decode_calls = [] |
| 2262 | prefill_calls = [] | 2453 | prefill_calls = [] |
| 2263 | - stop_calls = [] | ||
| 2264 | 2454 | ||
| 2265 | async def _forward_prefill(self, api, req_data, client, timeout): | 2455 | async def _forward_prefill(self, api, req_data, client, timeout): |
| 2266 | prefill_calls.append(req_data.copy()) | 2456 | prefill_calls.append(req_data.copy()) |
| 2267 | request = httpx.Request("POST", f"/{api}", json=req_data) | 2457 | request = httpx.Request("POST", f"/{api}", json=req_data) |
| 2268 | return httpx.Response( | 2458 | return httpx.Response( |
| 2269 | status_code=200, | 2459 | status_code=200, |
| 2270 | - json={"status": "cached", "id": req_data["request_id"]}, | 2460 | + json={"status": "cached", "id": req_data.get("request_id") or req_data.get("rid")}, |
| 2271 | request=request, | 2461 | request=request, |
| 2272 | ) | 2462 | ) |
| 2273 | 2463 | ||
| @@ -2281,16 +2471,8 @@ async def test_unified_pd_pool_node_fault_reschedules_and_stops_as_peer_failure( | |||
| 2281 | await asyncio.Event().wait() | 2471 | await asyncio.Event().wait() |
| 2282 | yield b'data: {"choices":[{"text":"B","index":0,"token_ids":[11],"finish_reason":"stop"}]}\n\n' | 2472 | yield b'data: {"choices":[{"text":"B","index":0,"token_ids":[11],"finish_reason":"stop"}]}\n\n' |
| 2283 | 2473 | ||
| 2284 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2285 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason)) | ||
| 2286 | - return None | ||
| 2287 | - | ||
| 2288 | monkeypatch.setattr(UnifiedPDRouter, "forward_request", _forward_prefill) | 2474 | monkeypatch.setattr(UnifiedPDRouter, "forward_request", _forward_prefill) |
| 2289 | monkeypatch.setattr(UnifiedPDRouter, "forward_stream_request", _forward_decode) | 2475 | monkeypatch.setattr(UnifiedPDRouter, "forward_stream_request", _forward_decode) |
| 2290 | - monkeypatch.setattr( | ||
| 2291 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2292 | - _stop, | ||
| 2293 | - ) | ||
| 2294 | 2476 | ||
| 2295 | pool = HTTPClientPool() | 2477 | pool = HTTPClientPool() |
| 2296 | p_endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) | 2478 | p_endpoint = next(iter(next(iter(scheduler.p.endpoints.values())).values())) |
| @@ -2325,10 +2507,6 @@ async def test_unified_pd_pool_node_fault_reschedules_and_stops_as_peer_failure( | |||
| 2325 | assert decode_calls[1]["prompt"] == [1, 2, 10] | 2507 | assert decode_calls[1]["prompt"] == [1, 2, 10] |
| 2326 | assert b'"text":"A"' in body | 2508 | assert b'"text":"A"' in body |
| 2327 | assert b'"text":"B"' in body | 2509 | assert b'"text":"B"' in body |
| 2328 | - assert len(stop_calls) == 2 | ||
| 2329 | - assert {call[0] for call in stop_calls} == {PDRole.ROLE_P, PDRole.ROLE_D} | ||
| 2330 | - assert all(call[1] == 1 for call in stop_calls) | ||
| 2331 | - assert all(call[2] == DispatchStopReason.PEER_FAILED for call in stop_calls) | ||
| 2332 | assert req_info.state == ReqState.DECODE_END | 2510 | assert req_info.state == ReqState.DECODE_END |
| 2333 | finally: | 2511 | finally: |
| 2334 | await pool.close_client( | 2512 | await pool.close_client( |
| @@ -2365,7 +2543,6 @@ async def test_unified_pd_stream_error_after_visible_chunk_without_replay_does_n | |||
| 2365 | ) | 2543 | ) |
| 2366 | p_client = _Client("prefill") | 2544 | p_client = _Client("prefill") |
| 2367 | d_client = _StreamClient("decode", exc_after_chunks=httpx.ReadError("after chunk")) | 2545 | d_client = _StreamClient("decode", exc_after_chunks=httpx.ReadError("after chunk")) |
| 2368 | - stop_calls = [] | ||
| 2369 | 2546 | ||
| 2370 | 2547 | ||
| 2371 | async def _client_for(resource: ScheduledResource): | 2548 | async def _client_for(resource: ScheduledResource): |
| @@ -2374,15 +2551,7 @@ async def test_unified_pd_stream_error_after_visible_chunk_without_replay_does_n | |||
| 2374 | else: | 2551 | else: |
| 2375 | yield d_client | 2552 | yield d_client |
| 2376 | 2553 | ||
| 2377 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2378 | - stop_calls.append((resource.instance.role, attempt.attempt_seq, reason.value)) | ||
| 2379 | - return None | ||
| 2380 | - | ||
| 2381 | monkeypatch.setattr(router, "_client_for", _client_for) | 2554 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 2382 | - monkeypatch.setattr( | ||
| 2383 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2384 | - _stop, | ||
| 2385 | - ) | ||
| 2386 | 2555 | ||
| 2387 | response = await router.handle_request() | 2556 | response = await router.handle_request() |
| 2388 | chunks = [chunk async for chunk in response.body_iterator] | 2557 | chunks = [chunk async for chunk in response.body_iterator] |
| @@ -2391,7 +2560,6 @@ async def test_unified_pd_stream_error_after_visible_chunk_without_replay_does_n | |||
| 2391 | error_chunk = chunks[1].decode("utf-8") if isinstance(chunks[1], bytes) else chunks[1] | 2560 | error_chunk = chunks[1].decode("utf-8") if isinstance(chunks[1], bytes) else chunks[1] |
| 2392 | assert "ReadError" in error_chunk | 2561 | assert "ReadError" in error_chunk |
| 2393 | assert len(d_client.requests) == 1 | 2562 | assert len(d_client.requests) == 1 |
| 2394 | - assert len(stop_calls) == 2 | ||
| 2395 | await router._drain_release_tasks() | 2563 | await router._drain_release_tasks() |
| 2396 | assert scheduler.update_workload.await_count == 2 | 2564 | assert scheduler.update_workload.await_count == 2 |
| 2397 | 2565 | ||
| @@ -2437,14 +2605,7 @@ async def test_unified_pd_stream_error_before_first_body_retries_without_token_r | |||
| 2437 | else: | 2605 | else: |
| 2438 | yield d_client | 2606 | yield d_client |
| 2439 | 2607 | ||
| 2440 | - async def _stop(self, resource, attempt, reason, timeout=1.0): | ||
| 2441 | - return None | ||
| 2442 | - | ||
| 2443 | monkeypatch.setattr(router, "_client_for", _client_for) | 2608 | monkeypatch.setattr(router, "_client_for", _client_for) |
| 2444 | - monkeypatch.setattr( | ||
| 2445 | - "motor.coordinator.router.stop_client.DispatchStopClient.stop", | ||
| 2446 | - _stop, | ||
| 2447 | - ) | ||
| 2448 | 2609 | ||
| 2449 | response = await router.handle_request() | 2610 | response = await router.handle_request() |
| 2450 | chunks = [chunk async for chunk in response.body_iterator] | 2611 | chunks = [chunk async for chunk in response.body_iterator] |
| @@ -2453,31 +2614,255 @@ async def test_unified_pd_stream_error_before_first_body_retries_without_token_r | |||
| 2453 | assert chunks == [b'data: {"choices":[{"delta":{"content":"B"},"index":0,"finish_reason":"stop"}]}\n\n'] | 2614 | assert chunks == [b'data: {"choices":[{"delta":{"content":"B"},"index":0,"finish_reason":"stop"}]}\n\n'] |
| 2454 | 2615 | ||
| 2455 | 2616 | ||
| 2456 | -def test_dispatch_carries_effective_output_budget_to_prefill_leg(): | 2617 | +@pytest.mark.asyncio |
| 2457 | - session = PDDispatchSession( | 2618 | +async def test_unified_pd_nonstream_falls_back_to_hybrid_when_decode_pool_exhausted( |
| 2458 | - "request-1", | 2619 | + monkeypatch, |
| 2459 | - prefill_context_budget=PrefillContextBudget( | 2620 | +): |
| 2460 | - max_output_tokens=24, | 2621 | + req_info = RequestInfo( |
| 2461 | - parameter="max_completion_tokens", | 2622 | + req_id="root-nonstream-fallback", |
| 2462 | - ), | 2623 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, |
| 2624 | + api="v1/completions", | ||
| 2625 | + entry_api="v1/completions", | ||
| 2626 | + req_len=10, | ||
| 2463 | ) | 2627 | ) |
| 2464 | - attempt = session.new_attempt(None, None, config=None, consumed_output_tokens=5) | 2628 | + scheduler = _Scheduler( |
| 2629 | + prefill_engine_type="vllm", | ||
| 2630 | + decode_engine_type="vllm", | ||
| 2631 | + ) | ||
| 2632 | + select_and_allocate = scheduler.select_and_allocate | ||
| 2465 | 2633 | ||
| 2466 | - dispatch = attempt.dispatch_for(PDRole.ROLE_P, "prefill_handoff_decode") | 2634 | + async def _select_and_allocate(role, request_info, **kwargs): |
| 2635 | + if role == PDRole.ROLE_D: | ||
| 2636 | + return None | ||
| 2637 | + return await select_and_allocate(role, request_info, **kwargs) | ||
| 2467 | 2638 | ||
| 2468 | - assert dispatch.prefill_context_budget == PrefillContextBudget( | 2639 | + async def _get_unblocked_instances(role): |
| 2469 | - max_output_tokens=19, | 2640 | + if role == PDRole.ROLE_D: |
| 2470 | - parameter="max_completion_tokens", | 2641 | + return [] |
| 2642 | + return [scheduler.p.id] | ||
| 2643 | + | ||
| 2644 | + scheduler.select_and_allocate = AsyncMock(side_effect=_select_and_allocate) | ||
| 2645 | + scheduler.get_unblocked_instances = _get_unblocked_instances | ||
| 2646 | + config = _config() | ||
| 2647 | + router = UnifiedPDRouter( | ||
| 2648 | + req_info, | ||
| 2649 | + config, | ||
| 2650 | + scheduler=scheduler, | ||
| 2651 | + request_manager=RequestManager(config), | ||
| 2652 | + ) | ||
| 2653 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 2654 | + fallback_calls = [] | ||
| 2655 | + | ||
| 2656 | + class _FallbackRouter: | ||
| 2657 | + async def handle_request(self, *, manage_request_context): | ||
| 2658 | + fallback_calls.append(manage_request_context) | ||
| 2659 | + return JSONResponse({"choices": [{"text": "hybrid"}]}) | ||
| 2660 | + | ||
| 2661 | + | ||
| 2662 | + async def _client_for(resource: ScheduledResource): | ||
| 2663 | + assert resource.instance.role == PDRole.ROLE_P | ||
| 2664 | + yield p_client | ||
| 2665 | + | ||
| 2666 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 2667 | + monkeypatch.setattr(router, "_build_hybrid_fallback_router", _FallbackRouter) | ||
| 2668 | + | ||
| 2669 | + response = await router.handle_request() | ||
| 2670 | + | ||
| 2671 | + assert json.loads(response.body)["choices"][0]["text"] == "hybrid" | ||
| 2672 | + assert fallback_calls == [False] | ||
| 2673 | + assert len(p_client.requests) == 1 | ||
| 2674 | + assert scheduler.select_and_allocate.await_count == 2 | ||
| 2675 | + await router._drain_release_tasks() | ||
| 2676 | + assert scheduler.update_workload.await_count == 1 | ||
| 2677 | + release = scheduler.update_workload.await_args.args[0] | ||
| 2678 | + assert release.role == PDRole.ROLE_P | ||
| 2679 | + assert release.workload_action == WorkloadAction.RELEASE_TOKENS | ||
| 2680 | + | ||
| 2681 | + | ||
| 2682 | + | ||
| 2683 | +async def test_unified_pd_runtime_fallback_respects_disabled_switch(monkeypatch): | ||
| 2684 | + req_info = RequestInfo( | ||
| 2685 | + req_id="root-fallback-disabled", | ||
| 2686 | + req_data={"model": "m", "prompt": "hello", "stream": False, "max_tokens": 8}, | ||
| 2687 | + api="v1/completions", | ||
| 2688 | + entry_api="v1/completions", | ||
| 2689 | + req_len=10, | ||
| 2690 | + ) | ||
| 2691 | + scheduler = _Scheduler( | ||
| 2692 | + prefill_engine_type="vllm", | ||
| 2693 | + decode_engine_type="vllm", | ||
| 2694 | + ) | ||
| 2695 | + select_and_allocate = scheduler.select_and_allocate | ||
| 2696 | + | ||
| 2697 | + async def _select_and_allocate(role, request_info, **kwargs): | ||
| 2698 | + if role == PDRole.ROLE_D: | ||
| 2699 | + return None | ||
| 2700 | + return await select_and_allocate(role, request_info, **kwargs) | ||
| 2701 | + | ||
| 2702 | + scheduler.select_and_allocate = AsyncMock(side_effect=_select_and_allocate) | ||
| 2703 | + scheduler.get_unblocked_instances = AsyncMock(return_value=[scheduler.p.id]) | ||
| 2704 | + config = _config() | ||
| 2705 | + config.scheduler_config.enable_pd_separation_fallback_to_hybrid = False | ||
| 2706 | + router = UnifiedPDRouter( | ||
| 2707 | + req_info, | ||
| 2708 | + config, | ||
| 2709 | + scheduler=scheduler, | ||
| 2710 | + request_manager=RequestManager(config), | ||
| 2711 | + ) | ||
| 2712 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 2713 | + | ||
| 2714 | + | ||
| 2715 | + async def _client_for(resource: ScheduledResource): | ||
| 2716 | + assert resource.instance.role == PDRole.ROLE_P | ||
| 2717 | + yield p_client | ||
| 2718 | + | ||
| 2719 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 2720 | + monkeypatch.setattr( | ||
| 2721 | + router, | ||
| 2722 | + "_build_hybrid_fallback_router", | ||
| 2723 | + lambda: pytest.fail("disabled runtime fallback must not build PDHybridRouter"), | ||
| 2471 | ) | 2724 | ) |
| 2472 | 2725 | ||
| 2726 | + with pytest.raises(HTTPException, match="No instance available for role"): | ||
| 2727 | + await router.handle_request() | ||
| 2473 | 2728 | ||
| 2474 | -def test_unified_pd_prefers_max_completion_tokens_and_preserves_parameter(): | 2729 | + scheduler.get_unblocked_instances.assert_not_awaited() |
| 2475 | - router = SimpleNamespace(req_info=SimpleNamespace(req_data={"max_tokens": 32, "max_completion_tokens": 24})) | ||
| 2476 | 2730 | ||
| 2477 | - assert UnifiedPDRouter._prefill_context_budget(router) == PrefillContextBudget( | 2731 | + |
| 2478 | - max_output_tokens=24, | 2732 | +@pytest.mark.asyncio |
| 2479 | - parameter="max_completion_tokens", | 2733 | +async def test_unified_pd_stream_restarts_on_hybrid_before_commit_when_decode_pool_exhausted( |
| 2734 | + monkeypatch, | ||
| 2735 | +): | ||
| 2736 | + req_info = RequestInfo( | ||
| 2737 | + req_id="root-stream-fallback-restart", | ||
| 2738 | + req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | ||
| 2739 | + api="v1/completions", | ||
| 2740 | + entry_api="v1/completions", | ||
| 2741 | + req_len=10, | ||
| 2480 | ) | 2742 | ) |
| 2743 | + scheduler = _Scheduler( | ||
| 2744 | + prefill_engine_type="vllm", | ||
| 2745 | + decode_engine_type="vllm", | ||
| 2746 | + ) | ||
| 2747 | + select_and_allocate = scheduler.select_and_allocate | ||
| 2748 | + | ||
| 2749 | + async def _select_and_allocate(role, request_info, **kwargs): | ||
| 2750 | + if role == PDRole.ROLE_D: | ||
| 2751 | + return None | ||
| 2752 | + return await select_and_allocate(role, request_info, **kwargs) | ||
| 2753 | + | ||
| 2754 | + async def _get_unblocked_instances(role): | ||
| 2755 | + if role == PDRole.ROLE_D: | ||
| 2756 | + return [] | ||
| 2757 | + return [scheduler.p.id] | ||
| 2758 | + | ||
| 2759 | + scheduler.select_and_allocate = AsyncMock(side_effect=_select_and_allocate) | ||
| 2760 | + scheduler.get_unblocked_instances = _get_unblocked_instances | ||
| 2761 | + config = _config() | ||
| 2762 | + router = UnifiedPDRouter( | ||
| 2763 | + req_info, | ||
| 2764 | + config, | ||
| 2765 | + scheduler=scheduler, | ||
| 2766 | + request_manager=RequestManager(config), | ||
| 2767 | + ) | ||
| 2768 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 2769 | + fallback_calls = [] | ||
| 2770 | + | ||
| 2771 | + class _FallbackRouter: | ||
| 2772 | + async def stream_fallback_from_existing_context(self, **kwargs): | ||
| 2773 | + fallback_calls.append(kwargs) | ||
| 2774 | + kwargs["mark_unified_ready"]() | ||
| 2775 | + yield b'data: {"choices":[{"text":"hybrid","finish_reason":"stop"}]}\n\n' | ||
| 2776 | + | ||
| 2777 | + | ||
| 2778 | + async def _client_for(resource: ScheduledResource): | ||
| 2779 | + assert resource.instance.role == PDRole.ROLE_P | ||
| 2780 | + yield p_client | ||
| 2781 | + | ||
| 2782 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 2783 | + monkeypatch.setattr(router, "_build_hybrid_fallback_router", _FallbackRouter) | ||
| 2784 | + | ||
| 2785 | + response = await router.handle_request() | ||
| 2786 | + chunks = [chunk async for chunk in response.body_iterator] | ||
| 2787 | + | ||
| 2788 | + assert chunks == [b'data: {"choices":[{"text":"hybrid","finish_reason":"stop"}]}\n\n'] | ||
| 2789 | + assert len(fallback_calls) == 1 | ||
| 2790 | + assert fallback_calls[0]["attempt_id"] == 2 | ||
| 2791 | + assert fallback_calls[0]["req_data"]["prompt"] == "hello" | ||
| 2792 | + assert fallback_calls[0]["mark_unified_ready"] is not None | ||
| 2793 | + | ||
| 2794 | + | ||
| 2795 | + | ||
| 2796 | +async def test_unified_pd_stream_resumes_on_hybrid_with_token_replay_after_commit( | ||
| 2797 | + monkeypatch, | ||
| 2798 | +): | ||
| 2799 | + req_info = RequestInfo( | ||
| 2800 | + req_id="root-stream-fallback-resume", | ||
| 2801 | + req_data={"model": "m", "prompt": "hello", "stream": True, "max_tokens": 8}, | ||
| 2802 | + api="v1/completions", | ||
| 2803 | + entry_api="v1/completions", | ||
| 2804 | + req_len=10, | ||
| 2805 | + ) | ||
| 2806 | + scheduler = _Scheduler( | ||
| 2807 | + prefill_engine_type="vllm", | ||
| 2808 | + decode_engine_type="vllm", | ||
| 2809 | + ) | ||
| 2810 | + | ||
| 2811 | + async def _get_unblocked_instances(role): | ||
| 2812 | + if role == PDRole.ROLE_D: | ||
| 2813 | + return [] | ||
| 2814 | + return [scheduler.p.id] | ||
| 2815 | + | ||
| 2816 | + scheduler.get_unblocked_instances = _get_unblocked_instances | ||
| 2817 | + config = _config() | ||
| 2818 | + config.exception_config.transport_max_retry = 2 | ||
| 2819 | + config.exception_config.reschedule_enabled = True | ||
| 2820 | + router = UnifiedPDRouter( | ||
| 2821 | + req_info, | ||
| 2822 | + config, | ||
| 2823 | + scheduler=scheduler, | ||
| 2824 | + request_manager=RequestManager(config), | ||
| 2825 | + ) | ||
| 2826 | + p_client = _NativeHandoffPrefillClient("prefill") | ||
| 2827 | + d_client = _SequenceStreamClient( | ||
| 2828 | + "decode", | ||
| 2829 | + [ | ||
| 2830 | + _StreamResponse( | ||
| 2831 | + [ | ||
| 2832 | + b'data: {"choices":[{"text":"A","index":0,"prompt_token_ids":[1,2],"token_ids":[10]}]}\n\n', | ||
| 2833 | + ], | ||
| 2834 | + exc_after_chunks=httpx.ReadError("decode disappeared"), | ||
| 2835 | + ) | ||
| 2836 | + ], | ||
| 2837 | + ) | ||
| 2838 | + fallback_calls = [] | ||
| 2839 | + | ||
| 2840 | + class _FallbackRouter: | ||
| 2841 | + async def stream_fallback_from_existing_context(self, **kwargs): | ||
| 2842 | + fallback_calls.append(kwargs) | ||
| 2843 | + yield b'data: {"choices":[{"text":"B","token_ids":[11],"finish_reason":"stop"}]}\n\n' | ||
| 2844 | + | ||
| 2845 | + | ||
| 2846 | + async def _client_for(resource: ScheduledResource): | ||
| 2847 | + if resource.instance.role == PDRole.ROLE_P: | ||
| 2848 | + yield p_client | ||
| 2849 | + else: | ||
| 2850 | + yield d_client | ||
| 2851 | + | ||
| 2852 | + monkeypatch.setattr(router, "_client_for", _client_for) | ||
| 2853 | + monkeypatch.setattr(router, "_build_hybrid_fallback_router", _FallbackRouter) | ||
| 2854 | + | ||
| 2855 | + response = await router.handle_request() | ||
| 2856 | + chunks = [chunk async for chunk in response.body_iterator] | ||
| 2857 | + body = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode() for chunk in chunks) | ||
| 2858 | + | ||
| 2859 | + assert b'"text":"A"' in body | ||
| 2860 | + assert b'"text":"B"' in body | ||
| 2861 | + assert len(fallback_calls) == 1 | ||
| 2862 | + assert fallback_calls[0]["is_resume"] is True | ||
| 2863 | + assert fallback_calls[0]["api"] == "v1/completions" | ||
| 2864 | + assert fallback_calls[0]["req_data"]["prompt"] == [1, 2, 10] | ||
| 2865 | + assert fallback_calls[0]["req_data"]["max_tokens"] == 7 | ||
| 2481 | 2866 | ||
| 2482 | 2867 | ||
| 2483 | def test_retry_plan_preserves_max_completion_tokens_precedence_for_completion_replay(): | 2868 | def test_retry_plan_preserves_max_completion_tokens_precedence_for_completion_replay(): |
| @@ -250,8 +250,7 @@ def _make_error(status_code: int, body: bytes = b"") -> UpstreamHTTPError: | |||
| 250 | 250 | ||
| 251 | def test_cb_not_reportable_for_4xx(): | 251 | def test_cb_not_reportable_for_4xx(): |
| 252 | """4xx errors are client errors (e.g. input too long, bad params); the instance | 252 | """4xx errors are client errors (e.g. input too long, bad params); the instance |
| 253 | - must not be penalised. The engine_server is responsible for mapping all known | 253 | + must not be penalised. Native engines return request-validation failures as 4xx. |
| 254 | - request-validation exceptions to 4xx before they reach the coordinator. | ||
| 255 | """ | 254 | """ |
| 256 | assert not is_cb_reportable_failure(_make_error(400)) | 255 | assert not is_cb_reportable_failure(_make_error(400)) |
| 257 | assert not is_cb_reportable_failure(_make_error(422)) | 256 | assert not is_cb_reportable_failure(_make_error(422)) |
| @@ -264,8 +263,8 @@ def test_cb_reportable_for_5xx(): | |||
| 264 | 263 | ||
| 265 | 264 | ||
| 266 | def test_cb_not_reportable_for_vllm_validation_error_returned_as_400(): | 265 | def test_cb_not_reportable_for_vllm_validation_error_returned_as_400(): |
| 267 | - """VLLMValidationError (max_tokens > max_model_len) must be mapped to HTTP 400 by | 266 | + """A native vLLM validation failure returned as HTTP 400 must not trip the |
| 268 | - the engine_server so the coordinator never sees it as a 5xx fault. | 267 | + coordinator circuit breaker. |
| 269 | """ | 268 | """ |
| 270 | body = json.dumps({"detail": "max_tokens=1024 cannot be greater than max_model_len=50"}).encode() | 269 | body = json.dumps({"detail": "max_tokens=1024 cannot be greater than max_model_len=50"}).encode() |
| 271 | error = _make_error(400, body) | 270 | error = _make_error(400, body) |
| @@ -14,7 +14,6 @@ from unittest.mock import AsyncMock, MagicMock, patch | |||
| 14 | 14 | ||
| 15 | import pytest | 15 | import pytest |
| 16 | 16 | ||
| 17 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 18 | from motor.common.resources.instance import Instance, PDRole | 17 | from motor.common.resources.instance import Instance, PDRole |
| 19 | from motor.common.resources.endpoint import Endpoint, EndpointStatus, Workload | 18 | from motor.common.resources.endpoint import Endpoint, EndpointStatus, Workload |
| 20 | from motor.coordinator.domain.scheduling import InstanceReadiness | 19 | from motor.coordinator.domain.scheduling import InstanceReadiness |
| @@ -192,8 +191,6 @@ async def test_has_required_instances_uses_cache_after_warmup(): | |||
| 192 | ("10.0.0.3", "8003", EndpointStatus.NORMAL), | 191 | ("10.0.0.3", "8003", EndpointStatus.NORMAL), |
| 193 | ], | 192 | ], |
| 194 | ) | 193 | ) |
| 195 | - p_inst.dispatch_capabilities = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | ||
| 196 | - d_inst.dispatch_capabilities = [DispatchPlan.PREFILL_HANDOFF_DECODE.value] | ||
| 197 | await client._cache.replace_all(PDRole.ROLE_E, [e_inst]) | 194 | await client._cache.replace_all(PDRole.ROLE_E, [e_inst]) |
| 198 | await client._cache.replace_all(PDRole.ROLE_P, [p_inst]) | 195 | await client._cache.replace_all(PDRole.ROLE_P, [p_inst]) |
| 199 | await client._cache.replace_all(PDRole.ROLE_D, [d_inst]) | 196 | await client._cache.replace_all(PDRole.ROLE_D, [d_inst]) |
| @@ -15,7 +15,6 @@ from unittest.mock import AsyncMock, Mock, call, patch | |||
| 15 | 15 | ||
| 16 | import pytest | 16 | import pytest |
| 17 | 17 | ||
| 18 | -from motor.common.resources.dispatch import DispatchPlan | ||
| 19 | from motor.common.resources.instance import Instance, PDRole | 18 | from motor.common.resources.instance import Instance, PDRole |
| 20 | from motor.common.resources.endpoint import Endpoint, Workload, WorkloadAction, EndpointStatus | 19 | from motor.common.resources.endpoint import Endpoint, Workload, WorkloadAction, EndpointStatus |
| 21 | from motor.coordinator.domain import InstanceReadiness, UpdateWorkloadParams | 20 | from motor.coordinator.domain import InstanceReadiness, UpdateWorkloadParams |
| @@ -67,7 +66,7 @@ def _make_instance( | |||
| 67 | instance_id: int = 1, | 66 | instance_id: int = 1, |
| 68 | role: str = "prefill", | 67 | role: str = "prefill", |
| 69 | endpoints: dict | None = None, | 68 | endpoints: dict | None = None, |
| 70 | - dispatch_capabilities: list[str] | None = None, | 69 | + engine_type: str | None = None, |
| 71 | ) -> Instance: | 70 | ) -> Instance: |
| 72 | """Create a real Instance (used by _SchedulerInstanceCache tests).""" | 71 | """Create a real Instance (used by _SchedulerInstanceCache tests).""" |
| 73 | if endpoints is None: | 72 | if endpoints is None: |
| @@ -76,10 +75,10 @@ def _make_instance( | |||
| 76 | return Instance( | 75 | return Instance( |
| 77 | job_name="test-job", | 76 | job_name="test-job", |
| 78 | model_name="test-model", | 77 | model_name="test-model", |
| 78 | + engine_type=engine_type, | ||
| 79 | id=instance_id, | 79 | id=instance_id, |
| 80 | role=role, | 80 | role=role, |
| 81 | endpoints=endpoints, | 81 | endpoints=endpoints, |
| 82 | - dispatch_capabilities=dispatch_capabilities or [], | ||
| 83 | ) | 82 | ) |
| 84 | 83 | ||
| 85 | 84 | ||
| @@ -419,6 +418,34 @@ class TestAsyncSchedulerClient: | |||
| 419 | ) | 418 | ) |
| 420 | assert result == [] | 419 | assert result == [] |
| 421 | 420 | ||
| 421 | + | ||
| 422 | + async def test_select_endpoint_candidates_filters_required_engine_type(self): | ||
| 423 | + """Decode candidate selection must not mix native engine protocols.""" | ||
| 424 | + vllm = _make_instance( | ||
| 425 | + instance_id=1, | ||
| 426 | + role="decode", | ||
| 427 | + endpoints={"pod1": {1: _make_endpoint(endpoint_id=1)}}, | ||
| 428 | + engine_type="vllm", | ||
| 429 | + ) | ||
| 430 | + sglang = _make_instance( | ||
| 431 | + instance_id=2, | ||
| 432 | + role="decode", | ||
| 433 | + endpoints={"pod2": {2: _make_endpoint(endpoint_id=2)}}, | ||
| 434 | + engine_type="sglang", | ||
| 435 | + ) | ||
| 436 | + self.mock_cache.get_instances.return_value = [vllm, sglang] | ||
| 437 | + req_info = Mock(spec=RequestInfo) | ||
| 438 | + req_info.req_id = "req-engine-filter" | ||
| 439 | + req_info.req_len = 10 | ||
| 440 | + | ||
| 441 | + candidates, _ = await self.client._select_endpoint_candidates_with_policy( | ||
| 442 | + req_info, | ||
| 443 | + PDRole.ROLE_D, | ||
| 444 | + required_engine_type=" SGLang ", | ||
| 445 | + ) | ||
| 446 | + | ||
| 447 | + assert [(instance.id, endpoint.id) for instance, endpoint, _ in candidates] == [(2, 2)] | ||
| 448 | + | ||
| 422 | # -- test_select_and_allocate ------------------------------------------- | 449 | # -- test_select_and_allocate ------------------------------------------- |
| 423 | 450 | ||
| 424 | 451 | ||
| @@ -666,9 +693,8 @@ class TestAsyncSchedulerClient: | |||
| 666 | 693 | ||
| 667 | async def test_has_required_instances_met(self): | 694 | async def test_has_required_instances_met(self): |
| 668 | """has_required_instances returns REQUIRED_MET when P and D present.""" | 695 | """has_required_instances returns REQUIRED_MET when P and D present.""" |
| 669 | - capability = [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | 696 | + mock_p = _make_instance(1, PDRole.ROLE_P) |
| 670 | - mock_p = _make_instance(1, PDRole.ROLE_P, dispatch_capabilities=capability) | 697 | + mock_d = _make_instance(2, PDRole.ROLE_D) |
| 671 | - mock_d = _make_instance(2, PDRole.ROLE_D, dispatch_capabilities=capability) | ||
| 672 | 698 | ||
| 673 | def _get_instances_side_effect(role): | 699 | def _get_instances_side_effect(role): |
| 674 | mapping = {PDRole.ROLE_P: [mock_p], PDRole.ROLE_D: [mock_d]} | 700 | mapping = {PDRole.ROLE_P: [mock_p], PDRole.ROLE_D: [mock_d]} |
| @@ -696,31 +722,6 @@ class TestAsyncSchedulerClient: | |||
| 696 | assert result.is_ready() is False | 722 | assert result.is_ready() is False |
| 697 | self.mock_transport.send_request.assert_not_awaited() | 723 | self.mock_transport.send_request.assert_not_awaited() |
| 698 | 724 | ||
| 699 | - | ||
| 700 | - async def test_has_required_instances_rejects_incompatible_pd_pair(self): | ||
| 701 | - prefill = _make_instance( | ||
| 702 | - 1, | ||
| 703 | - PDRole.ROLE_P, | ||
| 704 | - dispatch_capabilities=[DispatchPlan.CONCURRENT_ENGINE_SYNC.value], | ||
| 705 | - ) | ||
| 706 | - decode = _make_instance( | ||
| 707 | - 2, | ||
| 708 | - PDRole.ROLE_D, | ||
| 709 | - dispatch_capabilities=[DispatchPlan.PREFILL_HANDOFF_DECODE.value], | ||
| 710 | - ) | ||
| 711 | - | ||
| 712 | - def _get_instances_side_effect(role): | ||
| 713 | - mapping = {PDRole.ROLE_P: [prefill], PDRole.ROLE_D: [decode]} | ||
| 714 | - return mapping.get(role, []) | ||
| 715 | - | ||
| 716 | - self.mock_cache.get_instances.side_effect = _get_instances_side_effect | ||
| 717 | - | ||
| 718 | - result = await self.client.has_required_instances() | ||
| 719 | - | ||
| 720 | - assert result == InstanceReadiness.UNKNOWN | ||
| 721 | - assert result.is_run() is False | ||
| 722 | - self.mock_transport.send_request.assert_not_awaited() | ||
| 723 | - | ||
| 724 | 725 | ||
| 725 | async def test_has_required_instances_union_wins_over_partial_roles(self): | 726 | async def test_has_required_instances_union_wins_over_partial_roles(self): |
| 726 | decode = _make_instance(1, PDRole.ROLE_D) | 727 | decode = _make_instance(1, PDRole.ROLE_D) |
| @@ -95,12 +95,14 @@ def _make_instance( | |||
| 95 | instance_id: int, | 95 | instance_id: int, |
| 96 | endpoint_ids: tuple[int, ...], | 96 | endpoint_ids: tuple[int, ...], |
| 97 | role: PDRole = PDRole.ROLE_P, | 97 | role: PDRole = PDRole.ROLE_P, |
| 98 | + engine_type: str | None = None, | ||
| 98 | ) -> Instance: | 99 | ) -> Instance: |
| 99 | inst = Instance( | 100 | inst = Instance( |
| 100 | job_name=f"{role.value}-{instance_id}", | 101 | job_name=f"{role.value}-{instance_id}", |
| 101 | model_name="test_model", | 102 | model_name="test_model", |
| 102 | id=instance_id, | 103 | id=instance_id, |
| 103 | role=role, | 104 | role=role, |
| 105 | + engine_type=engine_type, | ||
| 104 | status=InsStatus.ACTIVE, | 106 | status=InsStatus.ACTIVE, |
| 105 | parallel_config=ParallelConfig(dp_size=len(endpoint_ids)), | 107 | parallel_config=ParallelConfig(dp_size=len(endpoint_ids)), |
| 106 | ) | 108 | ) |
| @@ -197,8 +199,8 @@ class TestSerializeInstanceMinimal: | |||
| 197 | assert result["job_name"] == inst.job_name | 199 | assert result["job_name"] == inst.job_name |
| 198 | assert result["model_name"] == "test_model" | 200 | assert result["model_name"] == "test_model" |
| 199 | assert result["engine_type"] is None | 201 | assert result["engine_type"] is None |
| 200 | - assert result["dispatch_capabilities"] == [] | 202 | + assert "dispatch_capabilities" not in result |
| 201 | - assert len(result) == 6 | 203 | + assert len(result) == 5 |
| 202 | 204 | ||
| 203 | 205 | ||
| 204 | class TestSerializeEndpointMinimal: | 206 | class TestSerializeEndpointMinimal: |
| @@ -224,6 +226,20 @@ class TestSerializeEndpointMinimal: | |||
| 224 | result = _serialize_endpoint_minimal(ep) | 226 | result = _serialize_endpoint_minimal(ep) |
| 225 | assert result["status"] == EndpointStatus.NORMAL.value | 227 | assert result["status"] == EndpointStatus.NORMAL.value |
| 226 | 228 | ||
| 229 | + def test_endpoint_serializes_bootstrap_port(self): | ||
| 230 | + ep = Endpoint( | ||
| 231 | + id=14, | ||
| 232 | + ip="10.0.0.14", | ||
| 233 | + business_port="8014", | ||
| 234 | + mgmt_port="9014", | ||
| 235 | + bootstrap_port=9114, | ||
| 236 | + ) | ||
| 237 | + | ||
| 238 | + result = _serialize_endpoint_minimal(ep) | ||
| 239 | + restored = Endpoint.model_validate(result) | ||
| 240 | + | ||
| 241 | + assert restored.bootstrap_port == 9114 | ||
| 242 | + | ||
| 227 | def test_endpoint_with_empty_mgmt_port_defaults_empty_string(self): | 243 | def test_endpoint_with_empty_mgmt_port_defaults_empty_string(self): |
| 228 | """mgmt_port='' → serializer returns empty string (falsy → or '' branch).""" | 244 | """mgmt_port='' → serializer returns empty string (falsy → or '' branch).""" |
| 229 | ep = Endpoint(id=13, ip="9.9.9.9", business_port="9000", mgmt_port="") | 245 | ep = Endpoint(id=13, ip="9.9.9.9", business_port="9000", mgmt_port="") |
| @@ -772,6 +788,30 @@ class TestHandleAllocateOnlyEdgeCases: | |||
| 772 | assert response.data["instance"] is None | 788 | assert response.data["instance"] is None |
| 773 | assert response.data["endpoint"] is None | 789 | assert response.data["endpoint"] is None |
| 774 | 790 | ||
| 791 | + | ||
| 792 | + async def test_required_engine_type_rejects_mismatched_candidate(self): | ||
| 793 | + dispatcher, instance_manager, *_ = _make_dispatcher(scheduler_type=SchedulerType.ROUND_ROBIN) | ||
| 794 | + instance = _make_instance(1, (10,), PDRole.ROLE_D, engine_type="vllm") | ||
| 795 | + await instance_manager.refresh_instances(EventType.ADD, [instance]) | ||
| 796 | + | ||
| 797 | + request = SchedulerRequest( | ||
| 798 | + request_type=SchedulerRequestType.ALLOCATE_ONLY, | ||
| 799 | + request_id="req-engine-filter", | ||
| 800 | + data={ | ||
| 801 | + "instance_id": 1, | ||
| 802 | + "endpoint_id": 10, | ||
| 803 | + "role": PDRole.ROLE_D.value, | ||
| 804 | + "required_engine_type": "sglang", | ||
| 805 | + "workload_active_tokens": 1, | ||
| 806 | + }, | ||
| 807 | + ) | ||
| 808 | + | ||
| 809 | + response = await dispatcher.dispatch(request) | ||
| 810 | + | ||
| 811 | + assert response.response_type == SchedulerResponseType.SUCCESS | ||
| 812 | + assert response.data["instance"] is None | ||
| 813 | + assert response.data["endpoint"] is None | ||
| 814 | + | ||
| 775 | 815 | ||
| 776 | class TestParseOptionalInt: | 816 | class TestParseOptionalInt: |
| 777 | def test_none_returns_none(self): | 817 | def test_none_returns_none(self): |
| @@ -172,18 +172,21 @@ class TestCoordinatorServer: | |||
| 172 | im_mock_cls.return_value = im_instance | 172 | im_mock_cls.return_value = im_instance |
| 173 | 173 | ||
| 174 | # Mock handle_request to return appropriate JSON response | 174 | # Mock handle_request to return appropriate JSON response |
| 175 | - async def mock_handle_request(request, config, scheduler=None, request_manager=None): | 175 | + async def mock_handle_request(request, config, scheduler=None, request_manager=None, request_json=None): |
| 176 | """Mock handle_request that returns JSON response matching test expectations""" | 176 | """Mock handle_request that returns JSON response matching test expectations""" |
| 177 | - try: | 177 | + if request_json is not None: |
| 178 | - # Try to get JSON from request (cached if already parsed) | 178 | + body_json = request_json |
| 179 | - body_json = await request.json() | 179 | + else: |
| 180 | - except Exception: | ||
| 181 | - # Fallback: try to read body directly | ||
| 182 | try: | 180 | try: |
| 183 | - request_body = await request.body() | 181 | + # Try to get JSON from request (cached if already parsed) |
| 184 | - body_json = json.loads(request_body.decode("utf-8")) | 182 | + body_json = await request.json() |
| 185 | except Exception: | 183 | except Exception: |
| 186 | - body_json = {} | 184 | + # Fallback: try to read body directly |
| 185 | + try: | ||
| 186 | + request_body = await request.body() | ||
| 187 | + body_json = json.loads(request_body.decode("utf-8")) | ||
| 188 | + except Exception: | ||
| 189 | + body_json = {} | ||
| 187 | 190 | ||
| 188 | # Extract input_data based on request type | 191 | # Extract input_data based on request type |
| 189 | input_data = "" | 192 | input_data = "" |
| @@ -800,18 +803,21 @@ class TestCoordinatorServerAdvanced: | |||
| 800 | im_mock_cls.return_value = im_instance | 803 | im_mock_cls.return_value = im_instance |
| 801 | 804 | ||
| 802 | # Mock handle_request to return appropriate JSON response | 805 | # Mock handle_request to return appropriate JSON response |
| 803 | - async def mock_handle_request(request, config, scheduler=None, request_manager=None): | 806 | + async def mock_handle_request(request, config, scheduler=None, request_manager=None, request_json=None): |
| 804 | """Mock handle_request that returns JSON response matching test expectations""" | 807 | """Mock handle_request that returns JSON response matching test expectations""" |
| 805 | - try: | 808 | + if request_json is not None: |
| 806 | - # Try to get JSON from request (cached if already parsed) | 809 | + body_json = request_json |
| 807 | - body_json = await request.json() | 810 | + else: |
| 808 | - except Exception: | ||
| 809 | - # Fallback: try to read body directly | ||
| 810 | try: | 811 | try: |
| 811 | - request_body = await request.body() | 812 | + # Try to get JSON from request (cached if already parsed) |
| 812 | - body_json = json.loads(request_body.decode("utf-8")) | 813 | + body_json = await request.json() |
| 813 | except Exception: | 814 | except Exception: |
| 814 | - body_json = {} | 815 | + # Fallback: try to read body directly |
| 816 | + try: | ||
| 817 | + request_body = await request.body() | ||
| 818 | + body_json = json.loads(request_body.decode("utf-8")) | ||
| 819 | + except Exception: | ||
| 820 | + body_json = {} | ||
| 815 | 821 | ||
| 816 | # Extract input_data based on request type | 822 | # Extract input_data based on request type |
| 817 | input_data = "" | 823 | input_data = "" |
| @@ -1850,11 +1856,14 @@ class TestAnthropicEndpoints: | |||
| 1850 | im_mock_cls.return_value = im_instance | 1856 | im_mock_cls.return_value = im_instance |
| 1851 | 1857 | ||
| 1852 | # Mock handle_request to return appropriate response | 1858 | # Mock handle_request to return appropriate response |
| 1853 | - async def mock_handle_request(request, config, scheduler=None, request_manager=None): | 1859 | + async def mock_handle_request(request, config, scheduler=None, request_manager=None, request_json=None): |
| 1854 | - try: | 1860 | + if request_json is not None: |
| 1855 | - body_json = await request.json() | 1861 | + body_json = request_json |
| 1856 | - except Exception: | 1862 | + else: |
| 1857 | - body_json = {} | 1863 | + try: |
| 1864 | + body_json = await request.json() | ||
| 1865 | + except Exception: | ||
| 1866 | + body_json = {} | ||
| 1858 | 1867 | ||
| 1859 | if "messages" in body_json: | 1868 | if "messages" in body_json: |
| 1860 | json.dumps(body_json["messages"], ensure_ascii=False) | 1869 | json.dumps(body_json["messages"], ensure_ascii=False) |
| @@ -0,0 +1,230 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""Opt-in smoke tests against a real native vLLM or SGLang P/D pair. | ||
| 12 | + | ||
| 13 | +Required environment: | ||
| 14 | + | ||
| 15 | +* ``MOTOR_RUN_NATIVE_PD_SMOKE=1`` | ||
| 16 | +* ``MOTOR_NATIVE_SMOKE_ENGINE_TYPE=vllm|sglang`` | ||
| 17 | +* ``MOTOR_NATIVE_SMOKE_PREFILL_URL=http[s]://host:port`` | ||
| 18 | +* ``MOTOR_NATIVE_SMOKE_DECODE_URL=http[s]://host:port`` | ||
| 19 | +* ``MOTOR_NATIVE_SMOKE_MODEL=<served model name>`` | ||
| 20 | +* ``MOTOR_NATIVE_SMOKE_BOOTSTRAP_PORT=<port>`` for SGLang | ||
| 21 | + | ||
| 22 | +HTTPS additionally requires ``MOTOR_NATIVE_SMOKE_CA_FILE``, | ||
| 23 | +``MOTOR_NATIVE_SMOKE_CERT_FILE``, and ``MOTOR_NATIVE_SMOKE_KEY_FILE``. | ||
| 24 | +""" | ||
| 25 | + | ||
| 26 | +import json | ||
| 27 | +import os | ||
| 28 | +from urllib.parse import urlparse | ||
| 29 | + | ||
| 30 | +import pytest | ||
| 31 | + | ||
| 32 | +from motor.common.http import HTTPClientPool | ||
| 33 | +from motor.common.resources.endpoint import Endpoint, EndpointStatus, Workload | ||
| 34 | +from motor.common.resources.instance import Instance, InsStatus, ParallelConfig, PDRole | ||
| 35 | +from motor.config.coordinator import CoordinatorConfig, ExceptionConfig | ||
| 36 | +from motor.config.tls_config import TLSConfig | ||
| 37 | +from motor.coordinator.domain.request_manager import RequestManager | ||
| 38 | +from motor.coordinator.models.request import RequestInfo, ReqState | ||
| 39 | +from motor.coordinator.router.strategies.unified_pd import UnifiedPDRouter | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +pytestmark = pytest.mark.integration | ||
| 43 | + | ||
| 44 | + | ||
| 45 | +def _required_env(name: str) -> str: | ||
| 46 | + value = os.getenv(name, "").strip() | ||
| 47 | + if not value: | ||
| 48 | + pytest.fail(f"{name} is required when MOTOR_RUN_NATIVE_PD_SMOKE=1") | ||
| 49 | + return value | ||
| 50 | + | ||
| 51 | + | ||
| 52 | +def _parse_endpoint_url(name: str) -> tuple[str, int, str]: | ||
| 53 | + raw = _required_env(name) | ||
| 54 | + parsed = urlparse(raw) | ||
| 55 | + if parsed.scheme not in {"http", "https"} or not parsed.hostname: | ||
| 56 | + pytest.fail(f"{name} must be an http[s]://host:port URL") | ||
| 57 | + if parsed.path not in {"", "/"} or parsed.query or parsed.fragment: | ||
| 58 | + pytest.fail(f"{name} must not contain a path, query, or fragment") | ||
| 59 | + port = parsed.port or (443 if parsed.scheme == "https" else 80) | ||
| 60 | + return parsed.hostname, port, parsed.scheme | ||
| 61 | + | ||
| 62 | + | ||
| 63 | +def _instance( | ||
| 64 | + instance_id: int, | ||
| 65 | + role: PDRole, | ||
| 66 | + engine_type: str, | ||
| 67 | + host: str, | ||
| 68 | + port: int, | ||
| 69 | + *, | ||
| 70 | + bootstrap_port: int | None = None, | ||
| 71 | +) -> Instance: | ||
| 72 | + endpoint = Endpoint( | ||
| 73 | + id=instance_id, | ||
| 74 | + ip=host, | ||
| 75 | + business_port=str(port), | ||
| 76 | + mgmt_port=str(port), | ||
| 77 | + bootstrap_port=bootstrap_port, | ||
| 78 | + status=EndpointStatus.NORMAL, | ||
| 79 | + ) | ||
| 80 | + return Instance( | ||
| 81 | + job_name=f"smoke-{engine_type}-{role.value}", | ||
| 82 | + model_name=_required_env("MOTOR_NATIVE_SMOKE_MODEL"), | ||
| 83 | + engine_type=engine_type, | ||
| 84 | + id=instance_id, | ||
| 85 | + role=role, | ||
| 86 | + status=InsStatus.ACTIVE, | ||
| 87 | + parallel_config=ParallelConfig(dp_size=1, tp_size=1), | ||
| 88 | + endpoints={host: {instance_id: endpoint}}, | ||
| 89 | + ) | ||
| 90 | + | ||
| 91 | + | ||
| 92 | +class _RealNativeScheduler: | ||
| 93 | + def __init__( | ||
| 94 | + self, | ||
| 95 | + engine_type: str, | ||
| 96 | + prefill: tuple[str, int], | ||
| 97 | + decode: tuple[str, int], | ||
| 98 | + bootstrap_port: int | None, | ||
| 99 | + ): | ||
| 100 | + self.p = _instance( | ||
| 101 | + 1, | ||
| 102 | + PDRole.ROLE_P, | ||
| 103 | + engine_type, | ||
| 104 | + *prefill, | ||
| 105 | + bootstrap_port=bootstrap_port, | ||
| 106 | + ) | ||
| 107 | + self.d = _instance(2, PDRole.ROLE_D, engine_type, *decode) | ||
| 108 | + | ||
| 109 | + async def select_and_allocate(self, role, _req_info, **_kwargs): | ||
| 110 | + instance = self.p if role == PDRole.ROLE_P else self.d | ||
| 111 | + endpoint = next(iter(next(iter(instance.endpoints.values())).values())) | ||
| 112 | + return instance, endpoint, Workload(active_tokens=1) | ||
| 113 | + | ||
| 114 | + async def update_workload(self, _params): | ||
| 115 | + return True | ||
| 116 | + | ||
| 117 | + async def report_cb_event(self, _instance_id: int, _event: str) -> None: | ||
| 118 | + return None | ||
| 119 | + | ||
| 120 | + async def get_unblocked_instances(self, role) -> list[int]: | ||
| 121 | + if role == PDRole.ROLE_P: | ||
| 122 | + return [self.p.id] | ||
| 123 | + if role == PDRole.ROLE_D: | ||
| 124 | + return [self.d.id] | ||
| 125 | + return [] | ||
| 126 | + | ||
| 127 | + | ||
| 128 | +def _tls_config(scheme: str) -> TLSConfig: | ||
| 129 | + if scheme == "http": | ||
| 130 | + return TLSConfig(enable_tls=False) | ||
| 131 | + return TLSConfig( | ||
| 132 | + enable_tls=True, | ||
| 133 | + ca_file=_required_env("MOTOR_NATIVE_SMOKE_CA_FILE"), | ||
| 134 | + cert_file=_required_env("MOTOR_NATIVE_SMOKE_CERT_FILE"), | ||
| 135 | + key_file=_required_env("MOTOR_NATIVE_SMOKE_KEY_FILE"), | ||
| 136 | + passwd_file=os.getenv("MOTOR_NATIVE_SMOKE_PASSWD_FILE", "").strip(), | ||
| 137 | + crl_file=os.getenv("MOTOR_NATIVE_SMOKE_CRL_FILE", "").strip(), | ||
| 138 | + ) | ||
| 139 | + | ||
| 140 | + | ||
| 141 | +async def _stream_body(response) -> bytes: | ||
| 142 | + chunks = [] | ||
| 143 | + async for chunk in response.body_iterator: | ||
| 144 | + chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode()) | ||
| 145 | + return b"".join(chunks) | ||
| 146 | + | ||
| 147 | + | ||
| 148 | + | ||
| 149 | + ("api", "request_payload", "stream"), | ||
| 150 | + [ | ||
| 151 | + ("v1/completions", {"prompt": "Reply with OK."}, False), | ||
| 152 | + ("v1/completions", {"prompt": "Reply with OK."}, True), | ||
| 153 | + ( | ||
| 154 | + "v1/chat/completions", | ||
| 155 | + {"messages": [{"role": "user", "content": "Reply with OK."}]}, | ||
| 156 | + False, | ||
| 157 | + ), | ||
| 158 | + ( | ||
| 159 | + "v1/chat/completions", | ||
| 160 | + {"messages": [{"role": "user", "content": "Reply with OK."}]}, | ||
| 161 | + True, | ||
| 162 | + ), | ||
| 163 | + ], | ||
| 164 | +) | ||
| 165 | + | ||
| 166 | +async def test_real_native_pd_openai_smoke(api, request_payload, stream): | ||
| 167 | + if os.getenv("MOTOR_RUN_NATIVE_PD_SMOKE", "").strip() != "1": | ||
| 168 | + pytest.skip("set MOTOR_RUN_NATIVE_PD_SMOKE=1 to run real native P/D smoke tests") | ||
| 169 | + | ||
| 170 | + engine_type = _required_env("MOTOR_NATIVE_SMOKE_ENGINE_TYPE").lower() | ||
| 171 | + if engine_type not in {"vllm", "sglang"}: | ||
| 172 | + pytest.fail("MOTOR_NATIVE_SMOKE_ENGINE_TYPE must be vllm or sglang") | ||
| 173 | + p_host, p_port, p_scheme = _parse_endpoint_url("MOTOR_NATIVE_SMOKE_PREFILL_URL") | ||
| 174 | + d_host, d_port, d_scheme = _parse_endpoint_url("MOTOR_NATIVE_SMOKE_DECODE_URL") | ||
| 175 | + if p_scheme != d_scheme: | ||
| 176 | + pytest.fail("Prefill and decode smoke endpoints must use the same HTTP scheme") | ||
| 177 | + bootstrap_port = None | ||
| 178 | + if engine_type == "sglang": | ||
| 179 | + bootstrap_port = int(_required_env("MOTOR_NATIVE_SMOKE_BOOTSTRAP_PORT")) | ||
| 180 | + if not 1 <= bootstrap_port <= 65535: | ||
| 181 | + pytest.fail("MOTOR_NATIVE_SMOKE_BOOTSTRAP_PORT must be in range 1..65535") | ||
| 182 | + | ||
| 183 | + scheduler = _RealNativeScheduler( | ||
| 184 | + engine_type, | ||
| 185 | + (p_host, p_port), | ||
| 186 | + (d_host, d_port), | ||
| 187 | + bootstrap_port, | ||
| 188 | + ) | ||
| 189 | + config = CoordinatorConfig() | ||
| 190 | + config.exception_config = ExceptionConfig(max_retry=1, retry_delay=0) | ||
| 191 | + config.infer_tls_config = _tls_config(p_scheme) | ||
| 192 | + req_data = { | ||
| 193 | + "model": _required_env("MOTOR_NATIVE_SMOKE_MODEL"), | ||
| 194 | + **request_payload, | ||
| 195 | + "max_tokens": 2, | ||
| 196 | + "stream": stream, | ||
| 197 | + } | ||
| 198 | + if stream: | ||
| 199 | + req_data["stream_options"] = {"include_usage": True} | ||
| 200 | + req_info = RequestInfo( | ||
| 201 | + req_id=f"native-smoke-{engine_type}-{api.replace('/', '-')}-{stream}", | ||
| 202 | + req_data=req_data, | ||
| 203 | + api=api, | ||
| 204 | + entry_api=api, | ||
| 205 | + req_len=8, | ||
| 206 | + ) | ||
| 207 | + router = UnifiedPDRouter( | ||
| 208 | + req_info, | ||
| 209 | + config, | ||
| 210 | + scheduler=scheduler, | ||
| 211 | + request_manager=RequestManager(config), | ||
| 212 | + ) | ||
| 213 | + | ||
| 214 | + try: | ||
| 215 | + response = await router.handle_request() | ||
| 216 | + if stream: | ||
| 217 | + body = await _stream_body(response) | ||
| 218 | + assert b"data:" in body | ||
| 219 | + assert b'"error"' not in body | ||
| 220 | + else: | ||
| 221 | + body = response.body | ||
| 222 | + payload = json.loads(body) | ||
| 223 | + assert payload.get("choices") | ||
| 224 | + assert b"kv_transfer_params" not in body | ||
| 225 | + assert b"bootstrap_host" not in body | ||
| 226 | + assert b"bootstrap_port" not in body | ||
| 227 | + assert b"bootstrap_room" not in body | ||
| 228 | + assert req_info.state == ReqState.DECODE_END | ||
| 229 | + finally: | ||
| 230 | + await HTTPClientPool().close_all() | ||
| @@ -39,6 +39,13 @@ from motor.common.resources.instance import ( # noqa: E402 | |||
| 39 | InsConditionEvent, | 39 | InsConditionEvent, |
| 40 | PDRole, | 40 | PDRole, |
| 41 | ) | 41 | ) |
| 42 | +from examples.deployer.prestop.prestop import ( # noqa: E402 | ||
| 43 | + build_curl_tls_args, | ||
| 44 | + get_engine_metrics, | ||
| 45 | + get_engine_type, | ||
| 46 | + get_infer_tls_config, | ||
| 47 | + wait_for_engine_drain, | ||
| 48 | +) | ||
| 42 | 49 | ||
| 43 | # --------------------------------------------------------------------------- | 50 | # --------------------------------------------------------------------------- |
| 44 | # Helpers | 51 | # Helpers |
| @@ -149,6 +156,81 @@ class TestHeartbeatManagerPrestop: | |||
| 149 | heartbeat_mgr.pause_all_endpoints() # no exception | 156 | heartbeat_mgr.pause_all_endpoints() # no exception |
| 150 | 157 | ||
| 151 | 158 | ||
| 159 | +def test_prestop_uses_inference_tls_for_native_metrics(): | ||
| 160 | + config = { | ||
| 161 | + "motor_deploy_config": { | ||
| 162 | + "tls_config": { | ||
| 163 | + "infer_tls_config": { | ||
| 164 | + "enable_tls": True, | ||
| 165 | + "ca_file": "/certs/ca.pem", | ||
| 166 | + "cert_file": "/certs/client.pem", | ||
| 167 | + "key_file": "/certs/client.key", | ||
| 168 | + "crl_file": "/certs/ca.crl", | ||
| 169 | + } | ||
| 170 | + } | ||
| 171 | + } | ||
| 172 | + } | ||
| 173 | + | ||
| 174 | + tls_config = get_infer_tls_config(config) | ||
| 175 | + | ||
| 176 | + assert build_curl_tls_args(tls_config) == [ | ||
| 177 | + "--cacert", | ||
| 178 | + "/certs/ca.pem", | ||
| 179 | + "--cert", | ||
| 180 | + "/certs/client.pem", | ||
| 181 | + "--key", | ||
| 182 | + "/certs/client.key", | ||
| 183 | + "--crlfile", | ||
| 184 | + "/certs/ca.crl", | ||
| 185 | + ] | ||
| 186 | + | ||
| 187 | + | ||
| 188 | + | ||
| 189 | + ("engine_type", "metrics_text"), | ||
| 190 | + [ | ||
| 191 | + ("vllm", "vllm:num_requests_waiting 2\nvllm:num_requests_running 3\n"), | ||
| 192 | + ("sglang", "sglang:num_waiting_reqs 4\nsglang:num_running_reqs 5\n"), | ||
| 193 | + ], | ||
| 194 | +) | ||
| 195 | +def test_prestop_parses_engine_specific_drain_metrics(engine_type, metrics_text): | ||
| 196 | + with patch("examples.deployer.prestop.prestop._http_get_text", return_value=metrics_text): | ||
| 197 | + assert get_engine_metrics("http://engine/metrics", {}, engine_type) == { | ||
| 198 | + "waiting": 2 if engine_type == "vllm" else 4, | ||
| 199 | + "running": 3 if engine_type == "vllm" else 5, | ||
| 200 | + } | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +def test_prestop_missing_required_metric_is_not_treated_as_drained(): | ||
| 204 | + with patch( | ||
| 205 | + "examples.deployer.prestop.prestop._http_get_text", | ||
| 206 | + return_value="sglang:num_waiting_reqs 0\n", | ||
| 207 | + ): | ||
| 208 | + assert get_engine_metrics("http://engine/metrics", {}, "sglang") is None | ||
| 209 | + | ||
| 210 | + | ||
| 211 | +def test_prestop_retries_when_metrics_are_temporarily_unavailable(): | ||
| 212 | + with ( | ||
| 213 | + patch( | ||
| 214 | + "examples.deployer.prestop.prestop.get_engine_metrics", | ||
| 215 | + side_effect=[None, {"waiting": 0, "running": 0}], | ||
| 216 | + ) as get_metrics, | ||
| 217 | + patch("examples.deployer.prestop.prestop.time.monotonic", side_effect=[0, 0, 1, 1]), | ||
| 218 | + patch("examples.deployer.prestop.prestop.time.sleep") as sleep, | ||
| 219 | + ): | ||
| 220 | + drained = wait_for_engine_drain(["http://engine/metrics"], {}, "sglang", 10, 3) | ||
| 221 | + | ||
| 222 | + assert drained is True | ||
| 223 | + assert get_metrics.call_count == 2 | ||
| 224 | + sleep.assert_called_once_with(3) | ||
| 225 | + | ||
| 226 | + | ||
| 227 | + | ||
| 228 | +def test_prestop_resolves_engine_type_from_active_role(): | ||
| 229 | + config = {"motor_engine_decode_config": {"engine_type": "sglang"}} | ||
| 230 | + | ||
| 231 | + assert get_engine_type(config) == "sglang" | ||
| 232 | + | ||
| 233 | + | ||
| 152 | # --------------------------------------------------------------------------- | 234 | # --------------------------------------------------------------------------- |
| 153 | # 2. Instance — is_all_endpoints_paused | 235 | # 2. Instance — is_all_endpoints_paused |
| 154 | # --------------------------------------------------------------------------- | 236 | # --------------------------------------------------------------------------- |
| @@ -70,7 +70,7 @@ def _install_engine_fakes(): | |||
| 70 | mock_endpoint_cfg_mod.EndpointConfig.init_endpoint_config.return_value = mock_ep | 70 | mock_endpoint_cfg_mod.EndpointConfig.init_endpoint_config.return_value = mock_ep |
| 71 | 71 | ||
| 72 | fakes = { | 72 | fakes = { |
| 73 | - "motor.engine_server.factory.config_factory": mock_cf_mod, | 73 | + "motor.node_manager.core.services.native_engine.config_factory": mock_cf_mod, |
| 74 | "motor.engine_server.factory.endpoint_factory": mock_ef_mod, | 74 | "motor.engine_server.factory.endpoint_factory": mock_ef_mod, |
| 75 | "motor.engine_server.core.infer_endpoint": mock_infer_mod, | 75 | "motor.engine_server.core.infer_endpoint": mock_infer_mod, |
| 76 | "motor.engine_server.core.mgmt_endpoint": mock_mgmt_mod, | 76 | "motor.engine_server.core.mgmt_endpoint": mock_mgmt_mod, |
| @@ -84,7 +84,7 @@ def _install_engine_fakes(): | |||
| 84 | def _uninstall_engine_fakes(): | 84 | def _uninstall_engine_fakes(): |
| 85 | """Remove engine-server module fakes that were added by _install_engine_fakes.""" | 85 | """Remove engine-server module fakes that were added by _install_engine_fakes.""" |
| 86 | for mod_name in ( | 86 | for mod_name in ( |
| 87 | - "motor.engine_server.factory.config_factory", | 87 | + "motor.node_manager.core.services.native_engine.config_factory", |
| 88 | "motor.engine_server.factory.endpoint_factory", | 88 | "motor.engine_server.factory.endpoint_factory", |
| 89 | "motor.engine_server.core.infer_endpoint", | 89 | "motor.engine_server.core.infer_endpoint", |
| 90 | "motor.engine_server.core.mgmt_endpoint", | 90 | "motor.engine_server.core.mgmt_endpoint", |
| @@ -179,7 +179,7 @@ class TestMainNativeLaunchRouting: | |||
| 179 | "motor.engine_server.core.infer_endpoint", | 179 | "motor.engine_server.core.infer_endpoint", |
| 180 | "motor.engine_server.core.mgmt_endpoint", | 180 | "motor.engine_server.core.mgmt_endpoint", |
| 181 | "motor.engine_server.factory.endpoint_factory", | 181 | "motor.engine_server.factory.endpoint_factory", |
| 182 | - "motor.engine_server.factory.config_factory", | 182 | + "motor.node_manager.core.services.native_engine.config_factory", |
| 183 | ): | 183 | ): |
| 184 | sys.modules[mod_name].reset_mock() | 184 | sys.modules[mod_name].reset_mock() |
| 185 | 185 | ||
| @@ -26,7 +26,7 @@ from motor.config.endpoint import ( | |||
| 26 | DECODE_PARALLEL_CONFIG_KEY, | 26 | DECODE_PARALLEL_CONFIG_KEY, |
| 27 | ENCODE_PARALLEL_CONFIG_KEY, | 27 | ENCODE_PARALLEL_CONFIG_KEY, |
| 28 | ) | 28 | ) |
| 29 | -from motor.engine_server.constants import constants | 29 | +from motor.common import engine_constants as constants |
| 30 | 30 | ||
| 31 | 31 | ||
| 32 | # --- ParallelConfig tests --- | 32 | # --- ParallelConfig tests --- |
| @@ -126,6 +126,7 @@ def test_health_check_config_defaults(): | |||
| 126 | config = HealthCheckConfig() | 126 | config = HealthCheckConfig() |
| 127 | assert config.health_collector_timeout == 5 | 127 | assert config.health_collector_timeout == 5 |
| 128 | assert config.health_collector_timeout_retry_attempts == 3 | 128 | assert config.health_collector_timeout_retry_attempts == 3 |
| 129 | + assert config.startup_timeout == 1800 | ||
| 129 | assert config.npu_usage_threshold == 3 | 130 | assert config.npu_usage_threshold == 3 |
| 130 | assert config.enable_virtual_inference is False | 131 | assert config.enable_virtual_inference is False |
| 131 | 132 | ||
| @@ -135,12 +136,14 @@ def test_health_check_config_from_dict(): | |||
| 135 | data = { | 136 | data = { |
| 136 | "health_collector_timeout": 10, | 137 | "health_collector_timeout": 10, |
| 137 | "health_collector_timeout_retry_attempts": 5, | 138 | "health_collector_timeout_retry_attempts": 5, |
| 139 | + "startup_timeout": 600, | ||
| 138 | "npu_usage_threshold": 20, | 140 | "npu_usage_threshold": 20, |
| 139 | "enable_virtual_inference": False, | 141 | "enable_virtual_inference": False, |
| 140 | } | 142 | } |
| 141 | config = HealthCheckConfig.from_dict(data) | 143 | config = HealthCheckConfig.from_dict(data) |
| 142 | assert config.health_collector_timeout == 10 | 144 | assert config.health_collector_timeout == 10 |
| 143 | assert config.health_collector_timeout_retry_attempts == 5 | 145 | assert config.health_collector_timeout_retry_attempts == 5 |
| 146 | + assert config.startup_timeout == 600 | ||
| 144 | assert config.npu_usage_threshold == 20 | 147 | assert config.npu_usage_threshold == 20 |
| 145 | assert config.enable_virtual_inference is False | 148 | assert config.enable_virtual_inference is False |
| 146 | 149 | ||
| @@ -667,13 +670,13 @@ def test_load_deploy_config_rejects_malformed_multi_connectors(valid_config_file | |||
| 667 | {constants.CONNECTORS: [{}, "not-a-dict"]}, | 670 | {constants.CONNECTORS: [{}, "not-a-dict"]}, |
| 668 | ) | 671 | ) |
| 669 | for bad_extra in bad_extras: | 672 | for bad_extra in bad_extras: |
| 670 | - with open(valid_config_file_for_endpoint) as f: | 673 | + with open(valid_config_file_for_endpoint, encoding="utf-8") as f: |
| 671 | data = json.load(f) | 674 | data = json.load(f) |
| 672 | kv = {constants.KV_CONNECTOR: constants.MULTI_CONNECTOR} | 675 | kv = {constants.KV_CONNECTOR: constants.MULTI_CONNECTOR} |
| 673 | if bad_extra is not None: | 676 | if bad_extra is not None: |
| 674 | kv[constants.KV_CONNECTOR_EXTRA_CONFIG] = bad_extra | 677 | kv[constants.KV_CONNECTOR_EXTRA_CONFIG] = bad_extra |
| 675 | data["engine_config"][constants.KV_TRANSFER_CONFIG] = kv | 678 | data["engine_config"][constants.KV_TRANSFER_CONFIG] = kv |
| 676 | - with open(valid_config_file_for_endpoint, "w") as f: | 679 | + with open(valid_config_file_for_endpoint, "w", encoding="utf-8") as f: |
| 677 | json.dump(data, f) | 680 | json.dump(data, f) |
| 678 | 681 | ||
| 679 | config = EndpointConfig( | 682 | config = EndpointConfig( |
| @@ -78,6 +78,7 @@ def _body(role="decode"): | |||
| 78 | "instance_id": 1, | 78 | "instance_id": 1, |
| 79 | "endpoint_id": 0, | 79 | "endpoint_id": 0, |
| 80 | "url": "http://127.0.0.1:8000", | 80 | "url": "http://127.0.0.1:8000", |
| 81 | + "bootstrap_port": 31000, | ||
| 81 | }, | 82 | }, |
| 82 | "decode": { | 83 | "decode": { |
| 83 | "instance_id": 2, | 84 | "instance_id": 2, |
| @@ -863,15 +864,14 @@ async def test_vllm_stream_normalization_preserves_done(): | |||
| 863 | 864 | ||
| 864 | 865 | ||
| 865 | 866 | ||
| 866 | -async def test_sglang_adapter_generates_stable_bootstrap(monkeypatch): | 867 | +async def test_sglang_adapter_generates_stable_bootstrap(): |
| 867 | - monkeypatch.setenv("DISAGGREGATION_BOOTSTRAP_PORT", "31000") | ||
| 868 | adapter = SGLangDispatchAdapter(_Config(engine_type="sglang", role="decode")) | 868 | adapter = SGLangDispatchAdapter(_Config(engine_type="sglang", role="decode")) |
| 869 | 869 | ||
| 870 | first, _ = await adapter.adapt_request_body(_body("decode")) | 870 | first, _ = await adapter.adapt_request_body(_body("decode")) |
| 871 | second, _ = await adapter.adapt_request_body(_body("decode")) | 871 | second, _ = await adapter.adapt_request_body(_body("decode")) |
| 872 | 872 | ||
| 873 | assert first["bootstrap_host"] == "127.0.0.1" | 873 | assert first["bootstrap_host"] == "127.0.0.1" |
| 874 | - assert first["bootstrap_port"] == "31000" | 874 | + assert first["bootstrap_port"] == 31000 |
| 875 | assert first["bootstrap_room"] == second["bootstrap_room"] | 875 | assert first["bootstrap_room"] == second["bootstrap_room"] |
| 876 | 876 | ||
| 877 | 877 | ||
| @@ -883,13 +883,14 @@ async def test_sglang_adapter_requires_bootstrap_port(monkeypatch): | |||
| 883 | peer_stops.append(dispatch) | 883 | peer_stops.append(dispatch) |
| 884 | return None | 884 | return None |
| 885 | 885 | ||
| 886 | - monkeypatch.delenv("DISAGGREGATION_BOOTSTRAP_PORT", raising=False) | ||
| 887 | monkeypatch.setattr(SGLangDispatchAdapter, "stop_peer", _stop_peer) | 886 | monkeypatch.setattr(SGLangDispatchAdapter, "stop_peer", _stop_peer) |
| 888 | adapter = SGLangDispatchAdapter(_Config(engine_type="sglang", role="decode")) | 887 | adapter = SGLangDispatchAdapter(_Config(engine_type="sglang", role="decode")) |
| 888 | + body = _body("decode") | ||
| 889 | + body[MOTOR_DISPATCH_KEY]["endpoints"]["prefill"].pop("bootstrap_port") | ||
| 889 | 890 | ||
| 890 | with pytest.raises(HTTPException) as exc_info: | 891 | with pytest.raises(HTTPException) as exc_info: |
| 891 | - await adapter.adapt_request_body(_body("decode")) | 892 | + await adapter.adapt_request_body(body) |
| 892 | 893 | ||
| 893 | assert exc_info.value.status_code == 500 | 894 | assert exc_info.value.status_code == 500 |
| 894 | - assert "DISAGGREGATION_BOOTSTRAP_PORT" in exc_info.value.detail | 895 | + assert "bootstrap_port" in exc_info.value.detail |
| 895 | assert peer_stops[0].engine_request_id == "req#a1" | 896 | assert peer_stops[0].engine_request_id == "req#a1" |
| @@ -13,9 +13,9 @@ from unittest.mock import MagicMock, patch | |||
| 13 | import pytest | 13 | import pytest |
| 14 | from fastapi.testclient import TestClient | 14 | from fastapi.testclient import TestClient |
| 15 | 15 | ||
| 16 | -from motor.engine_server.constants.constants import INIT_STATUS, STATUS_KEY | 16 | +from motor.common.engine_constants import INIT_STATUS, STATUS_KEY |
| 17 | from motor.engine_server.core.mgmt_endpoint import MgmtEndpoint | 17 | from motor.engine_server.core.mgmt_endpoint import MgmtEndpoint |
| 18 | -from tests.engine_server.core.test_vllm_config import _make_endpoint_config | 18 | +from tests.node_manager.core.services.native_engine.backends.vllm.test_config import _make_endpoint_config |
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | 21 | ||
| @@ -24,7 +24,7 @@ from motor.engine_server.core.sim_inference import ( | |||
| 24 | _VIRTUAL_WARMUP_TIMEOUT_SEC, | 24 | _VIRTUAL_WARMUP_TIMEOUT_SEC, |
| 25 | _is_virtual_metrics_request, | 25 | _is_virtual_metrics_request, |
| 26 | ) | 26 | ) |
| 27 | -from motor.engine_server.constants import constants | 27 | +from motor.common import engine_constants as constants |
| 28 | 28 | ||
| 29 | # pylint: disable=redefined-outer-name | 29 | # pylint: disable=redefined-outer-name |
| 30 | 30 | ||
| @@ -10,7 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | import pytest | 11 | import pytest |
| 12 | 12 | ||
| 13 | -from motor.engine_server.utils.ip import ip_valid_check, port_valid_check, is_valid_ipv6_address, build_endpoint | 13 | +from motor.common.utils.ip import build_endpoint, ip_valid_check, is_valid_ipv6_address, port_valid_check |
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | class TestIpUtils: | 16 | class TestIpUtils: |
| @@ -1,5 +1,3 @@ | |||
| 1 | -#!/usr/bin/env python3 | ||
| 2 | -# -*- coding: utf-8 -*- | ||
| 3 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | 1 | # Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. |
| 4 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 5 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| @@ -16,8 +14,16 @@ import tempfile | |||
| 16 | import pytest | 14 | import pytest |
| 17 | from pathlib import Path | 15 | from pathlib import Path |
| 18 | 16 | ||
| 19 | -from motor.engine_server.utils.validators import Validator, StringValidator, MapValidator, DirectoryValidator, \ | 17 | +from motor.common.utils.validators import ( |
| 20 | - RankSizeValidator, FileValidator, IntValidator, ClassValidator | 18 | + ClassValidator, |
| 19 | + DirectoryValidator, | ||
| 20 | + FileValidator, | ||
| 21 | + IntValidator, | ||
| 22 | + MapValidator, | ||
| 23 | + RankSizeValidator, | ||
| 24 | + StringValidator, | ||
| 25 | + Validator, | ||
| 26 | +) | ||
| 21 | 27 | ||
| 22 | BIN_PATH = "/usr/bin" | 28 | BIN_PATH = "/usr/bin" |
| 23 | DIRECTORY_BLACKLIST_PATH = "/abc/d/e" | 29 | DIRECTORY_BLACKLIST_PATH = "/abc/d/e" |
| @@ -59,8 +65,13 @@ def test_string_validator_can_be_transformed2int(): | |||
| 59 | 65 | ||
| 60 | 66 | ||
| 61 | def test_string_validator_contain_sensitive_words(): | 67 | def test_string_validator_contain_sensitive_words(): |
| 62 | - assert not StringValidator("passwordme").check_not_contain_black_element("pass") \ | 68 | + assert ( |
| 63 | - .check_string_length().check().is_valid() | 69 | + not StringValidator("passwordme") |
| 70 | + .check_not_contain_black_element("pass") | ||
| 71 | + .check_string_length() | ||
| 72 | + .check() | ||
| 73 | + .is_valid() | ||
| 74 | + ) | ||
| 64 | 75 | ||
| 65 | 76 | ||
| 66 | def test_map_validator_should_contain_inclusive_keys(): | 77 | def test_map_validator_should_contain_inclusive_keys(): |
| @@ -76,29 +87,37 @@ def test_directory_black_list(): | |||
| 76 | os.makedirs(test_path, exist_ok=True) | 87 | os.makedirs(test_path, exist_ok=True) |
| 77 | try: | 88 | try: |
| 78 | # Test exact match | 89 | # Test exact match |
| 79 | - assert not DirectoryValidator(test_path).with_blacklist( | 90 | + assert not DirectoryValidator(test_path).with_blacklist(lst=[test_path]).check().is_valid() |
| 80 | - lst=[test_path]).check().is_valid() | 91 | + assert DirectoryValidator(test_path).with_blacklist(lst=[""]).check().is_valid() |
| 81 | - assert DirectoryValidator(test_path).with_blacklist( | ||
| 82 | - lst=[""]).check().is_valid() | ||
| 83 | # Test parent path with exact_compare=True (should be valid) | 92 | # Test parent path with exact_compare=True (should be valid) |
| 84 | - assert DirectoryValidator(test_path).with_blacklist([temp_dir], exact_compare=True) \ | 93 | + assert DirectoryValidator(test_path).with_blacklist([temp_dir], exact_compare=True).check().is_valid() |
| 85 | - .check().is_valid() | ||
| 86 | # Test parent path with exact_compare=False (should be invalid as test_path is child of temp_dir) | 94 | # Test parent path with exact_compare=False (should be invalid as test_path is child of temp_dir) |
| 87 | - assert not DirectoryValidator(test_path) \ | 95 | + assert not DirectoryValidator(test_path).with_blacklist([temp_dir], exact_compare=False).check().is_valid() |
| 88 | - .with_blacklist([temp_dir], exact_compare=False).check().is_valid() | ||
| 89 | finally: | 96 | finally: |
| 90 | import shutil | 97 | import shutil |
| 98 | + | ||
| 91 | shutil.rmtree(temp_dir, ignore_errors=True) | 99 | shutil.rmtree(temp_dir, ignore_errors=True) |
| 92 | else: | 100 | else: |
| 93 | - assert not DirectoryValidator(DIRECTORY_BLACKLIST_PATH).with_blacklist( | 101 | + assert ( |
| 94 | - lst=[DIRECTORY_BLACKLIST_PATH]).check().is_valid() | 102 | + not DirectoryValidator(DIRECTORY_BLACKLIST_PATH) |
| 95 | - assert DirectoryValidator(DIRECTORY_BLACKLIST_PATH).with_blacklist( | 103 | + .with_blacklist(lst=[DIRECTORY_BLACKLIST_PATH]) |
| 96 | - lst=[""]).check().is_valid() | 104 | + .check() |
| 97 | - assert DirectoryValidator(DIRECTORY_BLACKLIST_PATH).with_blacklist(["/abc/d/"], exact_compare=True) \ | 105 | + .is_valid() |
| 98 | - .check().is_valid() | 106 | + ) |
| 107 | + assert DirectoryValidator(DIRECTORY_BLACKLIST_PATH).with_blacklist(lst=[""]).check().is_valid() | ||
| 108 | + assert ( | ||
| 109 | + DirectoryValidator(DIRECTORY_BLACKLIST_PATH) | ||
| 110 | + .with_blacklist(["/abc/d/"], exact_compare=True) | ||
| 111 | + .check() | ||
| 112 | + .is_valid() | ||
| 113 | + ) | ||
| 99 | # if not exact compare, the /abc/d/e is chirldren path of /abc/d/, so it is invalid | 114 | # if not exact compare, the /abc/d/e is chirldren path of /abc/d/, so it is invalid |
| 100 | - assert not DirectoryValidator(DIRECTORY_BLACKLIST_PATH) \ | 115 | + assert ( |
| 101 | - .with_blacklist(["/abc/d/"], exact_compare=False).check().is_valid() | 116 | + not DirectoryValidator(DIRECTORY_BLACKLIST_PATH) |
| 117 | + .with_blacklist(["/abc/d/"], exact_compare=False) | ||
| 118 | + .check() | ||
| 119 | + .is_valid() | ||
| 120 | + ) | ||
| 102 | assert DirectoryValidator("/usr/bin/bash").with_blacklist().check().is_valid() | 121 | assert DirectoryValidator("/usr/bin/bash").with_blacklist().check().is_valid() |
| 103 | assert not DirectoryValidator("/usr/bin/bash").with_blacklist(exact_compare=False).check().is_valid() | 122 | assert not DirectoryValidator("/usr/bin/bash").with_blacklist(exact_compare=False).check().is_valid() |
| 104 | 123 | ||
| @@ -121,21 +140,19 @@ def test_directory_soft_link(): | |||
| 121 | # Skip on Windows as creating symlinks requires admin privileges or developer mode | 140 | # Skip on Windows as creating symlinks requires admin privileges or developer mode |
| 122 | if IS_WINDOWS: | 141 | if IS_WINDOWS: |
| 123 | pytest.skip("Symlink creation requires admin privileges on Windows") | 142 | pytest.skip("Symlink creation requires admin privileges on Windows") |
| 124 | - | 143 | + |
| 125 | - tmp = tempfile.NamedTemporaryFile(delete=True) | ||
| 126 | temp_dir = tempfile.mkdtemp() | 144 | temp_dir = tempfile.mkdtemp() |
| 127 | path = os.path.join(temp_dir, "link.ink") | 145 | path = os.path.join(temp_dir, "link.ink") |
| 128 | - # make a soft link | 146 | + with tempfile.NamedTemporaryFile(delete=True) as tmp: |
| 129 | - os.symlink(tmp.name, path) | 147 | + # make a soft link |
| 130 | - | 148 | + os.symlink(tmp.name, path) |
| 131 | - try: | 149 | + try: |
| 132 | - # do stuff with temp | 150 | + # do stuff with temp |
| 133 | - tmp.write(b"stuff") | 151 | + tmp.write(b"stuff") |
| 134 | - assert not DirectoryValidator(path).check_not_soft_link().check().is_valid() | 152 | + assert not DirectoryValidator(path).check_not_soft_link().check().is_valid() |
| 135 | - finally: | 153 | + finally: |
| 136 | - tmp.close() | 154 | + os.remove(path) |
| 137 | - os.remove(path) | 155 | + os.removedirs(temp_dir) |
| 138 | - os.removedirs(temp_dir) | ||
| 139 | 156 | ||
| 140 | 157 | ||
| 141 | def test_directory_check(): | 158 | def test_directory_check(): |
| @@ -143,25 +160,36 @@ def test_directory_check(): | |||
| 143 | assert not DirectoryValidator("").check_is_not_none().check_dir_name().check().is_valid() | 160 | assert not DirectoryValidator("").check_is_not_none().check_dir_name().check().is_valid() |
| 144 | assert not DirectoryValidator(None).check_is_not_none().check_dir_name().check().is_valid() | 161 | assert not DirectoryValidator(None).check_is_not_none().check_dir_name().check().is_valid() |
| 145 | assert DirectoryValidator("a/bc/d").check_is_not_none().check_dir_name().check().is_valid() | 162 | assert DirectoryValidator("a/bc/d").check_is_not_none().check_dir_name().check().is_valid() |
| 146 | - assert DirectoryValidator("/user/restore/fault/config", max_len=255). \ | 163 | + assert ( |
| 147 | - check_is_not_none().check_dir_name(). \ | 164 | + DirectoryValidator("/user/restore/fault/config", max_len=255) |
| 148 | - path_should_exist(is_file=True, msg="can not find the fault ranks config file") \ | 165 | + .check_is_not_none() |
| 149 | - .should_not_contains_sensitive_words().with_blacklist().check() | 166 | + .check_dir_name() |
| 150 | - assert DirectoryValidator(os.path.dirname(__file__), max_len=255). \ | 167 | + .path_should_exist(is_file=True, msg="can not find the fault ranks config file") |
| 151 | - check_is_not_none().check_dir_name().check_dir_file_number() \ | 168 | + .should_not_contains_sensitive_words() |
| 152 | - .path_should_exist(is_file=False, msg="can not find the fault ranks config file") \ | 169 | + .with_blacklist() |
| 153 | - .should_not_contains_sensitive_words().with_blacklist().check() | 170 | + .check() |
| 171 | + ) | ||
| 172 | + assert ( | ||
| 173 | + DirectoryValidator(os.path.dirname(__file__), max_len=255) | ||
| 174 | + .check_is_not_none() | ||
| 175 | + .check_dir_name() | ||
| 176 | + .check_dir_file_number() | ||
| 177 | + .path_should_exist(is_file=False, msg="can not find the fault ranks config file") | ||
| 178 | + .should_not_contains_sensitive_words() | ||
| 179 | + .with_blacklist() | ||
| 180 | + .check() | ||
| 181 | + ) | ||
| 154 | assert not DirectoryValidator(os.path.dirname(__file__)).path_should_not_exist().check().is_valid() | 182 | assert not DirectoryValidator(os.path.dirname(__file__)).path_should_not_exist().check().is_valid() |
| 155 | 183 | ||
| 156 | 184 | ||
| 157 | def test_check_directory_permissions(): | 185 | def test_check_directory_permissions(): |
| 158 | - temp_dir = tempfile.TemporaryDirectory() | 186 | + with tempfile.TemporaryDirectory() as temp_dir: |
| 159 | - temp_path = Path(temp_dir.name) | 187 | + temp_path = Path(temp_dir) |
| 160 | - test_dir = temp_path / "test_dir" | 188 | + test_dir = temp_path / "test_dir" |
| 161 | - test_dir.mkdir() | 189 | + test_dir.mkdir() |
| 162 | - os.chmod(test_dir, 0o777) | 190 | + os.chmod(test_dir, 0o777) |
| 163 | - target_mode = 0o750 | 191 | + target_mode = 0o750 |
| 164 | - assert not DirectoryValidator(test_dir).check_directory_permissions(target_mode).check().is_valid() | 192 | + assert not DirectoryValidator(test_dir).check_directory_permissions(target_mode).check().is_valid() |
| 165 | 193 | ||
| 166 | 194 | ||
| 167 | def test_rank_size_check(): | 195 | def test_rank_size_check(): |
| @@ -0,0 +1,9 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| @@ -0,0 +1,9 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| @@ -0,0 +1,9 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
Rtests/engine_server/core/test_sglang_config.py→tests/node_manager/core/services/native_engine/backends/sglang/test_config.py+1-1
| @@ -10,7 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | 11 | ||
| 12 | from motor.config.endpoint import DeployConfig, EndpointConfig, EngineConfig, ModelConfig, ParallelConfig | 12 | from motor.config.endpoint import DeployConfig, EndpointConfig, EngineConfig, ModelConfig, ParallelConfig |
| 13 | -from motor.engine_server.core.sglang.sglang_config import SGLangConfig | 13 | +from motor.node_manager.core.services.native_engine.backends.sglang.config import SGLangConfig |
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | def _make_endpoint_config(engine_cfg: dict) -> EndpointConfig: | 16 | def _make_endpoint_config(engine_cfg: dict) -> EndpointConfig: |
| @@ -0,0 +1,475 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +"""Golden contracts for native-engine CLI configuration. | ||
| 12 | + | ||
| 13 | +These tests exercise the relocated runtime config builders and preserve the | ||
| 14 | +pre-move CLI contract. | ||
| 15 | +""" | ||
| 16 | + | ||
| 17 | +import copy | ||
| 18 | +import json | ||
| 19 | +import sys | ||
| 20 | + | ||
| 21 | +import pytest | ||
| 22 | + | ||
| 23 | +from motor.config.endpoint import DeployConfig, EndpointConfig, EngineConfig, ModelConfig, ParallelConfig | ||
| 24 | +from motor.common import engine_constants as constants | ||
| 25 | +from motor.node_manager.core.services.native_engine.config_factory import ConfigFactory | ||
| 26 | +from motor.node_manager.core.services.native_engine.backends.sglang.config import SGLangConfig | ||
| 27 | +from motor.node_manager.core.services.native_engine.backends.vllm.config import VLLMConfig | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +def _endpoint_config( | ||
| 31 | + *, | ||
| 32 | + engine_type: str, | ||
| 33 | + role: str = "union", | ||
| 34 | + engine_config: dict | None = None, | ||
| 35 | + host: str = "127.0.0.1", | ||
| 36 | + port: int = 8000, | ||
| 37 | + master_dp_ip: str = "10.0.0.1", | ||
| 38 | + dp_rank: int = 0, | ||
| 39 | + node_rank: int = 0, | ||
| 40 | + dp_size: int = 1, | ||
| 41 | + tp_size: int = 2, | ||
| 42 | + pp_size: int = 1, | ||
| 43 | + pcp_size: int = 1, | ||
| 44 | + dp_rpc_port: int = 9000, | ||
| 45 | + d2d_peer_ips: str | None = None, | ||
| 46 | +) -> EndpointConfig: | ||
| 47 | + prefill_parallel = ParallelConfig( | ||
| 48 | + dp_size=dp_size, | ||
| 49 | + tp_size=tp_size, | ||
| 50 | + pp_size=pp_size, | ||
| 51 | + pcp_size=pcp_size, | ||
| 52 | + dp_rpc_port=dp_rpc_port, | ||
| 53 | + ) | ||
| 54 | + decode_parallel = ParallelConfig( | ||
| 55 | + dp_size=dp_size, | ||
| 56 | + tp_size=tp_size, | ||
| 57 | + pp_size=pp_size, | ||
| 58 | + dp_rpc_port=dp_rpc_port, | ||
| 59 | + ) | ||
| 60 | + deploy_config = DeployConfig( | ||
| 61 | + engine_type=engine_type, | ||
| 62 | + model_config=ModelConfig( | ||
| 63 | + model_name="glm-test", | ||
| 64 | + model_path="/models/glm-test", | ||
| 65 | + npu_mem_utils=0.85, | ||
| 66 | + encode_parallel_config=ParallelConfig( | ||
| 67 | + dp_size=dp_size, | ||
| 68 | + tp_size=tp_size, | ||
| 69 | + pp_size=pp_size, | ||
| 70 | + dp_rpc_port=dp_rpc_port, | ||
| 71 | + ), | ||
| 72 | + prefill_parallel_config=prefill_parallel, | ||
| 73 | + decode_parallel_config=decode_parallel, | ||
| 74 | + ), | ||
| 75 | + engine_config=EngineConfig.from_dict(engine_config or {}), | ||
| 76 | + mgmt_tls_config=None, | ||
| 77 | + infer_tls_config=None, | ||
| 78 | + ) | ||
| 79 | + return EndpointConfig( | ||
| 80 | + engine_type=engine_type, | ||
| 81 | + role=role, | ||
| 82 | + host=host, | ||
| 83 | + port=port, | ||
| 84 | + mgmt_port=9001, | ||
| 85 | + master_dp_ip=master_dp_ip, | ||
| 86 | + dp_rank=dp_rank, | ||
| 87 | + node_rank=node_rank, | ||
| 88 | + d2d_peer_ips=d2d_peer_ips, | ||
| 89 | + deploy_config=deploy_config, | ||
| 90 | + ) | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +def test_vllm_union_cli_golden_preserves_order_types_and_precedence(): | ||
| 94 | + endpoint = _endpoint_config( | ||
| 95 | + engine_type="vllm", | ||
| 96 | + engine_config={ | ||
| 97 | + "dtype": "bfloat16", | ||
| 98 | + "trust_remote_code": True, | ||
| 99 | + "allowed_local_media_path": ["/data/a", "/data/b"], | ||
| 100 | + "rope_scaling": {"rope_type": "yarn", "factor": 2.0}, | ||
| 101 | + "tensor_parallel_size": 8, | ||
| 102 | + "enable_chunked_prefill": False, | ||
| 103 | + }, | ||
| 104 | + ) | ||
| 105 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 106 | + config.initialize() | ||
| 107 | + | ||
| 108 | + assert config.get_cli_args() == [ | ||
| 109 | + "--dtype", | ||
| 110 | + "bfloat16", | ||
| 111 | + "--trust-remote-code", | ||
| 112 | + "--allowed-local-media-path", | ||
| 113 | + "/data/a", | ||
| 114 | + "/data/b", | ||
| 115 | + "--rope-scaling", | ||
| 116 | + json.dumps({"rope_type": "yarn", "factor": 2.0}), | ||
| 117 | + "--tensor-parallel-size", | ||
| 118 | + "8", | ||
| 119 | + "--model", | ||
| 120 | + "/models/glm-test", | ||
| 121 | + "--served-model-name", | ||
| 122 | + "glm-test", | ||
| 123 | + "--gpu-memory-utilization", | ||
| 124 | + "0.85", | ||
| 125 | + "--data-parallel-size", | ||
| 126 | + "1", | ||
| 127 | + "--pipeline-parallel-size", | ||
| 128 | + "1", | ||
| 129 | + "--data-parallel-rpc-port", | ||
| 130 | + "9000", | ||
| 131 | + "--cp-kv-cache-interleave-size", | ||
| 132 | + "1", | ||
| 133 | + "--host", | ||
| 134 | + "127.0.0.1", | ||
| 135 | + "--port", | ||
| 136 | + "8000", | ||
| 137 | + ] | ||
| 138 | + | ||
| 139 | + | ||
| 140 | + | ||
| 141 | + ("role", "kv_role"), | ||
| 142 | + [ | ||
| 143 | + ("prefill", "kv_producer"), | ||
| 144 | + ("decode", "kv_consumer"), | ||
| 145 | + ], | ||
| 146 | +) | ||
| 147 | +def test_vllm_pd_cli_golden_injects_handoff_metadata(role, kv_role): | ||
| 148 | + endpoint = _endpoint_config( | ||
| 149 | + engine_type="vllm", | ||
| 150 | + role=role, | ||
| 151 | + engine_config={ | ||
| 152 | + "kv_transfer_config": { | ||
| 153 | + "kv_connector": "MooncakeConnectorV1", | ||
| 154 | + "kv_port": "36001", | ||
| 155 | + } | ||
| 156 | + }, | ||
| 157 | + dp_size=2, | ||
| 158 | + dp_rank=1, | ||
| 159 | + ) | ||
| 160 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 161 | + config.initialize() | ||
| 162 | + expected_kv = { | ||
| 163 | + "kv_connector": "MooncakeConnectorV1", | ||
| 164 | + "kv_port": "36001", | ||
| 165 | + "kv_role": kv_role, | ||
| 166 | + "engine_id": "0", | ||
| 167 | + "kv_connector_extra_config": { | ||
| 168 | + "prefill": {"dp_size": 2, "tp_size": 2, "pp_size": 1}, | ||
| 169 | + "decode": {"dp_size": 2, "tp_size": 2, "pp_size": 1}, | ||
| 170 | + }, | ||
| 171 | + } | ||
| 172 | + | ||
| 173 | + assert config.get_cli_args() == [ | ||
| 174 | + "--kv-transfer-config", | ||
| 175 | + json.dumps(expected_kv), | ||
| 176 | + "--model", | ||
| 177 | + "/models/glm-test", | ||
| 178 | + "--served-model-name", | ||
| 179 | + "glm-test", | ||
| 180 | + "--gpu-memory-utilization", | ||
| 181 | + "0.85", | ||
| 182 | + "--data-parallel-size", | ||
| 183 | + "2", | ||
| 184 | + "--tensor-parallel-size", | ||
| 185 | + "2", | ||
| 186 | + "--pipeline-parallel-size", | ||
| 187 | + "1", | ||
| 188 | + "--data-parallel-rpc-port", | ||
| 189 | + "9000", | ||
| 190 | + "--cp-kv-cache-interleave-size", | ||
| 191 | + "1", | ||
| 192 | + "--host", | ||
| 193 | + "127.0.0.1", | ||
| 194 | + "--port", | ||
| 195 | + "8000", | ||
| 196 | + "--data-parallel-address", | ||
| 197 | + "10.0.0.1", | ||
| 198 | + "--data-parallel-rank", | ||
| 199 | + "1", | ||
| 200 | + ] | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +def test_vllm_cross_node_pcp_headless_cli_golden(): | ||
| 204 | + endpoint = _endpoint_config( | ||
| 205 | + engine_type="vllm", | ||
| 206 | + engine_config={"nnodes": 2, "master_port": 7001}, | ||
| 207 | + node_rank=1, | ||
| 208 | + pcp_size=2, | ||
| 209 | + ) | ||
| 210 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 211 | + config.initialize() | ||
| 212 | + | ||
| 213 | + assert config.get_cli_args()[-11:] == [ | ||
| 214 | + "--prefill-context-parallel-size", | ||
| 215 | + "2", | ||
| 216 | + "--host", | ||
| 217 | + "127.0.0.1", | ||
| 218 | + "--port", | ||
| 219 | + "8000", | ||
| 220 | + "--node-rank", | ||
| 221 | + "1", | ||
| 222 | + "--master-addr", | ||
| 223 | + "10.0.0.1", | ||
| 224 | + "--headless", | ||
| 225 | + ] | ||
| 226 | + | ||
| 227 | + | ||
| 228 | +def test_vllm_multi_connector_and_d2d_cli_contract(): | ||
| 229 | + endpoint = _endpoint_config( | ||
| 230 | + engine_type="vllm", | ||
| 231 | + role="prefill", | ||
| 232 | + engine_config={ | ||
| 233 | + "kv_transfer_config": { | ||
| 234 | + "kv_connector": "MultiConnector", | ||
| 235 | + "kv_connector_extra_config": { | ||
| 236 | + "connectors": [ | ||
| 237 | + {"kv_connector": "NixlConnector"}, | ||
| 238 | + { | ||
| 239 | + "kv_connector": "AscendStoreConnector", | ||
| 240 | + "kv_connector_extra_config": {}, | ||
| 241 | + }, | ||
| 242 | + ] | ||
| 243 | + }, | ||
| 244 | + }, | ||
| 245 | + "model_loader_extra_config": { | ||
| 246 | + "source": "auto", | ||
| 247 | + "listen_port": 5000, | ||
| 248 | + }, | ||
| 249 | + }, | ||
| 250 | + d2d_peer_ips="2001:db8::1,10.0.0.2", | ||
| 251 | + ) | ||
| 252 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 253 | + config.initialize() | ||
| 254 | + cli_args = config.get_cli_args() | ||
| 255 | + | ||
| 256 | + kv_index = cli_args.index("--kv-transfer-config") | ||
| 257 | + kv_config = json.loads(cli_args[kv_index + 1]) | ||
| 258 | + assert kv_config == { | ||
| 259 | + "kv_connector": "MultiConnector", | ||
| 260 | + "kv_connector_extra_config": { | ||
| 261 | + "connectors": [ | ||
| 262 | + { | ||
| 263 | + "kv_connector": "NixlConnector", | ||
| 264 | + "kv_role": "kv_producer", | ||
| 265 | + "kv_connector_extra_config": { | ||
| 266 | + "prefill": {"dp_size": 1, "tp_size": 2, "pp_size": 1}, | ||
| 267 | + "decode": {"dp_size": 1, "tp_size": 2, "pp_size": 1}, | ||
| 268 | + }, | ||
| 269 | + }, | ||
| 270 | + { | ||
| 271 | + "kv_connector": "AscendStoreConnector", | ||
| 272 | + "kv_connector_extra_config": {"lookup_rpc_port": "0"}, | ||
| 273 | + "kv_role": "kv_producer", | ||
| 274 | + }, | ||
| 275 | + ] | ||
| 276 | + }, | ||
| 277 | + "engine_id": "0", | ||
| 278 | + "kv_role": "kv_producer", | ||
| 279 | + } | ||
| 280 | + loader_index = cli_args.index("--model-loader-extra-config") | ||
| 281 | + assert json.loads(cli_args[loader_index + 1]) == { | ||
| 282 | + "LISTEN_PORT": 5000, | ||
| 283 | + "SOURCE": [ | ||
| 284 | + { | ||
| 285 | + "device_id": 0, | ||
| 286 | + "sources": ["[2001:db8::1]:5000", "10.0.0.2:5000"], | ||
| 287 | + }, | ||
| 288 | + { | ||
| 289 | + "device_id": 1, | ||
| 290 | + "sources": ["[2001:db8::1]:5001", "10.0.0.2:5001"], | ||
| 291 | + }, | ||
| 292 | + ], | ||
| 293 | + "MODEL": "glm-test", | ||
| 294 | + } | ||
| 295 | + assert cli_args[cli_args.index("--load-format") + 1] == "netloader" | ||
| 296 | + | ||
| 297 | + | ||
| 298 | +def test_vllm_builder_output_is_stable_across_repeated_initialization(): | ||
| 299 | + endpoint = _endpoint_config( | ||
| 300 | + engine_type="vllm", | ||
| 301 | + role="prefill", | ||
| 302 | + engine_config={ | ||
| 303 | + "kv_transfer_config": { | ||
| 304 | + "kv_connector": "MooncakeHybridConnector", | ||
| 305 | + "kv_port": "36001", | ||
| 306 | + "kv_connector_extra_config": {}, | ||
| 307 | + } | ||
| 308 | + }, | ||
| 309 | + ) | ||
| 310 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 311 | + | ||
| 312 | + config.initialize() | ||
| 313 | + first_cli = config.get_cli_args() | ||
| 314 | + first_endpoint = copy.deepcopy(endpoint) | ||
| 315 | + config.initialize() | ||
| 316 | + | ||
| 317 | + assert config.get_cli_args() == first_cli | ||
| 318 | + assert endpoint == first_endpoint | ||
| 319 | + | ||
| 320 | + | ||
| 321 | +def test_vllm_convert_preserves_process_argv(monkeypatch): | ||
| 322 | + endpoint = _endpoint_config(engine_type="vllm") | ||
| 323 | + config = VLLMConfig(endpoint_config=endpoint) | ||
| 324 | + config.initialize() | ||
| 325 | + monkeypatch.setattr(sys, "argv", ["contract-test", "--original"]) | ||
| 326 | + | ||
| 327 | + config.convert() | ||
| 328 | + | ||
| 329 | + assert sys.argv == ["contract-test", "--original"] | ||
| 330 | + | ||
| 331 | + | ||
| 332 | +def test_sglang_union_cli_golden(): | ||
| 333 | + endpoint = _endpoint_config( | ||
| 334 | + engine_type="sglang", | ||
| 335 | + engine_config={ | ||
| 336 | + "dtype": "bfloat16", | ||
| 337 | + "trust_remote_code": True, | ||
| 338 | + "random_seed": 7, | ||
| 339 | + }, | ||
| 340 | + ) | ||
| 341 | + config = SGLangConfig(endpoint_config=endpoint) | ||
| 342 | + config.initialize() | ||
| 343 | + | ||
| 344 | + assert config.get_cli_args() == [ | ||
| 345 | + "--dtype", | ||
| 346 | + "bfloat16", | ||
| 347 | + "--trust-remote-code", | ||
| 348 | + "--random-seed", | ||
| 349 | + "7", | ||
| 350 | + "--enable-metrics", | ||
| 351 | + "--host", | ||
| 352 | + "127.0.0.1", | ||
| 353 | + "--port", | ||
| 354 | + "8000", | ||
| 355 | + "--disaggregation-mode", | ||
| 356 | + "null", | ||
| 357 | + ] | ||
| 358 | + | ||
| 359 | + | ||
| 360 | + | ||
| 361 | + ("role", "mode"), | ||
| 362 | + [ | ||
| 363 | + ("prefill", "prefill"), | ||
| 364 | + ("decode", "decode"), | ||
| 365 | + ], | ||
| 366 | +) | ||
| 367 | +def test_sglang_pd_multinode_cli_golden(role, mode): | ||
| 368 | + endpoint = _endpoint_config( | ||
| 369 | + engine_type="sglang", | ||
| 370 | + role=role, | ||
| 371 | + engine_config={"nnodes": 2}, | ||
| 372 | + host="::1", | ||
| 373 | + master_dp_ip="2001:db8::10", | ||
| 374 | + dp_rank=1, | ||
| 375 | + node_rank=7, | ||
| 376 | + dp_rpc_port=9100, | ||
| 377 | + ) | ||
| 378 | + config = SGLangConfig(endpoint_config=endpoint) | ||
| 379 | + | ||
| 380 | + assert config.get_cli_args() == [ | ||
| 381 | + "--nnodes", | ||
| 382 | + "2", | ||
| 383 | + "--enable-metrics", | ||
| 384 | + "--host", | ||
| 385 | + "::1", | ||
| 386 | + "--port", | ||
| 387 | + "8000", | ||
| 388 | + "--dist-init-addr", | ||
| 389 | + "[2001:db8::10]:9100", | ||
| 390 | + "--node-rank", | ||
| 391 | + "7", | ||
| 392 | + "--disaggregation-mode", | ||
| 393 | + mode, | ||
| 394 | + ] | ||
| 395 | + | ||
| 396 | + | ||
| 397 | +def test_sglang_multinode_accepts_string_nnodes_and_requires_master_address(): | ||
| 398 | + endpoint = _endpoint_config( | ||
| 399 | + engine_type="sglang", | ||
| 400 | + role="prefill", | ||
| 401 | + engine_config={"nnodes": "2"}, | ||
| 402 | + master_dp_ip="", | ||
| 403 | + ) | ||
| 404 | + | ||
| 405 | + with pytest.raises(ValueError, match="master_dp_ip is required"): | ||
| 406 | + SGLangConfig(endpoint_config=endpoint).get_cli_args() | ||
| 407 | + | ||
| 408 | + | ||
| 409 | + | ||
| 410 | +def test_sglang_rejects_invalid_nnodes(nnodes): | ||
| 411 | + endpoint = _endpoint_config(engine_type="sglang", engine_config={"nnodes": nnodes}) | ||
| 412 | + | ||
| 413 | + with pytest.raises(ValueError, match="nnodes"): | ||
| 414 | + SGLangConfig(endpoint_config=endpoint).get_cli_args() | ||
| 415 | + | ||
| 416 | + | ||
| 417 | +def test_native_cli_omits_none_values(): | ||
| 418 | + endpoint = _endpoint_config( | ||
| 419 | + engine_type="sglang", | ||
| 420 | + engine_config={"download-dir": None}, | ||
| 421 | + ) | ||
| 422 | + | ||
| 423 | + args = SGLangConfig(endpoint_config=endpoint).get_cli_args() | ||
| 424 | + | ||
| 425 | + assert "--download-dir" not in args | ||
| 426 | + assert "None" not in args | ||
| 427 | + | ||
| 428 | + | ||
| 429 | +def test_factory_rejects_unknown_engine_type_before_importing_builder(): | ||
| 430 | + endpoint = _endpoint_config(engine_type="unknown") | ||
| 431 | + | ||
| 432 | + with pytest.raises(ValueError, match="Unsupported engine type: unknown"): | ||
| 433 | + ConfigFactory(endpoint).parse() | ||
| 434 | + | ||
| 435 | + | ||
| 436 | + | ||
| 437 | + ("invalid_kv_config", "error_pattern"), | ||
| 438 | + [ | ||
| 439 | + (None, "kv_transfer_config is None in engine_config"), | ||
| 440 | + ( | ||
| 441 | + { | ||
| 442 | + "kv_connector": "MultiConnector", | ||
| 443 | + "kv_connector_extra_config": {"connectors": [{"kv_connector": "NixlConnector"}]}, | ||
| 444 | + }, | ||
| 445 | + "Failed to process kv_transfer_config", | ||
| 446 | + ), | ||
| 447 | + ( | ||
| 448 | + { | ||
| 449 | + "kv_connector": "MultiConnector", | ||
| 450 | + "kv_connector_extra_config": { | ||
| 451 | + "connectors": [ | ||
| 452 | + {"kv_connector": "NixlConnector"}, | ||
| 453 | + { | ||
| 454 | + "kv_connector": "UnsupportedStoreConnector", | ||
| 455 | + "kv_connector_extra_config": {}, | ||
| 456 | + }, | ||
| 457 | + ] | ||
| 458 | + }, | ||
| 459 | + }, | ||
| 460 | + "Failed to process kv_transfer_config", | ||
| 461 | + ), | ||
| 462 | + ], | ||
| 463 | +) | ||
| 464 | +def test_vllm_invalid_pd_connector_contract(invalid_kv_config, error_pattern): | ||
| 465 | + engine_config = {} | ||
| 466 | + if invalid_kv_config is not None: | ||
| 467 | + engine_config[constants.KV_TRANSFER_CONFIG] = invalid_kv_config | ||
| 468 | + endpoint = _endpoint_config( | ||
| 469 | + engine_type="vllm", | ||
| 470 | + role="prefill", | ||
| 471 | + engine_config=engine_config, | ||
| 472 | + ) | ||
| 473 | + | ||
| 474 | + with pytest.raises(ValueError, match=error_pattern): | ||
| 475 | + VLLMConfig(endpoint_config=endpoint).initialize() | ||
Rtests/engine_server/core/test_ipv6_address_formatting.py→tests/node_manager/core/services/native_engine/backends/test_ipv6_address_formatting.py+3-3
| @@ -10,9 +10,9 @@ | |||
| 10 | 10 | ||
| 11 | from types import SimpleNamespace | 11 | from types import SimpleNamespace |
| 12 | 12 | ||
| 13 | -from motor.engine_server.constants import constants | 13 | +from motor.common import engine_constants as constants |
| 14 | -from motor.engine_server.core.sglang.sglang_config import SGLangConfig | 14 | +from motor.node_manager.core.services.native_engine.backends.sglang.config import SGLangConfig |
| 15 | -from motor.engine_server.core.vllm.vllm_config import VLLMConfig | 15 | +from motor.node_manager.core.services.native_engine.backends.vllm.config import VLLMConfig |
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | class _DeployConfig: | 18 | class _DeployConfig: |
| @@ -0,0 +1,9 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
Rtests/engine_server/core/test_vllm_config.py→tests/node_manager/core/services/native_engine/backends/vllm/test_config.py+1-1
| @@ -14,7 +14,7 @@ import json | |||
| 14 | import pytest | 14 | import pytest |
| 15 | 15 | ||
| 16 | from motor.config.endpoint import DeployConfig, EndpointConfig, EngineConfig, ModelConfig, ParallelConfig | 16 | from motor.config.endpoint import DeployConfig, EndpointConfig, EngineConfig, ModelConfig, ParallelConfig |
| 17 | -from motor.engine_server.core.vllm.vllm_config import VLLMConfig | 17 | +from motor.node_manager.core.services.native_engine.backends.vllm.config import VLLMConfig |
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | def _make_endpoint_config( | 20 | def _make_endpoint_config( |
| @@ -0,0 +1,295 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from types import SimpleNamespace | ||
| 12 | +from unittest.mock import MagicMock, patch | ||
| 13 | + | ||
| 14 | +import pytest | ||
| 15 | + | ||
| 16 | +from motor.common.resources.instance import PDRole | ||
| 17 | +from motor.config.endpoint import EndpointConfig | ||
| 18 | +from motor.config.tls_config import TLSConfig | ||
| 19 | +from motor.node_manager.core.services.native_engine.backends.base import build_endpoint_config | ||
| 20 | +from motor.node_manager.core.services.native_engine.backends.sglang.backend import SGLangBackend | ||
| 21 | +from motor.node_manager.core.services.native_engine.backends.vllm.backend import VllmBackend | ||
| 22 | +from motor.node_manager.core.services.native_engine.factory import get_backend | ||
| 23 | +from motor.node_manager.core.services.native_engine.models import LaunchContext | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +def _context( | ||
| 27 | + *, | ||
| 28 | + role: PDRole = PDRole.ROLE_P, | ||
| 29 | + headless: bool = False, | ||
| 30 | +) -> LaunchContext: | ||
| 31 | + return LaunchContext( | ||
| 32 | + role=role, | ||
| 33 | + instance_id=7, | ||
| 34 | + dp_rank=2, | ||
| 35 | + node_rank=1, | ||
| 36 | + host="10.0.0.2", | ||
| 37 | + business_port=8002, | ||
| 38 | + mgmt_port=9002, | ||
| 39 | + config_path="/config/user_config.json", | ||
| 40 | + master_dp_ip="10.0.0.1", | ||
| 41 | + kv_port=5002, | ||
| 42 | + lookup_rpc_port=6002, | ||
| 43 | + dp_rpc_port=7002, | ||
| 44 | + d2d_peer_ips=("10.0.0.3",), | ||
| 45 | + environment={"VLLM_HOST_IP": "10.0.0.2"}, | ||
| 46 | + headless=headless, | ||
| 47 | + ) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +def _endpoint( | ||
| 51 | + *, | ||
| 52 | + engine_type: str, | ||
| 53 | + role: str, | ||
| 54 | + connector: str | None = None, | ||
| 55 | +): | ||
| 56 | + engine_config = {} | ||
| 57 | + if connector is not None: | ||
| 58 | + engine_config["kv_transfer_config"] = {"kv_connector": connector} | ||
| 59 | + return SimpleNamespace( | ||
| 60 | + engine_type=engine_type, | ||
| 61 | + role=role, | ||
| 62 | + deploy_config=SimpleNamespace( | ||
| 63 | + engine_config=engine_config, | ||
| 64 | + dispatch_profile=None, | ||
| 65 | + health_check_config=SimpleNamespace( | ||
| 66 | + health_collector_timeout=5, | ||
| 67 | + health_collector_timeout_retry_attempts=3, | ||
| 68 | + startup_timeout=1800, | ||
| 69 | + ), | ||
| 70 | + infer_tls_config=None, | ||
| 71 | + ), | ||
| 72 | + ) | ||
| 73 | + | ||
| 74 | + | ||
| 75 | +def _prepare_with_config(backend, context, endpoint): | ||
| 76 | + native_config = MagicMock() | ||
| 77 | + native_config.get_cli_args.return_value = ["--model", "/models/glm-test", "--port", "8002"] | ||
| 78 | + with ( | ||
| 79 | + patch( | ||
| 80 | + "motor.node_manager.core.services.native_engine.backends.base.build_endpoint_config", return_value=endpoint | ||
| 81 | + ), | ||
| 82 | + patch("motor.node_manager.core.services.native_engine.config_factory.ConfigFactory") as factory, | ||
| 83 | + ): | ||
| 84 | + factory.return_value.build_cli_config.return_value = native_config | ||
| 85 | + return backend.prepare(context) | ||
| 86 | + | ||
| 87 | + | ||
| 88 | +def _build_with_config(backend, context, endpoint): | ||
| 89 | + return _prepare_with_config(backend, context, endpoint).command | ||
| 90 | + | ||
| 91 | + | ||
| 92 | +def test_vllm_backend_builds_native_command_and_preserves_environment(): | ||
| 93 | + context = _context() | ||
| 94 | + spec = _build_with_config( | ||
| 95 | + VllmBackend(), | ||
| 96 | + context, | ||
| 97 | + _endpoint(engine_type="vllm", role="prefill", connector="MooncakeConnectorV1"), | ||
| 98 | + ) | ||
| 99 | + | ||
| 100 | + assert spec.argv == ("vllm", "serve", "--model", "/models/glm-test", "--port", "8002") | ||
| 101 | + assert dict(spec.env) == {"VLLM_HOST_IP": "10.0.0.2"} | ||
| 102 | + assert "engine_server" not in spec.argv | ||
| 103 | + | ||
| 104 | + | ||
| 105 | +def test_backend_prepares_command_and_probe_from_one_endpoint_config(): | ||
| 106 | + context = _context() | ||
| 107 | + endpoint = _endpoint(engine_type="vllm", role="prefill", connector="MooncakeConnectorV1") | ||
| 108 | + native_config = MagicMock() | ||
| 109 | + native_config.get_cli_args.return_value = ["--model", "/models/glm-test"] | ||
| 110 | + | ||
| 111 | + with ( | ||
| 112 | + patch( | ||
| 113 | + "motor.node_manager.core.services.native_engine.backends.base.build_endpoint_config", | ||
| 114 | + return_value=endpoint, | ||
| 115 | + ) as build_endpoint, | ||
| 116 | + patch("motor.node_manager.core.services.native_engine.config_factory.ConfigFactory") as factory, | ||
| 117 | + ): | ||
| 118 | + factory.return_value.build_cli_config.return_value = native_config | ||
| 119 | + launch_spec = VllmBackend().prepare(context) | ||
| 120 | + | ||
| 121 | + build_endpoint.assert_called_once_with(context, "vllm") | ||
| 122 | + assert launch_spec.command.argv == ("vllm", "serve", "--model", "/models/glm-test") | ||
| 123 | + assert launch_spec.probe.path == "/health" | ||
| 124 | + native_config.convert.assert_called_once_with() | ||
| 125 | + native_config.validate.assert_called_once_with() | ||
| 126 | + | ||
| 127 | + | ||
| 128 | + | ||
| 129 | + "connector", | ||
| 130 | + ["MooncakeConnectorV1", "MooncakeHybridConnector", "NixlConnector"], | ||
| 131 | +) | ||
| 132 | +def test_vllm_backend_accepts_frozen_handoff_connector_whitelist(connector): | ||
| 133 | + context = _context() | ||
| 134 | + | ||
| 135 | + spec = _build_with_config( | ||
| 136 | + VllmBackend(), | ||
| 137 | + context, | ||
| 138 | + _endpoint(engine_type="vllm", role="prefill", connector=connector), | ||
| 139 | + ) | ||
| 140 | + | ||
| 141 | + assert spec.argv[:2] == ("vllm", "serve") | ||
| 142 | + | ||
| 143 | + | ||
| 144 | +def test_vllm_backend_accepts_multi_connector_with_handoff_transport(): | ||
| 145 | + context = _context() | ||
| 146 | + endpoint = _endpoint(engine_type="vllm", role="prefill", connector="MooncakeConnectorV1") | ||
| 147 | + endpoint.deploy_config.engine_config = { | ||
| 148 | + "kv_transfer_config": { | ||
| 149 | + "kv_connector": "MultiConnector", | ||
| 150 | + "kv_connector_extra_config": { | ||
| 151 | + "connectors": [ | ||
| 152 | + {"kv_connector": "NixlConnector"}, | ||
| 153 | + {"kv_connector": "AscendStoreConnector"}, | ||
| 154 | + ] | ||
| 155 | + }, | ||
| 156 | + } | ||
| 157 | + } | ||
| 158 | + | ||
| 159 | + spec = _build_with_config(VllmBackend(), context, endpoint) | ||
| 160 | + | ||
| 161 | + assert spec.argv[:2] == ("vllm", "serve") | ||
| 162 | + | ||
| 163 | + | ||
| 164 | + | ||
| 165 | + "connector,profile", | ||
| 166 | + [ | ||
| 167 | + ("MooncakeLayerwiseConnector", "trigger"), | ||
| 168 | + ("UnknownConnector", "unknown"), | ||
| 169 | + ], | ||
| 170 | +) | ||
| 171 | +def test_vllm_backend_rejects_non_handoff_pd_connector(connector, profile): | ||
| 172 | + context = _context() | ||
| 173 | + endpoint = _endpoint(engine_type="vllm", role="prefill", connector=connector) | ||
| 174 | + | ||
| 175 | + with ( | ||
| 176 | + patch( | ||
| 177 | + "motor.node_manager.core.services.native_engine.backends.base.build_endpoint_config", return_value=endpoint | ||
| 178 | + ), | ||
| 179 | + pytest.raises(ValueError, match=rf"resolved dispatch profile is {profile}"), | ||
| 180 | + ): | ||
| 181 | + VllmBackend().prepare(context) | ||
| 182 | + | ||
| 183 | + | ||
| 184 | +def test_vllm_backend_rejects_explicit_trigger_profile(): | ||
| 185 | + context = _context() | ||
| 186 | + endpoint = _endpoint( | ||
| 187 | + engine_type="vllm", | ||
| 188 | + role="prefill", | ||
| 189 | + connector="MooncakeConnectorV1", | ||
| 190 | + ) | ||
| 191 | + endpoint.deploy_config.dispatch_profile = "trigger" | ||
| 192 | + | ||
| 193 | + with ( | ||
| 194 | + patch( | ||
| 195 | + "motor.node_manager.core.services.native_engine.backends.base.build_endpoint_config", return_value=endpoint | ||
| 196 | + ), | ||
| 197 | + pytest.raises(ValueError, match="resolved dispatch profile is trigger"), | ||
| 198 | + ): | ||
| 199 | + VllmBackend().prepare(context) | ||
| 200 | + | ||
| 201 | + | ||
| 202 | +def test_vllm_backend_allows_union_without_kv_connector(): | ||
| 203 | + context = _context(role=PDRole.ROLE_U) | ||
| 204 | + | ||
| 205 | + spec = _build_with_config( | ||
| 206 | + VllmBackend(), | ||
| 207 | + context, | ||
| 208 | + _endpoint(engine_type="vllm", role="union"), | ||
| 209 | + ) | ||
| 210 | + | ||
| 211 | + assert spec.argv[:2] == ("vllm", "serve") | ||
| 212 | + | ||
| 213 | + | ||
| 214 | +def test_sglang_backend_builds_native_module_command(): | ||
| 215 | + context = _context(role=PDRole.ROLE_D) | ||
| 216 | + | ||
| 217 | + spec = _build_with_config( | ||
| 218 | + SGLangBackend(), | ||
| 219 | + context, | ||
| 220 | + _endpoint(engine_type="sglang", role="decode"), | ||
| 221 | + ) | ||
| 222 | + | ||
| 223 | + assert spec.argv[:3] == ("python3", "-m", "sglang.launch_server") | ||
| 224 | + assert "engine_server" not in spec.argv | ||
| 225 | + | ||
| 226 | + | ||
| 227 | +def test_sglang_backend_rejects_encode_before_loading_config(): | ||
| 228 | + context = _context(role=PDRole.ROLE_E) | ||
| 229 | + | ||
| 230 | + with ( | ||
| 231 | + patch("motor.node_manager.core.services.native_engine.backends.base.build_endpoint_config") as build_endpoint, | ||
| 232 | + pytest.raises(ValueError, match="SGLang encode role is not supported"), | ||
| 233 | + ): | ||
| 234 | + SGLangBackend().prepare(context) | ||
| 235 | + | ||
| 236 | + build_endpoint.assert_not_called() | ||
| 237 | + | ||
| 238 | + | ||
| 239 | +def test_get_backend_rejects_unknown_engine(): | ||
| 240 | + with pytest.raises(ValueError, match="Unsupported engine type"): | ||
| 241 | + get_backend("unknown") | ||
| 242 | + | ||
| 243 | + | ||
| 244 | +def testbuild_endpoint_config_maps_launch_context_without_cli_globals(): | ||
| 245 | + context = _context() | ||
| 246 | + | ||
| 247 | + with ( | ||
| 248 | + patch.object(EndpointConfig, "validate") as validate, | ||
| 249 | + patch.object(EndpointConfig, "load_deploy_config") as load_deploy_config, | ||
| 250 | + ): | ||
| 251 | + endpoint_config = build_endpoint_config(context, "vllm") | ||
| 252 | + | ||
| 253 | + validate.assert_called_once_with() | ||
| 254 | + load_deploy_config.assert_called_once_with() | ||
| 255 | + assert endpoint_config.role == "prefill" | ||
| 256 | + assert endpoint_config.dp_rank == 2 | ||
| 257 | + assert endpoint_config.node_rank == 1 | ||
| 258 | + assert endpoint_config.d2d_peer_ips == "10.0.0.3" | ||
| 259 | + assert endpoint_config.snapshot_metadata is None | ||
| 260 | + | ||
| 261 | + | ||
| 262 | +def test_backend_builds_native_health_probe_from_engine_config(): | ||
| 263 | + context = _context() | ||
| 264 | + tls_config = TLSConfig(enable_tls=True, ca_file="/certs/ca.crt") | ||
| 265 | + endpoint = _endpoint(engine_type="vllm", role="prefill", connector="MooncakeConnectorV1") | ||
| 266 | + endpoint.deploy_config.health_check_config = SimpleNamespace( | ||
| 267 | + health_collector_timeout=7, | ||
| 268 | + health_collector_timeout_retry_attempts=4, | ||
| 269 | + startup_timeout=900, | ||
| 270 | + ) | ||
| 271 | + endpoint.deploy_config.infer_tls_config = tls_config | ||
| 272 | + | ||
| 273 | + probe = _prepare_with_config(VllmBackend(), context, endpoint).probe | ||
| 274 | + | ||
| 275 | + assert probe.path == "/health" | ||
| 276 | + assert probe.timeout_seconds == 7 | ||
| 277 | + assert probe.max_attempts == 4 | ||
| 278 | + assert probe.startup_timeout_seconds == 900 | ||
| 279 | + assert probe.tls_config is tls_config | ||
| 280 | + assert probe.process_only is False | ||
| 281 | + | ||
| 282 | + | ||
| 283 | +def test_headless_backend_uses_process_only_probe(): | ||
| 284 | + context = _context(headless=True) | ||
| 285 | + endpoint = _endpoint(engine_type="vllm", role="prefill", connector="MooncakeConnectorV1") | ||
| 286 | + endpoint.deploy_config.health_check_config = SimpleNamespace( | ||
| 287 | + health_collector_timeout=5, | ||
| 288 | + health_collector_timeout_retry_attempts=3, | ||
| 289 | + startup_timeout=1800, | ||
| 290 | + ) | ||
| 291 | + endpoint.deploy_config.infer_tls_config = None | ||
| 292 | + | ||
| 293 | + probe = _prepare_with_config(VllmBackend(), context, endpoint).probe | ||
| 294 | + | ||
| 295 | + assert probe.process_only is True | ||
| @@ -0,0 +1,62 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from types import SimpleNamespace | ||
| 12 | +from unittest.mock import patch | ||
| 13 | + | ||
| 14 | +from motor.common.resources.endpoint import Endpoint | ||
| 15 | +from motor.common.resources.instance import PDRole | ||
| 16 | +from motor.node_manager.core.services.native_engine.service import NativeEngineService | ||
| 17 | +from motor.node_manager.core.services.native_engine.models import CommandSpec, LaunchSpec, ProbeSpec | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +def _native_engine_service() -> NativeEngineService: | ||
| 21 | + with ( | ||
| 22 | + patch("motor.node_manager.core.services.native_engine.service.get_backend") as get_backend, | ||
| 23 | + patch("motor.node_manager.core.services.native_engine.service.ProcessSupervisor") as supervisor_class, | ||
| 24 | + ): | ||
| 25 | + service = NativeEngineService( | ||
| 26 | + engine_type="sglang", | ||
| 27 | + config_path="/tmp/config.json", | ||
| 28 | + device_num=1, | ||
| 29 | + parallel_config=SimpleNamespace(local_world_size=1), | ||
| 30 | + enable_multi_endpoints=False, | ||
| 31 | + ) | ||
| 32 | + service.backend = get_backend.return_value | ||
| 33 | + service.supervisor = supervisor_class.return_value | ||
| 34 | + return service | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +def test_successful_pull_clears_recovery_latch(): | ||
| 38 | + service = _native_engine_service() | ||
| 39 | + service._recovery_requested = True | ||
| 40 | + service.backend.prepare.return_value = LaunchSpec( | ||
| 41 | + command=CommandSpec(argv=("python", "-m", "sglang.launch_server"), env={}), | ||
| 42 | + probe=ProbeSpec(path="/health", timeout_seconds=1, startup_timeout_seconds=10), | ||
| 43 | + ) | ||
| 44 | + service.supervisor.start.return_value = True | ||
| 45 | + endpoint = Endpoint(id=0, ip="127.0.0.1", business_port="8000", mgmt_port="8001") | ||
| 46 | + | ||
| 47 | + service.pull(PDRole.ROLE_D, [endpoint], instance_id=1, master_dp_ip="127.0.0.1") | ||
| 48 | + | ||
| 49 | + assert service._recovery_requested is False | ||
| 50 | + service.supervisor.start.assert_called_once() | ||
| 51 | + | ||
| 52 | + | ||
| 53 | + | ||
| 54 | +def test_health_check_requests_recovery_only_once(mock_kill): | ||
| 55 | + service = _native_engine_service() | ||
| 56 | + service.restart_on_failure = True | ||
| 57 | + service.supervisor.dead_pids.side_effect = [[101], [102]] | ||
| 58 | + | ||
| 59 | + service.health_check() | ||
| 60 | + service.health_check() | ||
| 61 | + | ||
| 62 | + mock_kill.assert_called_once() | ||
| @@ -0,0 +1,354 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +import os | ||
| 12 | +import signal | ||
| 13 | +import subprocess | ||
| 14 | +import threading | ||
| 15 | +from concurrent.futures import ThreadPoolExecutor | ||
| 16 | +from unittest.mock import MagicMock, call, patch | ||
| 17 | + | ||
| 18 | +import pytest | ||
| 19 | +import requests | ||
| 20 | + | ||
| 21 | +from motor.node_manager.core.services.native_engine.models import CommandSpec, ProbeSpec, RuntimeState | ||
| 22 | +from motor.node_manager.core.services.native_engine.supervisor import ProcessSupervisor | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +def _command() -> CommandSpec: | ||
| 26 | + return CommandSpec(argv=("vllm", "serve"), env={"KEY": "value"}) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def _probe( | ||
| 30 | + *, | ||
| 31 | + startup_timeout: float = 1800, | ||
| 32 | + max_attempts: int = 1, | ||
| 33 | + process_only: bool = False, | ||
| 34 | +) -> ProbeSpec: | ||
| 35 | + return ProbeSpec( | ||
| 36 | + path="/health", | ||
| 37 | + timeout_seconds=5, | ||
| 38 | + startup_timeout_seconds=startup_timeout, | ||
| 39 | + max_attempts=max_attempts, | ||
| 40 | + process_only=process_only, | ||
| 41 | + ) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | +def test_probe_rejects_non_positive_attempt_limit(): | ||
| 45 | + with pytest.raises(ValueError, match="max_attempts must be a positive integer"): | ||
| 46 | + _probe(max_attempts=0) | ||
| 47 | + | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +def test_start_owns_native_process_group(mock_popen): | ||
| 51 | + process = MagicMock(pid=12345) | ||
| 52 | + process.poll.return_value = None | ||
| 53 | + mock_popen.return_value = process | ||
| 54 | + supervisor = ProcessSupervisor() | ||
| 55 | + | ||
| 56 | + started = supervisor.start(3, _command(), _probe()) | ||
| 57 | + | ||
| 58 | + assert started is True | ||
| 59 | + assert supervisor.pid_list() == [12345] | ||
| 60 | + assert mock_popen.call_args.kwargs["start_new_session"] is True | ||
| 61 | + | ||
| 62 | + | ||
| 63 | + | ||
| 64 | +def test_repeated_start_is_idempotent(mock_popen): | ||
| 65 | + process = MagicMock(pid=12345) | ||
| 66 | + process.poll.return_value = None | ||
| 67 | + mock_popen.return_value = process | ||
| 68 | + supervisor = ProcessSupervisor() | ||
| 69 | + | ||
| 70 | + assert supervisor.start(3, _command(), _probe()) is True | ||
| 71 | + assert supervisor.start(3, _command(), _probe()) is False | ||
| 72 | + | ||
| 73 | + mock_popen.assert_called_once() | ||
| 74 | + assert supervisor.pid_list() == [12345] | ||
| 75 | + | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +def test_repeated_start_rejects_different_launch_spec_without_stopping_existing(mock_popen): | ||
| 79 | + process = MagicMock(pid=12345) | ||
| 80 | + process.poll.return_value = None | ||
| 81 | + mock_popen.return_value = process | ||
| 82 | + supervisor = ProcessSupervisor() | ||
| 83 | + supervisor.start(3, _command(), _probe()) | ||
| 84 | + | ||
| 85 | + different_command = CommandSpec(argv=("vllm", "serve", "--port", "9000"), env={"KEY": "value"}) | ||
| 86 | + with pytest.raises(RuntimeError, match="different launch spec"): | ||
| 87 | + supervisor.start(3, different_command, _probe()) | ||
| 88 | + | ||
| 89 | + mock_popen.assert_called_once() | ||
| 90 | + assert supervisor.pid_list() == [12345] | ||
| 91 | + | ||
| 92 | + | ||
| 93 | + | ||
| 94 | +def test_concurrent_start_creates_only_one_process(mock_popen): | ||
| 95 | + process = MagicMock(pid=12345) | ||
| 96 | + process.poll.return_value = None | ||
| 97 | + second_popen_entered = threading.Event() | ||
| 98 | + popen_calls = 0 | ||
| 99 | + calls_lock = threading.Lock() | ||
| 100 | + | ||
| 101 | + def popen(*_args, **_kwargs): | ||
| 102 | + nonlocal popen_calls | ||
| 103 | + with calls_lock: | ||
| 104 | + popen_calls += 1 | ||
| 105 | + call_number = popen_calls | ||
| 106 | + if call_number == 1: | ||
| 107 | + second_popen_entered.wait(timeout=0.1) | ||
| 108 | + else: | ||
| 109 | + second_popen_entered.set() | ||
| 110 | + return process | ||
| 111 | + | ||
| 112 | + mock_popen.side_effect = popen | ||
| 113 | + supervisor = ProcessSupervisor() | ||
| 114 | + with ThreadPoolExecutor(max_workers=2) as executor: | ||
| 115 | + results = list(executor.map(lambda _: supervisor.start(3, _command(), _probe()), range(2))) | ||
| 116 | + | ||
| 117 | + assert sorted(results) == [False, True] | ||
| 118 | + assert mock_popen.call_count == 1 | ||
| 119 | + assert supervisor.pid_list() == [12345] | ||
| 120 | + | ||
| 121 | + | ||
| 122 | + | ||
| 123 | + | ||
| 124 | +def test_failed_probe_stays_starting_until_startup_timeout(mock_popen, mock_client): | ||
| 125 | + process = MagicMock(pid=12345) | ||
| 126 | + process.poll.return_value = None | ||
| 127 | + mock_popen.return_value = process | ||
| 128 | + mock_client.return_value.__enter__.return_value.do_get.side_effect = RuntimeError("not listening") | ||
| 129 | + supervisor = ProcessSupervisor() | ||
| 130 | + supervisor.start(3, _command(), _probe(startup_timeout=600)) | ||
| 131 | + | ||
| 132 | + with patch("motor.node_manager.core.services.native_engine.supervisor.time.monotonic", return_value=100): | ||
| 133 | + supervisor._processes[3].started_at = 0 | ||
| 134 | + state = supervisor.state(3, "10.0.0.1", 8000) | ||
| 135 | + | ||
| 136 | + assert state == RuntimeState.STARTING | ||
| 137 | + | ||
| 138 | + | ||
| 139 | + | ||
| 140 | + | ||
| 141 | +def test_failed_probe_after_startup_timeout_is_unhealthy(mock_popen, mock_client): | ||
| 142 | + process = MagicMock(pid=12345) | ||
| 143 | + process.poll.return_value = None | ||
| 144 | + mock_popen.return_value = process | ||
| 145 | + mock_client.return_value.__enter__.return_value.do_get.side_effect = RuntimeError("not listening") | ||
| 146 | + supervisor = ProcessSupervisor() | ||
| 147 | + supervisor.start(3, _command(), _probe(startup_timeout=10)) | ||
| 148 | + | ||
| 149 | + with patch("motor.node_manager.core.services.native_engine.supervisor.time.monotonic", return_value=11): | ||
| 150 | + supervisor._processes[3].started_at = 0 | ||
| 151 | + state = supervisor.state(3, "10.0.0.1", 8000) | ||
| 152 | + | ||
| 153 | + assert state == RuntimeState.UNHEALTHY | ||
| 154 | + | ||
| 155 | + | ||
| 156 | + | ||
| 157 | + | ||
| 158 | +def test_successful_native_probe_marks_ready(mock_popen, mock_client): | ||
| 159 | + process = MagicMock(pid=12345) | ||
| 160 | + process.poll.return_value = None | ||
| 161 | + mock_popen.return_value = process | ||
| 162 | + supervisor = ProcessSupervisor() | ||
| 163 | + supervisor.start(3, _command(), _probe()) | ||
| 164 | + | ||
| 165 | + state = supervisor.state(3, "10.0.0.1", 8000) | ||
| 166 | + | ||
| 167 | + assert state == RuntimeState.READY | ||
| 168 | + mock_client.return_value.__enter__.return_value.do_get.assert_called_once_with("/health") | ||
| 169 | + | ||
| 170 | + | ||
| 171 | + | ||
| 172 | + | ||
| 173 | +def test_native_probe_retries_timeout_within_attempt_limit(mock_popen, mock_client): | ||
| 174 | + process = MagicMock(pid=12345) | ||
| 175 | + process.poll.return_value = None | ||
| 176 | + mock_popen.return_value = process | ||
| 177 | + timeout = requests.exceptions.ReadTimeout("timed out") | ||
| 178 | + wrapped_timeout = RuntimeError("send request failed") | ||
| 179 | + wrapped_timeout.__cause__ = timeout | ||
| 180 | + client = mock_client.return_value.__enter__.return_value | ||
| 181 | + client.do_get.side_effect = [wrapped_timeout, None] | ||
| 182 | + supervisor = ProcessSupervisor() | ||
| 183 | + supervisor.start(3, _command(), _probe(max_attempts=2)) | ||
| 184 | + | ||
| 185 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.READY | ||
| 186 | + assert client.do_get.call_count == 2 | ||
| 187 | + | ||
| 188 | + | ||
| 189 | + | ||
| 190 | + | ||
| 191 | +def test_native_probe_does_not_retry_non_timeout_error(mock_popen, mock_client): | ||
| 192 | + process = MagicMock(pid=12345) | ||
| 193 | + process.poll.return_value = None | ||
| 194 | + mock_popen.return_value = process | ||
| 195 | + client = mock_client.return_value.__enter__.return_value | ||
| 196 | + client.do_get.side_effect = RuntimeError("connection refused") | ||
| 197 | + supervisor = ProcessSupervisor() | ||
| 198 | + supervisor.start(3, _command(), _probe(startup_timeout=10, max_attempts=3)) | ||
| 199 | + | ||
| 200 | + with patch("motor.node_manager.core.services.native_engine.supervisor.time.monotonic", return_value=11): | ||
| 201 | + supervisor._processes[3].started_at = 0 | ||
| 202 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.UNHEALTHY | ||
| 203 | + | ||
| 204 | + client.do_get.assert_called_once_with("/health") | ||
| 205 | + | ||
| 206 | + | ||
| 207 | + | ||
| 208 | + | ||
| 209 | +def test_native_probe_marks_unhealthy_after_exhausting_timeout_retries(mock_popen, mock_client): | ||
| 210 | + process = MagicMock(pid=12345) | ||
| 211 | + process.poll.return_value = None | ||
| 212 | + mock_popen.return_value = process | ||
| 213 | + mock_client.return_value.__enter__.return_value.do_get.side_effect = requests.exceptions.ReadTimeout("timed out") | ||
| 214 | + supervisor = ProcessSupervisor() | ||
| 215 | + supervisor.start(3, _command(), _probe(startup_timeout=10, max_attempts=2)) | ||
| 216 | + | ||
| 217 | + with patch("motor.node_manager.core.services.native_engine.supervisor.time.monotonic", return_value=11): | ||
| 218 | + supervisor._processes[3].started_at = 0 | ||
| 219 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.UNHEALTHY | ||
| 220 | + | ||
| 221 | + assert mock_client.return_value.__enter__.return_value.do_get.call_count == 2 | ||
| 222 | + | ||
| 223 | + | ||
| 224 | + | ||
| 225 | + | ||
| 226 | +def test_headless_probe_reports_running_without_claiming_readiness(mock_popen, mock_client): | ||
| 227 | + process = MagicMock(pid=12345) | ||
| 228 | + process.poll.return_value = None | ||
| 229 | + mock_popen.return_value = process | ||
| 230 | + supervisor = ProcessSupervisor() | ||
| 231 | + supervisor.start(3, _command(), _probe(process_only=True)) | ||
| 232 | + | ||
| 233 | + assert supervisor.state(3, "10.0.0.2", 8000) == RuntimeState.RUNNING | ||
| 234 | + assert supervisor._processes[3].ready_at is None | ||
| 235 | + mock_client.assert_not_called() | ||
| 236 | + | ||
| 237 | + | ||
| 238 | + | ||
| 239 | +def test_dead_process_is_stopped_without_http_probe(mock_popen): | ||
| 240 | + process = MagicMock(pid=12345) | ||
| 241 | + process.poll.side_effect = [None, 1, 1] | ||
| 242 | + mock_popen.return_value = process | ||
| 243 | + supervisor = ProcessSupervisor() | ||
| 244 | + supervisor.start(3, _command(), _probe()) | ||
| 245 | + | ||
| 246 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.STOPPED | ||
| 247 | + with patch.object(supervisor, "_kill_group") as kill_group: | ||
| 248 | + assert supervisor.dead_pids() == [12345] | ||
| 249 | + kill_group.assert_called_once() | ||
| 250 | + assert supervisor.dead_pids() == [] | ||
| 251 | + assert supervisor.pid_list() == [] | ||
| 252 | + | ||
| 253 | + | ||
| 254 | + | ||
| 255 | +def test_stop_uses_graceful_then_forced_process_group_cleanup(mock_popen): | ||
| 256 | + process = MagicMock(pid=12345) | ||
| 257 | + process.poll.return_value = None | ||
| 258 | + process.wait.side_effect = subprocess.TimeoutExpired(cmd="vllm", timeout=0) | ||
| 259 | + mock_popen.return_value = process | ||
| 260 | + supervisor = ProcessSupervisor(stop_grace_seconds=0) | ||
| 261 | + supervisor.start(3, _command(), _probe()) | ||
| 262 | + | ||
| 263 | + with ( | ||
| 264 | + patch.object(supervisor, "_terminate_group") as terminate_group, | ||
| 265 | + patch.object(supervisor, "_kill_group") as kill_group, | ||
| 266 | + ): | ||
| 267 | + stopped = supervisor.stop_all() | ||
| 268 | + | ||
| 269 | + assert stopped == [12345] | ||
| 270 | + runtime = terminate_group.call_args.args[0] | ||
| 271 | + assert runtime.process is process | ||
| 272 | + kill_group.assert_called_once_with(runtime) | ||
| 273 | + assert supervisor.pid_list() == [] | ||
| 274 | + | ||
| 275 | + | ||
| 276 | + | ||
| 277 | + | ||
| 278 | +def test_stop_kills_remaining_group_after_leader_exits(mock_popen): | ||
| 279 | + process = MagicMock(pid=12345) | ||
| 280 | + process.poll.return_value = None | ||
| 281 | + process.wait.return_value = 0 | ||
| 282 | + mock_popen.return_value = process | ||
| 283 | + supervisor = ProcessSupervisor() | ||
| 284 | + supervisor.start(3, _command(), _probe()) | ||
| 285 | + | ||
| 286 | + with patch("motor.node_manager.core.services.native_engine.supervisor.os.killpg") as killpg: | ||
| 287 | + assert supervisor.stop_all() == [12345] | ||
| 288 | + | ||
| 289 | + assert killpg.call_args_list == [ | ||
| 290 | + call(12345, signal.SIGTERM), | ||
| 291 | + call(12345, 0), | ||
| 292 | + call(12345, signal.SIGKILL), | ||
| 293 | + ] | ||
| 294 | + | ||
| 295 | + | ||
| 296 | + | ||
| 297 | +def test_stop_one_endpoint_preserves_other_processes(mock_popen): | ||
| 298 | + first = MagicMock(pid=12345) | ||
| 299 | + second = MagicMock(pid=12346) | ||
| 300 | + first.poll.return_value = None | ||
| 301 | + second.poll.return_value = None | ||
| 302 | + mock_popen.side_effect = [first, second] | ||
| 303 | + supervisor = ProcessSupervisor() | ||
| 304 | + supervisor.start(3, _command(), _probe()) | ||
| 305 | + supervisor.start(4, _command(), _probe()) | ||
| 306 | + | ||
| 307 | + with patch.object(supervisor, "_terminate_group"): | ||
| 308 | + assert supervisor.stop(4) == 12346 | ||
| 309 | + | ||
| 310 | + assert supervisor.pid_list() == [12345] | ||
| 311 | + | ||
| 312 | + | ||
| 313 | + | ||
| 314 | +def test_stop_state_remains_visible_until_process_cleanup_finishes(mock_popen): | ||
| 315 | + process = MagicMock(pid=12345) | ||
| 316 | + process.poll.return_value = None | ||
| 317 | + entered_wait = threading.Event() | ||
| 318 | + release_wait = threading.Event() | ||
| 319 | + | ||
| 320 | + def wait(*_args, **_kwargs): | ||
| 321 | + entered_wait.set() | ||
| 322 | + release_wait.wait(timeout=1) | ||
| 323 | + return 0 | ||
| 324 | + | ||
| 325 | + process.wait.side_effect = wait | ||
| 326 | + mock_popen.return_value = process | ||
| 327 | + supervisor = ProcessSupervisor() | ||
| 328 | + supervisor.start(3, _command(), _probe()) | ||
| 329 | + | ||
| 330 | + with patch.object(supervisor, "_terminate_group"): | ||
| 331 | + thread = threading.Thread(target=supervisor.stop_all) | ||
| 332 | + thread.start() | ||
| 333 | + assert entered_wait.wait(timeout=1) | ||
| 334 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.STOPPING | ||
| 335 | + release_wait.set() | ||
| 336 | + thread.join(timeout=1) | ||
| 337 | + | ||
| 338 | + assert supervisor.state(3, "10.0.0.1", 8000) == RuntimeState.STOPPED | ||
| 339 | + | ||
| 340 | + | ||
| 341 | + | ||
| 342 | + | ||
| 343 | +def test_cleanup_uses_cached_process_group_when_leader_has_exited(mock_popen): | ||
| 344 | + process = MagicMock(pid=12345) | ||
| 345 | + process.poll.side_effect = [None, 1] | ||
| 346 | + mock_popen.return_value = process | ||
| 347 | + supervisor = ProcessSupervisor() | ||
| 348 | + supervisor.start(3, _command(), _probe()) | ||
| 349 | + runtime = supervisor._processes[3] | ||
| 350 | + | ||
| 351 | + with patch("motor.node_manager.core.services.native_engine.supervisor.os.killpg") as killpg: | ||
| 352 | + supervisor._terminate_group(runtime) | ||
| 353 | + | ||
| 354 | + killpg.assert_called_once_with(12345, signal.SIGTERM) | ||
| @@ -10,6 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | import os | 11 | import os |
| 12 | import json | 12 | import json |
| 13 | +import signal | ||
| 13 | import pytest | 14 | import pytest |
| 14 | from unittest.mock import patch, MagicMock, mock_open | 15 | from unittest.mock import patch, MagicMock, mock_open |
| 15 | 16 | ||
| @@ -18,6 +19,35 @@ from motor.node_manager.core.services.registry import SERVICE_ENGINE | |||
| 18 | from motor.config.node_manager import NodeManagerConfig | 19 | from motor.config.node_manager import NodeManagerConfig |
| 19 | from motor.common.resources.endpoint import Endpoint | 20 | from motor.common.resources.endpoint import Endpoint |
| 20 | from motor.common.resources.instance import PDRole, ParallelConfig | 21 | from motor.common.resources.instance import PDRole, ParallelConfig |
| 22 | +from motor.node_manager.core.services.native_engine.models import CommandSpec, LaunchSpec, ProbeSpec | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +def _recording_backend(): | ||
| 26 | + backend = MagicMock() | ||
| 27 | + | ||
| 28 | + def prepare(context): | ||
| 29 | + return LaunchSpec( | ||
| 30 | + command=CommandSpec( | ||
| 31 | + argv=( | ||
| 32 | + "vllm", | ||
| 33 | + "serve", | ||
| 34 | + "--host", | ||
| 35 | + context.host, | ||
| 36 | + "--port", | ||
| 37 | + str(context.business_port), | ||
| 38 | + ), | ||
| 39 | + env=context.environment, | ||
| 40 | + ), | ||
| 41 | + probe=ProbeSpec(path="/health", timeout_seconds=5, startup_timeout_seconds=1800), | ||
| 42 | + ) | ||
| 43 | + | ||
| 44 | + backend.prepare.side_effect = prepare | ||
| 45 | + return backend | ||
| 46 | + | ||
| 47 | + | ||
| 48 | +def _last_launch_context(daemon): | ||
| 49 | + backend = daemon._services[SERVICE_ENGINE].backend | ||
| 50 | + return backend.prepare.call_args.args[0] | ||
| 21 | 51 | ||
| 22 | 52 | ||
| 23 | def create_config_mock(config_data): | 53 | def create_config_mock(config_data): |
| @@ -57,14 +87,18 @@ def daemon(config_data): | |||
| 57 | tp_size=config_data["parallel_config"]["tp_size"], pp_size=config_data["parallel_config"]["pp_size"] | 87 | tp_size=config_data["parallel_config"]["tp_size"], pp_size=config_data["parallel_config"]["pp_size"] |
| 58 | ) | 88 | ) |
| 59 | config.basic_config.job_name = config_data.get("model_name", "test_job") | 89 | config.basic_config.job_name = config_data.get("model_name", "test_job") |
| 90 | + config.basic_config.model_name = config_data.get("model_name", "test-model") | ||
| 91 | + config.basic_config.engine_type = "vllm" | ||
| 60 | config.basic_config.role = PDRole(config_data.get("role", "both")) | 92 | config.basic_config.role = PDRole(config_data.get("role", "both")) |
| 61 | config.api_config.node_manager_port = config_data.get("node_manager_port", 8080) | 93 | config.api_config.node_manager_port = config_data.get("node_manager_port", 8080) |
| 62 | 94 | ||
| 63 | # Set device_num for testing (simulating visible devices) | 95 | # Set device_num for testing (simulating visible devices) |
| 64 | config.basic_config.device_num = 8 # 8 devices for testing | 96 | config.basic_config.device_num = 8 # 8 devices for testing |
| 65 | 97 | ||
| 66 | - daemon_instance = Daemon(config) | 98 | + backend = _recording_backend() |
| 67 | - yield daemon_instance | 99 | + with patch("motor.node_manager.core.services.native_engine.service.get_backend", return_value=backend): |
| 100 | + daemon_instance = Daemon(config) | ||
| 101 | + yield daemon_instance | ||
| 68 | 102 | ||
| 69 | 103 | ||
| 70 | 104 | ||
| @@ -76,6 +110,48 @@ def endpoints(): | |||
| 76 | 110 | ||
| 77 | 111 | ||
| 78 | class TestDaemon: | 112 | class TestDaemon: |
| 113 | + | ||
| 114 | + ("launch_env", "expected_host_ip", "expected_mc_ipv6"), | ||
| 115 | + [ | ||
| 116 | + ( | ||
| 117 | + {"POD_IP": "10.0.0.8", "MOONCAKE_ASCEND_IPV6_EXPERIMENT": "1"}, | ||
| 118 | + "10.0.0.8", | ||
| 119 | + "1", | ||
| 120 | + ), | ||
| 121 | + ( | ||
| 122 | + { | ||
| 123 | + "POD_IP": "10.0.0.8", | ||
| 124 | + "VLLM_HOST_IP": "10.0.0.9", | ||
| 125 | + "MOONCAKE_ASCEND_IPV6_EXPERIMENT": "1", | ||
| 126 | + "MC_USE_IPV6": "0", | ||
| 127 | + }, | ||
| 128 | + "10.0.0.9", | ||
| 129 | + "0", | ||
| 130 | + ), | ||
| 131 | + ], | ||
| 132 | + ) | ||
| 133 | + | ||
| 134 | + def test_engine_launch_environment_contract( | ||
| 135 | + self, | ||
| 136 | + mock_popen, | ||
| 137 | + daemon, | ||
| 138 | + launch_env, | ||
| 139 | + expected_host_ip, | ||
| 140 | + expected_mc_ipv6, | ||
| 141 | + ): | ||
| 142 | + mock_process = MagicMock(pid=12345) | ||
| 143 | + mock_process.poll.return_value = None | ||
| 144 | + mock_popen.return_value = mock_process | ||
| 145 | + endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") | ||
| 146 | + | ||
| 147 | + with patch.dict(os.environ, launch_env, clear=True): | ||
| 148 | + daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") | ||
| 149 | + | ||
| 150 | + child_env = mock_popen.call_args.kwargs["env"] | ||
| 151 | + assert child_env["VLLM_HOST_IP"] == expected_host_ip | ||
| 152 | + assert child_env["MC_USE_IPV6"] == expected_mc_ipv6 | ||
| 153 | + assert child_env["ASCEND_RT_VISIBLE_DEVICES"] == "0,1" | ||
| 154 | + | ||
| 79 | 155 | ||
| 80 | def test_pull_engine_success(self, mock_popen, daemon, endpoints): | 156 | def test_pull_engine_success(self, mock_popen, daemon, endpoints): |
| 81 | mock_process = MagicMock(pid=12345) | 157 | mock_process = MagicMock(pid=12345) |
| @@ -85,8 +161,19 @@ class TestDaemon: | |||
| 85 | master_dp_ip = "192.168.1.100" | 161 | master_dp_ip = "192.168.1.100" |
| 86 | daemon.pull_engine(PDRole.ROLE_P, endpoints, instance_id, master_dp_ip) | 162 | daemon.pull_engine(PDRole.ROLE_P, endpoints, instance_id, master_dp_ip) |
| 87 | # Verify that process was added to engine_pids | 163 | # Verify that process was added to engine_pids |
| 88 | - assert len(daemon._services[SERVICE_ENGINE].engine_pids) > 0 | 164 | + assert len(daemon.engine_pids) > 0 |
| 89 | - assert 12345 in daemon._services[SERVICE_ENGINE].engine_pids | 165 | + assert 12345 in daemon.engine_pids |
| 166 | + | ||
| 167 | + def test_pull_engine_failure_rolls_back_only_new_endpoints(self, daemon, endpoints): | ||
| 168 | + engine = daemon._services[SERVICE_ENGINE] | ||
| 169 | + with ( | ||
| 170 | + patch.object(engine.supervisor, "start", side_effect=[False, True, RuntimeError("start failed")]), | ||
| 171 | + patch.object(engine.supervisor, "stop") as stop, | ||
| 172 | + ): | ||
| 173 | + with pytest.raises(RuntimeError, match="start failed"): | ||
| 174 | + daemon.pull_engine(PDRole.ROLE_P, endpoints, instance_id=1, master_dp_ip="192.168.1.100") | ||
| 175 | + | ||
| 176 | + stop.assert_called_once_with(endpoints[1].id) | ||
| 90 | 177 | ||
| 91 | 178 | ||
| 92 | "invalid_endpoint,error_msg", | 179 | "invalid_endpoint,error_msg", |
| @@ -99,41 +186,36 @@ class TestDaemon: | |||
| 99 | with pytest.raises(RuntimeError, match=error_msg): | 186 | with pytest.raises(RuntimeError, match=error_msg): |
| 100 | daemon.pull_engine(PDRole.ROLE_U, [invalid_endpoint], instance_id=1, master_dp_ip="192.168.1.100") | 187 | daemon.pull_engine(PDRole.ROLE_U, [invalid_endpoint], instance_id=1, master_dp_ip="192.168.1.100") |
| 101 | 188 | ||
| 102 | - @pytest.mark.parametrize( | 189 | + def test_exit_daemon_delegates_to_process_supervisor(self, daemon): |
| 103 | - "exception,should_not_raise", | 190 | + supervisor = daemon._services[SERVICE_ENGINE].supervisor |
| 104 | - [ | 191 | + with patch.object(supervisor, "stop_all", return_value=[1001, 1002]) as stop_all: |
| 105 | - (None, True), | 192 | + daemon.stop() |
| 106 | - (ProcessLookupError("No such process"), True), | 193 | + stop_all.assert_called_once_with() |
| 107 | - (PermissionError("Permission denied"), True), | ||
| 108 | - (Exception("Unexpected error"), True), | ||
| 109 | - ], | ||
| 110 | - ) | ||
| 111 | - | ||
| 112 | - def test_exit_daemon(self, mock_kill, daemon, exception, should_not_raise): | ||
| 113 | - # Mock SIGKILL for Windows compatibility | ||
| 114 | - with patch('motor.node_manager.core.services.engine.signal.SIGKILL', 9, create=True): | ||
| 115 | - daemon._services[SERVICE_ENGINE].engine_pids = [1001, 1002] | ||
| 116 | - if exception: | ||
| 117 | - mock_kill.side_effect = exception | ||
| 118 | - daemon.stop() # Method is called 'stop', not 'exit_daemon' | ||
| 119 | - assert mock_kill.call_count == len([1001, 1002]) | ||
| 120 | 194 | ||
| 121 | - @pytest.mark.parametrize( | 195 | + def test_engine_exit_requests_pod_recovery_only_once(self, daemon): |
| 122 | - "ip,port,expected", | 196 | + engine = daemon._services[SERVICE_ENGINE] |
| 123 | - [ | 197 | + engine.restart_on_failure = True |
| 124 | - ("192.168.1.100", "8080", True), | 198 | + with ( |
| 125 | - ("2001:db8::1", "8080", True), | 199 | + patch.object(engine.supervisor, "dead_pids", return_value=[12345]), |
| 126 | - ("invalid_ip", "8080", False), | 200 | + patch("motor.node_manager.core.services.native_engine.service.os.kill") as kill, |
| 127 | - ("192.168.1.100", "not_number", False), | 201 | + ): |
| 128 | - ("192.168.1.100", "0", False), | 202 | + engine.health_check() |
| 129 | - ("192.168.1.100", "99999", False), | 203 | + engine.health_check() |
| 130 | - ("192.168.1.100", "1", False), | 204 | + |
| 131 | - ("192.168.1.100", "65535", True), | 205 | + kill.assert_called_once_with(os.getpid(), signal.SIGTERM) |
| 132 | - ], | 206 | + |
| 133 | - ) | 207 | + @patch('subprocess.Popen') |
| 134 | - def test_check_params(self, daemon, ip, port, expected): | 208 | + def test_native_metrics_target_uses_business_port(self, mock_popen, daemon): |
| 135 | - endpoint = Endpoint(id=1, ip=ip, business_port=port, mgmt_port="9090") | 209 | + process = MagicMock(pid=12345) |
| 136 | - assert daemon._services[SERVICE_ENGINE]._check_params(endpoint) == expected | 210 | + process.poll.return_value = None |
| 211 | + mock_popen.return_value = process | ||
| 212 | + endpoint = Endpoint(id=0, ip="2001:db8::8", business_port="8000", mgmt_port="9000") | ||
| 213 | + daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") | ||
| 214 | + | ||
| 215 | + assert daemon.get_engine_metrics_target(endpoint) == "http://[2001:db8::8]:8000/metrics" | ||
| 216 | + | ||
| 217 | + endpoint.headless = True | ||
| 218 | + assert daemon.get_engine_metrics_target(endpoint) is None | ||
| 137 | 219 | ||
| 138 | 220 | ||
| 139 | 221 | ||
| @@ -148,10 +230,12 @@ class TestDaemon: | |||
| 148 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id, master_dp_ip) | 230 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id, master_dp_ip) |
| 149 | 231 | ||
| 150 | # Verify that process was added to engine_pids | 232 | # Verify that process was added to engine_pids |
| 151 | - assert len(daemon._services[SERVICE_ENGINE].engine_pids) > 0 | 233 | + assert len(daemon.engine_pids) > 0 |
| 152 | - assert 12345 in daemon._services[SERVICE_ENGINE].engine_pids | 234 | + assert 12345 in daemon.engine_pids |
| 153 | # Verify Popen was called | 235 | # Verify Popen was called |
| 154 | mock_popen.assert_called_once() | 236 | mock_popen.assert_called_once() |
| 237 | + assert mock_popen.call_args.args[0][:2] == ["vllm", "serve"] | ||
| 238 | + assert "engine_server" not in mock_popen.call_args.args[0] | ||
| 155 | 239 | ||
| 156 | 240 | ||
| 157 | def test_hybrid_role_starts_union_engine(self, mock_popen, daemon): | 241 | def test_hybrid_role_starts_union_engine(self, mock_popen, daemon): |
| @@ -162,9 +246,8 @@ class TestDaemon: | |||
| 162 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") | 246 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") |
| 163 | daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") | 247 | daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") |
| 164 | 248 | ||
| 165 | - cmd = mock_popen.call_args.args[0] | 249 | + context = _last_launch_context(daemon) |
| 166 | - role_arg_index = cmd.index("--role") + 1 | 250 | + assert context.role == PDRole.ROLE_U |
| 167 | - assert cmd[role_arg_index] == "union" | ||
| 168 | 251 | ||
| 169 | 252 | ||
| 170 | def test_single_container_hybrid_omits_kv_port_when_unset(self, mock_popen, config_data): | 253 | def test_single_container_hybrid_omits_kv_port_when_unset(self, mock_popen, config_data): |
| @@ -180,11 +263,14 @@ class TestDaemon: | |||
| 180 | ) | 263 | ) |
| 181 | config.basic_config.device_num = 8 | 264 | config.basic_config.device_num = 8 |
| 182 | config.basic_config.enable_multi_endpoints = False | 265 | config.basic_config.enable_multi_endpoints = False |
| 266 | + config.basic_config.engine_type = "vllm" | ||
| 183 | config.single_container_config.single_container_flag = True | 267 | config.single_container_config.single_container_flag = True |
| 184 | config.single_container_config.device_offset = 0 | 268 | config.single_container_config.device_offset = 0 |
| 185 | config.single_container_config.kv_port = None | 269 | config.single_container_config.kv_port = None |
| 186 | config.single_container_config.dp_rpc_port = 9000 | 270 | config.single_container_config.dp_rpc_port = 9000 |
| 187 | - daemon = Daemon(config) | 271 | + backend = _recording_backend() |
| 272 | + with patch("motor.node_manager.core.services.native_engine.service.get_backend", return_value=backend): | ||
| 273 | + daemon = Daemon(config) | ||
| 188 | 274 | ||
| 189 | mock_process = MagicMock(pid=12345) | 275 | mock_process = MagicMock(pid=12345) |
| 190 | mock_process.poll.return_value = None | 276 | mock_process.poll.return_value = None |
| @@ -193,9 +279,9 @@ class TestDaemon: | |||
| 193 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") | 279 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") |
| 194 | daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") | 280 | daemon.pull_engine(PDRole.ROLE_U, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") |
| 195 | 281 | ||
| 196 | - cmd = mock_popen.call_args.args[0] | 282 | + context = _last_launch_context(daemon) |
| 197 | - assert "--kv-port" not in cmd | 283 | + assert context.kv_port is None |
| 198 | - assert cmd[cmd.index("--dp-rpc-port") + 1] == "9000" | 284 | + assert context.dp_rpc_port == 9000 |
| 199 | 285 | ||
| 200 | # ===== D2D Weight Transfer Tests ===== | 286 | # ===== D2D Weight Transfer Tests ===== |
| 201 | 287 | ||
| @@ -218,10 +304,8 @@ class TestDaemon: | |||
| 218 | ) | 304 | ) |
| 219 | 305 | ||
| 220 | mock_popen.assert_called_once() | 306 | mock_popen.assert_called_once() |
| 221 | - cmd = mock_popen.call_args.args[0] | 307 | + context = _last_launch_context(daemon) |
| 222 | - assert '--d2d-peer-ips' in cmd | 308 | + assert context.d2d_peer_ips == ("192.168.1.10", "192.168.1.11") |
| 223 | - idx = cmd.index('--d2d-peer-ips') | ||
| 224 | - assert cmd[idx + 1] == "192.168.1.10,192.168.1.11" | ||
| 225 | 309 | ||
| 226 | 310 | ||
| 227 | def test_pull_engine_without_d2d_peer_ips(self, mock_popen, daemon): | 311 | def test_pull_engine_without_d2d_peer_ips(self, mock_popen, daemon): |
| @@ -240,8 +324,7 @@ class TestDaemon: | |||
| 240 | ) | 324 | ) |
| 241 | 325 | ||
| 242 | mock_popen.assert_called_once() | 326 | mock_popen.assert_called_once() |
| 243 | - cmd = mock_popen.call_args.args[0] | 327 | + assert _last_launch_context(daemon).d2d_peer_ips == () |
| 244 | - assert '--d2d-peer-ips' not in cmd | ||
| 245 | 328 | ||
| 246 | 329 | ||
| 247 | def test_pull_engine_with_empty_d2d_peer_ips(self, mock_popen, daemon): | 330 | def test_pull_engine_with_empty_d2d_peer_ips(self, mock_popen, daemon): |
| @@ -263,8 +346,7 @@ class TestDaemon: | |||
| 263 | ) | 346 | ) |
| 264 | 347 | ||
| 265 | mock_popen.assert_called_once() | 348 | mock_popen.assert_called_once() |
| 266 | - cmd = mock_popen.call_args.args[0] | 349 | + assert _last_launch_context(daemon).d2d_peer_ips == () |
| 267 | - assert '--d2d-peer-ips' not in cmd | ||
| 268 | 350 | ||
| 269 | 351 | ||
| 270 | def test_pull_engine_with_d2d_peer_ips_rank_encoded(self, mock_popen, daemon): | 352 | def test_pull_engine_with_d2d_peer_ips_rank_encoded(self, mock_popen, daemon): |
| @@ -288,10 +370,9 @@ class TestDaemon: | |||
| 288 | ) | 370 | ) |
| 289 | 371 | ||
| 290 | assert mock_popen.call_count == 2 | 372 | assert mock_popen.call_count == 2 |
| 291 | - first_cmd = mock_popen.call_args_list[0].args[0] | 373 | + contexts = [call.args[0] for call in daemon._services[SERVICE_ENGINE].backend.prepare.call_args_list] |
| 292 | - second_cmd = mock_popen.call_args_list[1].args[0] | 374 | + assert contexts[0].d2d_peer_ips == ("192.168.1.10",) |
| 293 | - assert first_cmd[first_cmd.index('--d2d-peer-ips') + 1] == "192.168.1.10" | 375 | + assert contexts[1].d2d_peer_ips == ("192.168.1.11",) |
| 294 | - assert second_cmd[second_cmd.index('--d2d-peer-ips') + 1] == "192.168.1.11" | ||
| 295 | 376 | ||
| 296 | 377 | ||
| 297 | def test_pull_engine_d2d_peer_ips_no_match_for_endpoint(self, mock_popen, daemon): | 378 | def test_pull_engine_d2d_peer_ips_no_match_for_endpoint(self, mock_popen, daemon): |
| @@ -312,12 +393,11 @@ class TestDaemon: | |||
| 312 | ) | 393 | ) |
| 313 | 394 | ||
| 314 | mock_popen.assert_called_once() | 395 | mock_popen.assert_called_once() |
| 315 | - cmd = mock_popen.call_args.args[0] | 396 | + assert _last_launch_context(daemon).d2d_peer_ips == () |
| 316 | - assert '--d2d-peer-ips' not in cmd | ||
| 317 | 397 | ||
| 318 | 398 | ||
| 319 | def test_pull_engine_includes_node_rank(self, mock_popen, daemon): | 399 | def test_pull_engine_includes_node_rank(self, mock_popen, daemon): |
| 320 | - """Test that --node-rank is included in the engine_server CLI with default value""" | 400 | + """Test that the default node rank is passed to the backend.""" |
| 321 | mock_process = MagicMock(pid=12345) | 401 | mock_process = MagicMock(pid=12345) |
| 322 | mock_process.poll.return_value = None | 402 | mock_process.poll.return_value = None |
| 323 | mock_popen.return_value = mock_process | 403 | mock_popen.return_value = mock_process |
| @@ -325,14 +405,11 @@ class TestDaemon: | |||
| 325 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") | 405 | endpoint = Endpoint(id=0, ip="10.0.0.1", business_port="9000", mgmt_port="9090") |
| 326 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") | 406 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id=1, master_dp_ip="192.168.1.100") |
| 327 | 407 | ||
| 328 | - cmd = mock_popen.call_args.args[0] | 408 | + assert _last_launch_context(daemon).node_rank == 0 |
| 329 | - assert "--node-rank" in cmd | ||
| 330 | - node_rank_index = cmd.index("--node-rank") | ||
| 331 | - assert cmd[node_rank_index + 1] == "0" | ||
| 332 | 409 | ||
| 333 | 410 | ||
| 334 | def test_pull_engine_custom_node_rank(self, mock_popen, daemon): | 411 | def test_pull_engine_custom_node_rank(self, mock_popen, daemon): |
| 335 | - """Test that --node-rank value matches the node_rank parameter""" | 412 | + """Test that the configured node rank is passed to the backend.""" |
| 336 | mock_process = MagicMock(pid=12345) | 413 | mock_process = MagicMock(pid=12345) |
| 337 | mock_process.poll.return_value = None | 414 | mock_process.poll.return_value = None |
| 338 | mock_popen.return_value = mock_process | 415 | mock_popen.return_value = mock_process |
| @@ -340,6 +417,4 @@ class TestDaemon: | |||
| 340 | endpoint = Endpoint(id=1, ip="10.0.0.1", business_port="9000", mgmt_port="9090") | 417 | endpoint = Endpoint(id=1, ip="10.0.0.1", business_port="9000", mgmt_port="9090") |
| 341 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id=1, master_dp_ip="192.168.1.100", node_rank=2) | 418 | daemon.pull_engine(PDRole.ROLE_P, [endpoint], instance_id=1, master_dp_ip="192.168.1.100", node_rank=2) |
| 342 | 419 | ||
| 343 | - cmd = mock_popen.call_args.args[0] | 420 | + assert _last_launch_context(daemon).node_rank == 2 |
| 344 | - node_rank_index = cmd.index("--node-rank") | ||
| 345 | - assert cmd[node_rank_index + 1] == "2" | ||
| @@ -175,6 +175,23 @@ class TestEngineManager: | |||
| 175 | assert msg is not None | 175 | assert msg is not None |
| 176 | assert msg.is_master is True | 176 | assert msg.is_master is True |
| 177 | 177 | ||
| 178 | + def test_gen_register_msg_includes_sglang_bootstrap_port(self, engine_manager): | ||
| 179 | + engine_manager._config.basic_config.job_name = "test_job" | ||
| 180 | + engine_manager._config.basic_config.model_name = "test_model" | ||
| 181 | + engine_manager._config.basic_config.role = PDRole.ROLE_P | ||
| 182 | + engine_manager._config.api_config.pod_ip = "192.168.1.101" | ||
| 183 | + engine_manager._config.endpoint_config.service_ports = ["8080"] | ||
| 184 | + engine_manager._config.endpoint_config.mgmt_ports = ["8081"] | ||
| 185 | + engine_manager._config.endpoint_config.bootstrap_port = 9100 | ||
| 186 | + engine_manager._config.api_config.node_manager_port = 8088 | ||
| 187 | + engine_manager._config.basic_config.parallel_config = ParallelConfig(tp_size=2, pp_size=1) | ||
| 188 | + engine_manager._config.basic_config.device_num = 2 | ||
| 189 | + | ||
| 190 | + msg = engine_manager._gen_register_msg() | ||
| 191 | + | ||
| 192 | + assert msg is not None | ||
| 193 | + assert msg.bootstrap_port == 9100 | ||
| 194 | + | ||
| 178 | def test_gen_register_msg_failure(self, engine_manager): | 195 | def test_gen_register_msg_failure(self, engine_manager): |
| 179 | """Test _gen_register_msg with invalid config""" | 196 | """Test _gen_register_msg with invalid config""" |
| 180 | engine_manager._config.basic_config.job_name = None | 197 | engine_manager._config.basic_config.job_name = None |
| @@ -153,41 +153,95 @@ class TestHeartBeatManager: | |||
| 153 | assert heart_beat_manager._endpoints[0].id == 1 | 153 | assert heart_beat_manager._endpoints[0].id == 1 |
| 154 | assert heart_beat_manager._endpoints[1].id == 2 | 154 | assert heart_beat_manager._endpoints[1].id == 2 |
| 155 | 155 | ||
| 156 | - def test_update_endpoint_resets_grace_period(self, heart_beat_manager, sample_start_cmd_msg): | 156 | + def test_update_endpoint_invalidates_inflight_native_probe(self, heart_beat_manager, sample_start_cmd_msg): |
| 157 | - """update_endpoint should restart grace period for engine cold start""" | 157 | + """update_endpoint advances the generation used to reject stale probe results.""" |
| 158 | - heart_beat_manager._thread_started = True | 158 | + before = heart_beat_manager._endpoints_generation |
| 159 | - heart_beat_manager._is_within_grace_period = False | ||
| 160 | - heart_beat_manager._engine_status_thread_start_time = time.time() - 300 | ||
| 161 | - | ||
| 162 | - before = heart_beat_manager._engine_status_thread_start_time | ||
| 163 | heart_beat_manager.update_endpoint(sample_start_cmd_msg) | 159 | heart_beat_manager.update_endpoint(sample_start_cmd_msg) |
| 164 | 160 | ||
| 165 | - assert heart_beat_manager._is_within_grace_period is True | 161 | + assert heart_beat_manager._endpoints_generation == before + 1 |
| 166 | - assert heart_beat_manager._engine_status_thread_start_time > before | ||
| 167 | 162 | ||
| 168 | - @patch('motor.node_manager.core.heartbeat_manager.EngineServerApiClient.query_status') | 163 | + @patch('motor.node_manager.core.heartbeat_manager.Daemon') |
| 169 | - def test_get_engine_server_status_success(self, mock_query_status, heart_beat_manager, sample_endpoints): | 164 | + def test_engine_metrics_targets_exclude_headless(self, mock_daemon, heart_beat_manager): |
| 170 | - """test get engine server status success""" | 165 | + routable = Endpoint(id=1, ip="10.0.0.1", business_port="8001", mgmt_port="9001") |
| 171 | - # Mock query_status to return normal status for each endpoint | 166 | + headless = Endpoint( |
| 172 | - mock_query_status.return_value = {"status": "normal"} | 167 | + id=2, |
| 168 | + ip="10.0.0.2", | ||
| 169 | + business_port="8002", | ||
| 170 | + mgmt_port="9002", | ||
| 171 | + headless=True, | ||
| 172 | + ) | ||
| 173 | + mock_daemon.return_value.get_engine_metrics_target.return_value = "https://10.0.0.1:8001/metrics" | ||
| 174 | + with heart_beat_manager._endpoint_lock: | ||
| 175 | + heart_beat_manager._endpoints = [routable, headless] | ||
| 176 | + | ||
| 177 | + targets = heart_beat_manager.get_engine_metrics_targets() | ||
| 178 | + | ||
| 179 | + assert targets == ["https://10.0.0.1:8001/metrics"] | ||
| 180 | + mock_daemon.return_value.get_engine_metrics_target.assert_called_once_with(routable) | ||
| 181 | + | ||
| 182 | + | ||
| 183 | + def test_refresh_native_engine_status_success(self, mock_daemon, heart_beat_manager, sample_endpoints): | ||
| 184 | + """READY native runtimes map to normal endpoint status.""" | ||
| 185 | + from motor.node_manager.core.services.native_engine.models import RuntimeState | ||
| 186 | + | ||
| 187 | + mock_daemon.return_value.get_engine_runtime_state.return_value = RuntimeState.READY | ||
| 173 | 188 | ||
| 174 | with heart_beat_manager._endpoint_lock: | 189 | with heart_beat_manager._endpoint_lock: |
| 175 | heart_beat_manager._endpoints = sample_endpoints.copy() | 190 | heart_beat_manager._endpoints = sample_endpoints.copy() |
| 176 | 191 | ||
| 177 | - heart_beat_manager._get_engine_server_status() | 192 | + heart_beat_manager._refresh_native_engine_status() |
| 178 | 193 | ||
| 179 | - # Verify that query_status was called for each endpoint | 194 | + assert mock_daemon.return_value.get_engine_runtime_state.call_count == 2 |
| 180 | - assert mock_query_status.call_count == 2 | ||
| 181 | 195 | ||
| 182 | - # Verify that status was updated correctly | ||
| 183 | assert heart_beat_manager._endpoints[0].status == EndpointStatus.NORMAL | 196 | assert heart_beat_manager._endpoints[0].status == EndpointStatus.NORMAL |
| 184 | assert heart_beat_manager._endpoints[1].status == EndpointStatus.NORMAL | 197 | assert heart_beat_manager._endpoints[1].status == EndpointStatus.NORMAL |
| 185 | 198 | ||
| 186 | - @patch('motor.node_manager.core.heartbeat_manager.EngineServerApiClient.query_status') | 199 | + @patch('motor.node_manager.core.heartbeat_manager.Daemon') |
| 187 | - def test_get_engine_server_status_discards_stale_probe_write_back( | 200 | + def test_refresh_native_engine_status_keeps_initial_while_loading(self, mock_daemon, heart_beat_manager): |
| 188 | - self, mock_query_status, heart_beat_manager, sample_start_cmd_msg | 201 | + from motor.node_manager.core.services.native_engine.models import RuntimeState |
| 202 | + | ||
| 203 | + mock_daemon.return_value.get_engine_runtime_state.return_value = RuntimeState.STARTING | ||
| 204 | + endpoint = Endpoint( | ||
| 205 | + id=1, | ||
| 206 | + ip="192.168.1.1", | ||
| 207 | + business_port="8080", | ||
| 208 | + mgmt_port="9090", | ||
| 209 | + status=EndpointStatus.INITIAL, | ||
| 210 | + ) | ||
| 211 | + with heart_beat_manager._endpoint_lock: | ||
| 212 | + heart_beat_manager._endpoints = [endpoint] | ||
| 213 | + | ||
| 214 | + heart_beat_manager._refresh_native_engine_status() | ||
| 215 | + | ||
| 216 | + assert heart_beat_manager._endpoints[0].status == EndpointStatus.INITIAL | ||
| 217 | + | ||
| 218 | + | ||
| 219 | + def test_refresh_headless_process_liveness_reports_wait2start(self, mock_daemon, heart_beat_manager): | ||
| 220 | + from motor.node_manager.core.services.native_engine.models import RuntimeState | ||
| 221 | + | ||
| 222 | + mock_daemon.return_value.get_engine_runtime_state.return_value = RuntimeState.RUNNING | ||
| 223 | + endpoint = Endpoint( | ||
| 224 | + id=1, | ||
| 225 | + ip="192.168.1.2", | ||
| 226 | + business_port="8080", | ||
| 227 | + mgmt_port="9090", | ||
| 228 | + status=EndpointStatus.INITIAL, | ||
| 229 | + headless=True, | ||
| 230 | + ) | ||
| 231 | + with heart_beat_manager._endpoint_lock: | ||
| 232 | + heart_beat_manager._endpoints = [endpoint] | ||
| 233 | + | ||
| 234 | + heart_beat_manager._refresh_native_engine_status() | ||
| 235 | + | ||
| 236 | + assert heart_beat_manager._endpoints[0].status == EndpointStatus.WAIT2START | ||
| 237 | + | ||
| 238 | + | ||
| 239 | + def test_refresh_native_engine_status_discards_stale_probe_write_back( | ||
| 240 | + self, mock_daemon, heart_beat_manager, sample_start_cmd_msg | ||
| 189 | ): | 241 | ): |
| 190 | """stale probe result must not overwrite endpoints refreshed by update_endpoint""" | 242 | """stale probe result must not overwrite endpoints refreshed by update_endpoint""" |
| 243 | + from motor.node_manager.core.services.native_engine.models import RuntimeState | ||
| 244 | + | ||
| 191 | stale_endpoint = Endpoint( | 245 | stale_endpoint = Endpoint( |
| 192 | id=0, | 246 | id=0, |
| 193 | ip="10.0.0.28", | 247 | ip="10.0.0.28", |
| @@ -198,15 +252,15 @@ class TestHeartBeatManager: | |||
| 198 | 252 | ||
| 199 | def probe_and_update_during_probe(*args, **kwargs): | 253 | def probe_and_update_during_probe(*args, **kwargs): |
| 200 | heart_beat_manager.update_endpoint(sample_start_cmd_msg) | 254 | heart_beat_manager.update_endpoint(sample_start_cmd_msg) |
| 201 | - return {"status": "abnormal"} | 255 | + return RuntimeState.UNHEALTHY |
| 202 | 256 | ||
| 203 | - mock_query_status.side_effect = probe_and_update_during_probe | 257 | + mock_daemon.return_value.get_engine_runtime_state.side_effect = probe_and_update_during_probe |
| 204 | 258 | ||
| 205 | with heart_beat_manager._endpoint_lock: | 259 | with heart_beat_manager._endpoint_lock: |
| 206 | heart_beat_manager._endpoints = [stale_endpoint] | 260 | heart_beat_manager._endpoints = [stale_endpoint] |
| 207 | heart_beat_manager._endpoints_generation = 0 | 261 | heart_beat_manager._endpoints_generation = 0 |
| 208 | 262 | ||
| 209 | - heart_beat_manager._get_engine_server_status() | 263 | + heart_beat_manager._refresh_native_engine_status() |
| 210 | 264 | ||
| 211 | assert len(heart_beat_manager._endpoints) == 2 | 265 | assert len(heart_beat_manager._endpoints) == 2 |
| 212 | assert heart_beat_manager._endpoints[0].ip == "192.168.1.1" | 266 | assert heart_beat_manager._endpoints[0].ip == "192.168.1.1" |
| @@ -498,7 +552,6 @@ class TestHeartBeatManager: | |||
| 498 | # Set endpoint info with abnormal status | 552 | # Set endpoint info with abnormal status |
| 499 | heart_beat_manager._job_name = "test_job" | 553 | heart_beat_manager._job_name = "test_job" |
| 500 | heart_beat_manager._instance_id = 1 | 554 | heart_beat_manager._instance_id = 1 |
| 501 | - heart_beat_manager._is_within_grace_period = False # Ensure we're past grace period | ||
| 502 | heart_beat_manager.stop_event.clear() | 555 | heart_beat_manager.stop_event.clear() |
| 503 | 556 | ||
| 504 | with heart_beat_manager._endpoint_lock: | 557 | with heart_beat_manager._endpoint_lock: |
| @@ -535,7 +588,6 @@ class TestHeartBeatManager: | |||
| 535 | 588 | ||
| 536 | heart_beat_manager._job_name = "test_job" | 589 | heart_beat_manager._job_name = "test_job" |
| 537 | heart_beat_manager._instance_id = 1 | 590 | heart_beat_manager._instance_id = 1 |
| 538 | - heart_beat_manager._is_within_grace_period = False | ||
| 539 | heart_beat_manager.stop_event.clear() | 591 | heart_beat_manager.stop_event.clear() |
| 540 | 592 | ||
| 541 | # Start with abnormal status | 593 | # Start with abnormal status |
| @@ -583,7 +635,6 @@ class TestHeartBeatManager: | |||
| 583 | 635 | ||
| 584 | heart_beat_manager._job_name = "test_job" | 636 | heart_beat_manager._job_name = "test_job" |
| 585 | heart_beat_manager._instance_id = 1 | 637 | heart_beat_manager._instance_id = 1 |
| 586 | - heart_beat_manager._is_within_grace_period = False | ||
| 587 | heart_beat_manager.stop_event.clear() | 638 | heart_beat_manager.stop_event.clear() |
| 588 | 639 | ||
| 589 | with heart_beat_manager._endpoint_lock: | 640 | with heart_beat_manager._endpoint_lock: |
| @@ -617,7 +668,6 @@ class TestHeartBeatManager: | |||
| 617 | 668 | ||
| 618 | heart_beat_manager._job_name = "test_job" | 669 | heart_beat_manager._job_name = "test_job" |
| 619 | heart_beat_manager._instance_id = 1 | 670 | heart_beat_manager._instance_id = 1 |
| 620 | - heart_beat_manager._is_within_grace_period = False | ||
| 621 | heart_beat_manager.stop_event.clear() | 671 | heart_beat_manager.stop_event.clear() |
| 622 | 672 | ||
| 623 | # Set multiple endpoints, one abnormal | 673 | # Set multiple endpoints, one abnormal |
| @@ -651,7 +701,6 @@ class TestHeartBeatManager: | |||
| 651 | 701 | ||
| 652 | heart_beat_manager._job_name = "test_job" | 702 | heart_beat_manager._job_name = "test_job" |
| 653 | heart_beat_manager._instance_id = 1 | 703 | heart_beat_manager._instance_id = 1 |
| 654 | - heart_beat_manager._is_within_grace_period = False | ||
| 655 | heart_beat_manager.stop_event.clear() | 704 | heart_beat_manager.stop_event.clear() |
| 656 | 705 | ||
| 657 | with heart_beat_manager._endpoint_lock: | 706 | with heart_beat_manager._endpoint_lock: |
| @@ -774,20 +823,22 @@ class TestHeartBeatManager: | |||
| 774 | mock_report_heartbeat.assert_called_once() | 823 | mock_report_heartbeat.assert_called_once() |
| 775 | 824 | ||
| 776 | 825 | ||
| 777 | - @patch("motor.node_manager.core.heartbeat_manager.EngineServerApiClient.query_status") | 826 | + @patch("motor.node_manager.core.heartbeat_manager.Daemon") |
| 778 | - def test_get_engine_server_status_keeps_status_before_start_after_restore( | 827 | + def test_refresh_native_engine_status_keeps_status_before_start_after_restore( |
| 779 | - self, mock_query_status, _mock_restored, heart_beat_manager, sample_endpoints | 828 | + self, mock_daemon, _mock_restored, heart_beat_manager, sample_endpoints |
| 780 | ): | 829 | ): |
| 781 | - mock_query_status.return_value = {"status": "abnormal"} | 830 | + from motor.node_manager.core.services.native_engine.models import RuntimeState |
| 831 | + | ||
| 832 | + mock_daemon.return_value.get_engine_runtime_state.return_value = RuntimeState.UNHEALTHY | ||
| 782 | 833 | ||
| 783 | with heart_beat_manager._endpoint_lock: | 834 | with heart_beat_manager._endpoint_lock: |
| 784 | heart_beat_manager._endpoints = sample_endpoints.copy() | 835 | heart_beat_manager._endpoints = sample_endpoints.copy() |
| 785 | 836 | ||
| 786 | - heart_beat_manager._get_engine_server_status() | 837 | + heart_beat_manager._refresh_native_engine_status() |
| 787 | 838 | ||
| 788 | assert heart_beat_manager._endpoints[0].status == EndpointStatus.NORMAL | 839 | assert heart_beat_manager._endpoints[0].status == EndpointStatus.NORMAL |
| 789 | assert heart_beat_manager._endpoints[1].status == EndpointStatus.NORMAL | 840 | assert heart_beat_manager._endpoints[1].status == EndpointStatus.NORMAL |
| 790 | - assert mock_query_status.call_count == 2 | 841 | + assert mock_daemon.return_value.get_engine_runtime_state.call_count == 2 |
| 791 | 842 | ||
| 792 | 843 | ||
| 793 | 844 | ||
| @@ -169,6 +169,35 @@ def test_logging_config_defaults(nm_config_data): | |||
| 169 | assert config.logging_config.log_date_format == '%m-%d %H:%M:%S' | 169 | assert config.logging_config.log_date_format == '%m-%d %H:%M:%S' |
| 170 | 170 | ||
| 171 | 171 | ||
| 172 | + | ||
| 173 | +def test_native_sglang_bootstrap_port_from_engine_config(): | ||
| 174 | + config = create_config_object() | ||
| 175 | + raw = { | ||
| 176 | + "motor_engine_prefill_config": { | ||
| 177 | + "engine_type": "sglang", | ||
| 178 | + "engine_config": {"disaggregation_bootstrap_port": 9100}, | ||
| 179 | + } | ||
| 180 | + } | ||
| 181 | + | ||
| 182 | + NodeManagerConfig._set_native_bootstrap_port(config, raw) | ||
| 183 | + | ||
| 184 | + assert config.endpoint_config.bootstrap_port == 9100 | ||
| 185 | + | ||
| 186 | + | ||
| 187 | + | ||
| 188 | +def test_native_bootstrap_port_rejects_invalid_range(): | ||
| 189 | + config = create_config_object() | ||
| 190 | + raw = { | ||
| 191 | + "motor_engine_prefill_config": { | ||
| 192 | + "engine_type": "sglang", | ||
| 193 | + "engine_config": {"disaggregation_bootstrap_port": 70000}, | ||
| 194 | + } | ||
| 195 | + } | ||
| 196 | + | ||
| 197 | + with pytest.raises(ValueError, match="disaggregation_bootstrap_port must be in range 1-65535"): | ||
| 198 | + NodeManagerConfig._set_native_bootstrap_port(config, raw) | ||
| 199 | + | ||
| 200 | + | ||
| 172 | 201 | ||
| 173 | "invalid_config,expected_error", | 202 | "invalid_config,expected_error", |
| 174 | [ | 203 | [ |
| @@ -542,6 +571,14 @@ def test_from_json_loads_union_config_for_hybrid(): | |||
| 542 | assert config.endpoint_config.endpoint_num == 2 | 571 | assert config.endpoint_config.endpoint_num == 2 |
| 543 | 572 | ||
| 544 | 573 | ||
| 574 | +def test_native_runtime_rejects_snapshot_configuration(): | ||
| 575 | + config = NodeManagerConfig() | ||
| 576 | + config.snapshot_config.enable_snapshot = True | ||
| 577 | + | ||
| 578 | + with pytest.raises(ValueError, match="Native engine runtime does not support snapshot yet"): | ||
| 579 | + config.validate_config() | ||
| 580 | + | ||
| 581 | + | ||
| 545 | def _single_container_hybrid_user_config(): | 582 | def _single_container_hybrid_user_config(): |
| 546 | return { | 583 | return { |
| 547 | "motor_deploy_config": { | 584 | "motor_deploy_config": { |
| @@ -654,9 +691,7 @@ def test_vllm_multi_connector_infers_transport_connector_capability(): | |||
| 654 | assert NodeManagerConfig._infer_dispatch_capabilities(_multi_connector("NixlConnector")) == [ | 691 | assert NodeManagerConfig._infer_dispatch_capabilities(_multi_connector("NixlConnector")) == [ |
| 655 | DispatchPlan.PREFILL_HANDOFF_DECODE.value | 692 | DispatchPlan.PREFILL_HANDOFF_DECODE.value |
| 656 | ] | 693 | ] |
| 657 | - assert NodeManagerConfig._infer_dispatch_capabilities(_multi_connector("MooncakeLayerwiseConnector")) == [ | 694 | + assert NodeManagerConfig._infer_dispatch_capabilities(_multi_connector("MooncakeLayerwiseConnector")) == [] |
| 658 | - DispatchPlan.CONCURRENT_ENGINE_SYNC.value | ||
| 659 | - ] | ||
| 660 | 695 | ||
| 661 | 696 | ||
| 662 | def test_vllm_multi_connector_ignores_non_transport_connector_profiles(): | 697 | def test_vllm_multi_connector_ignores_non_transport_connector_profiles(): |
| @@ -676,7 +711,7 @@ def test_vllm_multi_connector_ignores_non_transport_connector_profiles(): | |||
| 676 | }, | 711 | }, |
| 677 | } | 712 | } |
| 678 | 713 | ||
| 679 | - assert NodeManagerConfig._infer_dispatch_capabilities(engine_config) == [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | 714 | + assert NodeManagerConfig._infer_dispatch_capabilities(engine_config) == [] |
| 680 | 715 | ||
| 681 | 716 | ||
| 682 | def test_vllm_multi_connector_requires_transport_and_store_connectors(): | 717 | def test_vllm_multi_connector_requires_transport_and_store_connectors(): |
| @@ -742,7 +777,7 @@ def test_user_dispatch_capabilities_cannot_enable_unknown_connector(): | |||
| 742 | assert config_data["basic_config"]["dispatch_capabilities"] == [] | 777 | assert config_data["basic_config"]["dispatch_capabilities"] == [] |
| 743 | 778 | ||
| 744 | 779 | ||
| 745 | -def test_vllm_layerwise_connector_infers_concurrent_capability(): | 780 | +def test_vllm_layerwise_connector_does_not_advertise_unsupported_native_capability(): |
| 746 | capabilities = NodeManagerConfig._infer_dispatch_capabilities( | 781 | capabilities = NodeManagerConfig._infer_dispatch_capabilities( |
| 747 | { | 782 | { |
| 748 | "engine_type": "vllm", | 783 | "engine_type": "vllm", |
| @@ -754,7 +789,7 @@ def test_vllm_layerwise_connector_infers_concurrent_capability(): | |||
| 754 | } | 789 | } |
| 755 | ) | 790 | ) |
| 756 | 791 | ||
| 757 | - assert capabilities == [DispatchPlan.CONCURRENT_ENGINE_SYNC.value] | 792 | + assert capabilities == [] |
| 758 | 793 | ||
| 759 | 794 | ||
| 760 | def test_sglang_infers_concurrent_capability(): | 795 | def test_sglang_infers_concurrent_capability(): |
大幅度改动，要刷新skill，建议使用motor-dev的skill进行刷新