chore: many bug fixes and improvements when testing planner (#2776)

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuan <hongkuanz@nvidia.com>

Author: tedzhouhk

benchmarks/nixl/nixl-benchmark-deployment.yaml

@@ -15,7 +15,7 @@ spec:
         app: nixl-benchmark
     spec:
       imagePullSecrets:
-        - name: nvcrimagepullsecret
+        - name: nvcr-imagepullsecret
       containers:
       - name: nixl-benchmark
         image: nvcr.io/nvidian/nim-llm-dev/vllm-runtime:nixlbench-e42c07a8

benchmarks/sin_load_generator/sin_synth.py

@@ -51,7 +51,7 @@ def get_isl_osl(t):
             isl, osl = get_isl_osl(t_req)
             output_data.append(
                 {
-                    "timestamp": t_req * 1000,  # in ms
+                    "timestamp": int(t_req * 1000),  # in ms, integer
                     "input_length": isl,
                     "output_length": osl,
                     "hash_ids": np.random.choice(

components/backends/sglang/deploy/disagg_planner.yaml

@@ -111,19 +111,24 @@ spec:
           image: nvcr.io/nvidian/nim-llm-dev/sglang-runtime:hzhou-0811-1
           workingDir: /workspace/components/backends/sglang
           command:
-            - /bin/sh
-            - -c
+            - python3
           args:
-            - >-
-              python3 -m dynamo.sglang
-              --model-path Qwen/Qwen3-0.6B
-              --served-model-name Qwen/Qwen3-0.6B
-              --page-size 16
-              --tp 1
-              --trust-remote-code
-              --skip-tokenizer-init
-              --disaggregation-mode decode
-              --disaggregation-transfer-backend nixl
+            - -m
+            - dynamo.sglang
+            - --model-path
+            - Qwen/Qwen3-0.6B
+            - --served-model-name
+            - Qwen/Qwen3-0.6B
+            - --page-size
+            - "16"
+            - --tp
+            - "1"
+            - --trust-remote-code
+            - --skip-tokenizer-init
+            - --disaggregation-mode
+            - decode
+            - --disaggregation-transfer-backend
+            - nixl
     SGLangPrefillWorker:
       dynamoNamespace: dynamo
       envFromSecret: hf-token-secret
@@ -137,16 +142,21 @@ spec:
           image: nvcr.io/nvidian/nim-llm-dev/sglang-runtime:hzhou-0811-1
           workingDir: /workspace/components/backends/sglang
           command:
-            - /bin/sh
-            - -c
+            - python3
           args:
-            - >-
-              python3 -m dynamo.sglang
-              --model-path Qwen/Qwen3-0.6B
-              --served-model-name Qwen/Qwen3-0.6B
-              --page-size 16
-              --tp 1
-              --trust-remote-code
-              --skip-tokenizer-init
-              --disaggregation-mode prefill
-              --disaggregation-transfer-backend nixl
+            - -m
+            - dynamo.sglang
+            - --model-path
+            - Qwen/Qwen3-0.6B
+            - --served-model-name
+            - Qwen/Qwen3-0.6B
+            - --page-size
+            - "16"
+            - --tp
+            - "1"
+            - --trust-remote-code
+            - --skip-tokenizer-init
+            - --disaggregation-mode
+            - prefill
+            - --disaggregation-transfer-backend
+            - nixl

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

@@ -117,7 +117,9 @@ async def register_model():
         # Requests queue until ready_event is set
         await asyncio.gather(
             generate_endpoint.serve_endpoint(
-                handler.generate, graceful_shutdown=False, metrics_labels=metrics_labels
+                handler.generate,
+                graceful_shutdown == config.migration_limit <= 0,
+                metrics_labels=metrics_labels,
             ),
             register_model(),
         )

components/backends/vllm/deploy/disagg_planner.yaml

@@ -121,10 +121,12 @@ spec:
           image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.4.1
           workingDir: /workspace/components/backends/vllm
           command:
-            - /bin/sh
-            - -c
+            - python3
           args:
-            - "python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --migration-limit=3"
+            - -m
+            - dynamo.vllm
+            - --model
+            - Qwen/Qwen3-0.6B
     VllmPrefillWorker:
       dynamoNamespace: vllm-disagg-planner
       envFromSecret: hf-token-secret
@@ -144,7 +146,10 @@ spec:
           image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.4.1
           workingDir: /workspace/components/backends/vllm
           command:
-            - /bin/sh
-            - -c
+            - python3
           args:
-            - python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --is-prefill-worker --migration-limit=3
+            - -m
+            - dynamo.vllm
+            - --model
+            - Qwen/Qwen3-0.6B
+            - --is-prefill-worker

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

@@ -84,8 +84,12 @@ def signal_handler():
 
     if config.is_prefill_worker:
         await init_prefill(runtime, config)
+        logger.debug("init_prefill completed")
     else:
         await init(runtime, config)
+        logger.debug("init completed")
+
+    logger.debug("Worker function completed, exiting...")
 
 
 def setup_vllm_engine(config, stat_logger=None):
@@ -147,6 +151,7 @@ async def init_prefill(runtime: DistributedRuntime, config: Config):
     )
 
     try:
+        logger.debug("Starting serve_endpoint for prefill worker")
         await asyncio.gather(
             # for prefill, we want to shutdown the engine after all prefill requests are finished because
             #     (temp reason): we don't support re-routing prefill requests
@@ -161,10 +166,12 @@ async def init_prefill(runtime: DistributedRuntime, config: Config):
                 handler.clear_kv_blocks, metrics_labels=[("model", config.model)]
             ),
         )
+        logger.debug("serve_endpoint completed for prefill worker")
     except Exception as e:
         logger.error(f"Failed to serve endpoints: {e}")
         raise
     finally:
+        logger.debug("Cleaning up prefill worker")
         handler.cleanup()
 
 
@@ -254,22 +261,25 @@ async def init(runtime: DistributedRuntime, config: Config):
         )
 
     try:
+        logger.debug("Starting serve_endpoint for decode worker")
         await asyncio.gather(
             # for decode, we want to transfer the in-flight requests to other decode engines,
             # because waiting them to finish can take a long time for long OSLs
             generate_endpoint.serve_endpoint(
                 handler.generate,
-                graceful_shutdown=False,
+                graceful_shutdown=config.migration_limit <= 0,
                 metrics_labels=[("model", config.model)],
             ),
             clear_endpoint.serve_endpoint(
                 handler.clear_kv_blocks, metrics_labels=[("model", config.model)]
             ),
         )
+        logger.debug("serve_endpoint completed for decode worker")
     except Exception as e:
         logger.error(f"Failed to serve endpoints: {e}")
         raise
     finally:
+        logger.debug("Cleaning up decode worker")
         # Cleanup background tasks
         handler.cleanup()
 

components/planner/src/dynamo/planner/utils/load_predictor.py

@@ -42,10 +42,14 @@ def __init__(self, minimum_data_points=5):
 
     def add_data_point(self, value):
         """Add new data point to the buffer"""
-        if not math.isnan(value):
-            self.data_buffer.append(value)
+        if math.isnan(value):
+            value = 0
+
+        if len(self.data_buffer) == 0 and value == 0:
+            # skip the beginning idle period
+            return
         else:
-            self.data_buffer.append(0)
+            self.data_buffer.append(value)
 
     def get_last_value(self):
         """Get the last value from the buffer"""
@@ -126,6 +130,10 @@ def add_data_point(self, value):
         # Use proper datetime for Prophet
         timestamp = self.start_date + timedelta(seconds=self.curr_step)
         value = 0 if math.isnan(value) else value
+
+        if len(self.data_buffer) == 0 and value == 0:
+            # skip the beginning idle period
+            return
         self.data_buffer.append({"ds": timestamp, "y": value})
         self.curr_step += 1
 

components/planner/src/dynamo/planner/utils/planner_core.py

@@ -259,18 +259,24 @@ def _compute_replica_requirements(
         # compute how many replicas are needed for prefill
         # here we assume the prefill bias is purely due to request queueing
         # and we increase the number of prefill replicas linearly to account for the queueing delay
-        pred_prefill_load_per_gpu = (
+        pred_prefill_throughput = (
             next_num_req
             * next_isl
             / self.args.adjustment_interval
             * min(1, self.p_correction_factor)
         )
         next_num_p = math.ceil(
-            pred_prefill_load_per_gpu
+            pred_prefill_throughput
             / self.prefill_interpolator.interpolate_thpt_per_gpu(next_isl)
             / self.args.prefill_engine_num_gpu
         )
 
+        logger.info(
+            f"Prefill calculation: {pred_prefill_throughput:.2f}(p_thpt) / "
+            f"{self.prefill_interpolator.interpolate_thpt_per_gpu(next_isl) * self.args.prefill_engine_num_gpu:.2f}(p_engine_cap) = "
+            f"{next_num_p}(num_p)"
+        )
+
         # compute how many replicas are needed for decode
         # 1. apply d_correction_factor to the ITL SLA
         # Prevent divide by zero when d_correction_factor is 0 (no metrics yet)
@@ -290,14 +296,19 @@ def _compute_replica_requirements(
             itl=corrected_itl, context_length=next_isl + next_osl / 2
         )
         # 3. compute number of decode replicas needed
+        pred_decode_throughput = next_num_req * next_osl / self.args.adjustment_interval
         next_num_d = math.ceil(
-            next_num_req
-            * next_osl
-            / self.args.adjustment_interval
+            pred_decode_throughput
             / pred_decode_thpt_per_gpu
             / self.args.decode_engine_num_gpu
         )
 
+        logger.info(
+            f"Decode calculation: {pred_decode_throughput:.2f}(d_thpt) / "
+            f"{pred_decode_thpt_per_gpu * self.args.decode_engine_num_gpu:.2f}(d_engine_cap) = "
+            f"{next_num_d}(num_d)"
+        )
+
         # correct num_p and num_d based on the gpu budget
         next_num_p = max(next_num_p, self.args.min_endpoint)
         next_num_d = max(next_num_d, self.args.min_endpoint)

lib/llm/src/discovery/watcher.rs

@@ -176,46 +176,6 @@ impl ModelWatcher {
             .await
             .with_context(|| model_name.clone())?;
         if !active_instances.is_empty() {
-            let mut update_tx = true;
-            let mut model_type: ModelType = model_entry.model_type;
-            if model_entry.model_type == ModelType::Chat
-                && self.manager.list_chat_completions_models().is_empty()
-            {
-                self.manager.remove_chat_completions_model(&model_name).ok();
-                model_type = ModelType::Chat;
-            } else if model_entry.model_type == ModelType::Completion
-                && self.manager.list_completions_models().is_empty()
-            {
-                self.manager.remove_completions_model(&model_name).ok();
-                model_type = ModelType::Completion;
-            } else if model_entry.model_type == ModelType::Embedding
-                && self.manager.list_embeddings_models().is_empty()
-            {
-                self.manager.remove_embeddings_model(&model_name).ok();
-                model_type = ModelType::Embedding;
-            } else if model_entry.model_type == ModelType::Backend {
-                if self.manager.list_chat_completions_models().is_empty() {
-                    self.manager.remove_chat_completions_model(&model_name).ok();
-                    model_type = ModelType::Chat;
-                }
-                if self.manager.list_completions_models().is_empty() {
-                    self.manager.remove_completions_model(&model_name).ok();
-                    if model_type == ModelType::Chat {
-                        model_type = ModelType::Backend;
-                    } else {
-                        model_type = ModelType::Completion;
-                    }
-                }
-            } else {
-                tracing::debug!(
-                    "Model {} is still active in other instances, not removing",
-                    model_name
-                );
-                update_tx = false;
-            }
-            if update_tx && let Some(tx) = &self.model_update_tx {
-                tx.send(ModelUpdate::Removed(model_type)).await.ok();
-            }
             return Ok(None);
         }