jbilcke-hf's picture
jbilcke-hf HF Staff
more logs
9b8c4e3
#!/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())