Add available_device_memory to fetch free amount of memory on a GPU

python/rmm/rmm/_lib/memory_resource.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -111,6 +111,7 @@ cdef extern from "rmm/mr/device/cuda_async_memory_resource.hpp" \
 
 cdef extern from "rmm/cuda_device.hpp" namespace "rmm" nogil:
     size_t percent_of_free_device_memory(int percent) except +
+    size_t available_device_memory() expect +
 
 cdef extern from "rmm/mr/device/pool_memory_resource.hpp" \
         namespace "rmm::mr" nogil:
@@ -1188,3 +1189,10 @@ def get_log_filenames():
         else None
         for i, each_mr in _per_device_mrs.items()
     }
+
+
+def get_free_device_memory():
+    """
+    Return the number of bytes of available device memory.
+    """
+    return available_device_memory()

python/rmm/rmm/mr.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -32,6 +32,7 @@
     enable_logging,
     get_current_device_resource,
     get_current_device_resource_type,
+    get_free_device_memory,
     get_log_filenames,
     get_per_device_resource,
     get_per_device_resource_type,
@@ -65,6 +66,7 @@
     "get_current_device_resource",
     "get_per_device_resource_type",
     "get_current_device_resource_type",
+    "get_free_device_memory",
     "get_log_filenames",
     "is_initialized",
 ]

python/rmm/rmm/tests/test_rmm.py

@@ -1002,3 +1002,10 @@ def test_invalid_logging_level(level):
         rmm.set_flush_level(level)
     with pytest.raises(TypeError):
         rmm.should_log(level)
+
+
+def test_get_free_device_memory():
+    from rmm._lib.memory_resource import get_free_device_memory
+
+    memory = get_free_device_memory()
+    assert memory != 0