#!/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 = """ """ 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())