Merge branch 'INSTX-5919_make_auto_device_resolution_available' into …

…'master' INSTX-5919 Make auto device resolution available to server See merge request machine-learning/dorado!1137
nanoporetech · Jul 30, 2024 · df605b6 · df605b6
2 parents 7d74246 + c4da8b5
commit df605b6
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 14 deletions.
diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp
@@ -25,6 +25,7 @@
 #include "read_pipeline/ReadFilterNode.h"
 #include "read_pipeline/ReadToBamTypeNode.h"
 #include "read_pipeline/ResumeLoader.h"
+#include "torch_utils/auto_detect_device.h"
 #include "utils/SampleSheet.h"
 #include "utils/arg_parse_ext.h"
 #include "utils/bam_utils.h"
@@ -656,7 +657,7 @@ int basecaller(int argc, char* argv[]) {
         return EXIT_FAILURE;
     }
     if (device == cli::AUTO_DETECT_DEVICE) {
-        device = cli::get_auto_detected_device();
+        device = utils::get_auto_detected_device();
     }
 
     auto hts_file = cli::extract_hts_file(parser);

diff --git a/dorado/cli/cli_utils.h b/dorado/cli/cli_utils.h
@@ -150,16 +150,6 @@ inline bool validate_device_string(const std::string& device) {
     return false;
 }
 
-inline std::string get_auto_detected_device() {
-#if DORADO_METAL_BUILD
-    return "metal";
-#elif DORADO_CUDA_BUILD
-    return torch::cuda::is_available() ? "cuda:all" : "cpu";
-#else
-    return "cpu";
-#endif
-}
-
 }  // namespace cli
 
 }  // namespace dorado
diff --git a/dorado/cli/correct.cpp b/dorado/cli/correct.cpp
@@ -5,6 +5,7 @@
 #include "read_pipeline/CorrectionNode.h"
 #include "read_pipeline/ErrorCorrectionMapperNode.h"
 #include "read_pipeline/HtsWriter.h"
+#include "torch_utils/auto_detect_device.h"
 #include "torch_utils/torch_utils.h"
 #include "utils/arg_parse_ext.h"
 #include "utils/fs_utils.h"
@@ -92,7 +93,7 @@ int correct(int argc, char* argv[]) {
 #if DORADO_METAL_BUILD
         device = "cpu";
 #else
-        device = cli::get_auto_detected_device();
+        device = utils::get_auto_detected_device();
 #endif
     }
 

diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp
@@ -19,11 +19,11 @@
 #include "read_pipeline/ProgressTracker.h"
 #include "read_pipeline/ReadFilterNode.h"
 #include "read_pipeline/ReadToBamTypeNode.h"
+#include "torch_utils/auto_detect_device.h"
 #include "utils/SampleSheet.h"
 #include "utils/arg_parse_ext.h"
 #include "utils/bam_utils.h"
 #include "utils/basecaller_utils.h"
-
 #if DORADO_CUDA_BUILD
 #include "torch_utils/cuda_utils.h"
 #endif
@@ -365,7 +365,7 @@ int duplex(int argc, char* argv[]) {
             return EXIT_FAILURE;
         }
         if (device == cli::AUTO_DETECT_DEVICE) {
-            device = cli::get_auto_detected_device();
+            device = utils::get_auto_detected_device();
         }
 
         auto model(parser.visible.get<std::string>("model"));

diff --git a/dorado/torch_utils/CMakeLists.txt b/dorado/torch_utils/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_library(dorado_torch_utils
+    auto_detect_device.h    
     duplex_utils.cpp
     duplex_utils.h
     gpu_monitor.cpp

diff --git a/dorado/torch_utils/auto_detect_device.h b/dorado/torch_utils/auto_detect_device.h
@@ -0,0 +1,21 @@
+#pragma once
+
+#if DORADO_CUDA_BUILD
+#include <torch/cuda.h>
+#endif
+
+#include <string>
+
+namespace dorado::utils {
+
+inline std::string get_auto_detected_device() {
+#if DORADO_METAL_BUILD
+    return "metal";
+#elif DORADO_CUDA_BUILD
+    return torch::cuda::is_available() ? "cuda:all" : "cpu";
+#else
+    return "cpu";
+#endif
+}
+
+}  // namespace dorado::utils