From 17edaa934fbea80df18f5cc48f17af46bac6c6be Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Thu, 20 Jun 2024 14:30:22 +0000 Subject: [PATCH 1/2] GH-41126: [Python] Test Buffer device/device_type access on CUDA --- python/pyarrow/tests/test_cuda.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 400db2643bd..1a287ecaff4 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -534,6 +534,25 @@ def put(*args, **kwargs): 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 + 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 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) From 44d430a474b12959b1e82a9cd1094fd2df0879e6 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Thu, 8 Aug 2024 16:18:07 +0200 Subject: [PATCH 2/2] add comment --- python/pyarrow/tests/test_cuda.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 5aa0a443506..36b97a62064 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -553,11 +553,14 @@ def test_buffer_device(): 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