cuml.accel: serve managed memory from an RMM pool

docs/source/cuml-accel/usage.rst

@@ -72,7 +72,10 @@ When the platform supports it and RMM has not already been configured,
 ``cuml.accel`` enables `managed memory`_. Managed memory can use host memory in
 addition to GPU memory and migrates data as needed, reducing the risk of GPU
 out-of-memory errors. It does not prevent exhaustion of combined host and
-device memory, and heavy oversubscription can slow execution. Managed memory
+device memory, and heavy oversubscription can slow execution. The managed
+allocations are served from an RMM memory pool (as in ``cudf.pandas``'s
+``managed_pool`` mode), so repeated estimator calls do not pay a CUDA
+allocation and free each; the pool starts small and grows on demand. Managed memory
 is not enabled on WSL 2 or when RMM was configured before ``cuml.accel``.
 If managed-memory oversubscription causes unexpectedly slow execution, use the
 CLI's ``--disable-uvm`` flag to disable it and compare performance.

python/cuml/cuml/accel/core.py

@@ -255,12 +255,38 @@ def install(
                     "skipping enabling managed memory."
                 )
             else:
-                rmm.mr.set_current_device_resource(
-                    rmm.mr.PrefetchResourceAdaptor(
-                        rmm.mr.ManagedMemoryResource()
+                # Pool the managed allocations. Without a pool every
+                # allocation is a cudaMallocManaged and every free a
+                # cudaFree, each costing milliseconds (up to a dozen frees
+                # per predict call), which dominated short estimator calls.
+                # Same construction as cudf.pandas's "managed_pool" mode,
+                # but with a small initial block that grows on demand:
+                # cuml.accel is also installed in every worker process a
+                # user spawns (joblib, multiprocessing), so reserving a
+                # large share of the device up front in each of them is
+                # not acceptable.
+                managed = rmm.mr.ManagedMemoryResource()
+                try:
+                    mr = rmm.mr.PoolMemoryResource(
+                        managed, initial_pool_size=256 << 20
+                    )
+                except MemoryError as exc:
+                    # The pool allocates its initial block on construction.
+                    # Never let that stop acceleration on a constrained
+                    # device; the unpooled resource is what ran before.
+                    logger.debug(
+                        "Could not allocate the initial managed memory "
+                        f"pool ({exc}); using unpooled managed memory."
                     )
+                    mr = managed
+                rmm.mr.set_current_device_resource(
+                    rmm.mr.PrefetchResourceAdaptor(mr)
+                )
+                logger.debug(
+                    "Enabled managed memory."
+                    if mr is managed
+                    else "Enabled pooled managed memory."
                 )
-                logger.debug("Enabled managed memory.")
         else:
             logger.debug("Could not enable managed memory on this platform.")
 

python/cuml/cuml_accel_tests/test_core.py

@@ -1,8 +1,11 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 import importlib
 import importlib.metadata
 import multiprocessing
+import os
+import subprocess
+import sys
 from inspect import Parameter, signature
 
 import pytest
@@ -67,6 +70,48 @@ def test_enabled_in_loky_executor():
         assert cuml.accel.is_proxy(remote)
 
 
+_INSTALL_SCRIPT = """
+import rmm
+import cuml.accel.core as core
+
+core._is_concurrent_managed_access_supported = lambda: True
+if {fail_pool}:
+    def fail(*args, **kwargs):
+        raise MemoryError("no room")
+    rmm.mr.PoolMemoryResource = fail
+core.install(log_level="debug")
+mr = rmm.mr.get_current_device_resource()
+assert isinstance(mr, rmm.mr.PrefetchResourceAdaptor)
+print("UPSTREAM", type(mr.upstream_mr).__name__)
+"""
+
+
+@pytest.mark.parametrize(
+    "fail_pool, upstream, message",
+    [
+        (False, "PoolMemoryResource", "Enabled pooled managed memory."),
+        (True, "ManagedMemoryResource", "Enabled managed memory."),
+    ],
+)
+def test_install_managed_memory_resource(fail_pool, upstream, message):
+    """install() pools managed memory, and falls back to (and reports)
+    unpooled managed memory if the initial pool block cannot be allocated."""
+    env = {
+        k: v for k, v in os.environ.items() if not k.startswith("CUML_ACCEL")
+    }
+    proc = subprocess.run(
+        [sys.executable, "-c", _INSTALL_SCRIPT.format(fail_pool=fail_pool)],
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    output = proc.stdout + proc.stderr
+    assert proc.returncode == 0, output
+    assert f"UPSTREAM {upstream}" in output
+    assert message in output
+    assert ("Enabled pooled managed memory." in output) is not fail_pool
+
+
 def get_level():
     return cuml.accel.core.logger.level.name.lower()