[Refactor] Clean up code by removing unnecessary lines and adding missing imports in async_omni_llm.py, omni_stage.py, omni.py, and api_server.py

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

vllm_omni/entrypoints/async_omni_llm.py

@@ -164,10 +164,7 @@ class AsyncOmniLLM(AsyncLLM):
             pass
 
         # Use profiler_config from vllm_config (new way, aligned with vllm v1)
-        if (
-            vllm_config.profiler_config.profiler == "torch"
-            and not vllm_config.profiler_config.ignore_frontend
-        ):
+        if vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend:
             profiler_dir = vllm_config.profiler_config.torch_profiler_dir
             logger.info(
                 "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s",

vllm_omni/entrypoints/omni.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
 import multiprocessing as mp
 import os
 import time

vllm_omni/entrypoints/omni_stage.py

@@ -626,7 +626,6 @@ def _stage_worker(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
-            next_wait_log = wait_start + 5.0
 
             while not lock_acquired:
                 try:
@@ -655,7 +654,6 @@ def _stage_worker(
                             )
                             break
 
-
                         # Wait a bit before retrying
                         _time.sleep(0.1)
                 except OSError as e:
@@ -1173,8 +1171,6 @@ async def _stage_worker_async(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
-            next_wait_log = wait_start + 5.0
-
             while not lock_acquired:
                 try:
                     # Open or create the lock file
@@ -1203,7 +1199,6 @@ async def _stage_worker_async(
                             )
                             break
 
-
                         # Wait a bit before retrying
                         _time.sleep(0.1)
                 except OSError as e:

vllm_omni/entrypoints/openai/api_server.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import base64
 import io
+import json
 import multiprocessing
 import multiprocessing.forkserver as forkserver
 import os