WIP: fix fn dp_rank, add TODO's

launch/dynamo-run/src/subprocess/vllm_v1_inc.py

@@ -23,7 +23,7 @@
 
 import uvloop
 from vllm.config import VllmConfig
-from vllm.distributed.kv_events import KVEventsConfig
+from vllm.distributed.kv_events import KVEventsConfig, ZmqEventPublisher
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.inputs import TokensPrompt
 from vllm.sampling_params import SamplingParams
@@ -68,7 +68,7 @@ class DynamoStatLoggerPublisher(StatLoggerBase):
 
     def __init__(self, component: Component, dp_rank: int) -> None:
         self.inner = WorkerMetricsPublisher()
-        self.inner.create_endpoint(component)
+        self.inner.create_endpoint(component, dp_rank=dp_rank)
         self.dp_rank = dp_rank
 
     def record(
@@ -246,12 +246,23 @@ async def init(runtime: DistributedRuntime, config: Config):
     )
 
     logger.info("VllmWorker has been initialized")
+    base_zmq_endpoint = "tcp://127.0.0.1:5557"
+    dp_rank_size = vllm_config.parallel_config.data_parallel_size
+
+    # TODO This isn't working still
+    for dp_rank in range(dp_rank_size):
+        print(f"DP_RANK in Dynamo: {dp_rank}")
+        zmq_endpoint = ZmqEventPublisher.offset_endpoint_port(
+            base_zmq_endpoint, data_parallel_rank=dp_rank
+        )
+        print(f"ZMQ_ENPDPOINT in Dynamo: {zmq_endpoint}")
+        zmq_config = ZmqKvEventPublisherConfig(
+            worker_id=endpoint.lease_id(),
+            kv_block_size=engine_args.block_size,
+            zmq_endpoint=zmq_endpoint,
+        )
 
-    zmq_config = ZmqKvEventPublisherConfig(
-        worker_id=endpoint.lease_id(), kv_block_size=engine_args.block_size
-    )
-
-    _ = ZmqKvEventPublisher(component=component, config=zmq_config)
+        _ = ZmqKvEventPublisher(component=component, config=zmq_config)
 
     handler = RequestHandler(component, engine_client, default_sampling_params)
 

lib/bindings/python/src/dynamo/_core.pyi

@@ -356,7 +356,7 @@ class WorkerMetricsPublisher:
         Create a `WorkerMetricsPublisher` object
         """
 
-    def create_service(self, component: Component) -> None:
+    def create_endpoint(self, component: Component, dp_rank: int) -> None:
         """
         Similar to Component.create_service, but only service created through
         this method will interact with KV router of the same component.

lib/llm/src/kv_router/publisher.rs

@@ -497,7 +497,7 @@ impl WorkerMetricsPublisher {
         let handler = Ingress::for_engine(handler)?;
 
         let endpoint_name = match suffix {
-            Some(s) => format!("{}_{}", KV_METRICS_ENDPOINT, s),
+            Some(s) => format!("{}-{}", KV_METRICS_ENDPOINT, s),
             None => KV_METRICS_ENDPOINT.to_string(),
         };
 

lib/llm/src/kv_router/scoring.rs

@@ -46,12 +46,14 @@ impl Endpoint {
     }
 
     pub fn dp_rank(&self) -> Option<DpRank> {
+        tracing::info!("Parsing dp_rank from subject: {}", self.subject);
         let parts: Vec<&str> = self.subject.split("-").collect();
-        if parts.len() < 2 {
+        if parts.len() < 3 {
             return None;
         }
         let second_to_last = parts[parts.len() - 2];
-        second_to_last.parse::<DpRank>().ok()
+        let result = second_to_last.parse::<DpRank>().ok();
+        result
     }
 } // TODO: make dp_rank
 
@@ -70,6 +72,7 @@ impl ProcessedEndpoints {
             .map(|endpoint| endpoint.data.kv_active_blocks as f64)
             .collect();
         if load_values.is_empty() {
+            // TODO we hit this panic while vLLM is starting the ranks up. Need to avoid this
             panic!("No endpoints to process!")
         };
         let load_avg = load_values.iter().copied().sum::<f64>() / load_values.len() as f64;