Add multi-GPU rendering support

examples/multigpu_rendering.py

@@ -0,0 +1,173 @@
+"""
+This is an example for rendering on multiple GPUs in parallel,
+using the multiprocessing module.
+"""
+import ctypes
+from multiprocessing import (
+    Process, set_start_method, freeze_support, Condition,
+    Value, Array)
+from time import perf_counter
+import numpy as np
+from mujoco_py import load_model_from_path, MjSim
+
+
+#
+# Experiment parameters
+#
+
+# Image size for rendering
+IMAGE_WIDTH = 255
+IMAGE_HEIGHT = 255
+# Number of frames to render per sim
+N_FRAMES = 2000
+# Number of sims to run in parallel (assumes one per GPU),
+# so N_SIMS=2 assumes there are 2 GPUs available.
+N_SIMS = 2
+
+
+class RenderProcess:
+    """
+    Wraps a multiprocessing.Process for rendering. Assumes there
+    is one MjSim per process.
+    """
+
+    def __init__(self, device_id, setup_sim, update_sim, output_var_shape):
+        """
+        Args:
+        - device_id (int): GPU device to use for rendering (0-indexed)
+        - setup_sim (callback): callback that is given a device_id and
+            returns a MjSim. It is responsible for making MjSim render
+            to given device.
+        - update_sim (callback): callback given a sim and device_id, and
+            should return a numpy array of shape `output_var_shape`.
+        - output_var_shape (tuple): shape of the synchronized output
+            array from `update_sim`.
+        """
+        self.device_id = device_id
+        self.setup_sim = setup_sim
+        self.update_sim = update_sim
+
+        # Create a synchronized output variable (numpy array)
+        self._shared_output_var = Array(
+            ctypes.c_double, int(np.prod(output_var_shape)))
+        self._output_var = np.frombuffer(
+            self._shared_output_var.get_obj())
+
+        # Number of variables used to communicate with process
+        self._cv = Condition()
+        self._ready = Value('b', 0)
+        self._start = Value('b', 0)
+        self._terminate = Value('b', 0)
+
+        # Start the actual process
+        self._process = Process(target=self._run)
+        self._process.start()
+
+    def wait(self):
+        """ Wait for process to be ready for another update call. """
+        with self._cv:
+            if self._start.value:
+                self._cv.wait()
+            if self._ready.value:
+                return
+            self._cv.wait()
+
+    def read(self, copy=False):
+        """ Reads the output variable. Returns a copy if copy=True. """
+        if copy:
+            with self._shared_output_var.get_lock():
+                return np.copy(self._output_var)
+        else:
+            return self._output_var
+
+    def update(self):
+        """ Calls update_sim asynchronously. """
+        with self._cv:
+            self._start.value = 1
+            self._cv.notify()
+
+    def stop(self):
+        """ Tells process to stop and waits for it to terminate. """
+        with self._cv:
+            self._terminate.value = 1
+            self._cv.notify()
+        self._process.join()
+
+    def _run(self):
+        sim = self.setup_sim(self.device_id)
+
+        while True:
+            with self._cv:
+                self._ready.value = 1
+                self._cv.notify_all()
+
+            with self._cv:
+                if not self._start.value and not self._terminate.value:
+                    self._cv.wait()
+                if self._terminate.value:
+                    break
+                assert self._start.value
+                self._start.value = 0
+
+            # Run the update and assign output variable
+            with self._shared_output_var.get_lock():
+                self._output_var[:] = self.update_sim(
+                    sim, self.device_id).ravel()
+
+
+def setup_sim(device_id):
+    model = load_model_from_path("xmls/fetch/main.xml")
+    sim = MjSim(model)
+
+    image = sim.render(
+        IMAGE_WIDTH, IMAGE_HEIGHT, device_id=device_id)
+    assert image.shape == (IMAGE_HEIGHT, IMAGE_WIDTH, 3)
+
+    return sim
+
+
+def update_sim(sim, device_id):
+    sim.step()
+    return sim.render(IMAGE_WIDTH, IMAGE_HEIGHT, device_id=device_id)
+
+
+def main():
+    print("main(): create processes", flush=True)
+    processes = []
+    for device_id in range(N_SIMS):
+        p = RenderProcess(
+            device_id, setup_sim, update_sim,
+            (IMAGE_HEIGHT, IMAGE_WIDTH, 3))
+        processes.append(p)
+
+    for p in processes:
+        p.wait()
+
+    print("main(): start benchmarking", flush=True)
+    start_t = perf_counter()
+
+    for _ in range(N_FRAMES):
+        for p in processes:
+            p.update()
+
+        for p in processes:
+            p.wait()
+
+        for p in processes:
+            p.read()
+
+    t = perf_counter() - start_t
+    print("Completed in %.1fs: %.3fms, %.1f FPS" % (
+            t, t / (N_FRAMES * N_SIMS) * 1000, (N_FRAMES * N_SIMS) / t),
+          flush=True)
+
+    print("main(): stopping processes", flush=True)
+    for p in processes:
+        p.stop()
+
+    print("main(): finished", flush=True)
+
+
+if __name__ == "__main__":
+    set_start_method('spawn')
+    main()

mujoco_py/__init__.py

@@ -9,11 +9,17 @@
 MjSim = cymj.MjSim
 MjSimState = cymj.MjSimState
 MjSimPool = cymj.MjSimPool
-PyMjModel = cymj.PyMjModel
-PyMjData = cymj.PyMjData
+MjRenderContext = cymj.MjRenderContext
+MjRenderContextOffscreen = cymj.MjRenderContextOffscreen
+MjRenderContextWindow = cymj.MjRenderContextWindow
+
 
 # Public API:
-__all__ = ['MjSim', 'MjSimState', 'MjSimPool', 'MjViewer', "MjViewerBasic", "MujocoException",
-           'load_model_from_path', 'load_model_from_xml', 'load_model_from_mjb',
+__all__ = ['MjSim', 'MjSimState', 'MjSimPool',
+           'MjRenderContextOffscreen', 'MjRenderContextWindow',
+           'MjRenderContext', 'MjViewer', 'MjViewerBasic',
+           'MujocoException',
+           'load_model_from_path', 'load_model_from_xml',
+           'load_model_from_mjb',
            'ignore_mujoco_warnings', 'const', "functions",
            "__version__", "get_version"]

mujoco_py/builder.py

@@ -8,6 +8,7 @@
 from distutils.dist import Distribution
 from distutils.sysconfig import customize_compiler
 from os.path import abspath, dirname, exists, join, getmtime
+from shutil import move
 
 import numpy as np
 from Cython.Build import cythonize
@@ -50,8 +51,9 @@ def load_cython_ext(mjpro_path):
         raise RuntimeError("Unsupported platform %s" % sys.platform)
 
     builder = Builder(mjpro_path)
-
-    cext_so_path = builder.build()
+    cext_so_path = builder.get_so_file_path()
+    if not exists(cext_so_path):
+        cext_so_path = builder.build()
     mod = imp.load_dynamic("cymj", cext_so_path)
     return mod
 
@@ -114,6 +116,12 @@ def __init__(self, mjpro_path):
             language='c')
 
     def build(self):
+        built_so_file_path = self._build_impl()
+        new_so_file_path = self.get_so_file_path()
+        move(built_so_file_path, new_so_file_path)
+        return new_so_file_path
+
+    def _build_impl(self):
         dist = Distribution({
             "script_name": None,
             "script_args": ["build_ext"]
@@ -129,8 +137,13 @@ def build(self):
         dist.parse_command_line()
         obj_build_ext = dist.get_command_obj("build_ext")
         dist.run_commands()
-        so_file_path, = obj_build_ext.get_outputs()
-        return so_file_path
+        built_so_file_path, = obj_build_ext.get_outputs()
+        return built_so_file_path
+
+    def get_so_file_path(self):
+        dir_path = abspath(dirname(__file__))
+        return join(dir_path, "generated", "cymj_%s.so" % (
+            self.__class__.__name__.lower()))
 
 
 class WindowsExtensionBuilder(MujocoExtensionBuilder):
@@ -162,8 +175,8 @@ def __init__(self, mjpro_path):
         self.extension.libraries.extend(['glewegl'])
         self.extension.runtime_library_dirs = [join(mjpro_path, 'bin')]
 
-    def build(self):
-        so_file_path = super().build()
+    def _build_impl(self):
+        so_file_path = super()._build_impl()
         nvidia_lib_dir = '/usr/local/nvidia/lib64/'
         fix_shared_library(so_file_path, 'libOpenGL.so',
                            join(nvidia_lib_dir, 'libOpenGL.so.0'))
@@ -182,7 +195,7 @@ def __init__(self, mjpro_path):
         self.extension.define_macros = [('ONMAC', None)]
         self.extension.runtime_library_dirs = [join(mjpro_path, 'bin')]
 
-    def build(self):
+    def _build_impl(self):
         # Prefer GCC 6 for now since GCC 7 may behave differently.
         c_compilers = ['/usr/local/bin/gcc-6', '/usr/local/bin/gcc-7']
         available_c_compiler = None
@@ -197,7 +210,7 @@ def build(self):
                 '`brew install gcc --without-multilib`.')
         os.environ['CC'] = available_c_compiler
 
-        so_file_path = super().build()
+        so_file_path = super()._build_impl()
         del os.environ['CC']
         return self.manually_link_libraries(so_file_path)
 

mujoco_py/cymj.pyx

@@ -28,7 +28,8 @@ cdef extern from "gl/glshim.h":
 
     cdef int initOpenGL(int device_id)
     cdef void closeOpenGL()
-    cdef int setOpenGLBufferSize(int width, int height)
+    cdef int makeOpenGLContextCurrent(int device_id)
+    cdef int setOpenGLBufferSize(int device_id, int width, int height)
 
 # TODO: make this function or class so these comments turn into doc strings:
 

mujoco_py/gl/dummyshim.c

@@ -4,7 +4,11 @@ int initOpenGL(int device_id) {
     return 1;
 }
 
-int setOpenGLBufferSize(int width, int height) {
+int setOpenGLBufferSize(int device_id, int width, int height) {
+    return 1;
+}
+
+int makeOpenGLContextCurrent(int device_id) {
     return 1;
 }
 

mujoco_py/gl/eglshim.c

@@ -6,10 +6,18 @@
 
 #define MAX_DEVICES 8
 
-int is_initialized = 0;
+int is_device_initialized[MAX_DEVICES] = {0};
+EGLDisplay eglDisplays[MAX_DEVICES];
+EGLContext eglContexts[MAX_DEVICES];
 
 int initOpenGL(int device_id)
 {
+    if (device_id < 0 || device_id > MAX_DEVICES) {
+        printf("Device id outside of range.\n");
+        return -1;
+    }
+    int is_initialized = is_device_initialized[device_id];
+
     if (is_initialized)
         return 1;
 
@@ -94,34 +102,59 @@ int initOpenGL(int device_id)
         return -8;
     }
 
-    is_initialized = 1;
+    is_device_initialized[device_id] = 1;
+    eglDisplays[device_id] = eglDpy;
+    eglContexts[device_id] = eglCtx;
     return 1;
 }
 
-int setOpenGLBufferSize(int width, int height) {
+int makeOpenGLContextCurrent(device_id) {
+    if (device_id < 0 || device_id > MAX_DEVICES) {
+        printf("Device id outside of range.\n");
+        return -1;
+    }
+    if (!is_device_initialized[device_id])
+        return -2;
+
+    if( eglMakeCurrent(eglDisplays[device_id],
+                       EGL_NO_SURFACE,
+                       EGL_NO_SURFACE,
+                       eglContexts[device_id]) != EGL_TRUE ) {
+        eglDestroyContext(eglDisplays[device_id],
+                          eglContexts[device_id]);
+        printf("Could not make EGL context current\n");
+        return -3;
+    } else {
+        return 1;
+    }
+}
+
+int setOpenGLBufferSize(int device_id, int width, int height) {
     // Noop since we don't need to change buffer here.
     return 1;
 }
 
 void closeOpenGL()
 {
-    if (!is_initialized)
-        return;
+    for (int device_id=0; device_id<MAX_DEVICES; device_id++) {
+        if (!is_device_initialized[device_id])
+            continue;
 
-    EGLDisplay eglDpy = eglGetCurrentDisplay();
-    if( eglDpy==EGL_NO_DISPLAY )
-        return;
+        EGLDisplay eglDpy = eglDisplays[device_id];
+        if( eglDpy==EGL_NO_DISPLAY )
+            continue;
 
-    // get current context
-    EGLContext eglCtx = eglGetCurrentContext();
+        // get current context
+        EGLContext eglCtx = eglContexts[device_id];
 
-    // release context
-    eglMakeCurrent(eglDpy, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+        // release context
+        eglMakeCurrent(eglDpy, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
 
-    // destroy context if valid
-    if( eglCtx!=EGL_NO_CONTEXT )
-        eglDestroyContext(eglDpy, eglCtx);
+        // destroy context if valid
+        if( eglCtx!=EGL_NO_CONTEXT )
+            eglDestroyContext(eglDpy, eglCtx);
 
-    // terminate display
-    eglTerminate(eglDpy);
+        // terminate display
+        eglTerminate(eglDpy);
+    }
 }