# Writing Python Workers in Dynamo This guide explains how to create your own Python worker in Dynamo and deploy it via `dynamo serve` or `dynamo deploy`, covering basic concepts as well as advanced features like enabling KV routing and disaggregated serving. For detailed information about `dynamo serve` infrastructure, see the [Dynamo SDK Docs](../deploy/dynamo/sdk/docs/sdk/README.md). For a guide that walks through how to launch a vLLM-based worker with implementation of Disaggregated Serving and KV-Aware Routing included, see the [Dynamo Serve Guide](../docs/guides/dynamo_serve.md). ## Basic Concepts When deploying a python-based worker with `dynamo serve` or `dynamo deploy`, it is a Python class based definition that requires a few key decorators to get going: - `@service`: used to define a worker class - `@dynamo_endpoint`: marks methods that can be called by other workers or clients For more detailed information on these concepts, see the [Dynamo SDK Docs](../deploy/dynamo/sdk/docs/sdk/README.md). ### Worker Skeleton Here is the rough outline of what a worker may look like in its simplest form: ```python from dynamo.sdk import dynamo_endpoint, service @service( dynamo={ "enabled": True, "namespace": "your_namespace", }, ) class YourWorker: # Worker implementation # ... @dynamo_endpoint() async def your_endpoint(self, request: RequestType) -> AsyncIterator[ResponseType]: # Endpoint Implementation pass ``` Workers in Dynamo are identified by a `namespace/component/endpoint` naming schema. When addressing this worker's endpoint with the `namespace/component/endpoint` schema based on the definitions above, it would be: `your_namespace/YourWorker/your_endpoint`: - `namespace="your_namespace"`: Defined in the `@service` decorator - `component="YourWorker"`: Defined by the Python Class name - `endpoint="your_endpoint"`: Defined by the `@dynamo_endpoint` decorator, or by default the name of the function being decorated. For more details about service configuration, resource management, and dynamo endpoints, see the [Dynamo SDK Docs](../deploy/dynamo/sdk/docs/README.md). ### Request/Response Types Request/Response types of endpoints can be defined arbitraily for your use case's needs, as long as the client calling your worker matches the expectations. Define your request and response types using Pydantic models: ```python from pydantic import BaseModel class RequestType(BaseModel): text: str # Add other fields as needed class ResponseType(BaseModel): text: str # Add other fields as needed ``` For example, if you deploy your worker directly behind an OpenAI HTTP (`http`) service using `llmctl`, you can define the request and response types to correspond to Chat Completions objects, such as the ones specified in the OpenAI API. For example: ```python from vllm.entrypoints.openai.protocol import ChatCompletionRequest class YourLLMWorker: @dynamo_endpoint(name="my_chat_completions_endpoint") async def generate(self, request: ChatCompletionRequest): # Endpoint Implementation pass ``` ## Basic Worker Example Here's a simple example of a worker that takes text in and streams text out via custom RequestType/ResponseType definitions: ```python # basic_worker.py # This can be run standalone with `dynamo serve basic_worker:YourWorker` from pydantic import BaseModel from dynamo.sdk import dynamo_endpoint, service logger = logging.getLogger(__name__) class RequestType(BaseModel): text: str class ResponseType(BaseModel): text: str @service( dynamo={ "enabled": True, "namespace": "your_namespace", } ) class YourWorker: def __init__(self) -> None: logger.info("Starting worker...") @dynamo_endpoint() async def generate(self, request: RequestType): """Generate tokens and stream them back""" logger.info(f"Worker endpoint received: {request.text}") for token in request.text.split(): yield ResponseType(text=token).model_dump_json() ``` To see a minimal worker example like the above used in a larger pipeline of components, see the `dynamo serve` [Hello World example](../examples/hello_world). ### Client Example Here's a simple example of a client that directly calls the example worker above through Dynamo without any intermediate services: ```python import asyncio from pydantic import BaseModel from dynamo.runtime import dynamo_worker, DistributedRuntime # These could also be imported from a shared file/definition class RequestType(BaseModel): text: str class ResponseType(BaseModel): text: str @dynamo_worker() async def client_worker(runtime: DistributedRuntime): # Get a client to the worker endpoint from the distributed runtime client = await runtime.namespace("your_namespace").component("YourWorker").endpoint("generate").client() # Create a request request = RequestType(text="Hello, Dynamo!") # Call the dynamo endpoint exposed by the worker responses = await client.generate(request.model_dump_json()) async for response in responses: print(response) if __name__ == "__main__": asyncio.run(client_worker()) ``` If there is an OpenAI `http` service in front of your worker and the worker is defined to accept ChatCompletions-like requests, you could also use an OpenAI-based client (or `curl`) that sends requests to the OpenAI HTTP Service, and internally these requests would be routed to the attached worker endpoints instead. In more advanced scenarios where your worker may operate on some other intermediate format that may not directly match an OpenAI-like format, you could setup a separate processor worker that does something like the following: - Take in OpenAI Chat Completions requests from the HTTP service - Convert requests from Chat Completions format to the RequestType format your worker expects - Forward requests to the worker(s) - Convert responses from the worker's ResponseType back into Chat Completions response format - Forward responses back to client This advanced scenario of a separate [OpenAI Processor worker](../examples/llm/components/processor.py) is demonstrated in this [vLLM example](../examples/llm/). For a more minimal example of deploying a pipeline of components with a custom API that your client can communicate with, see the [Hello World example](../examples/hello_world). ## Advanced Features ### KV Routing for LLMs KV-aware routing is a powerful feature of Dynamo that optimizes for routing requests to specific workers while minimizing a specific KV-cache based cost function. In its simplest form, all a worker needs to do to enable KV-aware routing is to publish KV metrics through the `KvMetricsPublisher`, which will be consumed by a Dynamo KV Router through the `KvMetricsAggregator`: ```python # kv_metrics_worker.py # This can be run standalone with `dynamo serve kv_metrics_worker:YourWorker` import logging import random from pydantic import BaseModel from dynamo.llm import KvMetricsPublisher from dynamo.sdk import dynamo_endpoint, service, dynamo_context logger = logging.getLogger(__name__) class RequestType(BaseModel): text: str class ResponseType(BaseModel): text: str @service( dynamo={ "enabled": True, "namespace": "your_namespace", } ) class YourWorker: def __init__(self): # Initialize metrics publisher from Dynamo self.component = dynamo_context["component"] self.metrics_publisher = KvMetricsPublisher() # Register an endpoint for consumers of the KV Metrics # (KvMetricsAggregator) to listen/gather on. self.metrics_publisher.create_endpoint(self.component) # Initialize some metrics for the worker/class to track self.request_active_slots = 0 self.request_total_slots = 1024 self.kv_active_blocks = 0 self.kv_total_blocks = 1024 self.num_requests_waiting = 0 self.gpu_cache_usage_perc = 0.0 self.gpu_prefix_cache_hit_rate = 0.0 # Publish some initial metrics to register # this worker as a candidate for KV Routing. self.metrics_publisher.publish( self.request_active_slots, self.request_total_slots, self.kv_active_blocks, self.kv_total_blocks, self.num_requests_waiting, self.gpu_cache_usage_perc, self.gpu_prefix_cache_hit_rate, ) def publish_kv_metrics(self): # Populate the frequently changing metrics with random data for # demonstration. These values should be tracked by the implementation, # or queried from the underlying inference framework. self.kv_active_blocks = random.randint(0, 1024) self.num_requests_waiting = random.randint(0, 100) self.gpu_cache_usage_perc = random.uniform(0, 1.0) self.gpu_prefix_cache_hit_rate = random.uniform(0, 1.0) # Publish the metrics with the current state self.metrics_publisher.publish( self.request_active_slots, self.request_total_slots, self.kv_active_blocks, self.kv_total_blocks, self.num_requests_waiting, self.gpu_cache_usage_perc, self.gpu_prefix_cache_hit_rate, ) @dynamo_endpoint() async def generate(self, request: RequestType): """Generate tokens, update KV Cache metrics, and stream the tokens back""" # Increment the number of active requests on receiving one self.request_active_slots += 1 logger.info(f"Worker endpoint received: {request.text}") # Simulate each step of token generation for token in request.text.split(): # Update the metrics with the current state self.publish_kv_metrics() print("Returning token:", token) yield ResponseType(text=token).model_dump_json() # Decrement the number of active requests when complete with one self.request_active_slots -= 1 ``` The granularity at which metrics are published is up to the backend/worker implementation. For example, you may want to update metrics on every single generation step during token generation, or you may only want to update once per request, depending on your use case. Assuming long generation time or long output token sequence lengths, it would be more accurate to publish metrics at every generation step. With the worker publishing KV metrics, you should now be able to connect it to a KV Router through the `KvMetricsAggregator`. These metrics can then be inputs to a cost function to determine which of the available worker's the request should be routed to. For a [python-based KV Router](../examples/llm/components/kv_router.py) implementation, the router is like any other worker, and it can expose an endpoint that can do arbitrary things based on your use case. For example, you can initialize the `KvMetricsAggregator` and `KvIndexer` in your class implementation: ```python @service( dynamo={ "enabled": True, "namespace": "your_namespace", }, ) class Router: # ... @async_on_start async def async_init(self): self.runtime = dynamo_context["runtime"] # Initialize a listener/aggregator for collecting KV metrics # from the specified component (workers) publishing them kv_listener = self.runtime.namespace("your_namespace").component("YourWorker") await kv_listener.create_service() self.indexer = KvIndexer(kv_listener, self.args.block_size) self.metrics_aggregator = KvMetricsAggregator(kv_listener) ``` With this flexibility, you can also define your own cost function that takes the KV metrics from the KvMetricsAggregator above and the set of available workers as inputs, and returns which worker ID that the request should be routed to. Since the router is like any other worker in this context, you can also track your own custom metrics and use them in your cost function: ```python @service( dynamo={ "enabled": True, "namespace": "your_namespace", }, ) class Router: # ... def _cost_function( self, scores: OverlapScores | None, metrics: AggregatedMetrics | None, token_length: int, custom_metrics: dict = {}, ): """ Args: scores (OverlapScores | None): The number of matching blocks between the request and the prefix cache of each worker. metrics (AggregatedMetrics | None): Several worker metrics polled by the `KvMetricsAggregator`, currently including the GPU cache usage, number of waiting requests, and the GPU prefix cache hit rate. token_length (int): The number of tokens in the request. custom_metrics (dict): Arbitrary (optional) data from your implementation. Returns: worker_id (str): The best worker ID based on the inputs. """ # This is a dummy implementation for demonstration purposes, see the # llm/tensorrt_llm/hello_world examples for more realistic implementations. worker_ids = [] # KV cache block hit scores for worker_id, score in scores.scores.items(): print(f"{worker_id} # of matching KV Blocks of size {self.indexer.block_size()}: {score}") worker_ids.append(worker_id) # Aggregated KVMetrics published by workers for endpoint_metrics in metrics.endpoints: print(f"Endpoint metrics: {endpoint_metrics}") # Replace this random choice with your custom criteria to # select the best worker ID. best_worker_id = random.choice(worker_ids) return best_worker_id @dynamo_endpoint() async def generate(self, request: Tokens) -> AsyncIterator[WorkerId]: try: # lora_id is a placeholder for lora support, but not used in this example lora_id = 0 scores = await self.indexer.find_matches_for_request( request.tokens, lora_id ) except Exception as e: scores = {} print(f"Error finding matches: {e}") # Get published/aggregated KV Metrics metrics = await self.metrics_aggregator.get_metrics() # (Optional) Add custom metrics to consider in the cost function # The types and data used here are fully up to your implementation custom_metrics = {"my_custom_metric": 42} # Call custom cost function worker_id = self._cost_function( scores, metrics, len(request.tokens), custom_metrics ) # Return worker ID selected from cost function yield f"{worker_id}" ``` Similarly, for running a Rust-based Router as a standalone binary rather than as a Python Worker, see the [WorkerSelector Trait](../lib/llm/src/kv_router.rs) trait, and the [Router Component](../components/router/src/main.rs). For more details on receiving and routing based on the worker's published KV metrics, see the [KV Cache Routing Guide](../docs/kv_cache_routing.md). ### Disaggregated Serving #### NIXL NIXL (NVIDIA Inter-process Link) enables efficient GPU memory sharing between processes. In Prefill/Decode disaggregation, we use NIXL to transfer computed KV cache blocks from prefill workers to decode workers. Here are the core concepts: 1. **NIXL Agent Setup** ```python from nixl._api import nixl_agent class NixlConnector: def __init__(self, engine_id: str, rank: int): # Create unique NIXL agent for this worker self.nixl_wrapper = nixl_agent(str(uuid.uuid4()), None) self.engine_id = engine_id self.rank = rank self.block_len = None # Will be set during registration ``` 2. **Memory Registration and Transfer Preparation** ```python def register_kv_caches(self, kv_cache: torch.Tensor): # Get block size from the KV cache tensor # Note: KV cache layout depends on specific attention implementation num_blocks, block_size, num_heads, head_dim = kv_cache.shape self.block_len = block_size * num_heads * head_dim * kv_cache.element_size() self.num_blocks = num_blocks # Register KV cache tensor with NIXL for sharing base_addr = kv_cache.data_ptr() region_len = num_blocks * self.block_len caches_data = [(base_addr, region_len, self.rank, "")] # Register memory regions with NIXL descs = self.nixl_wrapper.get_reg_descs(caches_data, "VRAM") self.nixl_wrapper.register_memory(descs) # Prepare local side of the transfer blocks_data = [] for block_id in range(num_blocks): block_offset = block_id * self.block_len blocks_data.append((base_addr + block_offset, self.block_len, self.rank)) # Create transfer descriptors and prepare for transfers self.local_blocks_descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM") # Create transfer handle with block descriptors for future transfers self.local_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist("", self.local_blocks_descs) ``` 3. **Remote Agent Communication** ```python def get_agent_metadata(self): # Get metadata for sharing with other agents return self.nixl_wrapper.get_agent_metadata(), self.local_blocks_descs def add_remote_agent(self, engine_id: str, agent_metadata: bytes, remote_blocks_descs: bytes): # Connect to remote NIXL agent agent_name = self.nixl_wrapper.add_remote_agent(agent_metadata) # Prepare remote side transfer handle using provided block descriptors self.remote_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(agent_name, remote_blocks_descs) return agent_name # Example usage: # On decode worker: decode_metadata, decode_blocks_descs = nixl_connector.get_agent_metadata() # Share metadata with prefill worker through your communication channel # On prefill worker: nixl_connector.add_remote_agent(decode_engine_id, decode_metadata, decode_blocks_descs) ``` 4. **KV Cache Transfer** ```python def write_blocks(self, local_block_ids, remote_block_ids, notify_msg): # Initiate asynchronous transfer using block IDs # Block descriptors were specified during transfer preparation handle = self.nixl_wrapper.make_prepped_xfer( "WRITE", self.local_xfer_side_handle, local_block_ids, self.remote_xfer_side_handle, remote_block_ids, notify_msg ) status = self.nixl_wrapper.transfer(handle) # Example usage: # On prefill worker: nixl_connector.write_blocks([0, 3], [12, 16], "kv_transfer") ``` The NIXL connector provides: - GPU memory registration for sharing between processes - Connection establishment between Prefill and Decode workers - Efficient block-based KV cache transfers - Asynchronous transfer notifications For a complete implementation example using NIXL for disaggregated serving, see the [vLLM example](../examples/llm/README.md). #### Disaggregation in Dynamo Aside from the NIXL specifics above, at its core, disaggregation in Dynamo builds on the same concepts used for any Dynamo client<->worker or worker<->worker interaction over the DistributedRuntime. First you can define a worker for each as usual: ```python class DecodeWorker: # ... class PrefillWorker: # ... ``` Second, you decide when/how the (Decode) worker should decide to Prefill remotely (by calling a separate Prefill worker), or decide to simply do the Prefill itself. In some scenarios, it may be more efficient for the Decode worker to just do the Prefill itself rather than do the extra communication, such as if the input sequence length is below some small threshold. If you wanted to disable disaggregation, the DecodeWorker could just always do the Prefill step as well. ```python @service( dynamo={ "enabled": True, "namespace": "your_namespace", }, ) class DecodeWorker: def __init__(self): self.runtime = dynamo_context["runtime"] # Whether the decode worker should call a separate Prefill worker or not self.do_remote_prefill = True # Initialize client to PrefillWorker self.prefill_client = await self.runtime .namespace("your_namespace") .component("PrefillWorker") .endpoint("generate") .client() @dynamo_endpoint() async def generate(self, request): if self.do_remote_prefill: # Forward the request to the prefill worker prefill_response = await self.prefill_client.generate(...) # ... framework-specific decode logic ... @service( dynamo={ "enabled": True, "namespace": "your_namespace", }, ) class PrefillWorker: def __init__(self): # ... @dynamo_endpoint() async def generate(self, request): # ... framework-specific prefill logic ... ``` Depending on the load distribution of requests and number of Prefill/Decode worker instances, instead of directly forwarding requests to the Prefill worker endpoint, it may be advantageous to send Prefill requests into a queue that the Prefill workers can pull from on-demand instead. You can see an example of that [here](../examples/hello_world/disagg_skeleton/components/prefill_worker.py). For an introductory example on doing disaggregation with Dynamo using simple models, see [this example](../examples/hello_world/disagg_skeleton). For more information on Disaggregated Serving, see the [general guide](../docs/disagg_serving.md) and [performance tuning guide](../docs/guides/disagg_perf_tuning.md). ## Best Practices 1. **Resource Management**: Configure resource requirements based on your needs: ```python @service( resources={ "cpu": "10", "memory": "20Gi", "gpu": "1", } ) ``` 2. **Async Operations**: Use async/await for I/O operations: ```python @dynamo_endpoint() async def generate(self, request): # Use async operations for better performance result = await self.some_async_operation() ``` ## Additional Resources - Check the [examples](../examples/) directory for more detailed implementations - Refer to the [Dynamo SDK Docs](../deploy/dynamo/sdk/docs/sdk/README.md) for API details. - For Disaggregated Serving, see the [general guide](../docs/disagg_serving.md) and [performance tuning guide](../docs/guides/disagg_perf_tuning.md).