mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-31 15:37:49 +00:00
WorldLabs download job (#1125)
This commit is contained in:
Brandon Thomas authored and 
GitHub committed
1 parent
f97def83e1
commit
ff268c4e4e
25 files changed
+1148
-7
No files matched your search
+3
-3
@@ -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"
|
||||
}
|
||||
+102
@@ -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"
|
||||
}
|
||||
Generated
+31
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 /
|
||||
|
||||
+98
@@ -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<UserToken>,
|
||||
pub maybe_creator_anonymous_visitor_token: Option<AnonymousVisitorTrackingToken>,
|
||||
pub creator_ip_address: String,
|
||||
pub creator_set_visibility: Visibility,
|
||||
|
||||
pub maybe_prompt_token: Option<PromptToken>,
|
||||
|
||||
pub maybe_wallet_ledger_entry_token: Option<WalletLedgerEntryToken>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RawRecord {
|
||||
job_token: InferenceJobToken,
|
||||
operation_id: Option<String>,
|
||||
maybe_creator_user_token: Option<UserToken>,
|
||||
maybe_creator_anonymous_visitor_token: Option<AnonymousVisitorTrackingToken>,
|
||||
creator_ip_address: String,
|
||||
creator_set_visibility: Visibility,
|
||||
maybe_prompt_token: Option<PromptToken>,
|
||||
maybe_wallet_ledger_entry_token: Option<WalletLedgerEntryToken>,
|
||||
}
|
||||
|
||||
/// Returns all non-terminal WorldLabs jobs that have an associated operation_id.
|
||||
pub async fn list_pending_worldlabs_jobs(pool: &MySqlPool) -> AnyhowResult<Vec<PendingWorldlabsJob>> {
|
||||
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)
|
||||
}
|
||||
@@ -1 +1,2 @@
|
||||
pub mod insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token;
|
||||
pub mod list_pending_worldlabs_jobs;
|
||||
+10
-3
@@ -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,
|
||||
|
||||
+1
@@ -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,
|
||||
|
||||
+10
@@ -73,6 +73,9 @@ pub struct MediaFileInsertBuilder {
|
||||
// // If batch generated, this is the batch token.
|
||||
maybe_batch_generation_token: Option<BatchGenerationToken>,
|
||||
|
||||
// Cover image (e.g. thumbnail for 3D splats)
|
||||
maybe_cover_image_media_file_token: Option<MediaFileToken>,
|
||||
|
||||
// Storage details
|
||||
public_bucket_directory_hash: Option<MediaFileBucketPath>, // 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
|
||||
|
||||
+1
@@ -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)
|
||||
|
||||
+11
-1
@@ -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);
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
[package]
|
||||
name = "worldlabs-job"
|
||||
edition = "2021"
|
||||
version = "0.0.1"
|
||||
authors = [
|
||||
"Brandon Thomas <bt@brand.io>",
|
||||
"Brandon Thomas <echelon@gmail.com>",
|
||||
]
|
||||
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
|
||||
@@ -0,0 +1,4 @@
|
||||
# default .env configurations (these can be overridden with k8s)
|
||||
|
||||
# Development
|
||||
SERVER_ENVIRONMENT='development'
|
||||
@@ -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<Arc<HttpServerSharedState>>,
|
||||
) -> Result<HttpResponse, HealthCheckError> {
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod health_check_handler;
|
||||
@@ -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,
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod endpoints;
|
||||
pub mod http_server_shared_state;
|
||||
pub mod run_http_server;
|
||||
@@ -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<Server> {
|
||||
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(())
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod process_failed_job;
|
||||
pub mod process_successful_job;
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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<u8> = 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<MediaFileToken> {
|
||||
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<u8> = 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)
|
||||
}
|
||||
Reference in new issue
Block a user