fix(vllm): serialise AsyncMPClient input_socket sends to prevent zmq race

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -184,6 +184,8 @@ def _create_engine(self, llm_kwargs: dict[str, Any]) -> None:
 
         self.server_thread, self.base_url, self.http_server = None, None, None
         if self.cfg["vllm_cfg"].get("expose_http_server"):
+            # Must run before _setup_vllm_server spawns the uvicorn thread.
+            self._install_engine_input_socket_lock()
             self.server_thread, self.base_url, self.http_server = (
                 self._setup_vllm_server()
             )
@@ -194,6 +196,37 @@ def _create_engine(self, llm_kwargs: dict[str, Any]) -> None:
         if self.cfg["vllm_cfg"].get("enable_vllm_metrics_logger", False):
             self._start_vllm_metrics_logger()
 
+    def _install_engine_input_socket_lock(self) -> None:
+        """Serialise sends on AsyncMPClient.input_socket across OS threads
+        to prevent race conditions that block the vLLM engine (e.g. during
+        in flight weight updates in async grpo).
+        """
+        try:
+            shadow_sock = self.llm.engine_core.input_socket._shadow_sock
+        except AttributeError as exc:
+            # ``self.llm`` already owns EngineCore subprocesses; tear them down
+            # before re-raising so we don't leak when actor __init__ aborts.
+            try:
+                self.llm.shutdown()
+            except Exception:
+                pass
+            raise RuntimeError(
+                "AsyncMPClient.input_socket._shadow_sock is unreachable; vLLM "
+                "internals have shifted. Update this guard before exposing the "
+                "vLLM HTTP server."
+            ) from exc
+
+        lock = threading.Lock()
+        original_send_multipart = shadow_sock.send_multipart
+
+        def locked_send_multipart(*args: Any, **kwargs: Any) -> Any:
+            with lock:
+                return original_send_multipart(*args, **kwargs)
+
+        # Replace the bound method on this socket instance only; other zmq
+        # sockets in the process are unaffected.
+        shadow_sock.send_multipart = locked_send_multipart  # type: ignore[assignment]
+
     def _start_vllm_metrics_logger(self) -> None:
         """Start a background thread that periodically collects vLLM logger metrics.