Set UCX_TCP_KEEPINTVL for Python benchmarks in CI

Workaround for rapidsai#15
pentschev · Jun 14, 2023 · 44f6d68 · 44f6d68
1 parent 874b4c3
commit 44f6d68
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/ci/test_python.sh b/ci/test_python.sh
@@ -60,15 +60,19 @@ run_py_benchmark() {
 
   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}"
 
+  # Workaround for https://github.com/rapidsai/ucxx/issues/15
+  CMD_LINE="UCX_KEEPALIVE_INTERVAL=1ms ${CMD_LINE}"
+
   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}
+
+  UCX_KEEPALIVE_INTERVAL=1ms 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}
 }
 
 rapids-logger "Downloading artifacts from previous jobs"
-CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+CPP_CHANNEL=${RAPIDS_CONDA_BLD_OUTPUT_DIR}
 
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \