|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# |
| 3 | +# Copyright (c) 2023-2026 CLOUDRISK Limited and FT Advisory LLC |
| 4 | +# SPDX-License-Identifier: Apache-2.0 |
| 5 | +# |
| 6 | +""" |
| 7 | +Profile memory usage of importing CDM TradeState. |
| 8 | +
|
| 9 | +Outputs: |
| 10 | + - Process RSS before/after import (requires psutil) |
| 11 | + - tracemalloc allocation breakdown by file and by top-level package |
| 12 | + - Count and size of Pydantic BaseModel subclasses loaded |
| 13 | + - Number of modules imported |
| 14 | +
|
| 15 | +Usage (from project root, with CDM venv activated): |
| 16 | + python python-test/cdm-tests/profile_import.py [--top N] |
| 17 | +
|
| 18 | +Options: |
| 19 | + --top N Show top N allocating files (default: 30) |
| 20 | +""" |
| 21 | + |
| 22 | +import argparse |
| 23 | +import gc |
| 24 | +import sys |
| 25 | +import time |
| 26 | + |
| 27 | +import psutil |
| 28 | + |
| 29 | +_PROC = psutil.Process() |
| 30 | + |
| 31 | + |
| 32 | +def rss_mb() -> float: |
| 33 | + return _PROC.memory_info().rss / 1024 / 1024 |
| 34 | + |
| 35 | + |
| 36 | +def _count_pydantic_models(): |
| 37 | + try: |
| 38 | + import pydantic |
| 39 | + base = pydantic.BaseModel |
| 40 | + except ImportError: |
| 41 | + return 0 |
| 42 | + count = 0 |
| 43 | + for obj in gc.get_objects(): |
| 44 | + try: |
| 45 | + if isinstance(obj, type) and issubclass(obj, base) and obj is not base: |
| 46 | + count += 1 |
| 47 | + except TypeError: |
| 48 | + pass |
| 49 | + return count |
| 50 | + |
| 51 | + |
| 52 | +def main(): |
| 53 | + parser = argparse.ArgumentParser(description=__doc__, |
| 54 | + formatter_class=argparse.RawDescriptionHelpFormatter) |
| 55 | + parser.add_argument("--top", type=int, default=15, |
| 56 | + help="Top N slowest model_rebuild calls to show (default: 15)") |
| 57 | + args = parser.parse_args() |
| 58 | + |
| 59 | + sep = "=" * 72 |
| 60 | + |
| 61 | + # ── baseline ────────────────────────────────────────────────────────── |
| 62 | + gc.collect() |
| 63 | + modules_before = set(sys.modules.keys()) |
| 64 | + rss_before = rss_mb() |
| 65 | + |
| 66 | + # Instrument model_rebuild to capture per-class timing |
| 67 | + import pydantic |
| 68 | + rebuild_times = [] |
| 69 | + _orig = pydantic.BaseModel.model_rebuild.__func__ |
| 70 | + |
| 71 | + def _timed(cls, **kwargs): |
| 72 | + t0 = time.perf_counter() |
| 73 | + result = _orig(cls, **kwargs) |
| 74 | + rebuild_times.append((time.perf_counter() - t0, cls.__name__)) |
| 75 | + return result |
| 76 | + |
| 77 | + pydantic.BaseModel.model_rebuild = classmethod(_timed) |
| 78 | + |
| 79 | + # ── import ──────────────────────────────────────────────────────────── |
| 80 | + print("Importing TradeState …", flush=True) |
| 81 | + t_start = time.perf_counter() |
| 82 | + from finos.cdm.event.common.TradeState import TradeState # noqa: F401 |
| 83 | + t_total = time.perf_counter() - t_start |
| 84 | + print("Import done.", flush=True) |
| 85 | + |
| 86 | + # ── first-use (model_validate on minimal data) ───────────────────────── |
| 87 | + rss_pre_validate = rss_mb() |
| 88 | + rebuild_times_before_validate = len(rebuild_times) |
| 89 | + _first_use_data = {"trade": None, "state": None, "resetHistory": None, |
| 90 | + "transferHistory": None, "observationHistory": None} |
| 91 | + t_validate_start = time.perf_counter() |
| 92 | + try: |
| 93 | + TradeState.model_validate(_first_use_data) |
| 94 | + except Exception: |
| 95 | + pass # validation error is fine — we only care about schema-build cost |
| 96 | + t_validate = time.perf_counter() - t_validate_start |
| 97 | + rss_post_validate = rss_mb() |
| 98 | + rebuild_calls_during_validate = len(rebuild_times) - rebuild_times_before_validate |
| 99 | + |
| 100 | + # ── snapshot ────────────────────────────────────────────────────────── |
| 101 | + gc.collect() |
| 102 | + rss_after = rss_mb() |
| 103 | + modules_after = set(sys.modules.keys()) |
| 104 | + model_count = _count_pydantic_models() |
| 105 | + |
| 106 | + total_rebuild = sum(t for t, _ in rebuild_times) |
| 107 | + |
| 108 | + # ── report ──────────────────────────────────────────────────────────── |
| 109 | + print(f"\n{sep}") |
| 110 | + print("MEMORY SUMMARY") |
| 111 | + print(sep) |
| 112 | + print(f" RSS before import : {rss_before:.1f} MB") |
| 113 | + print(f" RSS after import : {rss_after:.1f} MB") |
| 114 | + print(f" RSS delta : {rss_after - rss_before:.1f} MB") |
| 115 | + print(f" Modules imported : {len(modules_after - modules_before)}") |
| 116 | + print(f" Pydantic models loaded: {model_count}") |
| 117 | + |
| 118 | + print(f"\n{sep}") |
| 119 | + print("TIMING SUMMARY") |
| 120 | + print(sep) |
| 121 | + print(f" Total import time : {t_total:.2f}s") |
| 122 | + print(f" Time in model_rebuild : {total_rebuild:.2f}s ({total_rebuild/t_total*100:.0f}% of total)") |
| 123 | + print(f" model_rebuild calls : {len(rebuild_times)}") |
| 124 | + print(f" Average per rebuild : {total_rebuild/len(rebuild_times)*1000:.1f}ms" if rebuild_times else "") |
| 125 | + print(f"\n -- First use (model_validate) --") |
| 126 | + print(f" First-use time : {t_validate*1000:.0f}ms") |
| 127 | + print(f" RSS before first use : {rss_pre_validate:.1f} MB") |
| 128 | + print(f" RSS after first use : {rss_post_validate:.1f} MB") |
| 129 | + print(f" RSS delta (first use) : {rss_post_validate - rss_pre_validate:.1f} MB") |
| 130 | + print(f" model_rebuild during : {rebuild_calls_during_validate} calls") |
| 131 | + |
| 132 | + print(f"\n{sep}") |
| 133 | + print(f"TOP {args.top} SLOWEST model_rebuild CALLS") |
| 134 | + print(sep) |
| 135 | + for elapsed, name in sorted(rebuild_times, reverse=True)[: args.top]: |
| 136 | + print(f" {elapsed*1000:7.1f}ms {name}") |
| 137 | + |
| 138 | + print(f"\n{sep}\n") |
| 139 | + |
| 140 | + |
| 141 | +if __name__ == "__main__": |
| 142 | + main() |
0 commit comments