341 lines
14 KiB
Python
341 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
from collections.abc import Callable
|
|
from datetime import datetime, time, timedelta
|
|
from typing import Any
|
|
|
|
from datahub.db import HubDB
|
|
from datahub.logutil import get_logger
|
|
from datahub.pipeline import Pipeline
|
|
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd
|
|
|
|
LOGGER = get_logger()
|
|
|
|
JobFn = Callable[[str], Any]
|
|
|
|
EOD_JOB_IDS = {"eod_a", "eod_b", "eod_retry"}
|
|
|
|
|
|
def is_open_day(db: HubDB, day: str) -> bool:
|
|
row = db.fetchone(
|
|
"SELECT is_open FROM trade_calendar WHERE exchange = 'SSE' AND cal_date = ?",
|
|
(day,),
|
|
)
|
|
if row is None:
|
|
return True # unknown calendar: do not skip reference refresh
|
|
return int(row["is_open"]) == 1
|
|
|
|
|
|
def _hhmm(value: str) -> time:
|
|
return datetime.strptime(value, "%H:%M").time()
|
|
|
|
|
|
class Scheduler:
|
|
"""Calendar-driven in-process scheduler. Non-trading days skip EOD fetches.
|
|
|
|
EOD datasets that failed to publish (e.g. upstream not ready at 15:05)
|
|
are retried automatically every ``eod_retry_interval_minutes`` between
|
|
``eod_retry_start`` and ``eod_retry_cutoff``. Progress is persisted in
|
|
``eod_progress`` so a container restart catches up instead of waiting
|
|
for the next day, and completed days are never re-fetched.
|
|
"""
|
|
|
|
def __init__(self, db: HubDB, pipeline: Pipeline, jobs: dict[str, JobFn] | None = None) -> None:
|
|
self.db = db
|
|
self.pipeline = pipeline
|
|
self.jobs = jobs or {
|
|
"precheck": self._precheck,
|
|
"eod_a": self._eod_a,
|
|
"eod_b": self._eod_b,
|
|
"eod_retry": self._eod_retry,
|
|
"cleanup": self._cleanup,
|
|
"backup": self._backup,
|
|
"history_backfill": self._history_backfill,
|
|
}
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
self._fired: set[tuple[str, str, str]] = set()
|
|
self._eod_lock = threading.Lock()
|
|
|
|
def start(self, interval_seconds: float = 30.0) -> None:
|
|
if self._thread and self._thread.is_alive():
|
|
return
|
|
|
|
def loop() -> None:
|
|
while not self._stop.wait(interval_seconds):
|
|
try:
|
|
self.tick()
|
|
except Exception:
|
|
LOGGER.exception("scheduler tick failed")
|
|
|
|
self._thread = threading.Thread(target=loop, name="datahub-scheduler", daemon=True)
|
|
self._thread.start()
|
|
|
|
def stop(self, timeout: float = 5.0) -> None:
|
|
self._stop.set()
|
|
if self._thread and self._thread is not threading.current_thread():
|
|
self._thread.join(timeout)
|
|
|
|
def tick(self, clock: datetime | None = None) -> list[str]:
|
|
now = now_shanghai(clock)
|
|
day = yyyymmdd(now)
|
|
current = now.time()
|
|
ran: list[str] = []
|
|
plan = [
|
|
("precheck", time(8, 45)),
|
|
("eod_a", time(15, 5)),
|
|
("eod_b", time(15, 10)),
|
|
("cleanup", time(0, 30)),
|
|
("backup", time(0, 40)),
|
|
]
|
|
open_day = is_open_day(self.db, day)
|
|
for job_id, at in plan:
|
|
if current < at:
|
|
continue
|
|
key = (job_id, day, at.strftime("%H%M"))
|
|
if key in self._fired:
|
|
continue
|
|
if job_id in {"eod_a", "eod_b"} and not open_day:
|
|
self._fired.add(key)
|
|
continue
|
|
self._fired.add(key)
|
|
if job_id in {"eod_a", "eod_b"}:
|
|
# Record the attempt before running: even a crash must not
|
|
# hide that today's first EOD try already happened.
|
|
self._record_eod_attempt(day, now)
|
|
try:
|
|
self.run_job(job_id, day)
|
|
except Exception:
|
|
if job_id not in {"eod_a", "eod_b"}:
|
|
raise
|
|
# Keep the tick alive; evening retries take over.
|
|
LOGGER.exception("scheduled job %s failed for %s", job_id, day)
|
|
ran.append(job_id)
|
|
if job_id in {"eod_a", "eod_b"}:
|
|
self._settle_eod(day)
|
|
ran.extend(self._eod_retry_tick(now, day, open_day))
|
|
return ran
|
|
|
|
# ------------------------------------------------------------------
|
|
# EOD retry window
|
|
# ------------------------------------------------------------------
|
|
|
|
def _eod_retry_tick(self, now: datetime, day: str, open_day: bool) -> list[str]:
|
|
if not open_day:
|
|
return []
|
|
settings = self.pipeline.settings
|
|
current = now.time()
|
|
start = _hhmm(settings.eod_retry_start)
|
|
cutoff = _hhmm(settings.eod_retry_cutoff)
|
|
interval = timedelta(minutes=settings.eod_retry_interval_minutes)
|
|
missing = self.pipeline.missing_official_datasets(day)
|
|
row = self.eod_progress(day)
|
|
|
|
if not missing:
|
|
if row is None or row["state"] != "done":
|
|
self._save_eod_progress(day, state="done", finished_at=isoformat(now))
|
|
return []
|
|
if current < start:
|
|
return []
|
|
if row and row["state"] == "cutoff_failed":
|
|
return []
|
|
if current >= cutoff:
|
|
detail = "截止时间已到，缺失数据集: " + ",".join(missing)
|
|
self._save_eod_progress(day, state="cutoff_failed", finished_at=isoformat(now), detail=detail)
|
|
with self.db.write() as connection:
|
|
connection.execute(
|
|
"INSERT INTO job_runs(job_id, state, started_at, finished_at, error, attempt, detail)"
|
|
" VALUES ('eod_retry','failed',?,?,?,?,?)",
|
|
(isoformat(now), isoformat(now), detail, int((row or {}).get("attempts") or 0), "eod cutoff reached"),
|
|
)
|
|
LOGGER.warning(
|
|
"eod retry window closed without data",
|
|
extra={"hub": {"trade_date": day, "missing": missing, "reason": "eod_cutoff"}},
|
|
)
|
|
return []
|
|
last = None
|
|
if row and row["last_attempt_at"]:
|
|
try:
|
|
last = datetime.fromisoformat(str(row["last_attempt_at"]))
|
|
except ValueError:
|
|
last = None
|
|
if last is not None and now_shanghai(last).replace(tzinfo=None) + interval > now.replace(tzinfo=None):
|
|
return []
|
|
if "eod_retry" not in self.jobs:
|
|
return []
|
|
self._record_eod_attempt(day, now)
|
|
ran = []
|
|
try:
|
|
self.run_job("eod_retry", day)
|
|
except Exception:
|
|
# job_runs already carries the failure; the window keeps retrying.
|
|
LOGGER.warning("eod retry failed for %s", day, exc_info=True)
|
|
ran.append("eod_retry")
|
|
self._settle_eod(day)
|
|
return ran
|
|
|
|
def _settle_eod(self, day: str) -> None:
|
|
"""Flip the day to done as soon as every official dataset is published."""
|
|
if not self.pipeline.missing_official_datasets(day):
|
|
row = self.eod_progress(day)
|
|
if row is None or row["state"] != "done":
|
|
self._save_eod_progress(day, state="done", finished_at=isoformat())
|
|
|
|
def eod_progress(self, day: str) -> dict[str, Any] | None:
|
|
return self.db.fetchone("SELECT * FROM eod_progress WHERE trade_date = ?", (day,))
|
|
|
|
def eod_status(self, trade_date: str | None = None, clock: datetime | None = None) -> dict[str, Any]:
|
|
"""Human/admin facing view: 等待上游 / 下次重试 / 已成功 / 已截止失败."""
|
|
day = yyyymmdd(trade_date or now_shanghai(clock))
|
|
now = now_shanghai(clock)
|
|
row = self.eod_progress(day)
|
|
open_day = is_open_day(self.db, day)
|
|
missing = self.pipeline.missing_official_datasets(day)
|
|
if row and row["state"] == "done":
|
|
state = "done"
|
|
elif not open_day:
|
|
state = "closed_day"
|
|
elif not missing:
|
|
state = "done"
|
|
elif row and row["state"] == "cutoff_failed":
|
|
state = "cutoff_failed"
|
|
elif now.time() < _hhmm("15:05"):
|
|
state = "pending_first_attempt"
|
|
else:
|
|
state = "waiting_upstream"
|
|
return {
|
|
"trade_date": day,
|
|
"is_open_day": open_day,
|
|
"state": state,
|
|
"missing_datasets": missing,
|
|
"attempts": int((row or {}).get("attempts") or 0),
|
|
"last_attempt_at": (row or {}).get("last_attempt_at"),
|
|
"next_retry_at": (row or {}).get("next_retry_at") if state == "waiting_upstream" else None,
|
|
"finished_at": (row or {}).get("finished_at"),
|
|
"detail": (row or {}).get("detail"),
|
|
}
|
|
|
|
def _record_eod_attempt(self, day: str, now: datetime) -> None:
|
|
row = self.eod_progress(day)
|
|
attempts = int((row or {}).get("attempts") or 0) + 1
|
|
interval = self.pipeline.settings.eod_retry_interval_minutes
|
|
self._save_eod_progress(
|
|
day,
|
|
state="waiting_upstream",
|
|
attempts=attempts,
|
|
last_attempt_at=isoformat(now),
|
|
next_retry_at=isoformat(now + timedelta(minutes=interval)),
|
|
)
|
|
|
|
def _save_eod_progress(self, day: str, **fields: Any) -> None:
|
|
columns = [
|
|
"trade_date", "state", "attempts", "last_attempt_at",
|
|
"next_retry_at", "finished_at", "detail", "updated_at",
|
|
]
|
|
with self.db.write() as connection:
|
|
existing = connection.execute(
|
|
"SELECT trade_date FROM eod_progress WHERE trade_date = ?",
|
|
(day,),
|
|
).fetchone()
|
|
if existing is None:
|
|
payload = {name: None for name in columns}
|
|
payload.update({"trade_date": day, "state": "waiting_upstream", "attempts": 0})
|
|
payload.update(fields)
|
|
payload["updated_at"] = isoformat()
|
|
placeholders = ",".join("?" for _ in columns)
|
|
connection.execute(
|
|
f"INSERT INTO eod_progress({','.join(columns)}) VALUES ({placeholders})",
|
|
tuple(payload[name] for name in columns),
|
|
)
|
|
else:
|
|
assignments = ", ".join(f"{name} = ?" for name in fields)
|
|
connection.execute(
|
|
f"UPDATE eod_progress SET {assignments}, updated_at = ? WHERE trade_date = ?",
|
|
(*fields.values(), isoformat(), day),
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Job execution
|
|
# ------------------------------------------------------------------
|
|
|
|
def run_job(self, job_id: str, trade_date: str) -> dict[str, Any]:
|
|
fn = self.jobs.get(job_id)
|
|
if fn is None:
|
|
raise KeyError(job_id)
|
|
if job_id in EOD_JOB_IDS:
|
|
if not self._eod_lock.acquire(blocking=False):
|
|
return {
|
|
"job_id": job_id,
|
|
"state": "skipped",
|
|
"detail": "another EOD job is already running",
|
|
}
|
|
try:
|
|
return self._run_job(fn, job_id, trade_date)
|
|
finally:
|
|
self._eod_lock.release()
|
|
return self._run_job(fn, job_id, trade_date)
|
|
|
|
def _run_job(self, fn: JobFn, job_id: str, trade_date: str) -> dict[str, Any]:
|
|
started = isoformat()
|
|
run_id = None
|
|
with self.db.write() as connection:
|
|
cur = connection.execute(
|
|
"INSERT INTO job_runs(job_id, state, started_at, attempt) VALUES (?,?,?,1)",
|
|
(job_id, "running", started),
|
|
)
|
|
run_id = cur.lastrowid
|
|
try:
|
|
result = fn(trade_date) or {}
|
|
if isinstance(result, dict):
|
|
failures = self.pipeline.eod_failures(result) if job_id in EOD_JOB_IDS else []
|
|
if failures:
|
|
raise RuntimeError("; ".join(failures))
|
|
with self.db.write() as connection:
|
|
connection.execute(
|
|
"UPDATE job_runs SET state=?, finished_at=?, rows_out=?, detail=? WHERE id=?",
|
|
("ok", isoformat(), result.get("rows") if isinstance(result, dict) else None, str(result)[:2000], run_id),
|
|
)
|
|
return {"job_id": job_id, "result": result, "state": "ok"}
|
|
except Exception as exc:
|
|
with self.db.write() as connection:
|
|
connection.execute(
|
|
"UPDATE job_runs SET state=?, finished_at=?, error=? WHERE id=?",
|
|
("failed", isoformat(), str(exc), run_id),
|
|
)
|
|
raise
|
|
|
|
def _precheck(self, trade_date: str) -> dict[str, Any]:
|
|
return self.pipeline.ingest_reference(trade_date)
|
|
|
|
def _eod_a(self, trade_date: str) -> dict[str, Any]:
|
|
return self.pipeline.run_eod_batch_a(trade_date)
|
|
|
|
def _eod_b(self, trade_date: str) -> dict[str, Any]:
|
|
return self.pipeline.run_eod_batch_b(trade_date)
|
|
|
|
def _eod_retry(self, trade_date: str) -> dict[str, Any]:
|
|
return self.pipeline.run_eod_missing(trade_date)
|
|
|
|
def _history_backfill(self, trade_date: str) -> dict[str, Any]:
|
|
return self.pipeline.backfill_history(trade_date)
|
|
|
|
def _cleanup(self, trade_date: str) -> dict[str, Any]:
|
|
result = self.pipeline.cleanup()
|
|
if now_shanghai().weekday() == 6:
|
|
self.pipeline.db.vacuum()
|
|
result["vacuum"] = True
|
|
return result
|
|
|
|
def _backup(self, trade_date: str) -> dict[str, Any]:
|
|
from pathlib import Path
|
|
|
|
dest_dir = Path(self.pipeline.settings.backup_dir)
|
|
dest = dest_dir / f"datahub-{trade_date}.db"
|
|
self.pipeline.db.backup_to(dest)
|
|
keep = int(self.pipeline.settings.quality.get("backup_retain") or 14)
|
|
backups = sorted(dest_dir.glob("datahub-*.db"))
|
|
for old in backups[:-keep]:
|
|
old.unlink(missing_ok=True)
|
|
return {"path": str(dest.name), "kept": min(len(backups), keep)}
|