From 4bf1a25fbb4ec89c98be5f6d8c409ff1bd7a4f31 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 2 Dec 2023 23:50:13 -0500 Subject: [PATCH] Kill inference jobs (mod endpoint) --- .../web/kill_generic_inference_jobs.rs | 157 ++++++++++++++++++ .../src/queries/generic_inference/web/mod.rs | 1 + .../inference_job/kill_inference_jobs.rs | 144 ++++++++++++++++ .../endpoints/inference_job/mod.rs | 1 + .../service/web/storyteller_web/src/routes.rs | 6 + 5 files changed, 309 insertions(+) create mode 100644 crates/lib/mysql_queries/src/queries/generic_inference/web/kill_generic_inference_jobs.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/kill_inference_jobs.rs diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/kill_generic_inference_jobs.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/kill_generic_inference_jobs.rs new file mode 100644 index 0000000000..4928fcd6f4 --- /dev/null +++ b/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, + pub target: KillableTarget, + pub maybe_priority_or_lower: Option, + 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, + target: KillableTarget, + maybe_priority_or_lower: Option, +) -> QueryBuilder<'static, MySql> { + + // NB: Query cannot be statically checked by sqlx + let mut query_builder: QueryBuilder = 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::>(); + + // 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()); + } +} 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 71a2786c0b..53a89e6a78 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 @@ -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; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/kill_inference_jobs.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/kill_inference_jobs.rs new file mode 100644 index 0000000000..76957e221c --- /dev/null +++ b/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, + pub target: KillableTarget, + pub maybe_priority_or_lower: Option, +} + +#[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, + server_state: web::Data>) -> Result +{ + 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::>(); + + 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)) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/mod.rs index 2ee8605d27..e39b062b03 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/mod.rs +++ b/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; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 9b6f920061..071b76c46d 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/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 (app: App) -> App .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")