Skip to content

Commit ae45f41

Browse files
Merge pull request #271 from finos/fix-reduce-memory-usage
Fix reduce memory usage
2 parents 3e53420 + dade1a4 commit ae45f41

10 files changed

Lines changed: 365 additions & 124 deletions

File tree

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ Thumbs.db
4646
# Python
4747
__pycache__
4848
.pyenv/
49+
.profile-pyenv/
4950
.pytest_cache/
5051
.pydevproject
5152
# checkstyle
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
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()
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
#!/bin/bash
2+
#
3+
# Copyright (c) 2023-2026 CLOUDRISK Limited and FT Advisory LLC
4+
# SPDX-License-Identifier: Apache-2.0
5+
#
6+
# Profile CDM import memory usage.
7+
#
8+
# Creates (or reuses) a lightweight venv with the pre-built CDM wheel and
9+
# psutil, then runs profile_import.py.
10+
#
11+
# Usage (from project root):
12+
# python-test/cdm-tests/run_memory_profile.sh [options]
13+
#
14+
# Options:
15+
# -r, --reuse-env Reuse existing venv (skip install)
16+
# --top N Show top N allocating entries (passed to profile_import.py)
17+
# -h, --help Show this help
18+
19+
set -euo pipefail
20+
21+
MY_PATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
22+
PROJECT_ROOT="$( cd "$MY_PATH/../.." && pwd )"
23+
VENV_DIR="$MY_PATH/.profile-pyenv"
24+
WHEEL_DIR="$PROJECT_ROOT/target/python-cdm"
25+
PROFILE_SCRIPT="$MY_PATH/profile_import.py"
26+
27+
REUSE_ENV=0
28+
TOP_ARG=""
29+
30+
while [[ $# -gt 0 ]]; do
31+
case "$1" in
32+
-r|--reuse-env) REUSE_ENV=1; shift ;;
33+
--top) TOP_ARG="--top $2"; shift 2 ;;
34+
-h|--help)
35+
sed -n '/^# Options:/,/^[^#]/{/^# /{ s/^# //; p }}' "$0"
36+
exit 0
37+
;;
38+
*) echo "Unknown option: $1"; exit 1 ;;
39+
esac
40+
done
41+
42+
# ── locate wheel ────────────────────────────────────────────────────────────
43+
CDM_WHEEL=$(ls "$WHEEL_DIR"/*.whl 2>/dev/null | head -1)
44+
if [[ -z "$CDM_WHEEL" ]]; then
45+
echo "ERROR: no CDM wheel found in $WHEEL_DIR"
46+
echo "Run the CDM build first:"
47+
echo " python-test/cdm-tests/setup/build_cdm.sh ..."
48+
exit 1
49+
fi
50+
echo "Using wheel: $CDM_WHEEL"
51+
52+
# ── create or reuse venv ─────────────────────────────────────────────────────
53+
if [[ $REUSE_ENV -eq 0 || ! -d "$VENV_DIR" ]]; then
54+
echo "Creating venv at $VENV_DIR"
55+
python3 -m venv "$VENV_DIR"
56+
source "$VENV_DIR/bin/activate"
57+
pip install --quiet --upgrade pip
58+
pip install --quiet "$CDM_WHEEL"
59+
pip install --quiet psutil
60+
else
61+
echo "Reusing existing venv at $VENV_DIR"
62+
source "$VENV_DIR/bin/activate"
63+
fi
64+
65+
# ── run profiler ─────────────────────────────────────────────────────────────
66+
echo ""
67+
python "$PROFILE_SCRIPT" $TOP_ARG
68+
69+
deactivate

src/main/java/com/regnosys/rosetta/generator/python/PythonCodeGenerator.java

Lines changed: 64 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,55 @@ private void partitionClasses(PythonCodeGeneratorContext context) {
350350
}
351351
}
352352
}
353+
354+
// Reverse-promote to bundled: any own-type standalone that is a direct supertype of a
355+
// bundled own-type is pulled into the bundle. Python evaluates base-class expressions
356+
// at class-definition time, so a standalone supertype would require an inline import
357+
// placed mid-file just before its subclass — creating scattered imports. Bundling it
358+
// instead lets the subclass use the flattened bundle name with no import needed.
359+
// Only supertype edges are considered; attribute-type edges are not (those annotations
360+
// are lazy strings under PEP 563 and are handled by the consolidated deferred section).
361+
// The fixpoint propagates up inheritance chains: if B (bundled) extends A (standalone)
362+
// extends X (standalone), A is promoted first, then X is promoted in the next iteration.
363+
boolean anyReversePromotion = true;
364+
while (anyReversePromotion) {
365+
anyReversePromotion = false;
366+
for (String cls : new ArrayList<>(ownTypes)) {
367+
if (standaloneClasses.contains(cls)) {
368+
continue; // cls is standalone — only check bundled types
369+
}
370+
String parentFqn = superTypes.get(cls);
371+
if (parentFqn == null || !ownTypes.contains(parentFqn)) {
372+
continue; // no own-namespace supertype
373+
}
374+
if (standaloneClasses.contains(parentFqn)) {
375+
standaloneClasses.remove(parentFqn);
376+
anyReversePromotion = true;
377+
LOGGER.debug("Reverse-promoted {} to bundled (direct supertype of bundled {})", parentFqn, cls);
378+
}
379+
}
380+
}
381+
382+
// Second forward-promotion pass: reverse-promotion may have newly bundled some types
383+
// whose standalone children now also need bundling (to inherit Phase 2/3 treatment).
384+
anyPromotion = true;
385+
while (anyPromotion) {
386+
anyPromotion = false;
387+
for (String cls : new ArrayList<>(standaloneClasses)) {
388+
if (!ownTypes.contains(cls)) {
389+
continue;
390+
}
391+
String parentFqn = superTypes.get(cls);
392+
if (parentFqn == null || !ownTypes.contains(parentFqn)) {
393+
continue;
394+
}
395+
if (!standaloneClasses.contains(parentFqn)) {
396+
standaloneClasses.remove(cls);
397+
anyPromotion = true;
398+
LOGGER.debug("Promoted {} to bundled (parent {} newly bundled by reverse-promotion)", cls, parentFqn);
399+
}
400+
}
401+
}
353402
}
354403

355404
private Map<String, CharSequence> processDAG(
@@ -376,10 +425,21 @@ private Map<String, CharSequence> processDAG(
376425
headerResult.standaloneSupertypesOfBundled(),
377426
dataObjectsWriter, functionsWriter, annotationUpdateWriter, pendingRebuilds, result);
378427

379-
// Add deferred standalone imports into the rebuild graph so they are ordered
380-
// correctly relative to bundled classes: a standalone type S must rebuild before any
381-
// bundled class B whose Phase 2 annotations reference S, and S itself must rebuild
382-
// after the bundled types it depends on.
428+
// Phase 3: model_rebuild(force=True) calls emitted in dependency order.
429+
//
430+
// These are required despite defer_build=True on every bundled class. Pydantic
431+
// builds None-typed placeholder schemas eagerly at class-definition time (None is
432+
// a trivially resolvable type), so the deferred-build flag alone does not prevent
433+
// the wrong schema from being used. model_rebuild(force=True) forces Pydantic to
434+
// re-read the Phase 2-updated __annotations__ and build the correct schema.
435+
//
436+
// defer_build=True still saves memory/time: because it prevents any intermediate
437+
// schema compilation during the class-definition phase, all cyclic types are fully
438+
// defined by the time Phase 3 runs, and Pydantic's schema builder can resolve
439+
// cross-type references in a single pass (~4× faster than without defer_build).
440+
//
441+
// Standalone classes with deferred imports are integrated into the rebuild graph
442+
// so they are ordered correctly relative to bundled classes.
383443
integrateStandaloneRebuilds(context, headerResult.deferredStandaloneImports(), pendingRebuilds);
384444

385445
String rebuildContent = emitRebuildCallsInOrder(pendingRebuilds, context);

src/main/java/com/regnosys/rosetta/generator/python/object/PythonModelObjectGenerator.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -311,6 +311,16 @@ private String generateBody(Data rc, PythonCodeGeneratorContext context, boolean
311311
writer.appendLine("class " + classNameDefinition + "(" + superClassName + "):");
312312
writer.indent();
313313

314+
// Bundled classes defer initial schema compilation until after all class bodies are
315+
// defined. By the time Phase 3 model_rebuild(force=True) runs, all cyclic types are
316+
// fully available, so Pydantic can build schemas ~4× faster (~1.8 GB / ~5s for CDM
317+
// vs ~7.9 GB / ~17s without defer_build). Phase 3 is still required: Pydantic
318+
// builds None-typed placeholder schemas eagerly even with defer_build=True, so an
319+
// explicit model_rebuild is needed to pick up Phase 2 annotation updates.
320+
if (!isStandalone) {
321+
writer.appendLine("model_config = ConfigDict(defer_build=True)");
322+
}
323+
314324
String metaData = getClassMetaDataString(rc);
315325
if (!metaData.isEmpty()) {
316326
writer.appendBlock(metaData);

src/main/java/com/regnosys/rosetta/generator/python/util/PythonCodeGeneratorUtil.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ public static String createImports() {
8080
from decimal import Decimal
8181
from typing import Annotated, Optional
8282
83-
from pydantic import Field, validate_call, InstanceOf
83+
from pydantic import ConfigDict, Field, validate_call, InstanceOf
8484
8585
from rune.runtime.base_data_class import BaseDataClass
8686
from rune.runtime.cow import rune_cow, rune_unwrap

src/test/java/com/regnosys/rosetta/generator/python/PythonGeneratorTestUtils.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,4 +147,9 @@ public void assertBundleContainsExpectedString(String model, String expectedStri
147147
String allFiles = generatePythonAndExtractBundle(model);
148148
assertGeneratedContainsExpectedString(allFiles, expectedString);
149149
}
150+
151+
public void assertBundleDoesNotContain(String model, String unexpectedString) {
152+
String allFiles = generatePythonAndExtractBundle(model);
153+
assertGeneratedDoesNotContain(allFiles, unexpectedString);
154+
}
150155
}

0 commit comments

Comments
 (0)