From 6cebf2294ff3cb5ed4d4712ccbb26f6a27687ad5 Mon Sep 17 00:00:00 2001
From: GALI PREM SAGAR <sagarprem75@gmail.com>
Date: Fri, 2 Feb 2024 19:46:11 -0600
Subject: [PATCH] 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: https://github.com/rapidsai/cudf/pull/14958
---
 python/cudf/cudf/tests/test_spilling.py | 104 +++++++++++++++---------
 1 file changed, 65 insertions(+), 39 deletions(-)

diff --git a/python/cudf/cudf/tests/test_spilling.py b/python/cudf/cudf/tests/test_spilling.py
index 7e66a7ab4ba..f18cb32a091 100644
--- a/python/cudf/cudf/tests/test_spilling.py
+++ b/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(