Spaces:
Running
Running
File size: 2,788 Bytes
bf1fb5f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """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())
|