`core.inference.inference_client`#

Module Contents#

Classes#

InferenceClient

An asynchronous client for communicating with an inference coordinator service.

API#

class core.inference.inference_client.InferenceClient( inference_coordinator_address: str, deserialize: bool = False, block_size_tokens: Optional[int] = None, prefix_caching_coordinator_policy=None, )#

An asynchronous client for communicating with an inference coordinator service.

This client uses ZeroMQ (ZMQ) for messaging and MessagePack (msgpack) for serialization. It is designed to work within an asyncio event loop. It can submit inference requests, listen for completed results, and send control signals (e.g., pause, stop) to the inference engines.

The client operates by connecting a ZMQ DEALER socket to the inference coordinator’s ROUTER socket. Requests are sent with a unique ID, and an asyncio.Future is created for each request. A background task listens for replies from the coordinator, and when a reply is received, it resolves the corresponding future with the result.

.. attribute:: context

The ZeroMQ context.

Type:: zmq.Context

.. attribute:: socket

The ZMQ DEALER socket used for communication.

Type:: zmq.Socket

.. attribute:: completion_futures

A dictionary mapping request IDs to the asyncio Future objects that will hold the results.

Type:: dict[int, asyncio.Future]

.. attribute:: next_request_id

A counter for generating unique request IDs.

Type:: int

.. attribute:: listener_task

The background task that listens for completed requests.

Type:: asyncio.Task

Initialization

Initializes the InferenceClient.

Parameters:

inference_coordinator_address (str) – The address on which the inference coordinator is listening.
deserialize (bool) – If True, deserialize completed requests into DynamicInferenceRequest objects. If False (default), return the raw serialized dict for lower overhead.
block_size_tokens (Optional[int]) – Token block size to hash prompts on. Must match the engine’s KV block size, or the hashes name blocks the engine never cached. None leaves hashing to the coordinator.
prefix_caching_coordinator_policy – The coordinator’s routing policy, which decides whether anyone reads the hashes at all.

_block_hashes(prompt, media_meta)#

Hash the prompt into per-block routing hashes for the coordinator.

Hashed here rather than at the coordinator because clients are many against its one serial loop, and because this is the only place holding both inputs at once: the tokens, and – for multimodal – the media key that must salt them, which serialize_multimodal_data has just derived. Salting is what stops two requests whose media placeholders tokenize identically from sharing KV computed for different images.

Returns None to mean “this client did not hash”, which tells the coordinator to hash the prompt itself. Deliberately distinct from an empty list, which means it did hash and the prompt was shorter than one block: re-hashing that would pay the prompt decode for every short request, which is the cost this whole split exists to avoid.

A client told neither the block size nor the policy cannot hash, so it says None rather than “nothing matched”. Callers that build an InferenceClient directly configure neither, and reporting an empty list for them would silently turn prefix-affinity routing into load balancing, with nothing raising to say so.

add_request( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, *, multi_modal_data=None, offload_params: Optional[dict] = None, ) → asyncio.Future#

Submits a new inference request to the coordinator.

This method sends the prompt and sampling parameters to the inference coordinator. It immediately returns an asyncio.Future, which can be awaited to get the result of the inference request when it is complete.

Parameters:

prompt (str) – The input prompt to send to the language model.
sampling_params – An object containing the sampling parameters for text generation (e.g., temperature, top_p). It must have a serialize() method.
multi_modal_data –
Optional vLLM-style modality dictionary.

Images: "image" accepts raw image bytes, a list of raw image bytes, or a preprocessed image tensor dictionary. Video: "video" accepts raw video bytes, a list of raw video bytes, or a preprocessed video tensor dictionary. Audio: Audio does not yet have any supported data preprocessing or modeling formats.
offload_params – Opaque JSON/msgpack-compatible metadata forwarded to the engine’s payload stager.

Returns:

A future that will be resolved with a DynamicInferenceRequest object (if deserialize=True) or a raw serialized dict (if deserialize=False) containing the completed result.

Return type:

asyncio.Future

add_request_with_id( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, *, multi_modal_data=None, offload_params: Optional[dict] = None, ) → tuple[int, asyncio.Future]#

Submit a request and return its id alongside its completion future.

Same submission as add_request, which delegates here. The id is what abort_request takes, so a caller that may need to cancel – an HTTP handler whose client can disconnect mid-generation, for instance – has to use this form. With only the future in hand there is no way to name the request to the coordinator, and cancelling the future alone leaves the engine generating.

Parameters:

prompt – A string or list of token IDs.
sampling_params – Sampling parameters for the request.
multi_modal_data – Optional vLLM-style modality dictionary; see add_request.

Returns:

The request id and its completion future.

Return type:

tuple[int, asyncio.Future]

_pack_submit_frames( request_id, prompt, sampling_params, multi_modal_data, *, offload_params=None, )#

Build the multipart frames for a SUBMIT_REQUEST.

Shared by the blocking and streaming submit paths so the wire format is defined once.

Five frames, each with a different contract:

0 metadata Decoded and repacked by the coordinator on every request, so everything in it must be bounded. Carries the sampling params and the media descriptor, whose media_cache_key the coordinator reads for media affinity and, when nobody hashed upstream, as the salt for hashing there. 1 prompt Never decoded by the coordinator, forwarded to the engine verbatim. Skipping that decode is the point of the split: it is one serial loop shared by every rank, and the decode grows with prompt length. 2 block hashes Decoded by the coordinator, consumed for routing, and not forwarded – the engine computes its own KV hashes. One int per block rather than one per token, which is why it does not belong in frame 0. None means this client did not hash; see :meth:_block_hashes. 3 media Never decoded by the coordinator, forwarded to the engine verbatim. Same reasoning as the prompt but a larger payload: raw image or video bytes, or serialized preprocessed tensors. 4 offload params Never decoded by the coordinator, forwarded to the engine verbatim. Opaque client-supplied metadata for the engine’s prompt preparer and payload stager, so it is unbounded and cannot share frame 0. Decoded on MP rank 0 by the prompt preparer before the broadcast, and by every rank at admission.

The frame count is fixed rather than varying with media or offload params, so a malformed submission is caught by an arity check at the coordinator; a text-only request without params pays one byte each for a None media frame and a None offload frame.

Returns:: The frames to send, in wire order.
Return type:: list

static _pack_prompt(prompt)#

Pack a prompt into its own frame.

Coercion happens here rather than at the coordinator: the coordinator no longer decodes the prompt, and clients are many against its one serial loop, so this is both the only place that still sees the object and the cheaper place to normalize it.

_make_kv_handoff_request( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, kv_meta: dict, src_block_ids: List[int], ) → tuple[int, list]#

Allocate an ID and build a decode request carrying remote KV metadata.

Framed as [metadata, prompt, src_block_ids].

Nothing whose size follows the sequence length belongs in the metadata frame, since that frame is the only one the coordinator decodes. src_block_ids names one block per block_size_tokens of prompt, so it grows with the prompt and travels as its own body. kv_meta stays in the metadata: it is the peer’s NIXL agent/layout export, bounded by TP size and num_speculative_tokens, not by prompt length.

add_request_with_kv_handoff( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, kv_meta: dict, src_block_ids: List[int], ) → asyncio.Future#

Submit a request with remote KV metadata.

The decode engine allocates local blocks, pulls the KV from the prefill peer described by kv_meta, then begins generation.

Parameters:

prompt – A string or list of token IDs.
sampling_params – Sampling parameters for the decode request.
kv_meta – Metadata identifying the remote KV buffers.
src_block_ids – Remote block IDs containing the request’s KV state.

Returns:

A future that resolves to the completed request.

Return type:

asyncio.Future

add_request_with_kv_handoff_streaming( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, kv_meta: dict, src_block_ids: List[int], ) → megatron.core.inference.async_stream.AsyncStream[dict]#

Submit a streaming request with remote KV metadata.

Returns the same per-step partial/final iterator as

Meth:

add_request_streaming.

Parameters:

prompt – A string or list of token IDs.
sampling_params – Sampling parameters for the decode request.
kv_meta – Metadata identifying the remote KV buffers.
src_block_ids – Remote block IDs containing the request’s KV state.

Returns:

Per-step partial and final reply frames.

Return type:

AsyncStream[dict]

release_handoff(request_id: int) → None#

Tell the coordinator to release the KV blocks pinned for request_id.

Fire-and-forget. The coordinator broadcasts RELEASE_KV to every engine; engines without that request_id ignore the message.

abort_request(request_id: int) → None#: Cancel an in-flight request and close its local response stream.

add_request_streaming( prompt: Union[str, List[int]], sampling_params: megatron.core.inference.sampling_params.SamplingParams, *, multi_modal_data=None, offload_params: Optional[dict] = None, ) → megatron.core.inference.async_stream.AsyncStream[dict]#

Submit a streaming inference request.

Used by Dynamo directly and by the OpenAI-compatible HTTP frontend.

Returns an async iterator that yields incremental output dictionaries:

{"partial": {"request_id": int, "new_tokens": list[int]}} whenever the request’s streaming interval is reached, in order.
{"final": <full reply dict or DynamicInferenceRequest>} exactly once at the end. The iterator then stops.

sampling_params.streaming is forced to True before submission so the engine knows to emit ENGINE_REPLY_PARTIAL frames for this request.

Parameters:

prompt – A string or list of token IDs.
sampling_params – Sampling parameters. streaming is set to True in-place.
multi_modal_data –
Optional vLLM-style modality dictionary.

Images: "image" accepts raw image bytes, a list of raw image bytes, or a preprocessed image tensor dictionary. Video: "video" accepts raw video bytes, a list of raw video bytes, or a preprocessed video tensor dictionary. Audio: Audio does not yet have any supported data preprocessing or modeling formats.
offload_params – Opaque JSON/msgpack-compatible metadata forwarded to the engine’s payload stager.

Returns:

Per-step partial and final reply frames.

Return type:

AsyncStream[dict]

_submit_request(frames: list, request_id: int) → asyncio.Future#: Send a prepared request and register its completion future.

_submit_stream( frames: list, request_id: int, ) → megatron.core.inference.async_stream.AsyncStream[dict]#: Send a prepared streaming request and register its response stream.

async _recv_task()#

Listens for completed inference requests from the coordinator.

This coroutine runs in an infinite loop, continuously polling the socket for data. When a request reply is received, it unpacks the message, finds the corresponding Future using the request ID, and sets the result. Other control packets are handled appropriately.

This method is started as a background task by the start() method.

_connect_with_inference_coordinator( timeout_seconds: Optional[float] = None, )#

Performs the initial handshake with the inference coordinator.

Sends a CONNECT signal and waits for a CONNECT_ACK reply to ensure the connection is established and acknowledged by the coordinator.

start( loop: Optional[asyncio.AbstractEventLoop] = None, connect_timeout_seconds: Optional[float] = None, )#

Connects to the coordinator and starts the background listener task.

This must be called before submitting any requests. It handles the initial handshake and spawns the listen_for_completed_requests coroutine.

_send_signal_to_engines(signal, *args)#

Sends a generic control signal to the inference coordinator.

Parameters:

signal – The signal to send, typically a value from the Headers enum.
*args – Optional extra values to include in the payload.

pause_engines()#

Sends PAUSE to all engines via coordinator.

The coordinator broadcasts PAUSE. Each engine reaches EP consensus, then synchronizes via a world-wide barrier before transitioning to PAUSED. Callers should await engine.paused for confirmation.

unpause_engines() → None#: Sends UNPAUSE to all engines. No synchronization needed.

start_cuda_profiler() → None#

Sends START_CUDA_PROFILER to all engines via coordinator.

Each engine calls torch.cuda.profiler.start() (cudaProfilerStart) on its next loop iteration, so an outer nsys profile --capture-range= cudaProfilerApi begins recording. No synchronization needed.

stop_cuda_profiler() → None#: Sends STOP_CUDA_PROFILER to all engines (cudaProfilerStop).

set_generation_epoch(generation_epoch: int)#

Sends a signal to stamp all in-flight requests with the given generation epoch.

Parameters:: generation_epoch – The current generation epoch number.

suspend_engines()#

Sends SUSPEND to all engines via coordinator. Requires PAUSED.

Callers should await engine.suspended for confirmation.

resume_engines()#

Sends RESUME to all engines via coordinator. Requires SUSPENDED.

Callers should await engine.paused (or engine.running after UNPAUSE) for confirmation.

stop_engines()#

Sends STOP to all engines via coordinator. Requires PAUSED or SUSPENDED.

Callers should await engine.stopped for confirmation. Does not affect the coordinator.

shutdown_coordinator()#

Tells the coordinator process to exit its main loop.

Does not affect the engines.

stop()#

Stops the client and cleans up all resources.

This method cancels the background listener task, closes the ZMQ socket, and terminates the ZMQ context. It should be called when the client is no longer needed to ensure a graceful shutdown.

core.inference.inference_client#

Module Contents#

Classes#

API#

`core.inference.inference_client`#