From ff268c4e4e7ecab63ac7ea93c6800223c8870eb8 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 10 Mar 2026 08:44:03 -0400 Subject: [PATCH] WorldLabs download job (#1125) --- ...18d56b6672c07be082daa7c1a782400cb18d.json} | 6 +- ...0f16402e2273dbab9a15d7c17eb1f0da7a1b1.json | 102 +++++++ Cargo.lock | 31 +++ Cargo.toml | 1 + build/service_cpu.Dockerfile | 7 + .../worldlabs/list_pending_worldlabs_jobs.rs | 98 +++++++ .../generic_inference/worldlabs/mod.rs | 1 + .../insert_media_file_generic.rs | 13 +- .../insert_media_file_generic_from_job.rs | 1 + .../media_file_insert_builder.rs | 10 + ...ert_media_file_from_studio_scene_render.rs | 1 + .../media_file_origin_product_category.rs | 12 +- crates/service/job/worldlabs_job/Cargo.toml | 57 ++++ .../config/worldlabs-job.development.env | 4 + .../endpoints/health_check_handler.rs | 101 +++++++ .../src/http_server/endpoints/mod.rs | 1 + .../http_server/http_server_shared_state.rs | 7 + .../job/worldlabs_job/src/http_server/mod.rs | 3 + .../src/http_server/run_http_server.rs | 67 +++++ .../job/worldlabs_job/src/job_dependencies.rs | 26 ++ crates/service/job/worldlabs_job/src/main.rs | 143 ++++++++++ .../job/worldlabs_job/src/main_loop.rs | 96 +++++++ .../job/worldlabs_job/src/process_job/mod.rs | 2 + .../src/process_job/process_failed_job.rs | 116 ++++++++ .../src/process_job/process_successful_job.rs | 249 ++++++++++++++++++ 25 files changed, 1148 insertions(+), 7 deletions(-) rename .sqlx/{query-b4470e607c005421dd8504b9f1d695080a581936ac6ce00b6f437b2764e741be.json => query-87d5511700985daaabfc9f2972f018d56b6672c07be082daa7c1a782400cb18d.json} (81%) create mode 100644 .sqlx/query-d56a26a7300495fac9741e77e390f16402e2273dbab9a15d7c17eb1f0da7a1b1.json create mode 100644 crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/list_pending_worldlabs_jobs.rs create mode 100644 crates/service/job/worldlabs_job/Cargo.toml create mode 100644 crates/service/job/worldlabs_job/config/worldlabs-job.development.env create mode 100644 crates/service/job/worldlabs_job/src/http_server/endpoints/health_check_handler.rs create mode 100644 crates/service/job/worldlabs_job/src/http_server/endpoints/mod.rs create mode 100644 crates/service/job/worldlabs_job/src/http_server/http_server_shared_state.rs create mode 100644 crates/service/job/worldlabs_job/src/http_server/mod.rs create mode 100644 crates/service/job/worldlabs_job/src/http_server/run_http_server.rs create mode 100644 crates/service/job/worldlabs_job/src/job_dependencies.rs create mode 100644 crates/service/job/worldlabs_job/src/main.rs create mode 100644 crates/service/job/worldlabs_job/src/main_loop.rs create mode 100644 crates/service/job/worldlabs_job/src/process_job/mod.rs create mode 100644 crates/service/job/worldlabs_job/src/process_job/process_failed_job.rs create mode 100644 crates/service/job/worldlabs_job/src/process_job/process_successful_job.rs diff --git a/.sqlx/query-b4470e607c005421dd8504b9f1d695080a581936ac6ce00b6f437b2764e741be.json b/.sqlx/query-87d5511700985daaabfc9f2972f018d56b6672c07be082daa7c1a782400cb18d.json similarity index 81% rename from .sqlx/query-b4470e607c005421dd8504b9f1d695080a581936ac6ce00b6f437b2764e741be.json rename to .sqlx/query-87d5511700985daaabfc9f2972f018d56b6672c07be082daa7c1a782400cb18d.json index 72e13fd7a8..de7500c216 100644 --- a/.sqlx/query-b4470e607c005421dd8504b9f1d695080a581936ac6ce00b6f437b2764e741be.json +++ b/.sqlx/query-87d5511700985daaabfc9f2972f018d56b6672c07be082daa7c1a782400cb18d.json @@ -1,12 +1,12 @@ { "db_name": "MySQL", - "query": "\n INSERT INTO media_files\n SET\n token = ?,\n\n media_class = ?,\n media_type = ?,\n\n is_user_upload = ?,\n is_intermediate_system_file = ?,\n\n origin_category = ?, \n origin_product_category = ?, \n maybe_origin_model_type = ?, \n maybe_origin_model_token = ?, \n maybe_origin_filename = ?,\n\n is_batch_generated = ?,\n maybe_batch_token = ?,\n\n maybe_mime_type = ?,\n file_size_bytes = ?,\n maybe_duration_millis = ?,\n maybe_audio_encoding = ?, \n maybe_video_encoding = ?, \n maybe_frame_width = ?, \n maybe_frame_height = ?,\n maybe_prompt_token = ?,\n checksum_sha2 = ?,\n \n maybe_engine_category = ?,\n\n maybe_title = ?,\n maybe_text_transcript = ?,\n\n maybe_scene_source_media_file_token = ?,\n\n public_bucket_directory_hash = ?, \n maybe_public_bucket_prefix = ?, \n maybe_public_bucket_extension = ?, \n\n maybe_creator_user_token = ?, \n maybe_creator_anonymous_visitor_token = ?, \n\n creator_ip_address = ?, \n creator_set_visibility = ?, \n\n maybe_creator_file_synthetic_id = ?, \n maybe_creator_category_synthetic_id = ?,\n\n extra_file_modification_info = ?,\n\n maybe_mod_user_token = ?, \n is_generated_on_prem = ?, \n generated_by_worker = ?, \n generated_by_cluster = ?\n ", + "query": "\n INSERT INTO media_files\n SET\n token = ?,\n\n media_class = ?,\n media_type = ?,\n\n is_user_upload = ?,\n is_intermediate_system_file = ?,\n\n origin_category = ?, \n origin_product_category = ?, \n maybe_origin_model_type = ?, \n maybe_origin_model_token = ?, \n maybe_origin_filename = ?,\n\n is_batch_generated = ?,\n maybe_batch_token = ?,\n\n maybe_mime_type = ?,\n file_size_bytes = ?,\n maybe_duration_millis = ?,\n maybe_audio_encoding = ?, \n maybe_video_encoding = ?, \n maybe_frame_width = ?, \n maybe_frame_height = ?,\n maybe_prompt_token = ?,\n checksum_sha2 = ?,\n \n maybe_engine_category = ?,\n\n maybe_title = ?,\n maybe_text_transcript = ?,\n\n maybe_scene_source_media_file_token = ?,\n\n public_bucket_directory_hash = ?, \n maybe_public_bucket_prefix = ?, \n maybe_public_bucket_extension = ?, \n\n maybe_creator_user_token = ?, \n maybe_creator_anonymous_visitor_token = ?, \n\n creator_ip_address = ?, \n creator_set_visibility = ?, \n\n maybe_creator_file_synthetic_id = ?, \n maybe_creator_category_synthetic_id = ?,\n\n extra_file_modification_info = ?,\n\n maybe_cover_image_media_file_token = ?,\n\n maybe_mod_user_token = ?,\n is_generated_on_prem = ?,\n generated_by_worker = ?,\n generated_by_cluster = ?\n ", "describe": { "columns": [], "parameters": { - "Right": 39 + "Right": 40 }, "nullable": [] }, - "hash": "b4470e607c005421dd8504b9f1d695080a581936ac6ce00b6f437b2764e741be" + "hash": "87d5511700985daaabfc9f2972f018d56b6672c07be082daa7c1a782400cb18d" } diff --git a/.sqlx/query-d56a26a7300495fac9741e77e390f16402e2273dbab9a15d7c17eb1f0da7a1b1.json b/.sqlx/query-d56a26a7300495fac9741e77e390f16402e2273dbab9a15d7c17eb1f0da7a1b1.json new file mode 100644 index 0000000000..4337a59014 --- /dev/null +++ b/.sqlx/query-d56a26a7300495fac9741e77e390f16402e2273dbab9a15d7c17eb1f0da7a1b1.json @@ -0,0 +1,102 @@ +{ + "db_name": "MySQL", + "query": "\nSELECT\n jobs.token as `job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken`,\n jobs.maybe_external_third_party_id as `operation_id`,\n jobs.maybe_creator_user_token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n jobs.maybe_creator_anonymous_visitor_token as `maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken`,\n jobs.creator_ip_address,\n jobs.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n jobs.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`,\n jobs.maybe_wallet_ledger_entry_token as `maybe_wallet_ledger_entry_token: tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken`\n\nFROM generic_inference_jobs as jobs\n\nWHERE jobs.maybe_external_third_party = ?\n AND jobs.status NOT IN ('complete_success', 'complete_failure')\n AND jobs.maybe_external_third_party_id IS NOT NULL\n\nLIMIT 25000\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "operation_id", + "type_info": { + "type": "VarString", + "flags": "MULTIPLE_KEY | BINARY", + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 2, + "name": "maybe_creator_user_token: tokens::tokens::users::UserToken", + "type_info": { + "type": "VarString", + "flags": "MULTIPLE_KEY | BINARY", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken", + "type_info": { + "type": "VarString", + "flags": "MULTIPLE_KEY | BINARY", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 5, + "name": "creator_set_visibility: enums::common::visibility::Visibility", + "type_info": { + "type": "String", + "flags": "NOT_NULL | BINARY | ENUM", + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 6, + "name": "maybe_prompt_token: tokens::tokens::prompts::PromptToken", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "maybe_wallet_ledger_entry_token: tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 128 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + true, + true, + false, + false, + true, + true + ] + }, + "hash": "d56a26a7300495fac9741e77e390f16402e2273dbab9a15d7c17eb1f0da7a1b1" +} diff --git a/Cargo.lock b/Cargo.lock index de0cd07954..1a9fb002b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15067,6 +15067,37 @@ dependencies = [ "wreq-util", ] +[[package]] +name = "worldlabs-job" +version = "0.0.1" +dependencies = [ + "actix-web", + "actix_helpers", + "anyhow", + "bootstrap", + "bucket_paths", + "cloud_storage", + "concurrency", + "config", + "easyenv", + "enums", + "errors", + "hashing", + "jobs_common", + "log", + "mysql_queries", + "reqwest 0.12.23", + "serde", + "serde_derive", + "serde_json", + "server_environment", + "sqlx", + "tokens", + "tokio", + "workspace-build-acceleration", + "world_labs_api", +] + [[package]] name = "wreq" version = "6.0.0-rc.28" diff --git a/Cargo.toml b/Cargo.toml index 9688c57159..35066058b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -120,6 +120,7 @@ members = [ "crates/service/job/es-update-job", "crates/service/job/inference_job", "crates/service/job/seedance2_pro_job", + "crates/service/job/worldlabs_job", "crates/service/job/tts_download_job", "crates/service/job/tts_inference_job", "crates/service/job/w2l_download_job", diff --git a/build/service_cpu.Dockerfile b/build/service_cpu.Dockerfile index a651f452e4..8276a1c4b7 100644 --- a/build/service_cpu.Dockerfile +++ b/build/service_cpu.Dockerfile @@ -143,6 +143,12 @@ RUN SQLX_OFFLINE=true \ --release \ --bin seedance2-pro-job +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin worldlabs-job + RUN SQLX_OFFLINE=true \ LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ $HOME/.cargo/bin/cargo build \ @@ -203,6 +209,7 @@ RUN echo -n ${GIT_SHA} > GIT_SHA # Copy all the binaries (except those that need a GPU): COPY --from=builder /tmp/target/release/storyteller-web / COPY --from=builder /tmp/target/release/seedance2-pro-job / +COPY --from=builder /tmp/target/release/worldlabs-job / COPY --from=builder /tmp/target/release/dummy-service / COPY --from=builder /tmp/target/release/analytics-job / COPY --from=builder /tmp/target/release/email-sender-job / diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/list_pending_worldlabs_jobs.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/list_pending_worldlabs_jobs.rs new file mode 100644 index 0000000000..e3e9899e40 --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/list_pending_worldlabs_jobs.rs @@ -0,0 +1,98 @@ +use anyhow::anyhow; +use log::warn; +use sqlx::MySqlPool; + +use enums::by_table::generic_inference_jobs::inference_job_external_third_party::InferenceJobExternalThirdParty; +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::prompts::PromptToken; +use tokens::tokens::users::UserToken; +use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken; + +/// A WorldLabs job that is waiting for a result from the external API. +#[derive(Debug)] +pub struct PendingWorldlabsJob { + pub job_token: InferenceJobToken, + + /// The external operation_id from World Labs. + pub operation_id: String, + + pub maybe_creator_user_token: Option, + pub maybe_creator_anonymous_visitor_token: Option, + pub creator_ip_address: String, + pub creator_set_visibility: Visibility, + + pub maybe_prompt_token: Option, + + pub maybe_wallet_ledger_entry_token: Option, +} + +#[derive(Debug, Default)] +struct RawRecord { + job_token: InferenceJobToken, + operation_id: Option, + maybe_creator_user_token: Option, + maybe_creator_anonymous_visitor_token: Option, + creator_ip_address: String, + creator_set_visibility: Visibility, + maybe_prompt_token: Option, + maybe_wallet_ledger_entry_token: Option, +} + +/// Returns all non-terminal WorldLabs jobs that have an associated operation_id. +pub async fn list_pending_worldlabs_jobs(pool: &MySqlPool) -> AnyhowResult> { + let records = sqlx::query_as!( + RawRecord, + r#" +SELECT + jobs.token as `job_token: tokens::tokens::generic_inference_jobs::InferenceJobToken`, + jobs.maybe_external_third_party_id as `operation_id`, + jobs.maybe_creator_user_token as `maybe_creator_user_token: tokens::tokens::users::UserToken`, + jobs.maybe_creator_anonymous_visitor_token as `maybe_creator_anonymous_visitor_token: tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken`, + jobs.creator_ip_address, + jobs.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, + jobs.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`, + jobs.maybe_wallet_ledger_entry_token as `maybe_wallet_ledger_entry_token: tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken` + +FROM generic_inference_jobs as jobs + +WHERE jobs.maybe_external_third_party = ? + AND jobs.status NOT IN ('complete_success', 'complete_failure') + AND jobs.maybe_external_third_party_id IS NOT NULL + +LIMIT 25000 + "#, + InferenceJobExternalThirdParty::Worldlabs.to_str(), + ) + .fetch_all(pool) + .await + .map_err(|err| anyhow!("error querying pending worldlabs jobs: {:?}", err))?; + + let jobs = records + .into_iter() + .filter_map(|record| { + let operation_id = match record.operation_id { + Some(id) => id, + None => { + warn!("PendingWorldlabsJob has no operation_id, skipping"); + return None; + } + }; + + Some(PendingWorldlabsJob { + job_token: record.job_token, + operation_id, + maybe_creator_user_token: record.maybe_creator_user_token, + maybe_creator_anonymous_visitor_token: record.maybe_creator_anonymous_visitor_token, + creator_ip_address: record.creator_ip_address, + creator_set_visibility: record.creator_set_visibility, + maybe_prompt_token: record.maybe_prompt_token, + maybe_wallet_ledger_entry_token: record.maybe_wallet_ledger_entry_token, + }) + }) + .collect(); + + Ok(jobs) +} diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs index 259d0e2f25..f013474c15 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs @@ -1 +1,2 @@ pub mod insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token; +pub mod list_pending_worldlabs_jobs; diff --git a/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic.rs b/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic.rs index d132b59188..6480534ec1 100644 --- a/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic.rs +++ b/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic.rs @@ -88,6 +88,9 @@ pub struct InsertArgs<'a> { pub generated_by_worker: Option<&'a str>, pub generated_by_cluster: Option<&'a str>, + // Cover image (e.g. thumbnail for 3D splats) + pub maybe_cover_image_media_file_token: Option<&'a MediaFileToken>, + // Moderation details (deprecated) pub maybe_mod_user_token: Option<&'a UserToken>, } @@ -179,9 +182,11 @@ pub async fn insert_media_file_generic( extra_file_modification_info = ?, - maybe_mod_user_token = ?, - is_generated_on_prem = ?, - generated_by_worker = ?, + maybe_cover_image_media_file_token = ?, + + maybe_mod_user_token = ?, + is_generated_on_prem = ?, + generated_by_worker = ?, generated_by_cluster = ? "#, result_token, @@ -233,6 +238,8 @@ pub async fn insert_media_file_generic( extra_file_modification_info, + args.maybe_cover_image_media_file_token.map(|t| t.as_str()), + args.maybe_mod_user_token, args.is_generated_on_prem, args.generated_by_worker, diff --git a/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic_from_job.rs b/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic_from_job.rs index 46d9ccd561..60e51afbac 100644 --- a/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic_from_job.rs +++ b/crates/schema/database/mysql_queries/src/queries/media_files/create/generic_insert/insert_media_file_generic_from_job.rs @@ -120,6 +120,7 @@ pub async fn insert_media_file_generic_from_job( maybe_text_transcript: args.maybe_text_transcript, maybe_title: args.maybe_title, maybe_video_encoding: args.maybe_video_encoding, + maybe_cover_image_media_file_token: None, media_class: args.media_class, media_type: args.media_type, origin_category: args.origin_category, diff --git a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs index fc1cbb995f..f632c2b023 100644 --- a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs +++ b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs @@ -73,6 +73,9 @@ pub struct MediaFileInsertBuilder { // // If batch generated, this is the batch token. maybe_batch_generation_token: Option, + // Cover image (e.g. thumbnail for 3D splats) + maybe_cover_image_media_file_token: Option, + // Storage details public_bucket_directory_hash: Option, // NB: Non-nullable field(s) @@ -107,6 +110,7 @@ impl MediaFileInsertBuilder { maybe_origin_filename: None, maybe_engine_category: None, maybe_batch_generation_token: None, + maybe_cover_image_media_file_token: None, public_bucket_directory_hash: None, maybe_prompt_token: None, } @@ -250,6 +254,11 @@ impl MediaFileInsertBuilder { self } + pub fn maybe_cover_image_media_file_token(mut self, maybe_token: Option<&MediaFileToken>) -> Self { + self.maybe_cover_image_media_file_token = maybe_token.cloned(); + self + } + pub fn public_bucket_directory_hash(mut self, public_bucket_directory_hash: &MediaFileBucketPath) -> Self { self.public_bucket_directory_hash = Some(public_bucket_directory_hash.clone()); self @@ -308,6 +317,7 @@ impl MediaFileInsertBuilder { maybe_creator_file_synthetic_id_category: IdCategory::MediaFile, // TODO: Remove this maybe_creator_category_synthetic_id_category: IdCategory::FileUpload, // TODO: Remove this maybe_extra_media_info: None, // TODO + maybe_cover_image_media_file_token: self.maybe_cover_image_media_file_token.as_ref(), is_generated_on_prem: false, // TODO generated_by_worker: None, // TODO generated_by_cluster: None, // TODO diff --git a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_media_file_from_studio_scene_render.rs b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_media_file_from_studio_scene_render.rs index 846b187a34..459704418a 100644 --- a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_media_file_from_studio_scene_render.rs +++ b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_media_file_from_studio_scene_render.rs @@ -112,6 +112,7 @@ pub async fn insert_media_file_from_studio_scene_render( maybe_prompt_token: None, maybe_mod_user_token: None, maybe_engine_category: None, + maybe_cover_image_media_file_token: None, }).await?; Ok(new_media_token) diff --git a/crates/schema/public/enums/src/by_table/media_files/media_file_origin_product_category.rs b/crates/schema/public/enums/src/by_table/media_files/media_file_origin_product_category.rs index 062a4dc723..16824299bf 100644 --- a/crates/schema/public/enums/src/by_table/media_files/media_file_origin_product_category.rs +++ b/crates/schema/public/enums/src/by_table/media_files/media_file_origin_product_category.rs @@ -69,6 +69,9 @@ pub enum MediaFileOriginProductCategory { #[serde(rename = "video_gen")] VideoGeneration, + #[serde(rename = "world_gen")] + WorldGeneration, + // Media files for video filters #[deprecated(note = "This isn't relevant product surface area anymore")] #[serde(rename = "video_filter")] @@ -102,6 +105,7 @@ impl MediaFileOriginProductCategory { Self::Mocap => "mocap", Self::ImageGeneration => "image_gen", Self::VideoGeneration => "video_gen", + Self::WorldGeneration => "world_gen", Self::VideoFilter => "video_filter", Self::Workflow => "workflow", } @@ -122,6 +126,7 @@ impl MediaFileOriginProductCategory { "mocap" => Ok(Self::Mocap), "image_gen" => Ok(Self::ImageGeneration), "video_gen" => Ok(Self::VideoGeneration), + "world_gen" => Ok(Self::WorldGeneration), "video_filter" => Ok(Self::VideoFilter), "workflow" => Ok(Self::Workflow), _ => Err(format!("invalid value: {:?}", value)), @@ -145,6 +150,7 @@ impl MediaFileOriginProductCategory { Self::Mocap, Self::ImageGeneration, Self::VideoGeneration, + Self::WorldGeneration, Self::VideoFilter, Self::Workflow, ]) @@ -174,6 +180,7 @@ mod tests { assert_serialization(MediaFileOriginProductCategory::Mocap, "mocap"); assert_serialization(MediaFileOriginProductCategory::ImageGeneration, "image_gen"); assert_serialization(MediaFileOriginProductCategory::VideoGeneration, "video_gen"); + assert_serialization(MediaFileOriginProductCategory::WorldGeneration, "world_gen"); assert_serialization(MediaFileOriginProductCategory::VideoFilter, "video_filter"); assert_serialization(MediaFileOriginProductCategory::Workflow, "workflow"); } @@ -193,6 +200,7 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::Mocap.to_str(), "mocap"); assert_eq!(MediaFileOriginProductCategory::ImageGeneration.to_str(), "image_gen"); assert_eq!(MediaFileOriginProductCategory::VideoGeneration.to_str(), "video_gen"); + assert_eq!(MediaFileOriginProductCategory::WorldGeneration.to_str(), "world_gen"); assert_eq!(MediaFileOriginProductCategory::VideoFilter.to_str(), "video_filter"); assert_eq!(MediaFileOriginProductCategory::Workflow.to_str(), "workflow"); } @@ -212,6 +220,7 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::from_str("mocap").unwrap(), MediaFileOriginProductCategory::Mocap); assert_eq!(MediaFileOriginProductCategory::from_str("image_gen").unwrap(),MediaFileOriginProductCategory::ImageGeneration); assert_eq!(MediaFileOriginProductCategory::from_str("video_gen").unwrap(), MediaFileOriginProductCategory::VideoGeneration); + assert_eq!(MediaFileOriginProductCategory::from_str("world_gen").unwrap(), MediaFileOriginProductCategory::WorldGeneration); assert_eq!(MediaFileOriginProductCategory::from_str("video_filter").unwrap(), MediaFileOriginProductCategory::VideoFilter); assert_eq!(MediaFileOriginProductCategory::from_str("workflow").unwrap(), MediaFileOriginProductCategory::Workflow); } @@ -219,7 +228,7 @@ mod tests { #[test] fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); - assert_eq!(variants.len(), 15); + assert_eq!(variants.len(), 16); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceFusion)); @@ -233,6 +242,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Mocap)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ImageGeneration)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoGeneration)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::WorldGeneration)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoFilter)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Workflow)); assert_eq!(variants.pop_first(), None); diff --git a/crates/service/job/worldlabs_job/Cargo.toml b/crates/service/job/worldlabs_job/Cargo.toml new file mode 100644 index 0000000000..433158714c --- /dev/null +++ b/crates/service/job/worldlabs_job/Cargo.toml @@ -0,0 +1,57 @@ +[package] +name = "worldlabs-job" +edition = "2021" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[[bin]] +name = "worldlabs-job" +path = "src/main.rs" + +[package.metadata.cargo-udeps.ignore] +normal = ["workspace-build-acceleration"] + +[dependencies] +# Workspace hack (faster compile times) +workspace-build-acceleration = { workspace = true } + +# Internal +actix_helpers.workspace = true +bootstrap = { path = "../../../lib/bootstrap" } +bucket_paths = { path = "../../../lib/bucket_paths" } +cloud_storage.workspace = true +concurrency = { path = "../../../lib/concurrency" } +config = { path = "../../../lib/deprecated/config" } +easyenv = { path = "../../../lib/easyenv" } +enums.workspace = true +errors.workspace = true +hashing = { path = "../../../lib/files/hashing" } +jobs_common = { path = "../../../lib/jobs_common" } +mysql_queries = { path = "../../../schema/database/mysql_queries" } +world_labs_api = { path = "../../../api_clients/world_labs_api" } +server_environment = { path = "../../../lib/server_environment" } +tokens.workspace = true + +# External +anyhow.workspace = true +log.workspace = true + +sqlx = { version = "0.7.4", features = [ "mysql", "runtime-tokio-rustls", "chrono" ] } + +# Runtime +tokio = { version = "1.40.0", features = ["macros", "sync"] } + +# HTTP Client +reqwest.workspace = true + +# HTTP Server +actix-web.workspace = true + +# Serialization +serde.workspace = true +serde_derive.workspace = true +serde_json.workspace = true diff --git a/crates/service/job/worldlabs_job/config/worldlabs-job.development.env b/crates/service/job/worldlabs_job/config/worldlabs-job.development.env new file mode 100644 index 0000000000..91b973c216 --- /dev/null +++ b/crates/service/job/worldlabs_job/config/worldlabs-job.development.env @@ -0,0 +1,4 @@ +# default .env configurations (these can be overridden with k8s) + +# Development +SERVER_ENVIRONMENT='development' diff --git a/crates/service/job/worldlabs_job/src/http_server/endpoints/health_check_handler.rs b/crates/service/job/worldlabs_job/src/http_server/endpoints/health_check_handler.rs new file mode 100644 index 0000000000..419bcba0cd --- /dev/null +++ b/crates/service/job/worldlabs_job/src/http_server/endpoints/health_check_handler.rs @@ -0,0 +1,101 @@ +use std::sync::Arc; + +use actix_web::error::ResponseError; +use actix_web::http::header::ContentType; +use actix_web::http::StatusCode; +use actix_web::{web, HttpRequest, HttpResponse}; +use log::error; + +use actix_helpers::response_serializers::error_to_json_http_response::error_to_json_http_response; + +use crate::http_server::http_server_shared_state::HttpServerSharedState; + +#[derive(Serialize)] +pub struct HealthCheckResponse { + pub success: bool, + pub is_healthy: bool, + pub consecutive_failure_count: u64, + pub consecutive_success_count: u64, + pub total_failure_count: u64, + pub total_success_count: u64, + pub total_failure_ratio: f32, + pub total_success_ratio: f32, +} + +#[derive(Debug, Serialize)] +pub enum HealthCheckError { + ServerError, +} + +impl ResponseError for HealthCheckError { + fn status_code(&self) -> StatusCode { + match *self { + HealthCheckError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + error_to_json_http_response(self) + } +} + +impl std::fmt::Display for HealthCheckError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_health_check_handler( + _http_request: HttpRequest, + server_state: web::Data>, +) -> Result { + let job_stats = server_state + .job_stats + .get_status() + .map_err(|e| { + error!("Error serving health check status: {:?}", e); + HealthCheckError::ServerError + })?; + + let total_tries = job_stats + .total_failure_count + .saturating_add(job_stats.total_success_count); + + let total_success_ratio = if total_tries > 0 { + (job_stats.total_success_count as f32) / (total_tries as f32) + } else { + 0.0 + }; + + let total_failure_ratio = if total_tries > 0 { + 1.0 - total_success_ratio + } else { + 0.0 + }; + + let is_healthy = + job_stats.consecutive_failure_count < server_state.consecutive_failure_unhealthy_threshold; + + let response = HealthCheckResponse { + success: true, + is_healthy, + consecutive_failure_count: job_stats.consecutive_failure_count, + consecutive_success_count: job_stats.consecutive_success_count, + total_failure_count: job_stats.total_failure_count, + total_success_count: job_stats.total_success_count, + total_failure_ratio, + total_success_ratio, + }; + + let body = serde_json::to_string(&response).map_err(|_e| HealthCheckError::ServerError)?; + + if is_healthy { + Ok(HttpResponse::Ok() + .content_type(ContentType::json()) + .body(body)) + } else { + Ok(HttpResponse::InternalServerError() + .content_type(ContentType::json()) + .body(body)) + } +} diff --git a/crates/service/job/worldlabs_job/src/http_server/endpoints/mod.rs b/crates/service/job/worldlabs_job/src/http_server/endpoints/mod.rs new file mode 100644 index 0000000000..0f408270bf --- /dev/null +++ b/crates/service/job/worldlabs_job/src/http_server/endpoints/mod.rs @@ -0,0 +1 @@ +pub mod health_check_handler; diff --git a/crates/service/job/worldlabs_job/src/http_server/http_server_shared_state.rs b/crates/service/job/worldlabs_job/src/http_server/http_server_shared_state.rs new file mode 100644 index 0000000000..534526580f --- /dev/null +++ b/crates/service/job/worldlabs_job/src/http_server/http_server_shared_state.rs @@ -0,0 +1,7 @@ +use jobs_common::job_stats::JobStats; + +#[derive(Clone)] +pub struct HttpServerSharedState { + pub job_stats: JobStats, + pub consecutive_failure_unhealthy_threshold: u64, +} diff --git a/crates/service/job/worldlabs_job/src/http_server/mod.rs b/crates/service/job/worldlabs_job/src/http_server/mod.rs new file mode 100644 index 0000000000..5d888824cc --- /dev/null +++ b/crates/service/job/worldlabs_job/src/http_server/mod.rs @@ -0,0 +1,3 @@ +pub mod endpoints; +pub mod http_server_shared_state; +pub mod run_http_server; diff --git a/crates/service/job/worldlabs_job/src/http_server/run_http_server.rs b/crates/service/job/worldlabs_job/src/http_server/run_http_server.rs new file mode 100644 index 0000000000..b286560ff0 --- /dev/null +++ b/crates/service/job/worldlabs_job/src/http_server/run_http_server.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use actix_web::dev::Server; +use actix_web::middleware::Logger; +use actix_web::{web, App, HttpResponse, HttpServer}; +use log::info; + +use bootstrap::bootstrap::ContainerEnvironment; +use errors::AnyhowResult; +use jobs_common::job_stats::JobStats; + +use crate::http_server::endpoints::health_check_handler::get_health_check_handler; +use crate::http_server::http_server_shared_state::HttpServerSharedState; + +const DEFAULT_BIND_ADDRESS: &str = "0.0.0.0:12345"; +const DEFAULT_NUM_WORKERS: usize = 2; + +pub struct CreateServerArgs { + pub container_environment: ContainerEnvironment, + pub job_stats: JobStats, +} + +pub fn run_http_server(args: CreateServerArgs) -> AnyhowResult { + let bind_address = easyenv::get_env_string_or_default("HTTP_BIND_ADDRESS", DEFAULT_BIND_ADDRESS); + let num_workers = easyenv::get_env_num("HTTP_NUM_WORKERS", DEFAULT_NUM_WORKERS)?; + let _hostname = args.container_environment.hostname.clone(); + + let server_state = HttpServerSharedState { + job_stats: args.job_stats.clone(), + consecutive_failure_unhealthy_threshold: easyenv::get_env_num( + "CONSECUTIVE_FAILURE_UNHEALTHY_THRESHOLD", + 3, + )?, + }; + + let server_state_arc = web::Data::new(Arc::new(server_state)); + + info!("Starting HTTP service (for k8s health checking)."); + + let log_format = "[%{HOSTNAME}e] IP=[%{X-Forwarded-For}i] \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T"; + + let handle = HttpServer::new(move || { + App::new() + .app_data(server_state_arc.clone()) + .wrap(Logger::new(&log_format).exclude("/_status")) + .service( + web::resource("/") + .route(web::get().to(|| HttpResponse::Ok())) + .route(web::head().to(|| HttpResponse::Ok())), + ) + .service( + web::resource("/_status") + .route(web::get().to(get_health_check_handler)) + .route(web::head().to(|| HttpResponse::Ok())), + ) + }) + .bind(bind_address)? + .workers(num_workers) + .run(); + + Ok(handle) +} + +pub async fn launch_http_server(args: CreateServerArgs) -> AnyhowResult<()> { + run_http_server(args)?.await?; + Ok(()) +} diff --git a/crates/service/job/worldlabs_job/src/job_dependencies.rs b/crates/service/job/worldlabs_job/src/job_dependencies.rs new file mode 100644 index 0000000000..dcf8e148ba --- /dev/null +++ b/crates/service/job/worldlabs_job/src/job_dependencies.rs @@ -0,0 +1,26 @@ +use cloud_storage::bucket_client::BucketClient; +use concurrency::relaxed_atomic_bool::RelaxedAtomicBool; +use jobs_common::job_stats::JobStats; +use server_environment::ServerEnvironment; +use sqlx::MySqlPool; +use world_labs_api::credentials::world_labs_api_creds::WorldLabsApiCreds; + +pub struct JobDependencies { + pub mysql_pool: MySqlPool, + + /// Public GCS/S3 bucket for storing generated splats. + pub public_bucket_client: BucketClient, + + /// API credentials for World Labs. + pub worldlabs_creds: WorldLabsApiCreds, + + pub server_environment: ServerEnvironment, + + pub job_stats: JobStats, + + /// How long to sleep between poll iterations (milliseconds). + pub poll_interval_millis: u64, + + /// Set to `true` from another thread to trigger graceful shutdown. + pub application_shutdown: RelaxedAtomicBool, +} diff --git a/crates/service/job/worldlabs_job/src/main.rs b/crates/service/job/worldlabs_job/src/main.rs new file mode 100644 index 0000000000..da164f4e3c --- /dev/null +++ b/crates/service/job/worldlabs_job/src/main.rs @@ -0,0 +1,143 @@ +// Never allow these +#![forbid(private_bounds)] +#![forbid(private_interfaces)] +#![forbid(unused_must_use)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +#[macro_use] extern crate serde_derive; + +use std::time::Duration; + +use anyhow::anyhow; +use log::{info, warn}; +use sqlx::mysql::MySqlPoolOptions; + +use bootstrap::bootstrap::{bootstrap, BootstrapArgs}; +use cloud_storage::bucket_client::BucketClient; +use concurrency::relaxed_atomic_bool::RelaxedAtomicBool; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use errors::AnyhowResult; +use jobs_common::job_stats::JobStats; +use server_environment::ServerEnvironment; +use world_labs_api::credentials::world_labs_api_creds::WorldLabsApiCreds; + +use crate::http_server::run_http_server::{launch_http_server, CreateServerArgs}; +use crate::main_loop::main_loop; +use crate::job_dependencies::JobDependencies; + +pub mod http_server; +pub mod job_dependencies; +pub mod main_loop; +pub mod process_job; + +// Bucket config +const ENV_ACCESS_KEY: &str = "ACCESS_KEY"; +const ENV_SECRET_KEY: &str = "SECRET_KEY"; +const ENV_REGION_NAME: &str = "REGION_NAME"; +const ENV_PUBLIC_BUCKET_NAME: &str = "PUBLIC_BUCKET_NAME"; +const ENV_S3_ENDPOINT: &str = "S3_COMPATIBLE_ENDPOINT_URL"; +const ENV_WORLDLABS_API_KEY: &str = "WORLDLABS_API_KEY"; + +#[actix_web::main] +async fn main() -> AnyhowResult<()> { + + let container_environment = bootstrap(BootstrapArgs { + app_name: "worldlabs-job", + default_logging_override: Some(DEFAULT_RUST_LOG), + config_search_directories: &[".", "./config", "crates/service/job/worldlabs_job/config"], + ignore_legacy_dot_env_file: true, + })?; + + info!("Hostname: {}", &container_environment.hostname); + + let _k8s_node_name = easyenv::get_env_string_optional("K8S_NODE_NAME"); + let _k8s_pod_name = easyenv::get_env_string_optional("K8S_POD_NAME"); + + let db_connection_string = easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING, + ); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(2) + .connect(&db_connection_string) + .await?; + + info!("Connected to MySQL."); + + let server_environment = ServerEnvironment::from_str( + &easyenv::get_env_string_required("SERVER_ENVIRONMENT")?, + ) + .ok_or(anyhow!("invalid server environment"))?; + + // Bucket setup + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?; + let s3_compatible_endpoint_url = easyenv::get_env_string_required(ENV_S3_ENDPOINT)?; + + let bucket_timeout = easyenv::get_env_duration_seconds_or_default( + "BUCKET_TIMEOUT_SECONDS", + Duration::from_secs(60 * 5), + ); + + let public_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &public_bucket_name, + &s3_compatible_endpoint_url, + None, + Some(bucket_timeout), + )?; + + // World Labs API credentials + let worldlabs_api_key = easyenv::get_env_string_required(ENV_WORLDLABS_API_KEY)?; + let worldlabs_creds = WorldLabsApiCreds::new(worldlabs_api_key); + + // How often to poll for results (default: 15 seconds) + let poll_interval_millis: u64 = easyenv::get_env_num( + "WORLDLABS_POLL_INTERVAL_MILLIS", + 15_000, + )?; + + let application_shutdown = RelaxedAtomicBool::new(false); + let job_stats = JobStats::new(); + + let create_server_args = CreateServerArgs { + container_environment: container_environment.clone(), + job_stats: job_stats.clone(), + }; + + let job_dependencies = JobDependencies { + mysql_pool, + public_bucket_client, + worldlabs_creds, + server_environment, + job_stats, + poll_interval_millis, + application_shutdown: application_shutdown.clone(), + }; + + std::thread::spawn(move || { + let actix_runtime = actix_web::rt::System::new(); + let http_server_handle = launch_http_server(create_server_args); + + actix_runtime.block_on(http_server_handle) + .expect("HTTP server should not exit."); + + warn!("Server thread is shut down."); + application_shutdown.set(true); + }); + + main_loop(job_dependencies).await; + + Ok(()) +} diff --git a/crates/service/job/worldlabs_job/src/main_loop.rs b/crates/service/job/worldlabs_job/src/main_loop.rs new file mode 100644 index 0000000000..c3d7dea107 --- /dev/null +++ b/crates/service/job/worldlabs_job/src/main_loop.rs @@ -0,0 +1,96 @@ +use std::time::Duration; + +use log::{error, info, warn}; +use mysql_queries::queries::generic_inference::worldlabs::list_pending_worldlabs_jobs::list_pending_worldlabs_jobs; +use world_labs_api::api::api_types::operation_id::OperationId; +use world_labs_api::api::requests::get_operation::get_operation::{get_operation, GetOperationArgs}; + +use crate::process_job::process_failed_job::process_failed_job; +use crate::process_job::process_successful_job::process_successful_job; +use crate::job_dependencies::JobDependencies; + +pub async fn main_loop(job_dependencies: JobDependencies) { + while !job_dependencies.application_shutdown.get() { + let result = run_poll_iteration(&job_dependencies).await; + + if let Err(err) = result { + error!("Error in poll iteration: {:?}", err); + let _ = job_dependencies.job_stats.increment_failure_count(); + } + + tokio::time::sleep(Duration::from_millis(job_dependencies.poll_interval_millis)).await; + } + + warn!("WorldLabs job runner main loop is shut down."); +} + +async fn run_poll_iteration(deps: &JobDependencies) -> anyhow::Result<()> { + // 1. Query all non-terminal WorldLabs jobs from DB. + let pending_jobs = list_pending_worldlabs_jobs(&deps.mysql_pool).await?; + + if pending_jobs.is_empty() { + info!("No pending WorldLabs jobs."); + return Ok(()); + } + + info!("Found {} pending WorldLabs job(s).", pending_jobs.len()); + + // 2. For each pending job, poll its operation status from the World Labs API. + for job in &pending_jobs { + let operation_id = OperationId(job.operation_id.clone()); + + let operation = match get_operation(GetOperationArgs { + creds: &deps.worldlabs_creds, + operation_id: &operation_id, + request_timeout: Some(Duration::from_secs(30)), + }).await { + Ok(op) => op, + Err(err) => { + warn!( + "Error polling WorldLabs operation {} for job {}: {:?}", + job.operation_id, job.job_token.as_str(), err + ); + continue; + } + }; + + // 3. Check if the operation has an error. + if let Some(ref op_error) = operation.error { + let reason = op_error.message.as_deref().unwrap_or("unknown error"); + info!( + "Operation {} failed: {}. Processing job {} as failed.", + job.operation_id, reason, job.job_token.as_str() + ); + process_failed_job(deps, job, reason).await; + continue; + } + + // 4. Check if the operation is done. + if !operation.done { + // Still in progress — check again next poll. + info!( + "Operation {} for job {} is still in progress.", + job.operation_id, job.job_token.as_str() + ); + continue; + } + + // 5. Operation is done — process the result. + info!( + "Operation {} completed, processing job {}.", + job.operation_id, job.job_token.as_str() + ); + + if let Err(err) = process_successful_job(deps, job, &operation).await { + warn!( + "Error processing completed operation {} for job {}: {:?}", + job.operation_id, job.job_token.as_str(), err + ); + let _ = deps.job_stats.increment_failure_count(); + } else { + let _ = deps.job_stats.increment_success_count(); + } + } + + Ok(()) +} diff --git a/crates/service/job/worldlabs_job/src/process_job/mod.rs b/crates/service/job/worldlabs_job/src/process_job/mod.rs new file mode 100644 index 0000000000..357e8d5ae1 --- /dev/null +++ b/crates/service/job/worldlabs_job/src/process_job/mod.rs @@ -0,0 +1,2 @@ +pub mod process_failed_job; +pub mod process_successful_job; diff --git a/crates/service/job/worldlabs_job/src/process_job/process_failed_job.rs b/crates/service/job/worldlabs_job/src/process_job/process_failed_job.rs new file mode 100644 index 0000000000..76bd4976e8 --- /dev/null +++ b/crates/service/job/worldlabs_job/src/process_job/process_failed_job.rs @@ -0,0 +1,116 @@ +use log::{error, info, warn}; + +use enums::by_table::generic_inference_jobs::frontend_failure_category::FrontendFailureCategory; +use mysql_queries::queries::generic_inference::job::mark_job_failed_by_token::{mark_job_failed_by_token, MarkJobFailedByTokenArgs}; +use mysql_queries::queries::generic_inference::worldlabs::list_pending_worldlabs_jobs::PendingWorldlabsJob; +use mysql_queries::queries::wallets::refund::try_to_refund_ledger_entry::{try_to_refund_ledger_entry, WalletRefundOutcome}; + +use crate::job_dependencies::JobDependencies; + +pub async fn process_failed_job( + deps: &JobDependencies, + job: &PendingWorldlabsJob, + reason: &str, +) { + // --- Step 1: Attempt the refund before touching the job status. --- + // + // We do this first so that a crash or error between the refund and the status update + // can never result in a failed job with no refund. If the refund itself fails we bail + // out early: the job stays pending and will be retried on the next poll cycle. + + match &job.maybe_wallet_ledger_entry_token { + None => { + warn!( + "Job {} has no wallet ledger entry token; skipping refund.", + job.job_token.as_str() + ); + } + Some(ledger_token) => { + let mut transaction = match deps.mysql_pool.begin().await { + Ok(tx) => tx, + Err(err) => { + error!( + "Failed to begin refund transaction for job {} (ledger {}): {:?}. \ + Job will NOT be marked failed yet and will be retried next poll.", + job.job_token.as_str(), ledger_token.as_str(), err + ); + return; + } + }; + + match try_to_refund_ledger_entry(ledger_token, &mut transaction).await { + Ok(WalletRefundOutcome::Refunded(summary)) => { + info!( + "Refunded {} credits for failed job {} (ledger {} → refund ledger {}).", + summary.refund_amount, + job.job_token.as_str(), + ledger_token.as_str(), + summary.refund_ledger_entry_token.as_str(), + ); + if let Err(err) = transaction.commit().await { + error!( + "Failed to commit refund transaction for job {} (ledger {}): {:?}. \ + Job will NOT be marked failed yet and will be retried next poll.", + job.job_token.as_str(), ledger_token.as_str(), err + ); + return; + } + } + Ok(WalletRefundOutcome::AlreadyRefunded) => { + info!( + "Ledger entry {} for job {} was already refunded; proceeding to mark job failed.", + ledger_token.as_str(), + job.job_token.as_str(), + ); + let _ = transaction.rollback().await; + } + Err(err) => { + error!( + "Failed to refund ledger entry {} for job {}: {:?}. \ + Job will NOT be marked failed yet and will be retried next poll.", + ledger_token.as_str(), + job.job_token.as_str(), + err, + ); + let _ = transaction.rollback().await; + return; + } + } + } + } + + // --- Step 2: Mark the job record as failed. --- + + let reason_lower = reason.to_lowercase(); + + let platform_rules_violation = reason_lower.contains("violates") || + reason_lower.contains("platform rules") || + reason_lower.contains("please modify"); + + let frontend_failure_category = if platform_rules_violation { + Some(FrontendFailureCategory::ModelRulesViolation) + } else { + None + }; + + warn!( + "Operation for job {} failed: {}. Marking job failed.", + job.job_token.as_str(), reason + ); + + let mark_failed_result = mark_job_failed_by_token(MarkJobFailedByTokenArgs { + pool: &deps.mysql_pool, + job_token: &job.job_token, + maybe_public_failure_reason: Some(reason), + internal_debugging_failure_reason: reason, + maybe_frontend_failure_category: frontend_failure_category, + }).await; + + if let Err(err) = mark_failed_result { + error!( + "Error marking job {} as failed: {:?}", + job.job_token.as_str(), + err + ); + } +} diff --git a/crates/service/job/worldlabs_job/src/process_job/process_successful_job.rs b/crates/service/job/worldlabs_job/src/process_job/process_successful_job.rs new file mode 100644 index 0000000000..9afca8257c --- /dev/null +++ b/crates/service/job/worldlabs_job/src/process_job/process_successful_job.rs @@ -0,0 +1,249 @@ +use anyhow::anyhow; +use log::{error, info, warn}; + +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_class::MediaFileClass; +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use errors::AnyhowResult; +use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes; +use mysql_queries::queries::generic_inference::worldlabs::list_pending_worldlabs_jobs::PendingWorldlabsJob; +use mysql_queries::queries::media_files::create::insert_builder::media_file_insert_builder::MediaFileInsertBuilder; +use mysql_queries::queries::generic_inference::web::mark_generic_inference_job_successfully_done_by_token::mark_generic_inference_job_successfully_done_by_token; +use tokens::tokens::media_files::MediaFileToken; +use world_labs_api::api::requests::get_operation::get_operation::GetOperationResponse; +use crate::job_dependencies::JobDependencies; + +const SPLAT_PREFIX: &str = "artcraft_"; +const SPLAT_SUFFIX: &str = ".ceramic.spz"; // NB: "ceramic" triggers some heuristics in Artcraft, such as flipping the model upside down + +const THUMB_PREFIX: &str = "artcraft_"; +const THUMB_SUFFIX: &str = ".png"; + +/// Download the completed splat, upload to bucket, create media file record, and mark job done. +pub async fn process_successful_job( + deps: &JobDependencies, + job: &PendingWorldlabsJob, + operation: &GetOperationResponse, +) -> AnyhowResult<()> { + // Get the full-res splat URL from the completed world. + let world = operation.world.as_ref() + .ok_or_else(|| anyhow!("Completed operation {} has no world object", operation.operation_id))?; + + let splat_url = world.assets.as_ref() + .and_then(|a| a.splats.as_ref()) + .and_then(|s| s.spz_url_full_res.as_ref()) + .ok_or_else(|| anyhow!( + "Completed operation {} has no spz_url_full_res", + operation.operation_id + ))?; + + // --- Step 1: Download and upload the thumbnail as a cover image. --- + + let maybe_cover_token = match world.assets.as_ref().and_then(|a| a.thumbnail_url.as_ref()) { + Some(thumbnail_url) => { + match download_and_upload_thumbnail(deps, job, operation, thumbnail_url).await { + Ok(token) => Some(token), + Err(err) => { + warn!( + "Failed to create thumbnail cover image for operation {}: {:?}. Continuing without cover.", + operation.operation_id, err + ); + None + } + } + } + None => { + info!("No thumbnail_url for operation {}; skipping cover image.", operation.operation_id); + None + } + }; + + // --- Step 2: Download and upload the splat. --- + + info!( + "Downloading splat for operation {} from: {}", + operation.operation_id, splat_url + ); + + // Download the splat bytes. + let splat_bytes: Vec = reqwest::get(splat_url) + .await + .map_err(|err| anyhow!("reqwest error downloading splat: {:?}", err))? + .bytes() + .await + .map_err(|err| anyhow!("error reading splat bytes: {:?}", err))? + .to_vec(); + + info!( + "Downloaded {} bytes for operation {}", + splat_bytes.len(), + operation.operation_id + ); + + // Hash the splat. + let checksum = sha256_hash_bytes(&splat_bytes) + .map_err(|err| anyhow!("error hashing splat: {:?}", err))?; + + // Build the bucket path. + let bucket_path = MediaFileBucketPath::generate_new(Some(SPLAT_PREFIX), Some(SPLAT_SUFFIX)); + + let object_path = bucket_path.get_full_object_path_str(); + + info!( + "Uploading splat to public bucket at path: {}", + object_path + ); + + // Upload to public bucket. + deps + .public_bucket_client + .upload_file_with_content_type_process(object_path, &splat_bytes, "application/gzip") + .await + .map_err(|err| anyhow!("error uploading splat to bucket: {:?}", err))?; + + info!( + "Uploaded splat for operation {}. Creating media file record.", + operation.operation_id + ); + + // Insert media file record. + let media_file_token = MediaFileInsertBuilder::new() + .maybe_creator_user(job.maybe_creator_user_token.as_ref()) + .maybe_creator_anonymous_visitor(job.maybe_creator_anonymous_visitor_token.as_ref()) + .creator_ip_address(&job.creator_ip_address) + .creator_set_visibility(job.creator_set_visibility) + .media_file_class(MediaFileClass::Dimensional) + .media_file_type(MediaFileType::Spz) + .media_file_origin_category(MediaFileOriginCategory::Inference) + .media_file_origin_product_category(MediaFileOriginProductCategory::WorldGeneration) + .mime_type("application/gzip") + .file_size_bytes(splat_bytes.len() as u64) + .checksum_sha2(&checksum) + .maybe_prompt_token(job.maybe_prompt_token.as_ref()) + .maybe_cover_image_media_file_token(maybe_cover_token.as_ref()) + .public_bucket_directory_hash(&bucket_path) + .insert_pool(&deps.mysql_pool) + .await + .map_err(|err| anyhow!("error inserting media file record: {:?}", err))?; + + info!( + "Created media file {} for operation {}. Marking job {} complete.", + media_file_token.as_str(), + operation.operation_id, + job.job_token.as_str() + ); + + // Mark inference job as successfully completed. + mark_generic_inference_job_successfully_done_by_token( + &deps.mysql_pool, + &job.job_token, + Some(InferenceResultType::MediaFile), + Some(media_file_token.as_str()), + None, + None, + ) + .await + .map_err(|err| { + error!( + "Error marking job {} done: {:?}", + job.job_token.as_str(), + err + ); + anyhow!("error marking job done: {:?}", err) + })?; + + info!("Job {} completed successfully.", job.job_token.as_str()); + + Ok(()) +} + +/// Guess (suffix, mime_type, MediaFileType) from a URL's file extension. +/// Falls back to PNG if the extension is unrecognized or absent. +fn guess_image_format_from_url(url: &str) -> (&'static str, &'static str, MediaFileType) { + // Strip query string / fragment before checking the extension. + let path = url.split('?').next().unwrap_or(url); + let path = path.split('#').next().unwrap_or(path); + + if let Some(dot) = path.rfind('.') { + match &path[dot..] { + ".jpg" | ".jpeg" => (".jpg", "image/jpeg", MediaFileType::Jpg), + ".png" => (".png", "image/png", MediaFileType::Png), + _ => (THUMB_SUFFIX, "image/png", MediaFileType::Png), + } + } else { + (THUMB_SUFFIX, "image/png", MediaFileType::Png) + } +} + +/// Download a thumbnail image, upload it to the public bucket, and create a media file record for it. +async fn download_and_upload_thumbnail( + deps: &JobDependencies, + job: &PendingWorldlabsJob, + operation: &GetOperationResponse, + thumbnail_url: &str, +) -> AnyhowResult { + info!( + "Downloading thumbnail for operation {} from: {}", + operation.operation_id, thumbnail_url + ); + + let (suffix, mime_type, media_file_type) = guess_image_format_from_url(thumbnail_url); + + let thumb_bytes: Vec = reqwest::get(thumbnail_url) + .await + .map_err(|err| anyhow!("reqwest error downloading thumbnail: {:?}", err))? + .bytes() + .await + .map_err(|err| anyhow!("error reading thumbnail bytes: {:?}", err))? + .to_vec(); + + info!( + "Downloaded {} thumbnail bytes for operation {}", + thumb_bytes.len(), + operation.operation_id + ); + + let checksum = sha256_hash_bytes(&thumb_bytes) + .map_err(|err| anyhow!("error hashing thumbnail: {:?}", err))?; + + let bucket_path = MediaFileBucketPath::generate_new(Some(THUMB_PREFIX), Some(suffix)); + let object_path = bucket_path.get_full_object_path_str(); + + info!("Uploading thumbnail to public bucket at path: {}", object_path); + + deps + .public_bucket_client + .upload_file_with_content_type_process(object_path, &thumb_bytes, mime_type) + .await + .map_err(|err| anyhow!("error uploading thumbnail to bucket: {:?}", err))?; + + let thumb_token = MediaFileInsertBuilder::new() + .maybe_creator_user(job.maybe_creator_user_token.as_ref()) + .maybe_creator_anonymous_visitor(job.maybe_creator_anonymous_visitor_token.as_ref()) + .creator_ip_address(&job.creator_ip_address) + .creator_set_visibility(job.creator_set_visibility) + .media_file_class(MediaFileClass::Image) + .media_file_type(media_file_type) + .is_intermediate_system_file(true) + .media_file_origin_category(MediaFileOriginCategory::Inference) + .media_file_origin_product_category(MediaFileOriginProductCategory::WorldGeneration) + .mime_type(mime_type) + .file_size_bytes(thumb_bytes.len() as u64) + .checksum_sha2(&checksum) + .maybe_prompt_token(job.maybe_prompt_token.as_ref()) + .public_bucket_directory_hash(&bucket_path) + .insert_pool(&deps.mysql_pool) + .await + .map_err(|err| anyhow!("error inserting thumbnail media file record: {:?}", err))?; + + info!( + "Created thumbnail media file {} for operation {}.", + thumb_token.as_str(), + operation.operation_id + ); + + Ok(thumb_token) +}