From 10d7eb4d1e9e61f9efd7133f02fdd195b1e94156 Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Tue, 21 May 2024 14:34:01 +0000 Subject: [PATCH 1/8] Add get_free_device_memory --- python/rmm/rmm/_lib/memory_resource.pyx | 9 ++++++++- python/rmm/rmm/tests/test_rmm.py | 7 +++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx index 100d18b56..63dd4c7d6 100644 --- a/python/rmm/rmm/_lib/memory_resource.pyx +++ b/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. @@ -1188,3 +1188,10 @@ def get_log_filenames(): else None for i, each_mr in _per_device_mrs.items() } + + +def get_free_device_memory(percent): + """ + Returns the percentage of the free memory available + """ + return percent_of_free_device_memory(percent) diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index c37fe0298..b8e5aabb1 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/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(percent=10) + assert memory != 0 From 25b1d5db8fc494e87a2e66b3a63e372ac3f66501 Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Tue, 21 May 2024 18:03:35 +0000 Subject: [PATCH 2/8] update import --- python/rmm/rmm/mr.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/rmm/rmm/mr.py b/python/rmm/rmm/mr.py index 4f6b801f5..571a43274 100644 --- a/python/rmm/rmm/mr.py +++ b/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", ] From 807d7e0632955887bee77e7307ac7b81f0f4cb8f Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Tue, 28 May 2024 19:19:22 +0000 Subject: [PATCH 3/8] Drop percentage --- python/rmm/rmm/_lib/memory_resource.pyx | 4 ++-- python/rmm/rmm/tests/test_rmm.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx index 63dd4c7d6..0cccbd1e3 100644 --- a/python/rmm/rmm/_lib/memory_resource.pyx +++ b/python/rmm/rmm/_lib/memory_resource.pyx @@ -1190,8 +1190,8 @@ def get_log_filenames(): } -def get_free_device_memory(percent): +def get_free_device_memory(): """ - Returns the percentage of the free memory available + Return the number of bytes of available device memory. """ return percent_of_free_device_memory(percent) diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index b8e5aabb1..f996725e1 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/python/rmm/rmm/tests/test_rmm.py @@ -1007,5 +1007,5 @@ def test_invalid_logging_level(level): def test_get_free_device_memory(): from rmm._lib.memory_resource import get_free_device_memory - memory = get_free_device_memory(percent=10) + memory = get_free_device_memory() assert memory != 0 From 9c5d48680add334268704d99095137534c276874 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Tue, 28 May 2024 15:22:15 -0500 Subject: [PATCH 4/8] Update python/rmm/rmm/_lib/memory_resource.pyx --- python/rmm/rmm/_lib/memory_resource.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx index 0cccbd1e3..1862c4a74 100644 --- a/python/rmm/rmm/_lib/memory_resource.pyx +++ b/python/rmm/rmm/_lib/memory_resource.pyx @@ -1194,4 +1194,4 @@ def get_free_device_memory(): """ Return the number of bytes of available device memory. """ - return percent_of_free_device_memory(percent) + return percent_of_free_device_memory(100) From e215b44efcd0113722ca36c619dad980bc65f496 Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Tue, 28 May 2024 22:57:40 +0000 Subject: [PATCH 5/8] Use available_device_memory --- python/rmm/rmm/_lib/memory_resource.pyx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx index 1862c4a74..11c862f35 100644 --- a/python/rmm/rmm/_lib/memory_resource.pyx +++ b/python/rmm/rmm/_lib/memory_resource.pyx @@ -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: @@ -1194,4 +1195,4 @@ def get_free_device_memory(): """ Return the number of bytes of available device memory. """ - return percent_of_free_device_memory(100) + return available_device_memory() From cbc7898a920fda1be8ea2f7ca343e5cab84efda4 Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Wed, 29 May 2024 14:30:51 +0000 Subject: [PATCH 6/8] address reviews --- python/rmm/rmm/_lib/memory_resource.pxd | 7 ++++++- python/rmm/rmm/_lib/memory_resource.pyx | 20 +++++++++++++------- python/rmm/rmm/mr.py | 4 ++-- python/rmm/rmm/tests/test_rmm.py | 13 +++++++++---- 4 files changed, 30 insertions(+), 14 deletions(-) diff --git a/python/rmm/rmm/_lib/memory_resource.pxd b/python/rmm/rmm/_lib/memory_resource.pxd index f9c2e91de..15ddc84f6 100644 --- a/python/rmm/rmm/_lib/memory_resource.pxd +++ b/python/rmm/rmm/_lib/memory_resource.pxd @@ -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. @@ -14,6 +14,7 @@ from libc.stdint cimport int8_t from libcpp.memory cimport shared_ptr +from libcpp.pair cimport pair from libcpp.string cimport string from libcpp.vector cimport vector @@ -32,6 +33,10 @@ cdef extern from "rmm/mr/device/device_memory_resource.hpp" \ cuda_stream_view stream ) except + +cdef extern from "rmm/cuda_device.hpp" namespace "rmm" nogil: + size_t percent_of_free_device_memory(int percent) except + + pair[size_t, size_t] available_device_memory() except + + cdef class DeviceMemoryResource: cdef shared_ptr[device_memory_resource] c_obj cdef device_memory_resource* get_mr(self) noexcept nogil diff --git a/python/rmm/rmm/_lib/memory_resource.pyx b/python/rmm/rmm/_lib/memory_resource.pyx index 11c862f35..992203c27 100644 --- a/python/rmm/rmm/_lib/memory_resource.pyx +++ b/python/rmm/rmm/_lib/memory_resource.pyx @@ -32,9 +32,16 @@ from libcpp.string cimport string from cuda.cudart import cudaError_t from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice + from rmm._cuda.stream cimport Stream + from rmm._cuda.stream import DEFAULT_STREAM + from rmm._lib.cuda_stream_view cimport cuda_stream_view +from rmm._lib.memory_resource cimport ( + available_device_memory as c_available_device_memory, + percent_of_free_device_memory as c_percent_of_free_device_memory, +) from rmm._lib.per_device_resource cimport ( cuda_device_id, set_per_device_resource as cpp_set_per_device_resource, @@ -109,9 +116,6 @@ cdef extern from "rmm/mr/device/cuda_async_memory_resource.hpp" \ win32 win32_kmt -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: @@ -369,7 +373,7 @@ cdef class PoolMemoryResource(UpstreamResourceAdaptor): cdef size_t c_initial_pool_size cdef optional[size_t] c_maximum_pool_size c_initial_pool_size = ( - percent_of_free_device_memory(50) if + c_percent_of_free_device_memory(50) if initial_pool_size is None else initial_pool_size ) @@ -1191,8 +1195,10 @@ def get_log_filenames(): } -def get_free_device_memory(): +def available_device_memory(): """ - Return the number of bytes of available device memory. + Returns a tuple of free and total device memory memory. """ - return available_device_memory() + cdef pair[size_t, size_t] res + res = c_available_device_memory() + return (res.first, res.second) diff --git a/python/rmm/rmm/mr.py b/python/rmm/rmm/mr.py index 571a43274..b6ae4e6cd 100644 --- a/python/rmm/rmm/mr.py +++ b/python/rmm/rmm/mr.py @@ -28,11 +28,11 @@ UpstreamResourceAdaptor, _flush_logs, _initialize, + available_device_memory, disable_logging, 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, @@ -58,6 +58,7 @@ "UpstreamResourceAdaptor", "_flush_logs", "_initialize", + "available_device_memory", "set_per_device_resource", "enable_logging", "disable_logging", @@ -66,7 +67,6 @@ "get_current_device_resource", "get_per_device_resource_type", "get_current_device_resource_type", - "get_free_device_memory", "get_log_filenames", "is_initialized", ] diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index f996725e1..f754f50c9 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/python/rmm/rmm/tests/test_rmm.py @@ -1004,8 +1004,13 @@ def test_invalid_logging_level(level): rmm.should_log(level) -def test_get_free_device_memory(): - from rmm._lib.memory_resource import get_free_device_memory +def test_available_device_memory(): + from rmm._lib.memory_resource import available_device_memory - memory = get_free_device_memory() - assert memory != 0 + initial_memory = available_device_memory() + device_buffer = rmm.DeviceBuffer.to_device( # noqa: F841 + np.zeros(10000, dtype="u1") + ) + final_memory = available_device_memory() + assert initial_memory[1] == final_memory[1] + assert initial_memory[0] > final_memory[0] From 3a6fcfad550c56e38f1e60c487ba05789cc719f6 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Thu, 30 May 2024 10:23:37 -0500 Subject: [PATCH 7/8] Apply suggestions from code review Co-authored-by: Bradley Dice --- python/rmm/rmm/tests/test_rmm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index f754f50c9..d0637630a 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/python/rmm/rmm/tests/test_rmm.py @@ -1005,7 +1005,7 @@ def test_invalid_logging_level(level): def test_available_device_memory(): - from rmm._lib.memory_resource import available_device_memory + from rmm.mr import available_device_memory initial_memory = available_device_memory() device_buffer = rmm.DeviceBuffer.to_device( # noqa: F841 From e3dffea83800e6b088422ffbe303178c5c51220a Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Thu, 30 May 2024 15:24:03 +0000 Subject: [PATCH 8/8] assert only +ve values --- python/rmm/rmm/tests/test_rmm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py index d0637630a..62adcd4a5 100644 --- a/python/rmm/rmm/tests/test_rmm.py +++ b/python/rmm/rmm/tests/test_rmm.py @@ -1013,4 +1013,5 @@ def test_available_device_memory(): ) final_memory = available_device_memory() assert initial_memory[1] == final_memory[1] - assert initial_memory[0] > final_memory[0] + assert initial_memory[0] > 0 + assert final_memory[0] > 0