Router Examples

1
import asyncio
2
from dynamollm import DistributedRuntime, KvPushRouter, KvRouterConfig
3
4
async def main():
5
    # Get runtime and create endpoint
6
    runtime = DistributedRuntime.detached()
7
    namespace = runtime.namespace("dynamo")
8
    component = namespace.component("backend")
9
    endpoint = component.endpoint("generate")
10
11
    # Create KV router
12
    kv_router_config = KvRouterConfig()
13
    router = KvPushRouter(
14
        endpoint=endpoint,
15
        block_size=16,
16
        kv_router_config=kv_router_config
17
    )
18
19
    # Your input tokens
20
    token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
21
22
    # Generate with per-request routing override
23
    stream = await router.generate(
24
        token_ids=token_ids,
25
        model="meta-llama/Llama-2-7b-hf",
26
        stop_conditions={
27
            "max_tokens": 20,        # Generate exactly 20 tokens
28
            "ignore_eos": True,      # Don't stop at EOS token
29
        },
30
        sampling_options={
31
            "temperature": 0.7,
32
            "top_p": 0.9,
33
        },
34
        router_config_override={
35
            "overlap_score_weight": 2.0,    # Prioritize cache hits for this request
36
            "router_temperature": 0.5,       # Add routing randomness
37
        }
38
    )
39
40
    # Collect generated tokens
41
    generated_tokens = []
42
    async for response in stream:
43
        if isinstance(response, dict) and "token_ids" in response:
44
            generated_tokens.extend(response["token_ids"])
45
46
    print(f"Generated {len(generated_tokens)} tokens: {generated_tokens}")
47
48
if __name__ == "__main__":
49
    asyncio.run(main())

1
stream = await router.generate(token_ids=tokens, model="model-name")

1
worker_id, _dp_rank, overlap = await router.best_worker(tokens, request_id="req-123")
2
response = await client.generate(tokens, request_id="req-123")
3
# await anext(response)  # Get first token
4
await router.mark_prefill_complete("req-123")  # After first token
5
# async for _ in response:  # Continue generating
6
#     ...
7
await router.free("req-123")  # After completion

1
# Probe multiple routers without updating state
2
worker_id_1, dp_rank, overlap_1 = await router_1.best_worker(tokens)  # No request_id
3
worker_id_2, dp_rank, overlap_2 = await router_2.best_worker(tokens)
4
5
# Pick the best router based on results
6
chosen_router = router_1 if overlap_1 > overlap_2 else router_2
7
stream = await chosen_router.generate(tokens, model="model-name", worker_id=worker_id)

1
loads = await router.get_potential_loads(tokens)
2
# Apply custom logic (e.g., weighted scoring, constraints)
3
best_worker = min(loads, key=lambda x: custom_cost_fn(x))
4
stream = await router.generate(tokens, model="model-name", worker_id=best_worker['worker_id'])

1
import asyncio
2
from dynamo.llm import DistributedRuntime, KvPushRouter, KvRouterConfig
3
4
async def minimize_ttft_routing():
5
    # Setup router
6
    runtime = DistributedRuntime.detached()
7
    namespace = runtime.namespace("dynamo")
8
    component = namespace.component("backend")
9
    endpoint = component.endpoint("generate")
10
11
    router = KvPushRouter(
12
        endpoint=endpoint,
13
        block_size=16,
14
        kv_router_config=KvRouterConfig()
15
    )
16
17
    # Your input tokens
18
    token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
19
20
    # Get potential loads for all workers
21
    potential_loads = await router.get_potential_loads(token_ids)
22
23
    # Find worker with minimum prefill tokens (best for TTFT)
24
    best_worker = min(potential_loads, key=lambda x: x['potential_prefill_tokens'])
25
26
    print(f"Worker loads: {potential_loads}")
27
    print(f"Selected worker {best_worker['worker_id']} with {best_worker['potential_prefill_tokens']} prefill tokens")
28
29
    # Route directly to the selected worker
30
    stream = await router.generate(
31
        token_ids=token_ids,
32
        model="meta-llama/Llama-2-7b-hf",
33
        worker_id=best_worker['worker_id'],  # Force routing to optimal worker
34
        stop_conditions={"max_tokens": 20}
35
    )
36
37
    # Process response
38
    async for response in stream:
39
        if isinstance(response, dict) and "token_ids" in response:
40
            print(f"Generated tokens: {response['token_ids']}")
41
42
if __name__ == "__main__":
43
    asyncio.run(minimize_ttft_routing())

1
from dynamo.llm import KvEventPublisher
2
3
class CustomEnginePublisher:
4
    def __init__(self, component, worker_id: int, block_size: int, dp_rank: int = 0):
5
        self.block_size = block_size
6
        self.event_id = 0
7
        self.kv_publisher = KvEventPublisher(
8
            component=component,
9
            worker_id=worker_id,
10
            kv_block_size=block_size,
11
            dp_rank=dp_rank,
12
            enable_local_indexer=False,
13
        )
14
15
    def on_blocks_stored(self, token_ids: list[int], block_hashes: list[int],
16
                         lora_id: int = 0, parent_hash: int | None = None):
17
        """Call after KV cache blocks are allocated."""
18
        self.event_id += 1
19
        num_block_tokens = [self.block_size] * len(block_hashes)
20
        self.kv_publisher.publish_stored(
21
            event_id=self.event_id,
22
            token_ids=token_ids,
23
            num_block_tokens=num_block_tokens,
24
            block_hashes=block_hashes,
25
            lora_id=lora_id,
26
            parent_hash=parent_hash,
27
        )
28
29
    def on_blocks_removed(self, block_hashes: list[int]):
30
        """Call when KV cache blocks are evicted."""
31
        self.event_id += 1
32
        self.kv_publisher.publish_removed(event_id=self.event_id, block_hashes=block_hashes)

1
from dynamo.llm import register_llm
2
3
async def main():
4
    # Register your engine with Dynamo
5
    component, endpoint = await register_llm(
6
        model="my-model",
7
        generator=my_generate_fn,
8
    )
9
10
    # Initialize publisher
11
    publisher = CustomEnginePublisher(
12
        component=component,
13
        worker_id=endpoint.connection_id(),
14
        block_size=16,  # Match your engine's block size
15
    )
16
17
    # Hook into your engine's cache events
18
    def on_prefill_complete(request_id, token_ids, blocks):
19
        block_hashes = [block.hash for block in blocks]
20
        publisher.on_blocks_stored(token_ids=token_ids, block_hashes=block_hashes)
21
22
    def on_cache_eviction(evicted_blocks):
23
        block_hashes = [block.hash for block in evicted_blocks]
24
        publisher.on_blocks_removed(block_hashes=block_hashes)

1
from dynamo.llm import ZmqKvEventPublisher, ZmqKvEventPublisherConfig
2
3
# Configure the ZMQ subscriber
4
config = ZmqKvEventPublisherConfig(
5
    worker_id=endpoint.connection_id(),
6
    kv_block_size=block_size,
7
    zmq_endpoint="tcp://127.0.0.1:5557",  # Where your engine publishes
8
    zmq_topic="",                          # Subscribe to all topics
9
    enable_local_indexer=False,
10
)
11
12
# Create publisher - it automatically subscribes to ZMQ and forwards to NATS
13
kv_publisher = ZmqKvEventPublisher(
14
    component=component,
15
    config=config,
16
)

1
import zmq
2
import msgpack
3
import time
4
5
class ZmqKvEventPublisher:
6
    """Pure Python ZMQ publisher for KV events (vLLM-compatible format)."""
7
8
    def __init__(self, zmq_endpoint: str, kv_block_size: int, topic: str = ""):
9
        self.kv_block_size = kv_block_size
10
        self.topic = topic
11
        self.ctx = zmq.Context()
12
        self.socket = self.ctx.socket(zmq.PUB)
13
        self.socket.bind(zmq_endpoint)
14
        self.sequence = 0
15
        self.data_parallel_rank = 0
16
17
    def _to_signed_i64(self, value: int | None) -> int | None:
18
        if value is None:
19
            return None
20
        return value - 0x10000000000000000 if value > 0x7FFFFFFFFFFFFFFF else value
21
22
    def publish_stored(self, event_id: int, token_ids: list[int], num_block_tokens: list[int],
23
                       block_hashes: list[int], lora_id: int = 0, parent_hash: int | None = None):
24
        event = {
25
            "type": "BlockStored",
26
            "block_hashes": [self._to_signed_i64(h) for h in block_hashes],
27
            "parent_block_hash": self._to_signed_i64(parent_hash),
28
            "token_ids": token_ids,
29
            "block_size": self.kv_block_size,
30
            "lora_id": lora_id if lora_id != 0 else None,
31
        }
32
        self._publish_event(event)
33
34
    def publish_removed(self, event_id: int, block_hashes: list[int]):
35
        event = {"type": "BlockRemoved", "block_hashes": [self._to_signed_i64(h) for h in block_hashes]}
36
        self._publish_event(event)
37
38
    def publish_all_cleared(self):
39
        self._publish_event({"type": "AllBlocksCleared"})
40
41
    def _publish_event(self, event: dict):
42
        batch = [time.time(), [event], self.data_parallel_rank]
43
        payload = msgpack.packb(batch, use_bin_type=True)
44
        sequence_bytes = self.sequence.to_bytes(8, byteorder="big")
45
        self.sequence += 1
46
        self.socket.send_multipart([self.topic.encode(), sequence_bytes, payload])
47
48
    def shutdown(self):
49
        self.socket.close()
50
        self.ctx.term()