diff --git a/ci/test_common.sh b/ci/test_common.sh index 2fbc2c17..b7b32eab 100755 --- a/ci/test_common.sh +++ b/ci/test_common.sh @@ -160,6 +160,20 @@ 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 @@ -167,9 +181,19 @@ run_distributed_ucxx_tests() { 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/ +} diff --git a/ci/test_python.sh b/ci/test_python.sh index c45c2cbb..5dbbbd57 100755 --- a/ci/test_python.sh +++ b/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 diff --git a/ci/test_wheel_distributed_ucxx.sh b/ci/test_wheel_distributed_ucxx.sh index c8c0bbca..c624b365 100755 --- a/ci/test_wheel_distributed_ucxx.sh +++ b/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 diff --git a/python/distributed-ucxx/distributed_ucxx/tests/test_ucxx.py b/python/distributed-ucxx/distributed_ucxx/tests/test_ucxx.py index 1f820629..2135b3ba 100644 --- a/python/distributed-ucxx/distributed_ucxx/tests/test_ucxx.py +++ b/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", [ diff --git a/python/distributed-ucxx/distributed_ucxx/tests_internal/conftest.py b/python/distributed-ucxx/distributed_ucxx/tests_internal/conftest.py new file mode 100644 index 00000000..ca71f84e --- /dev/null +++ b/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 diff --git a/python/distributed-ucxx/distributed_ucxx/tests_internal/test_distributed_comms.py b/python/distributed-ucxx/distributed_ucxx/tests_internal/test_distributed_comms.py new file mode 100644 index 00000000..9bc7afa1 --- /dev/null +++ b/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://")