vent_agent_2.0/algorithms/characteristic_curve_recognition/batch.py

"""有界批量队列、落盘状态及结果保留；由 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)