From d4f353242196ab8b87faf56763f5210501f1be2c Mon Sep 17 00:00:00 2001
From: Gangzheng Tong <gtong@nvidia.com>
Date: Thu, 6 Aug 2026 23:21:55 +0000
Subject: [PATCH] Fix CUDA context leak during NVENC probing

---
 .../flashdreams/serving/webrtc/encoders.py    | 33 +++++++++++++++++++
 flashdreams/tests/test_encoders.py            | 16 +++++++++
 2 files changed, 49 insertions(+)

diff --git a/flashdreams/flashdreams/serving/webrtc/encoders.py b/flashdreams/flashdreams/serving/webrtc/encoders.py
index 1b1d50f7..bb27c7f6 100644
--- a/flashdreams/flashdreams/serving/webrtc/encoders.py
+++ b/flashdreams/flashdreams/serving/webrtc/encoders.py
@@ -202,10 +202,43 @@ def select_encoder(
       — driver bug, session-pool exhaustion, hardware fault, or
       misconfiguration. Log with traceback and re-raise; do not silently
       degrade, because doing so would hide a real problem.
+
+    PyNvVideoCodec may replace the CUDA context current on the calling
+    thread, including when ``GetEncoderCaps`` fails. Preserve an already
+    initialized PyTorch device context so subsequent model work does not
+    inherit the encoder probe's context.
     """
     if backend == "default":
         return DefaultRTCEncoder(fps=fps)
 
+    restore_device = (
+        torch.cuda.current_device() if torch.cuda.is_initialized() else None
+    )
+    try:
+        return _select_hardware_encoder(
+            backend=backend,
+            width=width,
+            height=height,
+            fps=fps,
+            bitrate=bitrate,
+            gpu_id=gpu_id,
+            gop=gop,
+        )
+    finally:
+        if restore_device is not None:
+            torch.cuda.set_device(restore_device)
+
+
+def _select_hardware_encoder(
+    *,
+    backend: Literal["auto", "nvenc"],
+    width: int,
+    height: int,
+    fps: int,
+    bitrate: int,
+    gpu_id: int,
+    gop: int,
+) -> VideoEncoder:
     if not _pynvvideocodec_installed():
         reason = "PyNvVideoCodec library is not installed"
         if backend == "nvenc":
diff --git a/flashdreams/tests/test_encoders.py b/flashdreams/tests/test_encoders.py
index 11cdcf42..20d2183c 100644
--- a/flashdreams/tests/test_encoders.py
+++ b/flashdreams/tests/test_encoders.py
@@ -210,6 +210,22 @@ def test_caps_failure_nvenc_raises_with_reason(
         with pytest.raises(EncoderInitError, match="driver comms error"):
             select_encoder(backend="nvenc", **_SELECT_KW)
 
+    def test_caps_probe_restores_initialized_torch_cuda_context(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        fake_nvc = MagicMock()
+        fake_nvc.GetEncoderCaps.side_effect = RuntimeError("unsupported GPU")
+        _install_fake_nvc(monkeypatch, fake_nvc)
+        set_device = MagicMock()
+        monkeypatch.setattr(enc_mod.torch.cuda, "is_initialized", lambda: True)
+        monkeypatch.setattr(enc_mod.torch.cuda, "current_device", lambda: 3)
+        monkeypatch.setattr(enc_mod.torch.cuda, "set_device", set_device)
+
+        enc = select_encoder(backend="auto", **_SELECT_KW)
+
+        assert isinstance(enc, DefaultRTCEncoder)
+        set_device.assert_called_once_with(3)
+
 
 # ---------------------------------------------------------------------------
 # select_encoder: Stage-1 deferred-import failure (package present but