UCX tests broken after `distributed.utils_test.cleanup` import removed

After #4888 which removed from distributed.utils_test import cleanup from files, the UCX test now fails as below:

________________________________________________________ ERROR at setup of test_large_cupy[1000000000] _________________________________________________________file /datasets/pentschev/miniconda3/envs/ucx-1.9-112-21.08.210608/lib/python3.8/site-packages/distributed/comm/tests/test_ucx.py, line 228
  @pytest.mark.slow
  @pytest.mark.asyncio
  @pytest.mark.parametrize("n", [int(1e9), int(2.5e9)])
  async def test_large_cupy(n, cleanup):
      cupy = pytest.importorskip("cupy")
      com, serv_com = await get_comm_pair()

      arr = cupy.ones(n, dtype="u1")                                                                                                                                  msg = {"op": "ping", "data": to_serialize(arr)}

      _, result = await asyncio.gather(com.write(msg), serv_com.read())                                                                                               data2 = result.pop("data")

      assert result["op"] == "ping"                                                                                                                                   assert len(data2) == len(arr)
      await com.close()
      await serv_com.close()                                                                                                                                    E       fixture 'cleanup' not found
>       available fixtures: cache, capfd, capfdbinary, caplog, capsys, capsysbinary, doctest_namespace, event_loop, monkeypatch, pytestconfig, record_property,
record_testsuite_property, record_xml_attribute, recwarn, tmp_path, tmp_path_factory, tmpdir, tmpdir_factory, unused_tcp_port, unused_tcp_port_factory          >       use 'pytest --fixtures [testpath]' for help on them.

/datasets/pentschev/miniconda3/envs/ucx-1.9-112-21.08.210608/lib/python3.8/site-packages/distributed/comm/tests/test_ucx.py:228

This is a file which isn’t tested by Distributed’s CI currently, which is supposedly the reason the error above wasn’t catched. However, it seems that https://github.com/dask/distributed/pull/4888/files#diff-a31c7ed5d35f5ed8233994868c54d625b18e6bacb6794344c4531e62bd9dde59R19-R20 should make the fixture available for all tests, I’m not sure though why that isn’t the case for that file in particular.

@fjetter do you happen to have any ideas how to resolve the above without adding the import back? I verified that if I add the import back it then works.

[jakirkham]

Peter is currently OOTO. Let’s wait until he’s back

[pentschev]

Yup, all good from my side. Thanks @jakirkham and @jrbourbeau !