From 061ada449b7b11fd39789aca59a65a663f631a3e Mon Sep 17 00:00:00 2001
From: Joseph Nke <76006812+jnke2016@users.noreply.github.com>
Date: Wed, 4 Oct 2023 14:41:24 -0500
Subject: [PATCH] Increase dask-related timeouts for CI testing (#3907)

This PR increases the minimum timeout when waiting for the workers to complete their tasks.

Authors:
  - Joseph Nke (https://github.com/jnke2016)

Approvers:
  - Brad Rees (https://github.com/BradReesWork)
  - Vibhu Jawa (https://github.com/VibhuJawa)
  - Rick Ratzel (https://github.com/rlratzel)
  - Jake Awe (https://github.com/AyodeAwe)

URL: https://github.com/rapidsai/cugraph/pull/3907
---
 ci/test_python.sh | 5 ++++-
 ci/test_wheel.sh  | 6 +++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 825d5b242d5..df0f34377a3 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -65,7 +65,10 @@ popd
 
 rapids-logger "pytest cugraph"
 pushd python/cugraph/cugraph
-export DASK_WORKER_DEVICES="0"
+DASK_WORKER_DEVICES="0" \
+DASK_DISTRIBUTED__SCHEDULER__WORKER_TTL="1000s" \
+DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT="1000s" \
+DASK_CUDA_WAIT_WORKERS_MIN_TIMEOUT=20 \
 pytest \
   -v \
   --benchmark-disable \
diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index 3ac3549f143..d6ec67cd9e9 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -21,5 +21,9 @@ arch=$(uname -m)
 if [[ "${arch}" == "aarch64" && ${RAPIDS_BUILD_TYPE} == "pull-request" ]]; then
     python ./ci/wheel_smoke_test_${package_name}.py
 else
-    RAPIDS_DATASET_ROOT_DIR=`pwd`/datasets python -m pytest ./python/${package_name}/${python_package_name}/tests
+    RAPIDS_DATASET_ROOT_DIR=`pwd`/datasets \
+    DASK_DISTRIBUTED__SCHEDULER__WORKER_TTL="1000s" \
+    DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT="1000s" \
+    DASK_CUDA_WAIT_WORKERS_MIN_TIMEOUT=20 \
+    python -m pytest ./python/${package_name}/${python_package_name}/tests
 fi