From ae899d8aca1bc19b24fa365877ef5e68262f1362 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Fri, 6 Oct 2023 15:21:38 -0400 Subject: [PATCH] Fixed most of the job and implemented it. Need to run the container this might be fun haha ... --- .../config/inference-job.development.env | 2 +- .../tts/vall_e_x/model_downloaders.rs | 8 +- .../job/job_types/tts/vall_e_x/process_job.rs | 13 +- .../vall_e_x/vall_e_x_inference_command.rs | 385 +++++++++++++++--- .../job/inference_job/src/job_dependencies.rs | 2 + crates/service/job/inference_job/src/main.rs | 4 +- 6 files changed, 352 insertions(+), 62 deletions(-) diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 05f1227d14..1b14ed66b9 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -20,7 +20,7 @@ REDIS_FOR_KEEPALIVE_URL="redis://localhost/0" VALL_E_X_INFERENCE_ROOT_DIRECTORY = "/home/tensor/code/Storyteller/storyteller-ml/tts/VALL-E-X" VALL_E_X_INFERENCE_COMMAND = "python3 main.py" VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND = "source .env/bin/activate" - +VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE = "" # fill this out for dev? # ================================= # # # # RVC (v2) # diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs index 1504b64206..28887ccce8 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/model_downloaders.rs @@ -1,10 +1,10 @@ use crate::util::model_downloader::ModelDownloader; pub struct VallEXDownloaders { - pub vall_e_x: String, - pub whisper_lg: String, - pub whisper_md: String, - pub vocos_encodec_24khz: String + pub vall_e_x: VallEXModelMapping, + pub whisper_lg: WhisperLargeModelMapping, + pub whisper_md: WhisperMediumModelMapping, + pub vocos_encodec_24khz: VocosEncodecModelMapping } crate::impl_model_downloader!( diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs index aa455d0e15..3230a6efe2 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs @@ -26,12 +26,14 @@ use crate::job::job_types::tts::vall_e_x::validate_job::validate_job; use crate::job_dependencies::JobDependencies; + +// This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked pub struct VALLEXProcessJobArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, } - pub async fn process_job(args: SadTalkerProcessJobArgs<'_>) -> Result { + pub async fn process_job(args: VALLEXProcessJobArgs<'_>) -> Result { // get job args and dependencies and execute them with the inputs let job = args.job; @@ -84,4 +86,13 @@ pub struct VALLEXProcessJobArgs<'a> { + } + + + #[cfg(test)] + mod test { + #[test] + fn values() { + print!("{}","Hello!"); + } } \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 38e272a3e5..7eb2f12737 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -24,69 +24,346 @@ pub enum ExecutableOrCommand { Command(String), } +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS : Lazy> = Lazy::new(|| { + let env_var_names= [ + "MYSQL_URL", + "ACCESS_KEY", + "SECRET_KEY", + "NEWRELIC_API_KEY", + ]; + env_var_names.iter() + .map(|value| value.to_string()) + .collect::>() +}); -pub struct InferenceArgs<'s, P: AsRef> { +// The inference command assumes everything has been setup to run inference +// so we have all the files downloaded for the inputs. +pub struct InferenceArgs> { /// --driven_audio: path to the input audio - pub input_audio: P, // Change to files - - /// --source_image: path to the input image (or video) - pub input_embedding: P, - pub input_text: P, - /// --result_dir: path to directory work is performed - pub work_dir: P, + pub input_embedding: P, // name of the embedding.npz in the prompts folder? + pub input_text: P, // name of the text /// --result_file: path to final file output - pub output_file: P, - pub stderr_output_file: P, + pub output_file_name: P, // output file name in the output folder + pub stderr_output_file: P } +#[derive(Clone)] pub struct VallEXInferenceCommand { - + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option } -// pub impl VallEXInferenceCommand { -// pub fn new>( -// root_code_directory: P, -// executable_or_command: ExecutableOrCommand, -// maybe_virtual_env_activation_command: Option<&str>, -// maybe_default_config_path: Option

, -// maybe_docker_options: Option, -// maybe_execution_timeout: Option, -// alternate_checkpoint_dir: Option -// ) -> Self { -// Self { -// root_code_directory: root_code_directory.as_ref().to_path_buf(), -// executable_or_command, -// maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), -// maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), -// maybe_docker_options, -// maybe_execution_timeout, -// alternate_checkpoint_dir, -// } -// } -// } +impl VallEXInferenceCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let root_code_directory = easyenv::get_env_pathbuf_required( + "VALL_E_X_INFERENCE_ROOT_DIRECTORY")?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_COMMAND"); + + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "VALL_E_X_INFERENCE_EXECUTABLE"); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => return Err(anyhow!("neither command nor executable passed")), + }, + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "VALL_E_X_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); + + let maybe_execution_timeout = + easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + + // Probably for local + let maybe_docker_options = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "SAD_TALKER_ALTERNATE_CHECKPOINT_PATH"); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>( + &self, + args: InferenceArgs

, + ) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

, + ) -> AnyhowResult { + + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" -text "); + command.push_str(&path_to_string(args.input_text)); + + command.push_str(" -use_prompt_name "); + command.push_str(&path_to_string(args.input_embedding)); + + command.push_str(" -output_file_name "); + command.push_str(&path_to_string(args.output_file_name)); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push(( + OsString::from(env_key), + OsString::from(env_value), + )); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!("Subprocess didn't end after timeout: {:?}; terminating...", &timeout); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} + +#[derive(Clone)] pub struct VallEXCreateEmbeddingCommand { - + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option } -// pub impl VallEXCreateEmbeddingCommand { -// pub fn new>( -// root_code_directory: P, -// executable_or_command: ExecutableOrCommand, -// maybe_virtual_env_activation_command: Option<&str>, -// maybe_default_config_path: Option

, -// maybe_docker_options: Option, -// maybe_execution_timeout: Option, -// alternate_checkpoint_dir: Option -// ) -> Self { -// Self { -// root_code_directory: root_code_directory.as_ref().to_path_buf(), -// executable_or_command, -// maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), -// maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), -// maybe_docker_options, -// maybe_execution_timeout, -// alternate_checkpoint_dir, -// } -// } -// } \ No newline at end of file +impl VallEXCreateEmbeddingCommand { + + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + // helper funciton to set up env + pub fn from_env() -> AnyhowResult { + // TODO: fix for now just for compilation + let root_code_directory = easyenv::get_env_pathbuf_required( + "VALL_E_X_INFERENCE_ROOT_DIRECTORY")?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_COMMAND"); + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "VALL_E_X_INFERENCE_EXECUTABLE"); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => return Err(anyhow!("neither command nor executable passed")), + }, + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "VALL_E_X_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); + + let maybe_execution_timeout = + easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + + // Probably for local + let maybe_docker_options = easyenv::get_env_string_optional( + "VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "SAD_TALKER_ALTERNATE_CHECKPOINT_PATH"); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + +} diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index 9f87447404..47f4320940 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/crates/service/job/inference_job/src/job_dependencies.rs @@ -22,6 +22,7 @@ use newrelic_telemetry::Client as NewRelicClient; use crate::job::job_types::lipsync::sad_talker::model_downloaders::SadTalkerDownloaders; use crate::job::job_types::lipsync::sad_talker::sad_talker_inference_command::SadTalkerInferenceCommand; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::Tacotron2InferenceCommand; +use crate::job::job_types::tts::vall_e_x::model_downloaders::VallEXDownloaders; use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand; use crate::job::job_types::tts::vall_e_x::vall_e_x_inference_command::VallEXInferenceCommand; @@ -205,6 +206,7 @@ pub struct SadTalkerDetails { // TODO: we will probably want a command type of some kind that implements // some kind of interface that we can just pass strings to. pub struct VallEXDetails { + pub downloaders: VallEXDownloaders, pub inference_command: VallEXInferenceCommand, pub create_embedding_command: VallEXCreateEmbeddingCommand } diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index f25760e7f0..676145f5ad 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -316,8 +316,8 @@ async fn main() -> AnyhowResult<()> { }, vall_e_x: VallEXDetails { downloaders: VallEXDownloaders::build_all_from_env(), - inference_command: VallEXInferenceCommand, - create_embedding_command: VallEXCreateEmbeddingCommand, + inference_command: VallEXInferenceCommand::from_env()?, + create_embedding_command: VallEXCreateEmbeddingCommand::from_env()?, }, vits: VitsDetails { inference_command: vits_inference_command()?,