Spaces:
Runtime error
Runtime error
File size: 7,604 Bytes
e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 9b8c4e3 e6e8290 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
#!/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()) |