rollout use threadpool as translation unit

Author: aftersomemath

python/mujoco/CMakeLists.txt

@@ -383,7 +383,7 @@ target_link_libraries(
           structs_header
 )
 
-mujoco_pybind11_module(_rollout rollout.cc)
+mujoco_pybind11_module(_rollout rollout.cc threadpool.cc)
 target_link_libraries(_rollout PRIVATE functions_header mujoco raw)
 
 mujoco_pybind11_module(

python/mujoco/rollout.cc

@@ -20,120 +20,18 @@
 #include "errors.h"
 #include "raw.h"
 #include "structs.h"
+#include "threadpool.h"
 #include <pybind11/buffer_info.h>
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
-#include <condition_variable>
-#include <cstdint>
-#include <functional>
-#include <mutex>
-#include <queue>
-#include <thread>
-#include <vector>
-#include <absl/base/attributes.h>
-
 namespace mujoco::python {
 
 namespace {
 
 namespace py = ::pybind11;
 
-// Copied from https://github.com/google-deepmind/mujoco_mpc/blob/main/mjpc/threadpool.h
-// ThreadPool class
-class ThreadPool {
- public:
-  // constructor
-  explicit ThreadPool(int num_threads);
-  // destructor
-  ~ThreadPool();
-  int NumThreads() const { return threads_.size(); }
-  // returns an ID between 0 and NumThreads() - 1. must be called within
-  // worker thread (returns -1 if not).
-  static int WorkerId() { return worker_id_; }
-  // ----- methods ----- //
-  // set task for threadpool
-  void Schedule(std::function<void()> task);
-  // return number of tasks completed
-  std::uint64_t GetCount() { return ctr_; }
-  // reset count to zero
-  void ResetCount() { ctr_ = 0; }
-  // wait for count, then return
-  void WaitCount(int value) {
-    std::unique_lock<std::mutex> lock(m_);
-    cv_ext_.wait(lock, [&]() { return this->GetCount() >= value; });
-  }
- private:
-  // ----- methods ----- //
-  // execute task with available thread
-  void WorkerThread(int i);
-  ABSL_CONST_INIT static thread_local int worker_id_;
-  // ----- members ----- //
-  std::vector<std::thread> threads_;
-  std::mutex m_;
-  std::condition_variable cv_in_;
-  std::condition_variable cv_ext_;
-  std::queue<std::function<void()>> queue_;
-  std::uint64_t ctr_;
-};
-
-// Copied from https://github.com/google-deepmind/mujoco_mpc/blob/main/mjpc/threadpool.cc
-ABSL_CONST_INIT thread_local int ThreadPool::worker_id_ = -1;
-// ThreadPool constructor
-ThreadPool::ThreadPool(int num_threads) : ctr_(0) {
-  for (int i = 0; i < num_threads; i++) {
-    threads_.push_back(std::thread(&ThreadPool::WorkerThread, this, i));
-  }
-}
-// ThreadPool destructor
-ThreadPool::~ThreadPool() {
-  {
-    std::unique_lock<std::mutex> lock(m_);
-    for (int i = 0; i < threads_.size(); i++) {
-      queue_.push(nullptr);
-    }
-    cv_in_.notify_all();
-  }
-  for (auto& thread : threads_) {
-    thread.join();
-  }
-}
-// ThreadPool scheduler
-void ThreadPool::Schedule(std::function<void()> task) {
-  std::unique_lock<std::mutex> lock(m_);
-  queue_.push(std::move(task));
-  cv_in_.notify_one();
-}
-// ThreadPool worker
-void ThreadPool::WorkerThread(int i) {
-  worker_id_ = i;
-  while (true) {
-    auto task = [&]() {
-      std::unique_lock<std::mutex> lock(m_);
-      cv_in_.wait(lock, [&]() { return !queue_.empty(); });
-      std::function<void()> task = std::move(queue_.front());
-      queue_.pop();
-      cv_in_.notify_one();
-      return task;
-    }();
-    if (task == nullptr) {
-      {
-        std::unique_lock<std::mutex> lock(m_);
-        ++ctr_;
-        cv_ext_.notify_one();
-      }
-      break;
-    }
-    task();
-    {
-      std::unique_lock<std::mutex> lock(m_);
-      ++ctr_;
-      cv_ext_.notify_one();
-    }
-  }
-}
-
 // NOLINTBEGIN(whitespace/line_length)
 
 const auto rollout_doc = R"(
@@ -265,19 +163,18 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
 static ThreadPool* pool = nullptr;
 void _unsafe_rollout_threaded(std::vector<const mjModel*>& m, std::vector<mjData*>& d,
                               int nroll, int nstep, unsigned int control_spec,
-                              const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control,
-                              mjtNum* state, mjtNum* sensordata,
+                              const mjtNum* state0, const mjtNum* warmstart0,
+                              const mjtNum* control, mjtNum* state, mjtNum* sensordata,
                               int nthread, int chunk_size) {
   int nfulljobs = nroll / chunk_size;
   int chunk_remainder = nroll % chunk_size;
-  int njobs = nfulljobs;
-  if (chunk_remainder > 0) njobs++;
+  int njobs = (chunk_remainder > 0) ? nfulljobs + 1 : nfulljobs;
 
   if (pool == nullptr) {
     pool = new ThreadPool(nthread);
   }
   else if (pool->NumThreads() != nthread) {
-    delete pool; // TODO make sure pool is shutdown correctly
+    delete pool;
     pool = new ThreadPool(nthread);
   } else {
     pool->ResetCount();

python/mujoco/threadpool.cc

@@ -1,4 +1,4 @@
-// Copyright 2022 DeepMind Technologies Limited
+// Copyright 2024 DeepMind Technologies Limited
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "mjpc/threadpool.h"
+#include "threadpool.h"
 
 #include <condition_variable>
 #include <functional>
@@ -22,7 +22,7 @@
 
 #include <absl/base/attributes.h>
 
-namespace mjpc {
+namespace mujoco::python {
 
 ABSL_CONST_INIT thread_local int ThreadPool::worker_id_ = -1;
 
@@ -84,4 +84,4 @@ void ThreadPool::WorkerThread(int i) {
   }
 }
 
-}  // namespace mjpc
+}  // namespace mujoco::python

python/mujoco/threadpool.h

@@ -1,4 +1,4 @@
-// Copyright 2022 DeepMind Technologies Limited
+// Copyright 2024 DeepMind Technologies Limited
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef MJPC_THREADPOOL_H_
-#define MJPC_THREADPOOL_H_
+#ifndef MUJOCO_PYTHON_THREADPOOL_H_
+#define MUJOCO_PYTHON_THREADPOOL_H_
 
 #include <condition_variable>
 #include <cstdint>
@@ -26,7 +26,7 @@
 
 #include <absl/base/attributes.h>
 
-namespace mjpc {
+namespace mujoco::python {
 
 // ThreadPool class
 class ThreadPool {
@@ -76,6 +76,6 @@ class ThreadPool {
   std::uint64_t ctr_;
 };
 
-}  // namespace mjpc
+}  // namespace mujoco::python
 
-#endif  // MJPC_THREADPOOL_H_
+#endif  // MUJOCO_PYTHON_THREADPOOL_H_