Fix CUDA context leak during NVENC probing

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":

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