diff --git a/CHANGES.md b/CHANGES.md index 7ea2cccde299..32d31ef0e3d0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -71,6 +71,7 @@ * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). * (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)). +* (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/sdks/python/apache_beam/runners/worker/worker_status.py b/sdks/python/apache_beam/runners/worker/worker_status.py index 1d54a3ee1764..28917c6d949b 100644 --- a/sdks/python/apache_beam/runners/worker/worker_status.py +++ b/sdks/python/apache_beam/runners/worker/worker_status.py @@ -19,6 +19,7 @@ import gc import logging +import os import queue import sys import threading @@ -108,6 +109,135 @@ def thread_dump(thread_prefix=None): return '\n'.join(all_traces) +def _process_memory_stats(): + """Process-level memory usage (RSS), to gauge total footprint over time.""" + lines = ['--- Process memory ---'] + try: + import resource + usage = resource.getrusage(resource.RUSAGE_SELF) + # ru_maxrss is reported in kilobytes on Linux but in bytes on macOS/BSD. + if sys.platform == 'darwin': + peak_bytes = usage.ru_maxrss + else: + peak_bytes = usage.ru_maxrss * 1024 + lines.append('peak RSS (ru_maxrss): %d bytes' % peak_bytes) + except Exception as e: # pylint: disable=broad-except + lines.append('resource stats unavailable: %s' % e) + try: + # /proc/self/statm reports counts in pages; field 1 is resident set size. + with open('/proc/self/statm') as f: + resident_pages = int(f.read().split()[1]) + lines.append( + 'current RSS (/proc/self/statm): %d bytes' % + (resident_pages * os.sysconf('SC_PAGE_SIZE'))) + except Exception: # pylint: disable=broad-except + # /proc is Linux-only; skip silently on other platforms. + pass + return '\n'.join(lines) + + +def _python_memory_stats(): + """CPython allocator and garbage-collector stats. + + Together with the native-heap stats these help distinguish memory growth in + Python objects from growth on the native (C) heap. + """ + lines = ['--- Python allocations ---'] + try: + lines.append('sys.getallocatedblocks: %d' % sys.getallocatedblocks()) + except Exception as e: # pylint: disable=broad-except + lines.append('sys.getallocatedblocks unavailable: %s' % e) + try: + lines.append('gc.get_count (gen0, gen1, gen2): %s' % (gc.get_count(), )) + for i, stat in enumerate(gc.get_stats()): + lines.append( + 'gc gen%d: collections=%s collected=%s uncollectable=%s' % ( + i, + stat.get('collections'), + stat.get('collected'), + stat.get('uncollectable'))) + except Exception as e: # pylint: disable=broad-except + lines.append('gc stats unavailable: %s' % e) + return '\n'.join(lines) + + +def _glibc_malloc_stats(): + """glibc allocator stats, useful for reasoning about native-heap growth and + fragmentation. + + ``fordblks`` (free space the allocator retains rather than returning to the + OS) growing relative to ``uordblks`` (space in use) is a signal of native + heap fragmentation. Only available with glibc (Linux); degrades gracefully + elsewhere. + """ + lines = ['--- glibc malloc (native heap) ---'] + if not sys.platform.startswith('linux'): + lines.append('unavailable: glibc malloc stats are only collected on Linux.') + return '\n'.join(lines) + try: + import ctypes + + class _MallInfo2(ctypes.Structure): + # Mirrors glibc's ``struct mallinfo2`` (all fields are size_t). + _fields_ = [ + ('arena', ctypes.c_size_t), + ('ordblks', ctypes.c_size_t), + ('smblks', ctypes.c_size_t), + ('hblks', ctypes.c_size_t), + ('hblkhd', ctypes.c_size_t), + ('usmblks', ctypes.c_size_t), + ('fsmblks', ctypes.c_size_t), + ('uordblks', ctypes.c_size_t), + ('fordblks', ctypes.c_size_t), + ('keepcost', ctypes.c_size_t), + ] + + libc = ctypes.CDLL('libc.so.6') + if not hasattr(libc, 'mallinfo2'): + # The older mallinfo() uses int fields that overflow past 2GB and would + # misreport exactly when memory is the concern, so we do not fall back. + lines.append('unavailable: mallinfo2 not found (needs glibc >= 2.33).') + return '\n'.join(lines) + libc.mallinfo2.restype = _MallInfo2 + libc.mallinfo2.argtypes = [] + info = libc.mallinfo2() + lines.append('arena (non-mmapped bytes from sbrk): %d' % info.arena) + lines.append('hblkhd (mmapped bytes): %d' % info.hblkhd) + lines.append('uordblks (in-use bytes): %d' % info.uordblks) + lines.append( + 'fordblks (free bytes retained by allocator): %d' % info.fordblks) + lines.append( + 'keepcost (releasable top-most free bytes): %d' % info.keepcost) + footprint = info.arena + info.hblkhd + if footprint: + lines.append( + 'fragmentation (fordblks / (arena + hblkhd)): %.2f%%' % + (100.0 * info.fordblks / footprint)) + except Exception as e: # pylint: disable=broad-except + lines.append('unavailable: %s' % e) + return '\n'.join(lines) + + +def memory_stats(): + """Collect process, Python and native-heap memory statistics. + + This complements the guppy heap dump with information that helps distinguish + memory growth on the native (C) heap from Python-object allocations, and with + glibc allocator stats that hint at native-heap fragmentation. Every collector + degrades gracefully and never raises when a source is unavailable on the + current platform. + """ + banner = '=' * 10 + ' MEMORY STATS ' + '=' * 10 + sections = [ + banner, + _process_memory_stats(), + _python_memory_stats(), + _glibc_malloc_stats(), + '=' * 30, + ] + return '\n'.join(sections) + + def heap_dump(): """Get a heap dump for the current SDK worker harness. """ banner = '=' * 10 + ' HEAP DUMP ' + '=' * 10 + '\n' @@ -116,7 +246,7 @@ def heap_dump(): else: heap = '%s\n' % hpy().heap() ending = '=' * 30 - return banner + heap + ending + return banner + heap + ending + '\n' + memory_stats() def _state_cache_stats(state_cache: StateCache) -> str: diff --git a/sdks/python/apache_beam/runners/worker/worker_status_test.py b/sdks/python/apache_beam/runners/worker/worker_status_test.py index 88543258250a..8ddd88813ac5 100644 --- a/sdks/python/apache_beam/runners/worker/worker_status_test.py +++ b/sdks/python/apache_beam/runners/worker/worker_status_test.py @@ -28,6 +28,7 @@ from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.worker_status import FnApiWorkerStatusHandler from apache_beam.runners.worker.worker_status import heap_dump +from apache_beam.runners.worker.worker_status import memory_stats from apache_beam.utils import thread_pool_executor from apache_beam.utils.counters import CounterName @@ -231,6 +232,30 @@ def test_skip_heap_dump(self): self.assertTrue( 'Unable to import guppy, the heap dump will be skipped' in result) + @mock.patch('apache_beam.runners.worker.worker_status.hpy', None) + def test_heap_dump_includes_memory_stats_without_guppy(self): + # Memory stats must be emitted even when guppy is unavailable. + result = '%s' % heap_dump() + self.assertIn('MEMORY STATS', result) + + +class MemoryStatsTest(unittest.TestCase): + def test_memory_stats_sections_present(self): + result = memory_stats() + self.assertIn('MEMORY STATS', result) + self.assertIn('Process memory', result) + self.assertIn('Python allocations', result) + self.assertIn('glibc malloc', result) + # Python allocator introspection is available on every platform. + self.assertIn('sys.getallocatedblocks', result) + + @mock.patch('sys.platform', 'darwin') + def test_glibc_stats_graceful_on_non_glibc(self): + # On non-glibc platforms the native-heap section must degrade gracefully + # rather than raise. + result = memory_stats() + self.assertIn('unavailable', result) + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO)