fix start/stop

sfc-gh-mkeralapura · Sep 6, 2024 · 098f90f · 098f90f
1 parent b6f8764
commit 098f90f
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py
@@ -135,14 +135,16 @@ def __init__(
             self.profiler = None
 
     def start_profile(self):
-        if self.is_driver_worker and self.profiler is None:
+        if self.profiler is None:
             raise RuntimeError("Profiler is not enabled.")
-        self.profiler.start()
+        if self.is_driver_worker:
+            self.profiler.start()
 
     def stop_profile(self):
-        if self.is_driver_worker and self.profiler is None:
+        if self.profiler is None:
             raise RuntimeError("Profiler is not enabled.")
-        self.profiler.stop()
+        if self.is_driver_worker:
+            self.profiler.stop()
 
     def _is_encoder_decoder_model(self):
         return self.model_config.is_encoder_decoder_model