guard.py

#!/usr/bin/env python3
"""
guard.py — защита от OOM и алерты о крашах в Telegram.

Зачем:
  На free-Space жёсткий OOM = SIGKILL от ядра, процесс умирает мгновенно,
  никакой except не отрабатывает, алерт не уходит. Поэтому НЕ ждём краха —
  сами проверяем память перед каждым тяжёлым шагом и, если близко к лимиту,
  бросаем MemoryError ПОКА ЖИВЫ. Его ловит широкий except в build(), шлёт
  раскладку в телегу и отдаёт юзеру внятный отказ. Контролируемая посадка
  вместо немого краха.

Телеграм:
  Токен и chat_id берём ТОЛЬКО из окружения (Secrets Space), не из кода.
  Локально их нет -> алерты молча выключены. Шлём прямой POST на Bot API,
  без зависимости от внешних сервисов.

Память: psutil. Снимок копим по этапам, при крахе прикладываем последний.
"""

import os
import time
import urllib.request
import urllib.parse

try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    _HAS_PSUTIL = False

# 0.5 ГБ — реальный край на любой среде. Выше ставить нельзя: на free-Space
# память тесная и фоново занята, резерв в 1-1.5 ГБ давал бы ложную тревогу на
# старте (что и случилось с percent-порогом). 0.5 ГБ срабатывает только когда
# памяти правда почти не осталось — перед самым OOM-kill.
MEM_RESERVE_GB = 0.5

TG_TOKEN = os.environ.get("TG_BOT_TOKEN")
TG_CHAT = os.environ.get("TG_CHAT_ID")
IS_SPACE = bool(os.environ.get("SPACE_ID"))

_GB = 1024 ** 3
_last_alert_ts = 0.0
_ALERT_COOLDOWN = 300  # не чаще раза в 5 минут — антиспам при цикличных падениях


class MemorySnapshot:
    """Копит раскладку памяти по этапам обработки."""

    def __init__(self):
        self.stages = []          # [(label, rss_gb)]
        self.static_gb = 0.0      # старт+модели+видео (статика до кадров)
        self.videos_mb = 0.0
        self.n_videos = 0
        self.current_stage = "старт"
        self.mark("старт")

    @staticmethod
    def _rss_gb():
        if not _HAS_PSUTIL:
            return 0.0
        return psutil.Process().memory_info().rss / _GB

    @staticmethod
    def avail_gb():
        if not _HAS_PSUTIL:
            return 0.0
        return psutil.virtual_memory().total / _GB

    @staticmethod
    def used_pct():
        if not _HAS_PSUTIL:
            return 0.0
        return psutil.virtual_memory().percent

    def mark(self, label):
        self.current_stage = label
        self.stages.append((label, self._rss_gb()))

    def note_videos(self, n, total_mb):
        self.n_videos = n
        self.videos_mb = total_mb

    def freeze_static(self):
        """Зафиксировать статику = текущий RSS (вызывать после загрузки моделей+видео)."""
        self.static_gb = self._rss_gb()

    def report(self) -> str:
        lines = [f"Доступно всего: {self.avail_gb():.1f} ГБ",
                 f"Занято процессом: {self._rss_gb():.1f} ГБ ({self.used_pct():.0f}%)",
                 f"Упали на этапе: {self.current_stage}"]
        # раскладка по этапам как дельты
        prev = 0.0
        for label, rss in self.stages:
            delta = rss - prev
            lines.append(f"  {label}: {rss:.2f} ГБ (+{delta:.2f})")
            prev = rss
        lines.append(f"Видео: {self.n_videos} шт, {self.videos_mb:.0f} МБ на диске")
        lines.append(f"Статика (старт+модели+видео): {self.static_gb:.2f} ГБ")
        return "\n".join(lines)


def memory_guard(snapshot: MemorySnapshot, stage: str):
    """Проверка перед тяжёлым шагом. Бросает MemoryError ПОКА ЖИВЫ, если
    ДОСТУПНОЙ памяти осталось меньше резерва — чтобы краш был перехватываемым.

    Мерим available (сколько ещё можно занять), а НЕ percent: на Spaces системный
    percent фоново высокий (контейнер делит память), и percent>=90 на старте — это
    норма, а не тревога. available же честно говорит, сколько реально осталось.
    """
    snapshot.mark(stage)
    if not _HAS_PSUTIL:
        return
    vm = psutil.virtual_memory()
    avail_gb = vm.available / _GB
    if avail_gb < MEM_RESERVE_GB:
        raise MemoryError(
            f"Свободно лишь {avail_gb:.2f} ГБ (< резерва {MEM_RESERVE_GB} ГБ) "
            f"перед этапом «{stage}». Аварийная остановка до OOM-kill."
        )


def _tg_send(text: str):
    """Прямой POST в Telegram Bot API. Тихо ничего не делает, если нет токена."""
    if not (TG_TOKEN and TG_CHAT):
        return
    global _last_alert_ts
    now = time.time()
    if now - _last_alert_ts < _ALERT_COOLDOWN:
        return  # антиспам
    _last_alert_ts = now
    try:
        url = f"https://api.telegram.org/bot{TG_TOKEN}/sendMessage"
        data = urllib.parse.urlencode({"chat_id": TG_CHAT, "text": text}).encode()
        urllib.request.urlopen(url, data=data, timeout=10)
    except Exception:
        pass  # алерт не критичен — если телега недоступна, не валим основную работу


def alert_crash(snapshot: MemorySnapshot, err: Exception, ctx: dict):
    """Послать в телегу раскладку по крашу. ctx — доп. факты (слотов, цель и т.п.)."""
    where = "Space" if IS_SPACE else "локально"
    head = f"⚠️ MakeFaces [{where}] — СБОЙ\n{type(err).__name__}: {err}"
    ctx_lines = "\n".join(f"{k}: {v}" for k, v in ctx.items())
    body = snapshot.report() if snapshot else "снимок памяти недоступен"
    _tg_send(f"{head}\n\n{ctx_lines}\n\n{body}")