gabriel / musehub public
memory.py python
135 lines 4.2 KB
Raw
sha256:34035d72cef530c1ab9d6a6f53be18d803dde705fc3157617d70352a96d0747b Merge branch 'fix/wire-push-external-parent-manifest' into dev Human 8 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
31 import logging
32 import os
33 import time
34 import tracemalloc
35 from contextlib import asynccontextmanager
36 from typing import AsyncGenerator
37
38 from starlette.types import ASGIApp, Receive, Scope, Send
39 from musehub.types.json_types import MemoryFrame
40
41 logger = logging.getLogger(__name__)
42
43 try:
44 import psutil as _psutil
45 _PROC = _psutil.Process(os.getpid())
46
47 def rss_mb() -> float:
48 """Return current RSS in MiB."""
49 return float(_PROC.memory_info().rss) / 1024 / 1024
50
51 except ImportError: # pragma: no cover — psutil optional in test envs
52 def rss_mb() -> float: # noqa: F811
53 return -1.0
54
55 @asynccontextmanager
56 async def profile_task(name: str) -> AsyncGenerator[None, None]:
57 """Log RSS before and after an async block.
58
59 Logs at INFO so it always appears in production logs. If the delta
60 exceeds 50 MiB it is also logged at WARNING so it's easy to grep.
61
62 Example log output::
63
64 [memory] START symbol-index repo=abc rss=210.3 MiB
65 [memory] END symbol-index repo=abc rss=310.7 MiB delta=+100.4 MiB elapsed=2.14s ← WARNING
66 """
67 rss_before = rss_mb()
68 t0 = time.monotonic()
69 try:
70 yield
71 finally:
72 elapsed = time.monotonic() - t0
73 rss_after = rss_mb()
74 delta = rss_after - rss_before
75 sign = "+" if delta >= 0 else ""
76 msg = (
77 f"[memory] END {name} "
78 f"rss={rss_after:.1f} MiB delta={sign}{delta:.1f} MiB elapsed={elapsed:.2f}s"
79 )
80 if abs(delta) >= 50:
81 logger.warning(msg)
82 else:
83 logger.info(msg)
84 logger.info(
85 "[memory] START %s rss=%.1f MiB",
86 name, rss_before,
87 )
88
89 def top_allocations(limit: int = 20) -> list[MemoryFrame]:
90 """Return the top *limit* tracemalloc frames by cumulative size.
91
92 Returns an empty list if tracemalloc was not started.
93 """
94 if not tracemalloc.is_tracing():
95 return []
96 snapshot = tracemalloc.take_snapshot()
97 stats = snapshot.statistics("lineno")
98 result = []
99 for stat in stats[:limit]:
100 frame = stat.traceback[0]
101 result.append({
102 "file": frame.filename,
103 "line": frame.lineno,
104 "size_kb": round(stat.size / 1024, 1),
105 "count": stat.count,
106 })
107 return result
108
109 class MemoryLogMiddleware:
110 """ASGI middleware that logs RSS after every response.
111
112 Only logs when RSS exceeds ``warn_above_mb`` (default 400 MiB) so that
113 normal traffic doesn't flood the logs. Set ``warn_above_mb=0`` to log
114 every request (noisy but useful during a profiling session).
115 """
116
117 def __init__(self, app: ASGIApp, warn_above_mb: float = 400.0) -> None:
118 self.app = app
119 self.warn_above_mb = warn_above_mb
120
121 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
122 if scope["type"] not in ("http", "websocket"):
123 await self.app(scope, receive, send)
124 return
125
126 await self.app(scope, receive, send)
127
128 rss = rss_mb()
129 if rss >= self.warn_above_mb:
130 path = scope.get("path", "?")
131 method = scope.get("method", "?")
132 logger.warning(
133 "[memory] HIGH RSS %.1f MiB after %s %s",
134 rss, method, path,
135 )
File History 14 commits
sha256:34035d72cef530c1ab9d6a6f53be18d803dde705fc3157617d70352a96d0747b Merge branch 'fix/wire-push-external-parent-manifest' into dev Human 8 days ago
sha256:fc04e4cae9e1774d6a21b65c45daeed0e6787eb581d13aa1b03bfe9384a34226 Merge branch 'fix/two-column-scroll-layout' into dev Human 8 days ago
sha256:408916fc5973ba59c6e4eebaa80ebdcc801c0a63205651e25009d11548f79454 chore: bump version to 0.2.0.dev2 — nightly.2, matching muse Sonnet 4.6 patch 11 days ago
sha256:d035733f21ccff27735fddebfbbe0ed24565a32a22db8de5885402262671ecd2 chore: bump version to 0.2.0rc15 for musehub#113 fix release Sonnet 4.6 patch 14 days ago
sha256:0032d6cfa33bc3c8367436ad768e7dd0e339b4332153160247da8266cb5fa352 Merge branch 'task/version-tags-phase3-server' into dev Human 17 days ago
sha256:4669620efda9ff41c55bdefd1f7bfe1c239d468428744c84ead9957e5a003a53 merge: rescue snapshot-recovery hardening (c00aa21d) into d… Opus 4.8 minor 29 days ago
sha256:a59da49c4611b970fc4b6ae48678ce4943261c213a07ddbd73ce9201df869b4a fix: remove false-positive proposal_comments index drop fro… Sonnet 4.6 patch 33 days ago
sha256:0a240d6dbff234f07d98a28a4a9a68db702f3f9ff9260196f24219bdb1c0b6f3 feat: render markdown mists as HTML with heading anchor links Sonnet 4.6 patch 34 days ago
sha256:24a7d47486ebc4ebd1832830580e177ec6f877b48dced8c000e198cdec4ce9d6 Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump … Human 35 days ago
sha256:b9ff931d147e0114a1f17060f415b89ed551c170a91ff226c70437aa5c85f9ee Merge 'task/bump-version-rc12' into 'dev' — proposal: Bump … Human 35 days ago
sha256:d1122d21e73471879b460037b22c0b50fded7c423444a176f248428f75dac39c Merge 'task/fix-issue-pagination-cursor' into 'dev' — propo… Human 35 days ago
sha256:01e18975e73d2b3cd5b6db7929c895bef9aa6e0d4391dc5b2adfc548b41318dd Merge 'feat/adding-debug-logs-to-staging' into 'dev' — prop… Human 35 days ago
sha256:6b1949fc2797ca4c1936a637a4cbfec828ef56cf52398a2e74ca3c4f494e728f fix: use wire_bytes not mpack_bytes_raw in compute_object_b… Sonnet 4.6 patch 48 days ago
sha256:b99f2455dc346966d040133f5203297e6e3ef5803a93728a2c30568d0a0f7583 rename: delta_add → delta_upsert across wire format, models… Sonnet 4.6 patch 50 days ago