vllm_omni.distributed.omni_coordinator ¶

Bases: LoadBalancer

Select the replica with the smallest queue_length.

If multiple replicas share the same minimum queue length, one of them is chosen uniformly at random.

Raises:

select(task: Task, replicas: list[ReplicaInfo]) -> int

Bases: ABC

Abstract base class for load balancers.

Subclasses implement :meth:select to choose a replica for a given task.

select(task: Task, replicas: list[ReplicaInfo]) -> int

Bases: str, Enum

Enumeration for load balancing policies.

These policies are used by :class:LoadBalancer implementations to route tasks to a subset of available replicas.

LEAST_QUEUE_LENGTH = 'least-queue-length'

RANDOM = 'random'

ROUND_ROBIN = 'round-robin'

Client for AsyncOmni side to receive replica list updates.

This client maintains a SUB socket connected to OmniCoordinator's PUB endpoint and caches the latest :class:ReplicaList in memory for use by the load balancer and routing logic.

close() -> None

get_replica_list() -> ReplicaList

get_replicas_for_stage(stage_id: int) -> ReplicaList

Client used by stage replicas to send events to OmniCoordinator.

This client maintains a DEALER socket connected to OmniCoordinator's ROUTER endpoint and sends JSON-encoded events describing replica status.

close() -> None

update_info(
    status: ReplicaStatus | None = None,
    queue_length: int | None = None,
) -> None

Coordinator for stage replicas and hub clients.

This service receives replica events from :class:OmniCoordClientForStage via a ZMQ ROUTER socket and publishes active replica lists to :class:OmniCoordClientForHub via a PUB socket.

The coordinator maintains an in-memory registry of all known replicas, including their status, queue length, and heartbeat timestamps. A background thread periodically checks for heartbeat timeouts and marks unhealthy replicas as ReplicaStatus.ERROR.

pub_zmq_addr = getsockopt_string(LAST_ENDPOINT)

router_zmq_addr = getsockopt_string(LAST_ENDPOINT)

add_new_replica(event: ReplicaEvent) -> None

close() -> None

get_active_replicas() -> ReplicaList

publish_replica_list_update() -> bool

remove_replica(event: ReplicaEvent) -> None

update_replica_info(event: ReplicaEvent) -> None

Own one :class:OmniCoordinator and the two ports it binds.

Constructor binds; :meth:close tears down. The class deliberately does not expose the coordinator instance — callers should consume the coordinator only via its wire protocol through :class:OmniCoordClientForStage and :class:OmniCoordClientForHub.

pub_address: str = f'tcp://{host}:{pub_port}'

router_address: str = f'tcp://{host}:{router_port}'

close() -> None

Bases: LoadBalancer

Load balancer that selects a replica uniformly at random.

select(task: Task, replicas: list[ReplicaInfo]) -> int

Wire payload from OmniCoordClientForStage to OmniCoordinator.

Schema for Stage → Coordinator events over ZMQ: input_addr, output_addr, stage_id, status, queue_length, event_type.

event_type: str

input_addr: str

output_addr: str

queue_length: int

stage_id: int

status: ReplicaStatus

Metadata for a single stage replica.

This type is stored in OmniCoordinator's internal registry and is also published to hubs via :class:ReplicaList.

input_addr: str

last_heartbeat: float

output_addr: str

queue_length: int

registered_at: float

stage_id: int

status: ReplicaStatus

Container for replica list updates.

OmniCoordinator publishes a :class:ReplicaList whenever its view of active replicas changes. OmniCoordClientForHub caches the latest value and exposes it to AsyncOmni and the load balancer.

replicas: list[ReplicaInfo]

timestamp: float

Bases: str, Enum

Enumeration for stage replica status.

DOWN = 'down'

ERROR = 'error'

UP = 'up'

Bases: LoadBalancer

Load balancer that selects replicas in a round-robin fashion.

This implementation keeps a running index modulo len(replicas). It therefore depends on the order and stable meaning of the replicas list between calls. If the list length or ordering changes, the sequence of picks may skip or repeat entries relative to a fixed set of backends.

Concurrency: a threading.Lock serializes updates to _next_index for callers that invoke select from multiple threads or alongside threaded infrastructure (e.g. ZMQ receive threads).

select(task: Task, replicas: list[ReplicaInfo]) -> int

Bases: TypedDict

Task structure passed to StagePool.pick / LoadBalancer.select.

Mirrors the dict built around a stage submission with request_id and any payload-related fields a future load-balancing policy might inspect.

engine_inputs: Any

request_id: str

sampling_params: Any