memory.py
python
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65
refactor: enforce gRPC framing on all MWP wire traffic
Sonnet 4.6
minor
⚠ breaking
156 days ago
| 1 | """Process memory profiling utilities. |
| 2 | |
| 3 | Three instruments: |
| 4 | |
| 5 | 1. ``rss_mb()`` — current RSS in MiB (cheap, single call). |
| 6 | |
| 7 | 2. ``profile_task(name)`` — async context manager that logs RSS delta before |
| 8 | and after a block. Use it to wrap every background task so we can see |
| 9 | exactly which one is the memory hog. |
| 10 | |
| 11 | 3. ``top_allocations(limit)`` — returns the top tracemalloc frames by |
| 12 | cumulative memory. Requires tracemalloc to be started at boot |
| 13 | (``tracemalloc.start()`` in ``main.py``). |
| 14 | |
| 15 | 4. ``MemoryLogMiddleware`` — ASGI middleware that logs RSS after every |
| 16 | request above a configurable threshold. |
| 17 | |
| 18 | Usage in background tasks:: |
| 19 | |
| 20 | from musehub.debug.memory import profile_task |
| 21 | |
| 22 | async with profile_task("symbol-index repo=abc123"): |
| 23 | await build_symbol_index(session, repo_id, head) |
| 24 | |
| 25 | Usage in route handlers for a one-off snapshot:: |
| 26 | |
| 27 | from musehub.debug.memory import rss_mb |
| 28 | logger.info("rss after push: %.1f MiB", rss_mb()) |
| 29 | """ |
| 30 | from __future__ import annotations |
| 31 | |
| 32 | import logging |
| 33 | import os |
| 34 | import time |
| 35 | import tracemalloc |
| 36 | from contextlib import asynccontextmanager |
| 37 | from typing import AsyncGenerator |
| 38 | |
| 39 | from starlette.types import ASGIApp, Receive, Scope, Send |
| 40 | from musehub.types.json_types import MemoryFrame |
| 41 | |
| 42 | logger = logging.getLogger(__name__) |
| 43 | |
| 44 | try: |
| 45 | import psutil as _psutil |
| 46 | _PROC = _psutil.Process(os.getpid()) |
| 47 | |
| 48 | def rss_mb() -> float: |
| 49 | """Return current RSS in MiB.""" |
| 50 | return float(_PROC.memory_info().rss) / 1024 / 1024 |
| 51 | |
| 52 | except ImportError: # pragma: no cover — psutil optional in test envs |
| 53 | def rss_mb() -> float: # noqa: F811 |
| 54 | return -1.0 |
| 55 | |
| 56 | |
| 57 | @asynccontextmanager |
| 58 | async def profile_task(name: str) -> AsyncGenerator[None, None]: |
| 59 | """Log RSS before and after an async block. |
| 60 | |
| 61 | Logs at INFO so it always appears in production logs. If the delta |
| 62 | exceeds 50 MiB it is also logged at WARNING so it's easy to grep. |
| 63 | |
| 64 | Example log output:: |
| 65 | |
| 66 | [memory] START symbol-index repo=abc rss=210.3 MiB |
| 67 | [memory] END symbol-index repo=abc rss=310.7 MiB delta=+100.4 MiB elapsed=2.14s ← WARNING |
| 68 | """ |
| 69 | rss_before = rss_mb() |
| 70 | t0 = time.monotonic() |
| 71 | try: |
| 72 | yield |
| 73 | finally: |
| 74 | elapsed = time.monotonic() - t0 |
| 75 | rss_after = rss_mb() |
| 76 | delta = rss_after - rss_before |
| 77 | sign = "+" if delta >= 0 else "" |
| 78 | msg = ( |
| 79 | f"[memory] END {name} " |
| 80 | f"rss={rss_after:.1f} MiB delta={sign}{delta:.1f} MiB elapsed={elapsed:.2f}s" |
| 81 | ) |
| 82 | if abs(delta) >= 50: |
| 83 | logger.warning(msg) |
| 84 | else: |
| 85 | logger.info(msg) |
| 86 | logger.info( |
| 87 | "[memory] START %s rss=%.1f MiB", |
| 88 | name, rss_before, |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def top_allocations(limit: int = 20) -> list[MemoryFrame]: |
| 93 | """Return the top *limit* tracemalloc frames by cumulative size. |
| 94 | |
| 95 | Returns an empty list if tracemalloc was not started. |
| 96 | """ |
| 97 | if not tracemalloc.is_tracing(): |
| 98 | return [] |
| 99 | snapshot = tracemalloc.take_snapshot() |
| 100 | stats = snapshot.statistics("lineno") |
| 101 | result = [] |
| 102 | for stat in stats[:limit]: |
| 103 | frame = stat.traceback[0] |
| 104 | result.append({ |
| 105 | "file": frame.filename, |
| 106 | "line": frame.lineno, |
| 107 | "size_kb": round(stat.size / 1024, 1), |
| 108 | "count": stat.count, |
| 109 | }) |
| 110 | return result |
| 111 | |
| 112 | |
| 113 | class MemoryLogMiddleware: |
| 114 | """ASGI middleware that logs RSS after every response. |
| 115 | |
| 116 | Only logs when RSS exceeds ``warn_above_mb`` (default 400 MiB) so that |
| 117 | normal traffic doesn't flood the logs. Set ``warn_above_mb=0`` to log |
| 118 | every request (noisy but useful during a profiling session). |
| 119 | """ |
| 120 | |
| 121 | def __init__(self, app: ASGIApp, warn_above_mb: float = 400.0) -> None: |
| 122 | self.app = app |
| 123 | self.warn_above_mb = warn_above_mb |
| 124 | |
| 125 | async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: |
| 126 | if scope["type"] not in ("http", "websocket"): |
| 127 | await self.app(scope, receive, send) |
| 128 | return |
| 129 | |
| 130 | await self.app(scope, receive, send) |
| 131 | |
| 132 | rss = rss_mb() |
| 133 | if rss >= self.warn_above_mb: |
| 134 | path = scope.get("path", "?") |
| 135 | method = scope.get("method", "?") |
| 136 | logger.warning( |
| 137 | "[memory] HIGH RSS %.1f MiB after %s %s", |
| 138 | rss, method, path, |
| 139 | ) |
File History
1 commit
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65
refactor: enforce gRPC framing on all MWP wire traffic
Sonnet 4.6
minor
⚠
156 days ago