From 4d75d71f302e10d1bf3cfc4a80eff2961792a592 Mon Sep 17 00:00:00 2001 From: Ari Lotter Date: Fri, 6 Feb 2026 14:14:45 -0500 Subject: [PATCH] feat: use jemallocator instead of default malloc closes #459 the default system allocator isn't designed for multithreaded workloads like psyche. similar programs like the agave solana validator use jemalloc via the tikv-jemallocator crate. we adopt the same in our production binaries. --- Cargo.lock | 12 ++++++++---- Cargo.toml | 4 ++++ architectures/centralized/client/Cargo.toml | 1 + architectures/centralized/client/src/main.rs | 3 +++ architectures/centralized/server/Cargo.toml | 1 + architectures/centralized/server/src/main.rs | 3 +++ architectures/decentralized/solana-client/Cargo.toml | 1 + .../decentralized/solana-client/src/main.rs | 3 +++ .../inference-only/inference-node/Cargo.toml | 2 ++ .../inference-only/inference-node/src/main.rs | 3 +++ 10 files changed, 29 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 54722dcee..cb41d9aa9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6883,6 +6883,7 @@ dependencies = [ "psyche-tui", "psyche-watcher", "rand 0.9.2", + "tikv-jemallocator", "time", "tokio", "tokio-util 0.7.16", @@ -6929,6 +6930,7 @@ dependencies = [ "rand 0.9.2", "serde", "sha2 0.10.9", + "tikv-jemallocator", "tokio", "tokio-stream", "tokio-util 0.7.16", @@ -7174,6 +7176,7 @@ dependencies = [ "psyche-metrics", "psyche-network", "serde_json", + "tikv-jemallocator", "tokio", "tokio-util 0.7.16", "tracing", @@ -7339,6 +7342,7 @@ dependencies = [ "rand_chacha 0.9.0", "serde", "serde_json", + "tikv-jemallocator", "time", "tokio", "tokio-util 0.7.16", @@ -11969,9 +11973,9 @@ dependencies = [ [[package]] name = "tikv-jemalloc-sys" -version = "0.6.0+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" +version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd3c60906412afa9c2b5b5a48ca6a5abe5736aec9eb48ad05037a677e52e4e2d" +checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b" dependencies = [ "cc", "libc", @@ -11979,9 +11983,9 @@ dependencies = [ [[package]] name = "tikv-jemallocator" -version = "0.6.0" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cec5ff18518d81584f477e9bfdf957f5bb0979b0bac3af4ca30b5b3ae2d2865" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" dependencies = [ "libc", "tikv-jemalloc-sys", diff --git a/Cargo.toml b/Cargo.toml index 01a1653d1..544d9592b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -106,6 +106,10 @@ ts-rs = { git = "https://github.com/arilotter/ts-rs.git", rev = "92ce1752227fec9 ] } rayon = "1.10.0" +tikv-jemallocator = { version = "0.6.1", features = [ + "unprefixed_malloc_on_supported_platforms", +] } + # dev only test-log = { version = "0.2.16", features = ["trace"] } pretty_assertions = "1.4.1" diff --git a/architectures/centralized/client/Cargo.toml b/architectures/centralized/client/Cargo.toml index fbc6ca39a..fc47dfa5a 100644 --- a/architectures/centralized/client/Cargo.toml +++ b/architectures/centralized/client/Cargo.toml @@ -26,6 +26,7 @@ bytemuck.workspace = true clap-markdown.workspace = true hex = "0.4.3" psyche-python-extension-impl = { workspace = true, optional = true } +tikv-jemallocator.workspace = true [features] parallelism = ["psyche-client/parallelism"] diff --git a/architectures/centralized/client/src/main.rs b/architectures/centralized/client/src/main.rs index fad5d3817..bfc381c2d 100644 --- a/architectures/centralized/client/src/main.rs +++ b/architectures/centralized/client/src/main.rs @@ -16,6 +16,9 @@ use tracing::info; mod app; +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Parser, Debug)] struct Args { #[command(subcommand)] diff --git a/architectures/centralized/server/Cargo.toml b/architectures/centralized/server/Cargo.toml index 4ce5c9d59..4558b2423 100644 --- a/architectures/centralized/server/Cargo.toml +++ b/architectures/centralized/server/Cargo.toml @@ -31,6 +31,7 @@ bytemuck.workspace = true toml.workspace = true clap-markdown.workspace = true psyche-python-extension-impl = { workspace = true, optional = true } +tikv-jemallocator.workspace = true [features] python = ["psyche-python-extension-impl"] diff --git a/architectures/centralized/server/src/main.rs b/architectures/centralized/server/src/main.rs index 58cf2cf00..a4738557b 100644 --- a/architectures/centralized/server/src/main.rs +++ b/architectures/centralized/server/src/main.rs @@ -16,6 +16,9 @@ use std::{ }; use tracing::{error, info}; +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Parser, Debug)] struct Args { #[command(subcommand)] diff --git a/architectures/decentralized/solana-client/Cargo.toml b/architectures/decentralized/solana-client/Cargo.toml index 752e3af4f..6e19d5cf3 100644 --- a/architectures/decentralized/solana-client/Cargo.toml +++ b/architectures/decentralized/solana-client/Cargo.toml @@ -30,6 +30,7 @@ tokio.workspace = true tokio-util.workspace = true tracing.workspace = true psyche-python-extension-impl = { workspace = true, optional = true } +tikv-jemallocator.workspace = true [features] parallelism = ["psyche-client/parallelism"] diff --git a/architectures/decentralized/solana-client/src/main.rs b/architectures/decentralized/solana-client/src/main.rs index 4c4849c19..c85c6714d 100644 --- a/architectures/decentralized/solana-client/src/main.rs +++ b/architectures/decentralized/solana-client/src/main.rs @@ -29,6 +29,9 @@ use tracing::info; mod app; mod network_identity; +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Parser, Debug)] struct CliArgs { #[command(subcommand)] diff --git a/architectures/inference-only/inference-node/Cargo.toml b/architectures/inference-only/inference-node/Cargo.toml index b1a18d74f..8559d1fef 100644 --- a/architectures/inference-only/inference-node/Cargo.toml +++ b/architectures/inference-only/inference-node/Cargo.toml @@ -33,3 +33,5 @@ iroh-blobs.workspace = true iroh-gossip.workspace = true serde_json.workspace = true + +tikv-jemallocator.workspace = true diff --git a/architectures/inference-only/inference-node/src/main.rs b/architectures/inference-only/inference-node/src/main.rs index 4776458ae..06da8acce 100644 --- a/architectures/inference-only/inference-node/src/main.rs +++ b/architectures/inference-only/inference-node/src/main.rs @@ -18,6 +18,9 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Parser, Debug)] #[command(name = "psyche-inference-node")] struct Cli {