246 lines
8.0 KiB
Python
246 lines
8.0 KiB
Python
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,
|
|
)
|