diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index a44596983..8e19f1242 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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": []} diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 540d6d3b4..acefc081b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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( diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 133e9ec84..a2070f89d 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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":