Restore CUDA_VISIBLE_DEVICES after test

Signed-off-by: Enrico Minack <github@enrico.minack.dev>

test/single/test_ray.py

@@ -5,7 +5,6 @@
 import os
 import socket
 import sys
-import time
 
 import pytest
 import ray
@@ -39,19 +38,27 @@ def ray_start_4_cpus():
 @pytest.fixture
 def ray_start_6_cpus():
     address_info = ray.init(num_cpus=6)
-    yield address_info
-    # The code after the yield will run as teardown code.
-    ray.shutdown()
+    try:
+        yield address_info
+    finally:
+        # The code after the yield will run as teardown code.
+        ray.shutdown()
 
 
 @pytest.fixture
 def ray_start_4_cpus_4_gpus():
+    orig_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
     os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
     address_info = ray.init(num_cpus=4, num_gpus=4)
-    yield address_info
-    # The code after the yield will run as teardown code.
-    ray.shutdown()
-    del os.environ["CUDA_VISIBLE_DEVICES"]
+    try:
+        yield address_info
+        # The code after the yield will run as teardown code.
+        ray.shutdown()
+    finally:
+        if orig_devices:
+            os.environ["CUDA_VISIBLE_DEVICES"] = orig_devices
+        else:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
 
 
 @pytest.fixture