import json import os import signal import threading import traceback from datetime import datetime, timezone from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any, Dict import modelhub_pipeline_resumable as pipeline HOST = "0.0.0.0" VERSION = "1.1.0" PORT = int(os.getenv("PORT", "8080")) RUN_INTERVAL_SECONDS = int(os.getenv("PIPELINE_RUN_INTERVAL_SECONDS", "3600")) RETRY_INTERVAL_SECONDS = int(os.getenv("PIPELINE_RETRY_INTERVAL_SECONDS", "60")) SHUTDOWN_TIMEOUT_SECONDS = 25 STRATEGY_ID = os.environ["STRATEGY_ID"] stop_event = threading.Event() def utc_now() -> str: return datetime.now(timezone.utc).isoformat() class RuntimeState: def __init__(self) -> None: self._lock = threading.Lock() self._data: Dict[str, Any] = { "pipeline_status": "starting", "run_count": 0, "last_started_at": None, "last_finished_at": None, "last_error": None, } def update(self, **values: Any) -> None: with self._lock: self._data.update(values) def snapshot(self) -> Dict[str, Any]: with self._lock: return dict(self._data) runtime_state = RuntimeState() class Handler(BaseHTTPRequestHandler): def do_GET(self) -> None: path = self.path.split("?", 1)[0] state = runtime_state.snapshot() if path == "/health": self._write_json( 200, { "status": "ok", "pipeline_status": state["pipeline_status"], }, ) return if path == "/": self._write_json( 200, { "name": "new-pipeline", "version": VERSION, "status": "running", "strategy_id": STRATEGY_ID, "pipeline_status": state["pipeline_status"], "run_count": state["run_count"], "last_started_at": state["last_started_at"], "last_finished_at": state["last_finished_at"], }, ) return self._write_json(404, {"error": "not found"}) def _write_json(self, status_code: int, payload: Dict[str, Any]) -> None: body = json.dumps(payload, ensure_ascii=False).encode("utf-8") self.send_response(status_code) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def log_message(self, message_format: str, *args: Any) -> None: print(f"[health] {self.address_string()} - {message_format % args}") def pipeline_worker() -> None: pipeline.reset_shutdown() while not stop_event.is_set(): previous = runtime_state.snapshot() runtime_state.update( pipeline_status="running", run_count=previous["run_count"] + 1, last_started_at=utc_now(), last_error=None, ) try: pipeline.main() except pipeline.ShutdownRequested: runtime_state.update(pipeline_status="stopping") break except Exception as exc: runtime_state.update( pipeline_status="error", last_finished_at=utc_now(), last_error=f"{type(exc).__name__}: {exc}", ) traceback.print_exc() wait_seconds = RETRY_INTERVAL_SECONDS else: runtime_state.update( pipeline_status="idle", last_finished_at=utc_now(), ) wait_seconds = RUN_INTERVAL_SECONDS if stop_event.wait(max(1, wait_seconds)): break runtime_state.update(pipeline_status="stopped") def handle_signal(signum: int, _frame: Any) -> None: print(f"收到信号 {signum},开始优雅停机") runtime_state.update(pipeline_status="stopping") stop_event.set() pipeline.request_shutdown() def main() -> None: signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal) worker = threading.Thread( target=pipeline_worker, name="modelhub-pipeline", daemon=True, ) worker.start() server = ThreadingHTTPServer((HOST, PORT), Handler) server.timeout = 0.5 print(f"健康检查服务已启动: http://{HOST}:{PORT}/health") try: while not stop_event.is_set(): server.handle_request() finally: stop_event.set() pipeline.request_shutdown() server.server_close() worker.join(timeout=SHUTDOWN_TIMEOUT_SECONDS) if worker.is_alive(): print("后台任务未在停机窗口内结束,主进程将退出") else: print("策略已完成资源清理并停止") if __name__ == "__main__": main()