diff --git a/megatron/core/inference/data_parallel_inference_coordinator/__init__.py b/megatron/core/inference/data_parallel_inference_coordinator/__init__.py new file mode 100644 index 00000000000..097b9f28f72 --- /dev/null +++ b/megatron/core/inference/data_parallel_inference_coordinator/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Data parallel inference coordinator package. + +The coordinator class itself lives in coordinator.py; message handlers are in +handlers.py and the control-signal state machine in state.py. This module +re-exports the public names so existing imports of +``megatron.core.inference.data_parallel_inference_coordinator`` keep working. +""" + +from .coordinator import DataParallelInferenceCoordinator +from .handlers import HANDLERS, message_handler +from .state import CONTROL_TRANSITIONS, ControlTransition, CoordinatorState + +__all__ = [ + "DataParallelInferenceCoordinator", + "CoordinatorState", + "ControlTransition", + "CONTROL_TRANSITIONS", + "HANDLERS", + "message_handler", +] diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py similarity index 68% rename from megatron/core/inference/data_parallel_inference_coordinator.py rename to megatron/core/inference/data_parallel_inference_coordinator/coordinator.py index 57aa9122c2e..4e1026206c6 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py @@ -7,7 +7,6 @@ import signal import socket from collections import deque -from enum import Enum, auto from multiprocessing import Event from multiprocessing.connection import Connection @@ -21,6 +20,9 @@ TextGenerationController, ) +from .handlers import HANDLERS +from .state import CoordinatorState + try: import zmq @@ -63,6 +65,10 @@ class DataParallelInferenceCoordinator: 5. **Control Signal Broadcasting**: It relays control signals (e.g., PAUSE, STOP) from a client to all connected data parallel ranks. + Message handling is split out into handlers.py: the event loop in start() + dispatches each message to the handler registered for its header, so + supporting a new message type requires no changes here. + Attributes: router_socket (zmq.Socket): The central ZMQ ROUTER socket for all communication. data_parallel_size (int): The number of data parallel workers to expect. @@ -75,13 +81,9 @@ class DataParallelInferenceCoordinator: next_request_id (int): A counter for generating unique server-side request IDs. """ - class CoordinatorState(Enum): - """State machine for the coordinator.""" - - RUNNING = auto() - PAUSED = auto() - SUSPENDED = auto() - STOPPING = auto() + # Exposed as a class attribute for backwards compatibility; the canonical + # definition lives in state.py. + CoordinatorState = CoordinatorState def __init__( self, @@ -189,7 +191,7 @@ def __init__( self.next_request_id = 0 self.tokenizer = tokenizer - self.state = self.CoordinatorState.RUNNING + self.state = CoordinatorState.RUNNING # Prefix caching state for routing. self.block_size_tokens = block_size_tokens @@ -220,6 +222,12 @@ def __init__( self._hash_table: dict[int, dict[int, int]] = {} self._hash_assignment_counter = 0 + # Clients that have completed the CONNECT handshake. + self.known_clients = set() + + # Header -> handler dispatch table, sourced from the handler registry. + self._handlers = dict(HANDLERS) + def get_least_loaded_data_parallel_rank(self): """ Selects the data parallel rank with the fewest in-flight requests. @@ -307,6 +315,12 @@ def _send_to_engine(self, identity, payload): return False raise + def _broadcast_to_engines(self, payload): + """Send a deserialized payload to every connected data parallel rank.""" + serialized = msgpack.packb(payload, use_bin_type=True) + for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks): + self._send_to_engine(data_parallel_rank_id, serialized) + def compute_request_hashes(self, prompt): """Compute block hashes for a prompt on CPU. @@ -408,212 +422,33 @@ def start(self): Starts the main event loop for the coordinator. This method runs an infinite loop, continuously listening for incoming - messages on the ZMQ ROUTER socket. It parses the message header to - determine the message type and takes appropriate action, such as - handling new client connections, forwarding requests, broadcasting - control signals, or processing replies from the engines. + messages on the ZMQ ROUTER socket. It reads the message header and + dispatches to the handler registered for it (see handlers.py). + A handler that returns a truthy value stops the loop. """ # Todo [Siddharth]: Make this more robust to handle invalid messages. - known_clients = set() while True: sender_identity, serialized_payload = self.router_socket.recv_multipart() - # Allow for re-registration if connecting to a running coordinator. + # An empty payload is a data parallel rank (re-)registering itself. if serialized_payload == b"": - if sender_identity not in self.identities_of_data_parallel_ranks: - self.identities_of_data_parallel_ranks.append(sender_identity) - self._register_rank_identity(sender_identity) + self._handle_rank_registration(sender_identity) continue deserialized_payload = msgpack.unpackb(serialized_payload, raw=False) header = Headers(deserialized_payload[0]) - if header == Headers.CONNECT: - if sender_identity in known_clients: - logging.info( - f"Client {sender_identity} sent a duplicate connect request. Ignoring .." - ) - continue - - # print(f"New client connected: {sender_identity}") - known_clients.add(sender_identity) - self.router_socket.send_multipart( - [sender_identity, msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)] - ) - - elif header == Headers.SUBMIT_REQUEST: - # ToDo [Siddharth]: We might want to tokenize the prompt on the - # assigned data parallel rank for this process instead - # of the coordinator. - - # Message from a known client - if sender_identity not in known_clients: - logging.info( - f"Received message from unknown client {sender_identity}. Ignoring." - ) - continue - # this is a message from a client. - # route it to a data parallel rank - client_request_id, prompt, sampling_params = deserialized_payload[1:] - # map client request_id to server request_id - # necessary because multiple clients might have the same request_id. - request_id = self.next_request_id - self.next_request_id += 1 - self.request_id_to_client_id[request_id] = sender_identity - self.request_id_to_client_request_id[request_id] = client_request_id - - # Serialize prompt. - if isinstance(prompt, (str, list)): - pass - elif isinstance(prompt, torch.Tensor): - prompt = prompt.tolist() - else: - raise Exception("specialize for <%s> prompt." % type(prompt).__name__) - - payload = msgpack.packb( - [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params], - use_bin_type=True, - ) - - request_hashes = self.compute_request_hashes(prompt) - if ( - self.prefix_caching_coordinator_policy - == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK - ): - request_hashes = request_hashes[:1] - - # Account for the fact that some engines may have died. - for _ in range(len(self.identities_of_data_parallel_ranks)): - next_identity = self.get_best_data_parallel_rank(request_hashes) - if self._send_to_engine(next_identity, payload): - break - else: - # If all engines have died, we are in an abnormal state, and must exit cleanly. - logging.error("Coordinator: no reachable engines for request %d", request_id) - del self.request_id_to_client_id[request_id] - del self.request_id_to_client_request_id[request_id] - return - - self.request_id_to_rank[request_id] = next_identity - self._pending_counts[self.identity_to_rank_index[next_identity]] += 1 - if request_hashes: - self._update_rank_hashes(next_identity, request_hashes) - if self.schedule_records is not None: - self.schedule_records.append( - { - "request_id": request_id, - "rank_index": self.identity_to_rank_index[next_identity], - "num_hashes": len(request_hashes), - } - ) - - elif header in ( - Headers.PAUSE, - Headers.UNPAUSE, - Headers.SUSPEND, - Headers.RESUME, - Headers.SET_GENERATION_EPOCH, - Headers.STOP, - ): - # Start by checking the current state against the control signal. - if sender_identity not in known_clients: - logging.warning("Coordinator: ignoring signal from unknown client.") - continue - - if header == Headers.PAUSE: - idem_states = (self.CoordinatorState.PAUSED, self.CoordinatorState.SUSPENDED) - if self.state == self.CoordinatorState.RUNNING: - self.state = self.CoordinatorState.PAUSED - elif self.state in idem_states: - # Already paused/suspended, ignore redundant PAUSE. - continue - else: - logging.warning("Coordinator: ignoring PAUSE in state %s", self.state) - continue - elif header == Headers.UNPAUSE: - if self.state != self.CoordinatorState.PAUSED: - logging.warning("Coordinator: ignoring UNPAUSE in state %s", self.state) - continue - self.state = self.CoordinatorState.RUNNING - elif header == Headers.SUSPEND: - if self.state != self.CoordinatorState.PAUSED: - logging.warning("Coordinator: ignoring SUSPEND in state %s", self.state) - continue - self.state = self.CoordinatorState.SUSPENDED - elif header == Headers.RESUME: - if self.state != self.CoordinatorState.SUSPENDED: - logging.warning("Coordinator: ignoring RESUME in state %s", self.state) - continue - self.state = self.CoordinatorState.PAUSED - elif header == Headers.STOP: - good_states = (self.CoordinatorState.PAUSED, self.CoordinatorState.SUSPENDED) - if self.state not in good_states: - logging.warning("Coordinator: ignoring STOP in state %s", self.state) - continue - self.state = self.CoordinatorState.STOPPING - - # Broadcast the control signal if we're in a good state. - # Forward the full deserialized payload so that data-bearing - # signals (e.g. SET_GENERATION_EPOCH) retain their arguments. - broadcast_payload = msgpack.packb(deserialized_payload, use_bin_type=True) - for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks): - self._send_to_engine(data_parallel_rank_id, broadcast_payload) - - # STOP affects engines; reset coordinator to RUNNING to allow future engines. - if header == Headers.STOP: - self.state = self.CoordinatorState.RUNNING - - elif header in (Headers.START_CUDA_PROFILER, Headers.STOP_CUDA_PROFILER): - # Profiler control: broadcast to every connected DP engine. Not a - # state transition, so no CoordinatorState checks — just forward. - if sender_identity not in known_clients: - logging.warning("Coordinator: ignoring profiler signal from unknown client.") - continue - broadcast_payload = msgpack.packb(deserialized_payload, use_bin_type=True) - for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks): - self._send_to_engine(data_parallel_rank_id, broadcast_payload) - - elif header == Headers.ENGINE_REPLY: - # This is the output of a single engine step on some data parallel rank. - assert sender_identity in self.identities_of_data_parallel_ranks - finished_requests = deserialized_payload[1] - - for finished_request in finished_requests: - self.detokenize(finished_request) - fid = finished_request["request_id"] - client_identity = self.request_id_to_client_id[fid] - client_request_identity = self.request_id_to_client_request_id[fid] - del self.request_id_to_client_id[fid] - del self.request_id_to_client_request_id[fid] - assigned_rank = self.request_id_to_rank.pop(fid, None) - if assigned_rank is not None: - idx = self.identity_to_rank_index.get(assigned_rank) - if idx is not None: - assert self._pending_counts[idx] >= 1 - self._pending_counts[idx] -= 1 - - self.router_socket.send_multipart( - [ - client_identity, - msgpack.packb( - [header.value, client_request_identity, finished_request], - use_bin_type=True, - ), - ] - ) - - elif header == Headers.SHUTDOWN: - if sender_identity not in known_clients: - logging.warning("Coordinator: ignoring signal from unknown client.") - continue + handler = self._handlers.get(header) + if handler is None: + raise UnknownHeaderError(header) + if handler(self, sender_identity, deserialized_payload): break - elif header == Headers.DISCONNECT: - if sender_identity in self.identities_of_data_parallel_ranks: - self._remove_engine(sender_identity) - - else: - raise UnknownHeaderError(header) + def _handle_rank_registration(self, sender_identity): + """Register a data parallel rank that connected to a running coordinator.""" + if sender_identity not in self.identities_of_data_parallel_ranks: + self.identities_of_data_parallel_ranks.append(sender_identity) + self._register_rank_identity(sender_identity) def detokenize(self, finished_request): """ diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py new file mode 100644 index 00000000000..34932825b34 --- /dev/null +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -0,0 +1,227 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Message handlers for the data parallel inference coordinator. + +Each handler is a free function decorated with @message_handler, which records +it in the module-level HANDLERS registry keyed by message header. The +coordinator builds its dispatch table from this registry, so a new message type +is supported simply by adding a decorated function here; the coordinator's event +loop never changes. + +Handlers have the signature ``(coordinator, sender_identity, payload) -> bool | None`` +where ``payload`` is the already-deserialized message. Returning a truthy value +signals the coordinator's event loop to stop. +""" + +import logging + +import torch + +from megatron.core.inference.config import PrefixCachingCoordinatorPolicy +from megatron.core.inference.headers import Headers + +from .state import CONTROL_TRANSITIONS, CoordinatorState + +try: + import msgpack +except ImportError: + msgpack = None + + +# Maps a message header value to the function that handles it. Populated by the +# @message_handler decorator at import time. +HANDLERS = {} + + +def message_handler(*headers): + """Register a function as the handler for one or more message headers. + + A new message type is supported by writing a handler function and decorating + it with the header(s) it serves; it is added to HANDLERS, which the + coordinator turns into its dispatch table. The event loop never needs to + change when a header is added. + """ + + def decorator(fn): + for header in headers: + assert header not in HANDLERS, f"duplicate handler for {header}" + HANDLERS[header] = fn + return fn + + return decorator + + +@message_handler(Headers.CONNECT) +def handle_connect(coordinator, sender_identity, payload): + """Handshake with a new client, replying with a CONNECT_ACK.""" + if sender_identity in coordinator.known_clients: + logging.info(f"Client {sender_identity} sent a duplicate connect request. Ignoring ..") + return + + coordinator.known_clients.add(sender_identity) + coordinator.router_socket.send_multipart( + [sender_identity, msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)] + ) + + +@message_handler(Headers.SUBMIT_REQUEST) +def handle_submit_request(coordinator, sender_identity, payload): + """Route a client request to a data parallel rank. + + Returns True (stopping the loop) if no engines are reachable. + """ + # ToDo [Siddharth]: We might want to tokenize the prompt on the + # assigned data parallel rank for this process instead + # of the coordinator. + + # Message from a known client + if sender_identity not in coordinator.known_clients: + logging.info(f"Received message from unknown client {sender_identity}. Ignoring.") + return + # this is a message from a client. + # route it to a data parallel rank + client_request_id, prompt, sampling_params = payload[1:] + # map client request_id to server request_id + # necessary because multiple clients might have the same request_id. + request_id = coordinator.next_request_id + coordinator.next_request_id += 1 + coordinator.request_id_to_client_id[request_id] = sender_identity + coordinator.request_id_to_client_request_id[request_id] = client_request_id + + # Serialize prompt. + if isinstance(prompt, (str, list)): + pass + elif isinstance(prompt, torch.Tensor): + prompt = prompt.tolist() + else: + raise Exception("specialize for <%s> prompt." % type(prompt).__name__) + + engine_payload = msgpack.packb( + [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params], use_bin_type=True + ) + + request_hashes = coordinator.compute_request_hashes(prompt) + if ( + coordinator.prefix_caching_coordinator_policy + == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK + ): + request_hashes = request_hashes[:1] + + # Account for the fact that some engines may have died. + for _ in range(len(coordinator.identities_of_data_parallel_ranks)): + next_identity = coordinator.get_best_data_parallel_rank(request_hashes) + if coordinator._send_to_engine(next_identity, engine_payload): + break + else: + # If all engines have died, we are in an abnormal state, and must exit cleanly. + logging.error("Coordinator: no reachable engines for request %d", request_id) + del coordinator.request_id_to_client_id[request_id] + del coordinator.request_id_to_client_request_id[request_id] + return True + + coordinator.request_id_to_rank[request_id] = next_identity + coordinator._pending_counts[coordinator.identity_to_rank_index[next_identity]] += 1 + if request_hashes: + coordinator._update_rank_hashes(next_identity, request_hashes) + if coordinator.schedule_records is not None: + coordinator.schedule_records.append( + { + "request_id": request_id, + "rank_index": coordinator.identity_to_rank_index[next_identity], + "num_hashes": len(request_hashes), + } + ) + + +@message_handler( + Headers.PAUSE, + Headers.UNPAUSE, + Headers.SUSPEND, + Headers.RESUME, + Headers.SET_GENERATION_EPOCH, + Headers.STOP, +) +def handle_control_signal(coordinator, sender_identity, payload): + """Validate a control signal against the transition table and broadcast it.""" + if sender_identity not in coordinator.known_clients: + logging.warning("Coordinator: ignoring signal from unknown client.") + return + + header = Headers(payload[0]) + transition = CONTROL_TRANSITIONS[header] + if coordinator.state not in transition.allowed_from: + # Silently ignore redundant signals; warn on genuinely invalid ones. + if coordinator.state not in transition.idempotent_in: + logging.warning("Coordinator: ignoring %s in state %s", header.name, coordinator.state) + return + if transition.new_state is not None: + coordinator.state = transition.new_state + + # Broadcast the control signal. Forward the full deserialized payload so + # that data-bearing signals (e.g. SET_GENERATION_EPOCH) retain their args. + coordinator._broadcast_to_engines(payload) + + # STOP affects engines; reset coordinator to RUNNING to allow future engines. + if header == Headers.STOP: + coordinator.state = CoordinatorState.RUNNING + + +@message_handler(Headers.START_CUDA_PROFILER, Headers.STOP_CUDA_PROFILER) +def handle_cuda_profiler_signal(coordinator, sender_identity, payload): + """Broadcast a CUDA profiler control signal to every connected DP engine. + + Profiler control is not a coordinator state transition, so there are no + CoordinatorState checks — the signal is simply forwarded to all engines. + """ + if sender_identity not in coordinator.known_clients: + logging.warning("Coordinator: ignoring profiler signal from unknown client.") + return + coordinator._broadcast_to_engines(payload) + + +@message_handler(Headers.ENGINE_REPLY) +def handle_engine_reply(coordinator, sender_identity, payload): + """Route completed requests from an engine back to their originating clients.""" + # This is the output of a single engine step on some data parallel rank. + assert sender_identity in coordinator.identities_of_data_parallel_ranks + finished_requests = payload[1] + + for finished_request in finished_requests: + coordinator.detokenize(finished_request) + fid = finished_request["request_id"] + client_identity = coordinator.request_id_to_client_id[fid] + client_request_identity = coordinator.request_id_to_client_request_id[fid] + del coordinator.request_id_to_client_id[fid] + del coordinator.request_id_to_client_request_id[fid] + assigned_rank = coordinator.request_id_to_rank.pop(fid, None) + if assigned_rank is not None: + idx = coordinator.identity_to_rank_index.get(assigned_rank) + if idx is not None: + assert coordinator._pending_counts[idx] >= 1 + coordinator._pending_counts[idx] -= 1 + + coordinator.router_socket.send_multipart( + [ + client_identity, + msgpack.packb( + [Headers.ENGINE_REPLY.value, client_request_identity, finished_request], + use_bin_type=True, + ), + ] + ) + + +@message_handler(Headers.SHUTDOWN) +def handle_shutdown(coordinator, sender_identity, payload): + """Stop the coordinator event loop on request from a known client.""" + if sender_identity not in coordinator.known_clients: + logging.warning("Coordinator: ignoring signal from unknown client.") + return + return True + + +@message_handler(Headers.DISCONNECT) +def handle_disconnect(coordinator, sender_identity, payload): + """Remove a disconnecting engine from the routing pool.""" + if sender_identity in coordinator.identities_of_data_parallel_ranks: + coordinator._remove_engine(sender_identity) diff --git a/megatron/core/inference/data_parallel_inference_coordinator/state.py b/megatron/core/inference/data_parallel_inference_coordinator/state.py new file mode 100644 index 00000000000..bb6c32104f8 --- /dev/null +++ b/megatron/core/inference/data_parallel_inference_coordinator/state.py @@ -0,0 +1,63 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""State machine definitions for the data parallel inference coordinator.""" + +from dataclasses import dataclass, field +from enum import Enum, auto + +from megatron.core.inference.headers import Headers + + +class CoordinatorState(Enum): + """State machine for the coordinator.""" + + RUNNING = auto() + PAUSED = auto() + SUSPENDED = auto() + STOPPING = auto() + + +_ALL_STATES = frozenset(CoordinatorState) + + +@dataclass(frozen=True) +class ControlTransition: + """A single rule in the control-signal state machine. + + Attributes: + allowed_from: States the signal may be applied from. + new_state: State to move to once applied, or None to leave the state + unchanged (e.g. a pure broadcast such as SET_GENERATION_EPOCH). + idempotent_in: States in which the signal is a silent no-op rather than + a logged rejection (e.g. a redundant PAUSE while already paused). + """ + + allowed_from: frozenset + new_state: CoordinatorState | None + idempotent_in: frozenset = field(default_factory=frozenset) + + +# Control-signal state machine, expressed declaratively and consumed by the +# control-signal handler. +CONTROL_TRANSITIONS = { + Headers.PAUSE: ControlTransition( + allowed_from=frozenset({CoordinatorState.RUNNING}), + new_state=CoordinatorState.PAUSED, + idempotent_in=frozenset({CoordinatorState.PAUSED, CoordinatorState.SUSPENDED}), + ), + Headers.UNPAUSE: ControlTransition( + allowed_from=frozenset({CoordinatorState.PAUSED}), new_state=CoordinatorState.RUNNING + ), + Headers.SUSPEND: ControlTransition( + allowed_from=frozenset({CoordinatorState.PAUSED}), new_state=CoordinatorState.SUSPENDED + ), + Headers.RESUME: ControlTransition( + allowed_from=frozenset({CoordinatorState.SUSPENDED}), new_state=CoordinatorState.PAUSED + ), + Headers.STOP: ControlTransition( + allowed_from=frozenset({CoordinatorState.PAUSED, CoordinatorState.SUSPENDED}), + new_state=CoordinatorState.STOPPING, + ), + # No state change; broadcast in any state so engines stay in sync. + Headers.SET_GENERATION_EPOCH: ControlTransition(allowed_from=_ALL_STATES, new_state=None), +} diff --git a/megatron/core/inference/engines/async_zmq_communicator.py b/megatron/core/inference/engines/async_zmq_communicator.py index aa13f659d40..bba7508e08b 100644 --- a/megatron/core/inference/engines/async_zmq_communicator.py +++ b/megatron/core/inference/engines/async_zmq_communicator.py @@ -41,6 +41,12 @@ def __init__( hostname (str | None): Hostname or IP address to use for ZMQ socket binding. If None, defaults to socket.gethostname(). """ + # Normalize None to the default (world) group. get_rank/get_world_size + # already treat None this way, but get_process_group_ranks below does + # not accept None, so resolve it once here for all three calls. + if process_group is None: + process_group = dist.group.WORLD + self.rank = dist.get_rank(process_group) self.world_size = dist.get_world_size(process_group) self.is_leader = self.rank == 0 diff --git a/megatron/core/inference/headers.py b/megatron/core/inference/headers.py index 107a200818c..4acef82bc1c 100644 --- a/megatron/core/inference/headers.py +++ b/megatron/core/inference/headers.py @@ -1,32 +1,137 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -from enum import Enum, auto +"""Message headers for inference coordinator/engine/client communication. +Headers are grouped into category IntEnum classes by concern. Each category +occupies a distinct numeric range so that wire values never collide across +categories. Adding a new class of headers means adding a new category enum in +its own range and listing it in HEADER_ENUMS; the existing categories are left +untouched. + +On the wire a header travels as its integer value. decode_header maps an +integer back to the originating category member. +""" + +from enum import IntEnum + + +class Connection(IntEnum): + """Client <-> coordinator handshake.""" + + CONNECT = 0 + CONNECT_ACK = 1 + + +class Request(IntEnum): + """Inference request submission and completion reply.""" + + SUBMIT_REQUEST = 20 + ENGINE_REPLY = 21 + + +class Control(IntEnum): + """Runtime control signals broadcast to engines.""" + + PAUSE = 40 + UNPAUSE = 41 + SUSPEND = 42 + RESUME = 43 + SET_GENERATION_EPOCH = 44 + STOP = 45 + START_CUDA_PROFILER = 46 + STOP_CUDA_PROFILER = 47 + + +class Lifecycle(IntEnum): + """Process lifecycle signals.""" + + DISCONNECT = 60 + SHUTDOWN = 61 + + +class Transport(IntEnum): + """Low-level transport framing.""" + + TP_BROADCAST = 80 -class Headers(Enum): - """ - Enum representing headers used for communication with the inference-coordinator. - """ - CONNECT = auto() - CONNECT_ACK = auto() - SUBMIT_REQUEST = auto() - ENGINE_REPLY = auto() - PAUSE = auto() - UNPAUSE = auto() - SUSPEND = auto() - RESUME = auto() - SET_GENERATION_EPOCH = auto() - STOP = auto() - DISCONNECT = auto() - SHUTDOWN = auto() - TP_BROADCAST = auto() - START_CUDA_PROFILER = auto() - STOP_CUDA_PROFILER = auto() +# All header categories. To add a new class of headers, define a new IntEnum in +# its own (disjoint) numeric range and append it here; nothing else needs to +# change. Ranges are spaced out to leave room for growth within each category. +HEADER_ENUMS = (Connection, Request, Control, Lifecycle, Transport) class UnknownHeaderError(Exception): - """A signal with an unrecognized header was received by the coordinator.""" + """A signal with an unrecognized header was received.""" def __init__(self, header): super().__init__(f"specialize for {header}.") + + +def _build_tables(): + """Index every header by wire value and by name, asserting no collisions.""" + by_value = {} + by_name = {} + for enum_cls in HEADER_ENUMS: + for member in enum_cls: + if member.value in by_value: + raise ValueError( + f"duplicate header wire value {member.value}: " + f"{by_value[member.value]!r} and {member!r}" + ) + if member.name in by_name: + raise ValueError( + f"duplicate header name {member.name!r}: " + f"{by_name[member.name]!r} and {member!r}" + ) + by_value[member.value] = member + by_name[member.name] = member + return by_value, by_name + + +_CODE_TO_MEMBER, _NAME_TO_MEMBER = _build_tables() + + +def decode_header(value): + """Resolve an integer wire value to its category header member. + + Args: + value (int): The integer header value read off the wire. + + Returns: + The corresponding category enum member (e.g. ``Control.PAUSE``). + + Raises: + UnknownHeaderError: if no header is registered for ``value``. + """ + try: + return _CODE_TO_MEMBER[value] + except KeyError: + raise UnknownHeaderError(value) + + +class _Headers: + """Flat, read-only union over every header category. + + Lets callers use a single name without caring which category a header lives + in: attribute access (``Headers.PAUSE``) resolves to the underlying category + member (``Control.PAUSE``), and calling it (``Headers(value)``) decodes a + wire value via decode_header. Because both return the canonical category + members, equality and dict-key lookups against the categories work + unchanged. + """ + + def __call__(self, value): + return decode_header(value) + + def __getattr__(self, name): + try: + return _NAME_TO_MEMBER[name] + except KeyError: + raise AttributeError(f"no header named {name!r}") + + def __iter__(self): + return iter(_NAME_TO_MEMBER.values()) + + +Headers = _Headers()