diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 69ce218530..bacb76e971 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -86,7 +86,7 @@ mime = "0.3.17" rand = { version = "0.8.5", features = [] } jsonpath_lib = "0.3.0" walkdir = "2.4.0" -reqwest = "0.11.22" +reqwest = { version = "0.11.22", features = ["blocking"] } opentelemetry = {version = "0.21.0", features = ["metrics"]} opentelemetry-otlp = { version = "0.14.0", features = ["metrics"] } opentelemetry_sdk = { version = "0.21.2", features = ["metrics", "rt-tokio"] } 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 9e1d5e5807..88dac0d340 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -147,11 +147,13 @@ MOCAPNET_INFERENCE_EXECUTABLE_OR_COMMAND="pythonVirtualEnvironment/bin/python -m # # # ================================= # COMFY_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/" -COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunner.py" +COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunnerForSidecar.py" +COMFY_LAUNCH_COMMAND="venv/bin/python /home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/main.py" # directory to where files should be downloaded to (default to just the comfy directory) COMFY_MOUNTS_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/" COMFY_VIDEO_PROCESSING_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/format_video.py" +COMFY_SETUP_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/install.py" # ================================= # # # @@ -186,3 +188,4 @@ HF_DATASETS_CACHE="/tmp/huggingface" # nltk also caches data. NLTK_DATA="/tmp/nltk_data" +SKIP_TELEMETRY="true" \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs index 5e80cd9c99..f2f527a2b7 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs @@ -1,6 +1,10 @@ use std::path::PathBuf; - +use tokio::process::Command; +use std::time::{Duration, Instant}; +use actix_web::http::StatusCode; use errors::AnyhowResult; +use reqwest::blocking::Client; +use tokio::time::sleep; use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand; use crate::util::common_commands::ffmpeg_logo_watermark_command::FfmpegLogoWatermarkCommand; @@ -12,9 +16,56 @@ pub struct ComfyDependencies { } impl ComfyDependencies { - pub fn setup() -> AnyhowResult { + pub async fn setup() -> AnyhowResult { + let inference_command = ComfyInferenceCommand::from_env()?; + let comfy_setup_script = inference_command.clone().comfy_setup_script; + let output = Command::new("python3") + .arg(&comfy_setup_script) + .output() + .await.expect("failed to run comfyui setup script"); + println!("stdout: {}", String::from_utf8_lossy(&output.stdout)); + + // Start the comfy start script as a daemon process + let launch_command = format!( + "cd {} && {}", + inference_command.clone().comfy_root_code_directory.display(), + inference_command.clone().comfy_launch_command.display() + ); + + let _ = Command::new("sh") + .arg("-c") + .arg(&launch_command) + .stdout(std::process::Stdio::inherit()) // Inherit the stdout to stream the output + .spawn() + .expect("failed to start comfy start script"); + + let mut success = false; + + let client = reqwest::Client::new(); + let start = Instant::now(); + let timeout = Duration::from_secs(60); + while Instant::now().duration_since(start) < timeout { + let response = client.get("http://127.0.0.1:8188/prompt").send().await; + match response { + Ok(r) if r.status() == StatusCode::OK => { + println!("Received 200 response from http://127.0.0.1:8188/prompt"); + success = true; + break; + }, + _ => { + println!("Did not receive 200 response, retrying..."); + sleep(Duration::from_secs(5)).await; + } + } + } + + if !success { + println!("Timeout reached without receiving a 200 response."); + panic!("Comfy start failed"); + } + Ok(Self { - inference_command: ComfyInferenceCommand::from_env()?, + inference_command, dependency_tokens: RequiredModels::init(), ffmpeg_watermark_command: FfmpegLogoWatermarkCommand::from_env()?, }) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index b5f187a02e..a0ec7c627f 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs @@ -44,6 +44,10 @@ pub struct ComfyInferenceCommand { // Video processing script pub(crate) processing_script: PathBuf, + pub(crate) comfy_setup_script: PathBuf, + + pub(crate) comfy_launch_command: PathBuf, + /// Config file to use config_path: Option, @@ -95,6 +99,12 @@ impl ComfyInferenceCommand { } }; + let comfy_setup_script = easyenv::get_env_pathbuf_required( + "COMFY_SETUP_SCRIPT")?; + + let comfy_launch_command = easyenv::get_env_pathbuf_required( + "COMFY_LAUNCH_COMMAND")?; + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( "COMFY_INFERENCE_MAYBE_VENV_COMMAND"); @@ -124,6 +134,8 @@ impl ComfyInferenceCommand { config_path, mounts_directory, processing_script, + comfy_setup_script, + comfy_launch_command, maybe_virtual_env_activation_command, maybe_docker_options, maybe_execution_timeout, diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index 758110d284..7d33b29665 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -34,7 +34,7 @@ pub struct JobSpecificDependencies { impl JobSpecificDependencies { - pub fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult { + pub async fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult { let mut maybe_rvc_v2_dependencies = None; let mut maybe_sad_talker_dependencies = None; let mut maybe_svc_dependencies = None; @@ -101,7 +101,7 @@ impl JobSpecificDependencies { if scoped_execution.can_run_job(InferenceModelType::ComfyUi) { print_with_space("Setting ComfyUI dependencies..."); - maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup()?); + maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup().await?); } if scoped_execution.can_run_job(InferenceModelType::ConvertFbxToGltf) { diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index f418d7f387..72d091154b 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -237,7 +237,7 @@ async fn main() -> AnyhowResult<()> { let meter = opentelemetry::global::meter("inference-job"); - let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution)?; + let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution).await?; let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory( easyenv::get_env_pathbuf_or_default( diff --git a/crates/service/job/inference_job/src/util/instrumentation.rs b/crates/service/job/inference_job/src/util/instrumentation.rs index d269a4d696..ee91e1b16a 100644 --- a/crates/service/job/inference_job/src/util/instrumentation.rs +++ b/crates/service/job/inference_job/src/util/instrumentation.rs @@ -21,6 +21,11 @@ pub fn init_otel_metrics_pipeline( ) -> Result<(), opentelemetry::metrics::MetricsError> { info!("Setting up otel metrics pipeline..."); + // check env var to see if we should skip telemetry + if std::env::var("SKIP_TELEMETRY").is_ok() { + warn!("Skipping telemetry setup due to SKIP_TELEMETRY env var"); + return Ok(()); + } let provider = opentelemetry_otlp::new_pipeline() .metrics(opentelemetry_sdk::runtime::Tokio) // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset?