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.""" | |
try: | |
# Try to create an EGL display | |
from OpenGL import EGL | |
import ctypes | |
# Get EGL display | |
display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY) | |
if display == EGL.EGL_NO_DISPLAY: | |
print("β Failed to get EGL display") | |
return False | |
# Initialize EGL | |
major = ctypes.c_long() | |
minor = ctypes.c_long() | |
if not EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)): | |
print("β Failed to initialize EGL") | |
return False | |
print(f"β EGL initialized successfully (version {major.value}.{minor.value})") | |
# Clean up | |
EGL.eglTerminate(display) | |
return True | |
except Exception as e: | |
print(f"β EGL test failed: {e}") | |
return False | |
def test_mujoco_rendering(): | |
"""Test MuJoCo rendering capability.""" | |
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 to create a renderer (this is where EGL issues usually surface) | |
try: | |
renderer = mujoco.Renderer(model, height=240, width=320) | |
print("β MuJoCo renderer created successfully") | |
return True | |
except Exception as e: | |
print(f"β MuJoCo renderer creation failed: {e}") | |
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'] = 'device' | |
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:") | |
all_passed = True | |
for name, passed in results: | |
status = "β PASS" if passed else "β FAIL" | |
print(f" {name}: {status}") | |
if not passed: | |
all_passed = False | |
if all_passed: | |
print("\nπ All checks passed! GPU rendering should work.") | |
return 0 | |
else: | |
print("\nβ οΈ Some checks failed. GPU rendering may not work properly.") | |
return 1 | |
if __name__ == "__main__": | |
sys.exit(main()) |