unset CUDF_SPILL after a pytest (#14958)

This PR updates spilling tests to limit the scope of all environment and cudf option modifications to the scope of the test to avoid interfering with other tests. This PR also temporary skips a test that is not currently safe to run in an environment where other tests may already have modified related global state (the rmm default memory resource).

Authors:
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #14958

python/cudf/cudf/tests/test_spilling.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2022-2024, NVIDIA CORPORATION.
 
+import contextlib
 import importlib
 import random
 import time
@@ -215,45 +216,66 @@ def test_spilling_buffer(manager: SpillManager):
         buf.spill(target="cpu")
 
 
-def test_environment_variables(monkeypatch):
-    def reload_options():
-        # In order to enabling monkey patching of the environment variables
-        # mark the global manager as uninitialized.
-        set_global_manager(None)
-        cudf.core.buffer.spill_manager._global_manager_uninitialized = True
-        importlib.reload(cudf.options)
-
-    monkeypatch.setenv("CUDF_SPILL_ON_DEMAND", "off")
-    monkeypatch.setenv("CUDF_SPILL", "off")
-    reload_options()
-    assert get_global_manager() is None
-
-    monkeypatch.setenv("CUDF_SPILL", "on")
-    reload_options()
-    manager = get_global_manager()
-    assert isinstance(manager, SpillManager)
-    assert manager._spill_on_demand is False
-    assert manager._device_memory_limit is None
-    assert manager.statistics.level == 0
-
-    monkeypatch.setenv("CUDF_SPILL_DEVICE_LIMIT", "1000")
-    reload_options()
-    manager = get_global_manager()
-    assert isinstance(manager, SpillManager)
-    assert manager._device_memory_limit == 1000
-    assert manager.statistics.level == 0
-
-    monkeypatch.setenv("CUDF_SPILL_STATS", "1")
-    reload_options()
-    manager = get_global_manager()
-    assert isinstance(manager, SpillManager)
-    assert manager.statistics.level == 1
-
-    monkeypatch.setenv("CUDF_SPILL_STATS", "2")
-    reload_options()
-    manager = get_global_manager()
-    assert isinstance(manager, SpillManager)
-    assert manager.statistics.level == 2
+def _reload_options():
+    # In order to enabling monkey patching of the environment variables
+    # mark the global manager as uninitialized.
+    set_global_manager(None)
+    cudf.core.buffer.spill_manager._global_manager_uninitialized = True
+    importlib.reload(cudf.options)
+
+
+@contextlib.contextmanager
+def _get_manager_in_env(monkeypatch, var_vals):
+    with monkeypatch.context() as m:
+        for var, val in var_vals:
+            m.setenv(var, val)
+        _reload_options()
+        yield get_global_manager()
+    _reload_options()
+
+
+def test_environment_variables_spill_off(monkeypatch):
+    with _get_manager_in_env(
+        monkeypatch,
+        [("CUDF_SPILL", "off"), ("CUDF_SPILL_ON_DEMAND", "off")],
+    ) as manager:
+        assert manager is None
+
+
+def test_environment_variables_spill_on(monkeypatch):
+    with _get_manager_in_env(
+        monkeypatch,
+        [("CUDF_SPILL", "on")],
+    ) as manager:
+        assert isinstance(manager, SpillManager)
+        assert manager._spill_on_demand is True
+        assert manager._device_memory_limit is None
+        assert manager.statistics.level == 0
+
+
+def test_environment_variables_device_limit(monkeypatch):
+    with _get_manager_in_env(
+        monkeypatch,
+        [("CUDF_SPILL", "on"), ("CUDF_SPILL_DEVICE_LIMIT", "1000")],
+    ) as manager:
+        assert isinstance(manager, SpillManager)
+        assert manager._device_memory_limit == 1000
+        assert manager.statistics.level == 0
+
+
+@pytest.mark.parametrize("level", (1, 2))
+def test_environment_variables_spill_stats(monkeypatch, level):
+    with _get_manager_in_env(
+        monkeypatch,
+        [
+            ("CUDF_SPILL", "on"),
+            ("CUDF_SPILL_DEVICE_LIMIT", "1000"),
+            ("CUDF_SPILL_STATS", f"{level}"),
+        ],
+    ) as manager:
+        assert isinstance(manager, SpillManager)
+        assert manager._device_memory_limit == 1000
+        assert manager.statistics.level == level
 
 
 def test_spill_device_memory(manager: SpillManager):
@@ -507,6 +529,10 @@ def test_serialize_cuda_dataframe(manager: SpillManager):
     assert_eq(df1, df2)
 
 
+@pytest.mark.skip(
+    reason="This test is not safe because other tests may have enabled"
+    "spilling and already modified rmm's global state"
+)
 def test_get_rmm_memory_resource_stack():
     mr1 = rmm.mr.get_current_device_resource()
     assert all(