update service

cbpred/prediction_service.py

@@ -0,0 +1,323 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import fcntl
+import json
+import os
+import re
+import smtplib
+import ssl
+import subprocess
+import sys
+import time
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from email.message import EmailMessage
+from pathlib import Path
+from typing import Any, Iterator
+from urllib.request import Request, urlopen
+
+import yaml
+
+from .pipeline import floor_to_8h
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+ENVIRONMENT_VARIABLE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
+
+
+def _utc_label(timestamp_ms: int) -> str:
+    return datetime.fromtimestamp(timestamp_ms / 1000, timezone.utc).strftime(
+        "%Y-%m-%d %H:%M:%S UTC"
+    )
+
+
+def _expand_environment(value: Any) -> Any:
+    if isinstance(value, str):
+        def replace(match: re.Match[str]) -> str:
+            name = match.group(1)
+            if name not in os.environ:
+                raise ValueError(f"environment variable {name} is required by the configuration")
+            return os.environ[name]
+
+        return ENVIRONMENT_VARIABLE.sub(replace, value)
+    if isinstance(value, list):
+        return [_expand_environment(item) for item in value]
+    if isinstance(value, dict):
+        return {str(key): _expand_environment(item) for key, item in value.items()}
+    return value
+
+
+def _as_mapping(value: Any, name: str) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if not isinstance(value, dict):
+        raise ValueError(f"{name} must be a mapping")
+    return value
+
+
+def load_config(path: Path) -> dict[str, Any]:
+    data = yaml.safe_load(path.read_text()) or {}
+    config = _as_mapping(_expand_environment(data), "root configuration")
+    config["_config_dir"] = path.resolve().parent
+    prediction = _as_mapping(config.get("prediction"), "prediction")
+    if not prediction:
+        raise ValueError("prediction configuration is required")
+    config["prediction"] = prediction
+    config["scheduler"] = _as_mapping(config.get("scheduler"), "scheduler")
+    config["notifications"] = _as_mapping(config.get("notifications"), "notifications")
+    return config
+
+
+def _config_path(config: dict[str, Any], value: str | Path) -> Path:
+    path = Path(value)
+    return path if path.is_absolute() else Path(config["_config_dir"]) / path
+
+
+def _state_path(config: dict[str, Any]) -> Path:
+    value = config["scheduler"].get("state_file", "prediction-service-state.json")
+    return _config_path(config, value)
+
+
+def _read_state(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    data = json.loads(path.read_text())
+    return _as_mapping(data, "service state")
+
+
+def _write_state(path: Path, state: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    temporary.write_text(json.dumps(state, indent=2, ensure_ascii=False) + "\n")
+    temporary.replace(path)
+
+
+@contextmanager
+def _run_lock(path: Path) -> Iterator[bool]:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a+") as handle:
+        try:
+            fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            yield False
+            return
+        try:
+            yield True
+        finally:
+            fcntl.flock(handle, fcntl.LOCK_UN)
+
+
+def _prediction_command(config: dict[str, Any], end_ms: int) -> tuple[list[str], Path]:
+    prediction = config["prediction"]
+    working_dir = _config_path(config, prediction.get("working_dir", REPO_ROOT)).resolve()
+    command = [sys.executable, "-m", "cbpred.latest_predict", "--end", str(end_ms)]
+    path_options = {
+        "banbot_dir": "--banbot-dir",
+        "banbot_config": "--banbot-config",
+        "data_dir": "--data-dir",
+        "artifact": "--artifact",
+        "output_dir": "--output-dir",
+    }
+    value_options = {
+        "history_months": "--history-months",
+        "holdout_months": "--holdout-months",
+        "min_age_days": "--min-age-days",
+        "chunk_size": "--chunk-size",
+        "dsn": "--dsn",
+        "db_type": "--db-type",
+        "exchange": "--exchange",
+    }
+    for key, option in path_options.items():
+        if key in prediction:
+            command.extend([option, str(_config_path(config, prediction[key]).resolve())])
+    for key, option in value_options.items():
+        if prediction.get(key) is not None:
+            command.extend([option, str(prediction[key])])
+    return command, working_dir
+
+
+def run_prediction(config: dict[str, Any], end_ms: int) -> tuple[dict[str, Any], Path]:
+    command, working_dir = _prediction_command(config, end_ms)
+    print(f"[prediction-service] running completed candle ending {_utc_label(end_ms)}", flush=True)
+    subprocess.run(command, cwd=working_dir, check=True)
+    output_dir = _config_path(
+        config, config["prediction"].get("output_dir", "predictions/latest_v1")
+    ).resolve()
+    report_path = output_dir / "report.json"
+    if not report_path.exists():
+        raise RuntimeError(f"prediction finished without {report_path}")
+    return _as_mapping(json.loads(report_path.read_text()), "prediction report"), output_dir / "predictions.csv"
+
+
+def _prediction_rows(path: Path, max_rows: int | None) -> list[dict[str, str]]:
+    if not path.exists():
+        raise RuntimeError(f"prediction finished without {path}")
+    with path.open(newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    return rows if max_rows is None else rows[:max_rows]
+
+
+def _notification_payload(
+    summary: dict[str, Any],
+    csv_path: Path,
+    end_ms: int,
+    include_predictions: bool,
+    max_rows: int | None,
+) -> dict[str, Any]:
+    payload: dict[str, Any] = {
+        "event": "cbpred.prediction.completed",
+        "candle_end_ms": end_ms,
+        "candle_end_utc": _utc_label(end_ms),
+        "summary": summary,
+    }
+    if include_predictions:
+        payload["predictions"] = _prediction_rows(csv_path, max_rows)
+    return payload
+
+
+def send_http_notifications(
+    config: dict[str, Any], summary: dict[str, Any], csv_path: Path, end_ms: int
+) -> int:
+    targets = config["notifications"].get("http", [])
+    if not isinstance(targets, list):
+        raise ValueError("notifications.http must be a list")
+    sent = 0
+    for index, value in enumerate(targets, 1):
+        target = _as_mapping(value, f"notifications.http[{index}]")
+        url = target.get("url")
+        if not isinstance(url, str) or not url:
+            raise ValueError(f"notifications.http[{index}].url is required")
+        max_rows = target.get("max_rows")
+        if max_rows is not None and (not isinstance(max_rows, int) or max_rows < 0):
+            raise ValueError(f"notifications.http[{index}].max_rows must be a non-negative integer")
+        payload = _notification_payload(
+            summary,
+            csv_path,
+            end_ms,
+            bool(target.get("include_predictions", True)),
+            max_rows,
+        )
+        headers = {str(key): str(item) for key, item in _as_mapping(target.get("headers"), "headers").items()}
+        headers.setdefault("Content-Type", "application/json")
+        request = Request(
+            url,
+            data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
+            headers=headers,
+            method=str(target.get("method", "POST")).upper(),
+        )
+        timeout = float(target.get("timeout_seconds", 15))
+        with urlopen(request, timeout=timeout) as response:
+            response.read()
+        sent += 1
+    return sent
+
+
+def send_email_notification(
+    config: dict[str, Any], summary: dict[str, Any], csv_path: Path, end_ms: int
+) -> bool:
+    email_config = config["notifications"].get("email")
+    if not email_config or not _as_mapping(email_config, "notifications.email").get("enabled", True):
+        return False
+    email_config = _as_mapping(email_config, "notifications.email")
+    required = ("host", "from", "to")
+    missing = [name for name in required if not email_config.get(name)]
+    if missing:
+        raise ValueError(f"notifications.email is missing {', '.join(missing)}")
+    recipients = email_config["to"]
+    if not isinstance(recipients, list) or not all(isinstance(item, str) and item for item in recipients):
+        raise ValueError("notifications.email.to must be a non-empty list of addresses")
+
+    subject = str(email_config.get("subject", "cbpred 8H prediction completed"))
+    message = EmailMessage()
+    message["Subject"] = subject
+    message["From"] = str(email_config["from"])
+    message["To"] = ", ".join(recipients)
+    message.set_content(
+        "\n".join(
+            (
+                f"Completed candle: {_utc_label(end_ms)}",
+                f"Latest bar: {summary.get('latest_bar_start_utc')}",
+                f"Predicted symbols: {summary.get('symbols')}",
+                f"Eligible predictions: {summary.get('eligible_predictions')}",
+                f"CSV: {summary.get('csv')}",
+            )
+        )
+        + "\n"
+    )
+    if email_config.get("attach_csv", True):
+        message.add_attachment(csv_path.read_bytes(), maintype="text", subtype="csv", filename=csv_path.name)
+
+    port = int(email_config.get("port", 465 if email_config.get("ssl") else 587))
+    smtp_class = smtplib.SMTP_SSL if email_config.get("ssl") else smtplib.SMTP
+    with smtp_class(str(email_config["host"]), port, timeout=float(email_config.get("timeout_seconds", 30))) as client:
+        if email_config.get("starttls", not email_config.get("ssl")):
+            client.starttls(context=ssl.create_default_context())
+        if email_config.get("username"):
+            client.login(str(email_config["username"]), str(email_config.get("password", "")))
+        client.send_message(message)
+    return True
+
+
+def _ready_end_ms(now_ms: int, grace_seconds: int) -> int | None:
+    end_ms = floor_to_8h(now_ms)
+    return end_ms if now_ms >= end_ms + grace_seconds * 1000 else None
+
+
+def run_ready_cycle(config: dict[str, Any], now_ms: int | None = None) -> str:
+    now_ms = now_ms if now_ms is not None else int(time.time() * 1000)
+    grace_seconds = int(config["scheduler"].get("grace_seconds", 120))
+    if grace_seconds < 0:
+        raise ValueError("scheduler.grace_seconds must be non-negative")
+    end_ms = _ready_end_ms(now_ms, grace_seconds)
+    if end_ms is None:
+        return "waiting for the current 8H candle to close"
+    state_path = _state_path(config)
+    with _run_lock(state_path.with_suffix(state_path.suffix + ".lock")) as acquired:
+        if not acquired:
+            return "another prediction run is in progress"
+        state = _read_state(state_path)
+        if state.get("last_successful_end_ms") == end_ms:
+            return f"completed candle {_utc_label(end_ms)} was already processed"
+        summary, csv_path = run_prediction(config, end_ms)
+        http_count = send_http_notifications(config, summary, csv_path, end_ms)
+        email_sent = send_email_notification(config, summary, csv_path, end_ms)
+        _write_state(
+            state_path,
+            {
+                "last_successful_end_ms": end_ms,
+                "last_successful_end_utc": _utc_label(end_ms),
+                "updated_at_utc": _utc_label(int(time.time() * 1000)),
+            },
+        )
+    return f"processed {_utc_label(end_ms)}; http={http_count}, email={email_sent}"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run cbpred after each completed 8H K-line")
+    parser.add_argument("--config", type=Path, required=True, help="YAML service configuration")
+    parser.add_argument("--once", action="store_true", help="check once, then exit")
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    config = load_config(args.config)
+    interval_seconds = int(config["scheduler"].get("poll_seconds", 60))
+    if interval_seconds <= 0:
+        raise ValueError("scheduler.poll_seconds must be positive")
+    while True:
+        try:
+            print(f"[prediction-service] {run_ready_cycle(config)}", flush=True)
+        except Exception as exc:
+            print(f"[prediction-service] failed: {exc}", file=sys.stderr, flush=True)
+            if args.once:
+                raise
+        if args.once:
+            return
+        time.sleep(interval_seconds)
+
+
+if __name__ == "__main__":
+    main()