1use std::collections::{HashMap, HashSet};
2use std::io::Cursor;
3use std::ops::ControlFlow;
4use std::path::{Component, Path, PathBuf};
5use std::sync::atomic::Ordering;
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8
9use anyhow::{anyhow, bail, Context};
10use image::codecs::jpeg::JpegEncoder;
11use image::{ImageEncoder, RgbImage};
12use mold_core::chain::{ChainRequest, ChainStage, TransitionMode};
13use mold_core::chain_job::{
14 effective_stage_seed, settled, ChainJobEvent, ChainJobManifest, ChainJobState, FinalizeRecord,
15 GcOutcome, JobDirLayout, RetakeAmendment, RetakeMode, RetakeRequest, StageState, STAGES_DIR,
16};
17use mold_core::{GenerateRequest, OutputFormat};
18use mold_db::chain_jobs::{self, ChainJobRow, ChainJobStageRow};
19use mold_db::{settings, MetadataDb};
20use mold_inference::audio::NativeAudioTrack;
21use mold_inference::chain::stitch::fade_boundary;
22use mold_inference::chain::{ChainTail, StageOutcome, StageProgressEvent};
23use mold_inference::ltx_video::video_enc;
24use sha2::{Digest, Sha256};
25
26use crate::gpu_pool::{ActiveGeneration, GpuPool, GpuWorker};
27use crate::gpu_worker;
28use crate::model_manager;
29use crate::queue::save_video_to_dir;
30use crate::state::QueueHandle;
31
32const EVENT_BUS_CAPACITY: usize = 256;
33const DEFAULT_FADE_FRAMES: u32 = 8;
34const AUDIO_SIDECAR_MAGIC: &[u8; 8] = b"MOLDPCM1";
35pub const EPHEMERAL_GRACE_SECS: u64 = 900;
36
37pub struct ChainJobRunnerHandle {
38 kick_tx: tokio::sync::mpsc::UnboundedSender<RunnerCmd>,
39 cancel: Arc<CancelRegistry>,
40 events: Arc<JobEventBus>,
41 job_locks: Arc<JobMutationLocks>,
42 claims: Arc<EphemeralClaims>,
43}
44
45pub struct CancelRegistry {
46 known: Mutex<HashSet<String>>,
47 cancelled: Mutex<HashSet<String>>,
48}
49
50pub struct EphemeralClaimGuard {
55 job_id: String,
56 claims: Arc<EphemeralClaims>,
57}
58
59#[derive(Default)]
60pub struct EphemeralClaims {
61 claimed: std::sync::Mutex<std::collections::HashSet<String>>,
62}
63
64pub struct JobEventBus {
65 senders: Mutex<HashMap<String, tokio::sync::broadcast::Sender<ChainJobEvent>>>,
66}
67
68pub struct JobMutationLocks {
75 locks: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
76}
77
78pub struct RunnerDeps {
79 pub db: Arc<Option<mold_db::MetadataDb>>,
80 pub jobs_root: PathBuf,
81 pub executor: Arc<dyn StageExecutor>,
82 pub queue_probe: Arc<dyn QueueProbe>,
83 pub events: Arc<JobEventBus>,
84 pub cancel: Arc<CancelRegistry>,
85 pub job_locks: Arc<JobMutationLocks>,
86 pub claims: Arc<EphemeralClaims>,
87 pub output_dir: Option<PathBuf>,
88}
89
90pub(crate) struct CreateJobParams {
91 pub id: String,
92 pub ephemeral: bool,
93 pub request: ChainRequest,
94}
95
96pub enum RunnerCmd {
97 Kick,
98 Gc {
99 reply: tokio::sync::oneshot::Sender<std::result::Result<GcOutcome, String>>,
100 },
101}
102
103pub trait StageExecutor: Send + Sync {
122 fn render_stage(
123 &self,
124 model: &str,
125 stage_req: &GenerateRequest,
126 carry: Option<&ChainTail>,
127 motion_tail_frames: u32,
128 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
129 ) -> anyhow::Result<StageRenderOutcome>;
130}
131
132pub enum StageRenderOutcome {
133 Done(StageOutcome),
134 Cancelled,
135}
136
137pub trait QueueProbe: Send + Sync {
138 fn small_jobs_waiting(&self) -> usize;
139}
140
141struct StageArtifactPaths {
142 segment_rel: String,
143 preview_written: bool,
144}
145
146impl CancelRegistry {
147 pub fn new() -> Self {
148 Self {
149 known: Mutex::new(HashSet::new()),
150 cancelled: Mutex::new(HashSet::new()),
151 }
152 }
153
154 fn register(&self, job_id: &str) {
155 self.known
156 .lock()
157 .unwrap_or_else(|poisoned| poisoned.into_inner())
158 .insert(job_id.to_string());
159 }
160
161 fn unregister(&self, job_id: &str) {
162 self.known
163 .lock()
164 .unwrap_or_else(|poisoned| poisoned.into_inner())
165 .remove(job_id);
166 self.cancelled
167 .lock()
168 .unwrap_or_else(|poisoned| poisoned.into_inner())
169 .remove(job_id);
170 }
171
172 fn request(&self, job_id: &str) -> bool {
173 let known = self
174 .known
175 .lock()
176 .unwrap_or_else(|poisoned| poisoned.into_inner())
177 .contains(job_id);
178 if known {
179 self.cancelled
180 .lock()
181 .unwrap_or_else(|poisoned| poisoned.into_inner())
182 .insert(job_id.to_string());
183 true
184 } else {
185 false
186 }
187 }
188
189 fn is_cancelled(&self, job_id: &str) -> bool {
190 self.cancelled
191 .lock()
192 .unwrap_or_else(|poisoned| poisoned.into_inner())
193 .contains(job_id)
194 }
195}
196
197impl Default for CancelRegistry {
198 fn default() -> Self {
199 Self::new()
200 }
201}
202
203impl EphemeralClaims {
204 pub fn claim(self: &Arc<Self>, job_id: &str) -> EphemeralClaimGuard {
207 self.claimed
208 .lock()
209 .unwrap_or_else(|poisoned| poisoned.into_inner())
210 .insert(job_id.to_string());
211 EphemeralClaimGuard {
212 job_id: job_id.to_string(),
213 claims: self.clone(),
214 }
215 }
216
217 pub fn is_claimed(&self, job_id: &str) -> bool {
218 self.claimed
219 .lock()
220 .unwrap_or_else(|poisoned| poisoned.into_inner())
221 .contains(job_id)
222 }
223}
224
225impl Drop for EphemeralClaimGuard {
226 fn drop(&mut self) {
227 self.claims
228 .claimed
229 .lock()
230 .unwrap_or_else(|poisoned| poisoned.into_inner())
231 .remove(&self.job_id);
232 }
233}
234
235impl JobMutationLocks {
236 pub fn new() -> Self {
237 Self {
238 locks: Mutex::new(HashMap::new()),
239 }
240 }
241
242 fn mutex_for(&self, job_id: &str) -> Arc<tokio::sync::Mutex<()>> {
243 let mut locks = self
244 .locks
245 .lock()
246 .unwrap_or_else(|poisoned| poisoned.into_inner());
247 locks
248 .entry(job_id.to_string())
249 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
250 .clone()
251 }
252
253 pub async fn lock(&self, job_id: &str) -> tokio::sync::OwnedMutexGuard<()> {
254 self.mutex_for(job_id).lock_owned().await
255 }
256
257 pub fn blocking_lock(&self, job_id: &str) -> tokio::sync::OwnedMutexGuard<()> {
258 self.mutex_for(job_id).blocking_lock_owned()
259 }
260
261 pub fn remove(&self, job_id: &str) {
262 self.locks
263 .lock()
264 .unwrap_or_else(|poisoned| poisoned.into_inner())
265 .remove(job_id);
266 }
267}
268
269impl Default for JobMutationLocks {
270 fn default() -> Self {
271 Self::new()
272 }
273}
274
275impl JobEventBus {
276 pub fn new() -> Self {
277 Self {
278 senders: Mutex::new(HashMap::new()),
279 }
280 }
281
282 pub fn subscribe_for_job(
283 &self,
284 db: &MetadataDb,
285 job_id: &str,
286 ) -> anyhow::Result<tokio::sync::broadcast::Receiver<ChainJobEvent>> {
287 let mut senders = self
288 .senders
289 .lock()
290 .unwrap_or_else(|poisoned| poisoned.into_inner());
291 let state = chain_jobs::get_job(db, job_id)?.map(|row| row.state);
292 if state.is_none_or(settled) {
293 let (_tx, rx) = tokio::sync::broadcast::channel(EVENT_BUS_CAPACITY);
294 return Ok(rx);
295 }
296 Ok(senders
297 .entry(job_id.to_string())
298 .or_insert_with(|| tokio::sync::broadcast::channel(EVENT_BUS_CAPACITY).0)
299 .subscribe())
300 }
301
302 #[cfg(test)]
303 pub(crate) fn subscribe_persistent_for_tests(
304 &self,
305 job_id: &str,
306 ) -> tokio::sync::broadcast::Receiver<ChainJobEvent> {
307 let mut senders = self
308 .senders
309 .lock()
310 .unwrap_or_else(|poisoned| poisoned.into_inner());
311 senders
312 .entry(job_id.to_string())
313 .or_insert_with(|| tokio::sync::broadcast::channel(EVENT_BUS_CAPACITY).0)
314 .subscribe()
315 }
316
317 pub fn publish(&self, job_id: &str, event: ChainJobEvent) {
318 let sender = {
319 let senders = self
320 .senders
321 .lock()
322 .unwrap_or_else(|poisoned| poisoned.into_inner());
323 senders.get(job_id).cloned()
324 };
325 if let Some(sender) = sender {
326 let _ = sender.send(event);
327 }
328 }
329
330 pub fn publish_then_remove(&self, job_id: &str, event: ChainJobEvent) {
331 let mut senders = self
332 .senders
333 .lock()
334 .unwrap_or_else(|poisoned| poisoned.into_inner());
335 if let Some(sender) = senders.get(job_id) {
336 let _ = sender.send(event);
337 }
338 senders.remove(job_id);
339 }
340
341 pub fn remove(&self, job_id: &str) {
342 self.senders
343 .lock()
344 .unwrap_or_else(|poisoned| poisoned.into_inner())
345 .remove(job_id);
346 }
347
348 #[cfg(test)]
349 pub(crate) fn contains_for_tests(&self, job_id: &str) -> bool {
350 self.senders
351 .lock()
352 .unwrap_or_else(|poisoned| poisoned.into_inner())
353 .contains_key(job_id)
354 }
355}
356
357impl Default for JobEventBus {
358 fn default() -> Self {
359 Self::new()
360 }
361}
362
363impl ChainJobRunnerHandle {
364 #[cfg(test)]
365 pub(crate) fn inert_for_tests() -> Self {
366 let (kick_tx, _kick_rx) = tokio::sync::mpsc::unbounded_channel();
367 Self {
368 kick_tx,
369 cancel: Arc::new(CancelRegistry::new()),
370 events: Arc::new(JobEventBus::new()),
371 job_locks: Arc::new(JobMutationLocks::new()),
372 claims: Arc::new(EphemeralClaims::default()),
373 }
374 }
375
376 pub fn kick(&self) {
378 let _ = self.kick_tx.send(RunnerCmd::Kick);
379 }
380
381 pub fn request_cancel(&self, job_id: &str) -> bool {
383 self.cancel.request(job_id)
384 }
385
386 pub fn unregister_cancel(&self, job_id: &str) {
387 self.cancel.unregister(job_id);
388 }
389
390 pub async fn lock_job(&self, job_id: &str) -> tokio::sync::OwnedMutexGuard<()> {
391 self.job_locks.lock(job_id).await
392 }
393
394 pub(crate) fn blocking_lock_job(&self, job_id: &str) -> tokio::sync::OwnedMutexGuard<()> {
395 self.job_locks.blocking_lock(job_id)
396 }
397
398 pub fn remove_job_lock(&self, job_id: &str) {
399 self.job_locks.remove(job_id);
400 }
401
402 pub(crate) fn claim_ephemeral(&self, job_id: &str) -> EphemeralClaimGuard {
403 self.claims.claim(job_id)
404 }
405
406 pub async fn request_gc(&self) -> anyhow::Result<GcOutcome> {
407 let (reply, rx) = tokio::sync::oneshot::channel();
408 self.kick_tx
409 .send(RunnerCmd::Gc { reply })
410 .map_err(|_| anyhow!("chain job runner stopped before GC request could be sent"))?;
411 rx.await
412 .map_err(|_| anyhow!("chain job runner stopped before replying to GC request"))?
413 .map_err(|msg| anyhow!(msg))
414 }
415
416 pub fn cleanup_deleted(&self, job_id: &str) {
417 self.cancel.unregister(job_id);
418 self.events.remove(job_id);
419 }
420
421 pub fn publish_settled_state(&self, job_id: &str, state: ChainJobState, error: Option<String>) {
422 self.events
423 .publish_then_remove(job_id, ChainJobEvent::StateChanged { state, error });
424 self.cancel.unregister(job_id);
425 }
426
427 pub fn subscribe(
431 &self,
432 db: &MetadataDb,
433 job_id: &str,
434 ) -> anyhow::Result<tokio::sync::broadcast::Receiver<ChainJobEvent>> {
435 self.events.subscribe_for_job(db, job_id)
436 }
437
438 #[cfg(test)]
439 pub(crate) fn events_for_tests(&self) -> &JobEventBus {
440 &self.events
441 }
442}
443
444pub fn spawn_runner(deps: RunnerDeps) -> ChainJobRunnerHandle {
447 let (kick_tx, kick_rx) = tokio::sync::mpsc::unbounded_channel();
448 let cancel = deps.cancel.clone();
449 let events = deps.events.clone();
450 let job_locks = deps.job_locks.clone();
451 let claims = deps.claims.clone();
452 let deps = Arc::new(deps);
453 tokio::spawn(run_loop(deps, kick_rx));
454 ChainJobRunnerHandle {
455 kick_tx,
456 cancel,
457 events,
458 job_locks,
459 claims,
460 }
461}
462
463pub(crate) fn create_job_with_params(
471 db: &MetadataDb,
472 jobs_root: &Path,
473 params: CreateJobParams,
474) -> anyhow::Result<ChainJobRow> {
475 std::fs::create_dir_all(jobs_root)
476 .with_context(|| format!("creating chain jobs root '{}'", jobs_root.display()))?;
477 let job_dir = jobs_root.join(¶ms.id);
478 let layout = JobDirLayout::new(job_dir.clone());
479 layout.ensure_root()?;
480
481 let now = now_ms_i64();
482 let mut manifest = ChainJobManifest::new(params.id.clone(), now.max(0) as u64, ¶ms.request)
483 .map_err(|e| anyhow!("{e:#}"))?;
484 manifest.ephemeral = params.ephemeral;
485 manifest
486 .write_atomic(&job_dir)
487 .map_err(|e| anyhow!("{e:#}"))?;
488 let request_json = serde_json::to_string(¶ms.request)?;
489 let row = ChainJobRow {
490 id: params.id.clone(),
491 state: ChainJobState::Queued,
492 model: params.request.model.clone(),
493 request_json,
494 job_dir,
495 stage_count: params.request.stages.len() as u32,
496 current_stage: 0,
497 error: None,
498 created_at_ms: now,
499 updated_at_ms: now,
500 finalized_at_ms: None,
501 };
502 chain_jobs::insert_job(db, &row)?;
503 for stage in &manifest.stage_status {
504 chain_jobs::upsert_stage(
505 db,
506 &ChainJobStageRow {
507 job_id: row.id.clone(),
508 stage_idx: stage.idx,
509 state: stage.state,
510 seed: stage.seed,
511 frames_emitted: None,
512 generation_time_ms: None,
513 segment_rel_path: None,
514 error: None,
515 updated_at_ms: now,
516 },
517 )?;
518 }
519 Ok(row)
520}
521
522pub fn startup_reconcile(db: &MetadataDb, jobs_root: &Path) -> anyhow::Result<(usize, usize)> {
529 let now = now_ms_i64();
530 let running = chain_jobs::jobs_in_state(db, ChainJobState::Running)?;
531 let mut flipped = 0;
532 for row in running {
533 if chain_jobs::update_job_state(
534 db,
535 &row.id,
536 ChainJobState::Interrupted,
537 Some("server restarted while chain job was running"),
538 now,
539 )? {
540 flipped += 1;
541 }
542 }
543
544 let mut repaired = 0;
545 for row in chain_jobs::list_jobs(db)? {
546 if row.state.is_terminal() {
547 continue;
548 }
549 let job_dir = if row.job_dir.is_absolute() {
550 row.job_dir.clone()
551 } else {
552 jobs_root.join(&row.job_dir)
553 };
554 let manifest = match ChainJobManifest::read_from_dir(&job_dir) {
555 Ok(manifest) => manifest,
556 Err(err) => {
557 tracing::warn!(job_id = %row.id, "chain job manifest missing/unreadable during reconcile: {err:#}");
558 continue;
559 }
560 };
561
562 let mut changed = false;
563 for stage in &manifest.stage_status {
564 let db_stage = ChainJobStageRow {
565 job_id: row.id.clone(),
566 stage_idx: stage.idx,
567 state: stage.state,
568 seed: stage.seed,
569 frames_emitted: stage.frames_emitted,
570 generation_time_ms: stage.generation_time_ms,
571 segment_rel_path: stage.segment.clone(),
572 error: stage.error.clone(),
573 updated_at_ms: now,
574 };
575 chain_jobs::upsert_stage(db, &db_stage)?;
576 }
577
578 let (state, current_stage, mut error) = manifest_index_state(&manifest, row.state);
579 if state == ChainJobState::Interrupted && error.is_none() {
580 error = row.error.clone();
581 }
582 let finalized_at = row.finalized_at_ms.or_else(|| {
583 manifest
584 .finalizes
585 .last()
586 .and_then(|record| i64::try_from(record.at_unix_ms).ok())
587 });
588 if row.state != state
589 || row.current_stage != current_stage
590 || row.error.as_deref() != error.as_deref()
591 || row.finalized_at_ms != finalized_at
592 {
593 changed = chain_jobs::repair_job_from_manifest(
594 db,
595 &row.id,
596 state,
597 current_stage,
598 error.as_deref(),
599 now,
600 finalized_at,
601 )?;
602 }
603 if changed {
604 repaired += 1;
605 }
606 }
607
608 Ok((flipped, repaired))
609}
610
611pub(crate) fn run_gc_pass(
619 deps: &RunnerDeps,
620 ttl_days: i64,
621 now_ms: i64,
622) -> anyhow::Result<GcOutcome> {
623 let db = deps
624 .db
625 .as_ref()
626 .as_ref()
627 .ok_or_else(|| anyhow!("chain job GC invoked without metadata DB"))?;
628 let ttl_ms = ttl_days.max(0).saturating_mul(86_400_000);
629 let grace_ms = (EPHEMERAL_GRACE_SECS as i64).saturating_mul(1_000);
630 let mut outcome = GcOutcome {
631 swept_ephemeral_jobs: 0,
632 pruned_artifact_dirs: 0,
633 };
634
635 for row in chain_jobs::list_jobs(db)? {
636 let manifest = ChainJobManifest::read_from_dir(&row.job_dir).ok();
637 if manifest.as_ref().is_some_and(|manifest| manifest.ephemeral) {
638 if deps.claims.is_claimed(&row.id) {
639 continue;
640 }
641 if settled(row.state) && now_ms.saturating_sub(row.updated_at_ms) < grace_ms {
642 continue;
643 }
644 let remove_lock = {
645 let _guard = deps.job_locks.blocking_lock(&row.id);
646 let current = match chain_jobs::get_job(db, &row.id)? {
647 Some(current) => current,
648 None => {
649 continue;
650 }
651 };
652 let within_grace = settled(current.state)
653 && now_ms.saturating_sub(current.updated_at_ms) < grace_ms;
654 if current.state == ChainJobState::Running
655 || deps.claims.is_claimed(&row.id)
656 || within_grace
657 {
658 false
659 } else {
660 if current.job_dir.exists() {
661 std::fs::remove_dir_all(¤t.job_dir).with_context(|| {
662 format!(
663 "removing ephemeral chain job '{}'",
664 current.job_dir.display()
665 )
666 })?;
667 }
668 chain_jobs::delete_job_not_running(db, ¤t.id)?
669 }
670 };
671 if remove_lock {
672 outcome.swept_ephemeral_jobs += 1;
673 deps.cancel.unregister(&row.id);
674 deps.events.remove(&row.id);
675 deps.job_locks.remove(&row.id);
676 }
677 continue;
678 }
679
680 if row.state == ChainJobState::Completed
681 && ttl_ms >= 0
682 && now_ms.saturating_sub(row.updated_at_ms) >= ttl_ms
683 {
684 let stages_dir = row.job_dir.join(STAGES_DIR);
685 if stages_dir.exists() {
686 let _guard = deps.job_locks.blocking_lock(&row.id);
687 let Some(current) = chain_jobs::get_job(db, &row.id)? else {
688 continue;
689 };
690 if current.state == ChainJobState::Completed && stages_dir.exists() {
691 std::fs::remove_dir_all(&stages_dir).with_context(|| {
692 format!("pruning chain job stages '{}'", stages_dir.display())
693 })?;
694 outcome.pruned_artifact_dirs += 1;
695 }
696 }
697 }
698 }
699
700 Ok(outcome)
701}
702
703pub(crate) fn startup_gc_sweep(db: &MetadataDb, jobs_root: &Path) -> anyhow::Result<GcOutcome> {
707 let mut outcome = GcOutcome {
708 swept_ephemeral_jobs: 0,
709 pruned_artifact_dirs: 0,
710 };
711 for row in chain_jobs::list_jobs(db)? {
712 let job_dir = if row.job_dir.is_absolute() {
713 row.job_dir.clone()
714 } else {
715 jobs_root.join(&row.job_dir)
716 };
717 let Ok(manifest) = ChainJobManifest::read_from_dir(&job_dir) else {
718 continue;
719 };
720 if manifest.ephemeral && row.state != ChainJobState::Running {
721 if job_dir.exists() {
722 std::fs::remove_dir_all(&job_dir).with_context(|| {
723 format!(
724 "removing startup ephemeral chain job '{}'",
725 job_dir.display()
726 )
727 })?;
728 }
729 if chain_jobs::delete_job_not_running(db, &row.id)? {
730 outcome.swept_ephemeral_jobs += 1;
731 }
732 }
733 }
734 Ok(outcome)
735}
736
737async fn run_loop(
738 deps: Arc<RunnerDeps>,
739 mut kick_rx: tokio::sync::mpsc::UnboundedReceiver<RunnerCmd>,
740) {
741 let mut daily = tokio::time::interval(Duration::from_secs(24 * 60 * 60));
742 daily.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
743 daily.tick().await;
744 loop {
745 while let Ok(cmd) = kick_rx.try_recv() {
746 if !handle_runner_cmd(deps.clone(), cmd).await {
747 return;
748 }
749 }
750
751 let mut ran_job = false;
752 if let Some(db) = deps.db.as_ref() {
753 let job = match next_queued_job(db) {
754 Ok(Some(job)) => Some(job),
755 Ok(None) => None,
756 Err(err) => {
757 tracing::warn!("chain-job queued lookup failed: {err:#}");
758 None
759 }
760 };
761 if let Some(job) = job {
762 let deps_for_job = deps.clone();
763 let job_id = job.id.clone();
764 let start_stage = job.current_stage;
765 let join = tokio::task::spawn_blocking(move || -> anyhow::Result<bool> {
766 if !claim_for_execution(&deps_for_job, &job)? {
767 return Ok(false);
768 }
769 execute_job(&deps_for_job, &job, start_stage)?;
770 Ok(true)
771 })
772 .await;
773 match join {
774 Ok(Ok(claimed)) => {
775 if claimed {
776 ran_job = true;
777 }
778 }
779 Ok(Err(err)) => {
780 tracing::warn!(job_id = %job_id, "chain job execution failed: {err:#}");
781 break;
782 }
783 Err(err) => {
784 tracing::warn!(job_id = %job_id, "chain job task join failed: {err}");
785 break;
786 }
787 }
788 }
789 }
790
791 if ran_job {
792 tokio::task::yield_now().await;
793 continue;
794 }
795
796 tokio::select! {
797 maybe_cmd = kick_rx.recv() => {
798 let Some(cmd) = maybe_cmd else { break };
799 if !handle_runner_cmd(deps.clone(), cmd).await {
800 break;
801 }
802 }
803 _ = daily.tick() => {
804 if let Err(err) = run_gc_for_runner(deps.clone()).await {
805 tracing::warn!("daily chain job GC failed: {err}");
806 }
807 }
808 }
809 }
810}
811
812async fn handle_runner_cmd(deps: Arc<RunnerDeps>, cmd: RunnerCmd) -> bool {
813 match cmd {
814 RunnerCmd::Kick => true,
815 RunnerCmd::Gc { reply } => {
816 let result = run_gc_for_runner(deps).await;
817 let _ = reply.send(result);
818 true
819 }
820 }
821}
822
823async fn run_gc_for_runner(deps: Arc<RunnerDeps>) -> std::result::Result<GcOutcome, String> {
824 tokio::task::spawn_blocking(move || {
825 let db = deps
826 .db
827 .as_ref()
828 .as_ref()
829 .ok_or_else(|| "chain job GC invoked without metadata DB".to_string())?;
830 let ttl_days = settings::Settings::new(db)
831 .get_int(settings::CHAIN_JOBS_ARTIFACT_TTL_DAYS)
832 .map_err(|e| format!("{e:#}"))?
833 .unwrap_or(settings::CHAIN_JOBS_ARTIFACT_TTL_DEFAULT);
834 run_gc_pass(&deps, ttl_days, now_ms_i64()).map_err(|e| format!("{e:#}"))
835 })
836 .await
837 .map_err(|e| format!("chain job GC task failed: {e}"))?
838}
839
840fn next_queued_job(db: &MetadataDb) -> anyhow::Result<Option<ChainJobRow>> {
842 chain_jobs::next_queued_job(db)
843}
844
845fn claim_for_execution(deps: &RunnerDeps, job: &ChainJobRow) -> anyhow::Result<bool> {
846 let db = deps
847 .db
848 .as_ref()
849 .as_ref()
850 .ok_or_else(|| anyhow!("chain job runner invoked without metadata DB"))?;
851 let _guard = deps.job_locks.blocking_lock(&job.id);
852 let claimed = chain_jobs::claim_job(db, &job.id)?;
853 if claimed {
854 deps.cancel.register(&job.id);
855 }
856 Ok(claimed)
857}
858
859fn execute_job(deps: &RunnerDeps, job: &ChainJobRow, start_stage: u32) -> anyhow::Result<()> {
860 let db = deps
861 .db
862 .as_ref()
863 .as_ref()
864 .ok_or_else(|| anyhow!("chain job runner invoked without metadata DB"))?;
865 let mut terminal = false;
866
867 let current = chain_jobs::get_job(db, &job.id)?
868 .ok_or_else(|| anyhow!("chain job '{}' disappeared before execution", job.id))?;
869 match current.state {
870 ChainJobState::Queued => {
871 if !claim_for_execution(deps, ¤t)? {
872 return Ok(());
873 }
874 }
875 ChainJobState::Running => {
876 deps.cancel.register(&job.id);
877 }
878 other => bail!(
879 "chain job '{}' is not executable from state {}",
880 job.id,
881 other.as_str()
882 ),
883 }
884 let run_result = (|| -> anyhow::Result<()> {
885 let mut manifest = {
886 let _guard = deps.job_locks.blocking_lock(&job.id);
887 match ChainJobManifest::read_from_dir(&job.job_dir) {
888 Ok(manifest) => manifest,
889 Err(err) => {
890 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
891 terminal = true;
892 return Ok(());
893 }
894 }
895 };
896 let layout = JobDirLayout::new(job.job_dir.clone());
897 let mut effective = {
898 let _guard = deps.job_locks.blocking_lock(&job.id);
899 match effective_request(&manifest) {
900 Ok(effective) => effective,
901 Err(err) => {
902 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
903 terminal = true;
904 return Ok(());
905 }
906 }
907 };
908 let mut carry: Option<ChainTail> = None;
909
910 deps.events.publish(
911 &job.id,
912 ChainJobEvent::StateChanged {
913 state: ChainJobState::Running,
914 error: None,
915 },
916 );
917
918 let mut stage_idx = first_incomplete_stage(&manifest).unwrap_or(start_stage);
919 if stage_idx < start_stage {
920 stage_idx = start_stage;
921 }
922
923 while stage_idx < manifest.stage_status.len() as u32 {
924 if deps.cancel.is_cancelled(&job.id) {
925 if let Err(err) = set_cancelled(db, deps, &job.id) {
926 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
927 }
928 terminal = true;
929 return Ok(());
930 }
931
932 if manifest.stage_status[stage_idx as usize].state == StageState::Completed {
933 match carry_after_skipping_completed_stage(&layout, &manifest, stage_idx) {
934 Ok(next_carry) => {
935 carry = next_carry;
936 stage_idx += 1;
937 continue;
938 }
939 Err(err) => {
940 let fail_stage = (stage_idx + 1)
941 .min(manifest.stage_status.len().saturating_sub(1) as u32);
942 fail_stage_job(
943 db,
944 deps,
945 &mut manifest,
946 &layout,
947 &job.id,
948 fail_stage,
949 format!("{err:#}"),
950 )?;
951 terminal = true;
952 return Ok(());
953 }
954 }
955 }
956
957 let stage_dir = layout.stage_dir(stage_idx);
958 if stage_dir.exists() {
959 if let Err(err) = std::fs::remove_dir_all(&stage_dir).with_context(|| {
960 format!(
961 "removing partial chain stage directory '{}'",
962 stage_dir.display()
963 )
964 }) {
965 fail_stage_job(
966 db,
967 deps,
968 &mut manifest,
969 &layout,
970 &job.id,
971 stage_idx,
972 format!("{err:#}"),
973 )?;
974 terminal = true;
975 return Ok(());
976 }
977 }
978
979 if carry.is_none() {
980 match resume_carry_from_disk(layout.root(), &manifest, stage_idx) {
981 Ok(next_carry) => carry = next_carry,
982 Err(err) => {
983 fail_stage_job(
984 db,
985 deps,
986 &mut manifest,
987 &layout,
988 &job.id,
989 stage_idx,
990 format!("{err:#}"),
991 )?;
992 terminal = true;
993 return Ok(());
994 }
995 }
996 }
997
998 deps.events
999 .publish(&job.id, ChainJobEvent::StageStart { stage_idx });
1000 if let Err(err) = mark_stage_running(db, &job.id, &manifest, stage_idx) {
1001 fail_stage_job(
1002 db,
1003 deps,
1004 &mut manifest,
1005 &layout,
1006 &job.id,
1007 stage_idx,
1008 format!("{err:#}"),
1009 )?;
1010 terminal = true;
1011 return Ok(());
1012 }
1013
1014 let stage_carry = match effective.stages[stage_idx as usize].transition {
1015 TransitionMode::Smooth => carry.as_ref(),
1016 TransitionMode::Cut | TransitionMode::Fade => None,
1017 };
1018 let render_outcome = match execute_stage(deps, job, &manifest, stage_idx, stage_carry) {
1019 Ok(outcome) => outcome,
1020 Err(err) => {
1021 let error = format!("{err:#}");
1022 {
1023 let _guard = deps.job_locks.blocking_lock(&job.id);
1024 let _ =
1025 mark_manifest_stage_failed(&mut manifest, &layout, stage_idx, &error);
1026 }
1027 fail_job(db, deps, &job.id, Some(stage_idx), error)?;
1028 terminal = true;
1029 return Ok(());
1030 }
1031 };
1032 match render_outcome {
1033 StageRenderOutcome::Cancelled => {
1034 set_cancelled(db, deps, &job.id)?;
1035 terminal = true;
1036 return Ok(());
1037 }
1038 StageRenderOutcome::Done(outcome) => {
1039 let stage_artifacts = {
1040 let _guard = deps.job_locks.blocking_lock(&job.id);
1041 write_stage_artifacts(
1042 &layout,
1043 &mut manifest,
1044 stage_idx,
1045 &outcome,
1046 &effective,
1047 )
1048 };
1049 let paths = match stage_artifacts {
1050 Ok(paths) => paths,
1051 Err(err) => {
1052 let error = format!("{err:#}");
1053 {
1054 let _guard = deps.job_locks.blocking_lock(&job.id);
1055 let _ = mark_manifest_stage_failed(
1056 &mut manifest,
1057 &layout,
1058 stage_idx,
1059 &error,
1060 );
1061 }
1062 fail_job(db, deps, &job.id, Some(stage_idx), error)?;
1063 terminal = true;
1064 return Ok(());
1065 }
1066 };
1067
1068 let now = now_ms_i64();
1069 let status = manifest.stage_status[stage_idx as usize].clone();
1070 if let Err(err) = chain_jobs::upsert_stage(
1071 db,
1072 &ChainJobStageRow {
1073 job_id: job.id.clone(),
1074 stage_idx,
1075 state: StageState::Completed,
1076 seed: status.seed,
1077 frames_emitted: status.frames_emitted,
1078 generation_time_ms: status.generation_time_ms,
1079 segment_rel_path: Some(paths.segment_rel.clone()),
1080 error: None,
1081 updated_at_ms: now,
1082 },
1083 ) {
1084 fail_stage_job(
1085 db,
1086 deps,
1087 &mut manifest,
1088 &layout,
1089 &job.id,
1090 stage_idx,
1091 format!("{err:#}"),
1092 )?;
1093 terminal = true;
1094 return Ok(());
1095 }
1096 if let Err(err) = chain_jobs::set_current_stage(db, &job.id, stage_idx + 1, now)
1097 {
1098 fail_stage_job(
1099 db,
1100 deps,
1101 &mut manifest,
1102 &layout,
1103 &job.id,
1104 stage_idx,
1105 format!("{err:#}"),
1106 )?;
1107 terminal = true;
1108 return Ok(());
1109 }
1110
1111 if let Err(err) =
1112 maybe_reencode_next_after_fade(&layout, &manifest, stage_idx, &effective)
1113 {
1114 let fail_stage = (stage_idx + 1)
1115 .min(manifest.stage_status.len().saturating_sub(1) as u32);
1116 fail_stage_job(
1117 db,
1118 deps,
1119 &mut manifest,
1120 &layout,
1121 &job.id,
1122 fail_stage,
1123 format!("{err:#}"),
1124 )?;
1125 terminal = true;
1126 return Ok(());
1127 }
1128
1129 deps.events.publish(
1130 &job.id,
1131 ChainJobEvent::StageDone {
1132 stage_idx,
1133 frames_emitted: status.frames_emitted.unwrap_or(0),
1134 has_preview: paths.preview_written,
1135 },
1136 );
1137
1138 carry = Some(outcome.tail);
1139
1140 publish_yield_if_contended(deps, &job.id);
1141 }
1142 }
1143 stage_idx += 1;
1144 effective = match effective_request(&manifest) {
1145 Ok(effective) => effective,
1146 Err(err) => {
1147 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
1148 terminal = true;
1149 return Ok(());
1150 }
1151 };
1152 }
1153
1154 let output = match finalize_job(deps, job, &mut manifest) {
1155 Ok(output) => output,
1156 Err(err) => {
1157 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
1158 terminal = true;
1159 return Ok(());
1160 }
1161 };
1162 let take = manifest.finalizes.len() as u32;
1163 deps.events
1164 .publish(&job.id, ChainJobEvent::Finalized { output, take });
1165 let completed = match chain_jobs::try_transition(
1166 db,
1167 &job.id,
1168 &[ChainJobState::Running],
1169 ChainJobState::Completed,
1170 None,
1171 now_ms_i64(),
1172 ) {
1173 Ok(completed) => completed,
1174 Err(err) => {
1175 fail_job(db, deps, &job.id, None, format!("{err:#}"))?;
1176 terminal = true;
1177 return Ok(());
1178 }
1179 };
1180 if !completed {
1181 tracing::warn!(job_id = %job.id, "chain job completed but running->completed CAS lost");
1182 terminal = true;
1183 return Ok(());
1184 }
1185 deps.events.publish_then_remove(
1186 &job.id,
1187 ChainJobEvent::StateChanged {
1188 state: ChainJobState::Completed,
1189 error: None,
1190 },
1191 );
1192 deps.cancel.unregister(&job.id);
1193 terminal = true;
1194 Ok(())
1195 })();
1196
1197 deps.cancel.unregister(&job.id);
1198 if terminal {
1199 deps.events.remove(&job.id);
1200 }
1201 run_result
1202}
1203
1204fn execute_stage(
1205 deps: &RunnerDeps,
1206 job: &ChainJobRow,
1207 manifest: &ChainJobManifest,
1208 stage_idx: u32,
1209 carry: Option<&ChainTail>,
1210) -> anyhow::Result<StageRenderOutcome> {
1211 let effective = effective_request(manifest)?;
1212 let stage = effective
1213 .stages
1214 .get(stage_idx as usize)
1215 .ok_or_else(|| anyhow!("stage index {stage_idx} out of bounds"))?;
1216 let stage_seed = manifest.stage_status[stage_idx as usize].seed;
1217 let stage_req = build_stage_generate_request(stage, &effective, stage_seed, stage_idx as usize);
1218 let executor = deps.executor.clone();
1219 let model = effective.model.clone();
1220 let carry_owned = carry.cloned();
1221 let job_id = job.id.clone();
1222 let events = deps.events.clone();
1223 let cancel = deps.cancel.clone();
1224 let motion_tail_frames = effective.motion_tail_frames;
1225 let progress: Box<dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync> =
1226 Box::new(move |step, total| {
1227 events.publish(
1228 &job_id,
1229 ChainJobEvent::DenoiseStep {
1230 stage_idx,
1231 step,
1232 total,
1233 },
1234 );
1235 if cancel.is_cancelled(&job_id) {
1236 ControlFlow::Break(())
1237 } else {
1238 ControlFlow::Continue(())
1239 }
1240 });
1241
1242 match tokio::runtime::Handle::try_current() {
1243 Ok(handle) => handle
1244 .block_on(tokio::task::spawn_blocking(move || {
1245 executor.render_stage(
1246 &model,
1247 &stage_req,
1248 carry_owned.as_ref(),
1249 motion_tail_frames,
1250 progress.as_ref(),
1251 )
1252 }))
1253 .map_err(|e| anyhow!("chain stage task failed: {e}"))?,
1254 Err(_) => deps.executor.render_stage(
1255 &model,
1256 &stage_req,
1257 carry,
1258 motion_tail_frames,
1259 progress.as_ref(),
1260 ),
1261 }
1262}
1263
1264fn write_stage_artifacts(
1265 layout: &JobDirLayout,
1266 manifest: &mut ChainJobManifest,
1267 stage_idx: u32,
1268 outcome: &StageOutcome,
1269 effective: &ChainRequest,
1270) -> anyhow::Result<StageArtifactPaths> {
1271 layout.ensure_stage_dirs(stage_idx)?;
1272 let stage = effective
1273 .stages
1274 .get(stage_idx as usize)
1275 .ok_or_else(|| anyhow!("stage index {stage_idx} out of bounds"))?;
1276 let next_transition = effective
1277 .stages
1278 .get(stage_idx as usize + 1)
1279 .map(|s| s.transition);
1280 let incoming = stage.transition;
1281 let mut frames = outcome.frames.clone();
1282 let mut audio = outcome.audio.clone();
1283
1284 if matches!(incoming, TransitionMode::Smooth) && stage_idx > 0 {
1285 let drop = effective.motion_tail_frames as usize;
1286 if frames.len() < drop {
1287 bail!(
1288 "stage {stage_idx} emitted {} frames, cannot drop {drop} smooth carry frames",
1289 frames.len()
1290 );
1291 }
1292 frames.drain(0..drop);
1293 trim_audio_front(&mut audio, effective.motion_tail_frames, effective.fps)?;
1294 }
1295
1296 if matches!(incoming, TransitionMode::Fade) && stage_idx > 0 {
1297 let fade_len = stage.fade_frames.unwrap_or(DEFAULT_FADE_FRAMES);
1298 let n = fade_len as usize;
1299 if frames.len() < n {
1300 bail!(
1301 "stage {stage_idx} emitted {} frames, cannot apply incoming fade_len {n}",
1302 frames.len()
1303 );
1304 }
1305 write_frames_to_dir(&layout.boundary_in_dir(stage_idx), &frames[..n])?;
1306 let previous_out = read_frames_from_dir(&layout.boundary_out_dir(stage_idx - 1), n)?;
1307 let blended = fade_boundary(&previous_out, &frames, fade_len);
1308 for (idx, frame) in blended.into_iter().enumerate() {
1309 frames[idx] = frame;
1310 }
1311 blend_audio_front_from_previous_boundary(
1312 layout,
1313 stage_idx,
1314 &mut audio,
1315 fade_len,
1316 effective.fps,
1317 )?;
1318 }
1319
1320 if matches!(next_transition, Some(TransitionMode::Fade)) {
1321 let next = &effective.stages[stage_idx as usize + 1];
1322 let fade_len = next.fade_frames.unwrap_or(DEFAULT_FADE_FRAMES);
1323 let n = fade_len as usize;
1324 if frames.len() < n {
1325 bail!(
1326 "stage {stage_idx} emitted {} frames, cannot reserve outgoing fade_len {n}",
1327 frames.len()
1328 );
1329 }
1330 write_frames_to_dir(
1331 &layout.boundary_out_dir(stage_idx),
1332 &frames[frames.len() - n..],
1333 )?;
1334 write_audio_boundary_out(layout, stage_idx, &audio, fade_len, effective.fps)?;
1335 frames.truncate(frames.len() - n);
1336 trim_audio_back(&mut audio, fade_len, effective.fps)?;
1337 }
1338
1339 if frames.is_empty() {
1340 bail!("stage {stage_idx} produced no frames after boundary handling");
1341 }
1342
1343 let segment_bytes = video_enc::encode_mp4(&frames, effective.fps)
1344 .with_context(|| format!("encoding chain stage {stage_idx} segment"))?;
1345 write_file(&layout.segment_path(stage_idx), &segment_bytes)?;
1346
1347 let tail_frames = if matches!(next_transition, Some(TransitionMode::Smooth)) {
1348 write_frames_to_dir(&layout.tail_dir(stage_idx), &outcome.tail.tail_rgb_frames)?;
1349 outcome.tail.frames
1350 } else {
1351 0
1352 };
1353
1354 let audio_rel = if let Some(track) = audio.as_ref() {
1355 write_audio_sidecar(&layout.audio_path(stage_idx), track)?;
1356 Some(layout.audio_rel(stage_idx))
1357 } else {
1358 None
1359 };
1360
1361 write_preview_jpeg(&layout.preview_path(stage_idx), frames.last().unwrap())?;
1362
1363 let status = manifest
1364 .stage_status
1365 .get_mut(stage_idx as usize)
1366 .ok_or_else(|| anyhow!("manifest missing status for stage {stage_idx}"))?;
1367 status.state = StageState::Completed;
1368 status.frames_emitted = Some(frames.len() as u32);
1369 status.generation_time_ms = Some(outcome.generation_time_ms);
1370 status.segment = Some(layout.segment_rel(stage_idx));
1371 status.tail_frames = Some(tail_frames);
1372 status.audio = audio_rel.clone();
1373 status.error = None;
1374
1375 manifest.write_atomic(layout.root())?;
1376
1377 Ok(StageArtifactPaths {
1378 segment_rel: layout.segment_rel(stage_idx),
1379 preview_written: true,
1380 })
1381}
1382
1383fn resume_carry_from_disk(
1384 job_dir: &Path,
1385 manifest: &ChainJobManifest,
1386 stage_idx: u32,
1387) -> anyhow::Result<Option<ChainTail>> {
1388 if stage_idx == 0 {
1389 return Ok(None);
1390 }
1391 let effective = effective_request(manifest)?;
1392 let Some(stage) = effective.stages.get(stage_idx as usize) else {
1393 bail!("stage index {stage_idx} out of bounds for resume carry");
1394 };
1395 if !matches!(stage.transition, TransitionMode::Smooth) {
1396 return Ok(None);
1397 }
1398 let prev_idx = stage_idx - 1;
1399 let prev_status = manifest
1400 .stage_status
1401 .get(prev_idx as usize)
1402 .ok_or_else(|| anyhow!("manifest missing previous stage {prev_idx}"))?;
1403 if let Some(segment) = &prev_status.segment {
1404 safe_join_manifest_rel(job_dir, segment)?;
1405 }
1406 let expected = prev_status.tail_frames.unwrap_or(0);
1407 if expected == 0 {
1408 return Ok(None);
1409 }
1410 let tail_dir = JobDirLayout::new(job_dir.to_path_buf()).tail_dir(prev_idx);
1411 let mut paths: Vec<PathBuf> = std::fs::read_dir(&tail_dir)
1412 .with_context(|| format!("reading chain tail directory '{}'", tail_dir.display()))?
1413 .map(|entry| entry.map(|entry| entry.path()))
1414 .collect::<std::io::Result<Vec<_>>>()?;
1415 paths.sort();
1416 if paths.len() != expected as usize {
1417 bail!(
1418 "chain tail for stage {prev_idx} has {} PNG(s), expected {expected}",
1419 paths.len()
1420 );
1421 }
1422 let mut frames = Vec::with_capacity(paths.len());
1423 for path in paths {
1424 frames.push(
1425 image::open(&path)
1426 .with_context(|| format!("decoding chain tail PNG '{}'", path.display()))?
1427 .to_rgb8(),
1428 );
1429 }
1430 Ok(Some(ChainTail {
1431 frames: expected,
1432 tail_rgb_frames: frames,
1433 }))
1434}
1435
1436fn finalize_job(
1440 deps: &RunnerDeps,
1441 job: &ChainJobRow,
1442 manifest: &mut ChainJobManifest,
1443) -> anyhow::Result<String> {
1444 let db = deps
1445 .db
1446 .as_ref()
1447 .as_ref()
1448 .ok_or_else(|| anyhow!("chain job finalizer invoked without metadata DB"))?;
1449 let effective = effective_request(manifest)?;
1450 let layout = JobDirLayout::new(job.job_dir.clone());
1451 let total_frames: u32 = manifest
1452 .stage_status
1453 .iter()
1454 .map(|stage| stage.frames_emitted.unwrap_or(0))
1455 .sum();
1456 deps.events
1457 .publish(&job.id, ChainJobEvent::Finalizing { total_frames });
1458
1459 let mut encoder: Option<video_enc::Mp4StreamEncoder> = None;
1460 let mut audio_samples = Vec::new();
1461 let mut audio_format: Option<(u32, u16)> = None;
1462 let mut frame_count = 0u32;
1463
1464 for stage in &manifest.stage_status {
1465 if stage.state != StageState::Completed {
1466 bail!(
1467 "cannot finalize chain job with non-completed stage {}",
1468 stage.idx
1469 );
1470 }
1471 let segment = stage
1472 .segment
1473 .as_ref()
1474 .ok_or_else(|| anyhow!("completed stage {} has no segment path", stage.idx))?;
1475 let segment_path = safe_join_manifest_rel(layout.root(), segment)?;
1476 let (metadata, frames) =
1477 mold_inference::ltx2::media::decode_video_frames_from_path(&segment_path)
1478 .with_context(|| format!("decoding chain segment '{}'", segment_path.display()))?;
1479 if encoder.is_none() {
1480 encoder = Some(video_enc::Mp4StreamEncoder::new(
1481 metadata.width,
1482 metadata.height,
1483 effective.fps,
1484 )?);
1485 }
1486 let enc = encoder.as_mut().unwrap();
1487 for frame in frames {
1488 enc.push(&frame)?;
1489 frame_count += 1;
1490 }
1491 if let Some(audio_rel) = stage.audio.as_ref() {
1492 let audio_path = safe_join_manifest_rel(layout.root(), audio_rel)?;
1493 let track = read_audio_sidecar(&audio_path)?;
1494 match audio_format {
1495 None => audio_format = Some((track.sample_rate, track.channels)),
1496 Some((sample_rate, channels))
1497 if sample_rate == track.sample_rate && channels == track.channels => {}
1498 Some((sample_rate, channels)) => {
1499 bail!(
1500 "stage {} audio format {} Hz/{} ch does not match previous {} Hz/{} ch",
1501 stage.idx,
1502 track.sample_rate,
1503 track.channels,
1504 sample_rate,
1505 channels
1506 );
1507 }
1508 }
1509 audio_samples.extend_from_slice(&track.interleaved_samples);
1510 }
1511 }
1512
1513 let encoder = encoder.ok_or_else(|| anyhow!("cannot finalize chain job with no frames"))?;
1514 let video_bytes = encoder.finish()?;
1515 let video_bytes = if !audio_samples.is_empty() {
1516 #[cfg(feature = "mp4")]
1517 {
1518 let (sample_rate, channels) = audio_format.expect("samples imply format");
1519 mold_inference::ltx2::media::attach_aac_track_to_mp4_bytes(
1520 &video_bytes,
1521 &audio_samples,
1522 sample_rate,
1523 channels,
1524 )?
1525 }
1526 #[cfg(not(feature = "mp4"))]
1527 {
1528 bail!("chain job finalization with audio requires the mp4 feature for AAC muxing");
1529 }
1530 } else {
1531 video_bytes
1532 };
1533
1534 let take = manifest.finalizes.len() as u32 + 1;
1535 let output_path = layout.final_output_path(take);
1536 if let Some(parent) = output_path.parent() {
1537 std::fs::create_dir_all(parent)
1538 .with_context(|| format!("creating chain final directory '{}'", parent.display()))?;
1539 }
1540 write_file(&output_path, &video_bytes)?;
1541
1542 if !manifest.ephemeral {
1543 if let Some(output_dir) = deps.output_dir.as_ref() {
1544 let metadata = effective.stitched_output_metadata(OutputFormat::Mp4, frame_count);
1545 save_video_to_dir(
1546 output_dir,
1547 &video_bytes,
1548 &[],
1549 OutputFormat::Mp4,
1550 &effective.model,
1551 &metadata,
1552 None,
1553 Some(db),
1554 );
1555 }
1556 }
1557
1558 let now = now_ms_u64();
1559 let output = format!("final/output-{take}.mp4");
1560 {
1561 let _guard = deps.job_locks.blocking_lock(&job.id);
1562 manifest.finalizes.push(FinalizeRecord {
1563 output: output.clone(),
1564 at_unix_ms: now,
1565 stage_seeds: manifest
1566 .stage_status
1567 .iter()
1568 .map(|stage| stage.seed)
1569 .collect(),
1570 });
1571 manifest.write_atomic(layout.root())?;
1572 chain_jobs::set_finalized_at(db, &job.id, i64::try_from(now).unwrap_or(i64::MAX))?;
1573 }
1574 Ok(output)
1575}
1576
1577pub fn apply_retake(
1578 db: &MetadataDb,
1579 jobs_root: &Path,
1580 job_id: &str,
1581 req: &RetakeRequest,
1582) -> anyhow::Result<ChainJobRow> {
1583 let job = chain_jobs::get_job(db, job_id)?.ok_or_else(|| anyhow!("chain job not found"))?;
1584 if job.state == ChainJobState::Running {
1585 bail!("CHAIN_JOB_RUNNING");
1586 }
1587 let allowed_from = [
1588 ChainJobState::Interrupted,
1589 ChainJobState::Failed,
1590 ChainJobState::Cancelled,
1591 ChainJobState::Completed,
1592 ];
1593 if !allowed_from.contains(&job.state) {
1594 bail!(
1595 "chain job is not retakeable from state {}",
1596 job.state.as_str()
1597 );
1598 }
1599 let job_dir = if job.job_dir.is_absolute() {
1600 job.job_dir.clone()
1601 } else {
1602 jobs_root.join(&job.job_dir)
1603 };
1604 let mut manifest = ChainJobManifest::read_from_dir(&job_dir)?;
1605 let effective = effective_request(&manifest)?;
1606 let stage_idx = req.stage_idx as usize;
1607 if stage_idx >= manifest.stage_status.len() {
1608 bail!("stage_idx {} out of bounds", req.stage_idx);
1609 }
1610 if req.mode == RetakeMode::Splice && stage_idx + 1 < effective.stages.len() {
1611 let next_transition = effective.stages[stage_idx + 1].transition;
1612 if matches!(next_transition, TransitionMode::Smooth) {
1613 bail!("RETAKE_SPLICE_REQUIRES_CUT_OR_FADE");
1614 }
1615 }
1616
1617 let now = now_ms_u64();
1618 let base_seed = effective.seed.unwrap_or(0);
1619 let old_seed = manifest.stage_status[stage_idx].seed;
1620 let new_seed = req
1621 .seed_offset
1622 .map(|offset| effective_stage_seed(base_seed, Some(offset)))
1623 .unwrap_or(old_seed);
1624 let old_prompt = effective.stages[stage_idx].prompt.clone();
1625 let new_prompt = req.prompt.clone().unwrap_or_else(|| old_prompt.clone());
1626 let now_i64 = i64::try_from(now).unwrap_or(i64::MAX);
1627
1628 if !chain_jobs::try_transition(
1629 db,
1630 job_id,
1631 &allowed_from,
1632 ChainJobState::Queued,
1633 None,
1634 now_i64,
1635 )? {
1636 let observed = chain_jobs::get_job(db, job_id)?
1637 .map(|row| row.state.as_str().to_string())
1638 .unwrap_or_else(|| "missing".to_string());
1639 bail!("chain job is not retakeable from current state {observed}");
1640 }
1641
1642 manifest.retakes.push(RetakeAmendment {
1643 stage_idx: req.stage_idx,
1644 mode: req.mode,
1645 old_seed,
1646 new_seed,
1647 old_prompt: (new_prompt != old_prompt).then_some(old_prompt),
1648 new_prompt: req.prompt.clone(),
1649 at_unix_ms: now,
1650 });
1651 manifest.stage_status[stage_idx].seed = new_seed;
1652
1653 let reset_end = match req.mode {
1654 RetakeMode::Cascade => manifest.stage_status.len(),
1655 RetakeMode::Splice => stage_idx + 1,
1656 };
1657 for idx in stage_idx..reset_end {
1658 let status = &mut manifest.stage_status[idx];
1659 status.state = StageState::Pending;
1660 status.frames_emitted = None;
1661 status.generation_time_ms = None;
1662 status.segment = None;
1663 status.tail_frames = None;
1664 status.audio = None;
1665 status.error = None;
1666 let stage_dir = JobDirLayout::new(job_dir.clone()).stage_dir(idx as u32);
1667 if stage_dir.exists() {
1668 std::fs::remove_dir_all(&stage_dir).with_context(|| {
1669 format!(
1670 "removing reset chain stage directory '{}'",
1671 stage_dir.display()
1672 )
1673 })?;
1674 }
1675 }
1676 manifest.write_atomic(&job_dir)?;
1677
1678 match req.mode {
1679 RetakeMode::Cascade => {
1680 chain_jobs::reset_stages_from(db, job_id, req.stage_idx, now_i64)?;
1681 }
1682 RetakeMode::Splice => {
1683 chain_jobs::reset_one_stage(db, job_id, req.stage_idx, now_i64)?;
1684 }
1685 }
1686 for idx in stage_idx..reset_end {
1687 chain_jobs::upsert_stage(
1688 db,
1689 &ChainJobStageRow {
1690 job_id: job_id.to_string(),
1691 stage_idx: idx as u32,
1692 state: StageState::Pending,
1693 seed: manifest.stage_status[idx].seed,
1694 frames_emitted: None,
1695 generation_time_ms: None,
1696 segment_rel_path: None,
1697 error: None,
1698 updated_at_ms: now_i64,
1699 },
1700 )?;
1701 }
1702 chain_jobs::set_current_stage(db, job_id, req.stage_idx, now_i64)?;
1703
1704 chain_jobs::get_job(db, job_id)?.ok_or_else(|| anyhow!("chain job disappeared after retake"))
1705}
1706
1707pub struct ProductionStageExecutor {
1708 gpu_pool: Arc<GpuPool>,
1709 config: mold_core::Config,
1710}
1711
1712impl ProductionStageExecutor {
1713 pub fn new(gpu_pool: Arc<GpuPool>, config: mold_core::Config) -> Self {
1714 Self { gpu_pool, config }
1715 }
1716}
1717
1718impl StageExecutor for ProductionStageExecutor {
1719 fn render_stage(
1720 &self,
1721 model: &str,
1722 stage_req: &GenerateRequest,
1723 carry: Option<&ChainTail>,
1724 motion_tail_frames: u32,
1725 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
1726 ) -> anyhow::Result<StageRenderOutcome> {
1727 let worker = select_worker_for_stage(&self.gpu_pool, model)
1728 .ok_or_else(|| anyhow!("no GPU worker available for chain stage model '{model}'"))?;
1729 let _in_flight = WorkerInFlightGuard::new(worker.clone());
1730 let _active = WorkerActiveGenerationGuard::new(worker.clone(), model, &stage_req.prompt)?;
1731 let hint = model_manager::family_for_model_sync(model, &self.config).map(|family| {
1732 model_manager::ActivationHint {
1733 width: stage_req.width,
1734 height: stage_req.height,
1735 batch: 1,
1736 dtype_bytes: 2,
1737 family: mold_inference::device::activation_family_for(&family),
1738 }
1739 });
1740 let carry_owned = carry.cloned();
1741 let stage_req = stage_req.clone();
1742 let prep = gpu_worker::run_stage_blocking(
1743 &worker,
1744 model,
1745 &self.config,
1746 hint,
1747 move |engine| -> anyhow::Result<StageRenderOutcome> {
1748 let renderer = engine.as_chain_renderer().ok_or_else(|| {
1749 anyhow!(
1750 "model '{}' does not support chained video generation",
1751 model
1752 )
1753 })?;
1754 let mut cancelled = false;
1755 let mut stage_progress = |event: StageProgressEvent| match event {
1756 StageProgressEvent::DenoiseStep { step, total } => {
1757 if progress(step, total).is_break() {
1758 cancelled = true;
1759 }
1760 }
1761 };
1762 let outcome = renderer.render_stage(
1763 &stage_req,
1764 carry_owned.as_ref(),
1765 motion_tail_frames,
1766 Some(&mut stage_progress),
1767 )?;
1768 if cancelled {
1769 Ok(StageRenderOutcome::Cancelled)
1770 } else {
1771 Ok(StageRenderOutcome::Done(outcome))
1772 }
1773 },
1774 )?;
1775 prep
1776 }
1777}
1778
1779pub struct ProductionQueueProbe {
1780 queue: QueueHandle,
1781 gpu_pool: Arc<GpuPool>,
1782}
1783
1784impl ProductionQueueProbe {
1785 pub fn new(queue: QueueHandle, gpu_pool: Arc<GpuPool>) -> Self {
1786 Self { queue, gpu_pool }
1787 }
1788}
1789
1790impl QueueProbe for ProductionQueueProbe {
1791 fn small_jobs_waiting(&self) -> usize {
1792 self.queue.pending()
1793 + self
1794 .gpu_pool
1795 .workers
1796 .iter()
1797 .map(|worker| worker.in_flight.load(Ordering::SeqCst))
1798 .sum::<usize>()
1799 }
1800}
1801
1802fn mark_stage_running(
1803 db: &MetadataDb,
1804 job_id: &str,
1805 manifest: &ChainJobManifest,
1806 stage_idx: u32,
1807) -> anyhow::Result<()> {
1808 let status = &manifest.stage_status[stage_idx as usize];
1809 chain_jobs::upsert_stage(
1810 db,
1811 &ChainJobStageRow {
1812 job_id: job_id.to_string(),
1813 stage_idx,
1814 state: StageState::Running,
1815 seed: status.seed,
1816 frames_emitted: None,
1817 generation_time_ms: None,
1818 segment_rel_path: None,
1819 error: None,
1820 updated_at_ms: now_ms_i64(),
1821 },
1822 )
1823}
1824
1825fn mark_manifest_stage_failed(
1826 manifest: &mut ChainJobManifest,
1827 layout: &JobDirLayout,
1828 stage_idx: u32,
1829 error: &str,
1830) -> anyhow::Result<()> {
1831 let status = manifest
1832 .stage_status
1833 .get_mut(stage_idx as usize)
1834 .ok_or_else(|| anyhow!("manifest missing status for stage {stage_idx}"))?;
1835 status.state = StageState::Failed;
1836 status.error = Some(error.to_string());
1837 status.frames_emitted = None;
1838 status.generation_time_ms = None;
1839 status.segment = None;
1840 status.tail_frames = None;
1841 status.audio = None;
1842 manifest.write_atomic(layout.root())?;
1843 Ok(())
1844}
1845
1846fn fail_job(
1847 db: &MetadataDb,
1848 deps: &RunnerDeps,
1849 job_id: &str,
1850 stage_idx: Option<u32>,
1851 error: String,
1852) -> anyhow::Result<()> {
1853 if let Some(stage_idx) = stage_idx {
1854 let stages = chain_jobs::stages_for_job(db, job_id)?;
1855 let seed = stages
1856 .iter()
1857 .find(|stage| stage.stage_idx == stage_idx)
1858 .map(|stage| stage.seed)
1859 .unwrap_or(0);
1860 chain_jobs::upsert_stage(
1861 db,
1862 &ChainJobStageRow {
1863 job_id: job_id.to_string(),
1864 stage_idx,
1865 state: StageState::Failed,
1866 seed,
1867 frames_emitted: None,
1868 generation_time_ms: None,
1869 segment_rel_path: None,
1870 error: Some(error.clone()),
1871 updated_at_ms: now_ms_i64(),
1872 },
1873 )?;
1874 }
1875 let changed = chain_jobs::try_transition(
1876 db,
1877 job_id,
1878 &[ChainJobState::Running],
1879 ChainJobState::Failed,
1880 Some(&error),
1881 now_ms_i64(),
1882 )?;
1883 if changed {
1884 deps.events.publish_then_remove(
1885 job_id,
1886 ChainJobEvent::StateChanged {
1887 state: ChainJobState::Failed,
1888 error: Some(error),
1889 },
1890 );
1891 deps.cancel.unregister(job_id);
1892 }
1893 Ok(())
1894}
1895
1896fn fail_stage_job(
1897 db: &MetadataDb,
1898 deps: &RunnerDeps,
1899 manifest: &mut ChainJobManifest,
1900 layout: &JobDirLayout,
1901 job_id: &str,
1902 stage_idx: u32,
1903 error: String,
1904) -> anyhow::Result<()> {
1905 {
1906 let _guard = deps.job_locks.blocking_lock(job_id);
1907 let _ = mark_manifest_stage_failed(manifest, layout, stage_idx, &error);
1908 }
1909 fail_job(db, deps, job_id, Some(stage_idx), error)
1910}
1911
1912fn set_cancelled(db: &MetadataDb, deps: &RunnerDeps, job_id: &str) -> anyhow::Result<()> {
1913 let changed = chain_jobs::try_transition(
1914 db,
1915 job_id,
1916 &[ChainJobState::Running],
1917 ChainJobState::Cancelled,
1918 None,
1919 now_ms_i64(),
1920 )?;
1921 if changed {
1922 deps.events.publish_then_remove(
1923 job_id,
1924 ChainJobEvent::StateChanged {
1925 state: ChainJobState::Cancelled,
1926 error: None,
1927 },
1928 );
1929 deps.cancel.unregister(job_id);
1930 }
1931 Ok(())
1932}
1933
1934fn carry_after_skipping_completed_stage(
1935 layout: &JobDirLayout,
1936 manifest: &ChainJobManifest,
1937 completed_stage_idx: u32,
1938) -> anyhow::Result<Option<ChainTail>> {
1939 let next_idx = completed_stage_idx + 1;
1940 let Some(next_status) = manifest.stage_status.get(next_idx as usize) else {
1941 return Ok(None);
1942 };
1943 if next_status.state == StageState::Completed {
1944 return Ok(None);
1945 }
1946 resume_carry_from_disk(layout.root(), manifest, next_idx)
1947}
1948
1949fn publish_yield_if_contended(deps: &RunnerDeps, job_id: &str) {
1956 let pending_small_jobs = deps.queue_probe.small_jobs_waiting();
1957 if pending_small_jobs > 0 {
1958 deps.events
1959 .publish(job_id, ChainJobEvent::Yielded { pending_small_jobs });
1960 }
1961}
1962
1963fn first_incomplete_stage(manifest: &ChainJobManifest) -> Option<u32> {
1964 manifest
1965 .stage_status
1966 .iter()
1967 .find(|stage| stage.state != StageState::Completed)
1968 .map(|stage| stage.idx)
1969}
1970
1971fn manifest_index_state(
1972 manifest: &ChainJobManifest,
1973 existing_state: ChainJobState,
1974) -> (ChainJobState, u32, Option<String>) {
1975 if let Some(stage) = manifest
1976 .stage_status
1977 .iter()
1978 .find(|stage| stage.state == StageState::Failed)
1979 {
1980 return (
1981 ChainJobState::Failed,
1982 stage.idx,
1983 stage.error.clone().or_else(|| Some("stage failed".into())),
1984 );
1985 }
1986 let current_stage =
1987 first_incomplete_stage(manifest).unwrap_or(manifest.stage_status.len() as u32);
1988 if current_stage == manifest.stage_status.len() as u32 && !manifest.finalizes.is_empty() {
1989 (ChainJobState::Completed, current_stage, None)
1990 } else if existing_state == ChainJobState::Running {
1991 (
1992 ChainJobState::Interrupted,
1993 current_stage,
1994 Some("server restarted while chain job was running".into()),
1995 )
1996 } else {
1997 (existing_state, current_stage, None)
1998 }
1999}
2000
2001pub(crate) fn effective_request(manifest: &ChainJobManifest) -> anyhow::Result<ChainRequest> {
2002 let mut request = manifest.request()?;
2003 for retake in &manifest.retakes {
2004 if let Some(stage) = request.stages.get_mut(retake.stage_idx as usize) {
2005 if let Some(prompt) = retake.new_prompt.as_ref() {
2006 stage.prompt = prompt.clone();
2007 }
2008 let base_seed = request.seed.unwrap_or(0);
2009 stage.seed_offset = Some(base_seed ^ retake.new_seed);
2010 }
2011 }
2012 Ok(request)
2013}
2014
2015fn build_stage_generate_request(
2016 stage: &ChainStage,
2017 chain: &ChainRequest,
2018 stage_seed: u64,
2019 idx: usize,
2020) -> GenerateRequest {
2021 GenerateRequest {
2022 prompt: stage.prompt.clone(),
2023 negative_prompt: stage.negative_prompt.clone(),
2024 model: chain.model.clone(),
2025 width: chain.width,
2026 height: chain.height,
2027 steps: chain.steps,
2028 guidance: chain.guidance,
2029 seed: Some(stage_seed),
2030 batch_size: 1,
2031 output_format: Some(OutputFormat::Mp4),
2032 embed_metadata: None,
2033 scheduler: None,
2034 cfg_plus: None,
2035 source_image: stage.source_image.clone(),
2036 edit_images: None,
2037 strength: if idx == 0 { chain.strength } else { 1.0 },
2038 mask_image: None,
2039 control_image: None,
2040 control_model: None,
2041 control_scale: 1.0,
2042 expand: None,
2043 original_prompt: None,
2044 lora: None,
2045 frames: Some(stage.frames),
2046 fps: Some(chain.fps),
2047 upscale_model: None,
2048 gif_preview: false,
2049 enable_audio: Some(chain.enable_audio.unwrap_or(false)),
2050 audio_file: None,
2051 audio_file_path: None,
2052 source_video: None,
2053 source_video_path: None,
2054 keyframes: None,
2055 pipeline: None,
2056 loras: None,
2057 retake_range: None,
2058 spatial_upscale: None,
2059 temporal_upscale: None,
2060 placement: chain.placement.clone(),
2061 }
2062}
2063
2064fn maybe_reencode_next_after_fade(
2065 layout: &JobDirLayout,
2066 manifest: &ChainJobManifest,
2067 stage_idx: u32,
2068 effective: &ChainRequest,
2069) -> anyhow::Result<()> {
2070 let next_idx = stage_idx + 1;
2071 let Some(next_stage) = effective.stages.get(next_idx as usize) else {
2072 return Ok(());
2073 };
2074 if !matches!(next_stage.transition, TransitionMode::Fade) {
2075 return Ok(());
2076 }
2077 let Some(next_status) = manifest.stage_status.get(next_idx as usize) else {
2078 return Ok(());
2079 };
2080 if next_status.state != StageState::Completed {
2081 return Ok(());
2082 }
2083 let fade_len = next_stage.fade_frames.unwrap_or(DEFAULT_FADE_FRAMES);
2084 let n = fade_len as usize;
2085 let next_segment = safe_join_manifest_rel(
2086 layout.root(),
2087 next_status
2088 .segment
2089 .as_ref()
2090 .ok_or_else(|| anyhow!("completed next stage {next_idx} has no segment"))?,
2091 )?;
2092 let (_metadata, mut frames) =
2093 mold_inference::ltx2::media::decode_video_frames_from_path(&next_segment)?;
2094 if frames.len() < n {
2095 bail!("next stage {next_idx} segment shorter than fade_len {n}");
2096 }
2097 let new_boundary_out = read_frames_from_dir(&layout.boundary_out_dir(stage_idx), n)?;
2098 let raw_boundary_in = read_frames_from_dir(&layout.boundary_in_dir(next_idx), n)?;
2099 let blended = fade_boundary(&new_boundary_out, &raw_boundary_in, fade_len);
2100 for (idx, frame) in blended.into_iter().enumerate() {
2101 frames[idx] = frame;
2102 }
2103 let bytes = video_enc::encode_mp4(&frames, effective.fps)?;
2104 write_file(&next_segment, &bytes)?;
2105 write_preview_jpeg(&layout.preview_path(next_idx), frames.last().unwrap())?;
2106 Ok(())
2107}
2108
2109fn safe_join_manifest_rel(root: &Path, rel: &str) -> anyhow::Result<PathBuf> {
2110 let path = Path::new(rel);
2111 if path.is_absolute()
2112 || rel.starts_with('/')
2113 || rel.starts_with('\\')
2114 || rel.as_bytes().get(0..3).is_some_and(|b| {
2115 b[0].is_ascii_alphabetic() && b[1] == b':' && matches!(b[2], b'/' | b'\\')
2116 })
2117 || path
2118 .components()
2119 .any(|component| matches!(component, Component::ParentDir))
2120 || rel.split(['/', '\\']).any(|part| part == "..")
2121 {
2122 bail!("manifest artifact path '{rel}' must be relative and must not contain '..'");
2123 }
2124 Ok(root.join(path))
2125}
2126
2127fn write_file(path: &Path, bytes: &[u8]) -> anyhow::Result<()> {
2128 if let Some(parent) = path.parent() {
2129 std::fs::create_dir_all(parent)
2130 .with_context(|| format!("creating directory '{}'", parent.display()))?;
2131 }
2132 std::fs::write(path, bytes).with_context(|| format!("writing '{}'", path.display()))
2133}
2134
2135fn write_frames_to_dir(dir: &Path, frames: &[RgbImage]) -> anyhow::Result<()> {
2136 std::fs::create_dir_all(dir).with_context(|| format!("creating '{}'", dir.display()))?;
2137 for entry in std::fs::read_dir(dir).with_context(|| format!("reading '{}'", dir.display()))? {
2138 let path = entry?.path();
2139 if path.extension().and_then(|s| s.to_str()) == Some("png") {
2140 std::fs::remove_file(&path)
2141 .with_context(|| format!("removing stale frame '{}'", path.display()))?;
2142 }
2143 }
2144 for (idx, frame) in frames.iter().enumerate() {
2145 let path = dir.join(format!("{idx:03}.png"));
2146 frame
2147 .save_with_format(&path, image::ImageFormat::Png)
2148 .with_context(|| format!("writing PNG frame '{}'", path.display()))?;
2149 }
2150 Ok(())
2151}
2152
2153fn read_frames_from_dir(dir: &Path, expected: usize) -> anyhow::Result<Vec<RgbImage>> {
2154 let mut paths: Vec<PathBuf> = std::fs::read_dir(dir)
2155 .with_context(|| format!("reading frame directory '{}'", dir.display()))?
2156 .map(|entry| entry.map(|entry| entry.path()))
2157 .collect::<std::io::Result<Vec<_>>>()?;
2158 paths.retain(|path| path.extension().and_then(|s| s.to_str()) == Some("png"));
2159 paths.sort();
2160 if paths.len() < expected {
2161 bail!(
2162 "frame directory '{}' has {} PNG(s), expected at least {expected}",
2163 dir.display(),
2164 paths.len()
2165 );
2166 }
2167 paths
2168 .into_iter()
2169 .take(expected)
2170 .map(|path| {
2171 image::open(&path)
2172 .with_context(|| format!("decoding PNG frame '{}'", path.display()))
2173 .map(|img| img.to_rgb8())
2174 })
2175 .collect()
2176}
2177
2178fn write_preview_jpeg(path: &Path, frame: &RgbImage) -> anyhow::Result<()> {
2179 let mut bytes = Vec::new();
2180 {
2181 let mut cursor = Cursor::new(&mut bytes);
2182 let encoder = JpegEncoder::new_with_quality(&mut cursor, 85);
2183 encoder.write_image(
2184 frame.as_raw(),
2185 frame.width(),
2186 frame.height(),
2187 image::ExtendedColorType::Rgb8,
2188 )?;
2189 }
2190 write_file(path, &bytes)
2191}
2192
2193fn trim_audio_front(
2194 audio: &mut Option<NativeAudioTrack>,
2195 frames: u32,
2196 fps: u32,
2197) -> anyhow::Result<()> {
2198 let Some(track) = audio.as_mut() else {
2199 return Ok(());
2200 };
2201 let samples = samples_for_frames(track, frames, fps);
2202 if track.interleaved_samples.len() < samples {
2203 bail!("audio sidecar too short for front trim");
2204 }
2205 track.interleaved_samples.drain(0..samples);
2206 Ok(())
2207}
2208
2209fn trim_audio_back(
2210 audio: &mut Option<NativeAudioTrack>,
2211 frames: u32,
2212 fps: u32,
2213) -> anyhow::Result<()> {
2214 let Some(track) = audio.as_mut() else {
2215 return Ok(());
2216 };
2217 let samples = samples_for_frames(track, frames, fps);
2218 if track.interleaved_samples.len() < samples {
2219 bail!("audio sidecar too short for back trim");
2220 }
2221 let keep = track.interleaved_samples.len() - samples;
2222 track.interleaved_samples.truncate(keep);
2223 Ok(())
2224}
2225
2226fn write_audio_boundary_out(
2227 layout: &JobDirLayout,
2228 stage_idx: u32,
2229 audio: &Option<NativeAudioTrack>,
2230 frames: u32,
2231 fps: u32,
2232) -> anyhow::Result<()> {
2233 let Some(track) = audio.as_ref() else {
2234 return Ok(());
2235 };
2236 let samples = samples_for_frames(track, frames, fps);
2237 if track.interleaved_samples.len() < samples {
2238 bail!("audio sidecar too short for boundary-out");
2239 }
2240 let mut boundary = track.clone();
2241 boundary.interleaved_samples =
2242 boundary.interleaved_samples[boundary.interleaved_samples.len() - samples..].to_vec();
2243 write_audio_sidecar(
2244 &layout.boundary_out_dir(stage_idx).join("audio.pcm"),
2245 &boundary,
2246 )
2247}
2248
2249fn blend_audio_front_from_previous_boundary(
2250 layout: &JobDirLayout,
2251 stage_idx: u32,
2252 audio: &mut Option<NativeAudioTrack>,
2253 frames: u32,
2254 fps: u32,
2255) -> anyhow::Result<()> {
2256 let Some(track) = audio.as_mut() else {
2257 return Ok(());
2258 };
2259 let samples = samples_for_frames(track, frames, fps);
2260 if samples == 0 {
2261 return Ok(());
2262 }
2263 if track.interleaved_samples.len() < samples {
2264 bail!("audio sidecar too short for boundary-in");
2265 }
2266 let mut boundary_in = track.clone();
2267 boundary_in.interleaved_samples = track.interleaved_samples[..samples].to_vec();
2268 write_audio_sidecar(
2269 &layout.boundary_in_dir(stage_idx).join("audio.pcm"),
2270 &boundary_in,
2271 )?;
2272 let previous = read_audio_sidecar(&layout.boundary_out_dir(stage_idx - 1).join("audio.pcm"))?;
2273 if previous.sample_rate != track.sample_rate || previous.channels != track.channels {
2274 bail!("fade audio boundary format mismatch");
2275 }
2276 if previous.interleaved_samples.len() < samples {
2277 bail!("previous fade audio boundary too short");
2278 }
2279 let channels = track.channels as usize;
2280 let frame_samples = channels.max(1);
2281 let denom = (samples / frame_samples).max(1) as f32;
2282 for sample_idx in 0..samples {
2283 let t = (sample_idx / frame_samples) as f32 / denom;
2284 let prior = previous.interleaved_samples[sample_idx];
2285 let next = track.interleaved_samples[sample_idx];
2286 track.interleaved_samples[sample_idx] = prior * (1.0 - t) + next * t;
2287 }
2288 Ok(())
2289}
2290
2291fn samples_for_frames(track: &NativeAudioTrack, frames: u32, fps: u32) -> usize {
2292 if fps == 0 {
2293 return 0;
2294 }
2295 ((track.sample_rate as u64 * frames as u64) / fps as u64) as usize * track.channels as usize
2296}
2297
2298fn write_audio_sidecar(path: &Path, track: &NativeAudioTrack) -> anyhow::Result<()> {
2299 let mut bytes = Vec::with_capacity(16 + track.interleaved_samples.len() * 4);
2300 bytes.extend_from_slice(AUDIO_SIDECAR_MAGIC);
2301 bytes.extend_from_slice(&track.sample_rate.to_le_bytes());
2302 bytes.extend_from_slice(&track.channels.to_le_bytes());
2303 bytes.extend_from_slice(&0u16.to_le_bytes());
2304 for sample in &track.interleaved_samples {
2305 bytes.extend_from_slice(&sample.to_le_bytes());
2306 }
2307 write_file(path, &bytes)
2308}
2309
2310fn read_audio_sidecar(path: &Path) -> anyhow::Result<NativeAudioTrack> {
2311 let bytes = std::fs::read(path).with_context(|| format!("reading '{}'", path.display()))?;
2312 if bytes.len() < 16 || &bytes[..8] != AUDIO_SIDECAR_MAGIC {
2313 bail!("audio sidecar '{}' has unsupported format", path.display());
2314 }
2315 let sample_rate = u32::from_le_bytes(bytes[8..12].try_into().unwrap());
2316 let channels = u16::from_le_bytes(bytes[12..14].try_into().unwrap());
2317 let sample_bytes = &bytes[16..];
2318 if sample_bytes.len() % 4 != 0 {
2319 bail!(
2320 "audio sidecar '{}' has truncated f32 sample",
2321 path.display()
2322 );
2323 }
2324 let interleaved_samples = sample_bytes
2325 .chunks_exact(4)
2326 .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
2327 .collect();
2328 Ok(NativeAudioTrack {
2329 interleaved_samples,
2330 sample_rate,
2331 channels,
2332 })
2333}
2334
2335fn select_worker_for_stage(gpu_pool: &GpuPool, model: &str) -> Option<Arc<GpuWorker>> {
2336 let est = crate::queue::estimate_model_vram(model);
2337 gpu_pool.select_worker(model, est)
2338}
2339
2340struct WorkerInFlightGuard {
2341 worker: Arc<GpuWorker>,
2342}
2343
2344impl WorkerInFlightGuard {
2345 fn new(worker: Arc<GpuWorker>) -> Self {
2346 worker.in_flight.fetch_add(1, Ordering::SeqCst);
2347 Self { worker }
2348 }
2349}
2350
2351impl Drop for WorkerInFlightGuard {
2352 fn drop(&mut self) {
2353 self.worker.in_flight.fetch_sub(1, Ordering::SeqCst);
2354 }
2355}
2356
2357struct WorkerActiveGenerationGuard {
2358 worker: Arc<GpuWorker>,
2359}
2360
2361impl WorkerActiveGenerationGuard {
2362 fn new(worker: Arc<GpuWorker>, model: &str, prompt: &str) -> anyhow::Result<Self> {
2363 let mut active = worker
2364 .active_generation
2365 .write()
2366 .map_err(|e| anyhow!("active_generation lock poisoned: {e}"))?;
2367 *active = Some(ActiveGeneration {
2368 model: model.to_string(),
2369 prompt_sha256: format!("{:x}", Sha256::digest(prompt.as_bytes())),
2370 started_at_unix_ms: now_ms_u64(),
2371 started_at: std::time::Instant::now(),
2372 });
2373 drop(active);
2374 Ok(Self { worker })
2375 }
2376}
2377
2378impl Drop for WorkerActiveGenerationGuard {
2379 fn drop(&mut self) {
2380 if let Ok(mut active) = self.worker.active_generation.write() {
2381 *active = None;
2382 }
2383 }
2384}
2385
2386fn now_ms_u64() -> u64 {
2387 mold_core::time::now_epoch_ms_u64()
2388}
2389
2390fn now_ms_i64() -> i64 {
2391 mold_core::time::now_epoch_ms()
2392}
2393
2394#[cfg(test)]
2395mod tests {
2396 use super::*;
2397 use image::{Rgb, RgbImage};
2398 use mold_core::chain::{ChainStage, TransitionMode};
2399 use mold_core::types::OutputFormat;
2400 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2401
2402 fn db() -> MetadataDb {
2403 MetadataDb::open_in_memory().unwrap()
2404 }
2405
2406 fn stage(prompt: &str, transition: TransitionMode) -> ChainStage {
2407 ChainStage {
2408 prompt: prompt.into(),
2409 frames: 9,
2410 source_image: None,
2411 negative_prompt: None,
2412 seed_offset: None,
2413 transition,
2414 fade_frames: Some(2),
2415 model: None,
2416 loras: vec![],
2417 references: vec![],
2418 }
2419 }
2420
2421 fn request(transitions: Vec<TransitionMode>) -> ChainRequest {
2422 ChainRequest {
2423 model: "ltx-2-19b-distilled:fp8".into(),
2424 stages: transitions
2425 .into_iter()
2426 .enumerate()
2427 .map(|(idx, transition)| stage(&format!("stage {idx}"), transition))
2428 .collect(),
2429 motion_tail_frames: 1,
2430 width: 64,
2431 height: 48,
2432 fps: 8,
2433 seed: Some(42),
2434 steps: 2,
2435 guidance: 1.0,
2436 strength: 1.0,
2437 output_format: OutputFormat::Mp4,
2438 placement: None,
2439 prompt: None,
2440 total_frames: None,
2441 clip_frames: None,
2442 source_image: None,
2443 enable_audio: None,
2444 }
2445 }
2446
2447 fn job(
2448 id: &str,
2449 state: ChainJobState,
2450 created_at_ms: i64,
2451 dir: PathBuf,
2452 req: &ChainRequest,
2453 ) -> ChainJobRow {
2454 ChainJobRow {
2455 id: id.into(),
2456 state,
2457 model: req.model.clone(),
2458 request_json: serde_json::to_string(req).unwrap(),
2459 job_dir: dir,
2460 stage_count: req.stages.len() as u32,
2461 current_stage: 0,
2462 error: None,
2463 created_at_ms,
2464 updated_at_ms: created_at_ms,
2465 finalized_at_ms: None,
2466 }
2467 }
2468
2469 fn frame(value: u8) -> RgbImage {
2470 RgbImage::from_pixel(64, 48, Rgb([value, value, value]))
2471 }
2472
2473 fn outcome(value: u8) -> StageOutcome {
2474 let frames = (0..9).map(|i| frame(value + i)).collect::<Vec<_>>();
2475 StageOutcome {
2476 tail: ChainTail {
2477 frames: 1,
2478 tail_rgb_frames: vec![frame(value + 8)],
2479 },
2480 frames,
2481 audio: None,
2482 generation_time_ms: 10,
2483 }
2484 }
2485
2486 struct FakeExecutor {
2487 calls: AtomicUsize,
2488 cancel_on_progress: AtomicBool,
2489 }
2490
2491 impl StageExecutor for FakeExecutor {
2492 fn render_stage(
2493 &self,
2494 _model: &str,
2495 _stage_req: &GenerateRequest,
2496 _carry: Option<&ChainTail>,
2497 _motion_tail_frames: u32,
2498 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2499 ) -> anyhow::Result<StageRenderOutcome> {
2500 let call = self.calls.fetch_add(1, Ordering::SeqCst) as u8;
2501 if self.cancel_on_progress.load(Ordering::SeqCst) && progress(1, 2).is_break() {
2502 return Ok(StageRenderOutcome::Cancelled);
2503 }
2504 let _ = progress(1, 2);
2505 Ok(StageRenderOutcome::Done(outcome(10 + call * 20)))
2506 }
2507 }
2508
2509 struct FakeProbe(AtomicUsize);
2510
2511 impl QueueProbe for FakeProbe {
2512 fn small_jobs_waiting(&self) -> usize {
2513 self.0.load(Ordering::SeqCst)
2514 }
2515 }
2516
2517 struct CarryInspectExecutor {
2518 seen: Mutex<Vec<Option<[u8; 3]>>>,
2519 }
2520
2521 impl StageExecutor for CarryInspectExecutor {
2522 fn render_stage(
2523 &self,
2524 _model: &str,
2525 _stage_req: &GenerateRequest,
2526 carry: Option<&ChainTail>,
2527 _motion_tail_frames: u32,
2528 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2529 ) -> anyhow::Result<StageRenderOutcome> {
2530 let pixel = carry
2531 .and_then(|tail| tail.tail_rgb_frames.first())
2532 .map(|frame| frame.get_pixel(0, 0).0);
2533 self.seen.lock().unwrap().push(pixel);
2534 let _ = progress(1, 2);
2535 Ok(StageRenderOutcome::Done(outcome(140)))
2536 }
2537 }
2538
2539 struct BoundaryCancelExecutor {
2540 calls: AtomicUsize,
2541 cancel: Arc<CancelRegistry>,
2542 job_id: String,
2543 }
2544
2545 impl StageExecutor for BoundaryCancelExecutor {
2546 fn render_stage(
2547 &self,
2548 _model: &str,
2549 _stage_req: &GenerateRequest,
2550 _carry: Option<&ChainTail>,
2551 _motion_tail_frames: u32,
2552 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2553 ) -> anyhow::Result<StageRenderOutcome> {
2554 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2555 let _ = progress(1, 2);
2556 if call == 0 {
2557 assert!(self.cancel.request(&self.job_id));
2558 }
2559 Ok(StageRenderOutcome::Done(outcome(30 + call as u8 * 20)))
2560 }
2561 }
2562
2563 fn deps(
2564 db: MetadataDb,
2565 root: PathBuf,
2566 executor: Arc<FakeExecutor>,
2567 probe: Arc<FakeProbe>,
2568 ) -> RunnerDeps {
2569 RunnerDeps {
2570 db: Arc::new(Some(db)),
2571 jobs_root: root,
2572 executor,
2573 queue_probe: probe,
2574 events: Arc::new(JobEventBus::new()),
2575 cancel: Arc::new(CancelRegistry::new()),
2576 job_locks: Arc::new(JobMutationLocks::new()),
2577 claims: Arc::new(EphemeralClaims::default()),
2578 output_dir: None,
2579 }
2580 }
2581
2582 fn persist_job(
2583 db: &MetadataDb,
2584 dir: &Path,
2585 id: &str,
2586 req: &ChainRequest,
2587 state: ChainJobState,
2588 ) -> ChainJobRow {
2589 let manifest = ChainJobManifest::new(id.into(), 1_000, req).unwrap();
2590 std::fs::create_dir_all(dir).unwrap();
2591 manifest.write_atomic(dir).unwrap();
2592 let row = job(id, state, 1_000, dir.to_path_buf(), req);
2593 chain_jobs::insert_job(db, &row).unwrap();
2594 for stage in &manifest.stage_status {
2595 chain_jobs::upsert_stage(
2596 db,
2597 &ChainJobStageRow {
2598 job_id: id.into(),
2599 stage_idx: stage.idx,
2600 state: stage.state,
2601 seed: stage.seed,
2602 frames_emitted: None,
2603 generation_time_ms: None,
2604 segment_rel_path: None,
2605 error: None,
2606 updated_at_ms: 1_000,
2607 },
2608 )
2609 .unwrap();
2610 }
2611 row
2612 }
2613
2614 #[test]
2615 fn corrupt_manifest_marks_job_failed_without_rendering() {
2616 let dir = tempfile::tempdir().unwrap();
2617 let db = db();
2618 let req = request(vec![TransitionMode::Smooth]);
2619 let job_dir = dir.path().join("job");
2620 let row = persist_job(&db, &job_dir, "01JBR55CORRUPT", &req, ChainJobState::Queued);
2621 std::fs::write(job_dir.join("manifest.toml"), "not valid toml").unwrap();
2622 let executor = Arc::new(FakeExecutor {
2623 calls: AtomicUsize::new(0),
2624 cancel_on_progress: AtomicBool::new(false),
2625 });
2626 let deps = deps(
2627 db,
2628 dir.path().join("jobs"),
2629 executor.clone(),
2630 Arc::new(FakeProbe(AtomicUsize::new(0))),
2631 );
2632
2633 execute_job(&deps, &row, 0).unwrap();
2634
2635 let db = deps.db.as_ref().as_ref().unwrap();
2636 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2637 assert_eq!(failed.state, ChainJobState::Failed);
2638 assert!(
2639 failed
2640 .error
2641 .as_deref()
2642 .is_some_and(|err| err.contains("manifest TOML parse failed")),
2643 "parse error must be persisted, got {:?}",
2644 failed.error
2645 );
2646 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
2647 }
2648
2649 #[test]
2650 fn resume_carry_error_after_running_marks_job_failed_and_unlocks_row() {
2651 let dir = tempfile::tempdir().unwrap();
2652 let db = db();
2653 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
2654 let job_dir = dir.path().join("job");
2655 let row = persist_job(&db, &job_dir, "01JBR55BADTAIL", &req, ChainJobState::Queued);
2656 let layout = JobDirLayout::new(job_dir.clone());
2657 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2658 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
2659 std::fs::remove_dir_all(layout.tail_dir(0)).unwrap();
2660 chain_jobs::upsert_stage(
2661 &db,
2662 &ChainJobStageRow {
2663 job_id: row.id.clone(),
2664 stage_idx: 0,
2665 state: StageState::Completed,
2666 seed: manifest.stage_status[0].seed,
2667 frames_emitted: manifest.stage_status[0].frames_emitted,
2668 generation_time_ms: manifest.stage_status[0].generation_time_ms,
2669 segment_rel_path: manifest.stage_status[0].segment.clone(),
2670 error: None,
2671 updated_at_ms: 1_000,
2672 },
2673 )
2674 .unwrap();
2675 let executor = Arc::new(FakeExecutor {
2676 calls: AtomicUsize::new(0),
2677 cancel_on_progress: AtomicBool::new(false),
2678 });
2679 let deps = deps(
2680 db,
2681 dir.path().join("jobs"),
2682 executor.clone(),
2683 Arc::new(FakeProbe(AtomicUsize::new(0))),
2684 );
2685
2686 execute_job(&deps, &row, 0).unwrap();
2687
2688 let db = deps.db.as_ref().as_ref().unwrap();
2689 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2690 assert_eq!(failed.state, ChainJobState::Failed);
2691 assert!(
2692 failed
2693 .error
2694 .as_deref()
2695 .is_some_and(|err| err.contains("chain tail directory")),
2696 "resume error must be persisted, got {:?}",
2697 failed.error
2698 );
2699 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
2700 assert!(chain_jobs::set_job_queued(db, &row.id, 1, now_ms_i64()).unwrap());
2701 assert_eq!(
2702 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2703 ChainJobState::Queued
2704 );
2705 assert!(chain_jobs::delete_job(db, &row.id).unwrap());
2706 }
2707
2708 #[test]
2709 fn skipping_completed_stage_reloads_tail_for_next_smooth_stage() {
2710 let dir = tempfile::tempdir().unwrap();
2711 let db = db();
2712 let req = request(vec![
2713 TransitionMode::Smooth,
2714 TransitionMode::Smooth,
2715 TransitionMode::Smooth,
2716 ]);
2717 let job_dir = dir.path().join("job");
2718 let row = persist_job(
2719 &db,
2720 &job_dir,
2721 "01JBR55SKIPCARRY",
2722 &req,
2723 ChainJobState::Queued,
2724 );
2725 let layout = JobDirLayout::new(job_dir.clone());
2726 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2727 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
2728 write_stage_artifacts(&layout, &mut manifest, 1, &outcome(80), &req).unwrap();
2729 for idx in 0..2 {
2730 let status = &manifest.stage_status[idx];
2731 chain_jobs::upsert_stage(
2732 &db,
2733 &ChainJobStageRow {
2734 job_id: row.id.clone(),
2735 stage_idx: idx as u32,
2736 state: status.state,
2737 seed: status.seed,
2738 frames_emitted: status.frames_emitted,
2739 generation_time_ms: status.generation_time_ms,
2740 segment_rel_path: status.segment.clone(),
2741 error: None,
2742 updated_at_ms: 1_000,
2743 },
2744 )
2745 .unwrap();
2746 }
2747 let executor = Arc::new(CarryInspectExecutor {
2748 seen: Mutex::new(Vec::new()),
2749 });
2750 let deps = RunnerDeps {
2751 db: Arc::new(Some(db)),
2752 jobs_root: dir.path().join("jobs"),
2753 executor: executor.clone(),
2754 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
2755 events: Arc::new(JobEventBus::new()),
2756 cancel: Arc::new(CancelRegistry::new()),
2757 job_locks: Arc::new(JobMutationLocks::new()),
2758 claims: Arc::new(EphemeralClaims::default()),
2759 output_dir: None,
2760 };
2761
2762 execute_job(&deps, &row, 0).unwrap();
2763
2764 assert_eq!(
2765 executor.seen.lock().unwrap().as_slice(),
2766 &[Some([88, 88, 88])],
2767 "stage 2 must receive stage 1's persisted smooth tail"
2768 );
2769 }
2770
2771 #[test]
2772 fn runner_next_queued_job_uses_fifo_order() {
2773 let db = db();
2774 let req = request(vec![TransitionMode::Smooth]);
2775 let newer = job(
2776 "01JBR55NEWER",
2777 ChainJobState::Queued,
2778 2_000,
2779 std::env::temp_dir().join("newer"),
2780 &req,
2781 );
2782 let older = job(
2783 "01JBR55OLDER",
2784 ChainJobState::Queued,
2785 1_000,
2786 std::env::temp_dir().join("older"),
2787 &req,
2788 );
2789 chain_jobs::insert_job(&db, &newer).unwrap();
2790 chain_jobs::insert_job(&db, &older).unwrap();
2791
2792 let got = next_queued_job(&db).unwrap().expect("queued job");
2793 assert_eq!(got.id, older.id);
2794 }
2795
2796 #[test]
2797 fn safe_join_manifest_rel_rejects_windows_and_backslash_traversal() {
2798 let root = Path::new("/tmp/mold-chain-job");
2799 for rel in [
2800 r"C:\tmp\segment.mp4",
2801 r"stages\000\..\segment.mp4",
2802 r"\stages\000\segment.mp4",
2803 ] {
2804 let err = safe_join_manifest_rel(root, rel).unwrap_err();
2805 assert!(
2806 err.to_string().contains("must be relative"),
2807 "expected traversal rejection for {rel:?}, got {err:#}"
2808 );
2809 }
2810 }
2811
2812 #[test]
2813 fn completed_job_subscriptions_are_non_persistent() {
2814 let db = db();
2815 let req = request(vec![TransitionMode::Smooth]);
2816 let row = job(
2817 "01JBR55DONEATTACH",
2818 ChainJobState::Completed,
2819 1_000,
2820 std::env::temp_dir().join("done-attach"),
2821 &req,
2822 );
2823 chain_jobs::insert_job(&db, &row).unwrap();
2824 let events = JobEventBus::new();
2825
2826 let _rx1 = events.subscribe_for_job(&db, &row.id).unwrap();
2827 let _rx2 = events.subscribe_for_job(&db, &row.id).unwrap();
2828
2829 assert!(!events.contains_for_tests(&row.id));
2830 }
2831
2832 #[tokio::test]
2833 async fn subscribe_receives_live_event_buffered_after_attach() {
2834 let (kick_tx, _kick_rx) = tokio::sync::mpsc::unbounded_channel();
2835 let handle = ChainJobRunnerHandle {
2836 kick_tx,
2837 cancel: Arc::new(CancelRegistry::new()),
2838 events: Arc::new(JobEventBus::new()),
2839 job_locks: Arc::new(JobMutationLocks::new()),
2840 claims: Arc::new(EphemeralClaims::default()),
2841 };
2842
2843 let mut rx = handle
2844 .events
2845 .subscribe_persistent_for_tests("01JBR55EVENTS");
2846 handle.events.publish(
2847 "01JBR55EVENTS",
2848 ChainJobEvent::StateChanged {
2849 state: ChainJobState::Running,
2850 error: None,
2851 },
2852 );
2853
2854 let event = rx.recv().await.unwrap();
2855 assert!(matches!(
2856 event,
2857 ChainJobEvent::StateChanged {
2858 state: ChainJobState::Running,
2859 error: None
2860 }
2861 ));
2862 }
2863
2864 #[test]
2865 fn execute_job_persists_stage_artifacts_manifest_then_db_and_finalizes() {
2866 let dir = tempfile::tempdir().unwrap();
2867 let db = db();
2868 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
2869 let job_dir = dir.path().join("job");
2870 let row = persist_job(&db, &job_dir, "01JBR55ORDER", &req, ChainJobState::Queued);
2871 let executor = Arc::new(FakeExecutor {
2872 calls: AtomicUsize::new(0),
2873 cancel_on_progress: AtomicBool::new(false),
2874 });
2875 let deps = deps(
2876 db,
2877 dir.path().join("jobs"),
2878 executor,
2879 Arc::new(FakeProbe(AtomicUsize::new(0))),
2880 );
2881
2882 execute_job(&deps, &row, 0).unwrap();
2883
2884 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2885 let db = deps.db.as_ref().as_ref().unwrap();
2886 let job_after = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2887 assert!(
2888 job_dir.join("stages/000/segment.mp4").exists(),
2889 "job state {:?}, error {:?}",
2890 job_after.state,
2891 job_after.error
2892 );
2893 assert!(job_dir.join("stages/000/tail/000.png").exists());
2894 assert!(job_dir.join("stages/001/segment.mp4").exists());
2895 assert_eq!(manifest.stage_status[0].state, StageState::Completed);
2896 assert_eq!(manifest.stage_status[1].state, StageState::Completed);
2897 assert_eq!(manifest.finalizes.len(), 1);
2898
2899 let stages = chain_jobs::stages_for_job(db, &row.id).unwrap();
2900 assert_eq!(stages[0].state, StageState::Completed);
2901 assert_eq!(
2902 stages[0].segment_rel_path.as_deref(),
2903 Some("stages/000/segment.mp4")
2904 );
2905 assert_eq!(
2906 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2907 ChainJobState::Completed
2908 );
2909 }
2910
2911 #[test]
2912 fn ephemeral_execute_job_defers_gallery_record_to_legacy_shim() {
2913 let dir = tempfile::tempdir().unwrap();
2914 let db = db();
2915 let req = request(vec![TransitionMode::Smooth]);
2916 let job_dir = dir.path().join("job");
2917 let row = persist_job(
2918 &db,
2919 &job_dir,
2920 "01JBR55EPHGALLERY",
2921 &req,
2922 ChainJobState::Queued,
2923 );
2924 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2925 manifest.ephemeral = true;
2926 manifest.write_atomic(&job_dir).unwrap();
2927 let executor = Arc::new(FakeExecutor {
2928 calls: AtomicUsize::new(0),
2929 cancel_on_progress: AtomicBool::new(false),
2930 });
2931 let output_dir = dir.path().join("gallery");
2932 let mut deps = deps(
2933 db,
2934 dir.path().join("jobs"),
2935 executor,
2936 Arc::new(FakeProbe(AtomicUsize::new(0))),
2937 );
2938 deps.output_dir = Some(output_dir.clone());
2939
2940 execute_job(&deps, &row, 0).unwrap();
2941
2942 let db = deps.db.as_ref().as_ref().unwrap();
2943 assert_eq!(
2944 db.list(Some(&output_dir)).unwrap().len(),
2945 0,
2946 "ephemeral legacy shim jobs must not write runner-side gallery rows"
2947 );
2948 assert_eq!(
2949 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2950 ChainJobState::Completed
2951 );
2952 }
2953
2954 #[test]
2955 fn durable_execute_job_records_exactly_one_runner_gallery_row() {
2956 let dir = tempfile::tempdir().unwrap();
2957 let db = db();
2958 let req = request(vec![TransitionMode::Smooth]);
2959 let job_dir = dir.path().join("job");
2960 let row = persist_job(
2961 &db,
2962 &job_dir,
2963 "01JBR55DURGALLERY",
2964 &req,
2965 ChainJobState::Queued,
2966 );
2967 let executor = Arc::new(FakeExecutor {
2968 calls: AtomicUsize::new(0),
2969 cancel_on_progress: AtomicBool::new(false),
2970 });
2971 let output_dir = dir.path().join("gallery");
2972 let mut deps = deps(
2973 db,
2974 dir.path().join("jobs"),
2975 executor,
2976 Arc::new(FakeProbe(AtomicUsize::new(0))),
2977 );
2978 deps.output_dir = Some(output_dir.clone());
2979
2980 execute_job(&deps, &row, 0).unwrap();
2981
2982 let db = deps.db.as_ref().as_ref().unwrap();
2983 let rows = db.list(Some(&output_dir)).unwrap();
2984 assert_eq!(rows.len(), 1, "durable jobs save exactly one gallery row");
2985 assert_eq!(rows[0].format, OutputFormat::Mp4);
2986 assert_eq!(
2987 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2988 ChainJobState::Completed
2989 );
2990 }
2991
2992 #[test]
2993 fn cancel_via_progress_marks_job_cancelled_and_keeps_completed_artifacts() {
2994 let dir = tempfile::tempdir().unwrap();
2995 let db = db();
2996 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
2997 let job_dir = dir.path().join("job");
2998 let row = persist_job(&db, &job_dir, "01JBR55CANCEL", &req, ChainJobState::Queued);
2999 let executor = Arc::new(FakeExecutor {
3000 calls: AtomicUsize::new(0),
3001 cancel_on_progress: AtomicBool::new(true),
3002 });
3003 let deps = deps(
3004 db,
3005 dir.path().join("jobs"),
3006 executor,
3007 Arc::new(FakeProbe(AtomicUsize::new(0))),
3008 );
3009 deps.cancel.register(&row.id);
3010 assert!(deps.cancel.request(&row.id));
3011
3012 execute_job(&deps, &row, 0).unwrap();
3013
3014 let db = deps.db.as_ref().as_ref().unwrap();
3015 assert_eq!(
3016 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3017 ChainJobState::Cancelled
3018 );
3019 assert!(!job_dir.join("stages/000/segment.mp4").exists());
3020 }
3021
3022 #[test]
3023 fn cancel_then_retake_runs_to_completion_without_stale_cancel_mark() {
3024 let dir = tempfile::tempdir().unwrap();
3025 let db = db();
3026 let req = request(vec![TransitionMode::Smooth]);
3027 let job_dir = dir.path().join("job");
3028 let row = persist_job(
3029 &db,
3030 &job_dir,
3031 "01JBR55CANCELRETAKE",
3032 &req,
3033 ChainJobState::Queued,
3034 );
3035 let executor = Arc::new(FakeExecutor {
3036 calls: AtomicUsize::new(0),
3037 cancel_on_progress: AtomicBool::new(false),
3038 });
3039 let deps = deps(
3040 db,
3041 dir.path().join("jobs"),
3042 executor.clone(),
3043 Arc::new(FakeProbe(AtomicUsize::new(0))),
3044 );
3045 deps.cancel.register(&row.id);
3046 assert!(deps.cancel.request(&row.id));
3047 execute_job(&deps, &row, 0).unwrap();
3048
3049 let db = deps.db.as_ref().as_ref().unwrap();
3050 assert_eq!(
3051 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3052 ChainJobState::Cancelled
3053 );
3054 let updated = apply_retake(
3055 db,
3056 dir.path(),
3057 &row.id,
3058 &RetakeRequest {
3059 stage_idx: 0,
3060 mode: RetakeMode::Cascade,
3061 seed_offset: Some(9),
3062 prompt: None,
3063 },
3064 )
3065 .unwrap();
3066
3067 execute_job(&deps, &updated, 0).unwrap();
3068
3069 assert_eq!(
3070 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3071 ChainJobState::Completed
3072 );
3073 assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
3074 }
3075
3076 #[test]
3077 fn cancel_mark_after_claim_is_honored_before_first_stage() {
3078 let dir = tempfile::tempdir().unwrap();
3079 let db = db();
3080 let req = request(vec![TransitionMode::Smooth]);
3081 let job_dir = dir.path().join("job");
3082 let row = persist_job(
3083 &db,
3084 &job_dir,
3085 "01JBR55CLAIMCANCEL",
3086 &req,
3087 ChainJobState::Queued,
3088 );
3089 assert!(chain_jobs::claim_job(&db, &row.id).unwrap());
3090 let executor = Arc::new(FakeExecutor {
3091 calls: AtomicUsize::new(0),
3092 cancel_on_progress: AtomicBool::new(false),
3093 });
3094 let deps = deps(
3095 db,
3096 dir.path().join("jobs"),
3097 executor.clone(),
3098 Arc::new(FakeProbe(AtomicUsize::new(0))),
3099 );
3100 deps.cancel.register(&row.id);
3101 assert!(deps.cancel.request(&row.id));
3102 let running = chain_jobs::get_job(deps.db.as_ref().as_ref().unwrap(), &row.id)
3103 .unwrap()
3104 .unwrap();
3105
3106 execute_job(&deps, &running, 0).unwrap();
3107
3108 let db = deps.db.as_ref().as_ref().unwrap();
3109 assert_eq!(
3110 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3111 ChainJobState::Cancelled
3112 );
3113 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
3114 }
3115
3116 #[test]
3117 fn retake_aborts_after_resume_wins_without_manifest_corruption() {
3118 let dir = tempfile::tempdir().unwrap();
3119 let db = db();
3120 let req = request(vec![TransitionMode::Smooth, TransitionMode::Cut]);
3121 let job_dir = dir.path().join("job");
3122 let row = persist_job(
3123 &db,
3124 &job_dir,
3125 "01JBR55RETAKERACE",
3126 &req,
3127 ChainJobState::Failed,
3128 );
3129 assert!(chain_jobs::try_transition(
3130 &db,
3131 &row.id,
3132 &[ChainJobState::Failed],
3133 ChainJobState::Queued,
3134 None,
3135 2_000,
3136 )
3137 .unwrap());
3138
3139 let err = apply_retake(
3140 &db,
3141 dir.path(),
3142 &row.id,
3143 &RetakeRequest {
3144 stage_idx: 0,
3145 mode: RetakeMode::Cascade,
3146 seed_offset: Some(5),
3147 prompt: Some("should not land".into()),
3148 },
3149 )
3150 .unwrap_err();
3151
3152 assert!(
3153 err.to_string().contains("not retakeable")
3154 || err.to_string().contains("CHAIN_JOB_RUNNING"),
3155 "retake loser must abort with a current-state error, got {err:#}"
3156 );
3157 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3158 assert!(manifest.retakes.is_empty());
3159 assert_eq!(manifest.stage_status[0].seed, 42);
3160 }
3161
3162 #[test]
3163 fn apply_retake_rejects_splice_before_smooth_boundary() {
3164 let dir = tempfile::tempdir().unwrap();
3165 let db = db();
3166 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3167 let job_dir = dir.path().join("job");
3168 let _row = persist_job(
3169 &db,
3170 &job_dir,
3171 "01JBR55RETAKE",
3172 &req,
3173 ChainJobState::Completed,
3174 );
3175
3176 let err = apply_retake(
3177 &db,
3178 dir.path(),
3179 "01JBR55RETAKE",
3180 &RetakeRequest {
3181 stage_idx: 0,
3182 mode: RetakeMode::Splice,
3183 seed_offset: Some(9),
3184 prompt: None,
3185 },
3186 )
3187 .unwrap_err();
3188 assert!(err
3189 .to_string()
3190 .contains("RETAKE_SPLICE_REQUIRES_CUT_OR_FADE"));
3191 }
3192
3193 #[test]
3194 fn apply_retake_cascade_resets_target_through_end_and_records_amendment() {
3195 let dir = tempfile::tempdir().unwrap();
3196 let db = db();
3197 let req = request(vec![
3198 TransitionMode::Smooth,
3199 TransitionMode::Cut,
3200 TransitionMode::Cut,
3201 ]);
3202 let job_dir = dir.path().join("job");
3203 let row = persist_job(
3204 &db,
3205 &job_dir,
3206 "01JBR55CASCADE",
3207 &req,
3208 ChainJobState::Completed,
3209 );
3210 let layout = JobDirLayout::new(job_dir.clone());
3211 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3212 for stage in &mut manifest.stage_status {
3213 stage.state = StageState::Completed;
3214 stage.segment = Some(layout.segment_rel(stage.idx));
3215 layout.ensure_stage_dirs(stage.idx).unwrap();
3216 std::fs::write(layout.stage_dir(stage.idx).join("stale.txt"), b"stale").unwrap();
3217 chain_jobs::upsert_stage(
3218 &db,
3219 &ChainJobStageRow {
3220 job_id: row.id.clone(),
3221 stage_idx: stage.idx,
3222 state: StageState::Completed,
3223 seed: stage.seed,
3224 frames_emitted: Some(9),
3225 generation_time_ms: Some(10),
3226 segment_rel_path: stage.segment.clone(),
3227 error: None,
3228 updated_at_ms: 1_000,
3229 },
3230 )
3231 .unwrap();
3232 }
3233 manifest.write_atomic(&job_dir).unwrap();
3234
3235 let updated = apply_retake(
3236 &db,
3237 dir.path(),
3238 &row.id,
3239 &RetakeRequest {
3240 stage_idx: 1,
3241 mode: RetakeMode::Cascade,
3242 seed_offset: Some(7),
3243 prompt: Some("new middle".into()),
3244 },
3245 )
3246 .unwrap();
3247
3248 assert_eq!(updated.state, ChainJobState::Queued);
3249 assert_eq!(updated.current_stage, 1);
3250 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3251 assert_eq!(manifest.retakes.len(), 1);
3252 assert_eq!(manifest.retakes[0].stage_idx, 1);
3253 assert_eq!(
3254 manifest.retakes[0].new_prompt.as_deref(),
3255 Some("new middle")
3256 );
3257 assert_eq!(manifest.stage_status[0].state, StageState::Completed);
3258 assert_eq!(manifest.stage_status[1].state, StageState::Pending);
3259 assert_eq!(manifest.stage_status[2].state, StageState::Pending);
3260 assert!(!layout.stage_dir(1).exists());
3261 assert!(!layout.stage_dir(2).exists());
3262 let stages = chain_jobs::stages_for_job(&db, &row.id).unwrap();
3263 assert_eq!(stages[1].state, StageState::Pending);
3264 assert_eq!(stages[2].state, StageState::Pending);
3265 let effective = effective_request(&manifest).unwrap();
3266 assert_eq!(effective.stages[1].prompt, "new middle");
3267 assert_eq!(
3268 manifest.stage_status[1].seed,
3269 effective_stage_seed(req.seed.unwrap(), Some(7))
3270 );
3271 }
3272
3273 #[test]
3274 fn finalize_versions_outputs_and_records_stage_seeds() {
3275 let dir = tempfile::tempdir().unwrap();
3276 let db = db();
3277 let req = request(vec![TransitionMode::Smooth]);
3278 let job_dir = dir.path().join("job");
3279 let row = persist_job(&db, &job_dir, "01JBR55VERSION", &req, ChainJobState::Queued);
3280 let executor = Arc::new(FakeExecutor {
3281 calls: AtomicUsize::new(0),
3282 cancel_on_progress: AtomicBool::new(false),
3283 });
3284 let deps = deps(
3285 db,
3286 dir.path().join("jobs"),
3287 executor,
3288 Arc::new(FakeProbe(AtomicUsize::new(0))),
3289 );
3290
3291 execute_job(&deps, &row, 0).unwrap();
3292 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3293 let first_seed = manifest.stage_status[0].seed;
3294 assert_eq!(manifest.finalizes[0].output, "final/output-1.mp4");
3295
3296 let second = finalize_job(&deps, &row, &mut manifest).unwrap();
3297
3298 assert_eq!(second, "final/output-2.mp4");
3299 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3300 assert_eq!(manifest.finalizes.len(), 2);
3301 assert_eq!(manifest.finalizes[0].stage_seeds, vec![first_seed]);
3302 assert_eq!(manifest.finalizes[1].stage_seeds, vec![first_seed]);
3303 assert!(job_dir.join("final/output-1.mp4").exists());
3304 assert!(job_dir.join("final/output-2.mp4").exists());
3305 }
3306
3307 #[test]
3308 fn completed_fade_successor_is_reencoded_after_previous_stage_changes_boundary() {
3309 let dir = tempfile::tempdir().unwrap();
3310 let job_dir = dir.path().join("job");
3311 let layout = JobDirLayout::new(job_dir.clone());
3312 layout.ensure_root().unwrap();
3313 let req = request(vec![TransitionMode::Smooth, TransitionMode::Fade]);
3314 let mut manifest = ChainJobManifest::new("01JBR55FADE".into(), 1_000, &req).unwrap();
3315
3316 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
3317 write_stage_artifacts(&layout, &mut manifest, 1, &outcome(80), &req).unwrap();
3318 let before = std::fs::read(layout.segment_path(1)).unwrap();
3319
3320 write_frames_to_dir(&layout.boundary_out_dir(0), &[frame(210), frame(220)]).unwrap();
3321 maybe_reencode_next_after_fade(&layout, &manifest, 0, &req).unwrap();
3322
3323 let after = std::fs::read(layout.segment_path(1)).unwrap();
3324 assert_ne!(
3325 before, after,
3326 "changing outgoing fade boundary must re-encode completed successor segment"
3327 );
3328 }
3329
3330 #[test]
3331 fn splice_retake_skips_completed_successors_and_only_reencodes_fade_boundary() {
3332 for (next_transition, successor_should_change, id) in [
3333 (TransitionMode::Fade, true, "01JBR55SPLICEFADE"),
3334 (TransitionMode::Cut, false, "01JBR55SPLICECUT"),
3335 ] {
3336 let dir = tempfile::tempdir().unwrap();
3337 let db = db();
3338 let req = request(vec![
3339 TransitionMode::Smooth,
3340 next_transition,
3341 TransitionMode::Cut,
3342 TransitionMode::Cut,
3343 ]);
3344 let job_dir = dir.path().join("job");
3345 let row = persist_job(&db, &job_dir, id, &req, ChainJobState::Queued);
3346 let executor = Arc::new(FakeExecutor {
3347 calls: AtomicUsize::new(0),
3348 cancel_on_progress: AtomicBool::new(false),
3349 });
3350 let deps = deps(
3351 db,
3352 dir.path().join("jobs"),
3353 executor.clone(),
3354 Arc::new(FakeProbe(AtomicUsize::new(0))),
3355 );
3356 execute_job(&deps, &row, 0).unwrap();
3357
3358 let layout = JobDirLayout::new(job_dir.clone());
3359 let before_1 = std::fs::read(layout.segment_path(1)).unwrap();
3360 let before_2 = std::fs::read(layout.segment_path(2)).unwrap();
3361 let before_3 = std::fs::read(layout.segment_path(3)).unwrap();
3362 let calls_before = executor.calls.load(Ordering::SeqCst);
3363 let db = deps.db.as_ref().as_ref().unwrap();
3364 let updated = apply_retake(
3365 db,
3366 dir.path(),
3367 &row.id,
3368 &RetakeRequest {
3369 stage_idx: 0,
3370 mode: RetakeMode::Splice,
3371 seed_offset: Some(9),
3372 prompt: None,
3373 },
3374 )
3375 .unwrap();
3376
3377 execute_job(&deps, &updated, 0).unwrap();
3378
3379 assert_eq!(
3380 executor.calls.load(Ordering::SeqCst) - calls_before,
3381 1,
3382 "splice retake must render exactly the target stage for {id}"
3383 );
3384 let after_1 = std::fs::read(layout.segment_path(1)).unwrap();
3385 let after_2 = std::fs::read(layout.segment_path(2)).unwrap();
3386 let after_3 = std::fs::read(layout.segment_path(3)).unwrap();
3387 if successor_should_change {
3388 assert_ne!(
3389 before_1, after_1,
3390 "fade successor must be re-encoded at the edited boundary"
3391 );
3392 } else {
3393 assert_eq!(
3394 before_1, after_1,
3395 "cut successor must not be touched by splice retake"
3396 );
3397 }
3398 assert_eq!(before_2, after_2, "stage N+2 must remain untouched");
3399 assert_eq!(before_3, after_3, "stage N+3 must remain untouched");
3400 }
3401 }
3402
3403 #[test]
3404 fn disk_write_failure_marks_job_and_stage_failed_with_error() {
3405 let dir = tempfile::tempdir().unwrap();
3406 let db = db();
3407 let req = request(vec![TransitionMode::Smooth]);
3408 let job_dir = dir.path().join("job");
3409 let row = persist_job(&db, &job_dir, "01JBR55DISK", &req, ChainJobState::Queued);
3410 std::fs::write(job_dir.join("stages"), b"not a directory").unwrap();
3411 let executor = Arc::new(FakeExecutor {
3412 calls: AtomicUsize::new(0),
3413 cancel_on_progress: AtomicBool::new(false),
3414 });
3415 let deps = deps(
3416 db,
3417 dir.path().join("jobs"),
3418 executor,
3419 Arc::new(FakeProbe(AtomicUsize::new(0))),
3420 );
3421
3422 execute_job(&deps, &row, 0).unwrap();
3423
3424 let db = deps.db.as_ref().as_ref().unwrap();
3425 let job_after = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3426 assert_eq!(job_after.state, ChainJobState::Failed);
3427 assert!(
3428 job_after.error.as_deref().is_some_and(
3429 |error| error.contains("stages/000") || error.contains("Not a directory")
3430 ),
3431 "explicit filesystem error should be persisted, got {:?}",
3432 job_after.error
3433 );
3434 let stage = chain_jobs::stages_for_job(db, &row.id).unwrap().remove(0);
3435 assert_eq!(stage.state, StageState::Failed);
3436 assert!(stage
3437 .error
3438 .as_deref()
3439 .unwrap_or_default()
3440 .contains("stages/000"));
3441 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3442 assert_eq!(manifest.stage_status[0].state, StageState::Failed);
3443 }
3444
3445 #[test]
3446 fn finalize_failure_then_resume_reenters_finalize_without_rerendering() {
3447 let dir = tempfile::tempdir().unwrap();
3448 let db = db();
3449 let req = request(vec![TransitionMode::Smooth]);
3450 let job_dir = dir.path().join("job");
3451 let row = persist_job(
3452 &db,
3453 &job_dir,
3454 "01JBR55FINALRETRY",
3455 &req,
3456 ChainJobState::Queued,
3457 );
3458 std::fs::write(job_dir.join("final"), b"not a directory").unwrap();
3459 let executor = Arc::new(FakeExecutor {
3460 calls: AtomicUsize::new(0),
3461 cancel_on_progress: AtomicBool::new(false),
3462 });
3463 let deps = deps(
3464 db,
3465 dir.path().join("jobs"),
3466 executor.clone(),
3467 Arc::new(FakeProbe(AtomicUsize::new(0))),
3468 );
3469
3470 execute_job(&deps, &row, 0).unwrap();
3471 let db = deps.db.as_ref().as_ref().unwrap();
3472 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3473 assert_eq!(failed.state, ChainJobState::Failed);
3474 assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
3475 assert!(ChainJobManifest::read_from_dir(&job_dir)
3476 .unwrap()
3477 .finalizes
3478 .is_empty());
3479
3480 std::fs::remove_file(job_dir.join("final")).unwrap();
3481 assert!(chain_jobs::try_transition(
3482 db,
3483 &row.id,
3484 &[ChainJobState::Failed],
3485 ChainJobState::Queued,
3486 None,
3487 now_ms_i64(),
3488 )
3489 .unwrap());
3490 let queued = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3491 execute_job(&deps, &queued, req.stages.len() as u32).unwrap();
3492
3493 let completed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3494 assert_eq!(completed.state, ChainJobState::Completed);
3495 assert_eq!(
3496 executor.calls.load(Ordering::SeqCst),
3497 1,
3498 "resume from completed stages must not render again"
3499 );
3500 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3501 assert_eq!(manifest.finalizes.len(), 1);
3502 assert_eq!(manifest.finalizes[0].output, "final/output-1.mp4");
3503 }
3504
3505 #[test]
3506 fn boundary_cancel_keeps_completed_stage_artifacts_and_stops_next_stage() {
3507 let dir = tempfile::tempdir().unwrap();
3508 let db = db();
3509 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3510 let job_dir = dir.path().join("job");
3511 let row = persist_job(
3512 &db,
3513 &job_dir,
3514 "01JBR55BOUNDARY",
3515 &req,
3516 ChainJobState::Queued,
3517 );
3518 let cancel = Arc::new(CancelRegistry::new());
3519 let executor = Arc::new(BoundaryCancelExecutor {
3520 calls: AtomicUsize::new(0),
3521 cancel: cancel.clone(),
3522 job_id: row.id.clone(),
3523 });
3524 let deps = RunnerDeps {
3525 db: Arc::new(Some(db)),
3526 jobs_root: dir.path().join("jobs"),
3527 executor,
3528 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
3529 events: Arc::new(JobEventBus::new()),
3530 cancel,
3531 job_locks: Arc::new(JobMutationLocks::new()),
3532 claims: Arc::new(EphemeralClaims::default()),
3533 output_dir: None,
3534 };
3535
3536 execute_job(&deps, &row, 0).unwrap();
3537
3538 let db = deps.db.as_ref().as_ref().unwrap();
3539 assert_eq!(
3540 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3541 ChainJobState::Cancelled
3542 );
3543 assert!(job_dir.join("stages/000/segment.mp4").exists());
3544 assert!(!job_dir.join("stages/001/segment.mp4").exists());
3545 }
3546
3547 #[test]
3548 fn yields_between_stages_when_small_jobs_are_waiting_and_cleans_bus_at_terminal() {
3549 let dir = tempfile::tempdir().unwrap();
3550 let db = db();
3551 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3552 let job_dir = dir.path().join("job");
3553 let row = persist_job(&db, &job_dir, "01JBR55YIELD", &req, ChainJobState::Queued);
3554 let executor = Arc::new(FakeExecutor {
3555 calls: AtomicUsize::new(0),
3556 cancel_on_progress: AtomicBool::new(false),
3557 });
3558 let deps = deps(
3559 db,
3560 dir.path().join("jobs"),
3561 executor,
3562 Arc::new(FakeProbe(AtomicUsize::new(2))),
3563 );
3564 let mut rx = deps.events.subscribe_persistent_for_tests(&row.id);
3565
3566 execute_job(&deps, &row, 0).unwrap();
3567
3568 let mut events = Vec::new();
3569 loop {
3570 match rx.try_recv() {
3571 Ok(event) => events.push(event),
3572 Err(tokio::sync::broadcast::error::TryRecvError::Empty)
3573 | Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
3574 Err(err) => panic!("unexpected broadcast receive error: {err}"),
3575 }
3576 }
3577 assert!(matches!(
3578 events.as_slice(),
3579 [
3580 ChainJobEvent::StateChanged {
3581 state: ChainJobState::Running,
3582 ..
3583 },
3584 ChainJobEvent::StageStart { stage_idx: 0 },
3585 ChainJobEvent::DenoiseStep { stage_idx: 0, .. },
3586 ChainJobEvent::StageDone { stage_idx: 0, .. },
3587 ..
3588 ]
3589 ));
3590 assert!(events.iter().any(|event| {
3591 matches!(
3592 event,
3593 ChainJobEvent::Yielded {
3594 pending_small_jobs: 2
3595 }
3596 )
3597 }));
3598 assert!(!deps.events.senders.lock().unwrap().contains_key(&row.id));
3599 }
3600
3601 #[test]
3602 fn resume_carry_from_disk_loads_smooth_tail_and_rejects_cut_or_fade_carry() {
3603 let dir = tempfile::tempdir().unwrap();
3604 let job_dir = dir.path().join("job");
3605 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3606 let mut manifest = ChainJobManifest::new("01JBR55CARRY".into(), 1_000, &req).unwrap();
3607 manifest.stage_status[0].state = StageState::Completed;
3608 manifest.stage_status[0].segment = Some("stages/000/segment.mp4".into());
3609 manifest.stage_status[0].tail_frames = Some(1);
3610 let layout = JobDirLayout::new(job_dir.clone());
3611 layout.ensure_stage_dirs(0).unwrap();
3612 frame(77).save(layout.tail_dir(0).join("000.png")).unwrap();
3613 manifest.write_atomic(&job_dir).unwrap();
3614
3615 let carry = resume_carry_from_disk(&job_dir, &manifest, 1)
3616 .unwrap()
3617 .expect("smooth continuation should reload tail");
3618 assert_eq!(carry.frames, 1);
3619 assert_eq!(carry.tail_rgb_frames[0].get_pixel(0, 0).0, [77, 77, 77]);
3620
3621 for transition in [TransitionMode::Cut, TransitionMode::Fade] {
3622 let req = request(vec![TransitionMode::Smooth, transition]);
3623 let mut manifest = ChainJobManifest::new("01JBR55CARRY".into(), 1_000, &req).unwrap();
3624 manifest.stage_status[0].state = StageState::Completed;
3625 manifest.stage_status[0].tail_frames = Some(1);
3626 assert!(resume_carry_from_disk(&job_dir, &manifest, 1)
3627 .unwrap()
3628 .is_none());
3629 }
3630 }
3631
3632 #[test]
3633 fn malicious_manifest_resume_fails_and_job_remains_deletable() {
3634 let dir = tempfile::tempdir().unwrap();
3635 let db = db();
3636 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3637 let job_dir = dir.path().join("job");
3638 let row = persist_job(
3639 &db,
3640 &job_dir,
3641 "01JBR55TRAVERSAL",
3642 &req,
3643 ChainJobState::Queued,
3644 );
3645 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3646 manifest.stage_status[0].state = StageState::Completed;
3647 manifest.stage_status[0].segment = Some("../outside.mp4".into());
3648 manifest.stage_status[0].tail_frames = Some(1);
3649 manifest.write_atomic(&job_dir).unwrap();
3650 let executor = Arc::new(FakeExecutor {
3651 calls: AtomicUsize::new(0),
3652 cancel_on_progress: AtomicBool::new(false),
3653 });
3654 let deps = deps(
3655 db,
3656 dir.path().join("jobs"),
3657 executor.clone(),
3658 Arc::new(FakeProbe(AtomicUsize::new(0))),
3659 );
3660
3661 execute_job(&deps, &row, 0).unwrap();
3662
3663 let db = deps.db.as_ref().as_ref().unwrap();
3664 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3665 assert_eq!(failed.state, ChainJobState::Failed);
3666 assert!(
3667 failed
3668 .error
3669 .as_deref()
3670 .is_some_and(|err| err.contains("must be relative")),
3671 "expected traversal error to be persisted, got {:?}",
3672 failed.error
3673 );
3674 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
3675 assert!(chain_jobs::delete_job_not_running(db, &row.id).unwrap());
3676 assert!(chain_jobs::get_job(db, &row.id).unwrap().is_none());
3677 }
3678
3679 #[test]
3680 fn startup_reconcile_flips_running_to_interrupted_and_repairs_from_manifest() {
3681 let dir = tempfile::tempdir().unwrap();
3682 let db = db();
3683 let req = request(vec![TransitionMode::Smooth]);
3684 let job_dir = dir.path().join("jobs/01JBR55RECON");
3685 let row = persist_job(&db, &job_dir, "01JBR55RECON", &req, ChainJobState::Running);
3686 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3687 manifest.stage_status[0].state = StageState::Completed;
3688 manifest.stage_status[0].frames_emitted = Some(9);
3689 manifest.stage_status[0].generation_time_ms = Some(123);
3690 manifest.stage_status[0].segment = Some("stages/000/segment.mp4".into());
3691 manifest.write_atomic(&job_dir).unwrap();
3692
3693 let (flipped, repaired) = startup_reconcile(&db, &dir.path().join("jobs")).unwrap();
3694
3695 assert_eq!(flipped, 1);
3696 assert!(repaired >= 1);
3697 let row_after = chain_jobs::get_job(&db, &row.id).unwrap().unwrap();
3698 assert_eq!(row_after.state, ChainJobState::Interrupted);
3699 assert_eq!(row_after.current_stage, 1);
3700 assert_eq!(
3701 row_after.error.as_deref(),
3702 Some("server restarted while chain job was running")
3703 );
3704 let stage = chain_jobs::stages_for_job(&db, &row.id).unwrap().remove(0);
3705 assert_eq!(stage.state, StageState::Completed);
3706 assert_eq!(
3707 stage.segment_rel_path.as_deref(),
3708 Some("stages/000/segment.mp4")
3709 );
3710 }
3711
3712 #[test]
3713 fn ephemeral_claim_guard_releases_on_drop() {
3714 let claims = Arc::new(EphemeralClaims::default());
3715 {
3716 let _guard = claims.claim("01JBR55CLAIM");
3717 assert!(claims.is_claimed("01JBR55CLAIM"));
3718 }
3719 assert!(!claims.is_claimed("01JBR55CLAIM"));
3720 }
3721
3722 #[test]
3723 fn create_job_with_params_persists_ephemeral_manifest_and_stage_rows() {
3724 let dir = tempfile::tempdir().unwrap();
3725 let db = db();
3726 let req = request(vec![TransitionMode::Smooth, TransitionMode::Cut]);
3727
3728 let row = create_job_with_params(
3729 &db,
3730 dir.path(),
3731 CreateJobParams {
3732 id: "01JBR55CREATE".into(),
3733 ephemeral: true,
3734 request: req.clone(),
3735 },
3736 )
3737 .unwrap();
3738
3739 assert_eq!(row.id, "01JBR55CREATE");
3740 assert_eq!(row.state, ChainJobState::Queued);
3741 assert_eq!(row.stage_count, 2);
3742 assert!(row.job_dir.ends_with("01JBR55CREATE"));
3743 let manifest = ChainJobManifest::read_from_dir(&row.job_dir).unwrap();
3744 assert!(manifest.ephemeral);
3745 assert_eq!(manifest.request().unwrap(), req);
3746 let stages = chain_jobs::stages_for_job(&db, &row.id).unwrap();
3747 assert_eq!(stages.len(), 2);
3748 assert_eq!(stages[0].state, StageState::Pending);
3749 }
3750
3751 #[test]
3752 fn startup_gc_sweep_removes_non_running_ephemerals_only() {
3753 let dir = tempfile::tempdir().unwrap();
3754 let db = db();
3755 let req = request(vec![TransitionMode::Smooth]);
3756 let jobs_root = dir.path().join("jobs");
3757 let ephemeral_done = persist_job(
3758 &db,
3759 &jobs_root.join("01JBR55EPHDONE"),
3760 "01JBR55EPHDONE",
3761 &req,
3762 ChainJobState::Completed,
3763 );
3764 let mut manifest = ChainJobManifest::read_from_dir(&ephemeral_done.job_dir).unwrap();
3765 manifest.ephemeral = true;
3766 manifest.write_atomic(&ephemeral_done.job_dir).unwrap();
3767 let ephemeral_running = persist_job(
3768 &db,
3769 &jobs_root.join("01JBR55EPHRUN"),
3770 "01JBR55EPHRUN",
3771 &req,
3772 ChainJobState::Running,
3773 );
3774 let mut manifest = ChainJobManifest::read_from_dir(&ephemeral_running.job_dir).unwrap();
3775 manifest.ephemeral = true;
3776 manifest.write_atomic(&ephemeral_running.job_dir).unwrap();
3777 let durable_done = persist_job(
3778 &db,
3779 &jobs_root.join("01JBR55DURABLE"),
3780 "01JBR55DURABLE",
3781 &req,
3782 ChainJobState::Completed,
3783 );
3784
3785 let outcome = startup_gc_sweep(&db, &jobs_root).unwrap();
3786
3787 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3788 assert!(chain_jobs::get_job(&db, &ephemeral_done.id)
3789 .unwrap()
3790 .is_none());
3791 assert!(!ephemeral_done.job_dir.exists());
3792 assert!(chain_jobs::get_job(&db, &ephemeral_running.id)
3793 .unwrap()
3794 .is_some());
3795 assert!(ephemeral_running.job_dir.exists());
3796 assert!(chain_jobs::get_job(&db, &durable_done.id)
3797 .unwrap()
3798 .is_some());
3799 assert!(durable_done.job_dir.exists());
3800 }
3801
3802 #[test]
3803 fn gc_sweeps_unclaimed_ephemeral_after_grace_but_preserves_live_claim() {
3804 let dir = tempfile::tempdir().unwrap();
3805 let db = db();
3806 let req = request(vec![TransitionMode::Smooth]);
3807 let jobs_root = dir.path().join("jobs");
3808 let old = persist_job(
3809 &db,
3810 &jobs_root.join("01JBR55OLDCLAIMLESS"),
3811 "01JBR55OLDCLAIMLESS",
3812 &req,
3813 ChainJobState::Completed,
3814 );
3815 let mut manifest = ChainJobManifest::read_from_dir(&old.job_dir).unwrap();
3816 manifest.ephemeral = true;
3817 manifest.write_atomic(&old.job_dir).unwrap();
3818 chain_jobs::update_job_state(&db, &old.id, ChainJobState::Completed, None, 1_000).unwrap();
3819 let claimed = persist_job(
3820 &db,
3821 &jobs_root.join("01JBR55CLAIMEDGC"),
3822 "01JBR55CLAIMEDGC",
3823 &req,
3824 ChainJobState::Completed,
3825 );
3826 let mut manifest = ChainJobManifest::read_from_dir(&claimed.job_dir).unwrap();
3827 manifest.ephemeral = true;
3828 manifest.write_atomic(&claimed.job_dir).unwrap();
3829
3830 let deps = deps(
3831 db,
3832 jobs_root,
3833 Arc::new(FakeExecutor {
3834 calls: AtomicUsize::new(0),
3835 cancel_on_progress: AtomicBool::new(false),
3836 }),
3837 Arc::new(FakeProbe(AtomicUsize::new(0))),
3838 );
3839 let _claim = deps.claims.claim(&claimed.id);
3840 let outcome =
3841 run_gc_pass(&deps, 7, 1_000 + (EPHEMERAL_GRACE_SECS as i64 + 1) * 1_000).unwrap();
3842
3843 let db = deps.db.as_ref().as_ref().unwrap();
3844 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3845 assert!(chain_jobs::get_job(db, &old.id).unwrap().is_none());
3846 assert!(chain_jobs::get_job(db, &claimed.id).unwrap().is_some());
3847 assert!(claimed.job_dir.exists());
3848 }
3849
3850 #[test]
3851 fn gc_respects_ephemeral_grace_and_prunes_only_completed_durable_stages() {
3852 let dir = tempfile::tempdir().unwrap();
3853 let db = db();
3854 let req = request(vec![TransitionMode::Smooth]);
3855 let jobs_root = dir.path().join("jobs");
3856 let recent_ephemeral = persist_job(
3857 &db,
3858 &jobs_root.join("01JBR55RECENTEPH"),
3859 "01JBR55RECENTEPH",
3860 &req,
3861 ChainJobState::Completed,
3862 );
3863 let mut manifest = ChainJobManifest::read_from_dir(&recent_ephemeral.job_dir).unwrap();
3864 manifest.ephemeral = true;
3865 manifest.write_atomic(&recent_ephemeral.job_dir).unwrap();
3866 let now = 1_000 + 8 * 86_400_000;
3867 chain_jobs::update_job_state(
3868 &db,
3869 &recent_ephemeral.id,
3870 ChainJobState::Completed,
3871 None,
3872 now - (EPHEMERAL_GRACE_SECS as i64 * 1_000) + 100,
3873 )
3874 .unwrap();
3875
3876 let durable = persist_job(
3877 &db,
3878 &jobs_root.join("01JBR55TTL"),
3879 "01JBR55TTL",
3880 &req,
3881 ChainJobState::Completed,
3882 );
3883 let layout = JobDirLayout::new(durable.job_dir.clone());
3884 std::fs::create_dir_all(layout.stage_dir(0)).unwrap();
3885 std::fs::write(layout.stage_dir(0).join("segment.mp4"), b"stage").unwrap();
3886 std::fs::create_dir_all(durable.job_dir.join("final")).unwrap();
3887 std::fs::write(durable.job_dir.join("final/output-1.mp4"), b"final").unwrap();
3888 chain_jobs::update_job_state(&db, &durable.id, ChainJobState::Completed, None, 1_000)
3889 .unwrap();
3890
3891 let deps = deps(
3892 db,
3893 jobs_root,
3894 Arc::new(FakeExecutor {
3895 calls: AtomicUsize::new(0),
3896 cancel_on_progress: AtomicBool::new(false),
3897 }),
3898 Arc::new(FakeProbe(AtomicUsize::new(0))),
3899 );
3900 let outcome = run_gc_pass(&deps, 7, now).unwrap();
3901 let db = deps.db.as_ref().as_ref().unwrap();
3902
3903 assert_eq!(outcome.swept_ephemeral_jobs, 0);
3904 assert_eq!(outcome.pruned_artifact_dirs, 1);
3905 assert!(chain_jobs::get_job(db, &recent_ephemeral.id)
3906 .unwrap()
3907 .is_some());
3908 assert!(recent_ephemeral.job_dir.exists());
3909 assert!(!durable.job_dir.join("stages").exists());
3910 assert!(durable.job_dir.join("final/output-1.mp4").exists());
3911 assert!(durable.job_dir.join("manifest.toml").exists());
3912 assert!(chain_jobs::get_job(db, &durable.id).unwrap().is_some());
3913 }
3914
3915 #[test]
3916 fn gc_exempts_durable_failed_interrupted_cancelled_but_ephemeral_overrides() {
3917 let dir = tempfile::tempdir().unwrap();
3918 let db = db();
3919 let req = request(vec![TransitionMode::Smooth]);
3920 let jobs_root = dir.path().join("jobs");
3921 let mut durable_rows = Vec::new();
3922 for (id, state) in [
3923 ("01JBR55FAILED", ChainJobState::Failed),
3924 ("01JBR55INTR", ChainJobState::Interrupted),
3925 ("01JBR55CANCELLED", ChainJobState::Cancelled),
3926 ] {
3927 let row = persist_job(&db, &jobs_root.join(id), id, &req, state);
3928 std::fs::create_dir_all(row.job_dir.join("stages/000")).unwrap();
3929 durable_rows.push(row);
3930 }
3931 let eph_failed = persist_job(
3932 &db,
3933 &jobs_root.join("01JBR55EPHFAILED"),
3934 "01JBR55EPHFAILED",
3935 &req,
3936 ChainJobState::Failed,
3937 );
3938 let mut manifest = ChainJobManifest::read_from_dir(&eph_failed.job_dir).unwrap();
3939 manifest.ephemeral = true;
3940 manifest.write_atomic(&eph_failed.job_dir).unwrap();
3941 chain_jobs::update_job_state(&db, &eph_failed.id, ChainJobState::Failed, None, 1_000)
3942 .unwrap();
3943
3944 let deps = deps(
3945 db,
3946 jobs_root,
3947 Arc::new(FakeExecutor {
3948 calls: AtomicUsize::new(0),
3949 cancel_on_progress: AtomicBool::new(false),
3950 }),
3951 Arc::new(FakeProbe(AtomicUsize::new(0))),
3952 );
3953 let outcome =
3954 run_gc_pass(&deps, 7, 1_000 + (EPHEMERAL_GRACE_SECS as i64 + 1) * 1_000).unwrap();
3955 let db = deps.db.as_ref().as_ref().unwrap();
3956
3957 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3958 assert!(chain_jobs::get_job(db, &eph_failed.id).unwrap().is_none());
3959 for row in durable_rows {
3960 assert!(chain_jobs::get_job(db, &row.id).unwrap().is_some());
3961 assert!(row.job_dir.join("stages").exists());
3962 }
3963 }
3964
3965 #[tokio::test]
3966 async fn runner_request_gc_replies_between_jobs() {
3967 let dir = tempfile::tempdir().unwrap();
3968 let db = db();
3969 let req = request(vec![TransitionMode::Smooth]);
3970 let jobs_root = dir.path().join("jobs");
3971 let _row = persist_job(
3972 &db,
3973 &jobs_root.join("01JBR55REQUESTGC"),
3974 "01JBR55REQUESTGC",
3975 &req,
3976 ChainJobState::Completed,
3977 );
3978 let deps = RunnerDeps {
3979 db: Arc::new(Some(db)),
3980 jobs_root,
3981 executor: Arc::new(FakeExecutor {
3982 calls: AtomicUsize::new(0),
3983 cancel_on_progress: AtomicBool::new(false),
3984 }),
3985 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
3986 events: Arc::new(JobEventBus::new()),
3987 cancel: Arc::new(CancelRegistry::new()),
3988 job_locks: Arc::new(JobMutationLocks::new()),
3989 claims: Arc::new(EphemeralClaims::default()),
3990 output_dir: None,
3991 };
3992 let handle = spawn_runner(deps);
3993
3994 let outcome = tokio::time::timeout(std::time::Duration::from_secs(1), handle.request_gc())
3995 .await
3996 .expect("GC request should be serviced without waiting for the whole queue")
3997 .unwrap();
3998
3999 assert_eq!(outcome.swept_ephemeral_jobs, 0);
4000 }
4001}