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