diff --git a/.cargo/config.toml b/.cargo/config.toml index 02ef6843e995..ef43a5d2a177 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -19,3 +19,9 @@ rustflags = ["-C", "target-cpu=neoverse-n1", "-C", "force-frame-pointers=yes", " [env] PCRE2_SYS_STATIC = "1" +# jemalloc (the Python extension's and the router benches' Rust allocator) fixes its page +# size at build time and fails to initialize on a kernel with larger pages. A 64 KiB page size +# works on both 4 KiB and 64 KiB aarch64 kernels (for example, 64K-page kernels on Grace). +# jemalloc 5.3.1's configure already defaults to it for aarch64 Linux; this pins it so that a +# jemalloc update that drops that default cannot make arm64 wheels depend on the build host. +AARCH64_UNKNOWN_LINUX_GNU_JEMALLOC_SYS_WITH_LG_PAGE = "16" diff --git a/Cargo.lock b/Cargo.lock index 3da0c2753829..389fcf4bce77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2454,6 +2454,8 @@ dependencies = [ "serde_json", "sha2 0.10.9", "tempfile", + "tikv-jemalloc-sys", + "tikv-jemallocator", "tokenizers 0.21.4", "tokio", "tokio-util", @@ -10007,6 +10009,26 @@ dependencies = [ "rustc-hash 2.1.2", ] +[[package]] +name = "tikv-jemalloc-sys" +version = "0.7.1+5.3.1-0-g81034ce1f1373e37dc865038e1bc8eeecf559ce8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a2825c78386b4ae0314074867860ba9577875de945f05992c38815cbec327f0" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "249f09e49ab1609436f34c776e84231bead18d6a955f119f939bdc1d847561bd" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.47" diff --git a/container/templates/wheel_builder.Dockerfile b/container/templates/wheel_builder.Dockerfile index 8ea1ac208650..6f9dc77f4a2a 100644 --- a/container/templates/wheel_builder.Dockerfile +++ b/container/templates/wheel_builder.Dockerfile @@ -680,6 +680,8 @@ RUN --mount=type=cache,id=uv-root-{{ context.dynamo.uv_version }},target=/root/. # to match generators/rust.py. Best-effort: unreadable/absent files are skipped # and the generator falls back to canonical SPDX text. cargo's registry lives # under CARGO_HOME and/or the cache-mounted /root/.cargo — scan both. +# tikv-jemalloc-sys compiles its vendored jemalloc (BSD-2-Clause) into the +# _core extension; that license sits in jemalloc/COPYING, below the crate root. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ for src in "${CARGO_HOME}/registry/src" /root/.cargo/registry/src; do \ [ -d "$src" ] || continue; \ @@ -689,6 +691,11 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=shared \ [ -e "$lf" ] || continue; \ mkdir -p "$dest" && cp "$lf" "$dest/" 2>/dev/null || true; \ done; \ + case "$(basename "$crate")" in \ + tikv-jemalloc-sys-*) \ + [ -f "$crate/jemalloc/COPYING" ] && mkdir -p "$dest" \ + && cp "$crate/jemalloc/COPYING" "$dest/jemalloc-COPYING" 2>/dev/null || true ;; \ + esac; \ done; \ done; \ echo "rust license harvest: $(find /opt/dynamo/rust-licenses -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l) crates with license files"; \ diff --git a/docs/fern/pages/kubernetes/operations/performance-tuning.md b/docs/fern/pages/kubernetes/operations/performance-tuning.md index 9e6a1afffde8..0654f1e5ddfb 100644 --- a/docs/fern/pages/kubernetes/operations/performance-tuning.md +++ b/docs/fern/pages/kubernetes/operations/performance-tuning.md @@ -112,12 +112,15 @@ measured traffic. ## Bound frontend host memory growth The frontend allocates host memory per request for tokenization, routing bookkeeping, and response -assembly. With the default glibc allocator, freed memory may be retained in allocator arenas instead -of being returned to the operating system, so frontend resident memory can climb to a high-water -mark under load and may not drop when the load stops. This retention is allocator behavior rather -than a leak. - -To bound the retained footprint, preload jemalloc with decay enabled on the frontend container. See +assembly. Its Rust code allocates from a jemalloc that the Dynamo Python extension links in, which +returns unused memory to the operating system in the background. Its Python and C/C++ code allocate +from glibc by default. glibc may retain freed memory in allocator arenas instead of returning it to +the operating system, so frontend resident memory can climb to a high-water mark under load and may +not drop when the load stops. This retention is allocator behavior rather than a leak. + +To bound the retained footprint, preload jemalloc with background threads on the frontend container, +either with `DYN_FRONTEND_JEMALLOC=1` or through `LD_PRELOAD`. The preload serves the frontend's +Python and C/C++ allocations; its Rust allocations stay on the extension's jemalloc. See [Host memory allocator](../../reference/components/frontend-configuration.mdx#host-memory-allocator) for the `LD_PRELOAD` and `MALLOC_CONF` values and for which container images ship `libjemalloc2`. @@ -125,7 +128,8 @@ To verify the change, run the same workload and track the frontend container's m (`memory.current` on cgroup v2, or `memory.usage_in_bytes` on cgroup v1) across repeated benchmark rounds and through an idle window after the load stops. -With glibc the value typically stays at its peak, while with jemalloc and decay it may -settle toward a floor. Those decay settings do not guarantee that purging occurs during an idle -window; include an idle period in the benchmark to confirm whether memory actually settles for -your workload. +With glibc the value typically stays at its peak. With jemalloc background threads, it should +settle toward a floor during the idle window once the decay time (10 seconds by default) passes. +Include an idle period in the benchmark to confirm that memory settles for your workload, and +record the peak too: the extension's jemalloc can hold freed memory for several seconds after a +burst before its background threads return it. diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 90bc457b1fe7..5643d050ce65 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -721,8 +721,10 @@ curl -s localhost:8000/metrics | grep dynamo_frontend_inter_token_latency_second ## Host memory allocator The frontend performs per-request host-memory allocations. Under high CPU load, allocator -contention can cap frontend throughput before the GPU workers saturate. The frontend does not -preload jemalloc by default. +contention can cap frontend throughput before the GPU workers saturate. The frontend's Rust code +allocates from a jemalloc that the Dynamo Python extension links in, in every process that imports +the extension. The frontend's Python and C/C++ code allocate from glibc unless jemalloc is +preloaded. The frontend does not preload jemalloc by default. Whether the container image ships jemalloc depends on how the image is built. Images built with @@ -734,7 +736,7 @@ dev targets), as well as the standalone frontend entrypoint image do not. Verify frontend image that includes it, or install it from an initContainer into a shared volume mounted by both containers and point `LD_PRELOAD` at the copied library path. Pointing `LD_PRELOAD` at a path that does not exist only prints a loader warning, and the process continues with the default -glibc allocator. +glibc allocator for its Python and C/C++ code. @@ -744,8 +746,19 @@ glibc allocator. volume with an initContainer or use a custom image. + + Preload jemalloc so that it serves the frontend's Python and C/C++ allocations. Accepts `1`, + `true`, or `yes` (case-insensitive). The frontend re-executes itself once with jemalloc added to + `LD_PRELOAD`, so child processes inherit the preload. If the frontend cannot find `libjemalloc`, it + prints a warning and continues with glibc. The Rust allocations stay on the Dynamo Python + extension's jemalloc, so the process then runs two independent jemalloc instances, each with its + own memory and options. + + - Preload jemalloc for the frontend process. The library path depends on the container architecture: + Preload jemalloc for the frontend process's Python and C/C++ allocations. The Rust allocations + stay on the Dynamo Python extension's jemalloc. The library path depends on the container + architecture: ```bash # x86_64 @@ -771,18 +784,51 @@ export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2:$LD_PRELOAD - Configure jemalloc when it is preloaded. The best values depend on the request workload and CPU - allocation. This is an empirical starting point, not a production default: + Configure the preloaded jemalloc, which serves the Python and C/C++ allocations. It does not + configure the Dynamo Python extension's jemalloc; see `_RJEM_MALLOC_CONF`. The best values depend + on the request workload and CPU allocation. This starting point has not been measured for Python + and C/C++ allocations: ```bash - export MALLOC_CONF="narenas:32,tcache:true,lg_tcache_max:15,dirty_decay_ms:5000,muzzy_decay_ms:5000" + export MALLOC_CONF="background_thread:true" ``` + Background threads return unused memory to the operating system off the request path. To return + it sooner after load drops, lower `dirty_decay_ms` from its default of 10000 at some CPU cost. + + Do not use `narenas:32,tcache:true,lg_tcache_max:15,dirty_decay_ms:5000,muzzy_decay_ms:5000`. + It was measured on the KV router indexer while a preloaded jemalloc also served the extension's + Rust allocations: compared with jemalloc's defaults, it lowered write throughput and raised CPU + time per update, while `background_thread:true` removed the update latency tail that both showed. + That result is why the extension's jemalloc enables background threads by default. `tcache:true` + and `lg_tcache_max:15` restate jemalloc 5.3 defaults, and on hosts with more than eight CPUs + `narenas:32` caps arenas below the default of four per CPU. + Validate allocator changes with the same AIPerf workload before and after enabling them. See [Bound frontend host memory growth](../../kubernetes/operations/performance-tuning.md#bound-frontend-host-memory-growth) for the symptom this addresses and how to measure the effect. + + Configure the jemalloc that the Dynamo Python extension links in, which serves the Rust + allocations of every process that imports the extension, including the frontend. It overrides the + built-in default and takes the same options as `MALLOC_CONF`, except the profiling options + (`prof*`): this jemalloc is built without profiling and does not collect statistics (`stats.*`). + jemalloc reads it once, when the extension loads. + + On Linux with glibc, the default starts background threads, named `jemalloc_bg_thd`, that return + unused memory to the operating system: the first when the extension loads, and up to four in + total. Other platforms run without them. Because of that first thread, a process is multi-threaded + as soon as it imports `dynamo`, and Python 3.12 and later emit a `DeprecationWarning` when such a + process calls `os.fork()`. A child process forked without `exec` runs without background threads. + + Background threads lower the memory that the extension's jemalloc keeps while idle. They purge on + their own schedule, though, so resident memory can stay above its idle level for several seconds + after an allocation burst. Check both peak and idle memory against the container's memory limit. + `_RJEM_MALLOC_CONF=background_thread:false` turns the threads off. The allocating threads then + purge unused memory themselves, so a process that goes idle can keep more memory. + + ## HTTP endpoints The frontend exposes the following HTTP endpoints. diff --git a/lib/bench/Cargo.toml b/lib/bench/Cargo.toml index 445a51b8904c..77a15b9a7621 100644 --- a/lib/bench/Cargo.toml +++ b/lib/bench/Cargo.toml @@ -60,7 +60,9 @@ multiturn = [ "dep:rand", ] offline-replay = ["dep:blake3", "dep:dynamo-mocker", "dep:libc"] -router-bench = ["dynamo-kv-router/bench"] +# Router benches link jemalloc as the global allocator, matching the Python +# extension that runs the router in production (`lib/bindings/python`). +router-bench = ["dynamo-kv-router/bench", "dep:tikv-jemallocator", "dep:tikv-jemalloc-sys"] router-metrics = ["dynamo-kv-router/metrics"] router-test-support = [ "router-bench", @@ -169,6 +171,9 @@ serde = { workspace = true } serde_json = { workspace = true } sha2 = { version = "0.10", optional = true } tempfile = { workspace = true, optional = true } +# Same build as the Python extension's allocator (`lib/bindings/python/Cargo.toml`). +tikv-jemalloc-sys = { version = "0.7.1", optional = true } +tikv-jemallocator = { version = "0.7.0", features = ["disable_initial_exec_tls"], optional = true } tokio = { workspace = true } tokio-util = { workspace = true, optional = true } tokenizers = { version = "0.21.4", default-features = false, features = [ @@ -184,3 +189,6 @@ dynamo-kv-router = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] libc = { workspace = true, optional = true } + +[target.'cfg(all(target_os = "linux", target_env = "gnu"))'.dependencies] +tikv-jemallocator = { version = "0.7.0", features = ["background_threads"], optional = true } diff --git a/lib/bench/kv_router/INDEXER_BENCH.md b/lib/bench/kv_router/INDEXER_BENCH.md index 37020d371f90..841fe68453c7 100644 --- a/lib/bench/kv_router/INDEXER_BENCH.md +++ b/lib/bench/kv_router/INDEXER_BENCH.md @@ -27,12 +27,12 @@ application. Preparation consumes and releases the parsed trace, generated replay artifacts, worker timelines, and intermediate payload owners. The timed process retains only the prepared schedule, one flattened query-hash slab, and backend state. -On Linux, the benchmark returns free preparation pages with `malloc_trim` and -waits a fixed five seconds before constructing backend workers. It then -page-touches the schedule once and performs the fixed lookup warm-up before -timing. This explicit quiescence prevents the parallel event-generation phase -and released allocator arenas from affecting the timed backend drain; JSON -records it as `pre_run_quiescence_ms=5000`. +On Linux, the benchmark returns free preparation pages with `malloc_trim` and a +purge of all jemalloc arenas and waits a fixed five seconds before constructing +backend workers. It then page-touches the schedule once and performs the fixed +lookup warm-up before timing. This explicit quiescence prevents the parallel +event-generation phase and released allocator arenas from affecting the timed +backend drain; JSON records it as `pre_run_quiescence_ms=5000`. Outside timing: @@ -89,6 +89,19 @@ the engine re-chunks synthesized prompt tokens at that size. Linux uses absolute `CLOCK_MONOTONIC` sleeps followed by the configured spin. macOS uses a portable sleep-plus-spin timer for correctness tests only. +The router benches (`mooncake_bench`, `active_sequences_bench`, +`approximate_lru_bench`) use jemalloc as the global allocator, built and +configured like the Python extension that runs the router in production +(`lib/bindings/python`): prefixed symbols and, on Linux, +`background_thread:true`, which `_RJEM_MALLOC_CONF` overrides. One difference +remains: the linker relaxes jemalloc's thread-local storage accesses in the +bench executables to the local-exec model, while the extension keeps the slower +global-dynamic model so that Python can `dlopen` it. +The allocator is part of the measured system. Under glibc malloc, CRTC event +workers serialize on arena locks: write capacity stops rising after about 16 +event workers and falls beyond that. Compare allocators or event-worker counts +only between builds and runs that differ in that one variable. + ### Single CRTC trial Select one logical CPU per physical core from the process's actual affinity diff --git a/lib/bench/kv_router/active_sequences_bench.rs b/lib/bench/kv_router/active_sequences_bench.rs index e7c6cd387c4f..1dcd96213528 100644 --- a/lib/bench/kv_router/active_sequences_bench.rs +++ b/lib/bench/kv_router/active_sequences_bench.rs @@ -1,6 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#[path = "jemalloc.rs"] +mod jemalloc; + #[path = "active_sequences_open_loop.rs"] mod active_sequences_open_loop; #[path = "active_sequences_shared.rs"] diff --git a/lib/bench/kv_router/approximate_lru_bench.rs b/lib/bench/kv_router/approximate_lru_bench.rs index 740b24e7c5dc..170c1d8bf798 100644 --- a/lib/bench/kv_router/approximate_lru_bench.rs +++ b/lib/bench/kv_router/approximate_lru_bench.rs @@ -1,6 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#[path = "jemalloc.rs"] +mod jemalloc; + use std::collections::{HashMap, HashSet}; use std::hint::black_box; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}; @@ -1150,6 +1153,7 @@ async fn quiesce(milliseconds: u64) { unsafe { libc::malloc_trim(0); } + jemalloc::purge(); if milliseconds > 0 { tokio::time::sleep(Duration::from_millis(milliseconds)).await; } diff --git a/lib/bench/kv_router/jemalloc.rs b/lib/bench/kv_router/jemalloc.rs new file mode 100644 index 000000000000..2451d2950dcf --- /dev/null +++ b/lib/bench/kv_router/jemalloc.rs @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The router benches' global allocator: jemalloc, built like the Python extension's that runs +//! the router in production (`lib/bindings/python/rust/allocator.rs`). + +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +/// Returns the unused pages of every jemalloc arena to the OS. +#[allow(dead_code, reason = "only the benches that quiesce their heap call it")] +pub fn purge() { + // SAFETY: the name (`arena..purge`) is NUL-terminated, and this + // mallctl requires null value pointers and a zero length. + unsafe { + tikv_jemalloc_sys::mallctl( + c"arena.4096.purge".as_ptr(), + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 0, + ); + } +} diff --git a/lib/bench/kv_router/mooncake_bench.rs b/lib/bench/kv_router/mooncake_bench.rs index bbb4022ac1df..9894a78b88a1 100644 --- a/lib/bench/kv_router/mooncake_bench.rs +++ b/lib/bench/kv_router/mooncake_bench.rs @@ -1,6 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#[path = "jemalloc.rs"] +mod jemalloc; + #[path = "mooncake_open_loop.rs"] mod mooncake_open_loop; #[path = "mooncake_shared.rs"] @@ -306,6 +309,7 @@ fn quiesce_prepared_heap() { unsafe { libc::malloc_trim(0); } + jemalloc::purge(); std::thread::sleep(std::time::Duration::from_millis(PRE_RUN_QUIESCENCE_MS)); } diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 63d1be2d3bbd..35b7069d1815 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2439,6 +2439,8 @@ dependencies = [ "serde-transcode", "serde_json", "thiserror 2.0.18", + "tikv-jemalloc-sys", + "tikv-jemallocator", "tokio", "tokio-stream", "tokio-util", @@ -8647,6 +8649,26 @@ dependencies = [ "rustc-hash 2.1.2", ] +[[package]] +name = "tikv-jemalloc-sys" +version = "0.7.1+5.3.1-0-g81034ce1f1373e37dc865038e1bc8eeecf559ce8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a2825c78386b4ae0314074867860ba9577875de945f05992c38815cbec327f0" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "249f09e49ab1609436f34c776e84231bead18d6a955f119f939bdc1d847561bd" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.49" diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index 55d99362a881..ec0bbf4b3b21 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -22,7 +22,18 @@ name = "_core" crate-type = ["cdylib", "rlib"] [features] -default = ["custom-policy"] +default = ["custom-policy", "jemalloc"] +# Serve the extension's Rust allocations from a jemalloc linked into it. glibc's arena +# locks serialize the KV router's multi-threaded event workers. The jemalloc symbols +# are prefixed and stay local to the extension, so Python and C/C++ libraries in the +# process keep their own allocator. +# - On Linux with glibc, `background_threads` builds in `background_thread:true`, so +# background threads purge unused pages instead of the allocating threads; +# `_RJEM_MALLOC_CONF` overrides it at run time. +# - `disable_initial_exec_tls` (global-dynamic TLS) keeps the extension loadable +# through `dlopen` without consuming the static TLS reserve. +# - `.cargo/config.toml` pins the jemalloc page size at 64 KiB on aarch64 Linux. +jemalloc = ["dep:tikv-jemallocator", "dep:tikv-jemalloc-sys"] # Links the router plugins Dynamo ships plus the replaceable catalog selected by the # dependency alias below. Enabled by default so a deployment can select a shipped policy through # router-policy YAML without rebuilding. A private image replaces only the catalog dependency, so @@ -83,6 +94,8 @@ serde = { version = "1" } serde_json = { version = "1.0.138" } serde-transcode = "1.1" thiserror = { version = "2.0" } +tikv-jemalloc-sys = { version = "0.7.1", optional = true } +tikv-jemallocator = { version = "0.7.0", features = ["disable_initial_exec_tls"], optional = true } tokio = { version = "1.46.0", features = ["full"] } tokio-stream = { version = "0" } tokio-util = { version = "0.7", features = ["rt"] } @@ -126,6 +139,11 @@ pythonize = "0.23" [target.'cfg(target_os = "linux")'.dependencies] dynamo-llm = { path = "../../llm" } +# jemalloc fails to initialize with `background_thread:true` where it has no background +# threads (macOS), and tikv-jemalloc-sys does not support them on musl. +[target.'cfg(all(target_os = "linux", target_env = "gnu"))'.dependencies] +tikv-jemallocator = { version = "0.7.0", features = ["background_threads"], optional = true } + [target.'cfg(not(target_os = "linux"))'.dependencies] dynamo-llm = { path = "../../llm", default-features = false } diff --git a/lib/bindings/python/rust/allocator.rs b/lib/bindings/python/rust/allocator.rs new file mode 100644 index 000000000000..2c9d21603861 --- /dev/null +++ b/lib/bindings/python/rust/allocator.rs @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The extension's Rust global allocator: a jemalloc linked into the extension. +//! +//! Its symbols carry the `_rjem_` prefix and stay local to the extension, so it serves only +//! the extension's Rust allocations. Python and C/C++ libraries in the process keep their +//! allocator: glibc, or a preloaded jemalloc, which is then a second, independent instance. +//! Its options are built in (see the `jemalloc` feature) and `_RJEM_MALLOC_CONF` overrides them. + +use std::ptr; + +use tikv_jemallocator::Jemalloc; + +#[global_allocator] +static GLOBAL: Jemalloc = Jemalloc; + +/// Returns the unused pages of every arena of the extension's jemalloc to the OS. +pub fn purge() { + // `arena..purge`; MALLCTL_ARENAS_ALL is 4096. + const PURGE_ALL_ARENAS: &std::ffi::CStr = c"arena.4096.purge"; + // SAFETY: the name is NUL-terminated, and this mallctl neither reads nor writes a value, + // so it requires null value pointers and a zero length. + let status = unsafe { + tikv_jemalloc_sys::mallctl( + PURGE_ALL_ARENAS.as_ptr(), + ptr::null_mut(), + ptr::null_mut(), + ptr::null_mut(), + 0, + ) + }; + if status != 0 { + tracing::debug!(status, "jemalloc purge failed"); + } +} diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index d455c0e6c83e..5dd3aff7f97e 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -79,6 +79,8 @@ impl From for RsRouterMode { } } +#[cfg(feature = "jemalloc")] +mod allocator; mod backend; mod context; mod engine; @@ -285,6 +287,11 @@ fn wait_for_bridge_tasks_at_exit(py: Python<'_>) { } fn register_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Model teardown trims glibc arenas; jemalloc holds this extension's Rust allocations, + // so also return its unused pages. + #[cfg(feature = "jemalloc")] + dynamo_llm::discovery::register_allocator_trim_hook(allocator::purge); + // OTLP export no longer requires a pre-existing runtime, so initialize at import. if std::env::var_os(SKIP_PYTHON_LOG_INIT_ENV).is_none() { rs::logging::init(); @@ -1579,6 +1586,10 @@ impl DistributedRuntime { // recognize that both bindings use the same PyClass. #[pyo3(name = "to_capsule")] fn to_capsule<'py>(&self, py: Python<'py>) -> PyResult> { + // The runtime behind this pointer, and the memory it owns, live on this extension's + // jemalloc (the `jemalloc` feature). An extension with another global allocator, such + // as kvbm, that upgrades the pointer would free or reallocate that memory with its own + // allocator. let arc: Arc = Arc::new(self.inner.clone()); let weak: Weak = Arc::downgrade(&arc); diff --git a/lib/bindings/python/tests/test_allocator.py b/lib/bindings/python/tests/test_allocator.py new file mode 100644 index 000000000000..a3f308e4833d --- /dev/null +++ b/lib/bindings/python/tests/test_allocator.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""dynamo._core serves its Rust allocations from its own jemalloc (the default `jemalloc` feature). + +Each probe runs in a fresh interpreter, because jemalloc reads its options once, when the +extension loads (jemalloc initializes in its library constructor). A probe grows a RadixTree, whose nodes live on the Rust +heap, and reports how much of that growth glibc's malloc and a preloaded jemalloc saw. +""" + +import ctypes.util +import json +import os +import shutil +import subprocess +import sys + +import pytest + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="probes glibc and /proc" + ), +] + +BLOCKS = 400_000 +# The tree holds well over this much once it stores BLOCKS blocks, and the Python side of +# the probe allocates far less than this from malloc while the tree grows. +TREE_MIB = 8 +BACKGROUND_THREAD = "jemalloc_bg_thd" + +_PROBE = """ +import ctypes, json, os, sys, time + +from dynamo.llm import RadixTree + +BLOCKS = int(sys.argv[1]) +EXPECT_BACKGROUND_THREADS = sys.argv[2] == "1" +libc = ctypes.CDLL(None) + + +class Mallinfo2(ctypes.Structure): + _fields_ = [(name, ctypes.c_size_t) for name in ( + "arena", "ordblks", "smblks", "hblks", "hblkhd", + "usmblks", "fsmblks", "uordblks", "fordblks", "keepcost", + )] + + +def glibc_in_use(): + if not hasattr(libc, "mallinfo2"): + return None + libc.mallinfo2.restype = Mallinfo2 + info = libc.mallinfo2() + return info.uordblks + info.hblkhd + + +def preloaded_jemalloc_allocated(): + # The extension's jemalloc is prefixed and keeps its symbols local, so only a + # preloaded jemalloc answers here. + if not hasattr(libc, "mallctl"): + return None + epoch = ctypes.c_uint64(1) + libc.mallctl(b"epoch", None, None, ctypes.byref(epoch), ctypes.c_size_t(8)) + allocated = ctypes.c_size_t() + size = ctypes.c_size_t(ctypes.sizeof(allocated)) + if libc.mallctl(b"stats.allocated", ctypes.byref(allocated), ctypes.byref(size), None, 0): + return None + return allocated.value + + +def rss(): + with open("/proc/self/status") as status: + for line in status: + if line.startswith("VmRSS:"): + return int(line.split()[1]) << 10 + + +def sample(): + return {"glibc": glibc_in_use(), "jemalloc": preloaded_jemalloc_allocated(), "rss": rss()} + + +def background_threads(): + count = 0 + for task in os.listdir("/proc/self/task"): + try: + with open(f"/proc/self/task/{task}/comm") as comm: + count += comm.read().strip() == "jemalloc_bg_thd" + except FileNotFoundError: + pass + return count + + +blocks = [{"block_hash": i, "tokens_hash": i} for i in range(BLOCKS)] +event = json.dumps( + {"event_id": 1, "data": {"stored": {"parent_hash": None, "blocks": blocks}}} +).encode() +del blocks + +tree = RadixTree() +# Each find_matches waits for the tree's worker thread, so the tree is settled when sampled. +tree.find_matches([0]) +before = sample() +tree.apply_event(0, event) +assert tree.find_matches([0]).scores, "the tree did not store the event" +after = sample() + +# jemalloc starts its first background thread when the extension loads, and the thread +# names itself once it runs, so a short grace period suffices when none is expected. +deadline = time.monotonic() + (5 if EXPECT_BACKGROUND_THREADS else 0.5) +while not background_threads() and time.monotonic() < deadline: + time.sleep(0.05) + +growth = { + name: None if after[name] is None else (after[name] - before[name]) >> 20 + for name in after +} +print(json.dumps({"growth_mib": growth, "background_threads": background_threads()})) +""" + + +def _probe( + env_overrides: dict[str, str] | None = None, expect_background_threads: bool = True +) -> dict: + # Drop allocator settings inherited from the test environment. + env = { + k: v + for k, v in os.environ.items() + if k + not in ( + "DYN_FRONTEND_JEMALLOC", + "LD_PRELOAD", + "MALLOC_CONF", + "_RJEM_MALLOC_CONF", + ) + } + env.update(env_overrides or {}) + result = subprocess.run( + [ + sys.executable, + "-c", + _PROBE, + str(BLOCKS), + "1" if expect_background_threads else "0", + ], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def _require_mallinfo2(): + if not hasattr(ctypes.CDLL(None), "mallinfo2"): + pytest.skip("needs glibc 2.33 or later for mallinfo2") + + +def test_rust_allocations_bypass_glibc(): + _require_mallinfo2() + probe = _probe() + + assert probe["growth_mib"]["rss"] >= TREE_MIB, probe + assert probe["growth_mib"]["glibc"] < TREE_MIB, probe + + +def test_jemalloc_starts_background_threads(): + probe = _probe() + + assert probe["background_threads"] >= 1, probe + + +def test_rjem_malloc_conf_overrides_the_built_in_options(): + probe = _probe( + {"_RJEM_MALLOC_CONF": "background_thread:false"}, + expect_background_threads=False, + ) + + assert probe["background_threads"] == 0, probe + + +def test_preloaded_jemalloc_leaves_rust_allocations_on_the_extension_jemalloc(): + jemalloc = ctypes.util.find_library("jemalloc") + if not jemalloc: + pytest.skip("needs libjemalloc") + probe = _probe({"LD_PRELOAD": jemalloc, "MALLOC_CONF": "background_thread:false"}) + + assert probe["growth_mib"]["jemalloc"] is not None, "jemalloc did not load" + assert probe["growth_mib"]["rss"] >= TREE_MIB, probe + assert probe["growth_mib"]["jemalloc"] < TREE_MIB, probe + # The preload's own background threads are off, so these are the extension's. + assert probe["background_threads"] >= 1, probe + + +def _extension_path() -> str: + # Imported here, not at module scope, so a missing build fails only these tests. + import dynamo._core + + return dynamo._core.__file__ + + +def _readelf(*args: str) -> str: + readelf = shutil.which("readelf") + if not readelf: + pytest.skip("needs readelf") + return subprocess.run( + [readelf, "-W", *args, _extension_path()], + capture_output=True, + text=True, + check=True, + ).stdout + + +def test_extension_loads_without_static_tls(): + # A dlopen'ed library that uses initial-exec TLS draws on a small static TLS reserve + # and can fail with "cannot allocate memory in static TLS block". + assert "STATIC_TLS" not in _readelf("--dynamic") + + +def test_extension_exports_no_allocator_entry_points(): + entry_points = { + "malloc", + "calloc", + "realloc", + "free", + "posix_memalign", + "aligned_alloc", + "memalign", + "valloc", + "malloc_usable_size", + "mallctl", + } + defined = set() + for line in _readelf("--dyn-syms").splitlines(): + fields = line.split() + # Num: Value Size Type Bind Vis Ndx Name; undefined symbols have Ndx UND. + if len(fields) >= 8 and fields[6] != "UND": + defined.add(fields[7].split("@")[0]) + + exported = sorted( + name for name in defined if name in entry_points or name.startswith("_rjem_") + ) + assert not exported, exported diff --git a/lib/llm/src/discovery.rs b/lib/llm/src/discovery.rs index cc25c7377ca5..9ee037db01f3 100644 --- a/lib/llm/src/discovery.rs +++ b/lib/llm/src/discovery.rs @@ -25,6 +25,7 @@ pub use model_manager::{ModelManager, ModelManagerError, UNKNOWN_METRIC_MODEL}; mod controller; mod allocator; +pub use allocator::register_allocator_trim_hook; mod worker_set; pub use worker_set::WorkerSet; diff --git a/lib/llm/src/discovery/allocator.rs b/lib/llm/src/discovery/allocator.rs index b485d9179459..97dae759b523 100644 --- a/lib/llm/src/discovery/allocator.rs +++ b/lib/llm/src/discovery/allocator.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::sync::OnceLock; #[cfg(all(target_os = "linux", target_env = "gnu"))] use std::sync::{Arc, LazyLock, mpsc}; #[cfg(all(target_os = "linux", target_env = "gnu"))] @@ -52,11 +53,23 @@ static ALLOCATOR_TRIM_SCHEDULER: LazyLock> = LazyLoc Arc::new(|| { // SAFETY: malloc_trim is process-wide and internally synchronizes glibc arenas. let released_pages = unsafe { libc::malloc_trim(0) != 0 }; + if let Some(trim) = ALLOCATOR_TRIM_HOOK.get() { + trim(); + } tracing::debug!(released_pages, "Trimmed allocator after model teardown"); }), )) }); +static ALLOCATOR_TRIM_HOOK: OnceLock = OnceLock::new(); + +/// Registers a release step to run alongside `malloc_trim` after model teardown, for a +/// process whose Rust allocations go through another global allocator. Only the first +/// registration takes effect. +pub fn register_allocator_trim_hook(trim: fn()) { + let _ = ALLOCATOR_TRIM_HOOK.set(trim); +} + /// Schedules allocator trimming after the last WorkerSet, active request, and teardown task exits. pub(crate) struct AllocatorTrimOnDrop { #[cfg(all(target_os = "linux", target_env = "gnu"))]