"""有界批量队列、落盘状态及结果保留；由 HTTP lifespan 管理。""" from __future__ import annotations from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from datetime import datetime, timedelta, timezone import json import logging from pathlib import Path import re import shutil import threading from typing import BinaryIO from uuid import uuid4 import zipfile import config from . import service from .service import AlgorithmError MAX_BATCH_FILES = 20 MAX_BATCH_BYTES = 100 * 1024 * 1024 TERMINAL = {"completed", "partial_failed", "failed", "interrupted"} ITEM_TERMINAL = {"completed", "failed", "interrupted"} logger = logging.getLogger(__name__) def _now() -> str: return datetime.now(timezone.utc).isoformat() class BatchManager: def __init__(self, root: Path, *, parallelism: int | None = None, max_pending: int | None = None, ttl_hours: int | None = None): self.root = root.resolve() self.root.mkdir(parents=True, exist_ok=True) self.parallelism = parallelism or config.ALGORITHM_MAX_PARALLEL self.max_pending = max_pending or config.ALGORITHM_MAX_PENDING_IMAGES self.ttl_hours = ttl_hours or config.ALGORITHM_RESULT_TTL_HOURS self._lock = threading.RLock() self._jobs: dict[str, dict] = {} self._pending = 0 self._closed = False self._stop = threading.Event() self._restore() self.cleanup() self._pool = ThreadPoolExecutor(max_workers=self.parallelism, thread_name_prefix="curve-batch") self._cleaner = threading.Thread(target=self._maintain, name="curve-result-cleanup", daemon=True) self._cleaner.start() def _save(self, job: dict) -> None: path = self.root / job["id"] / "job.json" temporary = path.with_suffix(".tmp") temporary.write_text(json.dumps(job, ensure_ascii=False, allow_nan=False), encoding="utf-8") temporary.replace(path) def _restore(self) -> None: for directory in self.root.iterdir(): if directory.is_symlink() or not directory.is_dir() or not re.fullmatch(r"[0-9a-f]{32}", directory.name): continue path = directory / "job.json" if not path.exists(): # 创建阶段崩溃留下的目录，没有已提交任务。 shutil.rmtree(directory) continue job = json.loads(path.read_text(encoding="utf-8")) if job["id"] != directory.name: raise ValueError(f"算法任务目录与记录不符：{directory.name}") self._jobs[job["id"]] = job if job["status"] not in TERMINAL: for item in job["items"]: if item["status"] not in ITEM_TERMINAL: self._interrupt(item) shutil.rmtree(directory / "items" / item["id"], ignore_errors=True) shutil.rmtree(directory / "inputs", ignore_errors=True) self._finalize(job["id"], interrupted=True) @staticmethod def _interrupt(item: dict) -> None: item.update(status="interrupted", stage="interrupted", finished_at=_now(), error={"code": "service_interrupted", "message": "服务已停止，请重新提交此图片"}) def cleanup(self) -> None: with self._lock: now = _now() for job_id, job in list(self._jobs.items()): if job["status"] in TERMINAL and job.get("expires_at", "9999") <= now: try: shutil.rmtree(self.root / job_id) except FileNotFoundError: pass except OSError: logger.exception("清理算法任务失败：%s", job_id) continue del self._jobs[job_id] def _maintain(self) -> None: while not self._stop.wait(60): self.cleanup() def submit(self, files: list[tuple[str, bytes]], *, detailed: bool = False) -> dict: service.check_dependencies() if not 1 <= len(files) <= MAX_BATCH_FILES: raise AlgorithmError("invalid_batch_size", "每批需上传 1～20 张图片", 422) if any(len(content) > service.MAX_UPLOAD_BYTES for _, content in files): raise AlgorithmError("image_too_large", "单张图片不能超过 20 MiB", 413) if sum(len(content) for _, content in files) > MAX_BATCH_BYTES: raise AlgorithmError("batch_too_large", "每批图片合计不能超过 100 MiB", 413) with self._lock: self.cleanup() if self._closed: raise AlgorithmError("algorithm_unavailable", "算法队列正在关闭", 503) if self._pending + len(files) > self.max_pending: raise AlgorithmError("algorithm_queue_full", "算法队列已满，请稍后重试", 429) job_id = uuid4().hex directory = self.root / job_id inputs = directory / "inputs" inputs.mkdir(parents=True) job = {"id": job_id, "algorithm": "characteristic_curve_recognition", "status": "queued", "created_at": _now(), "finished_at": None, "expires_at": None, "detailed": detailed, "items": []} try: for index, (filename, content) in enumerate(files, 1): item_id = f"{index:04d}" try: suffix = service.validate_image(content) except AlgorithmError: suffix = ".bin" # 格式错误作为单张失败，不拒绝其他图片。 (inputs / (item_id + suffix)).write_bytes(content) job["items"].append({"id": item_id, "filename": filename.replace("\\", "/").rsplit("/", 1)[-1][:255], "status": "queued", "stage": "queued", "started_at": None, "finished_at": None, "error": None, "suffix": suffix}) self._save(job) except BaseException: shutil.rmtree(directory, ignore_errors=True) raise self._jobs[job_id] = job self._pending += len(files) for item in job["items"]: self._pool.submit(self._run, job_id, item["id"]) return self._snapshot(job) def _run(self, job_id: str, item_id: str) -> None: with self._lock: if self._closed: return # close() 等活动计算结束后统一将未开始项标记中断。 job = self._jobs[job_id] item = next(item for item in job["items"] if item["id"] == item_id) job["status"] = "running" item.update(status="running", stage="waiting_for_slot", started_at=_now()) self._save(job) directory = self.root / job_id image_path = directory / "inputs" / (item_id + item["suffix"]) result_dir = directory / "items" / item_id stage_path = directory / (item_id + ".stage.json") error = None try: service.validate_image(image_path.read_bytes()) service.execute_image(image_path, result_dir, detailed=job["detailed"], wait_for_slot=True, stage_path=stage_path) except AlgorithmError as exc: error = {"code": exc.code, "message": str(exc)} except Exception: logger.exception("批量图片处理失败：%s/%s", job_id, item_id) error = {"code": "algorithm_failed", "message": "算法处理失败，请通过服务日志排查"} finally: image_path.unlink(missing_ok=True) if error: shutil.rmtree(result_dir, ignore_errors=True) with self._lock: status = "failed" if error else "completed" item.update(status=status, stage=status, finished_at=_now(), error=error) self._pending -= 1 finalize = all(value["status"] in ITEM_TERMINAL for value in job["items"]) if finalize: job["status"] = "packaging" self._save(job) if finalize: self._finalize(job_id) def _finalize(self, job_id: str, *, interrupted: bool = False) -> None: with self._lock: job = deepcopy(self._jobs[job_id]) successes = sum(item["status"] == "completed" for item in job["items"]) job["status"] = ("interrupted" if interrupted else "completed" if successes == len(job["items"]) else "partial_failed" if successes else "failed") job["finished_at"] = _now() job["expires_at"] = (datetime.now(timezone.utc) + timedelta(hours=self.ttl_hours)).isoformat() directory = self.root / job_id temporary = directory / "results.zip.tmp" try: with zipfile.ZipFile(temporary, "w", zipfile.ZIP_DEFLATED) as archive: archive.writestr("summary.json", json.dumps(self._snapshot(job), ensure_ascii=False)) for item in job["items"]: if item["status"] == "completed": for path in sorted((directory / "items" / item["id"]).rglob("*")): if path.is_file(): archive.write(path, path.relative_to(directory).as_posix()) temporary.replace(directory / "results.zip") except Exception: logger.exception("批量结果打包失败：%s", job_id) temporary.unlink(missing_ok=True) job.update(status="failed", error={"code": "archive_failed", "message": "结果打包失败，可按图片读取成功结果"}) with self._lock: self._jobs[job_id] = job self._save(job) def _snapshot(self, job: dict) -> dict: result = deepcopy(job) counts = {state: sum(item["status"] == state for item in job["items"]) for state in ("queued", "running", "completed", "failed", "interrupted")} result["progress"] = {"total": len(job["items"]), **counts, "finished": counts["completed"] + counts["failed"] + counts["interrupted"]} for item in result["items"]: item.pop("suffix", None) if item["status"] == "running": try: item["stage"] = json.loads((self.root / job["id"] / (item["id"] + ".stage.json")).read_text())["stage"] except (FileNotFoundError, ValueError, KeyError): pass return result def _get(self, job_id: str) -> dict: self.cleanup() job = self._jobs.get(job_id) if job is None: raise AlgorithmError("batch_not_found", "批量任务不存在或已过期", 404) return job def status(self, job_id: str) -> dict: with self._lock: return self._snapshot(self._get(job_id)) def result(self, job_id: str, item_id: str) -> dict: with self._lock: job = self._get(job_id) item = next((item for item in job["items"] if item["id"] == item_id), None) if item is None: raise AlgorithmError("item_not_found", "图片任务不存在", 404) if item["status"] != "completed": raise AlgorithmError("result_unavailable", "该图片尚未成功完成，请查询任务状态", 409) return json.loads((self.root / job_id / "items" / item_id / "result.json").read_text()) def open_archive(self, job_id: str) -> BinaryIO: with self._lock: job = self._get(job_id) if job["status"] not in TERMINAL: raise AlgorithmError("batch_not_finished", "批量任务尚未结束", 409) path = self.root / job_id / "results.zip" if not path.is_file(): raise AlgorithmError("archive_unavailable", "结果包不可用，请按图片读取结果", 409) # 在清理锁内打开；POSIX 上过期删除也不影响已开始的下载。 return path.open("rb") def close(self) -> None: with self._lock: self._closed = True self._stop.set() self._cleaner.join() self._pool.shutdown(wait=True, cancel_futures=True) for job_id, job in list(self._jobs.items()): if job["status"] not in TERMINAL: for item in job["items"]: if item["status"] not in ITEM_TERMINAL: self._interrupt(item) shutil.rmtree(self.root / job_id / "inputs", ignore_errors=True) self._finalize(job_id, interrupted=True)