add run_benchmark_pytests.sh script

rapidsai · Feb 7, 2024 · 8d8faf3 · 8d8faf3
1 parent 3eeee55
commit 8d8faf3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/ci/run_benchmark_pytests.sh b/ci/run_benchmark_pytests.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+# Support invoking run_pytests.sh outside the script directory
+cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
+
+export UCX_TLS="${UCX_TLS:-tcp,cuda_copy}"
+
+# cd to root directory to prevent repo's `ucp` directory from being used
+# in subsequent commands
+pushd /
+timeout 1m python -m ucp.benchmarks.send_recv -o cupy --server-dev 0 --client-dev 0 --reuse-alloc --backend ucp-async
+timeout 1m python -m ucp.benchmarks.send_recv -o cupy --server-dev 0 --client-dev 0 --reuse-alloc --backend ucp-core
+timeout 1m python -m ucp.benchmarks.cudf_merge --chunks-per-dev 4 --chunk-size 10000 --rmm-init-pool-size 2097152
+popd
diff --git a/ci/test_python.sh b/ci/test_python.sh
@@ -45,11 +45,7 @@ run_tests() {
   rapids-logger "Run local benchmark"
   # cd to root directory to prevent repo's `ucp` directory from being used
   # in subsequent commands
-  pushd /
-  timeout 1m python -m ucp.benchmarks.send_recv -o cupy --server-dev 0 --client-dev 0 --reuse-alloc --backend ucp-async
-  timeout 1m python -m ucp.benchmarks.send_recv -o cupy --server-dev 0 --client-dev 0 --reuse-alloc --backend ucp-core
-  timeout 1m python -m ucp.benchmarks.cudf_merge --chunks-per-dev 4 --chunk-size 10000 --rmm-init-pool-size 2097152
-  popd
+  ./ci/run_benchmark_pytests.sh
 }
 
 rapids-logger "Downloading artifacts from previous jobs"