Do not reset CUDA context after UCX tests (#8201)

dask · Sep 25, 2023 · e1387a0 · e1387a0
1 parent b6333df
commit e1387a0
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 12 deletions.
diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py
@@ -179,7 +179,7 @@ async def test_ucx_deserialize(ucx_loop):
     "g",
     [
         lambda cudf: cudf.Series([1, 2, 3]),
-        lambda cudf: cudf.Series([]),
+        lambda cudf: cudf.Series([], dtype=object),
         lambda cudf: cudf.DataFrame([]),
         lambda cudf: cudf.DataFrame([1]).head(0),
         lambda cudf: cudf.DataFrame([1.0]).head(0),

diff --git a/distributed/utils_test.py b/distributed/utils_test.py
@@ -2156,17 +2156,6 @@ def ucx_loop():
     import distributed.comm.ucx
 
     distributed.comm.ucx.ucp = None
-    # If the test created a context, clean it up.
-    # TODO: should we check if there's already a context _before_ the test runs?
-    # I think that would be useful.
-    from distributed.diagnostics.nvml import has_cuda_context
-
-    ctx = has_cuda_context()
-    if ctx.has_context:
-        import numba.cuda
-
-        ctx = numba.cuda.current_context()
-        ctx.device.reset()
 
 
 def wait_for_log_line(