Spaces:
Runtime error
Runtime error
#!/usr/bin/env python3 | |
""" | |
GPU initialization script for MuJoCo EGL rendering in containerized environments. | |
This should be run before starting the notebook to ensure GPU is properly set up. | |
""" | |
import os | |
import subprocess | |
import sys | |
def check_nvidia_driver(): | |
"""Check if NVIDIA driver is accessible.""" | |
try: | |
result = subprocess.run(['nvidia-smi'], capture_output=True, text=True) | |
if result.returncode == 0: | |
print("β NVIDIA driver accessible") | |
print(result.stdout.split('\n')[2]) # Driver info line | |
return True | |
else: | |
print("β NVIDIA driver not accessible") | |
return False | |
except FileNotFoundError: | |
print("β nvidia-smi not found") | |
return False | |
def check_egl_libs(): | |
"""Check if EGL libraries are available.""" | |
import ctypes | |
libs_to_check = [ | |
'libEGL.so.1', | |
'libGL.so.1', | |
'libEGL_nvidia.so.0' | |
] | |
for lib in libs_to_check: | |
try: | |
ctypes.CDLL(lib) | |
print(f"β {lib} loaded successfully") | |
except OSError as e: | |
print(f"β Failed to load {lib}: {e}") | |
def test_egl_device(): | |
"""Test EGL device creation with multiple approaches.""" | |
try: | |
from OpenGL import EGL | |
import ctypes | |
# Method 1: Try platform device display (preferred for headless) | |
try: | |
display = EGL.eglGetPlatformDisplay(EGL.EGL_PLATFORM_DEVICE_EXT, | |
EGL.EGL_DEFAULT_DISPLAY, | |
None) | |
if display != EGL.EGL_NO_DISPLAY: | |
major = ctypes.c_long() | |
minor = ctypes.c_long() | |
if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)): | |
print(f"β EGL platform device initialized (version {major.value}.{minor.value})") | |
EGL.eglTerminate(display) | |
return True | |
except Exception as e: | |
print(f" Platform device method failed: {e}") | |
# Method 2: Try default display | |
try: | |
display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY) | |
if display != EGL.EGL_NO_DISPLAY: | |
major = ctypes.c_long() | |
minor = ctypes.c_long() | |
if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)): | |
print(f"β EGL default display initialized (version {major.value}.{minor.value})") | |
EGL.eglTerminate(display) | |
return True | |
except Exception as e: | |
print(f" Default display method failed: {e}") | |
# Method 3: Try surfaceless context (what MuJoCo likely uses) | |
try: | |
os.environ['EGL_PLATFORM'] = 'surfaceless' | |
display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY) | |
if display != EGL.EGL_NO_DISPLAY: | |
major = ctypes.c_long() | |
minor = ctypes.c_long() | |
if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)): | |
print(f"β EGL surfaceless display initialized (version {major.value}.{minor.value})") | |
EGL.eglTerminate(display) | |
return True | |
except Exception as e: | |
print(f" Surfaceless method failed: {e}") | |
print("β All EGL initialization methods failed") | |
return False | |
except Exception as e: | |
print(f"β EGL test failed: {e}") | |
return False | |
def test_mujoco_rendering(): | |
"""Test MuJoCo rendering capability with different approaches.""" | |
try: | |
import mujoco | |
# Create a simple model | |
xml = """ | |
<mujoco> | |
<worldbody> | |
<body> | |
<geom type="box" size="1 1 1"/> | |
</body> | |
</worldbody> | |
</mujoco> | |
""" | |
model = mujoco.MjModel.from_xml_string(xml) | |
# Try different rendering approaches | |
approaches = [ | |
("Small resolution", {"height": 64, "width": 64}), | |
("Default resolution", {"height": 240, "width": 320}), | |
("Large resolution", {"height": 480, "width": 640}), | |
] | |
for name, kwargs in approaches: | |
try: | |
renderer = mujoco.Renderer(model, **kwargs) | |
data = mujoco.MjData(model) | |
renderer.update_scene(data) | |
pixels = renderer.render() | |
print(f" β {name} ({kwargs['width']}x{kwargs['height']}): SUCCESS") | |
print(f" Image shape: {pixels.shape}, dtype: {pixels.dtype}") | |
# Test if we got actual rendered content (not all zeros) | |
if pixels.max() > 0: | |
print(f" β Non-zero pixels detected (max value: {pixels.max()})") | |
else: | |
print(f" β οΈ All pixels are zero - may indicate rendering issue") | |
# Clean up | |
del renderer | |
return True | |
except Exception as e: | |
print(f" β {name}: {e}") | |
continue | |
print("β All MuJoCo rendering approaches failed") | |
return False | |
except ImportError: | |
print("β MuJoCo not installed") | |
return False | |
except Exception as e: | |
print(f"β MuJoCo test failed: {e}") | |
return False | |
def main(): | |
"""Run all GPU initialization checks.""" | |
print("π§ Initializing GPU for MuJoCo rendering...") | |
print("=" * 50) | |
# Set environment variables | |
os.environ['MUJOCO_GL'] = 'egl' | |
os.environ['PYOPENGL_PLATFORM'] = 'egl' | |
os.environ['EGL_PLATFORM'] = 'surfaceless' # Better for headless | |
print("Environment variables set:") | |
print(f" MUJOCO_GL: {os.environ.get('MUJOCO_GL')}") | |
print(f" PYOPENGL_PLATFORM: {os.environ.get('PYOPENGL_PLATFORM')}") | |
print(f" EGL_PLATFORM: {os.environ.get('EGL_PLATFORM')}") | |
print() | |
# Run checks | |
checks = [ | |
("NVIDIA Driver", check_nvidia_driver), | |
("EGL Libraries", lambda: check_egl_libs() or True), # Always continue | |
("EGL Device", test_egl_device), | |
("MuJoCo Rendering", test_mujoco_rendering), | |
] | |
results = [] | |
for name, check_func in checks: | |
print(f"Checking {name}...") | |
try: | |
result = check_func() | |
results.append((name, result)) | |
except Exception as e: | |
print(f"β {name} check failed with exception: {e}") | |
results.append((name, False)) | |
print() | |
# Summary | |
print("=" * 50) | |
print("π Summary:") | |
# Check if the critical test (MuJoCo) passed | |
mujoco_passed = any(name == "MuJoCo Rendering" and passed for name, passed in results) | |
for name, passed in results: | |
status = "β PASS" if passed else "β FAIL" | |
print(f" {name}: {status}") | |
if mujoco_passed: | |
print("\nπ MuJoCo rendering works! The notebook should work even if some EGL tests fail.") | |
print("π‘ Note: EGL device tests may fail but MuJoCo can still render successfully.") | |
return 0 | |
else: | |
print("\nβ οΈ MuJoCo rendering failed. GPU rendering will not work properly.") | |
print("π‘ Try checking the container GPU configuration or driver compatibility.") | |
return 1 | |
if __name__ == "__main__": | |
sys.exit(main()) |