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 source_image_name: None,
2042 edit_images: None,
2043 strength: if idx == 0 { chain.strength } else { 1.0 },
2044 mask_image: None,
2045 control_image: None,
2046 control_model: None,
2047 control_scale: 1.0,
2048 expand: None,
2049 original_prompt: chain.original_prompt.clone(),
2050 batch_id: chain.batch_id.clone(),
2051 batch_index: chain.batch_index,
2052 batch_count: chain.batch_count,
2053 lora: None,
2054 frames: Some(stage.frames),
2055 fps: Some(chain.fps),
2056 upscale_model: None,
2057 gif_preview: false,
2058 enable_audio: Some(chain.enable_audio.unwrap_or(false)),
2059 audio_file: None,
2060 audio_file_path: None,
2061 source_video: None,
2062 source_video_path: None,
2063 keyframes: None,
2064 pipeline: None,
2065 loras: None,
2066 retake_range: None,
2067 spatial_upscale: None,
2068 temporal_upscale: None,
2069 placement: chain.placement.clone(),
2070 }
2071}
2072
2073fn maybe_reencode_next_after_fade(
2074 layout: &JobDirLayout,
2075 manifest: &ChainJobManifest,
2076 stage_idx: u32,
2077 effective: &ChainRequest,
2078) -> anyhow::Result<()> {
2079 let next_idx = stage_idx + 1;
2080 let Some(next_stage) = effective.stages.get(next_idx as usize) else {
2081 return Ok(());
2082 };
2083 if !matches!(next_stage.transition, TransitionMode::Fade) {
2084 return Ok(());
2085 }
2086 let Some(next_status) = manifest.stage_status.get(next_idx as usize) else {
2087 return Ok(());
2088 };
2089 if next_status.state != StageState::Completed {
2090 return Ok(());
2091 }
2092 let fade_len = next_stage.fade_frames.unwrap_or(DEFAULT_FADE_FRAMES);
2093 let n = fade_len as usize;
2094 let next_segment = safe_join_manifest_rel(
2095 layout.root(),
2096 next_status
2097 .segment
2098 .as_ref()
2099 .ok_or_else(|| anyhow!("completed next stage {next_idx} has no segment"))?,
2100 )?;
2101 let (_metadata, mut frames) =
2102 mold_inference::ltx2::media::decode_video_frames_from_path(&next_segment)?;
2103 if frames.len() < n {
2104 bail!("next stage {next_idx} segment shorter than fade_len {n}");
2105 }
2106 let new_boundary_out = read_frames_from_dir(&layout.boundary_out_dir(stage_idx), n)?;
2107 let raw_boundary_in = read_frames_from_dir(&layout.boundary_in_dir(next_idx), n)?;
2108 let blended = fade_boundary(&new_boundary_out, &raw_boundary_in, fade_len);
2109 for (idx, frame) in blended.into_iter().enumerate() {
2110 frames[idx] = frame;
2111 }
2112 let bytes = video_enc::encode_mp4(&frames, effective.fps)?;
2113 write_file(&next_segment, &bytes)?;
2114 write_preview_jpeg(&layout.preview_path(next_idx), frames.last().unwrap())?;
2115 Ok(())
2116}
2117
2118fn safe_join_manifest_rel(root: &Path, rel: &str) -> anyhow::Result<PathBuf> {
2119 let path = Path::new(rel);
2120 if path.is_absolute()
2121 || rel.starts_with('/')
2122 || rel.starts_with('\\')
2123 || rel.as_bytes().get(0..3).is_some_and(|b| {
2124 b[0].is_ascii_alphabetic() && b[1] == b':' && matches!(b[2], b'/' | b'\\')
2125 })
2126 || path
2127 .components()
2128 .any(|component| matches!(component, Component::ParentDir))
2129 || rel.split(['/', '\\']).any(|part| part == "..")
2130 {
2131 bail!("manifest artifact path '{rel}' must be relative and must not contain '..'");
2132 }
2133 Ok(root.join(path))
2134}
2135
2136fn write_file(path: &Path, bytes: &[u8]) -> anyhow::Result<()> {
2137 if let Some(parent) = path.parent() {
2138 std::fs::create_dir_all(parent)
2139 .with_context(|| format!("creating directory '{}'", parent.display()))?;
2140 }
2141 std::fs::write(path, bytes).with_context(|| format!("writing '{}'", path.display()))
2142}
2143
2144fn write_frames_to_dir(dir: &Path, frames: &[RgbImage]) -> anyhow::Result<()> {
2145 std::fs::create_dir_all(dir).with_context(|| format!("creating '{}'", dir.display()))?;
2146 for entry in std::fs::read_dir(dir).with_context(|| format!("reading '{}'", dir.display()))? {
2147 let path = entry?.path();
2148 if path.extension().and_then(|s| s.to_str()) == Some("png") {
2149 std::fs::remove_file(&path)
2150 .with_context(|| format!("removing stale frame '{}'", path.display()))?;
2151 }
2152 }
2153 for (idx, frame) in frames.iter().enumerate() {
2154 let path = dir.join(format!("{idx:03}.png"));
2155 frame
2156 .save_with_format(&path, image::ImageFormat::Png)
2157 .with_context(|| format!("writing PNG frame '{}'", path.display()))?;
2158 }
2159 Ok(())
2160}
2161
2162fn read_frames_from_dir(dir: &Path, expected: usize) -> anyhow::Result<Vec<RgbImage>> {
2163 let mut paths: Vec<PathBuf> = std::fs::read_dir(dir)
2164 .with_context(|| format!("reading frame directory '{}'", dir.display()))?
2165 .map(|entry| entry.map(|entry| entry.path()))
2166 .collect::<std::io::Result<Vec<_>>>()?;
2167 paths.retain(|path| path.extension().and_then(|s| s.to_str()) == Some("png"));
2168 paths.sort();
2169 if paths.len() < expected {
2170 bail!(
2171 "frame directory '{}' has {} PNG(s), expected at least {expected}",
2172 dir.display(),
2173 paths.len()
2174 );
2175 }
2176 paths
2177 .into_iter()
2178 .take(expected)
2179 .map(|path| {
2180 image::open(&path)
2181 .with_context(|| format!("decoding PNG frame '{}'", path.display()))
2182 .map(|img| img.to_rgb8())
2183 })
2184 .collect()
2185}
2186
2187fn write_preview_jpeg(path: &Path, frame: &RgbImage) -> anyhow::Result<()> {
2188 let mut bytes = Vec::new();
2189 {
2190 let mut cursor = Cursor::new(&mut bytes);
2191 let encoder = JpegEncoder::new_with_quality(&mut cursor, 85);
2192 encoder.write_image(
2193 frame.as_raw(),
2194 frame.width(),
2195 frame.height(),
2196 image::ExtendedColorType::Rgb8,
2197 )?;
2198 }
2199 write_file(path, &bytes)
2200}
2201
2202fn trim_audio_front(
2203 audio: &mut Option<NativeAudioTrack>,
2204 frames: u32,
2205 fps: u32,
2206) -> anyhow::Result<()> {
2207 let Some(track) = audio.as_mut() else {
2208 return Ok(());
2209 };
2210 let samples = samples_for_frames(track, frames, fps);
2211 if track.interleaved_samples.len() < samples {
2212 bail!("audio sidecar too short for front trim");
2213 }
2214 track.interleaved_samples.drain(0..samples);
2215 Ok(())
2216}
2217
2218fn trim_audio_back(
2219 audio: &mut Option<NativeAudioTrack>,
2220 frames: u32,
2221 fps: u32,
2222) -> anyhow::Result<()> {
2223 let Some(track) = audio.as_mut() else {
2224 return Ok(());
2225 };
2226 let samples = samples_for_frames(track, frames, fps);
2227 if track.interleaved_samples.len() < samples {
2228 bail!("audio sidecar too short for back trim");
2229 }
2230 let keep = track.interleaved_samples.len() - samples;
2231 track.interleaved_samples.truncate(keep);
2232 Ok(())
2233}
2234
2235fn write_audio_boundary_out(
2236 layout: &JobDirLayout,
2237 stage_idx: u32,
2238 audio: &Option<NativeAudioTrack>,
2239 frames: u32,
2240 fps: u32,
2241) -> anyhow::Result<()> {
2242 let Some(track) = audio.as_ref() else {
2243 return Ok(());
2244 };
2245 let samples = samples_for_frames(track, frames, fps);
2246 if track.interleaved_samples.len() < samples {
2247 bail!("audio sidecar too short for boundary-out");
2248 }
2249 let mut boundary = track.clone();
2250 boundary.interleaved_samples =
2251 boundary.interleaved_samples[boundary.interleaved_samples.len() - samples..].to_vec();
2252 write_audio_sidecar(
2253 &layout.boundary_out_dir(stage_idx).join("audio.pcm"),
2254 &boundary,
2255 )
2256}
2257
2258fn blend_audio_front_from_previous_boundary(
2259 layout: &JobDirLayout,
2260 stage_idx: u32,
2261 audio: &mut Option<NativeAudioTrack>,
2262 frames: u32,
2263 fps: u32,
2264) -> anyhow::Result<()> {
2265 let Some(track) = audio.as_mut() else {
2266 return Ok(());
2267 };
2268 let samples = samples_for_frames(track, frames, fps);
2269 if samples == 0 {
2270 return Ok(());
2271 }
2272 if track.interleaved_samples.len() < samples {
2273 bail!("audio sidecar too short for boundary-in");
2274 }
2275 let mut boundary_in = track.clone();
2276 boundary_in.interleaved_samples = track.interleaved_samples[..samples].to_vec();
2277 write_audio_sidecar(
2278 &layout.boundary_in_dir(stage_idx).join("audio.pcm"),
2279 &boundary_in,
2280 )?;
2281 let previous = read_audio_sidecar(&layout.boundary_out_dir(stage_idx - 1).join("audio.pcm"))?;
2282 if previous.sample_rate != track.sample_rate || previous.channels != track.channels {
2283 bail!("fade audio boundary format mismatch");
2284 }
2285 if previous.interleaved_samples.len() < samples {
2286 bail!("previous fade audio boundary too short");
2287 }
2288 let channels = track.channels as usize;
2289 let frame_samples = channels.max(1);
2290 let denom = (samples / frame_samples).max(1) as f32;
2291 for sample_idx in 0..samples {
2292 let t = (sample_idx / frame_samples) as f32 / denom;
2293 let prior = previous.interleaved_samples[sample_idx];
2294 let next = track.interleaved_samples[sample_idx];
2295 track.interleaved_samples[sample_idx] = prior * (1.0 - t) + next * t;
2296 }
2297 Ok(())
2298}
2299
2300fn samples_for_frames(track: &NativeAudioTrack, frames: u32, fps: u32) -> usize {
2301 if fps == 0 {
2302 return 0;
2303 }
2304 ((track.sample_rate as u64 * frames as u64) / fps as u64) as usize * track.channels as usize
2305}
2306
2307fn write_audio_sidecar(path: &Path, track: &NativeAudioTrack) -> anyhow::Result<()> {
2308 let mut bytes = Vec::with_capacity(16 + track.interleaved_samples.len() * 4);
2309 bytes.extend_from_slice(AUDIO_SIDECAR_MAGIC);
2310 bytes.extend_from_slice(&track.sample_rate.to_le_bytes());
2311 bytes.extend_from_slice(&track.channels.to_le_bytes());
2312 bytes.extend_from_slice(&0u16.to_le_bytes());
2313 for sample in &track.interleaved_samples {
2314 bytes.extend_from_slice(&sample.to_le_bytes());
2315 }
2316 write_file(path, &bytes)
2317}
2318
2319fn read_audio_sidecar(path: &Path) -> anyhow::Result<NativeAudioTrack> {
2320 let bytes = std::fs::read(path).with_context(|| format!("reading '{}'", path.display()))?;
2321 if bytes.len() < 16 || &bytes[..8] != AUDIO_SIDECAR_MAGIC {
2322 bail!("audio sidecar '{}' has unsupported format", path.display());
2323 }
2324 let sample_rate = u32::from_le_bytes(bytes[8..12].try_into().unwrap());
2325 let channels = u16::from_le_bytes(bytes[12..14].try_into().unwrap());
2326 let sample_bytes = &bytes[16..];
2327 if sample_bytes.len() % 4 != 0 {
2328 bail!(
2329 "audio sidecar '{}' has truncated f32 sample",
2330 path.display()
2331 );
2332 }
2333 let interleaved_samples = sample_bytes
2334 .chunks_exact(4)
2335 .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
2336 .collect();
2337 Ok(NativeAudioTrack {
2338 interleaved_samples,
2339 sample_rate,
2340 channels,
2341 })
2342}
2343
2344fn select_worker_for_stage(gpu_pool: &GpuPool, model: &str) -> Option<Arc<GpuWorker>> {
2345 let est = crate::queue::estimate_model_vram(model);
2346 gpu_pool.select_worker(model, est)
2347}
2348
2349struct WorkerInFlightGuard {
2350 worker: Arc<GpuWorker>,
2351}
2352
2353impl WorkerInFlightGuard {
2354 fn new(worker: Arc<GpuWorker>) -> Self {
2355 worker.in_flight.fetch_add(1, Ordering::SeqCst);
2356 Self { worker }
2357 }
2358}
2359
2360impl Drop for WorkerInFlightGuard {
2361 fn drop(&mut self) {
2362 self.worker.in_flight.fetch_sub(1, Ordering::SeqCst);
2363 }
2364}
2365
2366struct WorkerActiveGenerationGuard {
2367 worker: Arc<GpuWorker>,
2368}
2369
2370impl WorkerActiveGenerationGuard {
2371 fn new(worker: Arc<GpuWorker>, model: &str, prompt: &str) -> anyhow::Result<Self> {
2372 let mut active = worker
2373 .active_generation
2374 .write()
2375 .map_err(|e| anyhow!("active_generation lock poisoned: {e}"))?;
2376 *active = Some(ActiveGeneration {
2377 model: model.to_string(),
2378 prompt_sha256: format!("{:x}", Sha256::digest(prompt.as_bytes())),
2379 started_at_unix_ms: now_ms_u64(),
2380 started_at: std::time::Instant::now(),
2381 });
2382 drop(active);
2383 Ok(Self { worker })
2384 }
2385}
2386
2387impl Drop for WorkerActiveGenerationGuard {
2388 fn drop(&mut self) {
2389 if let Ok(mut active) = self.worker.active_generation.write() {
2390 *active = None;
2391 }
2392 }
2393}
2394
2395fn now_ms_u64() -> u64 {
2396 mold_core::time::now_epoch_ms_u64()
2397}
2398
2399fn now_ms_i64() -> i64 {
2400 mold_core::time::now_epoch_ms()
2401}
2402
2403#[cfg(test)]
2404mod tests {
2405 use super::*;
2406 use image::{Rgb, RgbImage};
2407 use mold_core::chain::{ChainStage, TransitionMode};
2408 use mold_core::types::OutputFormat;
2409 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2410
2411 fn db() -> MetadataDb {
2412 MetadataDb::open_in_memory().unwrap()
2413 }
2414
2415 fn stage(prompt: &str, transition: TransitionMode) -> ChainStage {
2416 ChainStage {
2417 prompt: prompt.into(),
2418 frames: 9,
2419 source_image: None,
2420 negative_prompt: None,
2421 seed_offset: None,
2422 transition,
2423 fade_frames: Some(2),
2424 model: None,
2425 loras: vec![],
2426 references: vec![],
2427 }
2428 }
2429
2430 fn request(transitions: Vec<TransitionMode>) -> ChainRequest {
2431 ChainRequest {
2432 model: "ltx-2-19b-distilled:fp8".into(),
2433 stages: transitions
2434 .into_iter()
2435 .enumerate()
2436 .map(|(idx, transition)| stage(&format!("stage {idx}"), transition))
2437 .collect(),
2438 motion_tail_frames: 1,
2439 width: 64,
2440 height: 48,
2441 fps: 8,
2442 seed: Some(42),
2443 steps: 2,
2444 guidance: 1.0,
2445 strength: 1.0,
2446 output_format: OutputFormat::Mp4,
2447 placement: None,
2448 original_prompt: None,
2449 batch_id: None,
2450 batch_index: None,
2451 batch_count: None,
2452 prompt: None,
2453 total_frames: None,
2454 clip_frames: None,
2455 source_image: None,
2456 enable_audio: None,
2457 }
2458 }
2459
2460 fn job(
2461 id: &str,
2462 state: ChainJobState,
2463 created_at_ms: i64,
2464 dir: PathBuf,
2465 req: &ChainRequest,
2466 ) -> ChainJobRow {
2467 ChainJobRow {
2468 id: id.into(),
2469 state,
2470 model: req.model.clone(),
2471 request_json: serde_json::to_string(req).unwrap(),
2472 job_dir: dir,
2473 stage_count: req.stages.len() as u32,
2474 current_stage: 0,
2475 error: None,
2476 created_at_ms,
2477 updated_at_ms: created_at_ms,
2478 finalized_at_ms: None,
2479 }
2480 }
2481
2482 fn frame(value: u8) -> RgbImage {
2483 RgbImage::from_pixel(64, 48, Rgb([value, value, value]))
2484 }
2485
2486 fn outcome(value: u8) -> StageOutcome {
2487 let frames = (0..9).map(|i| frame(value + i)).collect::<Vec<_>>();
2488 StageOutcome {
2489 tail: ChainTail {
2490 frames: 1,
2491 tail_rgb_frames: vec![frame(value + 8)],
2492 },
2493 frames,
2494 audio: None,
2495 generation_time_ms: 10,
2496 }
2497 }
2498
2499 struct FakeExecutor {
2500 calls: AtomicUsize,
2501 cancel_on_progress: AtomicBool,
2502 }
2503
2504 impl StageExecutor for FakeExecutor {
2505 fn render_stage(
2506 &self,
2507 _model: &str,
2508 _stage_req: &GenerateRequest,
2509 _carry: Option<&ChainTail>,
2510 _motion_tail_frames: u32,
2511 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2512 ) -> anyhow::Result<StageRenderOutcome> {
2513 let call = self.calls.fetch_add(1, Ordering::SeqCst) as u8;
2514 if self.cancel_on_progress.load(Ordering::SeqCst) && progress(1, 2).is_break() {
2515 return Ok(StageRenderOutcome::Cancelled);
2516 }
2517 let _ = progress(1, 2);
2518 Ok(StageRenderOutcome::Done(outcome(10 + call * 20)))
2519 }
2520 }
2521
2522 struct FakeProbe(AtomicUsize);
2523
2524 impl QueueProbe for FakeProbe {
2525 fn small_jobs_waiting(&self) -> usize {
2526 self.0.load(Ordering::SeqCst)
2527 }
2528 }
2529
2530 struct CarryInspectExecutor {
2531 seen: Mutex<Vec<Option<[u8; 3]>>>,
2532 }
2533
2534 impl StageExecutor for CarryInspectExecutor {
2535 fn render_stage(
2536 &self,
2537 _model: &str,
2538 _stage_req: &GenerateRequest,
2539 carry: Option<&ChainTail>,
2540 _motion_tail_frames: u32,
2541 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2542 ) -> anyhow::Result<StageRenderOutcome> {
2543 let pixel = carry
2544 .and_then(|tail| tail.tail_rgb_frames.first())
2545 .map(|frame| frame.get_pixel(0, 0).0);
2546 self.seen.lock().unwrap().push(pixel);
2547 let _ = progress(1, 2);
2548 Ok(StageRenderOutcome::Done(outcome(140)))
2549 }
2550 }
2551
2552 struct BoundaryCancelExecutor {
2553 calls: AtomicUsize,
2554 cancel: Arc<CancelRegistry>,
2555 job_id: String,
2556 }
2557
2558 impl StageExecutor for BoundaryCancelExecutor {
2559 fn render_stage(
2560 &self,
2561 _model: &str,
2562 _stage_req: &GenerateRequest,
2563 _carry: Option<&ChainTail>,
2564 _motion_tail_frames: u32,
2565 progress: &(dyn Fn(u32, u32) -> ControlFlow<()> + Send + Sync),
2566 ) -> anyhow::Result<StageRenderOutcome> {
2567 let call = self.calls.fetch_add(1, Ordering::SeqCst);
2568 let _ = progress(1, 2);
2569 if call == 0 {
2570 assert!(self.cancel.request(&self.job_id));
2571 }
2572 Ok(StageRenderOutcome::Done(outcome(30 + call as u8 * 20)))
2573 }
2574 }
2575
2576 fn deps(
2577 db: MetadataDb,
2578 root: PathBuf,
2579 executor: Arc<FakeExecutor>,
2580 probe: Arc<FakeProbe>,
2581 ) -> RunnerDeps {
2582 RunnerDeps {
2583 db: Arc::new(Some(db)),
2584 jobs_root: root,
2585 executor,
2586 queue_probe: probe,
2587 events: Arc::new(JobEventBus::new()),
2588 cancel: Arc::new(CancelRegistry::new()),
2589 job_locks: Arc::new(JobMutationLocks::new()),
2590 claims: Arc::new(EphemeralClaims::default()),
2591 output_dir: None,
2592 server_events: None,
2593 }
2594 }
2595
2596 fn persist_job(
2597 db: &MetadataDb,
2598 dir: &Path,
2599 id: &str,
2600 req: &ChainRequest,
2601 state: ChainJobState,
2602 ) -> ChainJobRow {
2603 let manifest = ChainJobManifest::new(id.into(), 1_000, req).unwrap();
2604 std::fs::create_dir_all(dir).unwrap();
2605 manifest.write_atomic(dir).unwrap();
2606 let row = job(id, state, 1_000, dir.to_path_buf(), req);
2607 chain_jobs::insert_job(db, &row).unwrap();
2608 for stage in &manifest.stage_status {
2609 chain_jobs::upsert_stage(
2610 db,
2611 &ChainJobStageRow {
2612 job_id: id.into(),
2613 stage_idx: stage.idx,
2614 state: stage.state,
2615 seed: stage.seed,
2616 frames_emitted: None,
2617 generation_time_ms: None,
2618 segment_rel_path: None,
2619 error: None,
2620 updated_at_ms: 1_000,
2621 },
2622 )
2623 .unwrap();
2624 }
2625 row
2626 }
2627
2628 #[test]
2629 fn corrupt_manifest_marks_job_failed_without_rendering() {
2630 let dir = tempfile::tempdir().unwrap();
2631 let db = db();
2632 let req = request(vec![TransitionMode::Smooth]);
2633 let job_dir = dir.path().join("job");
2634 let row = persist_job(&db, &job_dir, "01JBR55CORRUPT", &req, ChainJobState::Queued);
2635 std::fs::write(job_dir.join("manifest.toml"), "not valid toml").unwrap();
2636 let executor = Arc::new(FakeExecutor {
2637 calls: AtomicUsize::new(0),
2638 cancel_on_progress: AtomicBool::new(false),
2639 });
2640 let deps = deps(
2641 db,
2642 dir.path().join("jobs"),
2643 executor.clone(),
2644 Arc::new(FakeProbe(AtomicUsize::new(0))),
2645 );
2646
2647 execute_job(&deps, &row, 0).unwrap();
2648
2649 let db = deps.db.as_ref().as_ref().unwrap();
2650 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2651 assert_eq!(failed.state, ChainJobState::Failed);
2652 assert!(
2653 failed
2654 .error
2655 .as_deref()
2656 .is_some_and(|err| err.contains("manifest TOML parse failed")),
2657 "parse error must be persisted, got {:?}",
2658 failed.error
2659 );
2660 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
2661 }
2662
2663 #[test]
2664 fn resume_carry_error_after_running_marks_job_failed_and_unlocks_row() {
2665 let dir = tempfile::tempdir().unwrap();
2666 let db = db();
2667 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
2668 let job_dir = dir.path().join("job");
2669 let row = persist_job(&db, &job_dir, "01JBR55BADTAIL", &req, ChainJobState::Queued);
2670 let layout = JobDirLayout::new(job_dir.clone());
2671 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2672 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
2673 std::fs::remove_dir_all(layout.tail_dir(0)).unwrap();
2674 chain_jobs::upsert_stage(
2675 &db,
2676 &ChainJobStageRow {
2677 job_id: row.id.clone(),
2678 stage_idx: 0,
2679 state: StageState::Completed,
2680 seed: manifest.stage_status[0].seed,
2681 frames_emitted: manifest.stage_status[0].frames_emitted,
2682 generation_time_ms: manifest.stage_status[0].generation_time_ms,
2683 segment_rel_path: manifest.stage_status[0].segment.clone(),
2684 error: None,
2685 updated_at_ms: 1_000,
2686 },
2687 )
2688 .unwrap();
2689 let executor = Arc::new(FakeExecutor {
2690 calls: AtomicUsize::new(0),
2691 cancel_on_progress: AtomicBool::new(false),
2692 });
2693 let deps = deps(
2694 db,
2695 dir.path().join("jobs"),
2696 executor.clone(),
2697 Arc::new(FakeProbe(AtomicUsize::new(0))),
2698 );
2699
2700 execute_job(&deps, &row, 0).unwrap();
2701
2702 let db = deps.db.as_ref().as_ref().unwrap();
2703 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2704 assert_eq!(failed.state, ChainJobState::Failed);
2705 assert!(
2706 failed
2707 .error
2708 .as_deref()
2709 .is_some_and(|err| err.contains("chain tail directory")),
2710 "resume error must be persisted, got {:?}",
2711 failed.error
2712 );
2713 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
2714 assert!(chain_jobs::set_job_queued(db, &row.id, 1, now_ms_i64()).unwrap());
2715 assert_eq!(
2716 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2717 ChainJobState::Queued
2718 );
2719 assert!(chain_jobs::delete_job(db, &row.id).unwrap());
2720 }
2721
2722 #[test]
2723 fn skipping_completed_stage_reloads_tail_for_next_smooth_stage() {
2724 let dir = tempfile::tempdir().unwrap();
2725 let db = db();
2726 let req = request(vec![
2727 TransitionMode::Smooth,
2728 TransitionMode::Smooth,
2729 TransitionMode::Smooth,
2730 ]);
2731 let job_dir = dir.path().join("job");
2732 let row = persist_job(
2733 &db,
2734 &job_dir,
2735 "01JBR55SKIPCARRY",
2736 &req,
2737 ChainJobState::Queued,
2738 );
2739 let layout = JobDirLayout::new(job_dir.clone());
2740 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2741 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
2742 write_stage_artifacts(&layout, &mut manifest, 1, &outcome(80), &req).unwrap();
2743 for idx in 0..2 {
2744 let status = &manifest.stage_status[idx];
2745 chain_jobs::upsert_stage(
2746 &db,
2747 &ChainJobStageRow {
2748 job_id: row.id.clone(),
2749 stage_idx: idx as u32,
2750 state: status.state,
2751 seed: status.seed,
2752 frames_emitted: status.frames_emitted,
2753 generation_time_ms: status.generation_time_ms,
2754 segment_rel_path: status.segment.clone(),
2755 error: None,
2756 updated_at_ms: 1_000,
2757 },
2758 )
2759 .unwrap();
2760 }
2761 let executor = Arc::new(CarryInspectExecutor {
2762 seen: Mutex::new(Vec::new()),
2763 });
2764 let deps = RunnerDeps {
2765 db: Arc::new(Some(db)),
2766 jobs_root: dir.path().join("jobs"),
2767 executor: executor.clone(),
2768 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
2769 events: Arc::new(JobEventBus::new()),
2770 cancel: Arc::new(CancelRegistry::new()),
2771 job_locks: Arc::new(JobMutationLocks::new()),
2772 claims: Arc::new(EphemeralClaims::default()),
2773 output_dir: None,
2774 server_events: None,
2775 };
2776
2777 execute_job(&deps, &row, 0).unwrap();
2778
2779 assert_eq!(
2780 executor.seen.lock().unwrap().as_slice(),
2781 &[Some([88, 88, 88])],
2782 "stage 2 must receive stage 1's persisted smooth tail"
2783 );
2784 }
2785
2786 #[test]
2787 fn runner_next_queued_job_uses_fifo_order() {
2788 let db = db();
2789 let req = request(vec![TransitionMode::Smooth]);
2790 let newer = job(
2791 "01JBR55NEWER",
2792 ChainJobState::Queued,
2793 2_000,
2794 std::env::temp_dir().join("newer"),
2795 &req,
2796 );
2797 let older = job(
2798 "01JBR55OLDER",
2799 ChainJobState::Queued,
2800 1_000,
2801 std::env::temp_dir().join("older"),
2802 &req,
2803 );
2804 chain_jobs::insert_job(&db, &newer).unwrap();
2805 chain_jobs::insert_job(&db, &older).unwrap();
2806
2807 let got = next_queued_job(&db).unwrap().expect("queued job");
2808 assert_eq!(got.id, older.id);
2809 }
2810
2811 #[test]
2812 fn safe_join_manifest_rel_rejects_windows_and_backslash_traversal() {
2813 let root = Path::new("/tmp/mold-chain-job");
2814 for rel in [
2815 r"C:\tmp\segment.mp4",
2816 r"stages\000\..\segment.mp4",
2817 r"\stages\000\segment.mp4",
2818 ] {
2819 let err = safe_join_manifest_rel(root, rel).unwrap_err();
2820 assert!(
2821 err.to_string().contains("must be relative"),
2822 "expected traversal rejection for {rel:?}, got {err:#}"
2823 );
2824 }
2825 }
2826
2827 #[test]
2828 fn completed_job_subscriptions_are_non_persistent() {
2829 let db = db();
2830 let req = request(vec![TransitionMode::Smooth]);
2831 let row = job(
2832 "01JBR55DONEATTACH",
2833 ChainJobState::Completed,
2834 1_000,
2835 std::env::temp_dir().join("done-attach"),
2836 &req,
2837 );
2838 chain_jobs::insert_job(&db, &row).unwrap();
2839 let events = JobEventBus::new();
2840
2841 let _rx1 = events.subscribe_for_job(&db, &row.id).unwrap();
2842 let _rx2 = events.subscribe_for_job(&db, &row.id).unwrap();
2843
2844 assert!(!events.contains_for_tests(&row.id));
2845 }
2846
2847 #[tokio::test]
2848 async fn subscribe_receives_live_event_buffered_after_attach() {
2849 let (kick_tx, _kick_rx) = tokio::sync::mpsc::unbounded_channel();
2850 let handle = ChainJobRunnerHandle {
2851 kick_tx,
2852 cancel: Arc::new(CancelRegistry::new()),
2853 events: Arc::new(JobEventBus::new()),
2854 job_locks: Arc::new(JobMutationLocks::new()),
2855 claims: Arc::new(EphemeralClaims::default()),
2856 };
2857
2858 let mut rx = handle
2859 .events
2860 .subscribe_persistent_for_tests("01JBR55EVENTS");
2861 handle.events.publish(
2862 "01JBR55EVENTS",
2863 ChainJobEvent::StateChanged {
2864 state: ChainJobState::Running,
2865 error: None,
2866 },
2867 );
2868
2869 let event = rx.recv().await.unwrap();
2870 assert!(matches!(
2871 event,
2872 ChainJobEvent::StateChanged {
2873 state: ChainJobState::Running,
2874 error: None
2875 }
2876 ));
2877 }
2878
2879 #[test]
2880 fn execute_job_persists_stage_artifacts_manifest_then_db_and_finalizes() {
2881 let dir = tempfile::tempdir().unwrap();
2882 let db = db();
2883 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
2884 let job_dir = dir.path().join("job");
2885 let row = persist_job(&db, &job_dir, "01JBR55ORDER", &req, ChainJobState::Queued);
2886 let executor = Arc::new(FakeExecutor {
2887 calls: AtomicUsize::new(0),
2888 cancel_on_progress: AtomicBool::new(false),
2889 });
2890 let deps = deps(
2891 db,
2892 dir.path().join("jobs"),
2893 executor,
2894 Arc::new(FakeProbe(AtomicUsize::new(0))),
2895 );
2896
2897 execute_job(&deps, &row, 0).unwrap();
2898
2899 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2900 let db = deps.db.as_ref().as_ref().unwrap();
2901 let job_after = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
2902 assert!(
2903 job_dir.join("stages/000/segment.mp4").exists(),
2904 "job state {:?}, error {:?}",
2905 job_after.state,
2906 job_after.error
2907 );
2908 assert!(job_dir.join("stages/000/tail/000.png").exists());
2909 assert!(job_dir.join("stages/001/segment.mp4").exists());
2910 assert_eq!(manifest.stage_status[0].state, StageState::Completed);
2911 assert_eq!(manifest.stage_status[1].state, StageState::Completed);
2912 assert_eq!(manifest.finalizes.len(), 1);
2913
2914 let stages = chain_jobs::stages_for_job(db, &row.id).unwrap();
2915 assert_eq!(stages[0].state, StageState::Completed);
2916 assert_eq!(
2917 stages[0].segment_rel_path.as_deref(),
2918 Some("stages/000/segment.mp4")
2919 );
2920 assert_eq!(
2921 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2922 ChainJobState::Completed
2923 );
2924 }
2925
2926 #[test]
2927 fn ephemeral_execute_job_defers_gallery_record_to_legacy_shim() {
2928 let dir = tempfile::tempdir().unwrap();
2929 let db = db();
2930 let req = request(vec![TransitionMode::Smooth]);
2931 let job_dir = dir.path().join("job");
2932 let row = persist_job(
2933 &db,
2934 &job_dir,
2935 "01JBR55EPHGALLERY",
2936 &req,
2937 ChainJobState::Queued,
2938 );
2939 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
2940 manifest.ephemeral = true;
2941 manifest.write_atomic(&job_dir).unwrap();
2942 let executor = Arc::new(FakeExecutor {
2943 calls: AtomicUsize::new(0),
2944 cancel_on_progress: AtomicBool::new(false),
2945 });
2946 let output_dir = dir.path().join("gallery");
2947 let mut deps = deps(
2948 db,
2949 dir.path().join("jobs"),
2950 executor,
2951 Arc::new(FakeProbe(AtomicUsize::new(0))),
2952 );
2953 deps.output_dir = Some(output_dir.clone());
2954
2955 execute_job(&deps, &row, 0).unwrap();
2956
2957 let db = deps.db.as_ref().as_ref().unwrap();
2958 assert_eq!(
2959 db.list(Some(&output_dir)).unwrap().len(),
2960 0,
2961 "ephemeral legacy shim jobs must not write runner-side gallery rows"
2962 );
2963 assert_eq!(
2964 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
2965 ChainJobState::Completed
2966 );
2967 }
2968
2969 #[test]
2970 fn durable_execute_job_records_exactly_one_runner_gallery_row() {
2971 let dir = tempfile::tempdir().unwrap();
2972 let db = db();
2973 let req = request(vec![TransitionMode::Smooth]);
2974 let job_dir = dir.path().join("job");
2975 let row = persist_job(
2976 &db,
2977 &job_dir,
2978 "01JBR55DURGALLERY",
2979 &req,
2980 ChainJobState::Queued,
2981 );
2982 let executor = Arc::new(FakeExecutor {
2983 calls: AtomicUsize::new(0),
2984 cancel_on_progress: AtomicBool::new(false),
2985 });
2986 let output_dir = dir.path().join("gallery");
2987 let mut deps = deps(
2988 db,
2989 dir.path().join("jobs"),
2990 executor,
2991 Arc::new(FakeProbe(AtomicUsize::new(0))),
2992 );
2993 deps.output_dir = Some(output_dir.clone());
2994
2995 execute_job(&deps, &row, 0).unwrap();
2996
2997 let db = deps.db.as_ref().as_ref().unwrap();
2998 let rows = db.list(Some(&output_dir)).unwrap();
2999 assert_eq!(rows.len(), 1, "durable jobs save exactly one gallery row");
3000 assert_eq!(rows[0].format, OutputFormat::Mp4);
3001 assert_eq!(
3002 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3003 ChainJobState::Completed
3004 );
3005 }
3006
3007 #[test]
3008 fn cancel_via_progress_marks_job_cancelled_and_keeps_completed_artifacts() {
3009 let dir = tempfile::tempdir().unwrap();
3010 let db = db();
3011 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3012 let job_dir = dir.path().join("job");
3013 let row = persist_job(&db, &job_dir, "01JBR55CANCEL", &req, ChainJobState::Queued);
3014 let executor = Arc::new(FakeExecutor {
3015 calls: AtomicUsize::new(0),
3016 cancel_on_progress: AtomicBool::new(true),
3017 });
3018 let deps = deps(
3019 db,
3020 dir.path().join("jobs"),
3021 executor,
3022 Arc::new(FakeProbe(AtomicUsize::new(0))),
3023 );
3024 deps.cancel.register(&row.id);
3025 assert!(deps.cancel.request(&row.id));
3026
3027 execute_job(&deps, &row, 0).unwrap();
3028
3029 let db = deps.db.as_ref().as_ref().unwrap();
3030 assert_eq!(
3031 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3032 ChainJobState::Cancelled
3033 );
3034 assert!(!job_dir.join("stages/000/segment.mp4").exists());
3035 }
3036
3037 #[test]
3038 fn cancel_then_retake_runs_to_completion_without_stale_cancel_mark() {
3039 let dir = tempfile::tempdir().unwrap();
3040 let db = db();
3041 let req = request(vec![TransitionMode::Smooth]);
3042 let job_dir = dir.path().join("job");
3043 let row = persist_job(
3044 &db,
3045 &job_dir,
3046 "01JBR55CANCELRETAKE",
3047 &req,
3048 ChainJobState::Queued,
3049 );
3050 let executor = Arc::new(FakeExecutor {
3051 calls: AtomicUsize::new(0),
3052 cancel_on_progress: AtomicBool::new(false),
3053 });
3054 let deps = deps(
3055 db,
3056 dir.path().join("jobs"),
3057 executor.clone(),
3058 Arc::new(FakeProbe(AtomicUsize::new(0))),
3059 );
3060 deps.cancel.register(&row.id);
3061 assert!(deps.cancel.request(&row.id));
3062 execute_job(&deps, &row, 0).unwrap();
3063
3064 let db = deps.db.as_ref().as_ref().unwrap();
3065 assert_eq!(
3066 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3067 ChainJobState::Cancelled
3068 );
3069 let updated = apply_retake(
3070 db,
3071 dir.path(),
3072 &row.id,
3073 &RetakeRequest {
3074 stage_idx: 0,
3075 mode: RetakeMode::Cascade,
3076 seed_offset: Some(9),
3077 prompt: None,
3078 },
3079 )
3080 .unwrap();
3081
3082 execute_job(&deps, &updated, 0).unwrap();
3083
3084 assert_eq!(
3085 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3086 ChainJobState::Completed
3087 );
3088 assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
3089 }
3090
3091 #[test]
3092 fn cancel_mark_after_claim_is_honored_before_first_stage() {
3093 let dir = tempfile::tempdir().unwrap();
3094 let db = db();
3095 let req = request(vec![TransitionMode::Smooth]);
3096 let job_dir = dir.path().join("job");
3097 let row = persist_job(
3098 &db,
3099 &job_dir,
3100 "01JBR55CLAIMCANCEL",
3101 &req,
3102 ChainJobState::Queued,
3103 );
3104 assert!(chain_jobs::claim_job(&db, &row.id).unwrap());
3105 let executor = Arc::new(FakeExecutor {
3106 calls: AtomicUsize::new(0),
3107 cancel_on_progress: AtomicBool::new(false),
3108 });
3109 let deps = deps(
3110 db,
3111 dir.path().join("jobs"),
3112 executor.clone(),
3113 Arc::new(FakeProbe(AtomicUsize::new(0))),
3114 );
3115 deps.cancel.register(&row.id);
3116 assert!(deps.cancel.request(&row.id));
3117 let running = chain_jobs::get_job(deps.db.as_ref().as_ref().unwrap(), &row.id)
3118 .unwrap()
3119 .unwrap();
3120
3121 execute_job(&deps, &running, 0).unwrap();
3122
3123 let db = deps.db.as_ref().as_ref().unwrap();
3124 assert_eq!(
3125 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3126 ChainJobState::Cancelled
3127 );
3128 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
3129 }
3130
3131 #[test]
3132 fn retake_aborts_after_resume_wins_without_manifest_corruption() {
3133 let dir = tempfile::tempdir().unwrap();
3134 let db = db();
3135 let req = request(vec![TransitionMode::Smooth, TransitionMode::Cut]);
3136 let job_dir = dir.path().join("job");
3137 let row = persist_job(
3138 &db,
3139 &job_dir,
3140 "01JBR55RETAKERACE",
3141 &req,
3142 ChainJobState::Failed,
3143 );
3144 assert!(chain_jobs::try_transition(
3145 &db,
3146 &row.id,
3147 &[ChainJobState::Failed],
3148 ChainJobState::Queued,
3149 None,
3150 2_000,
3151 )
3152 .unwrap());
3153
3154 let err = apply_retake(
3155 &db,
3156 dir.path(),
3157 &row.id,
3158 &RetakeRequest {
3159 stage_idx: 0,
3160 mode: RetakeMode::Cascade,
3161 seed_offset: Some(5),
3162 prompt: Some("should not land".into()),
3163 },
3164 )
3165 .unwrap_err();
3166
3167 assert!(
3168 err.to_string().contains("not retakeable")
3169 || err.to_string().contains("CHAIN_JOB_RUNNING"),
3170 "retake loser must abort with a current-state error, got {err:#}"
3171 );
3172 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3173 assert!(manifest.retakes.is_empty());
3174 assert_eq!(manifest.stage_status[0].seed, 42);
3175 }
3176
3177 #[test]
3178 fn apply_retake_rejects_splice_before_smooth_boundary() {
3179 let dir = tempfile::tempdir().unwrap();
3180 let db = db();
3181 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3182 let job_dir = dir.path().join("job");
3183 let _row = persist_job(
3184 &db,
3185 &job_dir,
3186 "01JBR55RETAKE",
3187 &req,
3188 ChainJobState::Completed,
3189 );
3190
3191 let err = apply_retake(
3192 &db,
3193 dir.path(),
3194 "01JBR55RETAKE",
3195 &RetakeRequest {
3196 stage_idx: 0,
3197 mode: RetakeMode::Splice,
3198 seed_offset: Some(9),
3199 prompt: None,
3200 },
3201 )
3202 .unwrap_err();
3203 assert!(err
3204 .to_string()
3205 .contains("RETAKE_SPLICE_REQUIRES_CUT_OR_FADE"));
3206 }
3207
3208 #[test]
3209 fn apply_retake_cascade_resets_target_through_end_and_records_amendment() {
3210 let dir = tempfile::tempdir().unwrap();
3211 let db = db();
3212 let req = request(vec![
3213 TransitionMode::Smooth,
3214 TransitionMode::Cut,
3215 TransitionMode::Cut,
3216 ]);
3217 let job_dir = dir.path().join("job");
3218 let row = persist_job(
3219 &db,
3220 &job_dir,
3221 "01JBR55CASCADE",
3222 &req,
3223 ChainJobState::Completed,
3224 );
3225 let layout = JobDirLayout::new(job_dir.clone());
3226 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3227 for stage in &mut manifest.stage_status {
3228 stage.state = StageState::Completed;
3229 stage.segment = Some(layout.segment_rel(stage.idx));
3230 layout.ensure_stage_dirs(stage.idx).unwrap();
3231 std::fs::write(layout.stage_dir(stage.idx).join("stale.txt"), b"stale").unwrap();
3232 chain_jobs::upsert_stage(
3233 &db,
3234 &ChainJobStageRow {
3235 job_id: row.id.clone(),
3236 stage_idx: stage.idx,
3237 state: StageState::Completed,
3238 seed: stage.seed,
3239 frames_emitted: Some(9),
3240 generation_time_ms: Some(10),
3241 segment_rel_path: stage.segment.clone(),
3242 error: None,
3243 updated_at_ms: 1_000,
3244 },
3245 )
3246 .unwrap();
3247 }
3248 manifest.write_atomic(&job_dir).unwrap();
3249
3250 let updated = apply_retake(
3251 &db,
3252 dir.path(),
3253 &row.id,
3254 &RetakeRequest {
3255 stage_idx: 1,
3256 mode: RetakeMode::Cascade,
3257 seed_offset: Some(7),
3258 prompt: Some("new middle".into()),
3259 },
3260 )
3261 .unwrap();
3262
3263 assert_eq!(updated.state, ChainJobState::Queued);
3264 assert_eq!(updated.current_stage, 1);
3265 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3266 assert_eq!(manifest.retakes.len(), 1);
3267 assert_eq!(manifest.retakes[0].stage_idx, 1);
3268 assert_eq!(
3269 manifest.retakes[0].new_prompt.as_deref(),
3270 Some("new middle")
3271 );
3272 assert_eq!(manifest.stage_status[0].state, StageState::Completed);
3273 assert_eq!(manifest.stage_status[1].state, StageState::Pending);
3274 assert_eq!(manifest.stage_status[2].state, StageState::Pending);
3275 assert!(!layout.stage_dir(1).exists());
3276 assert!(!layout.stage_dir(2).exists());
3277 let stages = chain_jobs::stages_for_job(&db, &row.id).unwrap();
3278 assert_eq!(stages[1].state, StageState::Pending);
3279 assert_eq!(stages[2].state, StageState::Pending);
3280 let effective = effective_request(&manifest).unwrap();
3281 assert_eq!(effective.stages[1].prompt, "new middle");
3282 assert_eq!(
3283 manifest.stage_status[1].seed,
3284 effective_stage_seed(req.seed.unwrap(), Some(7))
3285 );
3286 }
3287
3288 #[test]
3289 fn finalize_versions_outputs_and_records_stage_seeds() {
3290 let dir = tempfile::tempdir().unwrap();
3291 let db = db();
3292 let req = request(vec![TransitionMode::Smooth]);
3293 let job_dir = dir.path().join("job");
3294 let row = persist_job(&db, &job_dir, "01JBR55VERSION", &req, ChainJobState::Queued);
3295 let executor = Arc::new(FakeExecutor {
3296 calls: AtomicUsize::new(0),
3297 cancel_on_progress: AtomicBool::new(false),
3298 });
3299 let deps = deps(
3300 db,
3301 dir.path().join("jobs"),
3302 executor,
3303 Arc::new(FakeProbe(AtomicUsize::new(0))),
3304 );
3305
3306 execute_job(&deps, &row, 0).unwrap();
3307 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3308 let first_seed = manifest.stage_status[0].seed;
3309 assert_eq!(manifest.finalizes[0].output, "final/output-1.mp4");
3310
3311 let second = finalize_job(&deps, &row, &mut manifest).unwrap();
3312
3313 assert_eq!(second, "final/output-2.mp4");
3314 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3315 assert_eq!(manifest.finalizes.len(), 2);
3316 assert_eq!(manifest.finalizes[0].stage_seeds, vec![first_seed]);
3317 assert_eq!(manifest.finalizes[1].stage_seeds, vec![first_seed]);
3318 assert!(job_dir.join("final/output-1.mp4").exists());
3319 assert!(job_dir.join("final/output-2.mp4").exists());
3320 }
3321
3322 #[test]
3323 fn completed_fade_successor_is_reencoded_after_previous_stage_changes_boundary() {
3324 let dir = tempfile::tempdir().unwrap();
3325 let job_dir = dir.path().join("job");
3326 let layout = JobDirLayout::new(job_dir.clone());
3327 layout.ensure_root().unwrap();
3328 let req = request(vec![TransitionMode::Smooth, TransitionMode::Fade]);
3329 let mut manifest = ChainJobManifest::new("01JBR55FADE".into(), 1_000, &req).unwrap();
3330
3331 write_stage_artifacts(&layout, &mut manifest, 0, &outcome(20), &req).unwrap();
3332 write_stage_artifacts(&layout, &mut manifest, 1, &outcome(80), &req).unwrap();
3333 let before = std::fs::read(layout.segment_path(1)).unwrap();
3334
3335 write_frames_to_dir(&layout.boundary_out_dir(0), &[frame(210), frame(220)]).unwrap();
3336 maybe_reencode_next_after_fade(&layout, &manifest, 0, &req).unwrap();
3337
3338 let after = std::fs::read(layout.segment_path(1)).unwrap();
3339 assert_ne!(
3340 before, after,
3341 "changing outgoing fade boundary must re-encode completed successor segment"
3342 );
3343 }
3344
3345 #[test]
3346 fn splice_retake_skips_completed_successors_and_only_reencodes_fade_boundary() {
3347 for (next_transition, successor_should_change, id) in [
3348 (TransitionMode::Fade, true, "01JBR55SPLICEFADE"),
3349 (TransitionMode::Cut, false, "01JBR55SPLICECUT"),
3350 ] {
3351 let dir = tempfile::tempdir().unwrap();
3352 let db = db();
3353 let req = request(vec![
3354 TransitionMode::Smooth,
3355 next_transition,
3356 TransitionMode::Cut,
3357 TransitionMode::Cut,
3358 ]);
3359 let job_dir = dir.path().join("job");
3360 let row = persist_job(&db, &job_dir, id, &req, ChainJobState::Queued);
3361 let executor = Arc::new(FakeExecutor {
3362 calls: AtomicUsize::new(0),
3363 cancel_on_progress: AtomicBool::new(false),
3364 });
3365 let deps = deps(
3366 db,
3367 dir.path().join("jobs"),
3368 executor.clone(),
3369 Arc::new(FakeProbe(AtomicUsize::new(0))),
3370 );
3371 execute_job(&deps, &row, 0).unwrap();
3372
3373 let layout = JobDirLayout::new(job_dir.clone());
3374 let before_1 = std::fs::read(layout.segment_path(1)).unwrap();
3375 let before_2 = std::fs::read(layout.segment_path(2)).unwrap();
3376 let before_3 = std::fs::read(layout.segment_path(3)).unwrap();
3377 let calls_before = executor.calls.load(Ordering::SeqCst);
3378 let db = deps.db.as_ref().as_ref().unwrap();
3379 let updated = apply_retake(
3380 db,
3381 dir.path(),
3382 &row.id,
3383 &RetakeRequest {
3384 stage_idx: 0,
3385 mode: RetakeMode::Splice,
3386 seed_offset: Some(9),
3387 prompt: None,
3388 },
3389 )
3390 .unwrap();
3391
3392 execute_job(&deps, &updated, 0).unwrap();
3393
3394 assert_eq!(
3395 executor.calls.load(Ordering::SeqCst) - calls_before,
3396 1,
3397 "splice retake must render exactly the target stage for {id}"
3398 );
3399 let after_1 = std::fs::read(layout.segment_path(1)).unwrap();
3400 let after_2 = std::fs::read(layout.segment_path(2)).unwrap();
3401 let after_3 = std::fs::read(layout.segment_path(3)).unwrap();
3402 if successor_should_change {
3403 assert_ne!(
3404 before_1, after_1,
3405 "fade successor must be re-encoded at the edited boundary"
3406 );
3407 } else {
3408 assert_eq!(
3409 before_1, after_1,
3410 "cut successor must not be touched by splice retake"
3411 );
3412 }
3413 assert_eq!(before_2, after_2, "stage N+2 must remain untouched");
3414 assert_eq!(before_3, after_3, "stage N+3 must remain untouched");
3415 }
3416 }
3417
3418 #[test]
3419 fn disk_write_failure_marks_job_and_stage_failed_with_error() {
3420 let dir = tempfile::tempdir().unwrap();
3421 let db = db();
3422 let req = request(vec![TransitionMode::Smooth]);
3423 let job_dir = dir.path().join("job");
3424 let row = persist_job(&db, &job_dir, "01JBR55DISK", &req, ChainJobState::Queued);
3425 std::fs::write(job_dir.join("stages"), b"not a directory").unwrap();
3426 let executor = Arc::new(FakeExecutor {
3427 calls: AtomicUsize::new(0),
3428 cancel_on_progress: AtomicBool::new(false),
3429 });
3430 let deps = deps(
3431 db,
3432 dir.path().join("jobs"),
3433 executor,
3434 Arc::new(FakeProbe(AtomicUsize::new(0))),
3435 );
3436
3437 execute_job(&deps, &row, 0).unwrap();
3438
3439 let db = deps.db.as_ref().as_ref().unwrap();
3440 let job_after = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3441 assert_eq!(job_after.state, ChainJobState::Failed);
3442 assert!(
3443 job_after.error.as_deref().is_some_and(
3444 |error| error.contains("stages/000") || error.contains("Not a directory")
3445 ),
3446 "explicit filesystem error should be persisted, got {:?}",
3447 job_after.error
3448 );
3449 let stage = chain_jobs::stages_for_job(db, &row.id).unwrap().remove(0);
3450 assert_eq!(stage.state, StageState::Failed);
3451 assert!(stage
3452 .error
3453 .as_deref()
3454 .unwrap_or_default()
3455 .contains("stages/000"));
3456 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3457 assert_eq!(manifest.stage_status[0].state, StageState::Failed);
3458 }
3459
3460 #[test]
3461 fn finalize_failure_then_resume_reenters_finalize_without_rerendering() {
3462 let dir = tempfile::tempdir().unwrap();
3463 let db = db();
3464 let req = request(vec![TransitionMode::Smooth]);
3465 let job_dir = dir.path().join("job");
3466 let row = persist_job(
3467 &db,
3468 &job_dir,
3469 "01JBR55FINALRETRY",
3470 &req,
3471 ChainJobState::Queued,
3472 );
3473 std::fs::write(job_dir.join("final"), b"not a directory").unwrap();
3474 let executor = Arc::new(FakeExecutor {
3475 calls: AtomicUsize::new(0),
3476 cancel_on_progress: AtomicBool::new(false),
3477 });
3478 let deps = deps(
3479 db,
3480 dir.path().join("jobs"),
3481 executor.clone(),
3482 Arc::new(FakeProbe(AtomicUsize::new(0))),
3483 );
3484
3485 execute_job(&deps, &row, 0).unwrap();
3486 let db = deps.db.as_ref().as_ref().unwrap();
3487 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3488 assert_eq!(failed.state, ChainJobState::Failed);
3489 assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
3490 assert!(ChainJobManifest::read_from_dir(&job_dir)
3491 .unwrap()
3492 .finalizes
3493 .is_empty());
3494
3495 std::fs::remove_file(job_dir.join("final")).unwrap();
3496 assert!(chain_jobs::try_transition(
3497 db,
3498 &row.id,
3499 &[ChainJobState::Failed],
3500 ChainJobState::Queued,
3501 None,
3502 now_ms_i64(),
3503 )
3504 .unwrap());
3505 let queued = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3506 execute_job(&deps, &queued, req.stages.len() as u32).unwrap();
3507
3508 let completed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3509 assert_eq!(completed.state, ChainJobState::Completed);
3510 assert_eq!(
3511 executor.calls.load(Ordering::SeqCst),
3512 1,
3513 "resume from completed stages must not render again"
3514 );
3515 let manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3516 assert_eq!(manifest.finalizes.len(), 1);
3517 assert_eq!(manifest.finalizes[0].output, "final/output-1.mp4");
3518 }
3519
3520 #[test]
3521 fn boundary_cancel_keeps_completed_stage_artifacts_and_stops_next_stage() {
3522 let dir = tempfile::tempdir().unwrap();
3523 let db = db();
3524 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3525 let job_dir = dir.path().join("job");
3526 let row = persist_job(
3527 &db,
3528 &job_dir,
3529 "01JBR55BOUNDARY",
3530 &req,
3531 ChainJobState::Queued,
3532 );
3533 let cancel = Arc::new(CancelRegistry::new());
3534 let executor = Arc::new(BoundaryCancelExecutor {
3535 calls: AtomicUsize::new(0),
3536 cancel: cancel.clone(),
3537 job_id: row.id.clone(),
3538 });
3539 let deps = RunnerDeps {
3540 db: Arc::new(Some(db)),
3541 jobs_root: dir.path().join("jobs"),
3542 executor,
3543 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
3544 events: Arc::new(JobEventBus::new()),
3545 cancel,
3546 job_locks: Arc::new(JobMutationLocks::new()),
3547 claims: Arc::new(EphemeralClaims::default()),
3548 output_dir: None,
3549 server_events: None,
3550 };
3551
3552 execute_job(&deps, &row, 0).unwrap();
3553
3554 let db = deps.db.as_ref().as_ref().unwrap();
3555 assert_eq!(
3556 chain_jobs::get_job(db, &row.id).unwrap().unwrap().state,
3557 ChainJobState::Cancelled
3558 );
3559 assert!(job_dir.join("stages/000/segment.mp4").exists());
3560 assert!(!job_dir.join("stages/001/segment.mp4").exists());
3561 }
3562
3563 #[test]
3564 fn yields_between_stages_when_small_jobs_are_waiting_and_cleans_bus_at_terminal() {
3565 let dir = tempfile::tempdir().unwrap();
3566 let db = db();
3567 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3568 let job_dir = dir.path().join("job");
3569 let row = persist_job(&db, &job_dir, "01JBR55YIELD", &req, ChainJobState::Queued);
3570 let executor = Arc::new(FakeExecutor {
3571 calls: AtomicUsize::new(0),
3572 cancel_on_progress: AtomicBool::new(false),
3573 });
3574 let deps = deps(
3575 db,
3576 dir.path().join("jobs"),
3577 executor,
3578 Arc::new(FakeProbe(AtomicUsize::new(2))),
3579 );
3580 let mut rx = deps.events.subscribe_persistent_for_tests(&row.id);
3581
3582 execute_job(&deps, &row, 0).unwrap();
3583
3584 let mut events = Vec::new();
3585 loop {
3586 match rx.try_recv() {
3587 Ok(event) => events.push(event),
3588 Err(tokio::sync::broadcast::error::TryRecvError::Empty)
3589 | Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
3590 Err(err) => panic!("unexpected broadcast receive error: {err}"),
3591 }
3592 }
3593 assert!(matches!(
3594 events.as_slice(),
3595 [
3596 ChainJobEvent::StateChanged {
3597 state: ChainJobState::Running,
3598 ..
3599 },
3600 ChainJobEvent::StageStart { stage_idx: 0 },
3601 ChainJobEvent::DenoiseStep { stage_idx: 0, .. },
3602 ChainJobEvent::StageDone { stage_idx: 0, .. },
3603 ..
3604 ]
3605 ));
3606 assert!(events.iter().any(|event| {
3607 matches!(
3608 event,
3609 ChainJobEvent::Yielded {
3610 pending_small_jobs: 2
3611 }
3612 )
3613 }));
3614 assert!(!deps.events.senders.lock().unwrap().contains_key(&row.id));
3615 }
3616
3617 #[test]
3618 fn resume_carry_from_disk_loads_smooth_tail_and_rejects_cut_or_fade_carry() {
3619 let dir = tempfile::tempdir().unwrap();
3620 let job_dir = dir.path().join("job");
3621 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3622 let mut manifest = ChainJobManifest::new("01JBR55CARRY".into(), 1_000, &req).unwrap();
3623 manifest.stage_status[0].state = StageState::Completed;
3624 manifest.stage_status[0].segment = Some("stages/000/segment.mp4".into());
3625 manifest.stage_status[0].tail_frames = Some(1);
3626 let layout = JobDirLayout::new(job_dir.clone());
3627 layout.ensure_stage_dirs(0).unwrap();
3628 frame(77).save(layout.tail_dir(0).join("000.png")).unwrap();
3629 manifest.write_atomic(&job_dir).unwrap();
3630
3631 let carry = resume_carry_from_disk(&job_dir, &manifest, 1)
3632 .unwrap()
3633 .expect("smooth continuation should reload tail");
3634 assert_eq!(carry.frames, 1);
3635 assert_eq!(carry.tail_rgb_frames[0].get_pixel(0, 0).0, [77, 77, 77]);
3636
3637 for transition in [TransitionMode::Cut, TransitionMode::Fade] {
3638 let req = request(vec![TransitionMode::Smooth, transition]);
3639 let mut manifest = ChainJobManifest::new("01JBR55CARRY".into(), 1_000, &req).unwrap();
3640 manifest.stage_status[0].state = StageState::Completed;
3641 manifest.stage_status[0].tail_frames = Some(1);
3642 assert!(resume_carry_from_disk(&job_dir, &manifest, 1)
3643 .unwrap()
3644 .is_none());
3645 }
3646 }
3647
3648 #[test]
3649 fn malicious_manifest_resume_fails_and_job_remains_deletable() {
3650 let dir = tempfile::tempdir().unwrap();
3651 let db = db();
3652 let req = request(vec![TransitionMode::Smooth, TransitionMode::Smooth]);
3653 let job_dir = dir.path().join("job");
3654 let row = persist_job(
3655 &db,
3656 &job_dir,
3657 "01JBR55TRAVERSAL",
3658 &req,
3659 ChainJobState::Queued,
3660 );
3661 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3662 manifest.stage_status[0].state = StageState::Completed;
3663 manifest.stage_status[0].segment = Some("../outside.mp4".into());
3664 manifest.stage_status[0].tail_frames = Some(1);
3665 manifest.write_atomic(&job_dir).unwrap();
3666 let executor = Arc::new(FakeExecutor {
3667 calls: AtomicUsize::new(0),
3668 cancel_on_progress: AtomicBool::new(false),
3669 });
3670 let deps = deps(
3671 db,
3672 dir.path().join("jobs"),
3673 executor.clone(),
3674 Arc::new(FakeProbe(AtomicUsize::new(0))),
3675 );
3676
3677 execute_job(&deps, &row, 0).unwrap();
3678
3679 let db = deps.db.as_ref().as_ref().unwrap();
3680 let failed = chain_jobs::get_job(db, &row.id).unwrap().unwrap();
3681 assert_eq!(failed.state, ChainJobState::Failed);
3682 assert!(
3683 failed
3684 .error
3685 .as_deref()
3686 .is_some_and(|err| err.contains("must be relative")),
3687 "expected traversal error to be persisted, got {:?}",
3688 failed.error
3689 );
3690 assert_eq!(executor.calls.load(Ordering::SeqCst), 0);
3691 assert!(chain_jobs::delete_job_not_running(db, &row.id).unwrap());
3692 assert!(chain_jobs::get_job(db, &row.id).unwrap().is_none());
3693 }
3694
3695 #[test]
3696 fn startup_reconcile_flips_running_to_interrupted_and_repairs_from_manifest() {
3697 let dir = tempfile::tempdir().unwrap();
3698 let db = db();
3699 let req = request(vec![TransitionMode::Smooth]);
3700 let job_dir = dir.path().join("jobs/01JBR55RECON");
3701 let row = persist_job(&db, &job_dir, "01JBR55RECON", &req, ChainJobState::Running);
3702 let mut manifest = ChainJobManifest::read_from_dir(&job_dir).unwrap();
3703 manifest.stage_status[0].state = StageState::Completed;
3704 manifest.stage_status[0].frames_emitted = Some(9);
3705 manifest.stage_status[0].generation_time_ms = Some(123);
3706 manifest.stage_status[0].segment = Some("stages/000/segment.mp4".into());
3707 manifest.write_atomic(&job_dir).unwrap();
3708
3709 let (flipped, repaired) = startup_reconcile(&db, &dir.path().join("jobs")).unwrap();
3710
3711 assert_eq!(flipped, 1);
3712 assert!(repaired >= 1);
3713 let row_after = chain_jobs::get_job(&db, &row.id).unwrap().unwrap();
3714 assert_eq!(row_after.state, ChainJobState::Interrupted);
3715 assert_eq!(row_after.current_stage, 1);
3716 assert_eq!(
3717 row_after.error.as_deref(),
3718 Some("server restarted while chain job was running")
3719 );
3720 let stage = chain_jobs::stages_for_job(&db, &row.id).unwrap().remove(0);
3721 assert_eq!(stage.state, StageState::Completed);
3722 assert_eq!(
3723 stage.segment_rel_path.as_deref(),
3724 Some("stages/000/segment.mp4")
3725 );
3726 }
3727
3728 #[test]
3729 fn ephemeral_claim_guard_releases_on_drop() {
3730 let claims = Arc::new(EphemeralClaims::default());
3731 {
3732 let _guard = claims.claim("01JBR55CLAIM");
3733 assert!(claims.is_claimed("01JBR55CLAIM"));
3734 }
3735 assert!(!claims.is_claimed("01JBR55CLAIM"));
3736 }
3737
3738 #[test]
3739 fn create_job_with_params_persists_ephemeral_manifest_and_stage_rows() {
3740 let dir = tempfile::tempdir().unwrap();
3741 let db = db();
3742 let req = request(vec![TransitionMode::Smooth, TransitionMode::Cut]);
3743
3744 let row = create_job_with_params(
3745 &db,
3746 dir.path(),
3747 CreateJobParams {
3748 id: "01JBR55CREATE".into(),
3749 ephemeral: true,
3750 request: req.clone(),
3751 },
3752 )
3753 .unwrap();
3754
3755 assert_eq!(row.id, "01JBR55CREATE");
3756 assert_eq!(row.state, ChainJobState::Queued);
3757 assert_eq!(row.stage_count, 2);
3758 assert!(row.job_dir.ends_with("01JBR55CREATE"));
3759 let manifest = ChainJobManifest::read_from_dir(&row.job_dir).unwrap();
3760 assert!(manifest.ephemeral);
3761 assert_eq!(manifest.request().unwrap(), req);
3762 let stages = chain_jobs::stages_for_job(&db, &row.id).unwrap();
3763 assert_eq!(stages.len(), 2);
3764 assert_eq!(stages[0].state, StageState::Pending);
3765 }
3766
3767 #[test]
3768 fn startup_gc_sweep_removes_non_running_ephemerals_only() {
3769 let dir = tempfile::tempdir().unwrap();
3770 let db = db();
3771 let req = request(vec![TransitionMode::Smooth]);
3772 let jobs_root = dir.path().join("jobs");
3773 let ephemeral_done = persist_job(
3774 &db,
3775 &jobs_root.join("01JBR55EPHDONE"),
3776 "01JBR55EPHDONE",
3777 &req,
3778 ChainJobState::Completed,
3779 );
3780 let mut manifest = ChainJobManifest::read_from_dir(&ephemeral_done.job_dir).unwrap();
3781 manifest.ephemeral = true;
3782 manifest.write_atomic(&ephemeral_done.job_dir).unwrap();
3783 let ephemeral_running = persist_job(
3784 &db,
3785 &jobs_root.join("01JBR55EPHRUN"),
3786 "01JBR55EPHRUN",
3787 &req,
3788 ChainJobState::Running,
3789 );
3790 let mut manifest = ChainJobManifest::read_from_dir(&ephemeral_running.job_dir).unwrap();
3791 manifest.ephemeral = true;
3792 manifest.write_atomic(&ephemeral_running.job_dir).unwrap();
3793 let durable_done = persist_job(
3794 &db,
3795 &jobs_root.join("01JBR55DURABLE"),
3796 "01JBR55DURABLE",
3797 &req,
3798 ChainJobState::Completed,
3799 );
3800
3801 let outcome = startup_gc_sweep(&db, &jobs_root).unwrap();
3802
3803 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3804 assert!(chain_jobs::get_job(&db, &ephemeral_done.id)
3805 .unwrap()
3806 .is_none());
3807 assert!(!ephemeral_done.job_dir.exists());
3808 assert!(chain_jobs::get_job(&db, &ephemeral_running.id)
3809 .unwrap()
3810 .is_some());
3811 assert!(ephemeral_running.job_dir.exists());
3812 assert!(chain_jobs::get_job(&db, &durable_done.id)
3813 .unwrap()
3814 .is_some());
3815 assert!(durable_done.job_dir.exists());
3816 }
3817
3818 #[test]
3819 fn gc_sweeps_unclaimed_ephemeral_after_grace_but_preserves_live_claim() {
3820 let dir = tempfile::tempdir().unwrap();
3821 let db = db();
3822 let req = request(vec![TransitionMode::Smooth]);
3823 let jobs_root = dir.path().join("jobs");
3824 let old = persist_job(
3825 &db,
3826 &jobs_root.join("01JBR55OLDCLAIMLESS"),
3827 "01JBR55OLDCLAIMLESS",
3828 &req,
3829 ChainJobState::Completed,
3830 );
3831 let mut manifest = ChainJobManifest::read_from_dir(&old.job_dir).unwrap();
3832 manifest.ephemeral = true;
3833 manifest.write_atomic(&old.job_dir).unwrap();
3834 chain_jobs::update_job_state(&db, &old.id, ChainJobState::Completed, None, 1_000).unwrap();
3835 let claimed = persist_job(
3836 &db,
3837 &jobs_root.join("01JBR55CLAIMEDGC"),
3838 "01JBR55CLAIMEDGC",
3839 &req,
3840 ChainJobState::Completed,
3841 );
3842 let mut manifest = ChainJobManifest::read_from_dir(&claimed.job_dir).unwrap();
3843 manifest.ephemeral = true;
3844 manifest.write_atomic(&claimed.job_dir).unwrap();
3845
3846 let deps = deps(
3847 db,
3848 jobs_root,
3849 Arc::new(FakeExecutor {
3850 calls: AtomicUsize::new(0),
3851 cancel_on_progress: AtomicBool::new(false),
3852 }),
3853 Arc::new(FakeProbe(AtomicUsize::new(0))),
3854 );
3855 let _claim = deps.claims.claim(&claimed.id);
3856 let outcome =
3857 run_gc_pass(&deps, 7, 1_000 + (EPHEMERAL_GRACE_SECS as i64 + 1) * 1_000).unwrap();
3858
3859 let db = deps.db.as_ref().as_ref().unwrap();
3860 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3861 assert!(chain_jobs::get_job(db, &old.id).unwrap().is_none());
3862 assert!(chain_jobs::get_job(db, &claimed.id).unwrap().is_some());
3863 assert!(claimed.job_dir.exists());
3864 }
3865
3866 #[test]
3867 fn gc_respects_ephemeral_grace_and_prunes_only_completed_durable_stages() {
3868 let dir = tempfile::tempdir().unwrap();
3869 let db = db();
3870 let req = request(vec![TransitionMode::Smooth]);
3871 let jobs_root = dir.path().join("jobs");
3872 let recent_ephemeral = persist_job(
3873 &db,
3874 &jobs_root.join("01JBR55RECENTEPH"),
3875 "01JBR55RECENTEPH",
3876 &req,
3877 ChainJobState::Completed,
3878 );
3879 let mut manifest = ChainJobManifest::read_from_dir(&recent_ephemeral.job_dir).unwrap();
3880 manifest.ephemeral = true;
3881 manifest.write_atomic(&recent_ephemeral.job_dir).unwrap();
3882 let now = 1_000 + 8 * 86_400_000;
3883 chain_jobs::update_job_state(
3884 &db,
3885 &recent_ephemeral.id,
3886 ChainJobState::Completed,
3887 None,
3888 now - (EPHEMERAL_GRACE_SECS as i64 * 1_000) + 100,
3889 )
3890 .unwrap();
3891
3892 let durable = persist_job(
3893 &db,
3894 &jobs_root.join("01JBR55TTL"),
3895 "01JBR55TTL",
3896 &req,
3897 ChainJobState::Completed,
3898 );
3899 let layout = JobDirLayout::new(durable.job_dir.clone());
3900 std::fs::create_dir_all(layout.stage_dir(0)).unwrap();
3901 std::fs::write(layout.stage_dir(0).join("segment.mp4"), b"stage").unwrap();
3902 std::fs::create_dir_all(durable.job_dir.join("final")).unwrap();
3903 std::fs::write(durable.job_dir.join("final/output-1.mp4"), b"final").unwrap();
3904 chain_jobs::update_job_state(&db, &durable.id, ChainJobState::Completed, None, 1_000)
3905 .unwrap();
3906
3907 let deps = deps(
3908 db,
3909 jobs_root,
3910 Arc::new(FakeExecutor {
3911 calls: AtomicUsize::new(0),
3912 cancel_on_progress: AtomicBool::new(false),
3913 }),
3914 Arc::new(FakeProbe(AtomicUsize::new(0))),
3915 );
3916 let outcome = run_gc_pass(&deps, 7, now).unwrap();
3917 let db = deps.db.as_ref().as_ref().unwrap();
3918
3919 assert_eq!(outcome.swept_ephemeral_jobs, 0);
3920 assert_eq!(outcome.pruned_artifact_dirs, 1);
3921 assert!(chain_jobs::get_job(db, &recent_ephemeral.id)
3922 .unwrap()
3923 .is_some());
3924 assert!(recent_ephemeral.job_dir.exists());
3925 assert!(!durable.job_dir.join("stages").exists());
3926 assert!(durable.job_dir.join("final/output-1.mp4").exists());
3927 assert!(durable.job_dir.join("manifest.toml").exists());
3928 assert!(chain_jobs::get_job(db, &durable.id).unwrap().is_some());
3929 }
3930
3931 #[test]
3932 fn gc_exempts_durable_failed_interrupted_cancelled_but_ephemeral_overrides() {
3933 let dir = tempfile::tempdir().unwrap();
3934 let db = db();
3935 let req = request(vec![TransitionMode::Smooth]);
3936 let jobs_root = dir.path().join("jobs");
3937 let mut durable_rows = Vec::new();
3938 for (id, state) in [
3939 ("01JBR55FAILED", ChainJobState::Failed),
3940 ("01JBR55INTR", ChainJobState::Interrupted),
3941 ("01JBR55CANCELLED", ChainJobState::Cancelled),
3942 ] {
3943 let row = persist_job(&db, &jobs_root.join(id), id, &req, state);
3944 std::fs::create_dir_all(row.job_dir.join("stages/000")).unwrap();
3945 durable_rows.push(row);
3946 }
3947 let eph_failed = persist_job(
3948 &db,
3949 &jobs_root.join("01JBR55EPHFAILED"),
3950 "01JBR55EPHFAILED",
3951 &req,
3952 ChainJobState::Failed,
3953 );
3954 let mut manifest = ChainJobManifest::read_from_dir(&eph_failed.job_dir).unwrap();
3955 manifest.ephemeral = true;
3956 manifest.write_atomic(&eph_failed.job_dir).unwrap();
3957 chain_jobs::update_job_state(&db, &eph_failed.id, ChainJobState::Failed, None, 1_000)
3958 .unwrap();
3959
3960 let deps = deps(
3961 db,
3962 jobs_root,
3963 Arc::new(FakeExecutor {
3964 calls: AtomicUsize::new(0),
3965 cancel_on_progress: AtomicBool::new(false),
3966 }),
3967 Arc::new(FakeProbe(AtomicUsize::new(0))),
3968 );
3969 let outcome =
3970 run_gc_pass(&deps, 7, 1_000 + (EPHEMERAL_GRACE_SECS as i64 + 1) * 1_000).unwrap();
3971 let db = deps.db.as_ref().as_ref().unwrap();
3972
3973 assert_eq!(outcome.swept_ephemeral_jobs, 1);
3974 assert!(chain_jobs::get_job(db, &eph_failed.id).unwrap().is_none());
3975 for row in durable_rows {
3976 assert!(chain_jobs::get_job(db, &row.id).unwrap().is_some());
3977 assert!(row.job_dir.join("stages").exists());
3978 }
3979 }
3980
3981 #[tokio::test]
3982 async fn runner_request_gc_replies_between_jobs() {
3983 let dir = tempfile::tempdir().unwrap();
3984 let db = db();
3985 let req = request(vec![TransitionMode::Smooth]);
3986 let jobs_root = dir.path().join("jobs");
3987 let _row = persist_job(
3988 &db,
3989 &jobs_root.join("01JBR55REQUESTGC"),
3990 "01JBR55REQUESTGC",
3991 &req,
3992 ChainJobState::Completed,
3993 );
3994 let deps = RunnerDeps {
3995 db: Arc::new(Some(db)),
3996 jobs_root,
3997 executor: Arc::new(FakeExecutor {
3998 calls: AtomicUsize::new(0),
3999 cancel_on_progress: AtomicBool::new(false),
4000 }),
4001 queue_probe: Arc::new(FakeProbe(AtomicUsize::new(0))),
4002 events: Arc::new(JobEventBus::new()),
4003 cancel: Arc::new(CancelRegistry::new()),
4004 job_locks: Arc::new(JobMutationLocks::new()),
4005 claims: Arc::new(EphemeralClaims::default()),
4006 output_dir: None,
4007 server_events: None,
4008 };
4009 let handle = spawn_runner(deps);
4010
4011 let outcome = tokio::time::timeout(std::time::Duration::from_secs(1), handle.request_gc())
4012 .await
4013 .expect("GC request should be serviced without waiting for the whole queue")
4014 .unwrap();
4015
4016 assert_eq!(outcome.swept_ephemeral_jobs, 0);
4017 }
4018}