Fix spill to device limit (#12252)

Fixing bug in `SpillManager.spill_device_memory` where it would spill more than the limit. 

#

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

Approvers:
  - Ashwin Srinath (https://github.com/shwina)
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #12252

python/cudf/cudf/core/buffer/spill_manager.py

@@ -359,7 +359,7 @@ def spill_device_memory(self, nbytes: int) -> int:
         return spilled
 
     def spill_to_device_limit(self, device_limit: int = None) -> int:
-        """Spill until device limit
+        """Try to spill device memory until device limit
 
         Notice, by default this is a no-op.
 
@@ -380,18 +380,10 @@ def spill_to_device_limit(self, device_limit: int = None) -> int:
         )
         if limit is None:
             return 0
-        ret = 0
-        while True:
-            unspilled = sum(
-                buf.size for buf in self.buffers() if not buf.is_spilled
-            )
-            if unspilled < limit:
-                break
-            nbytes = self.spill_device_memory(nbytes=limit - unspilled)
-            if nbytes == 0:
-                break  # No more to spill
-            ret += nbytes
-        return ret
+        unspilled = sum(
+            buf.size for buf in self.buffers() if not buf.is_spilled
+        )
+        return self.spill_device_memory(nbytes=unspilled - limit)
 
     def __repr__(self) -> str:
         spilled = sum(buf.size for buf in self.buffers() if buf.is_spilled)

python/cudf/cudf/tests/test_spilling.py

@@ -298,6 +298,21 @@ def test_zero_device_limit(manager: SpillManager):
     assert spilled_and_unspilled(manager) == (gen_df_data_nbytes * 2, 0)
 
 
+def test_spill_df_index(manager: SpillManager):
+    df = single_column_df()
+    df.index = [1, 3, 2]  # use a materialized index
+    assert spilled_and_unspilled(manager) == (0, gen_df_data_nbytes * 2)
+
+    manager.spill_to_device_limit(gen_df_data_nbytes)
+    assert spilled_and_unspilled(manager) == (
+        gen_df_data_nbytes,
+        gen_df_data_nbytes,
+    )
+
+    manager.spill_to_device_limit(0)
+    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,