pipu-music/services/recommender/pipu_rec/train/pipeline.py

"""Ночной пайплайн (docs/recommendations.md#8): события → implicit_score → ALS → проверка метрик →
векторы в track_embeddings.cf / user_profiles_rec.taste_cf и запись в model_versions."""

from __future__ import annotations

import datetime as dt
import json
import logging
import threading
import time
from collections import Counter
from typing import Any

import numpy as np
import psycopg

from pipu_rec.train.als import ALSModel, normalize, train_als
from pipu_rec.train.bandit import update_bandits
from pipu_rec.train.metrics import evaluate
from pipu_rec.train.ranker import train_and_save
from pipu_rec.train.signals import DISLIKE, LIKE, Interaction, event_weight, implicit_scores

log = logging.getLogger(__name__)

HOLDOUT_DAYS = 7
# Новая модель активируется, только если NDCG@10 не хуже активной более чем на 2%.
MAX_REGRESSION = 0.02
MIN_INTERACTIONS = 50
ADVISORY_LOCK = 0x70697075  # «pipu»: один прогон обучения на кластер


def load_interactions(conn: psycopg.Connection[Any]) -> list[Interaction]:
    out: list[Interaction] = []
    for user, track, ts, kind, played, duration, in_wave in conn.execute(
        """
        SELECT e.user_id::text, e.track_id::text, extract(epoch FROM e.occurred_at)::float8, e.type,
            e.played_ms, t.duration_ms, (e.wave_session_id IS NOT NULL OR e.context_type = 'wave')
        FROM play_events e JOIN tracks t ON t.id = e.track_id
        WHERE e.type IN ('play_end', 'skip') AND e.occurred_at > now() - interval '365 days'
        """
    ):
        w = event_weight(kind, played, duration, bool(in_wave))
        if w is not None:
            out.append(Interaction(user, track, float(ts), w))
    for user, track, ts, reaction in conn.execute(
        "SELECT user_id::text, track_id::text, extract(epoch FROM created_at)::float8, reaction::text "
        "FROM track_reactions"
    ):
        out.append(Interaction(user, track, float(ts), LIKE if reaction == "like" else DISLIKE))
    return out


def holdout(
    interactions: list[Interaction], now: float, days: int = HOLDOUT_DAYS
) -> tuple[dict[str, dict[str, float]], dict[str, set[str]]]:
    """Разбиение по времени: обучение — до отсечки, тест — треки, впервые понравившиеся после неё."""
    cutoff = now - days * 86400
    train = implicit_scores(interactions, now, before=cutoff)
    after: dict[str, dict[str, float]] = {}
    for it in interactions:
        if it.ts >= cutoff:
            row = after.setdefault(it.user, {})
            row[it.track] = row.get(it.track, 0.0) + it.weight
    relevant = {u: {t for t, w in c.items() if w > 0 and t not in train.get(u, {})} for u, c in after.items()}
    return train, {u: r for u, r in relevant.items() if r}


def offline_metrics(interactions: list[Interaction], now: float) -> dict[str, Any]:
    """ALS против популярности на отложенной неделе."""
    train, relevant = holdout(interactions, now)
    model = train_als(train)
    if model is None or not relevant:
        return {"users": 0}
    pop = Counter(t for row in train.values() for t, s in row.items() if s > 0)
    popular = [t for t, _ in pop.most_common()]

    def rank_als(u: str) -> list[str]:
        seen = train.get(u, {})
        order = np.argsort(-model.scores(u))
        return [model.items[i] for i in order if model.items[i] not in seen][:50]

    def rank_pop(u: str) -> list[str]:
        seen = train.get(u, {})
        return [t for t in popular if t not in seen][:50]

    users = list(relevant)
    return {"als": evaluate(users, relevant, rank_als), "popularity": evaluate(users, relevant, rank_pop)}


def vec_text(v: np.ndarray) -> str:
    return "[" + ",".join(f"{x:.6f}" for x in v.tolist()) + "]"


def active_metrics(conn: psycopg.Connection[Any], kind: str) -> dict[str, Any] | None:
    row = conn.execute(
        "SELECT metrics FROM model_versions WHERE kind = %s AND is_active ORDER BY created_at DESC LIMIT 1",
        (kind,),
    ).fetchone()
    return row[0] if row and row[0] else None


def should_activate(new: dict[str, Any], current: dict[str, Any] | None) -> bool:
    """Без активной модели или без тестовых данных — активируем; иначе NDCG@10 не хуже на > 2%."""
    if current is None:
        return True
    new_ndcg = new.get("als", {}).get("ndcg@10")
    cur_ndcg = (current.get("offline") or {}).get("als", {}).get("ndcg@10")
    if new_ndcg is None or not cur_ndcg:
        return True
    return float(new_ndcg) >= float(cur_ndcg) * (1 - MAX_REGRESSION)


# Размерность колонок track_embeddings.cf / user_profiles_rec.taste_cf (vector(64)).
CF_DIM = 64


def pad(v: np.ndarray, dim: int = CF_DIM) -> np.ndarray:
    """Маленькая библиотека даёт меньше факторов, чем колонка: добиваем нулями (косинус не меняется)."""
    if v.shape[1] >= dim:
        return v[:, :dim]
    return np.pad(v, ((0, 0), (0, dim - v.shape[1])))


def save_als(conn: psycopg.Connection[Any], model: ALSModel, version: str, metrics: dict[str, Any]) -> None:
    items = pad(normalize(model.item_factors))
    users = pad(normalize(model.user_factors))
    with conn.transaction():
        with conn.cursor() as cur:
            cur.executemany(
                """
                INSERT INTO track_embeddings (track_id, cf, updated_at) VALUES (%s, %s::vector, now())
                ON CONFLICT (track_id) DO UPDATE SET cf = EXCLUDED.cf, updated_at = now()
                """,
                [(t, vec_text(items[i])) for i, t in enumerate(model.items)],
            )
            cur.execute(
                "UPDATE track_embeddings SET cf = NULL "
                "WHERE cf IS NOT NULL AND NOT (track_id::text = ANY(%s))",
                (model.items,),
            )
            cur.executemany(
                """
                INSERT INTO user_profiles_rec (user_id, taste_cf, updated_at) VALUES (%s, %s::vector, now())
                ON CONFLICT (user_id) DO UPDATE SET taste_cf = EXCLUDED.taste_cf, updated_at = now()
                """,
                [(u, vec_text(users[i])) for i, u in enumerate(model.users)],
            )
        conn.execute("UPDATE model_versions SET is_active = false WHERE kind = 'als' AND is_active")
        conn.execute(
            "INSERT INTO model_versions (id, kind, metrics, is_active) VALUES (%s, 'als', %s, true)",
            (version, json.dumps(metrics)),
        )


def train_all(dsn: str, now: float | None = None) -> dict[str, Any]:
    """Один прогон обучения. Возвращает сводку (её же пишет в лог и model_versions)."""
    now = now or time.time()
    with psycopg.connect(dsn) as conn:
        if not conn.execute("SELECT pg_try_advisory_lock(%s)", (ADVISORY_LOCK,)).fetchone()[0]:  # type: ignore[index]
            return {"skipped": "another training is running"}
        try:
            # Бандит считается всегда — ему хватает и пары сессий волны, в отличие от ALS.
            bandit_users = update_bandits(conn, now)
            # Ранжировщик v2 учится на показах волны и от числа событий ALS не зависит.
            try:
                ranker = train_and_save(conn, now)
            except Exception as e:  # обучение ранжировщика не должно мешать ALS
                log.exception("ranker training failed")
                ranker = {"ranker": {"error": str(e)}}
            interactions = load_interactions(conn)
            if len(interactions) < MIN_INTERACTIONS:
                return {
                    "skipped": f"only {len(interactions)} interactions",
                    "bandit_users": bandit_users,
                    **ranker,
                }
            started = time.monotonic()
            offline = offline_metrics(interactions, now)
            model = train_als(implicit_scores(interactions, now))
            if model is None:
                return {"skipped": "no data"}
            version = "als-" + dt.datetime.fromtimestamp(now, dt.UTC).strftime("%Y%m%d%H%M")
            metrics = {
                "offline": offline,
                "users": len(model.users),
                "items": len(model.items),
                "interactions": len(interactions),
                "bandit_users": bandit_users,
                "train_s": round(time.monotonic() - started, 1),
            }
            if not should_activate(offline, active_metrics(conn, "als")):
                conn.execute(
                    "INSERT INTO model_versions (id, kind, metrics, is_active) VALUES (%s, 'als', %s, false) "
                    "ON CONFLICT (id) DO NOTHING",
                    (version, json.dumps(metrics)),
                )
                log.warning("als %s not activated: offline metrics regressed %s", version, offline)
                return {"version": version, "activated": False, **metrics, **ranker}
            save_als(conn, model, version, metrics)
            log.info("als %s activated: %s", version, metrics)
            return {"version": version, "activated": True, **metrics, **ranker}
        finally:
            conn.execute("SELECT pg_advisory_unlock(%s)", (ADVISORY_LOCK,))


class NightlyTrainer:
    """Обучение раз в сутки в hour:00 по местному времени сервера; при первом запуске без модели — сразу.
    on_done вызывается после успешной активации (горячая подмена каталога в памяти)."""

    def __init__(self, dsn: str, hour: int = 3, on_done: Any = None) -> None:
        self.dsn = dsn
        self.hour = hour
        self.on_done = on_done
        self._stop = threading.Event()

    def start(self) -> None:
        threading.Thread(target=self._loop, name="nightly-trainer", daemon=True).start()

    def stop(self) -> None:
        self._stop.set()

    def seconds_until_next(self, now: dt.datetime) -> float:
        target = now.replace(hour=self.hour, minute=0, second=0, microsecond=0)
        if target <= now:
            target += dt.timedelta(days=1)
        return (target - now).total_seconds()

    def _run(self) -> None:
        try:
            result = train_all(self.dsn)
            log.info("training finished: %s", result)
            ranker_on = bool((result.get("ranker") or {}).get("activated"))
            if (result.get("activated") or ranker_on) and self.on_done:
                self.on_done()
        except Exception:
            log.exception("training failed")

    def _loop(self) -> None:
        try:
            with psycopg.connect(self.dsn) as conn:
                never = active_metrics(conn, "als") is None
        except Exception:
            never = False
        if never:
            self._run()
        while not self._stop.wait(self.seconds_until_next(dt.datetime.now())):
            self._run()