This is an automated email from the ASF dual-hosted git repository.
jorisvandenbossche pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 4314fd7d84 GH-41126: [Python] Test Buffer device/device_type access on
CUDA (#42221)
4314fd7d84 is described below
commit 4314fd7d849497a638c2c8180a39c37ea0a88de1
Author: Joris Van den Bossche <[email protected]>
AuthorDate: Fri Aug 9 19:40:11 2024 +0200
GH-41126: [Python] Test Buffer device/device_type access on CUDA (#42221)
### Rationale for this change
Adding tests for the new Buffer properties added in
https://github.com/apache/arrow/pull/41685 but now testing that it works out of
the box with CUDA.
* GitHub Issue: #41126
Authored-by: Joris Van den Bossche <[email protected]>
Signed-off-by: Joris Van den Bossche <[email protected]>
---
python/pyarrow/tests/test_cuda.py | 22 ++++++++++++++++++++++
1 file changed, 22 insertions(+)
diff --git a/python/pyarrow/tests/test_cuda.py
b/python/pyarrow/tests/test_cuda.py
index 8749ab29d0..36b97a6206 100644
--- a/python/pyarrow/tests/test_cuda.py
+++ b/python/pyarrow/tests/test_cuda.py
@@ -545,6 +545,28 @@ def test_copy_from_host(size):
put(position=position, nbytes=nbytes)
+def test_buffer_device():
+ buf = cuda.new_host_buffer(10)
+ assert buf.device_type == pa.DeviceAllocationType.CUDA_HOST
+ assert isinstance(buf.device, pa.Device)
+ assert isinstance(buf.memory_manager, pa.MemoryManager)
+ assert buf.is_cpu
+ assert buf.device.is_cpu
+ assert buf.device == pa.default_cpu_memory_manager().device
+ # it is not entirely clear if CudaHostBuffer should use the default CPU
memory
+ # manager (as it does now), see https://github.com/apache/arrow/pull/42221
+ assert buf.memory_manager.is_cpu
+
+ _, buf = make_random_buffer(size=10, target='device')
+ assert buf.device_type == pa.DeviceAllocationType.CUDA
+ assert isinstance(buf.device, pa.Device)
+ assert buf.device == global_context.memory_manager.device
+ assert isinstance(buf.memory_manager, pa.MemoryManager)
+ assert not buf.is_cpu
+ assert not buf.device.is_cpu
+ assert not buf.memory_manager.is_cpu
+
+
def test_BufferWriter():
def allocate(size):
cbuf = global_context.new_buffer(size)
