Is this page helpful?

core.inference.data_parallel_inference_coordinator.coordinator#

Module Contents#

Classes#

DataParallelInferenceCoordinator

Coordinates inference requests between clients and distributed model engines.

API#

class core.inference.data_parallel_inference_coordinator.coordinator.DataParallelInferenceCoordinator(
pipe_connection: multiprocessing.connection.Connection,
data_parallel_size: int,
tokenizer,
max_requests,
inference_coordinator_port: int | None = None,
deterministic_mode: bool = False,
block_size_tokens: int | None = None,
enable_prefix_caching: bool = False,
prefix_caching_coordinator_policy: megatron.core.inference.config.PrefixCachingCoordinatorPolicy = PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK,
prefix_caching_routing_alpha: float = 0.5,
schedule_output_path: str | None = None,
hostname: str | None = None,
)#

Coordinates inference requests between clients and distributed model engines.

This class acts as a central server. It uses a ZMQ ROUTER socket to manage communication flows between multiple clients and multiple data parallel ranks.

The coordinator’s main responsibilities are:

  1. Worker Registration: It waits for a specified number of data parallel ranks (representing distributed model instances) to connect and register themselves.

  2. Client Connection: It accepts connections from external clients, like InferenceClient, and performs a simple handshake.

  3. Request Forwarding: It receives inference requests from clients, assigns a unique server-side request ID, tokenizes the prompt, and forwards the request to one of the available data parallel ranks using load-balanced (and, when prefix caching is enabled, prefix-affinity-aware) routing.

  4. Response Routing: It receives completed results from the data parallel ranks and routes them back to the original client that made the request.

  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.

.. attribute:: router_socket

The central ZMQ ROUTER socket for all communication.

Type:

zmq.Socket

.. attribute:: data_parallel_size

The number of data parallel workers to expect.

Type:

int

.. attribute:: identities_of_data_parallel_ranks

A deque holding the ZMQ identities of connected data parallel instances, used for request routing.

Type:

deque

.. attribute:: request_id_to_client_id

Maps server-side request IDs to the ZMQ identity of the client that initiated the request.

Type:

dict

.. attribute:: request_id_to_client_request_id

Maps server-side request IDs to the original request ID provided by the client.

Type:

dict

.. attribute:: next_request_id

A counter for generating unique server-side request IDs.

Type:

int

Initialization

Initializes the inference coordinator.

This sets up the ZMQ context and a ROUTER socket, binding it to the given port. It then enters a blocking loop to wait for all expected data parallel ranks to connect before proceeding.

Parameters:

  • pipe_connection (Connection) – A connecting pipe to the parent process.

  • data_parallel_size (int) – The number of data parallel instances that are expected to connect.

  • tokenizer – The tokenizer to use for prompt tokenization and detokenization.

  • inference_coordinator_port (Optional[int]) – The TCP port number to bind the server to.

  • prefix_caching_routing_alpha (float) – Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * normalized_load.

  • max_requests (int) – Max concurrent requests per rank, used to compute normalized_load for prefix-aware scoring.

CoordinatorState#

None

get_least_loaded_data_parallel_rank()#

Selects the data parallel rank with the fewest in-flight requests.

Ties are broken by lowest rank index for deterministic behavior.

Returns:

The ZMQ identity of the least-loaded data parallel rank.

Return type:

bytes

_register_rank_identity(identity)#

Register a new rank identity in the scoring data structures.

Called when a rank dynamically connects to a running coordinator (e.g. in tests that spawn the coordinator with data_parallel_size=0 and let engines register after the fact).

_remove_engine(identity)#

Remove a disconnected engine from all routing bookkeeping. Called both during shutdown and when an engine becomes unreachable mid-operation (e.g. zmq.EHOSTUNREACH in _send_to_engine). The O(n) index-shifting and hash-table rebuild are acceptable because the number of connected engines is small; optimize only if dynamic registration/deregistration at high engine counts becomes a use case.

_send_to_engine(identity, payload)#

Send payload to an engine, removing it from the pool if unreachable.

Returns:

True if the send succeeded, False if the engine was unreachable and removed.

_broadcast_to_engines(payload)#

Send a deserialized payload to every connected data parallel rank.

compute_request_hashes(prompt)#

Compute block hashes for a prompt on CPU.

Parameters:

prompt – Either a string (to be tokenized) or a list of token IDs.

Returns:

List of integer block hashes, or empty list if prefix caching is disabled.

get_best_data_parallel_rank(request_hashes)#

Select the best DP rank based on prefix cache affinity and load.

Uses a scoring function: score = alpha * match + (1 - alpha) * normalized_load where match is a policy-dependent affinity score in [0, 1] (binary for first_prefix_block, normalized prefix depth for longest_prefix) and normalized_load = free_slots / max_requests (higher means more free capacity).

Parameters:

request_hashes – List of block hashes for the request.

Returns:

The ZMQ identity of the selected data parallel rank.

Return type:

bytes

_update_rank_hashes(rank_identity, request_hashes)#

Record that a rank owns the given hashes.

Parameters:

  • rank_identity – ZMQ identity of the target rank.

  • request_hashes – List of block hashes assigned to this rank.

_match_vector(hashes)#

Return (match, recency) vectors of shape (n_ranks,).

match is binary depth: (depth + 1) / len(hashes) for ranks that have the deepest cached block, 0 otherwise. recency is the raw assignment timestamp for each matching rank (0 for non-matching ranks).

For FIRST_PREFIX_BLOCK the caller already truncates hashes to a single element, so the same logic yields a binary 0/1 match score.

start()#

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 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.

_handle_rank_registration(sender_identity)#

Register a data parallel rank that connected to a running coordinator.

detokenize(finished_request)#

Detokenizes the generated tokens in the finished request.

This method uses the coordinator’s tokenizer to convert the list of generated token IDs back into human-readable text.

Parameters:

finished_request (dict) – The serialized merged request containing the generated tokens to be detokenized. It is modified in place.

classmethod entrypoint(
pipe_connection: multiprocessing.connection.Connection,
ready_event: multiprocessing.Event,
data_parallel_size: int,
tokenizer,
max_requests,
inference_coordinator_port: int | None = None,
deterministic_mode: bool = False,
block_size_tokens: int | None = None,
enable_prefix_caching: bool = False,
prefix_caching_coordinator_policy: megatron.core.inference.config.PrefixCachingCoordinatorPolicy = PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK,
prefix_caching_routing_alpha: float = 0.5,
schedule_output_path: str | None = None,
hostname: str | None = None,
)#

Class method to instantiate and run the coordinator, for use in a separate process.

This method initializes the coordinator, signals a ready_event to indicate that it is fully initialized and listening, and then starts the main event loop.

Parameters:

  • pipe_connection (Connection) – A connecting pipe to the parent process.

  • ready_event (Event) – A threading or multiprocessing event object that is set() once the coordinator is ready to accept connections.

  • inference_coordinator_port (int) – The port to bind to.

  • data_parallel_size (int) – The number of expected data parallel instances.

  • deterministic_mode (bool) – Whether to enable deterministic scheduling.

  • block_size_tokens (Optional[int]) – Token block size for prefix caching hashing.

  • enable_prefix_caching (bool) – Whether prefix caching is enabled.

  • prefix_caching_coordinator_policy (PrefixCachingCoordinatorPolicy) – Routing policy.

  • schedule_output_path (Optional[str]) – Path to write scheduling decisions JSON.

  • prefix_caching_routing_alpha (float) – Weight for prefix-aware routing score.

  • max_requests (int) – Max concurrent requests per rank.

stop()#

Stops the inference coordinator, performing any necessary cleanup operations.