from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Literal WorkerRole = Literal["prefill", "decode", "direct"] @dataclass(frozen=True) class WorkerSpec: role: WorkerRole ordinal: int gpu_ids: tuple[int, ...] host: str port: int bootstrap_port: int | None = None @property def worker_id(self) -> str: return f"{self.role}-{self.ordinal}" @property def url(self) -> str: return f"http://{self.host}:{self.port}" @property def gpu_id(self) -> int: return self.gpu_ids[0] @property def tp_size(self) -> int: return len(self.gpu_ids) @dataclass(frozen=True) class SingleNodeTopology: model_path: str prefill_workers: tuple[WorkerSpec, ...] decode_workers: tuple[WorkerSpec, ...] direct_workers: tuple[WorkerSpec, ...] router_host: str router_port: int transfer_backend: str trust_remote_code: bool force_rdma: bool = False ib_device: str | None = None extra_server_args: tuple[str, ...] = () prefill_extra_server_args: tuple[str, ...] = () decode_extra_server_args: tuple[str, ...] = () direct_extra_server_args: tuple[str, ...] = () @property def model_name(self) -> str: return Path(self.model_path).name @property def router_url(self) -> str: return f"http://{self.router_host}:{self.router_port}" @property def route_workers(self) -> tuple[WorkerSpec, ...]: if self.decode_workers: return self.decode_workers return self.direct_workers def route_index(self, worker_id: str) -> int: for idx, worker in enumerate(self.route_workers): if worker.worker_id == worker_id: return idx raise KeyError(f"Unknown route worker: {worker_id}") def build_single_node_topology( *, model_path: str, prefill_worker_count: int, decode_worker_count: int, direct_worker_count: int = 0, prefill_tp_size: int = 1, decode_tp_size: int = 1, direct_tp_size: int = 1, prefill_gpu_ids: tuple[int, ...] | None = None, decode_gpu_ids: tuple[int, ...] | None = None, direct_gpu_ids: tuple[int, ...] | None = None, total_gpu_budget: int = 8, host: str = "127.0.0.1", router_port: int = 8000, prefill_port_base: int = 30000, decode_port_base: int = 31000, direct_port_base: int = 32000, bootstrap_port_base: int = 8998, transfer_backend: str = "nixl", force_rdma: bool = False, trust_remote_code: bool = True, ib_device: str | None = None, extra_server_args: tuple[str, ...] = (), prefill_extra_server_args: tuple[str, ...] = (), decode_extra_server_args: tuple[str, ...] = (), direct_extra_server_args: tuple[str, ...] = (), ) -> SingleNodeTopology: if prefill_worker_count < 0: raise ValueError("prefill_worker_count must be >= 0") if decode_worker_count < 0: raise ValueError("decode_worker_count must be >= 0") if direct_worker_count < 0: raise ValueError("direct_worker_count must be >= 0") if ( prefill_worker_count == 0 and decode_worker_count == 0 and direct_worker_count == 0 ): raise ValueError("At least one worker must be configured") if prefill_tp_size <= 0: raise ValueError("prefill_tp_size must be >= 1") if decode_tp_size <= 0: raise ValueError("decode_tp_size must be >= 1") if direct_tp_size <= 0: raise ValueError("direct_tp_size must be >= 1") if force_rdma and not ib_device: raise ValueError("force_rdma requires --ib-device to be set") if force_rdma and transfer_backend != "mooncake": raise ValueError( "force_rdma currently requires transfer_backend='mooncake' " "to guarantee an RDMA path" ) total_gpus_required = ( prefill_worker_count * prefill_tp_size + decode_worker_count * decode_tp_size + direct_worker_count * direct_tp_size ) if total_gpus_required > total_gpu_budget: raise ValueError( "Single-node GPU budget exceeded: " f"{prefill_worker_count} prefill x tp={prefill_tp_size} + " f"{decode_worker_count} decode x tp={decode_tp_size} + " f"{direct_worker_count} direct x tp={direct_tp_size} > " f"{total_gpu_budget} GPUs" ) if prefill_gpu_ids is None: prefill_gpu_ids = tuple(range(prefill_worker_count * prefill_tp_size)) if decode_gpu_ids is None: decode_gpu_ids = tuple( range( len(prefill_gpu_ids), len(prefill_gpu_ids) + decode_worker_count * decode_tp_size, ) ) if direct_gpu_ids is None: direct_gpu_ids = tuple( range( len(prefill_gpu_ids) + len(decode_gpu_ids), len(prefill_gpu_ids) + len(decode_gpu_ids) + direct_worker_count * direct_tp_size, ) ) if len(prefill_gpu_ids) != prefill_worker_count * prefill_tp_size: raise ValueError( "prefill_gpu_ids length must equal prefill_worker_count * prefill_tp_size: " f"{len(prefill_gpu_ids)} != {prefill_worker_count * prefill_tp_size}" ) if len(decode_gpu_ids) != decode_worker_count * decode_tp_size: raise ValueError( "decode_gpu_ids length must equal decode_worker_count * decode_tp_size: " f"{len(decode_gpu_ids)} != {decode_worker_count * decode_tp_size}" ) if len(direct_gpu_ids) != direct_worker_count * direct_tp_size: raise ValueError( "direct_gpu_ids length must equal direct_worker_count * direct_tp_size: " f"{len(direct_gpu_ids)} != {direct_worker_count * direct_tp_size}" ) assigned_gpu_ids = prefill_gpu_ids + decode_gpu_ids + direct_gpu_ids if len(set(assigned_gpu_ids)) != len(assigned_gpu_ids): raise ValueError("prefill/decode/direct GPU IDs must be unique") if any(gpu_id < 0 or gpu_id >= total_gpu_budget for gpu_id in assigned_gpu_ids): raise ValueError( "GPU IDs must fall within the single-node budget range " f"[0, {total_gpu_budget - 1}]" ) prefill_workers = tuple( WorkerSpec( role="prefill", ordinal=idx, gpu_ids=tuple( prefill_gpu_ids[ idx * prefill_tp_size : (idx + 1) * prefill_tp_size ] ), host=host, port=prefill_port_base + idx, bootstrap_port=bootstrap_port_base + idx, ) for idx in range(prefill_worker_count) ) decode_workers = tuple( WorkerSpec( role="decode", ordinal=idx, gpu_ids=tuple( decode_gpu_ids[ idx * decode_tp_size : (idx + 1) * decode_tp_size ] ), host=host, port=decode_port_base + idx, ) for idx in range(decode_worker_count) ) direct_workers = tuple( WorkerSpec( role="direct", ordinal=idx, gpu_ids=tuple( direct_gpu_ids[ idx * direct_tp_size : (idx + 1) * direct_tp_size ] ), host=host, port=direct_port_base + idx, ) for idx in range(direct_worker_count) ) return SingleNodeTopology( model_path=model_path, prefill_workers=prefill_workers, decode_workers=decode_workers, direct_workers=direct_workers, router_host=host, router_port=router_port, transfer_backend=transfer_backend, trust_remote_code=trust_remote_code, force_rdma=force_rdma, ib_device=ib_device, extra_server_args=extra_server_args, prefill_extra_server_args=prefill_extra_server_args, decode_extra_server_args=decode_extra_server_args, direct_extra_server_args=direct_extra_server_args, )