[Bugfix]fix proces close problem for omni class (#554)

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

vllm_omni/entrypoints/omni.py

@@ -430,7 +430,13 @@ class Omni(OmniBase):
                     per_stage_params.append(self.default_sampling_params_list[stage_id])
 
             sampling_params_list = per_stage_params
-        return self._run_generation(prompts, sampling_params_list)
+        try:
+            return self._run_generation(prompts, sampling_params_list)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to run generation: %s", e)
+            raise e
+        finally:
+            self.close()
 
     def _run_generation(
         self,

vllm_omni/entrypoints/omni_stage.py

@@ -630,7 +630,7 @@ def _stage_worker(
 
         _recv_dequeue_ts = _time.time()
         if task is None:
-            logger.error("Received shutdown signal")
+            logger.info("Received shutdown signal")
             break
 
         batch_tasks: list[dict[str, Any]] = [task]