CoolFace
Apppublic

build-small-hackathon/hackathon-advisor

sourceHugging Facemitupdated 3mo agoView on Hugging Face
16likes
profiling.py166 linesDownload Raw Back to hackathon_advisor
1"""Lightweight logging and per-turn profiling for the advisor runtime.2 3The numbers here are debug/operations signal only — they are written to logs, never to the4UI. Stage timings are measured by *observing the turn event stream from the main process*, so5they stay correct even when the model itself runs inside a ZeroGPU fork (where a module-global6counter would reset on every call).7"""8 9from __future__ import annotations10 11from dataclasses import dataclass, field12import logging13import os14import platform15import sys16import threading17import time18from typing import Any19 20logger = logging.getLogger("hackathon_advisor")21 22_counter_lock = threading.Lock()23_messages_processed = 024 25 26def configure_logging() -> None:27    """Attach a stream handler once, honoring ADVISOR_LOG_LEVEL (default INFO)."""28    level_name = os.environ.get("ADVISOR_LOG_LEVEL", "INFO").strip().upper()29    logger.setLevel(getattr(logging, level_name, logging.INFO))30    if not logger.handlers:31        handler = logging.StreamHandler()32        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s"))33        logger.addHandler(handler)34    logger.propagate = False35 36 37def next_message_index() -> int:38    """Increment and return the lifetime count of processed advisor messages (main process)."""39    global _messages_processed40    with _counter_lock:41        _messages_processed += 142        return _messages_processed43 44 45def messages_processed() -> int:46    return _messages_processed47 48 49def _ms(seconds: float) -> float:50    return round(seconds * 1000.0, 1)51 52 53def resource_snapshot() -> dict[str, Any]:54    """Best-effort process resource usage via the stdlib plus torch device memory if torch is55    already imported. Returns whatever could be sampled; never raises."""56    snapshot: dict[str, Any] = {}57    try:58        import resource59 60        usage = resource.getrusage(resource.RUSAGE_SELF)61        # ru_maxrss is bytes on macOS, kilobytes on Linux.62        divisor = 1024 * 1024 if platform.system() == "Darwin" else 102463        snapshot["rss_mb"] = round(usage.ru_maxrss / divisor, 1)64        snapshot["cpu_user_s"] = round(usage.ru_utime, 3)65        snapshot["cpu_sys_s"] = round(usage.ru_stime, 3)66    except Exception:  # pragma: no cover - platform dependent67        pass68    snapshot.update(_torch_memory_snapshot())69    return snapshot70 71 72def _torch_memory_snapshot() -> dict[str, Any]:73    out: dict[str, Any] = {}74    torch = sys.modules.get("torch")  # do not import torch just to profile75    if torch is None:76        return out77    try:78        if torch.cuda.is_available():79            out["cuda_alloc_mb"] = round(torch.cuda.memory_allocated() / 1e6, 1)80            out["cuda_peak_mb"] = round(torch.cuda.max_memory_allocated() / 1e6, 1)81    except Exception:  # pragma: no cover - device dependent82        pass83    try:84        mps = getattr(torch, "mps", None)85        current = getattr(mps, "current_allocated_memory", None)86        if current is not None:87            out["mps_alloc_mb"] = round(current() / 1e6, 1)88    except Exception:  # pragma: no cover - device dependent89        pass90    return out91 92 93@dataclass94class TurnProfiler:95    """Times a single advisor turn by observing its event stream. Drive it by calling96    ``observe(event)`` for every emitted event dict, then ``log_summary()`` when the turn97    ends (in a finally block, so partial turns still get logged)."""98 99    message_index: int100    compute: str101    backend: str102    device: str = ""103    message_chars: int = 0104    started: float = field(default_factory=time.perf_counter)105    stage_at: dict[str, float] = field(default_factory=dict)106    ended: float | None = None107    tokens: int = 0108    tool_count: int = 0109    fell_back: bool = False110    logged: bool = False111 112    def log_start(self) -> None:113        logger.info(114            "turn #%d start | compute=%s backend=%s message_chars=%d",115            self.message_index,116            self.compute,117            self.backend,118            self.message_chars,119        )120 121    def observe(self, event: dict[str, Any]) -> None:122        now = time.perf_counter()123        event_type = event.get("type")124        if event_type == "stage":125            self.stage_at.setdefault(str(event.get("stage")), now)126        elif event_type == "model_progress":127            self.tokens = max(self.tokens, int(event.get("tokens") or 0))128        elif event_type == "tool_event":129            self.tool_count += 1130        elif event_type == "fallback":131            self.fell_back = True132        elif event_type == "done":133            self.ended = now134 135    def durations(self) -> dict[str, float]:136        end = self.ended if self.ended is not None else time.perf_counter()137        out: dict[str, float] = {"total_ms": _ms(end - self.started)}138        planning = self.stage_at.get("planning")139        running = self.stage_at.get("running_tool")140        writing = self.stage_at.get("writing")141        if planning is not None and running is not None:142            out["decode_ms"] = _ms(running - planning)143        if running is not None and writing is not None:144            out["tools_ms"] = _ms(writing - running)145        if writing is not None:146            out["write_ms"] = _ms(end - writing)147        return out148 149    def log_summary(self, error: BaseException | None = None) -> None:150        if self.logged:151            return152        self.logged = True153        durations = self.durations()154        timing = " ".join(f"{key}={value}" for key, value in durations.items())155        resources = " ".join(f"{key}={value}" for key, value in resource_snapshot().items())156        status = "error" if error is not None else "done"157        message = (158            f"turn #{self.message_index} {status} | {timing} | "159            f"tokens={self.tokens} tools={self.tool_count} compute={self.compute} "160            f"device={self.device or '?'} backend={self.backend} fallback={self.fell_back} | {resources}"161        )162        if error is not None:163            logger.warning("%s | exception=%r", message, error)164        else:165            logger.info(message)166