From c4cd502d82c03e207adc869b4504cfacc8ea437d Mon Sep 17 00:00:00 2001
From: Gao Han <gaohan19@huawei.com>
Date: Wed, 17 Dec 2025 22:40:20 +0800
Subject: [PATCH] [Misc] Fix left comments for PR298 (#346)

Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>
---
 .buildkite/pipeline.yml           | 14 +++++++-------
 vllm_omni/entrypoints/omni_llm.py | 22 ++++++++--------------
 vllm_omni/entrypoints/utils.py    |  5 ++++-
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index a2f78f05c..648c00ffd 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -9,13 +9,13 @@ steps:
     agents:
       queue: "cpu_queue_premerge_us_east_1"
 
-  - label: "Test on NPU"
-    depends_on: ~
-    key: npu-test
-    commands:
-      - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
-    agents:
-      queue: "ascend"
+  # - label: "Test on NPU"
+  #   depends_on: ~
+  #   key: npu-test
+  #   commands:
+  #     - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
+  #   agents:
+  #     queue: "ascend"
 
   - label: "Simple Unit Test"
     commands:
diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index 9ca744d5d..de67bd618 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/vllm_omni/entrypoints/omni_llm.py
@@ -294,21 +294,15 @@ class OmniLLM:
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
         final_stage_id_to_prompt = {}
-        try:
-            for rid, prompt in request_id_to_prompt.items():
-                final_stage_id_for_e2e = get_final_stage_id_for_e2e(
-                    prompt.get("modalities", None), self.output_modalities, self.stage_list
-                )
-                final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
-
-        except Exception as e:
-            logger.debug(
-                "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
-                e,
-                exc_info=True,
+        for rid, prompt in request_id_to_prompt.items():
+            if isinstance(prompt, dict):
+                prompt_modalities = prompt.get("modalities", None)
+            else:
+                prompt_modalities = None
+            final_stage_id_for_e2e = get_final_stage_id_for_e2e(
+                prompt_modalities, self.output_modalities, self.stage_list
             )
-            for rid, prompt in request_id_to_prompt.items():
-                final_stage_id_to_prompt[rid] = len(self.stage_list) - 1
+            final_stage_id_to_prompt[rid] = final_stage_id_for_e2e
 
         # Metrics/aggregation helper
         metrics = OrchestratorMetrics(
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index 19715ef54..d86ded250 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -148,7 +148,9 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
     return stage_args
 
 
-def get_final_stage_id_for_e2e(output_modalities: list[str], default_modalities: list[str], stage_list: list) -> int:
+def get_final_stage_id_for_e2e(
+    output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
+) -> int:
     """Get the final stage id for e2e.
 
     Args:
@@ -163,6 +165,7 @@ def get_final_stage_id_for_e2e(output_modalities: list[str], default_modalities:
         for modality in output_modalities:
             if modality not in default_modalities:
                 logger.warning(f"Invalid output modality: {modality}, ignoring it")
+                # TODO: if user specifies unsupported modalities, invalid it and raise an error
                 continue
             prompt_modalities.append(modality)
         output_modalities = prompt_modalities