已合并
[fix] 修复 SGLang 单端口多 DP 时 KV 亲和选不到已缓存实例的问题 #1036
chaofengdev创建于 22 天前
[fix] 修复 SGLang 单端口多 DP 时 KV 亲和选不到已缓存实例的问题 #1036
已合并
共 5 个文件变更+459-66
| @@ -51,11 +51,24 @@ def decode_query_response_msgpack(payload: bytes) -> dict[str, Any]: | |||
| 51 | return msgspec.msgpack.decode(payload) | 51 | return msgspec.msgpack.decode(payload) |
| 52 | 52 | ||
| 53 | 53 | ||
| 54 | +def _conductor_engine_label(instance: Instance) -> str: | ||
| 55 | + """Return the conductor instance-id engine prefix. | ||
| 56 | + | ||
| 57 | + Only SGLang uses a distinct label. vLLM, missing, and unknown types keep | ||
| 58 | + the historical ``vllm`` prefix so existing registrations stay unchanged. | ||
| 59 | + """ | ||
| 60 | + raw = getattr(instance, "engine_type", None) | ||
| 61 | + if isinstance(raw, str) and raw.strip().lower() == "sglang": | ||
| 62 | + return "sglang" | ||
| 63 | + return "vllm" | ||
| 64 | + | ||
| 65 | + | ||
| 54 | def conductor_instance_id(instance: Instance) -> str: | 66 | def conductor_instance_id(instance: Instance) -> str: |
| 55 | """Return the Conductor tenant key for a KVA-eligible instance.""" | 67 | """Return the Conductor tenant key for a KVA-eligible instance.""" |
| 68 | + engine = _conductor_engine_label(instance) | ||
| 56 | if instance.role == PDRole.ROLE_U: | 69 | if instance.role == PDRole.ROLE_U: |
| 57 | - return f"vllm-union-{instance.id}" | 70 | + return f"{engine}-union-{instance.id}" |
| 58 | - return f"vllm-prefill-{instance.id}" | 71 | + return f"{engine}-prefill-{instance.id}" |
| 59 | 72 | ||
| 60 | 73 | ||
| 61 | class ConductorApiClient: | 74 | class ConductorApiClient: |
| @@ -114,10 +127,10 @@ class ConductorApiClient: | |||
| 114 | for instance in instances: | 127 | for instance in instances: |
| 115 | if instance.role not in _KVA_ROLES: | 128 | if instance.role not in _KVA_ROLES: |
| 116 | continue | 129 | continue |
| 117 | - for ep in instance.get_all_endpoints(): | 130 | + for ep, dp_rank in cls._hbm_event_targets(instance): |
| 118 | - cls._register_hbm_dp(reg, sb, instance, ep) | 131 | + cls._register_hbm_dp(reg, sb, instance, ep, dp_rank) |
| 119 | if mode == "per_node": | 132 | if mode == "per_node": |
| 120 | - cls._register_yuanrong_node(reg, sb, instance, ep) | 133 | + cls._register_yuanrong_node(reg, sb, instance, ep, dp_rank) |
| 121 | 134 | ||
| 122 | @classmethod | 135 | @classmethod |
| 123 | def unregister_kv_instance(cls, instances: list[Instance]) -> None: | 136 | def unregister_kv_instance(cls, instances: list[Instance]) -> None: |
| @@ -132,10 +145,10 @@ class ConductorApiClient: | |||
| 132 | for instance in instances: | 145 | for instance in instances: |
| 133 | if instance.role not in _KVA_ROLES: | 146 | if instance.role not in _KVA_ROLES: |
| 134 | continue | 147 | continue |
| 135 | - for ep in instance.get_all_endpoints(): | 148 | + for ep, dp_rank in cls._hbm_event_targets(instance): |
| 136 | - cls.unregister_post(instance, ep) | 149 | + cls.unregister_post(instance, ep, dp_rank) |
| 137 | if mode == "per_node": | 150 | if mode == "per_node": |
| 138 | - cls._unregister_yuanrong_node(cls._kv_reg(), sb, instance, ep) | 151 | + cls._unregister_yuanrong_node(cls._kv_reg(), sb, instance, ep, dp_rank) |
| 139 | 152 | ||
| 140 | # ── Pool registration (Mooncake / Memcache) ────────────────────── | 153 | # ── Pool registration (Mooncake / Memcache) ────────────────────── |
| 141 | 154 | ||
| @@ -195,23 +208,31 @@ class ConductorApiClient: | |||
| 195 | # ── HBM per-DP (Mooncake / Memcache / YuanRong) ───────────────────────────── | 208 | # ── HBM per-DP (Mooncake / Memcache / YuanRong) ───────────────────────────── |
| 196 | 209 | ||
| 197 | @classmethod | 210 | @classmethod |
| 198 | - def _register_hbm_dp(cls, reg, store_backend: str, instance: "Instance", endpoint: "Endpoint") -> None: | 211 | + def _register_hbm_dp( |
| 212 | + cls, | ||
| 213 | + reg, | ||
| 214 | + store_backend: str, | ||
| 215 | + instance: "Instance", | ||
| 216 | + endpoint: "Endpoint", | ||
| 217 | + dp_rank: int | None = None, | ||
| 218 | + ) -> None: | ||
| 199 | """Register a single DP's NPU/HBM endpoint (all backends).""" | 219 | """Register a single DP's NPU/HBM endpoint (all backends).""" |
| 220 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 200 | instance_id = conductor_instance_id(instance) | 221 | instance_id = conductor_instance_id(instance) |
| 201 | npu_url = cls._resolve_endpoint_url( | 222 | npu_url = cls._resolve_endpoint_url( |
| 202 | reg.npu_endpoint or reg.xpu_endpoint or reg.endpoint, | 223 | reg.npu_endpoint or reg.xpu_endpoint or reg.endpoint, |
| 203 | endpoint.ip, | 224 | endpoint.ip, |
| 204 | - endpoint.id, | 225 | + rank, |
| 205 | ) | 226 | ) |
| 206 | 227 | ||
| 207 | - replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, endpoint.id) | 228 | + replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, rank) |
| 208 | register_data: dict = { | 229 | register_data: dict = { |
| 209 | "instance_id": instance_id, | 230 | "instance_id": instance_id, |
| 210 | "type": reg.engine_type, | 231 | "type": reg.engine_type, |
| 211 | "store_backend": store_backend, | 232 | "store_backend": store_backend, |
| 212 | "modelname": instance.model_name, | 233 | "modelname": instance.model_name, |
| 213 | "block_size": reg.block_size, | 234 | "block_size": reg.block_size, |
| 214 | - "dp_rank": endpoint.id, | 235 | + "dp_rank": rank, |
| 215 | } | 236 | } |
| 216 | if npu_url: | 237 | if npu_url: |
| 217 | register_data["medium_endpoints"] = {"npu": npu_url} | 238 | register_data["medium_endpoints"] = {"npu": npu_url} |
| @@ -230,11 +251,11 @@ class ConductorApiClient: | |||
| 230 | mode, | 251 | mode, |
| 231 | store_backend, | 252 | store_backend, |
| 232 | instance_id, | 253 | instance_id, |
| 233 | - endpoint.id, | 254 | + rank, |
| 234 | replay_url or "none", | 255 | replay_url or "none", |
| 235 | ) | 256 | ) |
| 236 | except Exception as e: | 257 | except Exception as e: |
| 237 | - logger.error("HBM DP registration failed for %s dp=%d: %s", instance_id, endpoint.id, e) | 258 | + logger.error("HBM DP registration failed for %s dp=%d: %s", instance_id, rank, e) |
| 238 | 259 | ||
| 239 | # ── YuanRong: CPU/Disk per-node ─────────────────────────────────── | 260 | # ── YuanRong: CPU/Disk per-node ─────────────────────────────────── |
| 240 | 261 | ||
| @@ -261,7 +282,14 @@ class ConductorApiClient: | |||
| 261 | return ip, model_name | 282 | return ip, model_name |
| 262 | 283 | ||
| 263 | @classmethod | 284 | @classmethod |
| 264 | - def _register_yuanrong_node(cls, reg, store_backend: str, instance: "Instance", endpoint: "Endpoint") -> None: | 285 | + def _register_yuanrong_node( |
| 286 | + cls, | ||
| 287 | + reg, | ||
| 288 | + store_backend: str, | ||
| 289 | + instance: "Instance", | ||
| 290 | + endpoint: "Endpoint", | ||
| 291 | + dp_rank: int | None = None, | ||
| 292 | + ) -> None: | ||
| 265 | """Register YuanRong CPU/Disk PUB once per (node IP, model). | 293 | """Register YuanRong CPU/Disk PUB once per (node IP, model). |
| 266 | 294 | ||
| 267 | Shared by every DP on the node; the node pool is registered by whichever | 295 | Shared by every DP on the node; the node pool is registered by whichever |
| @@ -269,10 +297,11 @@ class ConductorApiClient: | |||
| 269 | registration is tracked by :attr:`_yuanrong_nodes_registered`, which | 297 | registration is tracked by :attr:`_yuanrong_nodes_registered`, which |
| 270 | :meth:`re_register_kv_instances` re-syncs against ``GET /workers``. | 298 | :meth:`re_register_kv_instances` re-syncs against ``GET /workers``. |
| 271 | """ | 299 | """ |
| 300 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 272 | node_key = cls._node_pool_key(endpoint.ip, instance.model_name) | 301 | node_key = cls._node_pool_key(endpoint.ip, instance.model_name) |
| 273 | # Record the reference before the dedup check so that every DP on the | 302 | # Record the reference before the dedup check so that every DP on the |
| 274 | # node is accounted for, including the ones that skip the POST below. | 303 | # node is accounted for, including the ones that skip the POST below. |
| 275 | - cls._yuanrong_node_refs.setdefault(node_key, set()).add((conductor_instance_id(instance), endpoint.id)) | 304 | + cls._yuanrong_node_refs.setdefault(node_key, set()).add((conductor_instance_id(instance), rank)) |
| 276 | if node_key in cls._yuanrong_nodes_registered: | 305 | if node_key in cls._yuanrong_nodes_registered: |
| 277 | return | 306 | return |
| 278 | 307 | ||
| @@ -318,7 +347,14 @@ class ConductorApiClient: | |||
| 318 | logger.error("YuanRong node pool registration failed for ip=%s: %s", endpoint.ip, e) | 347 | logger.error("YuanRong node pool registration failed for ip=%s: %s", endpoint.ip, e) |
| 319 | 348 | ||
| 320 | @classmethod | 349 | @classmethod |
| 321 | - def _unregister_yuanrong_node(cls, reg, store_backend: str, instance: "Instance", endpoint: "Endpoint") -> None: | 350 | + def _unregister_yuanrong_node( |
| 351 | + cls, | ||
| 352 | + reg, | ||
| 353 | + store_backend: str, | ||
| 354 | + instance: "Instance", | ||
| 355 | + endpoint: "Endpoint", | ||
| 356 | + dp_rank: int | None = None, | ||
| 357 | + ) -> None: | ||
| 322 | """Drop one DP's reference to the node pool, unregistering it when it is the last. | 358 | """Drop one DP's reference to the node pool, unregistering it when it is the last. |
| 323 | 359 | ||
| 324 | The CPU/Disk PUB serves the whole node, so a single DP leaving must not | 360 | The CPU/Disk PUB serves the whole node, so a single DP leaving must not |
| @@ -326,11 +362,12 @@ class ConductorApiClient: | |||
| 326 | instance and its ZMQ SUB forever (reconnecting against a PUB that no | 362 | instance and its ZMQ SUB forever (reconnecting against a PUB that no |
| 327 | longer exists), and ``GET /workers`` keeps advertising a dead pool. | 363 | longer exists), and ``GET /workers`` keeps advertising a dead pool. |
| 328 | """ | 364 | """ |
| 365 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 329 | node_key = cls._node_pool_key(endpoint.ip, instance.model_name) | 366 | node_key = cls._node_pool_key(endpoint.ip, instance.model_name) |
| 330 | refs = cls._yuanrong_node_refs.get(node_key) | 367 | refs = cls._yuanrong_node_refs.get(node_key) |
| 331 | if refs is None: | 368 | if refs is None: |
| 332 | return | 369 | return |
| 333 | - refs.discard((conductor_instance_id(instance), endpoint.id)) | 370 | + refs.discard((conductor_instance_id(instance), rank)) |
| 334 | if refs: | 371 | if refs: |
| 335 | return | 372 | return |
| 336 | cls._yuanrong_node_refs.pop(node_key, None) | 373 | cls._yuanrong_node_refs.pop(node_key, None) |
| @@ -367,6 +404,43 @@ class ConductorApiClient: | |||
| 367 | 404 | ||
| 368 | # ── Shared helpers ──────────────────────────────────────────────── | 405 | # ── Shared helpers ──────────────────────────────────────────────── |
| 369 | 406 | ||
| 407 | + @staticmethod | ||
| 408 | + def _instance_dp_size(instance: "Instance") -> int: | ||
| 409 | + """Return engine ``dp_size`` for KV-event expansion; 1 when unknown. | ||
| 410 | + | ||
| 411 | + ``parallel_config`` on mocks may be a MagicMock; reject non-ints so | ||
| 412 | + existing tests stay on a single HTTP endpoint mapping. | ||
| 413 | + """ | ||
| 414 | + pc = getattr(instance, "parallel_config", None) | ||
| 415 | + if pc is None: | ||
| 416 | + return 1 | ||
| 417 | + raw = getattr(pc, "dp_size", 1) | ||
| 418 | + if isinstance(raw, bool) or not isinstance(raw, int): | ||
| 419 | + try: | ||
| 420 | + raw = int(raw) | ||
| 421 | + except (TypeError, ValueError): | ||
| 422 | + return 1 | ||
| 423 | + return raw if raw > 1 else 1 | ||
| 424 | + | ||
| 425 | + @classmethod | ||
| 426 | + def _hbm_event_targets(cls, instance: "Instance") -> list[tuple["Endpoint", int]]: | ||
| 427 | + """Return ``(endpoint, dp_rank)`` pairs the conductor should subscribe. | ||
| 428 | + | ||
| 429 | + HTTP routing endpoints and engine KV-event DPs are not always 1:1. | ||
| 430 | + SGLang with ``enable_multi_endpoints=false`` exposes one HTTP port but | ||
| 431 | + publishes one ZMQ port per ``dp-size`` (base + rank). Expand only | ||
| 432 | + when there is a single HTTP so it matches affinity fold | ||
| 433 | + (``len(endpoints) == 1``). Mixed instances (1 < HTTP < dp_size) | ||
| 434 | + stay 1:1 on the existing endpoints and do not invent extra ranks. | ||
| 435 | + """ | ||
| 436 | + endpoints = list(instance.get_all_endpoints()) | ||
| 437 | + if not endpoints: | ||
| 438 | + return [] | ||
| 439 | + dp_size = cls._instance_dp_size(instance) | ||
| 440 | + if len(endpoints) == 1 and dp_size > 1: | ||
| 441 | + return [(endpoints[0], rank) for rank in range(dp_size)] | ||
| 442 | + return [(ep, ep.id) for ep in endpoints] | ||
| 443 | + | ||
| 370 | @staticmethod | 444 | @staticmethod |
| 371 | def _resolve_endpoint_url(pattern: str, ip: str, dp_rank: int) -> str | None: | 445 | def _resolve_endpoint_url(pattern: str, ip: str, dp_rank: int) -> str | None: |
| 372 | """Resolve an endpoint pattern like 'tcp://*:5557' with the given IP and dp_rank offset.""" | 446 | """Resolve an endpoint pattern like 'tcp://*:5557' with the given IP and dp_rank offset.""" |
| @@ -380,35 +454,45 @@ class ConductorApiClient: | |||
| 380 | 454 | ||
| 381 | @classmethod | 455 | @classmethod |
| 382 | def _build_medium_endpoints(cls, config, ip: str, dp_rank: int) -> dict[str, str]: | 456 | def _build_medium_endpoints(cls, config, ip: str, dp_rank: int) -> dict[str, str]: |
| 383 | - """Build the medium_endpoints map from per-medium endpoint patterns.""" | 457 | + """Build the medium_endpoints map from per-medium endpoint patterns. |
| 458 | + | ||
| 459 | + The generic ``endpoint`` / ``npu_endpoint`` pattern is the engine ZMQ | ||
| 460 | + publisher (vLLM/SGLang). Copying it onto cpu/disk makes kv-conductor | ||
| 461 | + treat the stream as a pool source and drop engine BlockStored events. | ||
| 462 | + CPU/Disk are included only when those patterns are set explicitly. | ||
| 463 | + """ | ||
| 384 | npu_url = cls._resolve_endpoint_url(config.npu_endpoint or config.xpu_endpoint, ip, dp_rank) | 464 | npu_url = cls._resolve_endpoint_url(config.npu_endpoint or config.xpu_endpoint, ip, dp_rank) |
| 385 | cpu_url = cls._resolve_endpoint_url(config.cpu_endpoint, ip, dp_rank) | 465 | cpu_url = cls._resolve_endpoint_url(config.cpu_endpoint, ip, dp_rank) |
| 386 | disk_url = cls._resolve_endpoint_url(config.disk_endpoint, ip, dp_rank) | 466 | disk_url = cls._resolve_endpoint_url(config.disk_endpoint, ip, dp_rank) |
| 387 | fallback = cls._resolve_endpoint_url(config.endpoint, ip, dp_rank) | 467 | fallback = cls._resolve_endpoint_url(config.endpoint, ip, dp_rank) |
| 388 | - return { | 468 | + endpoints: dict[str, str] = {} |
| 389 | - "npu": npu_url or fallback or "", | 469 | + if npu_url or fallback: |
| 390 | - "cpu": cpu_url or fallback or "", | 470 | + endpoints["npu"] = npu_url or fallback or "" |
| 391 | - "disk": disk_url or fallback or "", | 471 | + if cpu_url: |
| 392 | - } | 472 | + endpoints["cpu"] = cpu_url |
| 473 | + if disk_url: | ||
| 474 | + endpoints["disk"] = disk_url | ||
| 475 | + return endpoints | ||
| 393 | 476 | ||
| 394 | @classmethod | 477 | @classmethod |
| 395 | - def register_post(cls, instance: "Instance", endpoint: "Endpoint") -> None: | 478 | + def register_post(cls, instance: "Instance", endpoint: "Endpoint", dp_rank: int | None = None) -> None: |
| 396 | """Legacy single-DP registration (used by re-registration path).""" | 479 | """Legacy single-DP registration (used by re-registration path).""" |
| 480 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 397 | reg = cls._kv_reg() | 481 | reg = cls._kv_reg() |
| 398 | instance_id = conductor_instance_id(instance) | 482 | instance_id = conductor_instance_id(instance) |
| 399 | sb = cls._resolve_store_backend() | 483 | sb = cls._resolve_store_backend() |
| 400 | 484 | ||
| 401 | if cls._resolve_backend_mode() == "per_node": | 485 | if cls._resolve_backend_mode() == "per_node": |
| 402 | - cls._register_hbm_dp(reg, sb, instance, endpoint) | 486 | + cls._register_hbm_dp(reg, sb, instance, endpoint, rank) |
| 403 | - cls._register_yuanrong_node(reg, sb, instance, endpoint) | 487 | + cls._register_yuanrong_node(reg, sb, instance, endpoint, rank) |
| 404 | return | 488 | return |
| 405 | 489 | ||
| 406 | - medium_endpoints = cls._build_medium_endpoints(reg, endpoint.ip, endpoint.id) | 490 | + medium_endpoints = cls._build_medium_endpoints(reg, endpoint.ip, rank) |
| 407 | if all(v == "" for v in medium_endpoints.values()): | 491 | if all(v == "" for v in medium_endpoints.values()): |
| 408 | logger.debug("no endpoint configured for kv events, skipping registration") | 492 | logger.debug("no endpoint configured for kv events, skipping registration") |
| 409 | return | 493 | return |
| 410 | 494 | ||
| 411 | - replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, endpoint.id) | 495 | + replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, rank) |
| 412 | register_data: dict = { | 496 | register_data: dict = { |
| 413 | "medium_endpoints": medium_endpoints, | 497 | "medium_endpoints": medium_endpoints, |
| 414 | "type": reg.engine_type, | 498 | "type": reg.engine_type, |
| @@ -416,7 +500,7 @@ class ConductorApiClient: | |||
| 416 | "modelname": instance.model_name, | 500 | "modelname": instance.model_name, |
| 417 | "block_size": reg.block_size, | 501 | "block_size": reg.block_size, |
| 418 | "instance_id": instance_id, | 502 | "instance_id": instance_id, |
| 419 | - "dp_rank": endpoint.id, | 503 | + "dp_rank": rank, |
| 420 | } | 504 | } |
| 421 | if TENANT_ID != "default": | 505 | if TENANT_ID != "default": |
| 422 | register_data["tenant_id"] = TENANT_ID | 506 | register_data["tenant_id"] = TENANT_ID |
| @@ -435,12 +519,13 @@ class ConductorApiClient: | |||
| 435 | logger.info("register_data : %s", register_data) | 519 | logger.info("register_data : %s", register_data) |
| 436 | 520 | ||
| 437 | @classmethod | 521 | @classmethod |
| 438 | - def unregister_post(cls, instance: Instance, endpoint: Endpoint) -> None: | 522 | + def unregister_post(cls, instance: Instance, endpoint: Endpoint, dp_rank: int | None = None) -> None: |
| 439 | """ | 523 | """ |
| 440 | unregister_kv_instance. | 524 | unregister_kv_instance. |
| 441 | 525 | ||
| 442 | :returns: | 526 | :returns: |
| 443 | """ | 527 | """ |
| 528 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 444 | reg = cls._kv_reg() | 529 | reg = cls._kv_reg() |
| 445 | instance_id = conductor_instance_id(instance) | 530 | instance_id = conductor_instance_id(instance) |
| 446 | register_data: dict = { | 531 | register_data: dict = { |
| @@ -448,7 +533,7 @@ class ConductorApiClient: | |||
| 448 | "modelname": instance.model_name, | 533 | "modelname": instance.model_name, |
| 449 | "block_size": reg.block_size, | 534 | "block_size": reg.block_size, |
| 450 | "instance_id": instance_id, | 535 | "instance_id": instance_id, |
| 451 | - "dp_rank": endpoint.id, | 536 | + "dp_rank": rank, |
| 452 | } | 537 | } |
| 453 | if TENANT_ID != "default": | 538 | if TENANT_ID != "default": |
| 454 | register_data["tenant_id"] = TENANT_ID | 539 | register_data["tenant_id"] = TENANT_ID |
| @@ -587,17 +672,20 @@ class ConductorApiClient: | |||
| 587 | return {} | 672 | return {} |
| 588 | 673 | ||
| 589 | @classmethod | 674 | @classmethod |
| 590 | - def _build_register_payload(cls, instance: Instance, endpoint: Endpoint) -> dict[str, Any]: | 675 | + def _build_register_payload( |
| 676 | + cls, instance: Instance, endpoint: Endpoint, dp_rank: int | None = None | ||
| 677 | + ) -> dict[str, Any]: | ||
| 591 | """Build registration payload using the unified kv_conductor_config config. | 678 | """Build registration payload using the unified kv_conductor_config config. |
| 592 | 679 | ||
| 593 | Produces the same payload format as :meth:`register_post` so the | 680 | Produces the same payload format as :meth:`register_post` so the |
| 594 | re-registration comparison is consistent. | 681 | re-registration comparison is consistent. |
| 595 | """ | 682 | """ |
| 683 | + rank = endpoint.id if dp_rank is None else dp_rank | ||
| 596 | reg = cls._kv_reg() | 684 | reg = cls._kv_reg() |
| 597 | instance_id = conductor_instance_id(instance) | 685 | instance_id = conductor_instance_id(instance) |
| 598 | sb = cls._resolve_store_backend() | 686 | sb = cls._resolve_store_backend() |
| 599 | 687 | ||
| 600 | - medium_endpoints = cls._build_medium_endpoints(reg, endpoint.ip, endpoint.id) | 688 | + medium_endpoints = cls._build_medium_endpoints(reg, endpoint.ip, rank) |
| 601 | # Keep only non-empty endpoints. YuanRong CPU/Disk are per-node, not per-DP. | 689 | # Keep only non-empty endpoints. YuanRong CPU/Disk are per-node, not per-DP. |
| 602 | filtered = {k: v for k, v in medium_endpoints.items() if v} | 690 | filtered = {k: v for k, v in medium_endpoints.items() if v} |
| 603 | if cls._resolve_backend_mode() == "per_node": | 691 | if cls._resolve_backend_mode() == "per_node": |
| @@ -605,7 +693,7 @@ class ConductorApiClient: | |||
| 605 | if not filtered: | 693 | if not filtered: |
| 606 | return {} | 694 | return {} |
| 607 | 695 | ||
| 608 | - replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, endpoint.id) | 696 | + replay_url = cls._resolve_endpoint_url(reg.replay_endpoint, endpoint.ip, rank) |
| 609 | payload: dict[str, Any] = { | 697 | payload: dict[str, Any] = { |
| 610 | "medium_endpoints": filtered, | 698 | "medium_endpoints": filtered, |
| 611 | "type": reg.engine_type, | 699 | "type": reg.engine_type, |
| @@ -613,7 +701,7 @@ class ConductorApiClient: | |||
| 613 | "modelname": instance.model_name, | 701 | "modelname": instance.model_name, |
| 614 | "block_size": reg.block_size, | 702 | "block_size": reg.block_size, |
| 615 | "instance_id": instance_id, | 703 | "instance_id": instance_id, |
| 616 | - "dp_rank": endpoint.id, | 704 | + "dp_rank": rank, |
| 617 | } | 705 | } |
| 618 | if TENANT_ID != "default": | 706 | if TENANT_ID != "default": |
| 619 | payload["tenant_id"] = TENANT_ID | 707 | payload["tenant_id"] = TENANT_ID |
| @@ -752,18 +840,19 @@ class ConductorApiClient: | |||
| 752 | for instance in instances: | 840 | for instance in instances: |
| 753 | if instance.role not in _KVA_ROLES: | 841 | if instance.role not in _KVA_ROLES: |
| 754 | continue | 842 | continue |
| 755 | - for ep in instance.get_all_endpoints(): | 843 | + for ep, dp_rank in cls._hbm_event_targets(instance): |
| 756 | - payload = cls._build_register_payload(instance, ep) | 844 | + payload = cls._build_register_payload(instance, ep, dp_rank) |
| 757 | if not payload: | 845 | if not payload: |
| 758 | logger.debug( | 846 | logger.debug( |
| 759 | - "skip re-register because payload build failed for instance=%s endpoint=%s", | 847 | + "skip re-register because payload build failed for instance=%s endpoint=%s dp=%s", |
| 760 | instance.id, | 848 | instance.id, |
| 761 | ep.id, | 849 | ep.id, |
| 850 | + dp_rank, | ||
| 762 | ) | 851 | ) |
| 763 | continue | 852 | continue |
| 764 | 853 | ||
| 765 | instance_id = conductor_instance_id(instance) | 854 | instance_id = conductor_instance_id(instance) |
| 766 | - dp_registered = (instance_id, ep.id) in registered_dps | 855 | + dp_registered = (instance_id, dp_rank) in registered_dps |
| 767 | # The cache was just reconciled with the conductor, so a missing | 856 | # The cache was just reconciled with the conductor, so a missing |
| 768 | # key here really means the pool is gone. | 857 | # key here really means the pool is gone. |
| 769 | node_key = cls._node_pool_key(ep.ip, instance.model_name) | 858 | node_key = cls._node_pool_key(ep.ip, instance.model_name) |
| @@ -780,12 +869,12 @@ class ConductorApiClient: | |||
| 780 | ep.ip, | 869 | ep.ip, |
| 781 | instance.model_name, | 870 | instance.model_name, |
| 782 | ) | 871 | ) |
| 783 | - cls._register_yuanrong_node(cls._kv_reg(), sb, instance, ep) | 872 | + cls._register_yuanrong_node(cls._kv_reg(), sb, instance, ep, dp_rank) |
| 784 | continue | 873 | continue |
| 785 | 874 | ||
| 786 | logger.info( | 875 | logger.info( |
| 787 | "service missing in conductor, re-registering instance=%s dp_rank=%s", | 876 | "service missing in conductor, re-registering instance=%s dp_rank=%s", |
| 788 | instance_id, | 877 | instance_id, |
| 789 | - ep.id, | 878 | + dp_rank, |
| 790 | ) | 879 | ) |
| 791 | - cls.register_post(instance, ep) | 880 | + cls.register_post(instance, ep, dp_rank) |
| @@ -16,7 +16,12 @@ from motor.common.resources.instance import Instance, PDRole | |||
| 16 | from motor.common.utils.singleton import ThreadSafeSingleton | 16 | from motor.common.utils.singleton import ThreadSafeSingleton |
| 17 | from motor.common.resources.endpoint import Endpoint | 17 | from motor.common.resources.endpoint import Endpoint |
| 18 | from motor.coordinator.domain import InstanceProvider | 18 | from motor.coordinator.domain import InstanceProvider |
| 19 | -from motor.coordinator.domain.block_offset_translator import attach_block_offsets | 19 | +from motor.coordinator.domain.block_offset_translator import ( |
| 20 | + _DSV4_ASSISTANT_MARKER, | ||
| 21 | + _DSV4_EOS_MARKER, | ||
| 22 | + _DSV4_USER_MARKER, | ||
| 23 | + attach_block_offsets, | ||
| 24 | +) | ||
| 20 | from motor.coordinator.domain.responses_input import responses_scheduling_messages | 25 | from motor.coordinator.domain.responses_input import responses_scheduling_messages |
| 21 | from motor.coordinator.scheduler.policy.base import BaseSchedulingPolicy | 26 | from motor.coordinator.scheduler.policy.base import BaseSchedulingPolicy |
| 22 | from motor.config.coordinator import ( | 27 | from motor.config.coordinator import ( |
| @@ -353,6 +358,25 @@ class KvCacheAffinityPolicy(BaseSchedulingPolicy): | |||
| 353 | logger.debug("Could not read conductor block_size: %s", e) | 358 | logger.debug("Could not read conductor block_size: %s", e) |
| 354 | return 0 | 359 | return 0 |
| 355 | 360 | ||
| 361 | + @staticmethod | ||
| 362 | + def _matched_raw_for_endpoint(dp_map: dict, endpoint: Endpoint, fold_engine_dps: bool) -> object: | ||
| 363 | + """Return the conductor DP payload used to score one HTTP endpoint. | ||
| 364 | + | ||
| 365 | + When the instance exposes a single HTTP port, take the best match | ||
| 366 | + across every engine DP the conductor reported. Multi-endpoint | ||
| 367 | + instances keep the 1:1 ``endpoint.id`` mapping. | ||
| 368 | + """ | ||
| 369 | + if not fold_engine_dps: | ||
| 370 | + return dp_map.get(f"{endpoint.id}", 0) | ||
| 371 | + best_raw: object = dp_map.get(f"{endpoint.id}", 0) | ||
| 372 | + best_tokens = KvCacheAffinityPolicy._weighted_matched_tokens(best_raw, 1, 1.0, 1.0, 1.0) | ||
| 373 | + for raw in dp_map.values(): | ||
| 374 | + tokens = KvCacheAffinityPolicy._weighted_matched_tokens(raw, 1, 1.0, 1.0, 1.0) | ||
| 375 | + if tokens > best_tokens: | ||
| 376 | + best_tokens = tokens | ||
| 377 | + best_raw = raw | ||
| 378 | + return best_raw | ||
| 379 | + | ||
| 356 | @staticmethod | 380 | @staticmethod |
| 357 | def _weighted_matched_tokens( | 381 | def _weighted_matched_tokens( |
| 358 | matched_raw: object, | 382 | matched_raw: object, |
| @@ -436,8 +460,13 @@ class KvCacheAffinityPolicy(BaseSchedulingPolicy): | |||
| 436 | dp_map = instance_data.get("DP", {}) | 460 | dp_map = instance_data.get("DP", {}) |
| 437 | # get_all_endpoints() is the canonical accessor: it flattens the per-DP map and | 461 | # get_all_endpoints() is the canonical accessor: it flattens the per-DP map and |
| 438 | # already excludes headless endpoints / respects enable_multi_endpoints. | 462 | # already excludes headless endpoints / respects enable_multi_endpoints. |
| 439 | - for ep in instance.get_all_endpoints(): | 463 | + endpoints = list(instance.get_all_endpoints()) |
| 440 | - matched_raw = dp_map.get(f"{ep.id}", 0) | 464 | + # SGLang (and any single-HTTP / multi-DP engine) publishes one ZMQ |
| 465 | + # stream per engine DP. Fold those ranks onto the only routable | ||
| 466 | + # HTTP endpoint so instance pick can see hits on DP>0. | ||
| 467 | + fold_engine_dps = len(endpoints) == 1 | ||
 | |||
| 468 | + for ep in endpoints: | ||
| 469 | + matched_raw = KvCacheAffinityPolicy._matched_raw_for_endpoint(dp_map, ep, fold_engine_dps) | ||
| 441 | matched = KvCacheAffinityPolicy._weighted_matched_tokens(matched_raw, block_size, w_npu, w_cpu, w_disk) | 470 | matched = KvCacheAffinityPolicy._weighted_matched_tokens(matched_raw, block_size, w_npu, w_cpu, w_disk) |
| 442 | # Cap at the prompt length as a safety bound, since a matched prefix | 471 | # Cap at the prompt length as a safety bound, since a matched prefix |
| 443 | # cannot be longer than the prompt itself. | 472 | # cannot be longer than the prompt itself. |
| @@ -749,11 +778,22 @@ class TokenizerManager(ThreadSafeSingleton): | |||
| 749 | if not getattr(self, "model_path", ""): | 778 | if not getattr(self, "model_path", ""): |
| 750 | return None | 779 | return None |
| 751 | try: | 780 | try: |
| 752 | - if self.engine_type == "vllm" and self._is_deepseek_v4_model(self.model_path): | 781 | + if self._is_deepseek_v4_model(self.model_path): |
| 753 | - from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer # pylint: disable=import-error,no-name-in-module | 782 | + try: |
| 783 | + from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer # pylint: disable=import-error,no-name-in-module | ||
| 784 | + except ImportError as exc: | ||
| 785 | + logger.warning( | ||
| 786 | + "DeepseekV4Tokenizer unavailable for engine_type=%s (%s); falling back to AutoTokenizer", | ||
| 787 | + self.engine_type, | ||
| 788 | + exc, | ||
| 789 | + ) | ||
| 790 | + from transformers import AutoTokenizer | ||
| 754 | 791 | ||
| 755 | - self.tokenizer = DeepseekV4Tokenizer.from_pretrained(self.model_path, trust_remote_code=True) | 792 | + self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) |
| 756 | - self._is_dsv4 = True | 793 | + self._is_dsv4 = True |
| 794 | + else: | ||
| 795 | + self.tokenizer = DeepseekV4Tokenizer.from_pretrained(self.model_path, trust_remote_code=True) | ||
| 796 | + self._is_dsv4 = True | ||
| 757 | else: | 797 | else: |
| 758 | from transformers import AutoTokenizer | 798 | from transformers import AutoTokenizer |
| 759 | 799 | ||
| @@ -862,11 +902,33 @@ class TokenizerManager(ThreadSafeSingleton): | |||
| 862 | 902 | ||
| 863 | def _apply_chat_template_dsv4(self, messages: list, tools: list | None, req_data: dict | None) -> list[int]: | 903 | def _apply_chat_template_dsv4(self, messages: list, tools: list | None, req_data: dict | None) -> list[int]: |
| 864 | messages, tools = preprocess_messages_for_dsv4(messages, tools) | 904 | messages, tools = preprocess_messages_for_dsv4(messages, tools) |
| 905 | + if not getattr(self.tokenizer, "chat_template", None): | ||
| 906 | + return self._encode_dsv4_messages(messages) | ||
| 865 | result = self.tokenizer.apply_chat_template( | 907 | result = self.tokenizer.apply_chat_template( |
| 866 | messages, tools=tools, **self._build_dsv4_chat_template_kwargs(req_data) | 908 | messages, tools=tools, **self._build_dsv4_chat_template_kwargs(req_data) |
| 867 | ) | 909 | ) |
| 868 | return result if isinstance(result, list) else self.tokenizer.encode(result, add_special_tokens=False) | 910 | return result if isinstance(result, list) else self.tokenizer.encode(result, add_special_tokens=False) |
| 869 | 911 | ||
| 912 | + def _encode_dsv4_messages(self, messages: list) -> list[int]: | ||
| 913 | + """DeepSeek V4 chat markers when vLLM tokenizer / chat_template is absent.""" | ||
| 914 | + parts: list[str] = [] | ||
| 915 | + for message in messages: | ||
| 916 | + if not isinstance(message, dict): | ||
| 917 | + continue | ||
| 918 | + role = message.get("role") | ||
| 919 | + content = message.get("content") or "" | ||
| 920 | + if not isinstance(content, str): | ||
| 921 | + content = str(content) | ||
| 922 | + if role in ("user", "system", "developer"): | ||
| 923 | + parts.append(f"{_DSV4_USER_MARKER}{content}") | ||
| 924 | + elif role == "assistant": | ||
| 925 | + parts.append(f"{_DSV4_ASSISTANT_MARKER}{content}{_DSV4_EOS_MARKER}") | ||
| 926 | + elif role == "tool": | ||
| 927 | + parts.append(f"{_DSV4_USER_MARKER}{content}") | ||
| 928 | + if not messages or (isinstance(messages[-1], dict) and messages[-1].get("role") != "assistant"): | ||
| 929 | + parts.append(_DSV4_ASSISTANT_MARKER) | ||
| 930 | + return self.tokenizer.encode("".join(parts), add_special_tokens=False) | ||
| 931 | + | ||
| 870 | def _apply_chat_template_standard(self, messages: list, tools: list | None, req_data: dict | None) -> list[int]: | 932 | def _apply_chat_template_standard(self, messages: list, tools: list | None, req_data: dict | None) -> list[int]: |
| 871 | return self.tokenizer.apply_chat_template( | 933 | return self.tokenizer.apply_chat_template( |
| 872 | conversation=preprocess_messages_for_standard(messages), | 934 | conversation=preprocess_messages_for_standard(messages), |
| @@ -113,6 +113,26 @@ class TestConductorInstanceId: | |||
| 113 | inst = _make_instance(inst_id=9, role=PDRole.ROLE_D) | 113 | inst = _make_instance(inst_id=9, role=PDRole.ROLE_D) |
| 114 | assert conductor_instance_id(inst) == "vllm-prefill-9" | 114 | assert conductor_instance_id(inst) == "vllm-prefill-9" |
| 115 | 115 | ||
| 116 | + def test_sglang_prefill_uses_sglang_prefix(self): | ||
| 117 | + inst = _make_instance(inst_id=3, role=PDRole.ROLE_P) | ||
| 118 | + inst.engine_type = "sglang" | ||
| 119 | + assert conductor_instance_id(inst) == "sglang-prefill-3" | ||
| 120 | + | ||
| 121 | + def test_sglang_union_uses_sglang_prefix(self): | ||
| 122 | + inst = _make_instance(inst_id=7, role=PDRole.ROLE_U) | ||
| 123 | + inst.engine_type = "SGLang" | ||
| 124 | + assert conductor_instance_id(inst) == "sglang-union-7" | ||
| 125 | + | ||
| 126 | + def test_unknown_engine_keeps_vllm_prefix(self): | ||
| 127 | + inst = _make_instance(inst_id=4, role=PDRole.ROLE_P) | ||
| 128 | + inst.engine_type = "trtllm" | ||
| 129 | + assert conductor_instance_id(inst) == "vllm-prefill-4" | ||
| 130 | + | ||
| 131 | + def test_explicit_vllm_keeps_vllm_prefix(self): | ||
| 132 | + inst = _make_instance(inst_id=1, role=PDRole.ROLE_P) | ||
| 133 | + inst.engine_type = "vllm" | ||
| 134 | + assert conductor_instance_id(inst) == "vllm-prefill-1" | ||
| 135 | + | ||
| 116 | 136 | ||
| 117 | # ------------------------------------------------------------------ | 137 | # ------------------------------------------------------------------ |
| 118 | # _build_register_payload | 138 | # _build_register_payload |
| @@ -186,9 +206,9 @@ class TestBuildRegisterPayload: | |||
| 186 | with patch.object(ConductorApiClient, "coordinator_config", cfg): | 206 | with patch.object(ConductorApiClient, "coordinator_config", cfg): |
| 187 | payload = ConductorApiClient._build_register_payload(inst, ep) | 207 | payload = ConductorApiClient._build_register_payload(inst, ep) |
| 188 | 208 | ||
| 189 | - # Fallback endpoint fills gpu, cpu, disk | 209 | + # Engine ZMQ fallback is NPU-only. Copying it onto cpu/disk would make |
| 190 | - meps = payload["medium_endpoints"] | 210 | + # kv-conductor parse the stream as a pool and drop SGLang/vLLM events. |
| 191 | - assert "npu" in meps | 211 | + assert payload["medium_endpoints"] == {"npu": "tcp://10.0.0.4:15557"} |
| 192 | 212 | ||
| 193 | def test_yuanrong_payload_is_npu_only(self): | 213 | def test_yuanrong_payload_is_npu_only(self): |
| 194 | """YuanRong DP payload omits CPU/Disk (those are registered per node).""" | 214 | """YuanRong DP payload omits CPU/Disk (those are registered per node).""" |
| @@ -467,7 +487,7 @@ class TestReRegisterKvInstances: | |||
| 467 | ): | 487 | ): |
| 468 | ConductorApiClient.re_register_kv_instances([inst]) | 488 | ConductorApiClient.re_register_kv_instances([inst]) |
| 469 | 489 | ||
| 470 | - mock_register.assert_called_once_with(inst, ep) | 490 | + mock_register.assert_called_once_with(inst, ep, 0) |
| 471 | 491 | ||
| 472 | def test_skips_when_already_registered(self): | 492 | def test_skips_when_already_registered(self): |
| 473 | """Instance already in Conductor → register_post NOT called.""" | 493 | """Instance already in Conductor → register_post NOT called.""" |
| @@ -555,7 +575,7 @@ class TestReRegisterKvInstances: | |||
| 555 | ): | 575 | ): |
| 556 | ConductorApiClient.re_register_kv_instances([inst]) | 576 | ConductorApiClient.re_register_kv_instances([inst]) |
| 557 | 577 | ||
| 558 | - mock_register.assert_called_once_with(inst, ep) | 578 | + mock_register.assert_called_once_with(inst, ep, 0) |
| 559 | 579 | ||
| 560 | 580 | ||
| 561 | # ── Registration dispatch tests ────────────────────────────────────── | 581 | # ── Registration dispatch tests ────────────────────────────────────── |
| @@ -1142,3 +1162,99 @@ def test_re_register_node_pool_only_leaves_live_dps_alone(mock_http): | |||
| 1142 | payloads = [c[0][1] for c in mock_client.post.call_args_list] | 1162 | payloads = [c[0][1] for c in mock_client.post.call_args_list] |
| 1143 | assert len(payloads) == 1 | 1163 | assert len(payloads) == 1 |
| 1144 | assert payloads[0]["instance_id"] == "yuanrong-pool-10.0.0.1-test-model" | 1164 | assert payloads[0]["instance_id"] == "yuanrong-pool-10.0.0.1-test-model" |
| 1165 | + | ||
| 1166 | + | ||
| 1167 | +# ── SGLang single-HTTP / multi-DP event expansion ──────────────────── | ||
| 1168 | + | ||
| 1169 | + | ||
| 1170 | +def _make_single_http_dp_instance(instance_id: int, dp_size: int, ip: str = "10.0.0.1"): | ||
| 1171 | + """One HTTP endpoint, engine dp_size > 1 (SGLang dp-attention).""" | ||
| 1172 | + endpoint = Mock() | ||
| 1173 | + endpoint.id = 0 | ||
| 1174 | + endpoint.ip = ip | ||
| 1175 | + instance = Mock() | ||
| 1176 | + instance.id = instance_id | ||
| 1177 | + instance.model_name = "test-model" | ||
| 1178 | + instance.role = "prefill" | ||
| 1179 | + instance.engine_type = "sglang" | ||
| 1180 | + instance.endpoints = {"pod-0": {0: endpoint}} | ||
| 1181 | + instance.get_all_endpoints.return_value = (endpoint,) | ||
| 1182 | + instance.parallel_config = Mock() | ||
| 1183 | + instance.parallel_config.dp_size = dp_size | ||
| 1184 | + return instance | ||
| 1185 | + | ||
| 1186 | + | ||
| 1187 | +def test_hbm_event_targets_expands_single_http_to_dp_size(): | ||
| 1188 | + instance = _make_single_http_dp_instance(1, dp_size=16) | ||
| 1189 | + targets = ConductorApiClient._hbm_event_targets(instance) | ||
| 1190 | + assert [rank for _ep, rank in targets] == list(range(16)) | ||
| 1191 | + assert {ep.ip for ep, _rank in targets} == {"10.0.0.1"} | ||
| 1192 | + | ||
| 1193 | + | ||
| 1194 | +def test_hbm_event_targets_keeps_one_to_one_when_endpoints_cover_dp(): | ||
| 1195 | + instance = _make_node_instance(1, "10.0.0.1", [0, 1, 2]) | ||
| 1196 | + instance.parallel_config = Mock() | ||
| 1197 | + instance.parallel_config.dp_size = 3 | ||
| 1198 | + targets = ConductorApiClient._hbm_event_targets(instance) | ||
| 1199 | + assert [(ep.id, rank) for ep, rank in targets] == [(0, 0), (1, 1), (2, 2)] | ||
| 1200 | + | ||
| 1201 | + | ||
| 1202 | +def test_hbm_event_targets_does_not_pad_mixed_http_less_than_dp(): | ||
| 1203 | + """1 < HTTP < dp_size: do not invent ranks that scoring cannot fold.""" | ||
| 1204 | + instance = _make_node_instance(1, "10.0.0.1", [0, 1]) | ||
| 1205 | + instance.parallel_config = Mock() | ||
| 1206 | + instance.parallel_config.dp_size = 4 | ||
| 1207 | + targets = ConductorApiClient._hbm_event_targets(instance) | ||
| 1208 | + assert [(ep.id, rank) for ep, rank in targets] == [(0, 0), (1, 1)] | ||
| 1209 | + | ||
| 1210 | + | ||
| 1211 | +@patch("motor.coordinator.api_client.conductor_api_client.SafeHTTPSClient") | ||
| 1212 | +def test_register_kv_instance_subscribes_all_sglang_dp_ranks(mock_http): | ||
| 1213 | + """One HTTP port + dp_size=16 must register ZMQ 5557-5572 (ranks 0-15).""" | ||
| 1214 | + mock_client = Mock() | ||
| 1215 | + mock_client.post.return_value = {"status": "ok"} | ||
| 1216 | + mock_http.return_value.__enter__.return_value = mock_client | ||
| 1217 | + | ||
| 1218 | + instance = _make_single_http_dp_instance(1, dp_size=16) | ||
| 1219 | + with _setup_reg_config("YuanRong", npu_endpoint="tcp://*:5557", replay_endpoint="tcp://*:6667"): | ||
| 1220 | + ConductorApiClient.register_kv_instance([instance]) | ||
| 1221 | + | ||
| 1222 | + payloads = [c[0][1] for c in mock_client.post.call_args_list if "npu" in c[0][1].get("medium_endpoints", {})] | ||
| 1223 | + assert {p["dp_rank"] for p in payloads} == set(range(16)) | ||
| 1224 | + assert {p["instance_id"] for p in payloads} == {"sglang-prefill-1"} | ||
| 1225 | + assert {p["medium_endpoints"]["npu"] for p in payloads} == {f"tcp://10.0.0.1:{5557 + rank}" for rank in range(16)} | ||
| 1226 | + assert {p["replay_endpoint"] for p in payloads} == {f"tcp://10.0.0.1:{6667 + rank}" for rank in range(16)} | ||
| 1227 | + | ||
| 1228 | + | ||
| 1229 | +@patch("motor.coordinator.api_client.conductor_api_client.SafeHTTPSClient") | ||
| 1230 | +def test_unregister_kv_instance_drops_all_expanded_dp_ranks(mock_http): | ||
| 1231 | + mock_client = Mock() | ||
| 1232 | + mock_client.post.return_value = {"status": "ok"} | ||
| 1233 | + mock_http.return_value.__enter__.return_value = mock_client | ||
| 1234 | + | ||
| 1235 | + instance = _make_single_http_dp_instance(1, dp_size=4) | ||
| 1236 | + with _setup_reg_config("YuanRong", npu_endpoint="tcp://*:5557"): | ||
| 1237 | + ConductorApiClient.unregister_kv_instance([instance]) | ||
| 1238 | + | ||
| 1239 | + payloads = [c[0][1] for c in mock_client.post.call_args_list if c[0][0] == "/unregister"] | ||
| 1240 | + assert {p["dp_rank"] for p in payloads} == {0, 1, 2, 3} | ||
| 1241 | + | ||
| 1242 | + | ||
| 1243 | +@patch("motor.coordinator.api_client.conductor_api_client.SafeHTTPSClient") | ||
| 1244 | +def test_re_register_missing_expanded_sglang_dp_rank(mock_http): | ||
| 1245 | + mock_client = Mock() | ||
| 1246 | + mock_client.post.return_value = {"status": "ok"} | ||
| 1247 | + mock_http.return_value.__enter__.return_value = mock_client | ||
| 1248 | + | ||
| 1249 | + instance = _make_single_http_dp_instance(1, dp_size=4) | ||
| 1250 | + cfg = _mock_config(npu_endpoint="tcp://*:5557") | ||
| 1251 | + # Only dp 0 is known to conductor; ranks 1-3 must be rebuilt. | ||
| 1252 | + registered = [{"instance_id": "sglang-prefill-1", "endpoints": {"0": {}}}] | ||
| 1253 | + with ( | ||
| 1254 | + patch.object(ConductorApiClient, "coordinator_config", cfg), | ||
| 1255 | + patch.object(ConductorApiClient, "get_registered_services", return_value=registered), | ||
| 1256 | + patch.object(ConductorApiClient, "register_post") as mock_register, | ||
| 1257 | + ): | ||
| 1258 | + ConductorApiClient.re_register_kv_instances([instance]) | ||
| 1259 | + | ||
| 1260 | + assert {call.args[2] for call in mock_register.call_args_list} == {1, 2, 3} | ||
| @@ -364,6 +364,52 @@ class TestKvCacheAffinityPolicy(unittest.TestCase): | |||
| 364 | # Both have equal load; ep 0 has longer cached prefix (800 > 100 tokens) → lower prefill cost → chosen. | 364 | # Both have equal load; ep 0 has longer cached prefix (800 > 100 tokens) → lower prefill cost → chosen. |
| 365 | self.assertEqual(result[1].id, 0) | 365 | self.assertEqual(result[1].id, 0) |
| 366 | 366 | ||
| 367 | + @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.ConductorApiClient.query_conductor') | ||
| 368 | + @patch('motor.coordinator.scheduler.policy.kv_cache_affinity.TokenizerManager') | ||
| 369 | + def test_single_http_endpoint_uses_best_engine_dp_match(self, mock_tokenizer_manager, mock_query_conductor): | ||
| 370 | + """SGLang: one HTTP port, conductor hit on DP>0 still picks that instance.""" | ||
| 371 | + ep_hit = _make_endpoint(0, active_tokens=50.0) | ||
| 372 | + inst_hit = Mock() | ||
| 373 | + inst_hit.id = "hit" | ||
| 374 | + inst_hit.engine_type = "sglang" | ||
| 375 | + inst_hit.endpoints = {"group": {0: ep_hit}} | ||
| 376 | + inst_hit.get_all_endpoints.return_value = (ep_hit,) | ||
| 377 | + | ||
| 378 | + ep_miss = _make_endpoint(0, active_tokens=50.0) | ||
| 379 | + inst_miss = Mock() | ||
| 380 | + inst_miss.id = "miss" | ||
| 381 | + inst_miss.engine_type = "sglang" | ||
| 382 | + inst_miss.endpoints = {"group": {0: ep_miss}} | ||
| 383 | + inst_miss.get_all_endpoints.return_value = (ep_miss,) | ||
| 384 | + | ||
| 385 | + mock_req_info = Mock() | ||
| 386 | + mock_req_info.req_data = {"prompt": "hello"} | ||
| 387 | + mock_tokenizer = Mock() | ||
| 388 | + mock_tokenizer.encode.return_value = list(range(1000)) | ||
| 389 | + mock_tokenizer_manager.return_value = mock_tokenizer | ||
| 390 | + | ||
| 391 | + mock_query_conductor.return_value = { | ||
| 392 | + TENANT_ID: { | ||
| 393 | + "sglang-prefill-hit": { | ||
| 394 | + "DP": { | ||
| 395 | + "0": {"npu_blocks": 0, "cpu_blocks": 0, "disk_blocks": 0, "matched_tokens": 0}, | ||
| 396 | + "5": {"npu_blocks": 6, "cpu_blocks": 0, "disk_blocks": 0, "matched_tokens": 800}, | ||
| 397 | + } | ||
| 398 | + }, | ||
| 399 | + "sglang-prefill-miss": { | ||
| 400 | + "DP": { | ||
| 401 | + "0": {"npu_blocks": 0, "cpu_blocks": 0, "disk_blocks": 0, "matched_tokens": 0}, | ||
| 402 | + } | ||
| 403 | + }, | ||
| 404 | + } | ||
| 405 | + } | ||
| 406 | + | ||
| 407 | + result = KvCacheAffinityPolicy.select_endpoint_from_list([inst_miss, inst_hit], mock_req_info, load_weight=0.0) | ||
| 408 | + | ||
| 409 | + self.assertIsNotNone(result) | ||
| 410 | + self.assertIs(result[0], inst_hit) | ||
| 411 | + self.assertEqual(result[1].id, 0) | ||
| 412 | + | ||
| 367 | @patch.object(KvCacheAffinityPolicy, "_conductor_block_size", return_value=4) | 413 | @patch.object(KvCacheAffinityPolicy, "_conductor_block_size", return_value=4) |
| 368 | @patch("motor.coordinator.scheduler.policy.kv_cache_affinity.ConductorApiClient.query_conductor") | 414 | @patch("motor.coordinator.scheduler.policy.kv_cache_affinity.ConductorApiClient.query_conductor") |
| 369 | @patch("motor.coordinator.scheduler.policy.kv_cache_affinity.TokenizerManager") | 415 | @patch("motor.coordinator.scheduler.policy.kv_cache_affinity.TokenizerManager") |
| @@ -1001,6 +1047,19 @@ class TestTokenizerManagerDsv4(unittest.TestCase): | |||
| 1001 | self.assertEqual(out, []) | 1047 | self.assertEqual(out, []) |
| 1002 | self.assertEqual(tokenizer.apply_chat_template.call_count, 1) | 1048 | self.assertEqual(tokenizer.apply_chat_template.call_count, 1) |
| 1003 | 1049 | ||
| 1050 | + def test_apply_chat_template_dsv4_encodes_markers_without_chat_template(self): | ||
| 1051 | + tokenizer = Mock() | ||
| 1052 | + tokenizer.chat_template = None | ||
| 1053 | + tokenizer.encode.return_value = [11, 22, 33] | ||
| 1054 | + manager = self._make_manager(tokenizer, is_dsv4=True) | ||
| 1055 | + | ||
| 1056 | + out = manager._apply_chat_template_dsv4([{"role": "user", "content": "hi"}], None, None) | ||
| 1057 | + self.assertEqual(out, [11, 22, 33]) | ||
| 1058 | + tokenizer.apply_chat_template.assert_not_called() | ||
| 1059 | + encoded = tokenizer.encode.call_args[0][0] | ||
| 1060 | + self.assertIn("<｜User｜>hi", encoded) | ||
| 1061 | + self.assertTrue(encoded.endswith("<｜Assistant｜>")) | ||
| 1062 | + | ||
| 1004 | def test_is_deepseek_v4_model_detects_from_config(self): | 1063 | def test_is_deepseek_v4_model_detects_from_config(self): |
| 1005 | with tempfile.TemporaryDirectory() as tmp: | 1064 | with tempfile.TemporaryDirectory() as tmp: |
| 1006 | tmp_path = Path(tmp) | 1065 | tmp_path = Path(tmp) |
| @@ -1589,27 +1648,80 @@ class TestTokenizerManagerFunction(unittest.TestCase): | |||
| 1589 | # verification result | 1648 | # verification result |
| 1590 | self.assertEqual(result, []) | 1649 | self.assertEqual(result, []) |
| 1591 | 1650 | ||
| 1592 | - @patch('motor.config.coordinator.CoordinatorConfig') | 1651 | + def _mock_tokenizer_config(self, engine_type: str): |
| 1593 | - def test_dsv4_tokenizer_only_for_vllm_engine(self, mock_config_class): | ||
| 1594 | - """DeepSeek V4 vLLM tokenizer must only be used when engine_type=vllm.""" | ||
| 1595 | mock_config = Mock() | 1652 | mock_config = Mock() |
| 1596 | mock_config.scheduler_config.kv_conductor_config.conductor_service = "test_service" | 1653 | mock_config.scheduler_config.kv_conductor_config.conductor_service = "test_service" |
| 1597 | mock_config.scheduler_config.kv_conductor_config.model_path = "/path/to/model" | 1654 | mock_config.scheduler_config.kv_conductor_config.model_path = "/path/to/model" |
| 1598 | - mock_config.scheduler_config.kv_conductor_config.engine_type = "sglang" | 1655 | + mock_config.scheduler_config.kv_conductor_config.engine_type = engine_type |
| 1599 | mock_config.tracer_config.endpoint = "" | 1656 | mock_config.tracer_config.endpoint = "" |
| 1600 | - mock_config_class.return_value = mock_config | 1657 | + return mock_config |
| 1601 | 1658 | ||
| 1602 | - # If the code accidentally tries to import vllm.tokenizers.deepseek_v4 on sglang, | 1659 | + def _install_dsv4_module(self, mock_tokenizer): |
| 1603 | - # environments without vllm installed would crash. We assert we fall back to transformers. | 1660 | + import types |
| 1661 | + | ||
| 1662 | + mock_cls = Mock() | ||
| 1663 | + mock_cls.from_pretrained.return_value = mock_tokenizer | ||
| 1664 | + vllm_mod = types.ModuleType("vllm") | ||
| 1665 | + tokenizers_mod = types.ModuleType("vllm.tokenizers") | ||
| 1666 | + dsv4_mod = types.ModuleType("vllm.tokenizers.deepseek_v4") | ||
| 1667 | + dsv4_mod.DeepseekV4Tokenizer = mock_cls | ||
| 1668 | + return mock_cls, { | ||
| 1669 | + "vllm": vllm_mod, | ||
| 1670 | + "vllm.tokenizers": tokenizers_mod, | ||
| 1671 | + "vllm.tokenizers.deepseek_v4": dsv4_mod, | ||
| 1672 | + } | ||
| 1673 | + | ||
| 1674 | + @patch('motor.config.coordinator.CoordinatorConfig') | ||
| 1675 | + def test_dsv4_tokenizer_used_for_sglang_engine(self, mock_config_class): | ||
| 1676 | + """DeepSeek V4 chat template is engine-agnostic; SGLang must use it too.""" | ||
| 1677 | + mock_config = self._mock_tokenizer_config("sglang") | ||
| 1678 | + mock_config_class.return_value = mock_config | ||
| 1604 | mock_tokenizer = Mock() | 1679 | mock_tokenizer = Mock() |
| 1605 | - transformers_mod = Mock() | 1680 | + mock_cls, modules = self._install_dsv4_module(mock_tokenizer) |
| 1606 | - transformers_mod.AutoTokenizer.from_pretrained.return_value = mock_tokenizer | 1681 | + with patch.dict("sys.modules", modules): |
| 1607 | - with patch.dict("sys.modules", {"transformers": transformers_mod}): | ||
| 1608 | with patch.object(TokenizerManager, "_is_deepseek_v4_model", return_value=True): | 1682 | with patch.object(TokenizerManager, "_is_deepseek_v4_model", return_value=True): |
| 1609 | manager = TokenizerManager(mock_config) | 1683 | manager = TokenizerManager(mock_config) |
| 1610 | 1684 | ||
| 1611 | self.assertIs(manager.tokenizer, mock_tokenizer) | 1685 | self.assertIs(manager.tokenizer, mock_tokenizer) |
| 1612 | - self.assertFalse(manager._is_dsv4) | 1686 | + self.assertTrue(manager._is_dsv4) |
| 1687 | + mock_cls.from_pretrained.assert_called_once() | ||
| 1688 | + | ||
| 1689 | + @patch('motor.config.coordinator.CoordinatorConfig') | ||
| 1690 | + def test_dsv4_tokenizer_used_for_vllm_engine(self, mock_config_class): | ||
| 1691 | + mock_config = self._mock_tokenizer_config("vllm") | ||
| 1692 | + mock_config_class.return_value = mock_config | ||
| 1693 | + mock_tokenizer = Mock() | ||
| 1694 | + mock_cls, modules = self._install_dsv4_module(mock_tokenizer) | ||
| 1695 | + with patch.dict("sys.modules", modules): | ||
| 1696 | + with patch.object(TokenizerManager, "_is_deepseek_v4_model", return_value=True): | ||
| 1697 | + manager = TokenizerManager(mock_config) | ||
| 1698 | + | ||
| 1699 | + self.assertIs(manager.tokenizer, mock_tokenizer) | ||
| 1700 | + self.assertTrue(manager._is_dsv4) | ||
| 1701 | + mock_cls.from_pretrained.assert_called_once() | ||
| 1702 | + | ||
| 1703 | + @patch('motor.config.coordinator.CoordinatorConfig') | ||
| 1704 | + def test_dsv4_tokenizer_falls_back_without_vllm(self, mock_config_class): | ||
| 1705 | + """Missing vllm DeepseekV4Tokenizer must not crash SGLang coordinators.""" | ||
| 1706 | + mock_config = self._mock_tokenizer_config("sglang") | ||
| 1707 | + mock_config_class.return_value = mock_config | ||
| 1708 | + mock_tokenizer = Mock() | ||
| 1709 | + transformers_mod = Mock() | ||
| 1710 | + transformers_mod.AutoTokenizer.from_pretrained.return_value = mock_tokenizer | ||
| 1711 | + with patch.dict( | ||
| 1712 | + "sys.modules", | ||
| 1713 | + { | ||
| 1714 | + "vllm": None, | ||
| 1715 | + "vllm.tokenizers": None, | ||
| 1716 | + "vllm.tokenizers.deepseek_v4": None, | ||
| 1717 | + "transformers": transformers_mod, | ||
| 1718 | + }, | ||
| 1719 | + ): | ||
| 1720 | + with patch.object(TokenizerManager, "_is_deepseek_v4_model", return_value=True): | ||
| 1721 | + manager = TokenizerManager(mock_config) | ||
| 1722 | + | ||
| 1723 | + self.assertIs(manager.tokenizer, mock_tokenizer) | ||
| 1724 | + self.assertTrue(manager._is_dsv4) | ||
| 1613 | transformers_mod.AutoTokenizer.from_pretrained.assert_called_once() | 1725 | transformers_mod.AutoTokenizer.from_pretrained.assert_called_once() |
| 1614 | 1726 | ||
| 1615 | 1727 | ||
| @@ -53,6 +53,20 @@ def test_conductor_instance_id_role_p() -> None: | |||
| 53 | assert conductor_instance_id(instance) == "vllm-prefill-3" | 53 | assert conductor_instance_id(instance) == "vllm-prefill-3" |
| 54 | 54 | ||
| 55 | 55 | ||
| 56 | +def test_conductor_instance_id_sglang_role_p() -> None: | ||
| 57 | + instance = _build_instance(PDRole.ROLE_P) | ||
| 58 | + instance.id = 3 | ||
| 59 | + instance.engine_type = "sglang" | ||
| 60 | + assert conductor_instance_id(instance) == "sglang-prefill-3" | ||
| 61 | + | ||
| 62 | + | ||
| 63 | +def test_conductor_instance_id_sglang_role_u() -> None: | ||
| 64 | + instance = _build_instance(PDRole.ROLE_U) | ||
| 65 | + instance.id = 7 | ||
| 66 | + instance.engine_type = "sglang" | ||
| 67 | + assert conductor_instance_id(instance) == "sglang-union-7" | ||
| 68 | + | ||
| 69 | + | ||
| 56 | def test_register_post_uses_union_conductor_id_for_role_u() -> None: | 70 | def test_register_post_uses_union_conductor_id_for_role_u() -> None: |
| 57 | instance = _build_instance(PDRole.ROLE_U) | 71 | instance = _build_instance(PDRole.ROLE_U) |
| 58 | instance.id = 2 | 72 | instance.id = 2 |
注册侧 _hbm_event_targets 对 1 < len(endpoints) < dp_size 的实例也会把缺失的 rank 补上（挂在 endpoints[0]），但查询侧的 fold 只在 len(endpoints) == 1 时启用——这种混合实例补注册的 rank 即使有缓存命中，选实例时还是查不到（dp_map.get(ep.id) 只看 endpoint id）。现实中这个形态少见，不过两边的条件要不要对齐一下（比如都用 len(endpoints) < dp_size 来 fold），免得以后踩坑？