fix: python bindings for router should register to etcd as well (#3302)

Signed-off-by: PeaBrane <yanrpei@gmail.com>

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

@@ -1050,21 +1050,20 @@ impl KvPushRouter {
                 ));
             }
 
-            // Create KvRouter with a unique consumer UUID
-            let consumer_uuid = uuid::Uuid::new_v4().to_string();
-            let kv_router = llm_rs::kv_router::KvRouter::new(
-                component.clone(),
-                block_size as u32,
-                None, // default selector
-                Some(kv_router_config.inner()),
-                consumer_uuid,
-            )
-            .await
-            .map_err(to_pyerr)?;
+            // Create ModelManager and use it to create KvRouter (ensures etcd registration)
+            let model_manager = Arc::new(llm_rs::discovery::ModelManager::new());
+            let kv_router = model_manager
+                .kv_chooser_for(
+                    "dummy_name", // does not matter, never cached
+                    component,
+                    block_size as u32,
+                    Some(kv_router_config.inner()),
+                )
+                .await
+                .map_err(to_pyerr)?;
 
-            // Create KvPushRouter
-            let kv_push_router =
-                llm_rs::kv_router::KvPushRouter::new(push_router, Arc::new(kv_router));
+            // Create KvPushRouter (kv_router is already Arc<KvRouter>)
+            let kv_push_router = llm_rs::kv_router::KvPushRouter::new(push_router, kv_router);
 
             Ok(Self {
                 inner: Arc::new(kv_push_router),

lib/llm/src/discovery/model_manager.rs

@@ -227,22 +227,8 @@ impl ModelManager {
             }
             return Ok(kv_chooser);
         }
-        self.create_kv_chooser(model_name, component, kv_cache_block_size, kv_router_config)
-            .await
-    }
-
-    fn get_kv_chooser(&self, model_name: &str) -> Option<Arc<KvRouter>> {
-        self.kv_choosers.lock().get(model_name).cloned()
-    }
 
-    /// Create and return a KV chooser for this component and model
-    async fn create_kv_chooser(
-        &self,
-        model_name: &str,
-        component: &Component,
-        kv_cache_block_size: u32,
-        kv_router_config: Option<KvRouterConfig>,
-    ) -> anyhow::Result<Arc<KvRouter>> {
+        // Create new KV router with etcd registration
         let etcd_client = component
             .drt()
             .etcd_client()
@@ -278,6 +264,10 @@ impl ModelManager {
         Ok(new_kv_chooser)
     }
 
+    fn get_kv_chooser(&self, model_name: &str) -> Option<Arc<KvRouter>> {
+        self.kv_choosers.lock().get(model_name).cloned()
+    }
+
     pub fn get_model_tool_call_parser(&self, model: &str) -> Option<String> {
         self.cards
             .lock()