ltxc/runtime_reporting.py

from __future__ import annotations

import json
import logging
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

LOGGER_NAME = "ltx23"


def configure_logging() -> logging.Logger:
    logger = logging.getLogger(LOGGER_NAME)
    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(
            logging.Formatter(
                "%(asctime)s %(levelname)s %(name)s %(message)s",
                datefmt="%Y-%m-%dT%H:%M:%S",
            )
        )
        logger.addHandler(handler)
    logger.setLevel(logging.INFO)
    logger.propagate = False
    return logger


LOGGER = configure_logging()


def startup_event(event: str, message: str, **fields: Any) -> None:
    payload = {"event": event, "message": message, **fields}
    LOGGER.info(json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str))


class RunReporter:
    """One source for console logs, JSONL diagnostics, and Gradio progress."""

    def __init__(
        self,
        *,
        job_id: str,
        request_kind: str,
        event_path: str | Path,
        progress: Callable[..., Any] | None = None,
    ) -> None:
        self.job_id = str(job_id)
        self.request_kind = str(request_kind)
        self.event_path = Path(event_path)
        self.event_path.parent.mkdir(parents=True, exist_ok=True)
        self.progress = progress
        self.started = time.perf_counter()

    def emit(
        self,
        event: str,
        message: str,
        *,
        fraction: float | None = None,
        level: int = logging.INFO,
        **fields: Any,
    ) -> dict[str, Any]:
        record = {
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "elapsed_seconds": round(time.perf_counter() - self.started, 3),
            "event": str(event),
            "message": str(message),
            "job_id": self.job_id,
            "request_kind": self.request_kind,
            **fields,
        }
        with self.event_path.open("a", encoding="utf-8") as handle:
            handle.write(
                json.dumps(record, ensure_ascii=False, sort_keys=True, default=str)
                + "\n"
            )
        LOGGER.log(
            level, json.dumps(record, ensure_ascii=False, sort_keys=True, default=str)
        )
        if self.progress is not None and fraction is not None:
            self.progress(max(0.0, min(1.0, float(fraction))), desc=str(message))
        return record