Kill inference jobs (mod endpoint)

crates/lib/mysql_queries/src/queries/generic_inference/web/kill_generic_inference_jobs.rs

@@ -0,0 +1,157 @@
+use std::collections::HashSet;
+
+use anyhow::anyhow;
+use sqlx::{MySql, MySqlPool, QueryBuilder};
+
+use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
+use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
+use enums::common::job_status::JobStatus;
+use errors::AnyhowResult;
+
+/// Only certain job statuses should be modified.
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
+pub enum KillableStatus {
+  Failed,
+  Pending,
+  Started,
+}
+
+/// Target everything, by job category, or by model type.
+#[derive(Clone)]
+pub enum KillableTarget {
+  AllJobs,
+  Category(InferenceCategory),
+  ModelType(InferenceModelType),
+}
+
+pub struct KillGenericInferenceJobsArgs<'a> {
+  pub job_statuses: HashSet<KillableStatus>,
+  pub target: KillableTarget,
+  pub maybe_priority_or_lower: Option<u8>,
+  pub mysql_pool: &'a MySqlPool,
+}
+
+pub async fn kill_generic_inference_jobs(args: KillGenericInferenceJobsArgs<'_>) -> AnyhowResult<()> {
+  let mut query = query_builder(
+    args.job_statuses,
+    args.target,
+    args.maybe_priority_or_lower
+  );
+
+  let query = query.build();
+
+  let result = query.execute(args.mysql_pool).await;
+
+  match result {
+    Err(err) => {
+      Err(anyhow!("error with query: {:?}", err))
+    },
+    Ok(_r) => Ok(()),
+  }
+}
+
+fn query_builder(
+  job_statuses: HashSet<KillableStatus>,
+  target: KillableTarget,
+  maybe_priority_or_lower: Option<u8>,
+) -> QueryBuilder<'static, MySql> {
+
+  // NB: Query cannot be statically checked by sqlx
+  let mut query_builder: QueryBuilder<MySql> = QueryBuilder::new(
+    r#"
+UPDATE generic_inference_jobs
+SET status='cancelled_by_system'
+WHERE status IN(
+    "#
+  );
+
+  let job_statuses = job_statuses.into_iter()
+      .map(|status| match status {
+        KillableStatus::Failed => JobStatus::AttemptFailed.to_str(),
+        KillableStatus::Pending => JobStatus::Pending.to_str(),
+        KillableStatus::Started => JobStatus::Started.to_str(),
+      })
+      .collect::<Vec<_>>();
+
+  // NB: Syntax will be wrong if list has zero length
+  let mut separated = query_builder.separated(", ");
+  for value_type in job_statuses.into_iter() {
+    separated.push_bind(value_type);
+  }
+  separated.push_unseparated(") ");
+
+  if let Some(priority) = maybe_priority_or_lower {
+    query_builder.push(format!(" AND priority <= {priority} "));
+  }
+
+  match target {
+    KillableTarget::AllJobs => {
+      // No need to modify query.
+    }
+    KillableTarget::Category(category) => {
+      query_builder.push(" AND inference_category = ");
+      query_builder.push_bind(category);
+    }
+    KillableTarget::ModelType(model_type) => {
+      query_builder.push(" AND maybe_model_type = ");
+      query_builder.push_bind(model_type);
+    }
+  }
+
+  query_builder
+}
+
+#[cfg(test)]
+mod tests {
+  use std::collections::HashSet;
+
+  use regex::Regex;
+
+  use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
+  use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
+
+  use crate::queries::generic_inference::web::kill_generic_inference_jobs::{KillableStatus, KillableTarget, query_builder};
+
+  fn cleanup_query(query: &str) -> String {
+    let regex = Regex::new(r#"\s+"#).unwrap();
+    let result = regex.replace_all(query, " ").trim().to_string();
+    result
+  }
+
+  #[test]
+  fn test_statuses() {
+    let statuses = HashSet::from([KillableStatus::Failed]);
+    let builder = query_builder(statuses, KillableTarget::AllJobs, None);
+    let query = cleanup_query(&builder.into_sql());
+    assert_eq!(query, "UPDATE generic_inference_jobs SET status='cancelled_by_system' WHERE status IN( ?)".to_string());
+
+    let statuses = HashSet::from([KillableStatus::Failed, KillableStatus::Pending, KillableStatus::Started]);
+    let builder = query_builder(statuses, KillableTarget::AllJobs, None);
+    let query = cleanup_query(&builder.into_sql());
+    assert_eq!(query, "UPDATE generic_inference_jobs SET status='cancelled_by_system' WHERE status IN( ?, ?, ?)".to_string());
+  }
+
+  #[test]
+  fn test_priority() {
+    let statuses = HashSet::from([KillableStatus::Failed]);
+    let builder = query_builder(statuses, KillableTarget::AllJobs, Some(10));
+    let query = cleanup_query(&builder.into_sql());
+    assert_eq!(query, "UPDATE generic_inference_jobs SET status='cancelled_by_system' WHERE status IN( ?) AND priority <= 10".to_string());
+  }
+
+  #[test]
+  fn test_category_target() {
+    let statuses = HashSet::from([KillableStatus::Failed]);
+    let builder = query_builder(statuses, KillableTarget::Category(InferenceCategory::LipsyncAnimation), None);
+    let query = cleanup_query(&builder.into_sql());
+    assert_eq!(query, "UPDATE generic_inference_jobs SET status='cancelled_by_system' WHERE status IN( ?) AND inference_category = ?".to_string());
+  }
+
+  #[test]
+  fn test_model_target() {
+    let statuses = HashSet::from([KillableStatus::Failed]);
+    let builder = query_builder(statuses, KillableTarget::ModelType(InferenceModelType::RvcV2), None);
+    let query = cleanup_query(&builder.into_sql());
+    assert_eq!(query, "UPDATE generic_inference_jobs SET status='cancelled_by_system' WHERE status IN( ?) AND maybe_model_type = ?".to_string());
+  }
+}

crates/lib/mysql_queries/src/queries/generic_inference/web/mod.rs

@@ -1,3 +1,4 @@
 pub mod get_inference_job_status;
 pub mod get_pending_inference_job_count;
 pub mod insert_generic_inference_job;
+pub mod kill_generic_inference_jobs;

crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/kill_inference_jobs.rs

@@ -0,0 +1,144 @@
+use std::collections::HashSet;
+use std::fmt;
+use std::sync::Arc;
+
+use actix_web::{HttpRequest, HttpResponse, web};
+use actix_web::error::ResponseError;
+use actix_web::http::StatusCode;
+use log::warn;
+
+use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
+use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
+use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
+use mysql_queries::queries::generic_inference::web::kill_generic_inference_jobs;
+use mysql_queries::queries::generic_inference::web::kill_generic_inference_jobs::{kill_generic_inference_jobs, KillGenericInferenceJobsArgs};
+
+use crate::server_state::ServerState;
+
+/// Only certain job statuses should be modified.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum KillableStatus {
+  Failed,
+  Pending,
+  Started,
+}
+
+/// Target everything, by job category, or by model type.
+#[derive(Clone, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum KillableTarget {
+  AllJobs,
+  Category(InferenceCategory),
+  ModelType(InferenceModelType),
+}
+
+// NB: ONLY MODERATORS CAN EDIT CATEGORIES.
+// These are not sparse updates!
+#[derive(Deserialize)]
+pub struct KillInferenceJobsRequest {
+  pub job_statuses: HashSet<KillableStatus>,
+  pub target: KillableTarget,
+  pub maybe_priority_or_lower: Option<u8>,
+}
+
+#[derive(Serialize)]
+pub struct KillInferenceJobsResponse {
+  pub success: bool,
+}
+
+#[derive(Debug, Serialize)]
+pub enum KillInferenceJobsError {
+  BadInput(String),
+  NotFound,
+  NotAuthorized,
+  ServerError,
+}
+
+impl ResponseError for KillInferenceJobsError {
+  fn status_code(&self) -> StatusCode {
+    match *self {
+      KillInferenceJobsError::BadInput(_) => StatusCode::BAD_REQUEST,
+      KillInferenceJobsError::NotFound => StatusCode::NOT_FOUND,
+      KillInferenceJobsError::NotAuthorized => StatusCode::UNAUTHORIZED,
+      KillInferenceJobsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
+    }
+  }
+
+  fn error_response(&self) -> HttpResponse {
+    serialize_as_json_error(self)
+  }
+}
+
+// NB: Not using derive_more::Display since Clion doesn't understand it.
+impl fmt::Display for KillInferenceJobsError {
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    write!(f, "{:?}", self)
+  }
+}
+
+// =============== Handler ===============
+
+pub async fn kill_generic_inference_jobs_handler(
+  http_request: HttpRequest,
+  request: web::Json<KillInferenceJobsRequest>,
+  server_state: web::Data<Arc<ServerState>>) -> Result<HttpResponse, KillInferenceJobsError>
+{
+  let maybe_user_session = server_state
+      .session_checker
+      .maybe_get_user_session(&http_request, &server_state.mysql_pool)
+      .await
+      .map_err(|e| {
+        warn!("Session checker error: {:?}", e);
+        KillInferenceJobsError::ServerError
+      })?;
+
+  let user_session = match maybe_user_session {
+    Some(session) => session,
+    None => {
+      warn!("not logged in");
+      return Err(KillInferenceJobsError::NotAuthorized);
+    }
+  };
+
+  // TODO: We don't have a permission for this, so use this as a proxy permission
+  if !user_session.can_ban_users {
+    warn!("no permission to edit categories");
+    return Err(KillInferenceJobsError::NotAuthorized);
+  }
+
+  let job_statuses = request.job_statuses.iter()
+      .map(|status| match status {
+        KillableStatus::Failed => kill_generic_inference_jobs::KillableStatus::Failed,
+        KillableStatus::Pending => kill_generic_inference_jobs::KillableStatus::Pending,
+        KillableStatus::Started => kill_generic_inference_jobs::KillableStatus::Started,
+      })
+      .collect::<HashSet<_>>();
+
+  let target = match request.target {
+    KillableTarget::AllJobs => kill_generic_inference_jobs::KillableTarget::AllJobs,
+    KillableTarget::Category(category) => kill_generic_inference_jobs::KillableTarget::Category(category),
+    KillableTarget::ModelType(model_type) => kill_generic_inference_jobs::KillableTarget::ModelType(model_type),
+  };
+
+  kill_generic_inference_jobs(KillGenericInferenceJobsArgs {
+    job_statuses,
+    target,
+    maybe_priority_or_lower: request.maybe_priority_or_lower,
+    mysql_pool: &server_state.mysql_pool,
+  }).await.map_err(|err| {
+    warn!("Query Error: {:?}", err);
+    KillInferenceJobsError::ServerError
+  })?;
+
+  let response = KillInferenceJobsResponse {
+    success: true,
+  };
+
+  let body = serde_json::to_string(&response)
+      .map_err(|e| KillInferenceJobsError::ServerError)?;
+
+  Ok(HttpResponse::Ok()
+      .content_type("application/json")
+      .body(body))
+}

crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/mod.rs

@@ -1,3 +1,4 @@
 pub mod get_inference_job_status;
 pub mod get_pending_inference_job_count;
+pub mod kill_inference_jobs;
 pub mod terminate_inference_job_handler;

crates/service/web/storyteller_web/src/routes.rs

@@ -37,6 +37,7 @@ use crate::http_server::endpoints::flags::design_refresh_flag::disable_design_re
 use crate::http_server::endpoints::flags::design_refresh_flag::enable_design_refresh_flag_handler::enable_design_refresh_flag_handler;
 use crate::http_server::endpoints::inference_job::get_inference_job_status::get_inference_job_status_handler;
 use crate::http_server::endpoints::inference_job::get_pending_inference_job_count::get_pending_inference_job_count_handler;
+use crate::http_server::endpoints::inference_job::kill_inference_jobs::kill_generic_inference_jobs_handler;
 use crate::http_server::endpoints::inference_job::terminate_inference_job_handler::terminate_inference_job_handler;
 use crate::http_server::endpoints::investor_demo::disable_demo_mode_handler::disable_demo_mode_handler;
 use crate::http_server::endpoints::investor_demo::enable_demo_mode_handler::enable_demo_mode_handler;
@@ -355,6 +356,11 @@ fn add_moderator_routes<T, B> (app: App<T>) -> App<T>
                   .route(web::get().to(get_w2l_inference_queue_count_handler))
                   .route(web::head().to(|| HttpResponse::Ok()))
             )
+            .service(
+              web::resource("/kill_generic")
+                  .route(web::post().to(kill_generic_inference_jobs_handler))
+                  .route(web::head().to(|| HttpResponse::Ok()))
+            )
       )
       .service(
         web::scope("/pending")