new-pipeline/main.py

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"
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",
                    "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()