Skip to content

Topology-aware default pool sizing for PinnedMemoryResource#1012

Merged
rapids-bot[bot] merged 22 commits intorapidsai:mainfrom
nirandaperera:pinned_pool_init_max_size
May 7, 2026
Merged

Topology-aware default pool sizing for PinnedMemoryResource#1012
rapids-bot[bot] merged 22 commits intorapidsai:mainfrom
nirandaperera:pinned_pool_init_max_size

Conversation

@nirandaperera
Copy link
Copy Markdown
Contributor

@nirandaperera nirandaperera commented May 4, 2026

Replace the hardcoded pinned_initial_pool_size = 0 / pinned_max_pool_size = unbounded
defaults with values scaled to the system's GPU count.

Changes

  • get_host_memory_per_gpu() — new utility in system_info that returns
    current numa node total host memory/ N GPUs in current numa node

  • PinnedMemoryResource — two new named constants drive the defaults:

    • DefaultInitiPoolSizeFactor = "10%" → initial pool = 10% of per-GPU host memory
    • DefaultMaxPoolSizeFactor = "80%" → max pool = 80% of per-GPU host memory

    Also corrects the documented default from false to true.

  • Python bindings — exposes get_host_memory_per_gpu() via system_info.pyx/.pyi.

  • TestsPinnedResource.from_default_options verifies pool sizes match the
    factor-scaled values when no options are provided.

  • Benchmark -- Benchmark initialization time against initial pool size - Results

  • Docsconfiguration.md updated to reflect the new defaults.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested review from a team as code owners May 4, 2026 21:39
@nirandaperera nirandaperera added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels May 4, 2026
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera force-pushed the pinned_pool_init_max_size branch from ce335cb to 7bbd171 Compare May 5, 2026 05:39
@nirandaperera nirandaperera requested a review from a team as a code owner May 5, 2026 05:39
Comment thread cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp Outdated
Comment thread cpp/src/system_info.cpp Outdated

std::uint64_t get_host_memory_per_gpu() {
auto const num_gpus = get_topology().num_gpus;
return get_total_host_memory() / std::max<std::uint64_t>(1, num_gpus);
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you double-check how get_total_host_memory behaves on devices where GPU memory is also exposed as NUMA nodes? In those cases free will also show GPU memory as part of system memory, which may break the assumption get_total_host_memory presents only system memory.

IIRC, the conditions for GPU memory to be exposed as system memory are:

  1. Coherent system (e.g., Grace-Hopper/Grace-Blackwell)
  2. NVIDIA driver in NUMA mode (not CDMM)
  3. HMM enabled on the kernel

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah! this is a good point.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pentschev you were right Peter.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice that there was a simple solution. I also didn't think of the bigger picture of the problem, we cannot just take the whole host memory and divide equally among GPUs, since their NUMA nodes may have different amount of memory, and looks like your new solution already ensures that.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from pentschev May 5, 2026 20:37
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from madsbk May 5, 2026 21:12
Copy link
Copy Markdown
Member

@madsbk madsbk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, I only have minor stuff

Comment thread cpp/src/system_info.cpp Outdated
Comment thread cpp/src/memory/pinned_memory_resource.cpp Outdated
Comment thread cpp/src/memory/pinned_memory_resource.cpp Outdated
Comment thread cpp/include/rapidsmpf/system_info.hpp
Comment thread python/rapidsmpf/rapidsmpf/utils/system_info.pyx Outdated
nirandaperera and others added 3 commits May 6, 2026 07:40
Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera
Copy link
Copy Markdown
Contributor Author

/merge

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Copy link
Copy Markdown
Member

@pentschev pentschev left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please update the description to reflect recent changes @nirandaperera .

Temporarily blocking to ensure PR and Python get_host_memory_per_gpu() are updated.

Comment thread cpp/include/rapidsmpf/system_info.hpp
@pentschev
Copy link
Copy Markdown
Member

I think there's some partial GH UI breakage, I cannot comment on the files now for some reason , but get_host_memory_per_gpu description in python/rapidsmpf/rapidsmpf/utils/system_info.pyx is outdated, @nirandaperera please update with NUMA info.

Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from pentschev May 6, 2026 16:25
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
@nirandaperera nirandaperera requested a review from a team as a code owner May 6, 2026 23:07
@nirandaperera nirandaperera requested a review from bdice May 6, 2026 23:07

using rapidsmpf::safe_cast;

// When the RAPIDSMPF_SMOKE_TEST_MODE env var is set (to any value), each
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bad pattern. One will almost definitely one day attempt RAPIDSMPF_SMOKE_TEST_MODE=0 and will be confused. A better approach would be adding a CLI argument, like all other bench_* binaries we have.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pentschev Well, this is a bit tricky. Passing custom args might not work here, because in google benchmarks, Apply(...) runs during static initialization. So, to make it work, we have to parse args before registering benchmarks. That means, we need a custom main, parse smoke-test arg, and then move all bench registrations after that. I'd rather keep this simple like this.

@nirandaperera nirandaperera requested a review from pentschev May 7, 2026 16:30
Copy link
Copy Markdown
Member

@pentschev pentschev left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Niranda.

Comment on lines +28 to +32
// We use an env var rather than a CLI flag because google-benchmark's
// BENCHMARK(...)->Apply(...) macros run during static initialization, before
// main() has a chance to parse argv. A CLI-flag approach would require moving
// every benchmark registration into main() (via benchmark::RegisterBenchmark),
// which is more invasive. std::getenv works fine during static init.
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Whelp, this is lame. 😞

@nirandaperera
Copy link
Copy Markdown
Contributor Author

/merge

Copy link
Copy Markdown
Member

@jameslamb jameslamb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

approving for ci-codeowners / packaging-codeowners, the changes that pulled in those groups look non-controversial.

@rapids-bot rapids-bot Bot merged commit b72e232 into rapidsai:main May 7, 2026
127 of 129 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants