From 0c1f101922685ed7419da6cd80c68007240e4c26 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Thu, 27 Apr 2023 04:40:30 -0700
Subject: [PATCH] Temporarily increase Python benchmark timeout

Increase the timeout of the Python benchmarks with `thread` progress
mode to avoid hitting https://github.com/rapidsai/ucxx/issues/15 . This
can be reverted once that issue is resolved.
---
 ci/test_python.sh | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 82be0e8a..1b936770 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -58,13 +58,22 @@ run_py_benchmark() {
     ASYNCIO_WAIT=""
   fi
 
-  CMD_LINE="UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 2m python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}"
+  if [ $PROGRESS_MODE == "thread" ]; then
+    # Increased timeout to prevent issues from
+    # https://github.com/rapidsai/ucxx/issues/15 .
+    # TODO: keep the original 2m timeout when the issue above is resolved.
+    TIMEOUT=5m
+  else
+    TIMEOUT=2m
+  fi
+
+  CMD_LINE="UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout ${TIMEOUT} python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}"
 
   rapids-logger "Running: ${CMD_LINE}"
   if [ $SLOW -ne 0 ]; then
     rapids-logger "SLOW BENCHMARK: it may seem like a deadlock but will eventually complete."
   fi
-  UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout 2m python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}
+  UCXPY_ENABLE_DELAYED_SUBMISSION=${ENABLE_DELAYED_SUBMISSION} UCXPY_ENABLE_PYTHON_FUTURE=${ENABLE_PYTHON_FUTURE} timeout ${TIMEOUT} python -m ucxx.benchmarks.send_recv --backend ${BACKEND} -o cupy --reuse-alloc -n 8MiB --n-buffers $N_BUFFERS --progress-mode ${PROGRESS_MODE} ${ASYNCIO_WAIT}
 }
 
 rapids-logger "Downloading artifacts from previous jobs"