Record device after successful switch

Author: fallintoplace

python/cudf_polars/cudf_polars/callback.py

@@ -235,15 +235,14 @@ def set_device(device: int | None) -> Generator[int, None, None]:
                 f"A previous query used device-{SEEN_DEVICE}, "
                 f"the current query is using device-{to_use}."
             )
+        if to_use != current:
+            gpu.setDevice(to_use)
         SEEN_DEVICE = to_use
-    if to_use != current:
-        gpu.setDevice(to_use)
-        try:
-            yield to_use
-        finally:
-            gpu.setDevice(current)
-    else:
+    try:
         yield to_use
+    finally:
+        if to_use != current:
+            gpu.setDevice(current)
 
 
 @overload

python/cudf_polars/tests/test_config.py

@@ -89,6 +89,30 @@ def test_invalid_device_raises(device, monkeypatch):
             q.collect(engine=pl.GPUEngine(device=device))
 
 
+def test_failed_device_switch_does_not_set_seen_device(monkeypatch):
+    monkeypatch.setattr(cudf_polars.callback, "SEEN_DEVICE", None)
+    monkeypatch.setattr(gpu, "getDevice", lambda: 0)
+
+    def fail_for_invalid_device(device):
+        if device == -1:
+            raise RuntimeError("invalid device")
+
+    monkeypatch.setattr(gpu, "setDevice", fail_for_invalid_device)
+
+    with (
+        pytest.raises(RuntimeError, match="invalid device"),
+        cudf_polars.callback.set_device(-1),
+    ):
+        pass
+
+    assert cudf_polars.callback.SEEN_DEVICE is None
+
+    with cudf_polars.callback.set_device(0):
+        pass
+
+    assert cudf_polars.callback.SEEN_DEVICE == 0
+
+
 def test_multiple_devices_in_same_process_raise(monkeypatch):
     # A device we haven't already seen
     monkeypatch.setattr(cudf_polars.callback, "SEEN_DEVICE", 4)