1176 lines
45 KiB
Python
1176 lines
45 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""
|
||
Job Runner for FastVideo video generation jobs.
|
||
|
||
Manages job lifecycle, execution, logging, and generator caching.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import atexit
|
||
import collections
|
||
import contextlib
|
||
import copy
|
||
import enum
|
||
import json
|
||
import logging
|
||
import logging.handlers
|
||
import multiprocessing as mp
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
import yaml
|
||
|
||
from fastvideo.utils import get_mp_context
|
||
from fastvideo_studio.database import Database
|
||
from fastvideo_studio.training_config import (
|
||
build_training_config,
|
||
get_training_env,
|
||
)
|
||
|
||
logger = logging.getLogger("fastvideo.studio.job_runner")
|
||
|
||
# Regex patterns for parsing tqdm-style progress output.
|
||
# Matches e.g. " 40%|████ | 20/50 " or " 20/50 "
|
||
_TQDM_PCT_RE = re.compile(r"(\d+)%\|")
|
||
_TQDM_FRAC_RE = re.compile(r"\b(\d+)/(\d+)\b")
|
||
|
||
_MAX_LOG_LINES = 2000 # ring-buffer cap per job
|
||
|
||
|
||
class JobStatus(str, enum.Enum):
|
||
PENDING = "pending"
|
||
RUNNING = "running"
|
||
COMPLETED = "completed"
|
||
FAILED = "failed"
|
||
STOPPED = "stopped"
|
||
|
||
|
||
class JobLogBuffer:
|
||
"""Ring buffer for storing job log lines with progress tracking."""
|
||
|
||
def __init__(self, maxlen: int = _MAX_LOG_LINES):
|
||
self._lines: collections.deque[str] = collections.deque(maxlen=maxlen)
|
||
self._lock = threading.Lock()
|
||
self.progress: float = 0.0 # 0 – 100
|
||
self.progress_msg: str = "" # e.g. "20/50 steps"
|
||
self.phase: str = "initializing" # human-readable phase
|
||
|
||
def write(self, text: str):
|
||
"""Append *text* (may contain embedded newlines)."""
|
||
for line in text.splitlines():
|
||
stripped = line.strip()
|
||
if not stripped:
|
||
continue
|
||
with self._lock:
|
||
self._lines.append(stripped)
|
||
self._parse_progress(stripped)
|
||
|
||
def get_lines(self, after: int = 0) -> tuple[list[str], int]:
|
||
"""Return ``(lines[after:], total_len)``."""
|
||
with self._lock:
|
||
all_lines = list(self._lines)
|
||
return all_lines[after:], len(all_lines)
|
||
|
||
# -- internal helpers ---------------------------------------------------
|
||
|
||
def _parse_progress(self, line: str):
|
||
"""Try to extract a percentage / fraction from a tqdm line."""
|
||
m = _TQDM_PCT_RE.search(line)
|
||
if m:
|
||
self.progress = min(float(m.group(1)), 100.0)
|
||
m2 = _TQDM_FRAC_RE.search(line)
|
||
if m2:
|
||
cur, total = int(m2.group(1)), int(m2.group(2))
|
||
if total > 0:
|
||
self.progress = min(cur / total * 100.0, 100.0)
|
||
self.progress_msg = f"{cur}/{total} steps"
|
||
# Detect high-level phases from FastVideo log messages
|
||
low = line.lower()
|
||
if "loading model" in low or "loading" in low and "checkpoint" in low:
|
||
self.phase = "loading model"
|
||
elif "denoising" in low or "timestep" in low:
|
||
self.phase = "denoising"
|
||
elif "saving" in low or "saved" in low or "checkpoint" in low:
|
||
self.phase = "saving"
|
||
elif "encoding" in low or "vae" in low:
|
||
self.phase = "VAE encoding"
|
||
elif "training" in low or "step" in low or "loss" in low:
|
||
self.phase = "training"
|
||
|
||
|
||
class LogBufferHandler(logging.Handler):
|
||
"""Logging handler that writes to a JobLogBuffer."""
|
||
|
||
def __init__(self, buffer: JobLogBuffer):
|
||
super().__init__()
|
||
self.buffer = buffer
|
||
self.setFormatter(logging.Formatter("%(levelname)s %(asctime)s [%(name)s] %(message)s"))
|
||
|
||
def emit(self, record):
|
||
try:
|
||
msg = self.format(record)
|
||
self.buffer.write(msg)
|
||
except Exception:
|
||
self.handleError(record)
|
||
|
||
|
||
@dataclass
|
||
class Job:
|
||
id: str
|
||
model_id: str
|
||
name: str = ""
|
||
prompt: str = ""
|
||
workload_type: str = "t2v"
|
||
job_type: str = "inference"
|
||
image_path: str = ""
|
||
last_image_path: str = ""
|
||
references: list[dict[str, Any]] = field(default_factory=list)
|
||
status: JobStatus = JobStatus.PENDING
|
||
created_at: float = field(default_factory=time.time)
|
||
started_at: float | None = None
|
||
finished_at: float | None = None
|
||
error: str | None = None
|
||
output_path: str | None = None
|
||
log_file_path: str | None = None # Path to the job's log file
|
||
num_inference_steps: int = 50
|
||
num_frames: int = 81
|
||
height: int = 480
|
||
width: int = 832
|
||
guidance_scale: float = 5.0
|
||
guidance_rescale: float = 0.0
|
||
fps: int = 24
|
||
seed: int = 1024
|
||
negative_prompt: str = ""
|
||
num_gpus: int = 1
|
||
dit_cpu_offload: bool = False
|
||
dit_layerwise_offload: bool = False
|
||
text_encoder_cpu_offload: bool = False
|
||
vae_cpu_offload: bool = False
|
||
image_encoder_cpu_offload: bool = False
|
||
use_fsdp_inference: bool = False
|
||
enable_torch_compile: bool = False
|
||
vsa_sparsity: float = 0.0
|
||
tp_size: int = -1
|
||
sp_size: int = -1
|
||
# Training-specific (for finetuning, distillation, LoRA)
|
||
data_path: str = ""
|
||
max_train_steps: int = 1000
|
||
train_batch_size: int = 1
|
||
learning_rate: float = 5e-5
|
||
num_latent_t: int = 20
|
||
validation_dataset_file: str = ""
|
||
lora_rank: int = 32
|
||
# DMD options
|
||
dmd_use_vsa: bool = False
|
||
dmd_vsa_sparsity: float = 0.8
|
||
dmd_denoising_steps: str = "1000,757,522"
|
||
real_score_guidance_scale: float = 3.5
|
||
generator_update_interval: int = 5
|
||
real_score_model_path: str = ""
|
||
fake_score_model_path: str = ""
|
||
# Internal
|
||
_thread: threading.Thread | None = field(default=None, repr=False)
|
||
_stop_event: threading.Event = field(default_factory=threading.Event, repr=False)
|
||
_log_buf: JobLogBuffer = field(default_factory=JobLogBuffer, repr=False)
|
||
log_file_handler: logging.FileHandler | None = field(default=None, repr=False)
|
||
_process: subprocess.Popen | None = field(default=None, repr=False)
|
||
|
||
def to_dict(self) -> dict[str, Any]:
|
||
return {
|
||
"id": self.id,
|
||
"model_id": self.model_id,
|
||
"name": self.name,
|
||
"prompt": self.prompt,
|
||
"workload_type": self.workload_type,
|
||
"job_type": self.job_type,
|
||
"image_path": self.image_path,
|
||
"last_image_path": self.last_image_path,
|
||
"references": self.references,
|
||
"status": self.status.value,
|
||
"created_at": self.created_at,
|
||
"started_at": self.started_at,
|
||
"finished_at": self.finished_at,
|
||
"error": self.error,
|
||
"output_path": self.output_path,
|
||
"log_file_path": self.log_file_path,
|
||
"num_inference_steps": self.num_inference_steps,
|
||
"num_frames": self.num_frames,
|
||
"height": self.height,
|
||
"width": self.width,
|
||
"guidance_scale": self.guidance_scale,
|
||
"guidance_rescale": self.guidance_rescale,
|
||
"fps": self.fps,
|
||
"seed": self.seed,
|
||
"negative_prompt": self.negative_prompt,
|
||
"num_gpus": self.num_gpus,
|
||
"dit_cpu_offload": self.dit_cpu_offload,
|
||
"dit_layerwise_offload": self.dit_layerwise_offload,
|
||
"text_encoder_cpu_offload": self.text_encoder_cpu_offload,
|
||
"vae_cpu_offload": self.vae_cpu_offload,
|
||
"image_encoder_cpu_offload": self.image_encoder_cpu_offload,
|
||
"use_fsdp_inference": self.use_fsdp_inference,
|
||
"enable_torch_compile": self.enable_torch_compile,
|
||
"vsa_sparsity": self.vsa_sparsity,
|
||
"tp_size": self.tp_size,
|
||
"sp_size": self.sp_size,
|
||
"data_path": self.data_path,
|
||
"max_train_steps": self.max_train_steps,
|
||
"train_batch_size": self.train_batch_size,
|
||
"learning_rate": self.learning_rate,
|
||
"num_latent_t": self.num_latent_t,
|
||
"num_height": self.height,
|
||
"num_width": self.width,
|
||
"validation_dataset_file": self.validation_dataset_file,
|
||
"lora_rank": self.lora_rank,
|
||
"dmd_use_vsa": self.dmd_use_vsa,
|
||
"dmd_vsa_sparsity": self.dmd_vsa_sparsity,
|
||
"dmd_denoising_steps": self.dmd_denoising_steps,
|
||
"real_score_guidance_scale": self.real_score_guidance_scale,
|
||
"generator_update_interval": self.generator_update_interval,
|
||
"real_score_model_path": self.real_score_model_path or "",
|
||
"fake_score_model_path": self.fake_score_model_path or "",
|
||
"progress": self._log_buf.progress,
|
||
"progress_msg": self._log_buf.progress_msg,
|
||
"phase": self._log_buf.phase,
|
||
}
|
||
|
||
|
||
MINIMAX_H3_REF2VA_PIPELINE = "MiniMaxH3Ref2VAModularPipeline"
|
||
|
||
|
||
def _build_h3_references(raw: list[dict[str, Any]]) -> list[Any]:
|
||
"""Turn the API's reference dicts into MiniMaxH3Reference objects.
|
||
|
||
Imported lazily so the API server starts without pulling in fastvideo.
|
||
"""
|
||
from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
|
||
|
||
built = []
|
||
for i, ref in enumerate(raw):
|
||
source = (ref or {}).get("source")
|
||
if not source:
|
||
raise ValueError(f"reference {i} has no source")
|
||
if not os.path.isfile(source):
|
||
raise ValueError(f"reference {i} source not found: {source}")
|
||
kwargs: dict[str, Any] = {
|
||
"source": source,
|
||
"media_type": (ref.get("media_type") or "image"),
|
||
}
|
||
for opt in ("soundtrack", "fps", "sample_rate"):
|
||
if ref.get(opt) not in (None, ""):
|
||
kwargs[opt] = ref[opt]
|
||
built.append(MiniMaxH3Reference(**kwargs))
|
||
return built
|
||
|
||
|
||
JOB_LOG_FILENAME = "out.log"
|
||
|
||
|
||
def _job_log_path(output_dir: str, job_id: str) -> str:
|
||
"""Each job's log lives beside its outputs: <output_dir>/<job_id>/out.log."""
|
||
return os.path.join(output_dir, job_id, JOB_LOG_FILENAME)
|
||
|
||
|
||
def _decode_references(value: Any) -> list[dict[str, Any]]:
|
||
"""Reference lists round-trip through the DB as JSON text."""
|
||
if not value:
|
||
return []
|
||
if isinstance(value, list):
|
||
return list(value)
|
||
try:
|
||
decoded = json.loads(value)
|
||
except (TypeError, ValueError):
|
||
logger.warning("Could not decode stored references: %r", value)
|
||
return []
|
||
return list(decoded) if isinstance(decoded, list) else []
|
||
|
||
|
||
def _generator_is_alive(generator: Any) -> bool:
|
||
"""True if the generator's worker processes are all still running.
|
||
|
||
A cached VideoGenerator holds a MultiprocExecutor whose workers are separate
|
||
processes; nothing notices when they exit. Probing `proc.is_alive()` is what
|
||
the executor itself uses during shutdown. Anything unexpected in the object
|
||
graph is treated as alive so a probe failure can never wedge the cache.
|
||
"""
|
||
executor = getattr(generator, "executor", None)
|
||
workers = getattr(executor, "workers", None)
|
||
if not workers:
|
||
return True
|
||
try:
|
||
return all(w.proc.is_alive() for w in workers)
|
||
except Exception:
|
||
logger.debug("Worker liveness probe failed", exc_info=True)
|
||
return True
|
||
|
||
|
||
class JobRunner:
|
||
"""Manages video generation jobs, their execution, and generator caching."""
|
||
|
||
def __init__(
|
||
self,
|
||
output_dir: str,
|
||
log_dir: str,
|
||
database: Database,
|
||
verbose: bool = False,
|
||
):
|
||
"""Initialize the job runner.
|
||
|
||
Args:
|
||
output_dir: Directory where generated videos are saved
|
||
log_dir: Directory where job log files are saved
|
||
verbose: Whether to print full tracebacks in error messages
|
||
database: Optional SQLite database for job persistence
|
||
"""
|
||
self.output_dir = output_dir
|
||
self.log_dir = log_dir
|
||
self.verbose = verbose
|
||
self._db = database
|
||
|
||
self._jobs: dict[str, Job] = {}
|
||
self._jobs_lock = threading.Lock()
|
||
self._load_jobs()
|
||
|
||
# Cache of loaded generators keyed by model config so that we only pay
|
||
# the model-loading cost once per model configuration.
|
||
self._generators: dict[tuple, Any] = {}
|
||
self._generators_lock = threading.Lock()
|
||
|
||
# Shared Manager for log queues (avoids spawning a new process per job)
|
||
self._mp_manager = get_mp_context().Manager()
|
||
atexit.register(self._shutdown)
|
||
|
||
# Ensure directories exist
|
||
os.makedirs(self.output_dir, exist_ok=True)
|
||
os.makedirs(self.log_dir, exist_ok=True)
|
||
|
||
def _load_logs(self, job: Job) -> None:
|
||
"""Populate job's log buffer from its log file if it exists."""
|
||
path = job.log_file_path
|
||
if not path:
|
||
path = _job_log_path(self.output_dir, job.id)
|
||
if not os.path.isfile(path):
|
||
return
|
||
try:
|
||
with open(path, encoding="utf-8", errors="replace") as f:
|
||
for line in f:
|
||
stripped = line.rstrip("\n\r")
|
||
if stripped:
|
||
job._log_buf.write(stripped)
|
||
if not job.log_file_path:
|
||
job.log_file_path = path
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to load logs from %s for job %s: %s",
|
||
path,
|
||
job.id,
|
||
exc,
|
||
)
|
||
|
||
def _load_jobs(self) -> None:
|
||
"""Load jobs from database."""
|
||
try:
|
||
for row in self._db.get_all_jobs():
|
||
status = row["status"]
|
||
if status == "running":
|
||
status = JobStatus.FAILED
|
||
row["error"] = "Server restarted (job was running)"
|
||
row["finished_at"] = time.time()
|
||
self._db.update_job(row["id"], {
|
||
"status": "failed",
|
||
"error": row["error"],
|
||
"finished_at": row["finished_at"],
|
||
})
|
||
elif status == "pending":
|
||
status = JobStatus.PENDING
|
||
job = Job(
|
||
id=row["id"],
|
||
model_id=row["model_id"],
|
||
name=row.get("name", "") or "",
|
||
prompt=row["prompt"],
|
||
workload_type=row.get("workload_type", "t2v"),
|
||
job_type=row.get("job_type", "inference"),
|
||
image_path=row.get("image_path", "") or "",
|
||
last_image_path=row.get("last_image_path", "") or "",
|
||
references=_decode_references(row.get("references")),
|
||
data_path=row.get("data_path", "") or "",
|
||
max_train_steps=row.get("max_train_steps", 1000),
|
||
train_batch_size=row.get("train_batch_size", 1),
|
||
learning_rate=float(row.get("learning_rate", 5e-5)),
|
||
num_latent_t=row.get("num_latent_t", 20),
|
||
validation_dataset_file=row.get("validation_dataset_file", "") or "",
|
||
lora_rank=row.get("lora_rank", 32),
|
||
dmd_use_vsa=row.get("dmd_use_vsa", False),
|
||
dmd_vsa_sparsity=float(row.get("dmd_vsa_sparsity", 0.8)),
|
||
dmd_denoising_steps=row.get("dmd_denoising_steps", "1000,757,522") or "1000,757,522",
|
||
real_score_guidance_scale=float(row.get("real_score_guidance_scale", 3.5)),
|
||
generator_update_interval=int(row.get("generator_update_interval", 5)),
|
||
real_score_model_path=row.get("real_score_model_path", "") or "",
|
||
fake_score_model_path=row.get("fake_score_model_path", "") or "",
|
||
status=JobStatus(status),
|
||
created_at=row["created_at"],
|
||
started_at=row.get("started_at"),
|
||
finished_at=row.get("finished_at"),
|
||
error=row.get("error"),
|
||
output_path=row.get("output_path"),
|
||
log_file_path=row.get("log_file_path"),
|
||
num_inference_steps=row.get("num_inference_steps", 50),
|
||
num_frames=row.get("num_frames", 81),
|
||
height=row.get("height", 480),
|
||
width=row.get("width", 832),
|
||
guidance_scale=row.get("guidance_scale", 5.0),
|
||
guidance_rescale=row.get("guidance_rescale", 0.0),
|
||
fps=row.get("fps", 24),
|
||
seed=row.get("seed", 1024),
|
||
negative_prompt=row.get("negative_prompt", "") or "",
|
||
num_gpus=row.get("num_gpus", 1),
|
||
dit_cpu_offload=row.get("dit_cpu_offload", False),
|
||
dit_layerwise_offload=row.get("dit_layerwise_offload", False),
|
||
text_encoder_cpu_offload=row.get("text_encoder_cpu_offload", False),
|
||
vae_cpu_offload=row.get("vae_cpu_offload", False),
|
||
image_encoder_cpu_offload=row.get("image_encoder_cpu_offload", False),
|
||
use_fsdp_inference=row.get("use_fsdp_inference", False),
|
||
enable_torch_compile=row.get("enable_torch_compile", False),
|
||
vsa_sparsity=row.get("vsa_sparsity", 0.0),
|
||
tp_size=row.get("tp_size", -1),
|
||
sp_size=row.get("sp_size", -1),
|
||
)
|
||
self._load_logs(job)
|
||
with self._jobs_lock:
|
||
self._jobs[job.id] = job
|
||
if self._jobs:
|
||
logger.info("Loaded %d jobs from database", len(self._jobs))
|
||
except Exception as exc:
|
||
logger.warning("Failed to load jobs from database: %s", exc)
|
||
|
||
def _save_job(self, job: Job) -> None:
|
||
"""Persist job to database."""
|
||
try:
|
||
self._db.update_job(
|
||
job.id, {
|
||
"status": job.status.value,
|
||
"started_at": job.started_at,
|
||
"finished_at": job.finished_at,
|
||
"error": job.error,
|
||
"output_path": job.output_path,
|
||
"log_file_path": job.log_file_path,
|
||
})
|
||
except Exception as exc:
|
||
logger.warning("Failed to persist job %s: %s", job.id, exc)
|
||
|
||
def _shutdown(self) -> None:
|
||
"""Shutdown the shared multiprocessing manager on exit."""
|
||
try:
|
||
self._mp_manager.shutdown()
|
||
except Exception as exc:
|
||
logger.warning("Failed to shutdown mp manager: %s", exc)
|
||
|
||
def create_job(
|
||
self,
|
||
job_id: str,
|
||
model_id: str,
|
||
prompt: str,
|
||
name: str = "",
|
||
workload_type: str = "t2v",
|
||
job_type: str = "inference",
|
||
image_path: str = "",
|
||
last_image_path: str = "",
|
||
references: list[dict[str, Any]] | None = None,
|
||
data_path: str = "",
|
||
max_train_steps: int = 1000,
|
||
train_batch_size: int = 1,
|
||
learning_rate: float = 5e-5,
|
||
num_latent_t: int = 20,
|
||
validation_dataset_file: str = "",
|
||
lora_rank: int = 32,
|
||
dmd_use_vsa: bool = False,
|
||
dmd_vsa_sparsity: float = 0.8,
|
||
dmd_denoising_steps: str = "1000,757,522",
|
||
real_score_guidance_scale: float = 3.5,
|
||
generator_update_interval: int = 5,
|
||
real_score_model_path: str = "",
|
||
fake_score_model_path: str = "",
|
||
num_inference_steps: int = 50,
|
||
num_frames: int = 81,
|
||
height: int = 480,
|
||
width: int = 832,
|
||
guidance_scale: float = 5.0,
|
||
guidance_rescale: float = 0.0,
|
||
fps: int = 24,
|
||
seed: int = 1024,
|
||
num_gpus: int = 1,
|
||
negative_prompt: str = "",
|
||
dit_cpu_offload: bool = False,
|
||
dit_layerwise_offload: bool = False,
|
||
text_encoder_cpu_offload: bool = False,
|
||
vae_cpu_offload: bool = False,
|
||
image_encoder_cpu_offload: bool = False,
|
||
use_fsdp_inference: bool = False,
|
||
enable_torch_compile: bool = False,
|
||
vsa_sparsity: float = 0.0,
|
||
tp_size: int = -1,
|
||
sp_size: int = -1,
|
||
) -> Job:
|
||
"""Create a new job (does not start it automatically)."""
|
||
job = Job(
|
||
id=job_id,
|
||
model_id=model_id,
|
||
name=(name or "").strip(),
|
||
prompt=prompt.strip(),
|
||
workload_type=workload_type or "t2v",
|
||
job_type=job_type or "inference",
|
||
image_path=image_path or "",
|
||
last_image_path=last_image_path or "",
|
||
references=list(references or []),
|
||
data_path=data_path or "",
|
||
max_train_steps=max_train_steps,
|
||
train_batch_size=train_batch_size,
|
||
learning_rate=learning_rate,
|
||
num_latent_t=num_latent_t,
|
||
validation_dataset_file=validation_dataset_file or "",
|
||
lora_rank=lora_rank,
|
||
dmd_use_vsa=dmd_use_vsa,
|
||
dmd_vsa_sparsity=dmd_vsa_sparsity,
|
||
dmd_denoising_steps=dmd_denoising_steps,
|
||
real_score_guidance_scale=real_score_guidance_scale,
|
||
generator_update_interval=generator_update_interval,
|
||
real_score_model_path=real_score_model_path or "",
|
||
fake_score_model_path=fake_score_model_path or "",
|
||
num_inference_steps=num_inference_steps,
|
||
num_frames=num_frames,
|
||
height=height,
|
||
width=width,
|
||
guidance_scale=guidance_scale,
|
||
guidance_rescale=guidance_rescale,
|
||
fps=fps,
|
||
seed=seed,
|
||
negative_prompt=negative_prompt or "",
|
||
num_gpus=num_gpus,
|
||
dit_cpu_offload=dit_cpu_offload,
|
||
dit_layerwise_offload=dit_layerwise_offload,
|
||
text_encoder_cpu_offload=text_encoder_cpu_offload,
|
||
vae_cpu_offload=vae_cpu_offload,
|
||
image_encoder_cpu_offload=image_encoder_cpu_offload,
|
||
use_fsdp_inference=use_fsdp_inference,
|
||
enable_torch_compile=enable_torch_compile,
|
||
vsa_sparsity=vsa_sparsity,
|
||
tp_size=tp_size,
|
||
sp_size=sp_size,
|
||
)
|
||
with self._jobs_lock:
|
||
self._jobs[job.id] = job
|
||
try:
|
||
self._db.insert_job(job.to_dict())
|
||
except Exception as exc:
|
||
logger.warning("Failed to persist new job %s: %s", job.id, exc)
|
||
logger.info(
|
||
"Created job %s (model=%s, prompt=%s…)",
|
||
job.id,
|
||
job.model_id,
|
||
job.prompt[:60],
|
||
)
|
||
return job
|
||
|
||
def get_job(self, job_id: str) -> Job | None:
|
||
"""Get a job by ID."""
|
||
with self._jobs_lock:
|
||
return self._jobs.get(job_id)
|
||
|
||
def list_jobs(self, job_type: str | None = None) -> list[Job]:
|
||
"""Return jobs, sorted by creation time (newest first).
|
||
If job_type is set, filter to that type only."""
|
||
with self._jobs_lock:
|
||
jobs = list(self._jobs.values())
|
||
if job_type:
|
||
jobs = [j for j in jobs if j.job_type == job_type]
|
||
return sorted(jobs, key=lambda j: j.created_at, reverse=True)
|
||
|
||
def delete_job(self, job_id: str) -> bool:
|
||
"""Delete a job. Running jobs are stopped first.
|
||
|
||
Returns:
|
||
True if job was found and deleted, False otherwise
|
||
"""
|
||
with self._jobs_lock:
|
||
job = self._jobs.pop(job_id, None)
|
||
if job is None:
|
||
return False
|
||
|
||
# Best-effort stop
|
||
job._stop_event.set()
|
||
try:
|
||
self._db.delete_job(job_id)
|
||
except Exception as exc:
|
||
logger.warning("Failed to delete job %s from database: %s", job_id, exc)
|
||
logger.info("Deleted job %s", job.id)
|
||
return True
|
||
|
||
CONFIG_FIELDS: tuple[str, ...] = (
|
||
"model_id",
|
||
"name",
|
||
"prompt",
|
||
"workload_type",
|
||
"job_type",
|
||
"image_path",
|
||
"last_image_path",
|
||
"references",
|
||
"negative_prompt",
|
||
"num_inference_steps",
|
||
"num_frames",
|
||
"height",
|
||
"width",
|
||
"guidance_scale",
|
||
"guidance_rescale",
|
||
"fps",
|
||
"seed",
|
||
"num_gpus",
|
||
"dit_cpu_offload",
|
||
"dit_layerwise_offload",
|
||
"text_encoder_cpu_offload",
|
||
"vae_cpu_offload",
|
||
"image_encoder_cpu_offload",
|
||
"use_fsdp_inference",
|
||
"enable_torch_compile",
|
||
"vsa_sparsity",
|
||
"tp_size",
|
||
"sp_size",
|
||
"data_path",
|
||
"max_train_steps",
|
||
"train_batch_size",
|
||
"learning_rate",
|
||
"num_latent_t",
|
||
"validation_dataset_file",
|
||
"lora_rank",
|
||
"dmd_use_vsa",
|
||
"dmd_vsa_sparsity",
|
||
"dmd_denoising_steps",
|
||
"real_score_guidance_scale",
|
||
"generator_update_interval",
|
||
"real_score_model_path",
|
||
"fake_score_model_path",
|
||
)
|
||
|
||
def duplicate_job(self, job_id: str, new_job_id: str) -> Job:
|
||
"""Create a new pending job with an existing job's configuration.
|
||
|
||
Runtime state (status, timings, logs, outputs) is not carried over.
|
||
"""
|
||
with self._jobs_lock:
|
||
source = self._jobs.get(job_id)
|
||
if source is None:
|
||
raise ValueError(f"Job {job_id} not found")
|
||
config = {f: copy.deepcopy(getattr(source, f)) for f in self.CONFIG_FIELDS}
|
||
return self.create_job(job_id=new_job_id, **config)
|
||
|
||
#: Editable exactly when startable: the same set start_job() accepts.
|
||
EDITABLE_STATUSES = (JobStatus.PENDING, JobStatus.FAILED, JobStatus.STOPPED)
|
||
|
||
def update_job_config(self, job_id: str, updates: dict[str, Any]) -> Job:
|
||
"""Edit the configuration of a job that has not produced a result."""
|
||
with self._jobs_lock:
|
||
job = self._jobs.get(job_id)
|
||
if job is None:
|
||
raise ValueError(f"Job {job_id} not found")
|
||
if job.status not in self.EDITABLE_STATUSES:
|
||
allowed = ", ".join(s.value for s in self.EDITABLE_STATUSES)
|
||
raise ValueError(f"Job is {job.status.value}; only {allowed} jobs can be edited. "
|
||
"Duplicate it instead.")
|
||
unknown = set(updates) - set(self.CONFIG_FIELDS)
|
||
if unknown:
|
||
raise ValueError(f"Not editable: {', '.join(sorted(unknown))}")
|
||
for field_name, value in updates.items():
|
||
setattr(job, field_name, value)
|
||
self._save_job(job)
|
||
return job
|
||
|
||
def start_job(self, job_id: str) -> Job:
|
||
"""Start (or restart) a pending / stopped / failed job.
|
||
|
||
Raises:
|
||
ValueError: If job not found or cannot be started
|
||
"""
|
||
with self._jobs_lock:
|
||
job = self._jobs.get(job_id)
|
||
if job is None:
|
||
raise ValueError(f"Job {job_id} not found")
|
||
if job.status == JobStatus.RUNNING:
|
||
raise ValueError("Job is already running")
|
||
if job.status == JobStatus.COMPLETED:
|
||
raise ValueError("Job already completed. Delete and re-create to run again.")
|
||
# Reset state for re-run
|
||
job.status = JobStatus.PENDING
|
||
job.error = None
|
||
job.output_path = None
|
||
job.log_file_path = None # Will be set when job starts
|
||
job.started_at = None
|
||
job.finished_at = None
|
||
job._stop_event.clear()
|
||
job._log_buf = JobLogBuffer() # fresh log buffer
|
||
job.log_file_handler = None # Reset file handler
|
||
job._process = None # Reset subprocess handle
|
||
|
||
# Wrap _run_job in an additional safety layer to catch any exceptions
|
||
# that might escape (though they shouldn't with our comprehensive handling)
|
||
def safe_run_job(job: Job):
|
||
"""Wrapper to ensure _run_job never raises an unhandled exception."""
|
||
try:
|
||
self._run_job(job)
|
||
except BaseException as exc:
|
||
# This should never happen, but if it does, we catch it here
|
||
logger.critical("Unhandled exception escaped from _run_job for job %s: %s", job.id, exc, exc_info=True)
|
||
with contextlib.suppress(Exception):
|
||
job.status = JobStatus.FAILED
|
||
job.error = (f"Unhandled exception: {type(exc).__name__}: {str(exc)}")
|
||
job.finished_at = time.time()
|
||
|
||
thread = threading.Thread(target=safe_run_job, args=(job, ), daemon=True)
|
||
job._thread = thread
|
||
thread.start()
|
||
logger.info("Started job %s", job.id)
|
||
return job
|
||
|
||
def stop_job(self, job_id: str) -> Job:
|
||
"""Request a running job to stop.
|
||
|
||
For inference: cooperative stop between phases.
|
||
For training: terminates the subprocess.
|
||
|
||
Raises:
|
||
ValueError: If job not found or not running
|
||
"""
|
||
with self._jobs_lock:
|
||
job = self._jobs.get(job_id)
|
||
if job is None:
|
||
raise ValueError(f"Job {job_id} not found")
|
||
if job.status != JobStatus.RUNNING:
|
||
raise ValueError(f"Job is not running (status={job.status.value})")
|
||
|
||
job._stop_event.set()
|
||
if job._process is not None:
|
||
with contextlib.suppress(Exception):
|
||
job._process.terminate()
|
||
logger.info("Stop requested for job %s", job.id)
|
||
return job
|
||
|
||
def get_job_logs(self, job_id: str, after: int = 0) -> dict[str, Any]:
|
||
"""Return log lines for a job.
|
||
|
||
Args:
|
||
job_id: The job ID
|
||
after: Return only lines after this index (for incremental polling)
|
||
|
||
Returns:
|
||
Dictionary with 'lines', 'total', 'progress', 'progress_msg', 'phase'
|
||
|
||
Raises:
|
||
ValueError: If job not found
|
||
"""
|
||
with self._jobs_lock:
|
||
job = self._jobs.get(job_id)
|
||
if job is None:
|
||
raise ValueError(f"Job {job_id} not found")
|
||
|
||
lines, total = job._log_buf.get_lines(after=after)
|
||
return {
|
||
"lines": lines,
|
||
"total": total,
|
||
"progress": job._log_buf.progress,
|
||
"progress_msg": job._log_buf.progress_msg,
|
||
"phase": job._log_buf.phase,
|
||
}
|
||
|
||
def _get_or_create_generator(
|
||
self,
|
||
model_id: str,
|
||
workload_type: str,
|
||
num_gpus: int,
|
||
dit_cpu_offload: bool = False,
|
||
dit_layerwise_offload: bool = False,
|
||
override_pipeline_cls_name: str | None = None,
|
||
text_encoder_cpu_offload: bool = False,
|
||
vae_cpu_offload: bool = False,
|
||
image_encoder_cpu_offload: bool = False,
|
||
use_fsdp_inference: bool = False,
|
||
enable_torch_compile: bool = False,
|
||
vsa_sparsity: float = 0.0,
|
||
tp_size: int = -1,
|
||
sp_size: int = -1,
|
||
log_queue: mp.Queue | None = None,
|
||
) -> Any:
|
||
cache_key = (
|
||
model_id,
|
||
workload_type,
|
||
num_gpus,
|
||
dit_cpu_offload,
|
||
dit_layerwise_offload,
|
||
# Ref2VA loads different DiT weights (transformer_ref), so the
|
||
# override must key the cache or a t2v/i2v generator gets reused.
|
||
override_pipeline_cls_name,
|
||
text_encoder_cpu_offload,
|
||
vae_cpu_offload,
|
||
image_encoder_cpu_offload,
|
||
use_fsdp_inference,
|
||
enable_torch_compile,
|
||
vsa_sparsity,
|
||
tp_size,
|
||
sp_size,
|
||
)
|
||
|
||
# Generators are cached by model_id and configuration parameters
|
||
with self._generators_lock:
|
||
cached = self._generators.get(cache_key)
|
||
if cached is not None:
|
||
if _generator_is_alive(cached):
|
||
return cached
|
||
# Workers can exit while a generator sits idle in the cache;
|
||
# reusing it fails every later job with the same config.
|
||
logger.warning(
|
||
"Cached generator for %s has dead workers; reloading.",
|
||
model_id,
|
||
)
|
||
self._generators.pop(cache_key, None)
|
||
try:
|
||
cached.shutdown()
|
||
except Exception:
|
||
logger.debug("Shutdown of the dead generator failed", exc_info=True)
|
||
|
||
# Import lazily so starting the server is fast even without a GPU.
|
||
from fastvideo import VideoGenerator
|
||
|
||
logger.info(
|
||
"Loading model %s (workload=%s, num_gpus=%d, offloads: "
|
||
"dit=%s text_encoder=%s vae=%s image_encoder=%s, fsdp=%s, "
|
||
"torch_compile=%s, vsa_sparsity=%.2f, tp=%d sp=%d) …",
|
||
model_id,
|
||
workload_type,
|
||
num_gpus,
|
||
dit_cpu_offload,
|
||
text_encoder_cpu_offload,
|
||
vae_cpu_offload,
|
||
image_encoder_cpu_offload,
|
||
use_fsdp_inference,
|
||
enable_torch_compile,
|
||
vsa_sparsity,
|
||
tp_size,
|
||
sp_size,
|
||
)
|
||
|
||
gen = VideoGenerator.from_pretrained(
|
||
model_id,
|
||
workload_type=workload_type,
|
||
num_gpus=num_gpus,
|
||
dit_layerwise_offload=dit_layerwise_offload,
|
||
**({
|
||
"override_pipeline_cls_name": override_pipeline_cls_name
|
||
} if override_pipeline_cls_name else {}),
|
||
dit_cpu_offload=dit_cpu_offload,
|
||
text_encoder_cpu_offload=text_encoder_cpu_offload,
|
||
vae_cpu_offload=vae_cpu_offload,
|
||
image_encoder_cpu_offload=image_encoder_cpu_offload,
|
||
use_fsdp_inference=use_fsdp_inference,
|
||
enable_torch_compile=enable_torch_compile,
|
||
VSA_sparsity=vsa_sparsity,
|
||
tp_size=tp_size,
|
||
sp_size=sp_size,
|
||
log_queue=log_queue,
|
||
)
|
||
|
||
with self._generators_lock:
|
||
if cache_key not in self._generators:
|
||
self._generators[cache_key] = gen
|
||
else: # Another thread may have created it while we were loading.
|
||
gen.shutdown()
|
||
gen = self._generators[cache_key]
|
||
return gen
|
||
|
||
def _run_job(self, job: Job):
|
||
if job.job_type == "inference":
|
||
self._run_inference_job(job)
|
||
else:
|
||
self._run_training_job(job)
|
||
|
||
def _run_training_job(self, job: Job):
|
||
"""Run a finetuning, distillation, or LoRA job via subprocess."""
|
||
buf = job._log_buf
|
||
job_output_dir = os.path.join(self.output_dir, job.id)
|
||
os.makedirs(job_output_dir, exist_ok=True)
|
||
job.log_file_path = _job_log_path(self.output_dir, job.id)
|
||
|
||
if not job.data_path or not os.path.isdir(job.data_path):
|
||
job.status = JobStatus.FAILED
|
||
job.error = (f"Data path '{job.data_path}' is required and must be an "
|
||
"existing directory. Preprocess your dataset first.")
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
return
|
||
|
||
env = os.environ.copy()
|
||
env.update(get_training_env())
|
||
|
||
try:
|
||
# Job.to_dict() carries every key the config builder reads
|
||
# (extra keys are ignored by its .get() lookups).
|
||
train_config = build_training_config(job.to_dict(), job_output_dir)
|
||
except ValueError as exc:
|
||
job.status = JobStatus.FAILED
|
||
job.error = str(exc)
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
return
|
||
|
||
config_path = os.path.join(job_output_dir, "train_config.yaml")
|
||
with open(config_path, "w", encoding="utf-8") as f:
|
||
yaml.safe_dump(train_config, f, sort_keys=False)
|
||
|
||
repo_root = Path(__file__).resolve().parent.parent
|
||
torchrun_cmd = [
|
||
sys.executable,
|
||
"-m",
|
||
"torch.distributed.run",
|
||
"--nproc_per_node",
|
||
str(job.num_gpus),
|
||
"--nnodes",
|
||
"1",
|
||
"-m",
|
||
"fastvideo.train.entrypoint.train",
|
||
"--config",
|
||
config_path,
|
||
]
|
||
buf.write(f"Starting training: {' '.join(torchrun_cmd)}")
|
||
buf.phase = "starting"
|
||
|
||
try:
|
||
job.status = JobStatus.RUNNING
|
||
job.started_at = time.time()
|
||
self._save_job(job)
|
||
|
||
with open(job.log_file_path, "w", encoding="utf-8") as log_file:
|
||
job._process = subprocess.Popen(
|
||
torchrun_cmd,
|
||
cwd=str(repo_root),
|
||
env=env,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
bufsize=1,
|
||
)
|
||
assert job._process.stdout is not None
|
||
for line in iter(job._process.stdout.readline, ""):
|
||
if job._stop_event.is_set():
|
||
job._process.terminate()
|
||
try:
|
||
job._process.wait(timeout=30)
|
||
except subprocess.TimeoutExpired:
|
||
job._process.kill()
|
||
job.status = JobStatus.STOPPED
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
buf.phase = "stopped"
|
||
return
|
||
line = line.rstrip()
|
||
if line:
|
||
buf.write(line)
|
||
log_file.write(line + "\n")
|
||
log_file.flush()
|
||
|
||
job._process.wait()
|
||
exit_code = job._process.returncode or 0
|
||
|
||
if job._stop_event.is_set():
|
||
# Terminated by stop_job() without the reader loop observing the
|
||
# flag (e.g. the process died between log lines).
|
||
job.status = JobStatus.STOPPED
|
||
buf.phase = "stopped"
|
||
elif exit_code == 0:
|
||
job.status = JobStatus.COMPLETED
|
||
buf.progress = 100.0
|
||
buf.phase = "done"
|
||
# Training outputs checkpoints, not video. Sort by step number,
|
||
# not lexically ("checkpoint-1000" < "checkpoint-500" as strings).
|
||
ckpt_dirs = sorted(
|
||
Path(job_output_dir).glob("checkpoint-*"),
|
||
key=lambda p: int(m.group(1)) if (m := re.fullmatch(r"checkpoint-(\d+)", p.name)) else -1,
|
||
)
|
||
if ckpt_dirs:
|
||
job.output_path = str(ckpt_dirs[-1])
|
||
else:
|
||
job.status = JobStatus.FAILED
|
||
job.error = f"Training exited with code {exit_code}"
|
||
buf.phase = "failed"
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
|
||
except Exception as exc:
|
||
job.status = JobStatus.FAILED
|
||
job.error = f"{type(exc).__name__}: {exc}"
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
buf.phase = "failed"
|
||
logger.exception("Training job %s failed", job.id)
|
||
finally:
|
||
job._process = None
|
||
|
||
def _run_inference_job(self, job: Job):
|
||
buf = job._log_buf
|
||
os.makedirs(os.path.join(self.output_dir, job.id), exist_ok=True)
|
||
job.log_file_path = _job_log_path(self.output_dir, job.id)
|
||
|
||
# Add file handler to persist logs
|
||
file_handler = logging.FileHandler(job.log_file_path, mode='w', encoding='utf-8')
|
||
file_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s"))
|
||
job.log_file_handler = file_handler
|
||
|
||
# Hook logger output into job log buffer (main process logs)
|
||
fastvideo_logger = logging.getLogger("fastvideo")
|
||
buffer_handler = LogBufferHandler(buf)
|
||
fastvideo_logger.addHandler(buffer_handler)
|
||
fastvideo_logger.addHandler(file_handler)
|
||
|
||
# Queue for worker process logs (fsdp_load, cuda, etc.)
|
||
# Use Manager().Queue() so it can be shared with spawned workers (spawn
|
||
# does not inherit memory; mp.Queue only works through inheritance).
|
||
log_queue = self._mp_manager.Queue()
|
||
queue_listener = logging.handlers.QueueListener(log_queue,
|
||
buffer_handler,
|
||
file_handler,
|
||
respect_handler_level=True)
|
||
queue_listener.start()
|
||
|
||
# Set output directory, create if it doesn't exist
|
||
job_output_dir = os.path.join(self.output_dir, job.id)
|
||
os.makedirs(job_output_dir, exist_ok=True)
|
||
|
||
try:
|
||
job.status = JobStatus.RUNNING
|
||
job.started_at = time.time()
|
||
self._save_job(job)
|
||
buf.phase = "starting"
|
||
started = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(job.started_at))
|
||
logger.info("Job %s started at %s", job.id, started)
|
||
logger.info("Model: %s", job.model_id)
|
||
logger.info("Prompt: %s", job.prompt)
|
||
|
||
if job._stop_event.is_set():
|
||
job.status = JobStatus.STOPPED
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
logger.warning("Job stopped before execution started")
|
||
return
|
||
|
||
buf.phase = "loading model"
|
||
logger.info("Loading model...")
|
||
|
||
# The generator MUST be created on this thread: building it spawns
|
||
# the executor's worker processes, and they are torn down if the
|
||
# creating thread exits. Running it in a helper thread (to poll
|
||
# _stop_event during load) made every collective_rpc fail with
|
||
# ConnectionResetError.
|
||
if job._stop_event.is_set():
|
||
job.status = JobStatus.STOPPED
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
logger.warning("Job %s stopped before model loading", job.id)
|
||
buf.phase = "stopped"
|
||
return
|
||
|
||
generator = self._get_or_create_generator(
|
||
job.model_id,
|
||
job.workload_type,
|
||
job.num_gpus,
|
||
dit_cpu_offload=job.dit_cpu_offload,
|
||
dit_layerwise_offload=job.dit_layerwise_offload,
|
||
override_pipeline_cls_name=(MINIMAX_H3_REF2VA_PIPELINE if job.references else None),
|
||
text_encoder_cpu_offload=(job.text_encoder_cpu_offload),
|
||
vae_cpu_offload=job.vae_cpu_offload,
|
||
image_encoder_cpu_offload=(job.image_encoder_cpu_offload),
|
||
use_fsdp_inference=job.use_fsdp_inference,
|
||
enable_torch_compile=(job.enable_torch_compile),
|
||
vsa_sparsity=job.vsa_sparsity,
|
||
tp_size=job.tp_size,
|
||
sp_size=job.sp_size,
|
||
log_queue=log_queue,
|
||
)
|
||
buf.phase = "generating"
|
||
logger.info("Starting generation for job %s (model=%s)", job.id, job.model_id)
|
||
|
||
# Without a name FastVideo derives the filename from the prompt.
|
||
safe_name = re.sub(r'[\\/:*?"<>|]+', "", job.name).strip().strip(".")
|
||
output_target = (os.path.join(job_output_dir, f"{safe_name[:80]}.mp4") if safe_name else job_output_dir)
|
||
gen_kwargs: dict[str, Any] = {
|
||
"prompt": job.prompt,
|
||
"output_path": output_target,
|
||
"save_video": True,
|
||
"num_inference_steps": job.num_inference_steps,
|
||
"num_frames": job.num_frames,
|
||
"height": job.height,
|
||
"width": job.width,
|
||
"guidance_scale": job.guidance_scale,
|
||
"guidance_rescale": job.guidance_rescale,
|
||
"fps": job.fps,
|
||
"seed": job.seed,
|
||
"negative_prompt": job.negative_prompt or "",
|
||
"log_queue": log_queue,
|
||
}
|
||
if job.image_path:
|
||
gen_kwargs["image_path"] = job.image_path
|
||
if job.references:
|
||
gen_kwargs["references"] = _build_h3_references(job.references)
|
||
if job.last_image_path:
|
||
# _prepare_fl2va requires a PIL image, not a path.
|
||
from PIL import Image as _PILImage
|
||
gen_kwargs["last_image"] = _PILImage.open(job.last_image_path)
|
||
generator.generate_video(**gen_kwargs)
|
||
|
||
buf.phase = "saving"
|
||
logger.info("Generation completed, searching for output file...")
|
||
|
||
# Find the generated video file
|
||
video_files = sorted(Path(job_output_dir).glob("*.mp4"))
|
||
if video_files:
|
||
job.output_path = str(video_files[0])
|
||
logger.info("Found video output: %s", job.output_path)
|
||
else:
|
||
# Could be an image workload
|
||
image_files = sorted(Path(job_output_dir).glob("*.png"))
|
||
if image_files:
|
||
job.output_path = str(image_files[0])
|
||
logger.info("Found image output: %s", job.output_path)
|
||
else:
|
||
logger.warning("No output file found in job directory")
|
||
|
||
if job._stop_event.is_set():
|
||
job.status = JobStatus.STOPPED
|
||
logger.warning("Job was stopped during execution")
|
||
else:
|
||
job.status = JobStatus.COMPLETED
|
||
buf.progress = 100.0
|
||
logger.info("Job completed successfully")
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
buf.phase = "done"
|
||
|
||
except Exception as exception:
|
||
error_msg = str(exception)
|
||
logger.exception("Critical error in job thread: %s", error_msg)
|
||
job.status = JobStatus.FAILED
|
||
job.error = f"Critical error ({type(exception).__name__}): {error_msg}"
|
||
job.finished_at = time.time()
|
||
self._save_job(job)
|
||
buf.phase = "failed"
|
||
|
||
finally:
|
||
queue_listener.stop()
|
||
# Remove handlers and close file
|
||
fastvideo_logger.removeHandler(buffer_handler)
|
||
fastvideo_logger.removeHandler(file_handler)
|
||
file_handler.flush()
|
||
file_handler.close()
|
||
job.log_file_handler = None
|