Shrink live stats memory + harden profiler self-exclusion (#1043)

Two in-process RAM reductions and two correctness tightens identified
while auditing the overhead of the stitched-stacks timeline:

- __slots__ on PyFrameKey, NativeFrameKey, CombinedStackRun, StackFrame,
  StackStats, and _ResolvedNativeFrame. These types are instantiated
  per CPU / malloc sample and aggregated into long-lived dicts and the
  timeline list. Dropping the per-instance __dict__ takes PyFrameKey
  from 664 → 192 bytes and NativeFrameKey from 424 → 72 bytes. Manual
  slots tuples rather than slots=True because 3.8 is still supported.

- Per-process frame-key intern caches in scalene_utility (add_stack /
  add_combined_stack / add_async_await_run) and scalene_memory_profiler
  (malloc-sample StackFrame construction). Before this, each sample
  constructed fresh dataclass instances that compared equal; only the
  combined_stacks dict de-duplicated them, and the timeline retained
  all the redundant objects for the lifetime of the run. Caches are
  cleared in Scalene.start() so repeated magic-cell profiling sessions
  don't carry state across runs.

- TraceConfig now accepts the canonical path of the installed scalene
  package (from os.path.realpath(os.path.dirname(scalene.__file__)))
  and uses it as an absolute-path prefix check in should_trace, in
  addition to the legacy "immediate parent directory literally named
  scalene" heuristic. The prefix check is robust to vendoring, editable
  installs under symlinks, and any packaging layout that doesn't name
  the directory "scalene". The parent-name check stays as a defensive
  fallback. register_files_to_profile / setup_trace_config take the
  path as an optional 4th arg to preserve the old call signature.

- Scalene.stop() now always flips p_scalene_done = true, not only when
  args.memory is set, so any allocation made while we serialize the
  profile to JSON (or via an atypical shutdown path such as
  term_signal_handler) is guaranteed not to land in the profile.
  set_scalene_done_{true,false} in pywhere.cpp become no-ops when
  libscalene isn't preloaded (CPU-only runs), so the unconditional
  flip is safe regardless of mode.

Measured on big_workload.py (10s, 3 threads, --memory --stacks):

  ScaleneStatistics total     1616 KB -> 215 KB   -86.7%
  combined_stacks_timeline    1544 KB -> 123 KB   -92.0%
  combined_stacks              243 KB ->  23 KB   -90.5%
  avg CombinedStackRun        4688 B  -> 1920 B   -59.0%

Tests: 397 pytest passed + 19 skipped (full suite), ruff clean,
mypy clean on all 63 source files. tests/test_coverup_78.py updated
for the new 4-arg register_files_to_profile signature.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: emeryberger

scalene/scalene_json.py

@@ -65,6 +65,7 @@ class _ResolvedNativeFrame:
     pure cost here.
     """
 
+    __slots__ = ("module", "symbol", "ip", "offset")
     module: str
     symbol: str
     ip: int

scalene/scalene_memory_profiler.py

@@ -24,6 +24,7 @@
     ScaleneStatistics,
     StackFrame,
 )
+from scalene.scalene_utility import intern_stack_frame
 
 
 class ScaleneMemoryProfiler:
@@ -158,7 +159,7 @@ def process_malloc_free_samples(
                         except ValueError:
                             continue
                         frames.append(
-                            StackFrame(
+                            intern_stack_frame(
                                 filename=fn,
                                 function_name="",
                                 line_number=ln,

scalene/scalene_profiler.py

@@ -118,6 +118,7 @@
 from scalene.scalene_tracing import ScaleneTracing
 from scalene.scalene_utility import (
     add_stack,
+    clear_intern_caches,
     compute_frames_to_record,
     drain_native_stacks,
     enter_function_meta,
@@ -1203,6 +1204,11 @@ def start() -> None:
                 file=sys.stderr,
             )
             sys.exit(1)
+        # Drop any frame-key intern cache carried over from a prior session
+        # (Jupyter magic cells, repeated profile runs in the same process,
+        # etc.). Must happen before we start sampling so the cache reflects
+        # only the current run.
+        clear_intern_caches()
         Scalene.__stats.start_clock()
         Scalene.enable_signals()
         Scalene.__start_time = time.monotonic_ns()
@@ -1230,10 +1236,20 @@ def stop() -> None:
         if Scalene.__args.async_profile:
             ScaleneAsync.disable()
 
-        if Scalene.__args.memory:
+        # Flip p_scalene_done back to true so the allocator interposer
+        # stops recording — in particular, allocations made while we
+        # serialize the profile to JSON must not land in the profile
+        # itself. Called unconditionally (not gated on --memory) so the
+        # invariant "no allocations are tracked after stop()" holds even
+        # on code paths that exit through term_signal_handler or other
+        # non-standard shutdown routes. When libscalene isn't preloaded
+        # (CPU-only mode), set_scalene_done_true is a no-op by design.
+        try:
             from scalene import pywhere  # type: ignore
 
             pywhere.set_scalene_done_true()
+        except ImportError:
+            pass
 
         Scalene._disable_signals()
         Scalene.__stats.stop_clock()
@@ -1698,14 +1714,31 @@ def main() -> None:
     @staticmethod
     def _register_files_to_profile() -> None:
         """Tells the pywhere module, which tracks memory, which files to profile."""
+        import scalene as _scalene_pkg
         from scalene import pywhere  # type: ignore
 
         profile_only_list = Scalene.__args.profile_only.split(",")
 
+        # Pass the canonical path of the installed scalene package so the C++
+        # TraceConfig can exclude Scalene-internal frames via an absolute-path
+        # prefix match, not just a parent-directory-name heuristic. The
+        # fallback "scalene" parent-directory check stays in place; this is
+        # additive belt-and-suspenders.
+        scalene_pkg_path: str | None = None
+        try:
+            scalene_pkg_file = getattr(_scalene_pkg, "__file__", None)
+            if scalene_pkg_file:
+                scalene_pkg_path = os.path.realpath(
+                    os.path.dirname(scalene_pkg_file)
+                )
+        except (TypeError, ValueError, OSError):
+            scalene_pkg_path = None
+
         pywhere.register_files_to_profile(
             list(Scalene.__files_to_profile) + profile_only_list,
             Scalene.__program_path,
             Scalene.__args.profile_all,
+            scalene_pkg_path,
         )
 
     @staticmethod

scalene/scalene_statistics.py

@@ -42,8 +42,15 @@ class PyFrameKey:
     Frozen so instances are hashable and usable as parts of dict keys
     (specifically inside CombinedStackKey tuples used to key
     ``ScaleneStatistics.combined_stacks``).
+
+    __slots__ keeps each instance at ~200 bytes instead of ~650 bytes
+    (no per-instance __dict__). At the scale of the stitched-stacks
+    timeline this is the single biggest in-process memory win. Manual
+    slots tuple rather than ``slots=True`` because Scalene supports
+    Python 3.8.
     """
 
+    __slots__ = ("filename", "function", "line")
     filename: str
     function: str
     line: int
@@ -56,6 +63,7 @@ class NativeFrameKey:
     in the signal handler. Frozen for the same reason as PyFrameKey.
     """
 
+    __slots__ = ("ip",)
     ip: int
 
 
@@ -84,8 +92,14 @@ class CombinedStackRun:
         stack_key: The stitched stack — same shape as the keys of
             ``ScaleneStatistics.combined_stacks``.
         count: Number of consecutive CPU samples in this run.
+
+    Non-frozen so ``count`` can be incremented in place when the next
+    sample hits the same stack. __slots__ drops the per-run overhead
+    from ~590 bytes to ~150 bytes, which matters because the timeline
+    can hold up to ``combined_stacks_timeline_max_runs`` (100K) entries.
     """
 
+    __slots__ = ("timestamp", "stack_key", "count")
     timestamp: float
     stack_key: CombinedStackKey
     count: int
@@ -152,7 +166,16 @@ def __eq__(self, other: object) -> bool:
 
 
 class StackFrame:
-    """Represents a single frame in the stack."""
+    """Represents a single frame in the stack.
+
+    __slots__ drops each instance from ~660 bytes to ~80 bytes by
+    eliminating the per-instance __dict__. StackFrame is the element
+    type of ``ScaleneStatistics.memory_stacks`` keys, where every
+    malloc sample that flows through ``process_malloc_free_samples``
+    constructs fresh frames before the dict dedupes by equality.
+    """
+
+    __slots__ = ("filename", "function_name", "line_number")
 
     def __init__(self, filename: str, function_name: str, line_number: int) -> None:
         self.filename = filename
@@ -178,6 +201,8 @@ def __eq__(self, other: Any) -> bool:
 class StackStats:
     """Represents statistics for a stack."""
 
+    __slots__ = ("count", "python_time", "c_time", "cpu_samples")
+
     def __init__(
         self, count: int, python_time: float, c_time: float, cpu_samples: float
     ) -> None:

scalene/scalene_utility.py

@@ -62,6 +62,91 @@
     _native_unwind_available = False
 
 
+# Per-process intern caches for stitched-stack frame keys.
+#
+# Every CPU sample that captures a stitched stack constructs a fresh tuple of
+# PyFrameKey / NativeFrameKey instances. Without interning, two consecutive
+# samples on the same line allocate two distinct frame objects that compare
+# equal, and only the `combined_stacks` dict de-duplicates them (the one fed
+# into the timeline is retained for the lifetime of the profile). At high
+# sample rates the timeline list holds millions of these redundant objects.
+#
+# These caches ensure that equal frames share a single instance, so a tuple
+# captured twice at the same point is represented by the same Python objects
+# throughout. Combined with __slots__ on the frame dataclasses, this collapses
+# the timeline's per-run memory cost by roughly an order of magnitude.
+#
+# The caches are module-level because they belong to the profiler process (one
+# profile per process) and because add_combined_stack / add_async_await_run
+# are called from the CPU-sample sigqueue thread on a hot path — a single
+# dict lookup is cheaper than constructing a fresh dataclass instance.
+# Cleared on each ``clear_intern_caches()`` call (invoked from Scalene.start()
+# so that multiprocess runs / repeated magic-cell invocations don't leak state
+# across profiling sessions).
+_py_frame_intern: Dict[Tuple[str, str, int], PyFrameKey] = {}
+_native_frame_intern: Dict[int, NativeFrameKey] = {}
+
+
+def _intern_py_frame(filename: str, function: str, line: int) -> PyFrameKey:
+    """Return the shared PyFrameKey for (filename, function, line),
+    constructing and caching a new one on first sight. Not thread-safe by
+    design: callers are single-threaded (the sigqueue worker and the
+    CPU-sample Python handler, which are serialized by the GIL + sigqueue
+    lock). A racy double-insert would still be correctness-safe (both keys
+    hash-equal) but waste a slot.
+    """
+    key = (filename, function, line)
+    cached = _py_frame_intern.get(key)
+    if cached is None:
+        cached = PyFrameKey(filename=filename, function=function, line=line)
+        _py_frame_intern[key] = cached
+    return cached
+
+
+def _intern_native_frame(ip: int) -> NativeFrameKey:
+    """Return the shared NativeFrameKey for ``ip``, constructing one on
+    first sight. Same thread-safety caveat as _intern_py_frame."""
+    cached = _native_frame_intern.get(ip)
+    if cached is None:
+        cached = NativeFrameKey(ip=ip)
+        _native_frame_intern[ip] = cached
+    return cached
+
+
+def clear_intern_caches() -> None:
+    """Drop all cached frame-key instances. Safe to call any time the
+    profiler is between sessions. After a clear, subsequent samples
+    repopulate the caches from scratch."""
+    _py_frame_intern.clear()
+    _native_frame_intern.clear()
+    _stack_frame_intern.clear()
+
+
+# StackFrame intern cache for the memory-stacks path (parallel rationale to
+# the PyFrameKey cache above, but for the memory_stacks dict). Keyed by
+# (filename, function_name, line_number). Populated by ``intern_stack_frame``;
+# callers outside this module (notably scalene_memory_profiler) reuse it to
+# avoid reconstructing equal StackFrame objects for every malloc sample.
+_stack_frame_intern: Dict[Tuple[str, str, int], StackFrame] = {}
+
+
+def intern_stack_frame(
+    filename: str, function_name: str, line_number: int
+) -> StackFrame:
+    """Return a shared StackFrame for (filename, function_name, line_number).
+    Constructs and caches a new one on first sight."""
+    key = (filename, function_name, line_number)
+    cached = _stack_frame_intern.get(key)
+    if cached is None:
+        cached = StackFrame(
+            filename=filename,
+            function_name=function_name,
+            line_number=line_number,
+        )
+        _stack_frame_intern[key] = cached
+    return cached
+
+
 def enter_function_meta(
     frame: FrameType,
     should_trace: Callable[[Filename, str], bool],
@@ -190,14 +275,12 @@ def add_stack(
         if should_trace(Filename(f.f_code.co_filename), f.f_code.co_name):
             stk.insert(
                 0,
-                StackFrame(
-                    filename=str(f.f_code.co_filename),
-                    function_name=str(get_fully_qualified_name(f)),
-                    line_number=(
-                        int(f.f_lineno)
-                        if f.f_lineno is not None
-                        else int(f.f_code.co_firstlineno)
-                    ),
+                intern_stack_frame(
+                    str(f.f_code.co_filename),
+                    str(get_fully_qualified_name(f)),
+                    int(f.f_lineno)
+                    if f.f_lineno is not None
+                    else int(f.f_code.co_firstlineno),
                 ),
             )
         f = f.f_back
@@ -297,10 +380,10 @@ def add_combined_stack(
             )
             py_chain.insert(
                 0,
-                PyFrameKey(
-                    filename=str(f.f_code.co_filename),
-                    function=str(get_fully_qualified_name(f)),
-                    line=line,
+                _intern_py_frame(
+                    str(f.f_code.co_filename),
+                    str(get_fully_qualified_name(f)),
+                    line,
                 ),
             )
         f = f.f_back
@@ -312,7 +395,7 @@ def add_combined_stack(
         # native-entry -> native-leaf, which means reversing the
         # leaf-first tuple from the unwinder.
         native_segment: Tuple[NativeFrameKey, ...] = tuple(
-            NativeFrameKey(ip=ip) for ip in reversed(native_stk)
+            _intern_native_frame(ip) for ip in reversed(native_stk)
         )
         key: CombinedStackKey = py_chain_tuple + native_segment
         combined_stacks[key] += 1
@@ -398,10 +481,10 @@ def add_async_await_run(
                     else (cfunc or "<coroutine>")
                 )
                 py_frames.append(
-                    PyFrameKey(
-                        filename=str(cf_filename),
-                        function=func_label,
-                        line=int(cl),
+                    _intern_py_frame(
+                        str(cf_filename),
+                        func_label,
+                        int(cl),
                     )
                 )
 
@@ -414,10 +497,10 @@ def add_async_await_run(
                 continue
             func_name: str = task.func_name or "<await>"
             py_frames.append(
-                PyFrameKey(
-                    filename=str(filename),
-                    function=f"[await] {func_name}{task_suffix}",
-                    line=int(task.lineno),
+                _intern_py_frame(
+                    str(filename),
+                    f"[await] {func_name}{task_suffix}",
+                    int(task.lineno),
                 )
             )