Backport PR rapidsai#6359

Author: dantegd

python/cuml/cuml/experimental/accel/__init__.py

@@ -19,10 +19,13 @@
 
 from .magics import load_ipython_extension
 
+from cuda.bindings import runtime
 from cuml.internals import logger
 from cuml.internals.global_settings import GlobalSettings
 from cuml.internals.memory_utils import set_global_output_type
-from cuml.internals.safe_imports import UnavailableError
+from cuml.internals.safe_imports import UnavailableError, gpu_only_import
+
+rmm = gpu_only_import("rmm")
 
 __all__ = ["load_ipython_extension", "install"]
 
@@ -31,11 +34,38 @@ def _install_for_library(library_name):
     importlib.import_module(f"._wrappers.{library_name}", __name__)
 
 
-def install():
+def _is_concurrent_managed_access_supported():
+    """Check the availability of concurrent managed access (UVM).
+    Note that WSL2 does not support managed memory.
+    """
+
+    # Ensure CUDA is initialized before checking cudaDevAttrConcurrentManagedAccess
+    runtime.cudaFree(0)
+
+    device_id = 0
+    err, supports_managed_access = runtime.cudaDeviceGetAttribute(
+        runtime.cudaDeviceAttr.cudaDevAttrConcurrentManagedAccess, device_id
+    )
+    if err != runtime.cudaError_t.cudaSuccess:
+        logger.error(
+            f"Failed to check cudaDevAttrConcurrentManagedAccess with error {err}"
+        )
+        return False
+    return supports_managed_access != 0
+
+
+def install(disable_uvm=False):
     """Enable cuML Accelerator Mode."""
     logger.set_level(logger.level_enum.info)
     logger.set_pattern("%v")
 
+    if not disable_uvm:
+        if _is_concurrent_managed_access_supported():
+            logger.debug("cuML: Enabling managed memory...")
+            rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource())
+        else:
+            logger.warn("cuML: Could not enable managed memory.")
+
     logger.debug("cuML: Installing accelerator...")
     libraries_to_accelerate = ["sklearn", "umap", "hdbscan"]
     accelerated_libraries = []

python/cuml/cuml/experimental/accel/__main__.py

@@ -54,14 +54,22 @@
     default="converted_sklearn_model.pkl",
     help="Output path for the converted sklearn estimator file.",
 )
+@click.option(
+    "--disable-uvm",
+    is_flag=True,
+    default=False,
+    help="Disable UVM (managed memory) allocations.",
+)
 @click.option(
     "-v",
     "--verbose",
     count=True,
     help="Increase output verbosity (can be used multiple times, e.g. -vv). Default shows warnings only.",
 )
 @click.argument("args", nargs=-1)
-def main(module, strict, convert_to_sklearn, format, output, verbose, args):
+def main(
+    module, convert_to_sklearn, format, output, disable_uvm, verbose, args
+):
     default_logger_level_index = list(logger.level_enum).index(
         logger.level_enum.warn
     )
@@ -70,7 +78,7 @@ def main(module, strict, convert_to_sklearn, format, output, verbose, args):
     logger.set_level(logger_level)
     logger.set_pattern("%v")
 
-    install()
+    install(disable_uvm=disable_uvm)
 
     # If the user requested a conversion, handle it and exit
     if convert_to_sklearn: