diff --git a/crates/lib/google_drive_common/src/google_drive_download_command.rs b/crates/lib/google_drive_common/src/google_drive_download_command.rs index 486d988a1c..f95b59debc 100644 --- a/crates/lib/google_drive_common/src/google_drive_download_command.rs +++ b/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, + + /// If this is run under Docker (eg. in development), these are the options. maybe_docker_options: Option, } @@ -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); diff --git a/crates/lib/subprocess_common/src/docker_options.rs b/crates/lib/subprocess_common/src/docker_options.rs index b28d5bcd5a..012fb8706c 100644 --- a/crates/lib/subprocess_common/src/docker_options.rs +++ b/crates/lib/subprocess_common/src/docker_options.rs @@ -3,6 +3,7 @@ pub struct DockerOptions { pub image_name: String, pub maybe_bind_mount: Option, + pub maybe_gpu: Option, } #[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")); } } diff --git a/crates/service/download_job/src/job_types/hifigan/hifigan_model_check_command.rs b/crates/service/download_job/src/job_types/hifigan/hifigan_model_check_command.rs index 406b4e87ca..73a2dea8e0 100644 --- a/crates/service/download_job/src/job_types/hifigan/hifigan_model_check_command.rs +++ b/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, } 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, ) -> 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 = [ diff --git a/crates/service/download_job/src/main.rs b/crates/service/download_job/src/main.rs index d179eb7743..eef76d314a 100644 --- a/crates/service/download_job/src/main.rs +++ b/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 {