diff --git a/server.py b/server.py index 4b9e78a..1504ec2 100644 --- a/server.py +++ b/server.py @@ -5,26 +5,21 @@ POST /patch — upload image, returns {job_id} GET /jobs/{job_id} — returns {status, result_url?} GET /jobs/{job_id}/result — returns the PNG directly -Jobs are stored in Redis. Results are stored on disk (mount a PVC in k8s). +Jobs are queued via Redis (rq). Results are stored in Redis as blobs. """ import os import uuid -import shutil -import tempfile -from pathlib import Path from fastapi import FastAPI, UploadFile, HTTPException -from fastapi.responses import FileResponse +from fastapi.responses import Response import redis import rq -RESULTS_DIR = Path(os.environ.get("RESULTS_DIR", "/tmp/p2p_results")) -RESULTS_DIR.mkdir(parents=True, exist_ok=True) - REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") +RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) -app = FastAPI(title="pic-to-patch", version="0.1.0") +app = FastAPI(title="pic-to-patch", version="0.2.0") redis_conn = redis.from_url(REDIS_URL) queue = rq.Queue("patches", connection=redis_conn) @@ -37,30 +32,21 @@ async def create_patch( postprocess: bool = True, ): job_id = str(uuid.uuid4()) - job_dir = RESULTS_DIR / job_id - job_dir.mkdir() - - ext = Path(file.filename or "input.png").suffix or ".png" - input_path = job_dir / f"input{ext}" - with open(input_path, "wb") as f: - shutil.copyfileobj(file.file, f) - - output_path = job_dir / "patch.png" + input_bytes = await file.read() + ext = (file.filename or "input.png").rsplit(".", 1)[-1] or "png" - is_svg = ext.lower() == ".svg" + redis_conn.setex(f"job:{job_id}:input", RESULT_TTL, input_bytes) + redis_conn.setex(f"job:{job_id}:ext", RESULT_TTL, ext) + redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "processing") queue.enqueue( "worker.run_pipeline", job_id=job_id, - input_path=str(input_path), - output_path=str(output_path), border_color=border_color, color_precision=color_precision, postprocess=postprocess, - is_svg=is_svg, job_timeout=300, - result_ttl=3600, - meta={"job_id": job_id}, + result_ttl=RESULT_TTL, ) return {"job_id": job_id} @@ -68,27 +54,29 @@ async def create_patch( @app.get("/jobs/{job_id}") async def get_job(job_id: str): - job_dir = RESULTS_DIR / job_id - if not job_dir.exists(): + status = redis_conn.get(f"job:{job_id}:status") + if status is None: raise HTTPException(404, "job not found") - output = job_dir / "patch.png" - error_file = job_dir / "error.txt" + status = status.decode() + resp = {"status": status} - if error_file.exists(): - return {"status": "failed", "error": error_file.read_text()} - if output.exists() and output.stat().st_size > 0: - return {"status": "complete", "result_url": f"/jobs/{job_id}/result"} + if status == "complete": + resp["result_url"] = f"/jobs/{job_id}/result" + elif status == "failed": + error = redis_conn.get(f"job:{job_id}:error") + resp["error"] = error.decode() if error else "unknown error" - return {"status": "processing"} + return resp @app.get("/jobs/{job_id}/result") async def get_result(job_id: str): - output = RESULTS_DIR / job_id / "patch.png" - if not output.exists(): + data = redis_conn.get(f"job:{job_id}:result") + if data is None: raise HTTPException(404, "result not ready") - return FileResponse(output, media_type="image/png", filename=f"{job_id}.png") + return Response(content=data, media_type="image/png", + headers={"Content-Disposition": f'inline; filename="{job_id}.png"'}) @app.get("/health") diff --git a/worker.py b/worker.py index 11e4eb4..f397151 100644 --- a/worker.py +++ b/worker.py @@ -4,21 +4,50 @@ pic-to-patch worker. Picks jobs from Redis queue and runs the pipeline. Run with: rq worker patches --url redis://localhost:6379 """ +import os +import tempfile from pathlib import Path + +import redis + from pipeline.convert import convert, convert_svg +REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") +RESULT_TTL = int(os.environ.get("RESULT_TTL", "3600")) +redis_conn = redis.from_url(REDIS_URL) + -def run_pipeline(input_path, output_path, border_color="#0a0a14", - color_precision=8, postprocess=True, is_svg=False, job_id=None): - error_file = Path(output_path).parent / "error.txt" +def run_pipeline(job_id, border_color="#0a0a14", + color_precision=8, postprocess=True): try: - if is_svg: - convert_svg(input_path, output_path, - border_color=border_color, postprocess=postprocess) - else: - convert(input_path, output_path, - border_color=border_color, color_precision=color_precision, - postprocess=postprocess) + input_bytes = redis_conn.get(f"job:{job_id}:input") + ext = (redis_conn.get(f"job:{job_id}:ext") or b"png").decode() + if not input_bytes: + raise RuntimeError("input not found in redis") + + is_svg = ext.lower() == "svg" + + with tempfile.TemporaryDirectory(prefix="p2p_") as tmpdir: + tmpdir = Path(tmpdir) + input_path = tmpdir / f"input.{ext}" + input_path.write_bytes(input_bytes) + output_path = tmpdir / "patch.png" + + if is_svg: + convert_svg(str(input_path), str(output_path), + border_color=border_color, postprocess=postprocess) + else: + convert(str(input_path), str(output_path), + border_color=border_color, color_precision=color_precision, + postprocess=postprocess) + + result_bytes = output_path.read_bytes() + + redis_conn.setex(f"job:{job_id}:result", RESULT_TTL, result_bytes) + redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "complete") + redis_conn.delete(f"job:{job_id}:input") + except Exception as e: - error_file.write_text(str(e)) + redis_conn.setex(f"job:{job_id}:status", RESULT_TTL, "failed") + redis_conn.setex(f"job:{job_id}:error", RESULT_TTL, str(e)) raise