auto-kill jobs in dev

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<MySql>,
+) -> 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(())
+}

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;

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());
+  }
 }

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"
 
 # ================================= #
 #                                   #

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(" ");
 

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
       

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<u64>,
 
+  /// Kill old jobs (only in development)
+  pub debug_kill_old_jobs: Option<bool>,
+
   // TODO
   pub output_width: Option<u64>,
   pub output_height: Option<u64>,
@@ -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()),