More @acquire_spill_lock() and as_buffer(..., exposed=False) (#12535

)

Now that #12106 is in, wrapping calls to `libcudf` and setting  `as_buffer(..., exposed=False)` is an ongoing effort.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: #12535

python/cudf/cudf/_lib/transform.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2023, NVIDIA CORPORATION.
 
 from numba.np import numpy_support
 
@@ -53,6 +53,7 @@ def bools_to_mask(Column col):
     return buf
 
 
+@acquire_spill_lock()
 def mask_to_bools(object mask_buffer, size_type begin_bit, size_type end_bit):
     """
     Given a mask buffer, returns a boolean column representng bit 0 -> False
@@ -72,6 +73,7 @@ def mask_to_bools(object mask_buffer, size_type begin_bit, size_type end_bit):
     return Column.from_unique_ptr(move(result))
 
 
+@acquire_spill_lock()
 def nans_to_nulls(Column input):
     cdef column_view c_input = input.view()
     cdef pair[unique_ptr[device_buffer], size_type] c_output
@@ -84,9 +86,7 @@ def nans_to_nulls(Column input):
     if c_output.second == 0:
         return None
 
-    buffer = DeviceBuffer.c_from_unique_ptr(move(c_buffer))
-    buffer = as_buffer(buffer)
-    return buffer
+    return as_buffer(DeviceBuffer.c_from_unique_ptr(move(c_buffer)))
 
 
 @acquire_spill_lock()

python/cudf/cudf/core/column/column.py

@@ -1833,7 +1833,7 @@ def as_column(
         ):
             arbitrary = cupy.ascontiguousarray(arbitrary)
 
-        data = as_buffer(arbitrary, exposed=True)
+        data = as_buffer(arbitrary)
         col = build_column(data, dtype=current_dtype, mask=mask)
 
         if dtype is not None:
@@ -2290,7 +2290,7 @@ def _mask_from_cuda_array_interface_desc(obj) -> Union[Buffer, None]:
         typecode = typestr[1]
         if typecode == "t":
             mask_size = bitmask_allocation_size_bytes(nelem)
-            mask = as_buffer(data=ptr, size=mask_size, owner=obj, exposed=True)
+            mask = as_buffer(data=ptr, size=mask_size, owner=obj)
         elif typecode == "b":
             col = as_column(mask)
             mask = bools_to_mask(col)

python/cudf/cudf/tests/test_spilling.py

@@ -1,9 +1,10 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 import importlib
 import random
 import time
 import warnings
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Tuple
 
@@ -313,18 +314,16 @@ def test_spill_df_index(manager: SpillManager):
     assert spilled_and_unspilled(manager) == (gen_df_data_nbytes * 2, 0)
 
 
-def test_external_memory_never_spills(manager):
-    """
-    Test that external data, i.e., data not managed by RMM,
-    is never spilled
-    """
-
+def test_external_memory(manager):
     cupy.cuda.set_allocator()  # uses default allocator
-
-    a = cupy.asarray([1, 2, 3])
-    s = cudf.Series(a)
-    assert len(manager.buffers()) == 0
-    assert not s._data[None].data.spillable
+    cpy = cupy.asarray([1, 2, 3])
+    s = cudf.Series(cpy)
+    # Check that the cupy array is still alive after overwriting `cpy`
+    cpy = weakref.ref(cpy)
+    assert cpy() is not None
+    # Check that the series is spillable and known by the spill manager
+    assert len(manager.buffers()) == 1
+    assert s._data[None].data.spillable
 
 
 def test_spilling_df_views(manager):