|
|
|
|
|
""" |
|
|
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]) |
|
|
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: |
|
|
|
|
|
from OpenGL import EGL |
|
|
import ctypes |
|
|
|
|
|
|
|
|
display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY) |
|
|
if display == EGL.EGL_NO_DISPLAY: |
|
|
print("✗ Failed to get EGL display") |
|
|
return False |
|
|
|
|
|
|
|
|
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})") |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
xml = """ |
|
|
<mujoco> |
|
|
<worldbody> |
|
|
<body> |
|
|
<geom type="box" size="1 1 1"/> |
|
|
</body> |
|
|
</worldbody> |
|
|
</mujoco> |
|
|
""" |
|
|
|
|
|
model = mujoco.MjModel.from_xml_string(xml) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
|
|
|
checks = [ |
|
|
("NVIDIA Driver", check_nvidia_driver), |
|
|
("EGL Libraries", lambda: check_egl_libs() or True), |
|
|
("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() |
|
|
|
|
|
|
|
|
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()) |
