app/gallery_worker.py

"""Isolated, killable community-gallery publisher.

The mounted Hugging Face Bucket is S3-paced storage. Running its copy in a
child process lets the web worker enforce a publication deadline without
leaving an uncancellable Python thread that could commit after the caller has
already reported a timeout.
"""

from __future__ import annotations

import json
import os
import signal
import sys
from pathlib import Path
from typing import Any

from .gallery import GalleryError, GalleryStore


class _WorkerStopped(BaseException):
    pass


def _stop_worker(_signum, _frame) -> None:
    # Raising through GalleryStore.publish executes its staging cleanup before
    # the worker exits. The parent escalates to SIGKILL if a blocked FUSE call
    # does not return promptly enough to deliver this signal.
    raise _WorkerStopped


def _write_response(path: Path, payload: dict[str, Any]) -> None:
    temporary = path.with_name(f".{path.name}.tmp")
    with temporary.open("x", encoding="utf-8") as stream:
        json.dump(payload, stream, ensure_ascii=False, separators=(",", ":"))
        stream.write("\n")
        stream.flush()
        try:
            os.fsync(stream.fileno())
        except OSError:
            pass
    temporary.replace(path)


def main(argv: list[str] | None = None) -> int:
    arguments = list(sys.argv[1:] if argv is None else argv)
    if len(arguments) != 2:
        print("usage: python -m app.gallery_worker REQUEST_JSON RESPONSE_JSON", file=sys.stderr)
        return 2

    request_path = Path(arguments[0])
    response_path = Path(arguments[1])
    signal.signal(signal.SIGTERM, _stop_worker)
    try:
        request = json.loads(request_path.read_text(encoding="utf-8"))
        store = GalleryStore(Path(request["galleryDir"]))
        item = store.publish(
            item_id=request["itemId"],
            job_dir=Path(request["jobDir"]),
            result=request["result"],
            created_at=request.get("createdAt"),
            elapsed_started_at=request.get("elapsedStartedAt"),
            staging_token=request["stagingToken"],
        )
    except _WorkerStopped:
        return 143
    except GalleryError as exc:
        _write_response(response_path, {"ok": False, "error": str(exc)})
        return 1
    except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as exc:
        _write_response(
            response_path,
            {
                "ok": False,
                "error": (
                    "The isolated gallery publisher could not complete "
                    f"({type(exc).__name__})."
                ),
            },
        )
        return 1

    _write_response(response_path, {"ok": True, "item": item})
    return 0


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