runs.py

"""Local persistence for generation runs under /data.

Each run stores the uploaded input image, one output video per provider,
and its own run.json metadata file inside its run directory (rather than
an entry appended to one shared index file), so concurrent runs never
contend for the same file on disk.

On the actual Hugging Face Space, /data is the mounted persistent storage
directory. Locally (outside the Space container) /data usually doesn't
exist or isn't writable, so ensure_data_dirs() falls back to a ./data
directory next to this file for development.
"""

from __future__ import annotations

import json
import logging
import os
import shutil
import time
import uuid
from dataclasses import asdict, dataclass, field

logger = logging.getLogger(__name__)

DATA_DIR = "/data"
_LOCAL_FALLBACK_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
RUNS_DIR = os.path.join(DATA_DIR, "runs")


@dataclass
class RunResult:
    provider_name: str
    output_path: str | None
    status: str  # "ok" | "error"
    error: str | None
    duration_seconds: float
    params_used: dict


@dataclass
class Run:
    run_id: str
    timestamp: float
    input_path: str
    prompt: str | None = None
    results: list[RunResult] = field(default_factory=list)


def ensure_data_dirs() -> None:
    """Creates the runs directory, falling back to a local ./data dir if
    /data isn't writable (i.e. not running inside the Space container)."""
    global DATA_DIR, RUNS_DIR
    try:
        os.makedirs(RUNS_DIR, exist_ok=True)
    except OSError:
        logger.warning("%s is not writable; falling back to local dir %s", DATA_DIR, _LOCAL_FALLBACK_DIR)
        DATA_DIR = _LOCAL_FALLBACK_DIR
        RUNS_DIR = os.path.join(DATA_DIR, "runs")
        os.makedirs(RUNS_DIR, exist_ok=True)


def get_data_dir() -> str:
    """Returns the data directory actually in use (call after
    ensure_data_dirs(), since it may fall back to a local directory)."""
    return DATA_DIR


def create_run(image_path: str) -> tuple[str, str, str]:
    """Creates a new run directory and copies the uploaded image into it as the input.

    Returns (run_id, run_dir, input_path).
    """
    run_id = f"{int(time.time())}-{uuid.uuid4().hex[:8]}"
    run_dir = os.path.join(RUNS_DIR, run_id)
    os.makedirs(run_dir, exist_ok=True)

    _, ext = os.path.splitext(image_path)
    input_path = os.path.join(run_dir, f"input{ext or '.png'}")
    shutil.copyfile(image_path, input_path)

    return run_id, run_dir, input_path


def save_output_video(run_dir: str, provider_name: str, video_bytes: bytes) -> str:
    output_path = os.path.join(run_dir, f"{provider_name}.mp4")
    with open(output_path, "wb") as video_file:
        video_file.write(video_bytes)
    return output_path


def save_run(run: Run) -> None:
    """Writes a run's metadata to its own run.json inside its run directory
    (data/runs/{run_id}/run.json), instead of a shared index file, so
    concurrent runs never race on read-modify-write of the same file."""
    run_path = os.path.join(RUNS_DIR, run.run_id, "run.json")
    try:
        with open(run_path, "w") as run_file:
            json.dump(asdict(run), run_file, indent=2)
    except OSError:
        logger.exception("Failed to save run %s to %s", run.run_id, run_path)
        raise


def load_runs() -> list[Run]:
    """Loads every run by scanning each run directory for its run.json."""
    runs = []
    if not os.path.isdir(RUNS_DIR):
        return runs

    for run_id in sorted(os.listdir(RUNS_DIR)):
        run_path = os.path.join(RUNS_DIR, run_id, "run.json")
        try:
            with open(run_path, "r") as run_file:
                entry = json.load(run_file)
        except FileNotFoundError:
            # Legacy/partial run dirs can exist without metadata; skip quietly.
            continue
        except (NotADirectoryError, json.JSONDecodeError, OSError) as exc:
            logger.warning("Skipping run %s: could not read %s (%s)", run_id, run_path, exc)
            continue

        results = [RunResult(**result) for result in entry.get("results", [])]
        runs.append(
            Run(
                run_id=entry["run_id"],
                timestamp=entry["timestamp"],
                input_path=entry["input_path"],
                prompt=entry.get("prompt"),
                results=results,
            )
        )

    runs.sort(key=lambda run: run.timestamp)
    return runs