[Debug] Correct Unreasonable Long Timeout (#1175)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/diffusion/diffusion_engine.py

@@ -253,7 +253,7 @@ class DiffusionEngine:
 
         try:
             # Give worker enough time — export + compression + table can be slow
-            results = self.collective_rpc(method="stop_profile", timeout=60000)
+            results = self.collective_rpc(method="stop_profile", timeout=600)
         except Exception:
             logger.error("Failed to stop profiling on workers", exc_info=True)
             return {"traces": [], "tables": []}

vllm_omni/entrypoints/cli/serve.py

@@ -97,7 +97,7 @@ class OmniServeCommand(CLISubcommand):
         omni_config_group.add_argument(
             "--init-timeout",
             type=int,
-            default=60000,
+            default=600,
             help="The timeout for initializing the stages.",
         )
         omni_config_group.add_argument(

vllm_omni/entrypoints/omni_stage.py

@@ -221,8 +221,8 @@ class OmniStage:
 
         # Wait for result from worker
         try:
-            # Profiling stop might take time to flush files, give it 180s
-            response = self._out_q.get(timeout=60000)
+            # Profiling stop might take time to flush files, give it 600s
+            response = self._out_q.get(timeout=600)
 
             if isinstance(response, dict):
                 if response.get("type") == "profiler_result":