-
Notifications
You must be signed in to change notification settings - Fork 4k
GH-42112: [Python] Array gracefully fails on non-cpu device #42113
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
GH-42112: [Python] Array gracefully fails on non-cpu device #42113
Conversation
|
|
|
@github-actions crossbow submit test-cuda-python |
|
|
@github-actions crossbow submit test-cuda-python |
|
Revision: 534f33e9f1b09dabe482404173fb44d6e75f5d37 Submitted crossbow builds: ursacomputing/crossbow @ actions-ac535722c7
|
|
Revision: 9f3efef9c738808ac4976b930a215cefa67d1839 Submitted crossbow builds: ursacomputing/crossbow @ actions-8413b104ec
|
jorisvandenbossche
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks Dane, looking good!
python/pyarrow/array.pxi
Outdated
| self._assert_cpu() | ||
| cdef int64_t total_buffer_size | ||
| total_buffer_size = TotalBufferSize(deref(self.ap)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My first thought was that getting the buffer size should be possible without looking at the actual data. But so it seems that to avoid counting identical buffers twice (if they are reused in a single array, which is possible), it uses the buffer's address to distinguish buffers. Currently that uses buffer->data() in DoTotalBufferSize. data() will return null for non-cpu data, but since it doesn't actually use that address, I think this could also use buffer->address() which will return the address always even for non-cpu data.
(but that could be fine for a follow-up as well)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That checks out! I tested it locally and witnessed a segfault, but hadn't dug into the reasoning.
python/pyarrow/array.pxi
Outdated
| for i in range(len(self)): | ||
| yield self.getitem(i) | ||
|
|
||
| def __repr__(self): | ||
| self._assert_cpu() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This one should not be needed because this ends up calling to_string which already has it, I think
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point!
python/pyarrow/array.pxi
Outdated
| @@ -1307,6 +1319,8 @@ cdef class Array(_PandasConvertible): | |||
| ------- | |||
| bool | |||
| """ | |||
| self._assert_cpu() | |||
| self.other._assert_cpu() | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| self.other._assert_cpu() | |
| other._assert_cpu() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LOL, thank you. Nice catch.
python/pyarrow/array.pxi
Outdated
| @@ -1404,8 +1424,9 @@ cdef class Array(_PandasConvertible): | |||
| ------- | |||
| sliced : RecordBatch | |||
| """ | |||
| cdef: | |||
| shared_ptr[CArray] result | |||
| self._assert_cpu() | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think that slicing actually works?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You are right. I tried printing the resulting RecordBatch, which segfaults. But .slice() itself does not segfault.
python/pyarrow/array.pxi
Outdated
| return self.device_type == DeviceAllocationType.CPU | ||
|
|
||
| def _assert_cpu(self): | ||
| if not self.is_cpu: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we would want to speed this up a bit (it's called in many places, although typically should only give tiny overhead), you could also inline similar code as you used above for the python attributes, like if self.sp_array.get().device_type() != CDeviceAllocationType_kCPU: .., and make it a cdef instead of def (not entirely sure this is worth it)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think its worth doing, good suggestion
python/pyarrow/tests/test_array.py
Outdated
|
|
||
| # Supported | ||
| arr.validate() | ||
| arr.validate(full=True) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am wondering if this will actually be supported for more complex data types. For example for variable size binary, it will check the actual offsets if they are correct numbers (eg not negative, not out of bounds, increasing, etc)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great catch, let me experiment and find out.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You are right! There is an abort with validate(full=True) when using variable binary type.
Check failed: is_cpu() not a CPU buffer (device: CudaDevice(device_number=0, name="NVIDIA RTX A5000"))
Aborted
python/pyarrow/tests/test_array.py
Outdated
| with pytest.raises(NotImplementedError): | ||
| arr.fill_null(0) | ||
| with pytest.raises(NotImplementedError): | ||
| arr.__getitem__(0) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| arr.__getitem__(0) | |
| arr[0] |
python/pyarrow/tests/test_array.py
Outdated
| with pytest.raises(NotImplementedError): | ||
| arr.getitem(0) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This one can be removed (getitem is a cdef function which is not callable from python, so this will error for that reason), this is tested through testing __getitem__
python/pyarrow/tests/test_array.py
Outdated
| arr.__dlpack__() | ||
| with pytest.raises(NotImplementedError): | ||
| arr.__dlpack_device__() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe add a TODO comment here that this should be supported in the future
|
@github-actions crossbow submit test-cuda-python |
|
Revision: ace94eccef774383124243b77f68e5f1912ce42b Submitted crossbow builds: ursacomputing/crossbow @ actions-570bed1718
|
|
@github-actions crossbow submit test-cuda-python |
|
Revision: abe942e0e68b176144d8b156b77bfc330bcd37a0 Submitted crossbow builds: ursacomputing/crossbow @ actions-26aaa9352b
|
ff03596 to
d5a547f
Compare
|
@github-actions crossbow submit test-cuda-python |
|
Revision: d5a547f Submitted crossbow builds: ursacomputing/crossbow @ actions-b8e234a0f3
|
jorisvandenbossche
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Updates look good! There is still some test failure
Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
|
@github-actions crossbow submit test-cuda-python |
|
Revision: 1c903e2 Submitted crossbow builds: ursacomputing/crossbow @ actions-2a8fac0c0e
|
|
@github-actions crossbow submit test-cuda-python |
|
Revision: c12765a Submitted crossbow builds: ursacomputing/crossbow @ actions-599aea54a2
|
|
@github-actions crossbow submit test-cuda-python |
|
Revision: 9740a90 Submitted crossbow builds: ursacomputing/crossbow @ actions-6cbf11bdde
|
jorisvandenbossche
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good, thanks!
|
After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit a42ec1d. There were no benchmark performance regressions. 🎉 The full Conbench report has more details. It also includes information about 6 possible false positives for unstable benchmarks that are known to sometimes produce them. |
Rationale for this change
Common
ArrayAPIs should not segfault or abort on non-cpu devices.What changes are included in this PR?
device_typeandis_cpumethods added to theArrayclassAre these changes tested?
Are there any user-facing changes?
device_typeandis_cpumethods added to theArrayclass