fix lint error in vllm_worker_async.py

Removed unused import statements in vllm_worker_async.py

Signed-off-by: Youngeun Kwon <[email protected]>

fix lint error

Signed-off-by: Youngeun Kwon <[email protected]>

Remove unnecessary blank line in vllm_worker_async.py

Signed-off-by: Youngeun Kwon <[email protected]>

ci pipe clean

Signed-off-by: Youngeun Kwon <[email protected]>

add lock for thread safety

Signed-off-by: Youngeun Kwon <[email protected]>

fix the ci error

ci error

Signed-off-by: Youngeun Kwon <[email protected]>

Author: youngeunkwon0405

nemo_rl/algorithms/grpo.py

@@ -2154,6 +2154,7 @@ def async_grpo_train(
                     train_results = policy.train(train_data, loss_fn)
 
                 print("🔄 Synchronizing policy weights to trajectory collector…")
+                vllm_logger_metrics = None
                 if NEED_REFIT:
                     # Measure pending-generation wait as exposed_generation time
                     print("🔄 Coordinating with trajectory collector before refit...")
@@ -2257,7 +2258,8 @@ def async_grpo_train(
                     else:
                         metrics[k] = np.sum(v).item()
                 metrics.update(rollout_metrics)
-                metrics["vllm_logger_metrics"] = vllm_logger_metrics
+                if vllm_logger_metrics is not None:
+                    metrics["vllm_logger_metrics"] = vllm_logger_metrics
                 total_valid_tokens += metrics["global_valid_toks"]
 
                 # Checkpointing (same as sync version)

nemo_rl/algorithms/utils.py

@@ -498,11 +498,13 @@ def visualize_per_worker_timeline(
             else:
                 print(f"    - Generation Worker {dp_idx:3.0f}: {''.join(timeline)}")
 
-    vllm_logger_metrics = metrics["vllm_logger_metrics"]
-    is_vllm_metrics_logger_enabled = master_config["policy"]["generation"][
-        "vllm_cfg"
-    ].get("enable_vllm_metrics_logger", False)
+    is_vllm_metrics_logger_enabled = (
+        master_config["policy"]["generation"]
+        .get("vllm_cfg", {})
+        .get("enable_vllm_metrics_logger", False)
+    )
     if is_vllm_metrics_logger_enabled:
+        vllm_logger_metrics = metrics["vllm_logger_metrics"]
         # vllm_logger_me    trics: dict[str (metric_name), dict[int (dp_idx), list[int] (metric_values)]]
         # metric_name: "inflight_batch_sizes" or "num_pending_samples"
         vllm_metrics_logger_interval = master_config["policy"]["generation"][

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -366,6 +366,9 @@ def _start_vllm_metrics_logger(self) -> None:
         stop_event = threading.Event()
         self._vllm_metrics_logger_stop_event = stop_event
 
+        # Thread synchronization for metrics access
+        self._vllm_metrics_lock = threading.Lock()
+
         self.inflight_batch_sizes: list[int] = []
         self.num_pending_samples: list[int] = []
 
@@ -389,10 +392,12 @@ def _logger_loop():
                             if isinstance(m, Gauge):
                                 # Log the vllm inflight batch sizes
                                 if m.name == "vllm:num_requests_running":
-                                    self.inflight_batch_sizes.append(int(m.value))
+                                    with self._vllm_metrics_lock:
+                                        self.inflight_batch_sizes.append(int(m.value))
                                 # Log the vllm pending number of requests in the queue
                                 elif m.name == "vllm:num_requests_waiting":
-                                    self.num_pending_samples.append(int(m.value))
+                                    with self._vllm_metrics_lock:
+                                        self.num_pending_samples.append(int(m.value))
                         except Exception:
                             print(
                                 "⚠️[vLLM Metric Logger]⚠️ Exception in vLLM metrics logger",
@@ -426,16 +431,20 @@ def get_vllm_logger_metrics(self) -> dict[str, Any]:
         if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False):
             return {}
 
-        return {
-            "inflight_batch_sizes": copy.deepcopy(self.inflight_batch_sizes),
-            "num_pending_samples": copy.deepcopy(self.num_pending_samples),
-        }
+        with self._vllm_metrics_lock:
+            metric = {
+                "inflight_batch_sizes": copy.deepcopy(self.inflight_batch_sizes),
+                "num_pending_samples": copy.deepcopy(self.num_pending_samples),
+            }
+        return metric
 
     def clear_vllm_logger_metrics(self) -> None:
         if not self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False):
             return
-        self.inflight_batch_sizes = []
-        self.num_pending_samples = []
+
+        with self._vllm_metrics_lock:
+            self.inflight_batch_sizes = []
+            self.num_pending_samples = []
 
     def llm(self):
         return self.llm

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -165,10 +165,12 @@ def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI:
         from logging import LogRecord
         from typing import List, Optional, Union
 
-        from fastapi import Request
-        from fastapi.responses import JSONResponse, StreamingResponse
-
-        from vllm.entrypoints.openai.api_server import (
+        from fastapi import Request  # pyright: ignore[reportMissingImports]
+        from fastapi.responses import (  # pyright: ignore[reportMissingImports]
+            JSONResponse,
+            StreamingResponse,
+        )
+        from vllm.entrypoints.openai.api_server import (  # pyright: ignore[reportMissingImports]
             BaseModelPath,
             OpenAIServingChat,
             OpenAIServingModels,