GPU support for docker commands

crates/lib/google_drive_common/src/google_drive_download_command.rs

@@ -22,6 +22,8 @@ use tempdir::TempDir;
 pub struct GoogleDriveDownloadCommand {
   download_script: String,
   maybe_venv_activation_script: Option<String>,
+
+  /// If this is run under Docker (eg. in development), these are the options.
   maybe_docker_options: Option<DockerOptions>,
 }
 
@@ -75,18 +77,7 @@ impl GoogleDriveDownloadCommand {
     }
 
     if let Some(docker_options) = self.maybe_docker_options.as_ref() {
-      let fuse_command = match docker_options.maybe_bind_mount.as_ref()  {
-        None => "".to_string(),
-        Some(mount) => format!(" --mount type=bind,source={},target={}",
-          &mount.local_filesystem,
-          &mount.container_filesystem),
-      };
-
-      command = format!("docker run --rm {} {} /bin/bash -c \"{}\"",
-        &fuse_command,
-        &docker_options.image_name,
-        command
-      )
+      command = docker_options.to_command_string(&command);
     }
 
     info!("Running command: {}", command);

crates/lib/subprocess_common/src/docker_options.rs

@@ -3,6 +3,7 @@
 pub struct DockerOptions {
   pub image_name: String,
   pub maybe_bind_mount: Option<DockerFilesystemMount>,
+  pub maybe_gpu: Option<DockerGpu>,
 }
 
 #[derive(Clone)]
@@ -11,6 +12,13 @@ pub struct DockerFilesystemMount {
   pub container_filesystem: String,
 }
 
+#[derive(Clone)]
+pub enum DockerGpu {
+  All,
+  Named(String),
+}
+
+
 impl DockerFilesystemMount {
   pub fn tmp_to_tmp() -> Self {
     Self {
@@ -25,6 +33,16 @@ impl DockerFilesystemMount {
   }
 }
 
+impl DockerGpu {
+  pub fn to_option_string(&self) -> String {
+    match self {
+      DockerGpu::All => " --gpus all ".to_string(),
+      // TODO: Handle spaces and special characters. Make sure this can't lead to injection.
+      DockerGpu::Named(named) => format!(" --gpus {}", named),
+    }
+  }
+}
+
 impl DockerOptions {
   pub fn to_command_string(&self, container_command: &str) -> String {
     let fuse_command = self.maybe_bind_mount
@@ -32,9 +50,15 @@ impl DockerOptions {
         .map(|mount| mount.to_fuse_option_string())
         .unwrap_or("".to_string());
 
+    let gpu_command = self.maybe_gpu
+        .as_ref()
+        .map(|gpu| gpu.to_option_string())
+        .unwrap_or("".to_string());
+
     // TODO: Handle spaces and special characters. Make sure this can't lead to injection.
-    format!("docker run --rm {} {} /bin/bash -c \"{}\"",
+    format!("docker run --rm {} {} {} /bin/bash -c \"{}\"",
             &fuse_command,
+            &gpu_command,
             &self.image_name,
             container_command)
   }
@@ -42,7 +66,7 @@ impl DockerOptions {
 
 #[cfg(test)]
 mod tests {
-  use crate::docker_options::{DockerFilesystemMount, DockerOptions};
+  use crate::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
 
   #[test]
   fn test_command() {
@@ -52,9 +76,10 @@ mod tests {
         local_filesystem: "/local".to_string(),
         container_filesystem: "/container".to_string(),
       }),
+      maybe_gpu: Some(DockerGpu::All),
     };
 
-    assert_eq!("docker run --rm  --mount type=bind,source=/local,target=/container  MY_IMAAGE /bin/bash -c \"echo wat\"",
+    assert_eq!("docker run --rm  --mount type=bind,source=/local,target=/container   --gpus all  MY_IMAAGE /bin/bash -c \"echo wat\"",
                command.to_command_string("echo wat"));
   }
 }

crates/service/download_job/src/job_types/hifigan/hifigan_model_check_command.rs

@@ -2,6 +2,7 @@ use crate::AnyhowResult;
 use log::info;
 use subprocess::{Popen, PopenConfig};
 use std::path::Path;
+use subprocess_common::docker_options::DockerOptions;
 
 /// This command is used to check hifigan for being a real model
 #[derive(Clone)]
@@ -13,6 +14,9 @@ pub struct HifiGanModelCheckCommand {
   virtual_env_activation_command: String,
   
   hifigan_model_check_script_name: String,
+
+  /// If this is run under Docker (eg. in development), these are the options.
+  maybe_docker_options: Option<DockerOptions>,
 }
 
 impl HifiGanModelCheckCommand {
@@ -20,11 +24,13 @@ impl HifiGanModelCheckCommand {
     hifigan_root_code_directory: &str,
     virtual_env_activation_command: &str,
     hifigan_model_check_script_name: &str,
+    maybe_docker_options: Option<DockerOptions>,
   ) -> Self {
     Self {
       hifigan_root_code_directory: hifigan_root_code_directory.to_string(),
       virtual_env_activation_command: virtual_env_activation_command.to_string(),
       hifigan_model_check_script_name: hifigan_model_check_script_name.to_string(),
+      maybe_docker_options,
     }
   }
 
@@ -46,6 +52,10 @@ impl HifiGanModelCheckCommand {
     command.push_str(" --output_metadata_filename ");
     command.push_str(&output_metadata_filename.as_ref().display().to_string());
 
+    if let Some(docker_options) = self.maybe_docker_options.as_ref() {
+      command = docker_options.to_command_string(&command);
+    }
+
     info!("Command: {:?}", command);
 
     let command_parts = [

crates/service/download_job/src/main.rs

@@ -41,7 +41,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 use storage_buckets_common::bucket_client::BucketClient;
 use storage_buckets_common::bucket_path_unifier::BucketPathUnifier;
-use subprocess_common::docker_options::{DockerFilesystemMount, DockerOptions};
+use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
 
 // Buckets
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -99,11 +99,12 @@ async fn main() -> AnyhowResult<()> {
     "./download_internet_file.py",
     "./python/bin/activate",
     DockerOptions {
-      image_name: "d73f28ce3ff6".to_string(),
+      image_name: "d73f28ce3ff6".to_string(), // web-downloader
       maybe_bind_mount: Some(DockerFilesystemMount {
         local_filesystem: "/tmp".to_string(),
         container_filesystem: "/tmp".to_string()
-      })
+      }),
+      maybe_gpu: None,
     }
   );
 
@@ -143,13 +144,21 @@ async fn main() -> AnyhowResult<()> {
   };
 
   let hifigan_model_check_command= HifiGanModelCheckCommand::new(
-    &easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?,
-    &easyenv::get_env_string_or_default(
-    "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND",
-    "source python-tacotron/bin/activate"),
-  &easyenv::get_env_string_or_default(
-    "HIFIGAN_MODEL_CHECK_SCRIPT_NAME",
-    "vocodes_model_check_hifigan.py"),
+    //&easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?,
+    "/models/tts",
+    //&easyenv::get_env_string_or_default(
+    //"HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND",
+    //"source python-tacotron/bin/activate"),
+    "source python/bin/activate",
+    //&easyenv::get_env_string_or_default(
+    //"HIFIGAN_MODEL_CHECK_SCRIPT_NAME",
+    //"vocodes_model_check_hifigan.py"),
+    "./vocodes_model_check_hifigan.py",
+    Some(DockerOptions {
+      image_name: "5642d0fd7fc1".to_string(), // storyteller-ml
+      maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()),
+      maybe_gpu: Some(DockerGpu::All),
+    }),
   );
 
   let job_state = JobState {