wzy-warehouse
325 lines
10 KiB
Python
325 lines
10 KiB
Python
"""
|
||
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()
|