From 3cd04414f8e4ed8de0a02278784f18332051992e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 5 Feb 2025 19:37:02 -0500 Subject: [PATCH] auto-kill jobs in dev --- .../web/kill_jobs_in_development.rs | 48 +++++++++++++++++++ .../src/queries/generic_inference/web/mod.rs | 1 + crates/lib/server_environment/src/lib.rs | 14 ++++++ .../config/inference-job.development.env | 15 +++++- .../animate_x_process_frames_command.rs | 2 +- .../process_single_studio_gen2_job.rs | 5 +- .../enqueue_studio_gen2_handler.rs | 14 ++++++ 7 files changed, 94 insertions(+), 5 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/generic_inference/web/kill_jobs_in_development.rs diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/kill_jobs_in_development.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/kill_jobs_in_development.rs new file mode 100644 index 0000000000..b75a29337b --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/kill_jobs_in_development.rs @@ -0,0 +1,48 @@ +use anyhow::anyhow; +use sqlx; +use sqlx::pool::PoolConnection; +use sqlx::{MySql, MySqlPool}; + +use errors::AnyhowResult; + +/// Kill a handful of jobs when called. +/// This should only be used in development to keep the job queue empty. +pub async fn kill_jobs_in_development(mysql_pool: &MySqlPool) -> AnyhowResult<()> +{ + let mut connection = mysql_pool.acquire().await?; + kill_jobs_in_development_using_connection(&mut connection).await +} + +/// Kill a handful of jobs when called. +/// This should only be used in development to keep the job queue empty. +pub async fn kill_jobs_in_development_using_connection( + mysql_connection: &mut PoolConnection, +) -> AnyhowResult<()> { + // We cancel any non-complete / non-terminal jobs + // We also mark "started" jobs in the event that they're left in a bad state. + // Any started jobs that are actually running should still run to completion, and + // the user should still see the final result even if the job system UI doesn't + // indicate it. + let query_result = sqlx::query!( + r#" +UPDATE generic_inference_jobs +SET + status = 'cancelled_by_user', + retry_at = NULL +WHERE status IN ( + 'pending', + 'started', + 'attempt_failed' + ) +LIMIT 100 + "# + ) + .execute(&mut **mysql_connection) + .await; + + if let Err(err) = query_result { + return Err(anyhow!("error with terminate job query (1): {:?}", err)); + } + + Ok(()) +} diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/mod.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/mod.rs index 4daa07efdd..37b2b2d2bf 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/mod.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/mod.rs @@ -5,5 +5,6 @@ pub mod get_pending_inference_job_count; pub mod insert_generic_inference_job; pub mod job_status; pub mod kill_generic_inference_jobs; +pub mod kill_jobs_in_development; pub mod list_session_jobs; pub mod mark_generic_inference_job_cancelled_by_user; \ No newline at end of file diff --git a/crates/lib/server_environment/src/lib.rs b/crates/lib/server_environment/src/lib.rs index e444af7c9e..7d3b94a0da 100644 --- a/crates/lib/server_environment/src/lib.rs +++ b/crates/lib/server_environment/src/lib.rs @@ -52,6 +52,14 @@ impl ServerEnvironment { ServerEnvironment::Production => true, } } + + /// Reports whether the software is running in development. + pub fn is_development(&self) -> bool { + match self { + ServerEnvironment::Development => true, + ServerEnvironment::Production => false, + } + } } #[cfg(test)] @@ -87,4 +95,10 @@ mod tests { assert_eq!(true, ServerEnvironment::Production.is_deployed_in_production()); assert_eq!(false, ServerEnvironment::Development.is_deployed_in_production()); } + + #[test] + fn test_is_development() { + assert_eq!(false, ServerEnvironment::Production.is_development()); + assert_eq!(true, ServerEnvironment::Development.is_development()); + } } 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 1036c0b062..9dda597840 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -175,13 +175,24 @@ COMFY_MAPPINGS_DIRECTORY="/home/bt/dev/storyteller/V2V-StylesList/V2VWorkflows/m # # # ================================= # +# Feature Flags +STABLE_ANIMATOR_ENABLED="false" +ANIMATE_X_ENABLED="true" + +# Animate-X +ANIMATE_X_MODEL_DIRECTORY_PATH="checkpoints/" +ANIMATE_X_ROOT_CODE_DIRECTORY="/home/bt/dev/storyteller/storyteller-ml/animation/animate-x" +ANIMATE_X_VENV_ACTIVATION_COMMAND="source venv/bin/activate" + +# Stable Animator +STABLE_ANIMATOR_ROOT_CODE_DIRECTORY="/home/bt/dev/storyteller/storyteller-ml/animation/StableAnimator" +STABLE_ANIMATOR_VENV_ACTIVATION_COMMAND="source venv/bin/activate" + STABLE_ANIMATOR_PRETRAINED_MODEL_PATH="checkpoints/stable-video-diffusion-img2vid-xt" STABLE_ANIMATOR_POSENET_MODEL_PATH="checkpoints/Animation/pose_net.pth" STABLE_ANIMATOR_FACE_ENCODER_MODEL_PATH="checkpoints/Animation/face_encoder.pth" STABLE_ANIMATOR_UNET_MODEL_PATH="checkpoints/Animation/unet.pth" -STABLE_ANIMATOR_ROOT_CODE_DIRECTORY="/home/bt/dev/storyteller/storyteller-ml/animation/StableAnimator" -STABLE_ANIMATOR_VENV_ACTIVATION_COMMAND="source venv/bin/activate" # ================================= # # # diff --git a/crates/service/job/inference_job/src/job/job_types/studio_gen2/animate_x/animate_x_process_frames_command.rs b/crates/service/job/inference_job/src/job/job_types/studio_gen2/animate_x/animate_x_process_frames_command.rs index f2125aa51a..be7bb068e0 100644 --- a/crates/service/job/inference_job/src/job/job_types/studio_gen2/animate_x/animate_x_process_frames_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/studio_gen2/animate_x/animate_x_process_frames_command.rs @@ -98,7 +98,7 @@ impl AnimateXProcessFramesCommand { // TODO: Build command // NB: arg is "paths" not "path" - command.push_str(" --source_video_paths"); + command.push_str(" --source_video_paths "); command.push_str(&path_to_string(&args.source_video_path)); command.push_str(" "); diff --git a/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs b/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs index dae5c954c8..2061e32560 100644 --- a/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs @@ -250,7 +250,7 @@ pub async fn process_single_studio_gen2_job( return Err(ProcessSingleJobError::Other(anyhow!("Studio job type not set"))); } StudioModelPipeline::AnimateX(deps) => { - info!("Running Studio Gen2 inference (Animate-X)..."); + info!("Running Studio Gen2 pose frame generation (Animate-X)..."); let pose_pkl_file= work_paths.output_dir.path().join("pose_data.pkl"); @@ -273,7 +273,8 @@ pub async fn process_single_studio_gen2_job( saved_pose_frames_dir: &pose_frames_dir, saved_original_frames_dir: &original_frames_dir, }).await; - + + info!("Running Studio Gen2 inference (Animate-X)..."); // TODO diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs index 1025fcc3c3..d6fcc81645 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/studio_gen2/enqueue_studio_gen2_handler.rs @@ -24,6 +24,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{Ge use mysql_queries::payloads::generic_inference_args::inner_payloads::studio_gen2_payload::StudioGen2Payload; use mysql_queries::payloads::generic_inference_args::inner_payloads::workflow_payload::{WorkflowArgs, WorkflowType}; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; +use mysql_queries::queries::generic_inference::web::kill_jobs_in_development::kill_jobs_in_development; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use primitives::lazy_any_option_true::lazy_any_option_true; use primitives::str_to_bool::str_to_bool; @@ -66,6 +67,9 @@ pub struct EnqueueStudioGen2Request { /// Sleep for debugging pub debug_sleep_millis: Option, + /// Kill old jobs (only in development) + pub debug_kill_old_jobs: Option, + // TODO pub output_width: Option, pub output_height: Option, @@ -252,6 +256,16 @@ pub async fn enqueue_studio_gen2_handler( watermark_type = Some(WatermarkType::Storyteller); // Studio is only on Storyteller. } + if request.debug_kill_old_jobs.unwrap_or(false) + && server_state.server_environment.is_development() + { + info!("Killing old jobs before inserting new job (DEVELOPMENT MODE ONLY!)"); + let result = kill_jobs_in_development(&server_state.mysql_pool).await; + if let Err(err) = result { + warn!("Error killing old jobs: {:?}", err); + } + } + let inference_args = StudioGen2Payload { image_file: Some(request.image_file.clone()), video_file: Some(request.video_file.clone()),