job supports alternate mode

This commit is contained in:
Brandon Thomas committed 2026-05-18 22:21:11 -04:00
1 parent 89c2b9f088
commit 440dffc5d4
4 files changed
+51 -14

No files matched your search

@@ -41,7 +41,11 @@ struct RawRecord {
}
/// Returns all non-terminal Seedance2Pro video jobs that have an associated order_id.
pub async fn list_pending_seedance2pro_video_jobs(pool: &MySqlPool) -> Result<Vec<PendingSeedance2ProJob>, sqlx::Error> {
pub async fn list_pending_seedance2pro_video_jobs(
pool: &MySqlPool,
external_third_party: InferenceJobExternalThirdParty,
job_type: InferenceJobType,
) -> Result<Vec<PendingSeedance2ProJob>, sqlx::Error> {
let records = sqlx::query_as!(
RawRecord,
r#"
@@ -64,8 +68,8 @@ WHERE jobs.maybe_external_third_party = ?
LIMIT 25000
"#,
InferenceJobExternalThirdParty::Seedance2Pro.to_str(),
InferenceJobType::Seedance2ProQueue.to_str(),
external_third_party.to_str(),
job_type.to_str(),
)
.fetch_all(pool)
.await?;
@@ -20,6 +20,10 @@ pub struct JobDependencies {
/// Session credentials for polling seedance2-pro.com.
pub seedance2pro_session: Seedance2ProSession,
/// When true, this instance polls for alternate-account jobs
/// (Seedance2ProAltQueue / Seedance2ProAlt) instead of the primary ones.
pub is_alternate_mode: bool,
pub server_environment: ServerEnvironment,
pub job_stats: JobStats,
@@ -2,6 +2,8 @@ use std::collections::HashMap;
use std::time::{Duration, Instant};
use chrono::Utc;
use enums::by_table::generic_inference_jobs::inference_job_external_third_party::InferenceJobExternalThirdParty;
use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
use log::{error, info, warn};
use pager::notification::notification_details_builder::NotificationDetailsBuilder;
use pager::notification::notification_urgency::NotificationUrgency;
@@ -61,7 +63,13 @@ pub async fn video_polling_main_loop(job_dependencies: JobDependencies) {
async fn run_poll_iteration(deps: &JobDependencies) -> anyhow::Result<()> {
// 1. Query all (limit 25,000) non-terminal Seedance2Pro jobs from DB.
// This is all non-(complete_success, complete_failure) jobs.
let pending_jobs = match list_pending_seedance2pro_video_jobs(&deps.mysql_pool).await {
let (third_party, job_type) = if deps.is_alternate_mode {
(InferenceJobExternalThirdParty::Seedance2ProAlt, InferenceJobType::Seedance2ProAltQueue)
} else {
(InferenceJobExternalThirdParty::Seedance2Pro, InferenceJobType::Seedance2ProQueue)
};
let pending_jobs = match list_pending_seedance2pro_video_jobs(&deps.mysql_pool, third_party, job_type).await {
Ok(jobs) => jobs,
Err(err) => {
error!("Failed to list pending database jobs: {:?}", err);
@@ -45,7 +45,11 @@ 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_SEEDANCE2PRO_COOKIES : &str = "SEEDANCE2PRO_COOKIES";
const ENV_SEEDANCE2PRO_ALTERNATE: &str = "SEEDANCE2PRO_ALTERNATE";
const ENV_SEEDANCE2PRO_COOKIES: &str = "SEEDANCE2PRO_COOKIES";
const ENV_SEEDANCE2PRO_ALT_COOKIES: &str = "SEEDANCE2PRO_ALT_COOKIES";
const ENV_MAX_JOB_AGE_THRESHOLD_HOURS: &str = "MAX_JOB_AGE_THRESHOLD_HOURS";
#[tokio::main]
@@ -101,8 +105,20 @@ async fn main() -> AnyhowResult<()> {
Some(bucket_timeout),
)?;
// Alternate mode: uses a different Kinovi account, different job types, skips character polling.
let is_alternate_mode = easyenv::get_env_string_optional(ENV_SEEDANCE2PRO_ALTERNATE)
.map(|v| v.trim().to_lowercase())
.map(|v| v == "true" || v == "1")
.unwrap_or(false);
if is_alternate_mode {
info!("Running in ALTERNATE mode (SEEDANCE2PRO_ALTERNATE=true).");
}
// Seedance2Pro session from cookie string
let seedance2pro_cookies = easyenv::get_env_string_required(ENV_SEEDANCE2PRO_COOKIES)?;
let cookie_env_var = if is_alternate_mode { ENV_SEEDANCE2PRO_ALT_COOKIES } else { ENV_SEEDANCE2PRO_COOKIES };
info!("Reading Seedance2Pro cookies from environment variable: {}", cookie_env_var);
let seedance2pro_cookies = easyenv::get_env_string_required(cookie_env_var)?;
let seedance2pro_session = Seedance2ProSession::from_cookies_string(seedance2pro_cookies);
// How often to poll for results (default: 15 seconds)
@@ -171,6 +187,7 @@ async fn main() -> AnyhowResult<()> {
mysql_pool,
public_bucket_client,
seedance2pro_session,
is_alternate_mode,
server_environment,
job_stats,
poll_interval_millis,
@@ -214,23 +231,27 @@ async fn main() -> AnyhowResult<()> {
// Spawn all polling loops as concurrent tasks.
let video_deps = job_dependencies.clone();
let character_deps = job_dependencies.clone();
let credits_deps = job_dependencies;
let credits_deps = job_dependencies.clone();
let video_handle = tokio::spawn(async move {
video_polling_main_loop(video_deps).await;
});
let character_handle = tokio::spawn(async move {
character_polling_main_loop(character_deps).await;
});
let credits_handle = tokio::spawn(async move {
credits_checking_main_loop(credits_deps).await;
});
// Wait for all to finish (they exit when application_shutdown is set).
let _ = tokio::join!(video_handle, character_handle, credits_handle);
if is_alternate_mode {
// Alternate mode: skip character polling entirely.
info!("Alternate mode: character polling is disabled.");
let _ = tokio::join!(video_handle, credits_handle);
} else {
let character_deps = job_dependencies;
let character_handle = tokio::spawn(async move {
character_polling_main_loop(character_deps).await;
});
let _ = tokio::join!(video_handle, character_handle, credits_handle);
}
info!("Shutting down pager worker...");
pager_for_shutdown.shutdown_worker();