[BUG FIX] Fix CUDA runtime being initialized by 'get_device'. (Genesis-Embodied-AI#1634)

duburcqa · web-flow · commit 6e8c686c30c8 · 2025-08-26T17:22:20.000+02:00
* Fix CUDA runtime initialized by 'get_device', breaking device selection for distributed unit tests.
* Fix support of undefined device index.
* Ensure consistency between EGL_DEVICE_ID and CUDA_VISIBLE_DEVICES in unit tests.
* Fix flaky unit tests due to mutable fixture interference.
diff --git a/genesis/ext/pyrender/camera.py b/genesis/ext/pyrender/camera.py
@@ -109,7 +109,7 @@ class PerspectiveCamera(Camera):
     """
 
     def __init__(self, yfov, znear=DEFAULT_Z_NEAR, zfar=None, aspectRatio=None, name=None):
-        super(PerspectiveCamera, self).__init__(znear=znear, zfar=zfar, name=name)
+        super().__init__(znear=znear, zfar=zfar, name=name)
 
         self.yfov = yfov
         self.aspectRatio = aspectRatio
@@ -209,7 +209,7 @@ class OrthographicCamera(Camera):
     """
 
     def __init__(self, xmag, ymag, znear=DEFAULT_Z_NEAR, zfar=DEFAULT_Z_FAR, name=None):
-        super(OrthographicCamera, self).__init__(znear=znear, zfar=zfar, name=name)
+        super().__init__(znear=znear, zfar=zfar, name=name)
 
         self.xmag = xmag
         self.ymag = ymag
@@ -305,7 +305,7 @@ class IntrinsicsCamera(Camera):
     """
 
     def __init__(self, fx, fy, cx, cy, znear=DEFAULT_Z_NEAR, zfar=DEFAULT_Z_FAR, name=None):
-        super(IntrinsicsCamera, self).__init__(znear=znear, zfar=zfar, name=name)
+        super().__init__(znear=znear, zfar=zfar, name=name)
 
         self.fx = fx
         self.fy = fy
diff --git a/genesis/ext/pyrender/light.py b/genesis/ext/pyrender/light.py
@@ -135,11 +135,7 @@ class DirectionalLight(Light):
     """
 
     def __init__(self, color=None, intensity=None, name=None):
-        super(DirectionalLight, self).__init__(
-            color=color,
-            intensity=intensity,
-            name=name,
-        )
+        super().__init__(color=color, intensity=intensity, name=name)
 
     def _generate_shadow_texture(self, size=None):
         """Generate a shadow texture for this light.
@@ -191,11 +187,7 @@ class PointLight(Light):
     """
 
     def __init__(self, color=None, intensity=None, range=None, name=None):
-        super(PointLight, self).__init__(
-            color=color,
-            intensity=intensity,
-            name=name,
-        )
+        super().__init__(color=color, intensity=intensity, name=name)
         self.range = range
 
     @property
@@ -306,11 +298,7 @@ class SpotLight(Light):
     def __init__(
         self, color=None, intensity=None, range=None, innerConeAngle=0.0, outerConeAngle=(np.pi / 4.0), name=None
     ):
-        super(SpotLight, self).__init__(
-            name=name,
-            color=color,
-            intensity=intensity,
-        )
+        super().__init__(name=name, color=color, intensity=intensity)
         self.outerConeAngle = outerConeAngle
         self.innerConeAngle = innerConeAngle
         self.range = range
diff --git a/genesis/ext/pyrender/material.py b/genesis/ext/pyrender/material.py
@@ -377,7 +377,7 @@ def __init__(
         roughnessFactor=1.0,
         metallicRoughnessTexture=None,
     ):
-        super(MetallicRoughnessMaterial, self).__init__(
+        super().__init__(
             name=name,
             normalTexture=normalTexture,
             occlusionTexture=occlusionTexture,
@@ -586,7 +586,7 @@ def __init__(
         glossinessFactor=1.0,
         specularGlossinessTexture=None,
     ):
-        super(SpecularGlossinessMaterial, self).__init__(
+        super().__init__(
             name=name,
             normalTexture=normalTexture,
             occlusionTexture=occlusionTexture,
diff --git a/genesis/ext/pyrender/platforms/egl.py b/genesis/ext/pyrender/platforms/egl.py
@@ -117,7 +117,7 @@ class EGLPlatform(Platform):
     """Renders using EGL."""
 
     def __init__(self, viewport_width, viewport_height, device_id: int | None = None):
-        super(EGLPlatform, self).__init__(viewport_width, viewport_height)
+        super().__init__(viewport_width, viewport_height)
         if _eglQueryDevicesEXT is None and device_id not in (0, None):
             raise RuntimeError("EGL platform plugin is not available. Enforcing specific EGL device not supported.")
         self._egl_device_id = device_id
diff --git a/genesis/ext/pyrender/platforms/osmesa.py b/genesis/ext/pyrender/platforms/osmesa.py
@@ -10,7 +10,7 @@ class OSMesaPlatform(Platform):
     """
 
     def __init__(self, viewport_width, viewport_height):
-        super(OSMesaPlatform, self).__init__(viewport_width, viewport_height)
+        super().__init__(viewport_width, viewport_height)
         self._context = None
         self._buffer = None
 
diff --git a/genesis/ext/pyrender/platforms/pyglet_platform.py b/genesis/ext/pyrender/platforms/pyglet_platform.py
@@ -13,7 +13,7 @@ class PygletPlatform(Platform):
     """
 
     def __init__(self, viewport_width, viewport_height):
-        super(PygletPlatform, self).__init__(viewport_width, viewport_height)
+        super().__init__(viewport_width, viewport_height)
         self._window = None
 
     def init_context(self):
diff --git a/genesis/ext/pyrender/viewer.py b/genesis/ext/pyrender/viewer.py
@@ -1238,7 +1238,7 @@ def start(self, auto_refresh=True):
             # This approach avoids "flickering" when creating and closing an invalid context. Besides, it avoids
             # "frozen" graphical window during compilation that would be interpreted as as bug by the end-user.
             try:
-                super(Viewer, self).__init__(
+                super().__init__(
                     config=conf,
                     visible=False,
                     resizable=True,
diff --git a/genesis/logging/logger.py b/genesis/logging/logger.py
@@ -10,7 +10,7 @@
 
 class GenesisFormatter(logging.Formatter):
     def __init__(self, verbose_time=True):
-        super(GenesisFormatter, self).__init__()
+        super().__init__()
 
         self.mapping = {
             logging.DEBUG: colors.GREEN,
diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py
@@ -119,6 +119,7 @@ def __exit__(self, exc_type, exc_value, traceback):
 def assert_initialized(cls):
     original_init = cls.__init__
 
+    @functools.wraps(original_init)
     def new_init(self, *args, **kwargs):
         if not gs._initialized:
             raise RuntimeError("Genesis hasn't been initialized. Did you call `gs.init()`?")
@@ -180,8 +181,6 @@ def get_device(backend: gs_backend, device_idx: Optional[int] = None):
         if not torch.cuda.is_available():
             gs.raise_exception("torch cuda not available")
 
-        if device_idx is None:
-            device_idx = torch.cuda.current_device()
         device = torch.device("cuda", device_idx)
         device_property = torch.cuda.get_device_properties(device)
         device_name = device_property.name
diff --git a/genesis/utils/particle.py b/genesis/utils/particle.py
@@ -362,11 +362,9 @@ def parse_args(backend):
             compute_normals=True,
             enable_multi_threading=True,
         )
-        mesh = trimesh.Trimesh(
-            vertices=mesh_with_data.mesh.vertices,
-            faces=mesh_with_data.mesh.triangles,
-            face_normals=mesh_with_data.get_point_attribute("normals"),
-        )
+        normals = mesh_with_data.get_point_attribute("normals")
+        vertices, triangles = mesh_with_data.take_mesh().take_vertices_and_triangles()
+        mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, face_normals=normals)
         gs.logger.debug(f"[splashsurf]: reconstruct vertices: {mesh.vertices.shape}, {mesh.faces.shape}")
         return mesh
 
diff --git a/genesis/utils/usda.py b/genesis/utils/usda.py
@@ -309,7 +309,7 @@ def parse_mesh_usd(path, group_by_material, scale, surface, bake_cache=True):
             "--usd_material_paths",
             *map(str, baked_materials.values()),
             "--device",
-            str(device.index),
+            str(device.index if device.index is not None else 0),
             "--log_level",
             logging.getLevelName(gs.logger.level).lower(),
         ]
diff --git a/genesis/vis/batch_renderer.py b/genesis/vis/batch_renderer.py
@@ -90,7 +90,7 @@ def build(self):
         rigid = self._visualizer.scene.rigid_solver
         n_envs = max(self._visualizer.scene.n_envs, 1)
         res = self._cameras[0].res
-        gpu_id = gs.device.index
+        gpu_id = gs.device.index if gs.device.index is not None else 0
         use_rasterizer = self._renderer_options.use_rasterizer
 
         # Cameras
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,4 +1,5 @@
 import base64
+import ctypes
 import gc
 import os
 import re
@@ -115,6 +116,49 @@ def _get_gpu_indices():
     return (0,)
 
 
+def _get_egl_index(gpu_index):
+    from OpenGL import EGL
+    from OpenGL.EGL.NV.device_cuda import EGL_CUDA_DEVICE_NV
+
+    # Get the list of Nvidia GPU that are visible
+    nvidia_gpu_indices = _get_gpu_indices()
+
+    # Define some ctypes for convenience
+    EGLDeviceEXT = ctypes.c_void_p
+    EGLAttrib = ctypes.c_ssize_t
+    EGLint = ctypes.c_int
+    EGLuint = ctypes.c_uint
+
+    # Load EGL extension functions dynamically
+    EGLuint = ctypes.c_uint
+    eglQueryDevicesEXT_addr = EGL.eglGetProcAddress(b"eglQueryDevicesEXT")
+    if not eglQueryDevicesEXT_addr:
+        raise RuntimeError("eglQueryDevicesEXT not available")
+    eglQueryDevicesEXT = ctypes.CFUNCTYPE(EGLuint, EGLint, ctypes.POINTER(EGLDeviceEXT), ctypes.POINTER(EGLint))(
+        eglQueryDevicesEXT_addr
+    )
+    eglQueryDeviceAttribEXT_addr = EGL.eglGetProcAddress(b"eglQueryDeviceAttribEXT")
+    if not eglQueryDeviceAttribEXT_addr:
+        raise RuntimeError("eglQueryDeviceAttribEXT not available")
+    eglQueryDeviceAttribEXT = ctypes.CFUNCTYPE(EGLuint, EGLDeviceEXT, EGLint, ctypes.POINTER(EGLAttrib))(
+        eglQueryDeviceAttribEXT_addr
+    )
+
+    # Query EGL devices
+    num_devices = EGLint()
+    eglQueryDevicesEXT(0, None, ctypes.byref(num_devices))
+    devices = (EGLDeviceEXT * num_devices.value)()
+    eglQueryDevicesEXT(num_devices, devices, ctypes.byref(num_devices))
+    egl_map = {}
+    for i in range(num_devices.value):
+        dev = devices[i]
+        cuda_id = EGLAttrib()
+        if eglQueryDeviceAttribEXT(dev, EGL_CUDA_DEVICE_NV, ctypes.byref(cuda_id)):
+            egl_map[nvidia_gpu_indices[cuda_id.value]] = i
+
+    return egl_map[gpu_index]
+
+
 def pytest_xdist_auto_num_workers(config):
     import genesis as gs
 
@@ -161,14 +205,18 @@ def pytest_xdist_auto_num_workers(config):
 
 
 def pytest_runtest_setup(item):
-    # Enforce GPU affinity that distributed framework is enabled
+    # Enforce GPU affinity if distributed framework is enabled
     worker_id = os.environ.get("PYTEST_XDIST_WORKER")
     if worker_id and worker_id.startswith("gw"):
         worker_num = int(worker_id[2:])
         gpu_indices = _get_gpu_indices()
-        gpu_num = worker_num % len(gpu_indices)
-        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_indices[gpu_num])
-        os.environ["TI_VISIBLE_DEVICE"] = str(gpu_indices[gpu_num])
+        gpu_index = gpu_indices[worker_num % len(gpu_indices)]
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_index)
+        os.environ["TI_VISIBLE_DEVICE"] = str(gpu_index)
+        try:
+            os.environ["EGL_DEVICE_ID"] = str(_get_egl_index(gpu_index))
+        except Exception:
+            pass
 
 
 def pytest_addoption(parser):
@@ -331,6 +379,11 @@ def initialize_genesis(request, backend, precision, taichi_offline_cache):
             pytest.skip(f"Backend '{backend}' not available on this machine")
         gs.init(backend=backend, precision=precision, debug=debug, seed=0, logging_level=logging_level)
 
+        if gs.backend != gs.cpu:
+            device_index = gs.device.index
+            if device_index is not None and device_index not in _get_gpu_indices():
+                assert RuntimeError("Wrong CUDA GPU device.")
+
         import gstaichi as ti
 
         ti_runtime = ti.lang.impl.get_runtime()
diff --git a/tests/test_bvh.py b/tests/test_bvh.py
@@ -1,5 +1,3 @@
-import torch
-import gstaichi as ti
 import numpy as np
 import pytest
 
@@ -24,12 +22,14 @@ def lbvh():
 
     lbvh = LBVH(aabb, max_n_query_result_per_aabb=32)
     lbvh.build()
+
     return lbvh
 
 
 @pytest.mark.required
 def test_morton_code(lbvh):
     morton_codes = lbvh.morton_codes.to_numpy()
+
     # Check that the morton codes are sorted
     for i_b in range(morton_codes.shape[0]):
         for i in range(1, morton_codes.shape[1]):
@@ -121,14 +121,15 @@ def test_build_tree(lbvh):
                 assert_allclose(parent_max, parent_max_expected, atol=1e-6, rtol=1e-5)
 
 
-@ti.kernel
-def query_kernel(lbvh: ti.template(), aabbs: ti.template()):
-    lbvh.query(aabbs)
-
-
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu, gs.gpu])
 def test_query(lbvh):
+    import gstaichi as ti
+
+    @ti.kernel
+    def query_kernel(lbvh: ti.template(), aabbs: ti.template()):
+        lbvh.query(aabbs)
+
     aabbs = lbvh.aabbs
 
     # Query the tree
diff --git a/tests/test_fem.py b/tests/test_fem.py
diff --git a/tests/test_mesh.py b/tests/test_mesh.py
diff --git a/tests/test_pbd.py b/tests/test_pbd.py
diff --git a/tests/test_rigid_benchmarks.py b/tests/test_rigid_benchmarks.py
diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py
diff --git a/tests/test_utils.py b/tests/test_utils.py
