Spaces:
Configuration error
Configuration error
File size: 4,524 Bytes
3cbfa3c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """Lightweight RSS Memory Tracker for profiling repository analysis pipelines.
Provides cross-platform Resident Set Size (RSS) monitoring without heavy profiling
dependencies. Designed for high-frequency low-overhead telemetry in constrained
environments (such as 512 MB instances).
"""
from __future__ import annotations
import logging
import os
import sys
import time
from typing import Any, Optional
logger = logging.getLogger(__name__)
def get_current_rss_mb() -> float:
"""Return current process Resident Set Size (RSS) in MiB.
Resolution strategy:
1. psutil (cross-platform, fast C bindings)
2. Linux /proc/self/statm (zero overhead on Linux/Render)
3. Python standard library `resource` (Unix/macOS fallback)
4. 0.0 fallback if unavailable
"""
# 1. psutil
try:
import psutil
return round(psutil.Process().memory_info().rss / (1024 * 1024), 2)
except Exception:
pass
# 2. Linux /proc/self/statm
try:
with open("/proc/self/statm", "r") as f:
fields = f.read().split()
if len(fields) >= 2:
rss_pages = int(fields[1])
page_size = os.sysconf("SC_PAGE_SIZE")
return round((rss_pages * page_size) / (1024 * 1024), 2)
except Exception:
pass
# 3. resource.getrusage (Unix / macOS)
try:
import resource
ru = resource.getrusage(resource.RUSAGE_SELF)
if sys.platform == "darwin":
return round(ru.ru_maxrss / (1024 * 1024), 2)
# On Linux ru_maxrss is in KiB
return round(ru.ru_maxrss / 1024, 2)
except Exception:
pass
return 0.0
class MemoryTracker:
"""Tracks progression and logs process RSS at defined pipeline boundaries."""
def __init__(
self,
repo_name: str = "",
log_interval_files: int = 50,
log_interval_chunks: int = 200,
logger_instance: Optional[logging.Logger] = None,
) -> None:
self.repo_name = repo_name
self.log_interval_files = log_interval_files
self.log_interval_chunks = log_interval_chunks
self.log = logger_instance or logger
self.start_time = time.time()
self.files_processed: int = 0
self.bytes_processed: int = 0
self.chunks_processed: int = 0
self.embeddings_indexed: int = 0
self._last_logged_files: int = 0
self._last_logged_chunks: int = 0
def log_phase(self, phase: str, **kwargs: Any) -> float:
"""Log memory metrics for a given execution boundary.
Format:
MEMORY phase=<phase> rss_mb=<rss_mb> files_processed=<n> bytes_processed=<n>
chunks_processed=<n> embeddings_indexed=<n> elapsed_seconds=<sec> [key=val ...]
"""
rss_mb = get_current_rss_mb()
elapsed_seconds = round(time.time() - self.start_time, 2)
parts = [
f"MEMORY phase={phase}",
f"rss_mb={rss_mb:.1f}",
f"files_processed={self.files_processed}",
f"bytes_processed={self.bytes_processed}",
f"chunks_processed={self.chunks_processed}",
f"embeddings_indexed={self.embeddings_indexed}",
f"elapsed_seconds={elapsed_seconds:.1f}",
]
if self.repo_name:
parts.append(f"repo={self.repo_name}")
for k, v in kwargs.items():
if isinstance(v, float):
parts.append(f"{k}={v:.1f}")
else:
parts.append(f"{k}={v}")
log_line = " ".join(parts)
self.log.info(log_line)
return rss_mb
def record_file(self, num_bytes: int = 0) -> None:
"""Record an extracted file and log memory periodically."""
self.files_processed += 1
self.bytes_processed += num_bytes
if self.files_processed - self._last_logged_files >= self.log_interval_files:
self._last_logged_files = self.files_processed
self.log_phase("during_extraction")
def record_chunk(self, count: int = 1) -> None:
"""Record chunked tokens and log memory periodically."""
self.chunks_processed += count
if self.chunks_processed - self._last_logged_chunks >= self.log_interval_chunks:
self._last_logged_chunks = self.chunks_processed
self.log_phase("during_chunking")
def record_embeddings_indexed(self, count: int) -> None:
"""Record indexed embeddings count."""
self.embeddings_indexed += count
|