Merge branch 'cherry-pick-5392e88e' into 'release-v0.7'

DOR-741: Cherry pick modbase-gpu-mutex into 0.7 See merge request machine-learning/dorado!1112
nanoporetech · Jul 12, 2024 · 5ddfc2f · 5ddfc2f
2 parents 56d3e8e + 5a61234
commit 5ddfc2f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/dorado/modbase/ModBaseCaller.cpp b/dorado/modbase/ModBaseCaller.cpp
@@ -212,6 +212,7 @@ void ModBaseCaller::start_threads() {
 void ModBaseCaller::modbase_task_thread_fn(size_t model_id) {
     auto& caller_data = m_caller_data[model_id];
 #if DORADO_CUDA_BUILD
+    static std::vector<std::mutex> gpu_mutexes(torch::cuda::device_count());
     c10::cuda::OptionalCUDAStreamGuard stream_guard(caller_data->stream);
 #endif
     while (true) {
@@ -231,6 +232,14 @@ void ModBaseCaller::modbase_task_thread_fn(size_t model_id) {
         caller_data->input_queue.pop_back();
         input_lock.unlock();
 
+#if DORADO_CUDA_BUILD
+        auto gpu_lock = [&] {
+            if (m_options.device().is_cuda()) {
+                return std::unique_lock(gpu_mutexes[m_options.device().index()]);
+            }
+            return std::unique_lock<std::mutex>{};
+        }();
+#endif
         std::unique_lock<std::mutex> task_lock(task->mut);
         stats::Timer timer;
         task->out = caller_data->module_holder->forward(task->input_sigs, task->input_seqs);

diff --git a/dorado/utils/torch_utils.cpp b/dorado/utils/torch_utils.cpp
@@ -35,7 +35,7 @@ void set_torch_allocator_max_split_size() {
     const char *pytorch_cuda_alloc_conf = std::getenv("PYTORCH_CUDA_ALLOC_CONF");
     if (pytorch_cuda_alloc_conf != nullptr) {
         std::string_view str(pytorch_cuda_alloc_conf);
-        if (str.find("max_split_size_mb") != std::string::npos) {
+        if (str.find("max_split_size_mb") != std::string_view::npos) {
             // user has set this via env_var - let torch parse and use their value
             return;
         }