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