Run tests requiring Distributed internals separately (#183)

Some tests require Distributed internals and installed in developer mode. Since we have no good way to install it with `rapids-dask-dependency`, we now run those separately from the rest.

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Ray Douglass (https://github.com/raydouglass)
  - Charles Blackmon-Luca (https://github.com/charlesbluca)

URL: #183

ci/test_common.sh

@@ -160,16 +160,40 @@ run_py_benchmark() {
 }
 
 ################################## Distributed #################################
+install_distributed_dev_mode() {
+  # Running Distributed tests which access its internals requires installing it in
+  # developer mode. This isn't a great solution but it's what we can currently do
+  # to run non-public API tests in CI.
+
+  rapids-logger "Install Distributed in developer mode"
+  git clone https://github.com/dask/distributed /tmp/distributed
+  pip install -e /tmp/distributed
+  # `pip install -e` removes files under `distributed` but not the directory, later
+  # causing failures to import modules.
+  PYTHON_ENV_PATH=${CONDA_PREFIX:-/pyenv}
+  rm -rf $(find ${PYTHON_ENV_PATH} -type d -iname "site-packages")/distributed
+}
+
 run_distributed_ucxx_tests() {
   PROGRESS_MODE=$1
   ENABLE_DELAYED_SUBMISSION=$2
   ENABLE_PYTHON_FUTURE=$3
 
   CMD_LINE="UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m python -m pytest -vs python/distributed-ucxx/distributed_ucxx/tests/"
 
-  # Workaround for https://github.com/rapidsai/ucxx/issues/15
-  # CMD_LINE="UCX_KEEPALIVE_INTERVAL=1ms ${CMD_LINE}"
-
   log_command "${CMD_LINE}"
   UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m python -m pytest -vs python/distributed-ucxx/distributed_ucxx/tests/
 }
+
+run_distributed_ucxx_tests_internal() {
+  # Note that tests here require Distributed installed in developer mode!
+
+  PROGRESS_MODE=$1
+  ENABLE_DELAYED_SUBMISSION=$2
+  ENABLE_PYTHON_FUTURE=$3
+
+  CMD_LINE="UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m python -m pytest -vs python/distributed-ucxx/distributed_ucxx/tests_internal/"
+
+  log_command "${CMD_LINE}"
+  UCXPY_PROGRESS_MODE=${PROGRESS_MODE} UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 10m python -m pytest -vs python/distributed-ucxx/distributed_ucxx/tests_internal/
+}

ci/test_python.sh

@@ -29,12 +29,6 @@ rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   libucxx ucxx distributed-ucxx
 
-# TODO: Perhaps install from conda? We need distributed installed in developer
-# mode to provide test utils, but that's probably not doable from conda packages.
-rapids-logger "Install Distributed in developer mode"
-git clone https://github.com/dask/distributed /tmp/distributed
-pip install -e /tmp/distributed
-
 print_ucx_config
 
 rapids-logger "Run tests with conda package"
@@ -69,5 +63,14 @@ run_distributed_ucxx_tests      thread          0                           1
 run_distributed_ucxx_tests      thread          1                           0
 run_distributed_ucxx_tests      thread          1                           1
 
+install_distributed_dev_mode
+
+# run_distributed_ucxx_tests_internal   PROGRESS_MODE   ENABLE_DELAYED_SUBMISSION   ENABLE_PYTHON_FUTURE
+run_distributed_ucxx_tests_internal     polling         0                           0
+run_distributed_ucxx_tests_internal     thread          0                           0
+run_distributed_ucxx_tests_internal     thread          0                           1
+run_distributed_ucxx_tests_internal     thread          1                           0
+run_distributed_ucxx_tests_internal     thread          1                           1
+
 rapids-logger "C++ future -> Python future notifier example"
 python -m ucxx.examples.python_future_task_example

ci/test_wheel_distributed_ucxx.sh

@@ -18,12 +18,6 @@ python -m pip install ./local-ucxx-dep/ucxx*.whl
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install $(echo ./dist/${PROJECT_NAME}*.whl)[test]
 
-# TODO: We need distributed installed in developer mode to provide test utils,
-# we still need to match to the `rapids-dask-dependency` version.
-rapids-logger "Install Distributed in developer mode"
-git clone https://github.com/dask/distributed /tmp/distributed
-python -m pip install -e /tmp/distributed
-
 # Run smoke tests for aarch64 pull requests
 if [[ "$(arch)" == "aarch64" && "${RAPIDS_BUILD_TYPE}" == "pull-request" ]]; then
   rapids-logger "Distributed Smoke Tests"
@@ -33,4 +27,9 @@ else
 
   # run_distributed_ucxx_tests    PROGRESS_MODE   ENABLE_DELAYED_SUBMISSION   ENABLE_PYTHON_FUTURE
   run_distributed_ucxx_tests      thread          1                           1
+
+  install_distributed_dev_mode
+
+  # run_distributed_ucxx_tests_internal PROGRESS_MODE   ENABLE_DELAYED_SUBMISSION   ENABLE_PYTHON_FUTURE
+  run_distributed_ucxx_tests_internal   thread          1                           1
 fi

python/distributed-ucxx/distributed_ucxx/tests/test_ucxx.py

@@ -165,16 +165,6 @@ async def test_ping_pong_data(ucxx_loop):
     await serv_com.close()
 
 
-@gen_test()
-async def test_ucxx_deserialize(ucxx_loop):
-    # Note we see this error on some systems with this test:
-    # `socket.gaierror: [Errno -5] No address associated with hostname`
-    # This may be due to a system configuration issue.
-    from distributed.comm.tests.test_comms import check_deserialize
-
-    await check_deserialize("tcp://")
-
-
 @pytest.mark.parametrize(
     "g",
     [

python/distributed-ucxx/distributed_ucxx/tests_internal/conftest.py

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Make all fixtures available
+from distributed_ucxx.utils_test import *  # noqa

python/distributed-ucxx/distributed_ucxx/tests_internal/test_distributed_comms.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from distributed.comm.tests.test_comms import check_deserialize
+
+from distributed_ucxx.utils_test import gen_test
+
+
+@gen_test()
+async def test_ucxx_deserialize(ucxx_loop):
+    # Note we see this error on some systems with this test:
+    # `socket.gaierror: [Errno -5] No address associated with hostname`
+    # This may be due to a system configuration issue.
+
+    await check_deserialize("ucxx://")