Add model label to Component

Author: tzulingk

components/backends/vllm/src/dynamo/vllm/main.py

@@ -131,7 +131,10 @@ async def init_prefill(runtime: DistributedRuntime, config: Config):
     """
     Instantiate and serve
     """
-    component = runtime.namespace(config.namespace).component(config.component)
+
+    component = runtime.namespace(config.namespace).component(
+        config.component, config.model
+    )
     await component.create_service()
 
     generate_endpoint = component.endpoint(config.endpoint)
@@ -164,7 +167,9 @@ async def init(runtime: DistributedRuntime, config: Config):
     Instantiate and serve
     """
 
-    component = runtime.namespace(config.namespace).component(config.component)
+    component = runtime.namespace(config.namespace).component(
+        config.component, config.model
+    )
     await component.create_service()
 
     generate_endpoint = component.endpoint(config.endpoint)

components/metrics/src/bin/mock_worker.rs

@@ -146,7 +146,7 @@ async fn backend(runtime: DistributedRuntime) -> Result<()> {
     let namespace = runtime.namespace("dynamo")?;
     // we must first create a service, then we can attach one more more endpoints
     let component = namespace
-        .component("MyComponent")?
+        .component("MyComponent", Some("MyModel".to_string()))?
         .service_builder()
         .create()
         .await?;

components/metrics/src/lib.rs

@@ -145,6 +145,7 @@ impl MetricsMode {
 pub struct LLMWorkerLoadCapacityConfig {
     pub component_name: String,
     pub endpoint_name: String,
+    pub model_name: Option<String>,
 }
 
 /// Metrics collector for exposing metrics to prometheus/grafana

components/metrics/src/main.rs

@@ -60,6 +60,10 @@ struct Args {
     #[arg(long)]
     endpoint: String,
 
+    /// Model name for the target component (optional)
+    #[arg(long)]
+    model_name: String,
+
     /// Polling interval in seconds for scraping dynamo endpoint stats (minimum 1 second)
     #[arg(long, default_value = "1")]
     poll_interval: u64,
@@ -109,6 +113,7 @@ fn get_config(args: &Args) -> Result<LLMWorkerLoadCapacityConfig> {
     Ok(LLMWorkerLoadCapacityConfig {
         component_name: args.component.clone(),
         endpoint_name: args.endpoint.clone(),
+        model_name: Some(args.model_name.clone()),
     })
 }
 
@@ -120,7 +125,9 @@ async fn app(runtime: Runtime) -> Result<()> {
     let drt = DistributedRuntime::from_settings(runtime.clone()).await?;
 
     let namespace = drt.namespace(args.namespace)?;
-    let component = namespace.component("count")?;
+    // The metrics aggregator operates independently of any model,
+    // hence the model name is set to None.
+    let component = namespace.component("count", None)?;
 
     // Create unique instance of Count
     let key = format!("{}/instance", component.etcd_root());
@@ -131,7 +138,7 @@ async fn app(runtime: Runtime) -> Result<()> {
         .await
         .context("Unable to create unique instance of Count; possibly one already exists")?;
 
-    let target_component = namespace.component(&config.component_name)?;
+    let target_component = namespace.component(&config.component_name, config.model_name.clone())?;
     let target_endpoint = target_component.endpoint(&config.endpoint_name);
 
     let service_path = target_endpoint.path();

components/router/src/main.rs

@@ -60,9 +60,11 @@ async fn app(runtime: Runtime) -> Result<()> {
     let args = Args::parse();
     let runtime = DistributedRuntime::from_settings(runtime).await?;
 
+    // The metrics aggregator operates independently of any model,
+    // hence the model name is set to None.
     let component = runtime
         .namespace(&args.namespace)?
-        .component(&args.component)?;
+        .component(&args.component, None)?;
 
     let selector = Box::new(CustomWorkerSelector::default());
 

launch/dynamo-run/src/lib.rs

@@ -138,7 +138,7 @@ async fn engine_for(
             let endpoint = local_model.endpoint_id().clone();
 
             let engine =
-                dynamo_llm::mocker::engine::make_mocker_engine(drt, endpoint, args).await?;
+                dynamo_llm::mocker::engine::make_mocker_engine(drt, endpoint, Some(local_model.card().service_name.clone()), args).await?;
 
             Ok(EngineConfig::StaticCore {
                 engine,

lib/bindings/c/src/lib.rs

@@ -147,7 +147,9 @@ fn dynamo_create_kv_publisher(
         .ok_or(anyhow::Error::msg("Could not get Distributed Runtime"))
     {
         Ok(drt) => {
-            let backend = drt.namespace(namespace)?.component(component)?;
+            // The model publisher isn't tied to a specific model. Thus, the model name 
+            // is set to None.
+            let backend = drt.namespace(namespace)?.component(component, None)?;
             KvEventPublisher::new(backend, worker_id, kv_block_size, None)
         }
         Err(e) => Err(e),

lib/bindings/python/rust/lib.rs

@@ -535,8 +535,9 @@ impl Endpoint {
 
 #[pymethods]
 impl Namespace {
-    fn component(&self, name: String) -> PyResult<Component> {
-        let inner = self.inner.component(name).map_err(to_pyerr)?;
+    #[pyo3(signature = (name, model = None))]
+    fn component(&self, name: String, model: Option<String>) -> PyResult<Component> {
+        let inner = self.inner.component(name, model).map_err(to_pyerr)?;
         Ok(Component {
             inner,
             event_loop: self.event_loop.clone(),

lib/bindings/python/rust/llm/entrypoint.rs

@@ -211,6 +211,7 @@ async fn select_engine(
             let engine = dynamo_llm::mocker::engine::make_mocker_engine(
                 distributed_runtime.inner,
                 endpoint,
+                Some(local_model.card().service_name.clone()),
                 mocker_args,
             )
             .await?;

lib/llm/src/discovery/watcher.rs

@@ -169,7 +169,7 @@ impl ModelWatcher {
         let component = self
             .drt
             .namespace(&endpoint_id.namespace)?
-            .component(&endpoint_id.component)?;
+            .component(&endpoint_id.component, Some(model_entry.name.clone()))?;
         let client = component.endpoint(&endpoint_id.name).client().await?;
 
         let Some(etcd_client) = self.drt.etcd_client() else {