livef12rocks/app/worker.py

"""Inbox watcher + sequential job runner.

Polls `DATA_DIR/inbox` for new, size-stable image files, copies each one
into its own `jobs/<job_id>/original.*` and runs the compose pipeline on
it. Files are NEVER deleted or moved from the inbox — `processed.json`
tracks what has already been handled (by path + size + mtime) so restarts
don't reprocess everything.

Runs one job at a time (no threading) — this is intentional: the compose
pipeline is CPU heavy (gmic/rembg) and the worker container is capped at
`cpus: "1.0"` in compose.yml, so concurrency would only cause thrashing.
"""

from __future__ import annotations

import json
import logging
import shutil
import time
from pathlib import Path
from typing import Any

from . import config, pipeline

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
)
logger = logging.getLogger("livef12.worker")


def ensure_dirs() -> None:
    for path in (config.INBOX_DIR, config.JOBS_DIR):
        path.mkdir(parents=True, exist_ok=True)


def load_processed() -> dict[str, Any]:
    if not config.PROCESSED_FILE.exists():
        return {}
    try:
        return json.loads(config.PROCESSED_FILE.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        logger.warning("processed.json unreadable, starting fresh")
        return {}


def save_processed(processed: dict[str, Any]) -> None:
    tmp = config.PROCESSED_FILE.with_suffix(".json.tmp")
    tmp.write_text(json.dumps(processed, indent=2, ensure_ascii=False), encoding="utf-8")
    tmp.replace(config.PROCESSED_FILE)


def _file_key(path: Path) -> str:
    return str(path.relative_to(config.INBOX_DIR))


def _is_already_processed(processed: dict[str, Any], path: Path, stat: Any) -> bool:
    entry = processed.get(_file_key(path))
    if not entry:
        return False
    return entry.get("size") == stat.st_size and entry.get("mtime") == stat.st_mtime


def _is_stable(path: Path) -> bool:
    """A file is "stable" if its size doesn't change across a short wait —
    cheap way to avoid picking up a half-uploaded SFTP transfer."""
    try:
        size_before = path.stat().st_size
    except OSError:
        return False
    time.sleep(config.STABLE_WAIT_SECONDS)
    try:
        size_after = path.stat().st_size
    except OSError:
        return False
    return size_before == size_after and size_after > 0


def find_new_files(processed: dict[str, Any]) -> list[Path]:
    if not config.INBOX_DIR.exists():
        return []
    candidates: list[Path] = []
    for path in sorted(config.INBOX_DIR.rglob("*")):
        if not path.is_file():
            continue
        if path.suffix.lower() not in config.SUPPORTED_EXTENSIONS:
            continue
        try:
            stat = path.stat()
        except OSError:
            continue
        if _is_already_processed(processed, path, stat):
            continue
        candidates.append(path)
    return candidates


def handle_file(path: Path, processed: dict[str, Any]) -> None:
    if not _is_stable(path):
        logger.info("skip %s: still being written", path.name)
        return

    stat = path.stat()
    job_id = pipeline.new_job_id()
    suffix = path.suffix.lower()
    paths = pipeline.job_paths(job_id, suffix)
    paths.root.mkdir(parents=True, exist_ok=True)

    logger.info("new inbox file %s -> job %s", path.name, job_id)
    # Copy (never move) so the inbox stays untouched.
    shutil.copy2(path, paths.original)

    try:
        pipeline.process_job(job_id, paths.original, suffix)
    finally:
        # Mark as processed regardless of pipeline outcome so a permanently
        # broken image doesn't get retried forever; failures are visible in
        # jobs/<id>/status.json for manual follow-up.
        processed[_file_key(path)] = {
            "size": stat.st_size,
            "mtime": stat.st_mtime,
            "job_id": job_id,
            "processed_at": pipeline.now_iso(),
        }
        save_processed(processed)


def scan_once(processed: dict[str, Any]) -> None:
    for path in find_new_files(processed):
        try:
            handle_file(path, processed)
        except Exception:  # noqa: BLE001 - one bad file must not kill the loop
            logger.exception("failed to handle %s", path)


def main() -> int:
    ensure_dirs()
    processed = load_processed()
    logger.info(
        "worker started, watching %s (output_count=%d, nice=%d)",
        config.INBOX_DIR,
        config.OUTPUT_COUNT,
        config.NICE_LEVEL,
    )
    while True:
        try:
            scan_once(processed)
        except Exception:  # noqa: BLE001
            logger.exception("scan cycle failed")
        time.sleep(config.POLL_INTERVAL_SECONDS)


if __name__ == "__main__":
    raise SystemExit(main())