diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs index d8976d7094..b043ab53c0 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs @@ -23,10 +23,11 @@ use mysql_queries::queries::generic_inference::web::job_status::GenericInference use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; + use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; -use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; -use crate::http_server::endpoints::inference_job::utils::extract_live_portrait_details::extract_live_portrait_details; -use crate::http_server::endpoints::inference_job::utils::extract_polymorphic_inference_args::extract_polymorphic_inference_args; +use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::endpoints::media_files::get::batch_get_media_files_handler::BatchGetMediaFilesQueryParams; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; @@ -264,7 +265,7 @@ fn db_record_to_response_payload( .ok() .flatten(); - let progress_percentage = estimate_job_progress(&record); + let progress_percentage = estimate_job_progress(&record, maybe_polymorphic_args.as_ref()); BatchInferenceJobStatusResponsePayload { job_token: record.job_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs index b83a2c31cd..c554a99691 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs @@ -20,10 +20,11 @@ use mysql_queries::queries::generic_inference::web::get_inference_job_status::ge use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; + use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; -use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; -use crate::http_server::endpoints::inference_job::utils::extract_live_portrait_details::extract_live_portrait_details; -use crate::http_server::endpoints::inference_job::utils::extract_polymorphic_inference_args::extract_polymorphic_inference_args; +use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::state::server_state::ServerState; @@ -254,7 +255,7 @@ fn record_to_payload( .ok() .flatten(); - let progress_percentage = estimate_job_progress(&record); + let progress_percentage = estimate_job_progress(&record, maybe_polymorphic_args.as_ref()); InferenceJobStatusResponsePayload { job_token: record.job_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs index cf67a614e9..1cfbdefa4b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs @@ -26,9 +26,9 @@ use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::users::UserToken; use crate::http_server::endpoints::inference_job::common_responses::live_portrait::JobDetailsLivePortraitRequest; -use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; -use crate::http_server::endpoints::inference_job::utils::extract_live_portrait_details::extract_live_portrait_details; -use crate::http_server::endpoints::inference_job::utils::extract_polymorphic_inference_args::extract_polymorphic_inference_args; +use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_live_portrait_details::extract_live_portrait_details; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::state::server_state::ServerState; @@ -329,7 +329,7 @@ fn db_record_to_response_payload( .ok() .flatten(); - let progress_percentage = estimate_job_progress(&record); + let progress_percentage = estimate_job_progress(&record, maybe_polymorphic_args.as_ref()); ListSessionJobsItem { job_token: record.job_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/comfy_workflow_estimate.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/comfy_workflow_estimate.rs similarity index 79% rename from crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/comfy_workflow_estimate.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/comfy_workflow_estimate.rs index a5b92a31c4..dbdc508572 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/comfy_workflow_estimate.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/comfy_workflow_estimate.rs @@ -1,10 +1,10 @@ use std::cmp::max; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; -use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; -use crate::http_server::endpoints::inference_job::utils::extract_comfy_workflow_args::extract_comfy_workflow_args; -use crate::http_server::endpoints::inference_job::utils::percent::percent; +use crate::http_server::endpoints::inference_job::utils::estimates::percent::percent; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_comfy_workflow_args::extract_comfy_workflow_args; const COMFY_JOB_DEFAULT_VIDEO_LENGTH_SECONDS : u64 = 3; @@ -16,8 +16,13 @@ const COMFY_JOB_AVERAGE_FACE_FUSION_EXECUTION_SECONDS_PER_SECOND: u64 = 40; // TODO: These numbers are made up. We should measure the average job durations. const COMFY_FALLBACK_AVERAGE_JOB_DURATION_SECONDS: u64 = 60 * 5; -pub fn comfy_workflow_estimate(job: &GenericInferenceJobStatus, job_duration_seconds: u64) -> u8 { - let args = match extract_comfy_workflow_args(job) { +pub fn comfy_workflow_estimate(maybe_args: Option<&PolymorphicInferenceArgs>, job_duration_seconds: u64) -> u8 { + let args = match maybe_args { + Some(args) => args, + None => return percent(job_duration_seconds, COMFY_FALLBACK_AVERAGE_JOB_DURATION_SECONDS), + }; + + let args = match extract_comfy_workflow_args(args) { Some(args) => args, None => return percent(job_duration_seconds, COMFY_FALLBACK_AVERAGE_JOB_DURATION_SECONDS), }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimate_job_progress.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs similarity index 83% rename from crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimate_job_progress.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs index 7ea4634887..a87d9e84c9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimate_job_progress.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs @@ -2,17 +2,18 @@ use std::cmp::min; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; use enums::common::job_status_plus::JobStatusPlus; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; -use crate::http_server::endpoints::inference_job::utils::comfy_workflow_estimate::comfy_workflow_estimate; -use crate::http_server::endpoints::inference_job::utils::percent::percent; +use crate::http_server::endpoints::inference_job::utils::estimates::comfy_workflow_estimate::comfy_workflow_estimate; +use crate::http_server::endpoints::inference_job::utils::estimates::percent::percent; // TODO: These numbers are made up. We should measure the average job durations. const LIPSYNC_JOB_AVERAGE_SECONDS : u64 = 60 * 3; const TTS_JOB_AVERAGE_SECONDS : u64 = 7; const VC_JOB_AVERAGE_SECONDS : u64 = 90; -pub fn estimate_job_progress(job: &GenericInferenceJobStatus) -> u8 { +pub fn estimate_job_progress(job: &GenericInferenceJobStatus, maybe_args: Option<&PolymorphicInferenceArgs>) -> u8 { match job.status { // Jobs that haven't started JobStatusPlus::Pending @@ -53,7 +54,7 @@ pub fn estimate_job_progress(job: &GenericInferenceJobStatus) -> u8 { InferenceCategory::VoiceConversion => percent(duration_seconds, VC_JOB_AVERAGE_SECONDS), // TODO: Better estimate using video duration, params, etc. - InferenceCategory::Workflow => comfy_workflow_estimate(job, duration_seconds), + InferenceCategory::Workflow => comfy_workflow_estimate(maybe_args, duration_seconds), // NB: We don't run these job types anymore. InferenceCategory::FormatConversion => 0, @@ -75,20 +76,20 @@ mod tests { use enums::common::job_status_plus::JobStatusPlus; use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; - use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; + use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; #[test] fn test_pending() { let mut job = GenericInferenceJobStatus::default(); job.status = JobStatusPlus::Pending; - assert_eq!(0, estimate_job_progress(&job)); + assert_eq!(0, estimate_job_progress(&job, None)); } #[test] fn test_complete_success() { let mut job = GenericInferenceJobStatus::default(); job.status = JobStatusPlus::CompleteSuccess; - assert_eq!(100, estimate_job_progress(&job)); + assert_eq!(100, estimate_job_progress(&job, None)); } mod running_jobs { @@ -101,7 +102,7 @@ mod tests { job.created_at = Utc::now(); job.maybe_first_started_at = Some(Utc::now()); job.database_clock = Utc::now(); - assert_eq!(0, estimate_job_progress(&job)); + assert_eq!(0, estimate_job_progress(&job, None)); } #[test] @@ -111,7 +112,7 @@ mod tests { job.created_at = Utc::now() - Duration::seconds(10); job.maybe_first_started_at = Some(Utc::now() - Duration::seconds(10)); job.database_clock = Utc::now(); - assert_eq!(5, estimate_job_progress(&job)); + assert_eq!(5, estimate_job_progress(&job, None)); } #[test] @@ -121,7 +122,7 @@ mod tests { job.created_at = Utc::now() - Duration::days(10); job.maybe_first_started_at = Some(Utc::now() - Duration::days(10)); job.database_clock = Utc::now(); - assert_eq!(95, estimate_job_progress(&job)); + assert_eq!(95, estimate_job_progress(&job, None)); } } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/mod.rs new file mode 100644 index 0000000000..0456f4c669 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/mod.rs @@ -0,0 +1,3 @@ +pub mod comfy_workflow_estimate; +pub mod estimate_job_progress; +mod percent; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/percent.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/percent.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/percent.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/percent.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_comfy_workflow_args.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_comfy_workflow_args.rs deleted file mode 100644 index 3d228f3d55..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_comfy_workflow_args.rs +++ /dev/null @@ -1,26 +0,0 @@ -use anyhow::anyhow; - -use errors::AnyhowResult; -use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; -use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; -use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus; - -use crate::http_server::endpoints::inference_job::utils::extract_polymorphic_inference_args::extract_polymorphic_inference_args; - -pub fn extract_comfy_workflow_args(job: &GenericInferenceJobStatus) -> Option { - extract_comfy_workflow_args_fallible(job).ok().flatten() -} - -pub fn extract_comfy_workflow_args_fallible(job: &GenericInferenceJobStatus) -> AnyhowResult> { - let maybe_args = extract_polymorphic_inference_args(job)?; - - match maybe_args { - Some(args) => match args { - PolymorphicInferenceArgs::Cu(workflow_args) => Ok(Some(workflow_args.clone())), - _ => return Err(anyhow!("wrong inner args for job!")), - }, - None => { - return Err(anyhow!("no args for job!")); - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_comfy_workflow_args.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_comfy_workflow_args.rs new file mode 100644 index 0000000000..15f7eb5c30 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_comfy_workflow_args.rs @@ -0,0 +1,16 @@ +use anyhow::anyhow; + +use errors::AnyhowResult; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowArgs; + +pub fn extract_comfy_workflow_args(args: &PolymorphicInferenceArgs) -> Option { + extract_comfy_workflow_args_fallible(args).ok().flatten() +} + +pub fn extract_comfy_workflow_args_fallible(args: &PolymorphicInferenceArgs) -> AnyhowResult> { + match args { + PolymorphicInferenceArgs::Cu(workflow_args) => Ok(Some(workflow_args.clone())), + _ => return Err(anyhow!("wrong inner args for job!")), + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_live_portrait_details.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_live_portrait_details.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_live_portrait_details.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_live_portrait_details.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_polymorphic_inference_args.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_polymorphic_inference_args.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extract_polymorphic_inference_args.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/extract_polymorphic_inference_args.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs new file mode 100644 index 0000000000..8300fc5b96 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/extractors/mod.rs @@ -0,0 +1,3 @@ +pub mod extract_comfy_workflow_args; +pub mod extract_live_portrait_details; +pub mod extract_polymorphic_inference_args; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/mod.rs index 9cd9253e8c..c2c8fc0b14 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/mod.rs @@ -1,6 +1,2 @@ -mod comfy_workflow_estimate; -mod extract_comfy_workflow_args; -mod percent; -pub mod estimate_job_progress; -pub mod extract_live_portrait_details; -pub mod extract_polymorphic_inference_args; \ No newline at end of file +pub mod estimates; +pub mod extractors; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status.rs index 0954f51e7d..b7db780c27 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status.rs @@ -24,7 +24,8 @@ use redis_schema::keys::inference_job::style_transfer_progress_key::StyleTransfe use redis_schema::payloads::inference_job::style_transfer_progress_state::{InferenceProgressDetailsResponse, InferenceStageDetails}; use tokens::tokens::generic_inference_jobs::InferenceJobToken; -use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::state::server_state::ServerState; @@ -212,7 +213,12 @@ fn record_to_payload( ) -> InferenceJobStatusResponsePayload { let inference_category = record.request_details.inference_category; - let progress_percentage = estimate_job_progress(&record); + // NB: Fail open. We don't want to fail the request if we can't extract the args. + let maybe_polymorphic_args = extract_polymorphic_inference_args(&record) + .ok() + .flatten(); + + let progress_percentage = estimate_job_progress(&record, maybe_polymorphic_args.as_ref()); InferenceJobStatusResponsePayload { job_token: record.job_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status_ws.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status_ws.rs index 9d076beaf9..0a54a05cf4 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status_ws.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/get_inference_preview_status_ws.rs @@ -30,7 +30,8 @@ use redis_schema::keys::inference_job::style_transfer_progress_key::StyleTransfe use redis_schema::payloads::inference_job::style_transfer_progress_state::{InferenceProgressDetailsResponse, InferenceStageDetails}; use tokens::tokens::generic_inference_jobs::InferenceJobToken; -use crate::http_server::endpoints::inference_job::utils::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::estimates::estimate_job_progress::estimate_job_progress; +use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args; use crate::http_server::endpoints::workflows::enqueue::{progress_tracker_server, progress_tracker_session}; use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; @@ -221,7 +222,12 @@ fn record_to_payload( ) -> InferenceJobStatusResponsePayload { let inference_category = record.request_details.inference_category; - let progress_percentage = estimate_job_progress(&record); + // NB: Fail open. We don't want to fail the request if we can't extract the args. + let maybe_polymorphic_args = extract_polymorphic_inference_args(&record) + .ok() + .flatten(); + + let progress_percentage = estimate_job_progress(&record, maybe_polymorphic_args.as_ref()); InferenceJobStatusResponsePayload { job_token: record.job_token,