From 25a4a67cdcd8b70a046465ae3c333ae433cbeee1 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 3 Oct 2024 17:01:09 -0700 Subject: [PATCH 1/5] Fix rmm imports --- python/rmm/rmm/_lib/__init__.py | 2 -- python/rmm/rmm/_lib/device_buffer.py | 21 ++++++++++++++++ python/rmm/rmm/_lib/logger.py | 24 ++++++++++++++++++ python/rmm/rmm/_lib/memory_resource.py | 34 ++++++++++++++++++++++++++ 4 files changed, 79 insertions(+), 2 deletions(-) create mode 100644 python/rmm/rmm/_lib/device_buffer.py create mode 100644 python/rmm/rmm/_lib/logger.py create mode 100644 python/rmm/rmm/_lib/memory_resource.py diff --git a/python/rmm/rmm/_lib/__init__.py b/python/rmm/rmm/_lib/__init__.py index 7cfddab60..2a9c9939f 100644 --- a/python/rmm/rmm/_lib/__init__.py +++ b/python/rmm/rmm/_lib/__init__.py @@ -11,5 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -from rmm.pylibrmm import * diff --git a/python/rmm/rmm/_lib/device_buffer.py b/python/rmm/rmm/_lib/device_buffer.py new file mode 100644 index 000000000..ddaf6383c --- /dev/null +++ b/python/rmm/rmm/_lib/device_buffer.py @@ -0,0 +1,21 @@ +# Copyright (c) 2018-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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rmm.pylibrmm.device_buffer import ( # noqa: F401 + _copy_async, + copy_device_to_ptr, + copy_host_to_ptr, + copy_ptr_to_host, + to_device, +) diff --git a/python/rmm/rmm/_lib/logger.py b/python/rmm/rmm/_lib/logger.py new file mode 100644 index 000000000..1e9b519b8 --- /dev/null +++ b/python/rmm/rmm/_lib/logger.py @@ -0,0 +1,24 @@ +# Copyright (c) 2018-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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rmm.librmm._logger import logging_level # noqa: F401 +from rmm.pylibrmm.logger import ( # noqa: F401 + _validate_level_type, + flush_logger, + get_flush_level, + get_logging_level, + set_flush_level, + set_logging_level, + should_log, +) diff --git a/python/rmm/rmm/_lib/memory_resource.py b/python/rmm/rmm/_lib/memory_resource.py new file mode 100644 index 000000000..acd24bd2b --- /dev/null +++ b/python/rmm/rmm/_lib/memory_resource.py @@ -0,0 +1,34 @@ +# Copyright (c) 2018-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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rmm.pylibrmm.memory_resource import ( # noqa: F401 + BinningMemoryResource, + CallbackMemoryResource, + CudaAsyncMemoryResource, + CudaMemoryResource, + DeviceMemoryResource, + FailureCallbackResourceAdaptor, + FixedSizeMemoryResource, + LimitingResourceAdaptor, + LoggingResourceAdaptor, + ManagedMemoryResource, + PoolMemoryResource, + PrefetchResourceAdaptor, + SamHeadroomMemoryResource, + StatisticsResourceAdaptor, + SystemMemoryResource, + TrackingResourceAdaptor, + UpstreamResourceAdaptor, + get_current_device_resource, +) From 98b098f19f476817f681ac504f0c36d097957252 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 3 Oct 2024 20:42:48 -0700 Subject: [PATCH 2/5] add missing imports to memory_resource and remove _copy_async import from device buffer --- python/rmm/rmm/_lib/device_buffer.py | 1 - python/rmm/rmm/_lib/memory_resource.py | 10 ++++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/python/rmm/rmm/_lib/device_buffer.py b/python/rmm/rmm/_lib/device_buffer.py index ddaf6383c..78bd38b83 100644 --- a/python/rmm/rmm/_lib/device_buffer.py +++ b/python/rmm/rmm/_lib/device_buffer.py @@ -13,7 +13,6 @@ # limitations under the License. from rmm.pylibrmm.device_buffer import ( # noqa: F401 - _copy_async, copy_device_to_ptr, copy_host_to_ptr, copy_ptr_to_host, diff --git a/python/rmm/rmm/_lib/memory_resource.py b/python/rmm/rmm/_lib/memory_resource.py index acd24bd2b..0d47e8c9b 100644 --- a/python/rmm/rmm/_lib/memory_resource.py +++ b/python/rmm/rmm/_lib/memory_resource.py @@ -30,5 +30,15 @@ SystemMemoryResource, TrackingResourceAdaptor, UpstreamResourceAdaptor, + _flush_logs, + available_device_memory, + disable_logging, + enable_logging, get_current_device_resource, + get_current_device_resource_type, + get_log_filenames, + get_per_device_resource_type, + is_initialized, + set_current_device_resource, + set_per_device_resource, ) From 07bd75b184c4c01ed4b4289236cbd0402edfe04d Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 3 Oct 2024 20:54:33 -0700 Subject: [PATCH 3/5] add cuda_stream --- python/rmm/rmm/_lib/cuda_stream.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 python/rmm/rmm/_lib/cuda_stream.py diff --git a/python/rmm/rmm/_lib/cuda_stream.py b/python/rmm/rmm/_lib/cuda_stream.py new file mode 100644 index 000000000..1eb424e12 --- /dev/null +++ b/python/rmm/rmm/_lib/cuda_stream.py @@ -0,0 +1,15 @@ +# Copyright (c) 2018-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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from rmm.pylibrmm.cuda_stream import CudaStream # noqa: F401 From 61887b5684a163743e4c92a89e5e3d510ef69f5d Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 4 Oct 2024 10:25:42 -0700 Subject: [PATCH 4/5] from rmm._lib import Foo --- python/rmm/rmm/_lib/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/rmm/rmm/_lib/__init__.py b/python/rmm/rmm/_lib/__init__.py index 2a9c9939f..7cfddab60 100644 --- a/python/rmm/rmm/_lib/__init__.py +++ b/python/rmm/rmm/_lib/__init__.py @@ -11,3 +11,5 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +from rmm.pylibrmm import * From 543a5ed53f4e7b8da557a28df73b4850c50c43b1 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 4 Oct 2024 11:03:06 -0700 Subject: [PATCH 5/5] the last import, fingers crossed --- python/rmm/rmm/_lib/device_buffer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/rmm/rmm/_lib/device_buffer.py b/python/rmm/rmm/_lib/device_buffer.py index 78bd38b83..c531bca5f 100644 --- a/python/rmm/rmm/_lib/device_buffer.py +++ b/python/rmm/rmm/_lib/device_buffer.py @@ -13,6 +13,7 @@ # limitations under the License. from rmm.pylibrmm.device_buffer import ( # noqa: F401 + DeviceBuffer, copy_device_to_ptr, copy_host_to_ptr, copy_ptr_to_host,