""" QGIS Worker 进程池管理器 — 主机侧运行。 管理 N 个长驻 Docker Worker 进程，提供任务提交和结果回收。 架构: 主机 FastAPI ←→ qgis_pool (本模块) ←→ docker exec -i ←→ qgis_worker.py (容器内) 用法: from app.services.qgis.qgis_pool import qgis_pool # 提交任务（阻塞等待完成） results, summary = qgis_pool.submit_job(config, models, progress_callback) # 关闭池 qgis_pool.shutdown() """ import json import subprocess import threading import time from typing import Callable, Dict, List, Optional, Tuple from app.config.paths import get_logger logger = get_logger("qgis.pool") class _WorkerHandle: """单个 Worker 进程的句柄""" def __init__(self, worker_id: int, proc: subprocess.Popen): self.worker_id = worker_id self.proc = proc self.busy = False self.lock = threading.Lock() self._reader_thread: Optional[threading.Thread] = None def is_alive(self) -> bool: return self.proc.poll() is None def mark_busy(self): with self.lock: self.busy = True def mark_free(self): with self.lock: self.busy = False def is_free(self) -> bool: with self.lock: return not self.busy and self.is_alive() class QgisPool: """QGIS Worker 进程池""" def __init__(self, pool_size: int = None): """ Args: pool_size: 池中 Worker 数量，默认从配置读取 """ if pool_size is None: try: from config import settings pool_size = getattr(settings, "QGIS_POOL_SIZE", 2) except Exception: pool_size = 2 self._pool_size = pool_size self._workers: List[_WorkerHandle] = [] self._lock = threading.Lock() self._started = False def _ensure_started(self): """懒启动：首次使用时启动 Worker 进程""" if self._started: return with self._lock: if self._started: return self._start_workers() self._started = True def _start_workers(self): """启动所有 Worker 进程""" from app.services.qgis.qgis_env import ( get_docker_container, get_docker_project_dir, get_container_python_path, ) container = get_docker_container() project_dir = get_docker_project_dir() python_path = get_container_python_path() worker_script = f"{project_dir}/app/services/qgis/qgis_worker.py" # Qt 平台设置 try: from config import settings qt_platform = getattr(settings, "QGIS_DOCKER_QT_PLATFORM", "offscreen") except Exception: qt_platform = "offscreen" for i in range(self._pool_size): cmd = [ "docker", "exec", "-i", "-w", project_dir, "-e", f"QT_QPA_PLATFORM={qt_platform}", container, python_path, worker_script, ] logger.info(f"[Pool] 启动 Worker-{i}: {' '.join(cmd)}") proc = subprocess.Popen( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="replace", ) worker = _WorkerHandle(i, proc) self._workers.append(worker) # 后台线程消费 stderr，防止管道满阻塞 t = threading.Thread( target=self._drain_stderr, args=(worker,), daemon=True ) t.start() logger.info(f"[Pool] 已启动 {self._pool_size} 个 Worker") @staticmethod def _drain_stderr(worker: _WorkerHandle): """持续读取 Worker stderr 防止管道缓冲区满""" try: for line in worker.proc.stderr: line = line.strip() if line: logger.debug(f"[Worker-{worker.worker_id}] stderr: {line}") except Exception: pass def _find_free_worker(self, timeout: float = 30) -> Optional[_WorkerHandle]: """找到一个空闲 Worker，超时返回 None""" deadline = time.time() + timeout while time.time() < deadline: with self._lock: for w in self._workers: if w.is_free(): return w time.sleep(0.1) return None def submit_job( self, config: dict, models: list, progress_callback: Callable[[dict], None] = None, ) -> Tuple[list, dict]: """ 提交一个批次任务，阻塞等待完成。 Args: config: QGIS 配置字典 models: 模型参数列表 progress_callback: 每张图完成时的回调（接收 PROGRESS dict） Returns: (results_list, summary_dict) """ self._ensure_started() worker = self._find_free_worker(timeout=60) if worker is None: raise RuntimeError("[Pool] 无可用 Worker（等待 60s 超时）") worker.mark_busy() try: return self._run_job(worker, config, models, progress_callback) finally: # 检查 Worker 是否还活着，死了就重启 if not worker.is_alive(): logger.error(f"[Pool] Worker-{worker.worker_id} 已死亡，重启...") self._restart_worker(worker) worker.mark_free() def _run_job( self, worker: _WorkerHandle, config: dict, models: list, progress_callback: Callable[[dict], None] = None, ) -> Tuple[list, dict]: """向指定 Worker 发送任务并读取结果""" job = {"config": config, "models": models} job_json = json.dumps(job, ensure_ascii=False) + "\n" # 写入 stdin try: worker.proc.stdin.write(job_json) worker.proc.stdin.flush() except (BrokenPipeError, OSError) as e: raise RuntimeError(f"[Pool] 写入 Worker-{worker.worker_id} stdin 失败: {e}") # 读取输出（直到 JOB_DONE 或 FATAL） results = [] summary = {} while True: line = worker.proc.stdout.readline() if not line: # stdout 关闭，Worker 可能已死亡 logger.error(f"[Pool] Worker-{worker.worker_id} stdout 关闭") summary = {"error": "Worker stdout 关闭"} break line = line.strip() if not line: continue if line.startswith("PROGRESS:"): try: r = json.loads(line[len("PROGRESS:"):]) results.append(r) if progress_callback: progress_callback(r) except json.JSONDecodeError: logger.warning(f"[Pool] Worker-{worker.worker_id} PROGRESS 解析失败: {line[:200]}") elif line.startswith("JOB_DONE:"): try: summary = json.loads(line[len("JOB_DONE:"):]) except json.JSONDecodeError: summary = {"error": f"JOB_DONE 解析失败: {line[:200]}"} break elif line.startswith("FATAL:"): try: summary = json.loads(line[len("FATAL:"):]) except json.JSONDecodeError: summary = {"error": f"FATAL: {line[:200]}"} logger.error(f"[Pool] Worker-{worker.worker_id} FATAL: {summary}") break else: logger.debug(f"[Pool] Worker-{worker.worker_id} 未知输出: {line[:200]}") return results, summary def _restart_worker(self, dead_worker: _WorkerHandle): """重启一个死亡的 Worker""" from app.services.qgis.qgis_env import ( get_docker_container, get_docker_project_dir, get_container_python_path, ) container = get_docker_container() project_dir = get_docker_project_dir() python_path = get_container_python_path() worker_script = f"{project_dir}/app/services/qgis/qgis_worker.py" try: from config import settings qt_platform = getattr(settings, "QGIS_DOCKER_QT_PLATFORM", "offscreen") except Exception: qt_platform = "offscreen" cmd = [ "docker", "exec", "-i", "-w", project_dir, "-e", f"QT_QPA_PLATFORM={qt_platform}", container, python_path, worker_script, ] proc = subprocess.Popen( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="replace", ) dead_worker.proc = proc dead_worker.busy = False # 重启 stderr 消费线程 t = threading.Thread( target=self._drain_stderr, args=(dead_worker,), daemon=True ) t.start() logger.info(f"[Pool] Worker-{dead_worker.worker_id} 已重启") def shutdown(self): """关闭所有 Worker""" with self._lock: for w in self._workers: if w.proc.poll() is None: try: w.proc.stdin.close() except Exception: pass try: w.proc.terminate() w.proc.wait(timeout=5) except Exception: try: w.proc.kill() except Exception: pass self._workers.clear() self._started = False logger.info("[Pool] 已关闭") # 全局单例 qgis_pool = QgisPool()