Switch pytest-xdist algorithm to worksteal

ci/run_cuml_dask_pytests.sh

@@ -4,4 +4,4 @@
 # Support invoking run_cuml_dask_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cuml/tests/dask
 
-pytest --cache-clear "$@" .
+python -m pytest --cache-clear "$@" .

ci/run_cuml_singlegpu_memleak_pytests.sh

@@ -4,4 +4,4 @@
 # Support invoking run_cuml_singlegpu_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cuml/tests
 
-pytest --cache-clear --ignore=dask -m "memleak" "$@" .
+python -m pytest --cache-clear --ignore=dask -m "memleak" "$@" .

ci/run_cuml_singlegpu_pytests.sh

@@ -4,4 +4,4 @@
 # Support invoking run_cuml_singlegpu_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cuml/tests
 
-pytest --cache-clear --ignore=dask -m "not memleak" "$@" .
+python -m pytest --cache-clear --ignore=dask -m "not memleak" "$@" .

ci/test_python_singlegpu.sh

@@ -11,16 +11,17 @@ EXITCODE=0
 trap "EXITCODE=1" ERR
 set +e
 
-rapids-logger "pytest cuml single GPU..."
+rapids-logger "pytest cuml single GPU"
 ./ci/run_cuml_singlegpu_pytests.sh \
   --numprocesses=8 \
+  --dist=worksteal \
   --junitxml="${RAPIDS_TESTS_DIR}/junit-cuml.xml" \
   --cov-config=../../.coveragerc \
   --cov=cuml \
   --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cuml-coverage.xml" \
-  --cov-report=term \
+  --cov-report=term
 
-rapids-logger "memory leak pytests..."
+rapids-logger "memory leak pytests"
 
 ./ci/run_cuml_singlegpu_memleak_pytests.sh \
   --numprocesses=1 \
@@ -29,7 +30,7 @@ rapids-logger "memory leak pytests..."
   --cov=cuml \
   --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cuml-memleak-coverage.xml" \
   --cov-report=term \
-  -m "memleak" \
+  -m "memleak"
 
 rapids-logger "Test script exiting with value: $EXITCODE"
 exit ${EXITCODE}

ci/test_wheel.sh

@@ -15,9 +15,30 @@ fi
 # echo to expand wildcard before adding `[extra]` requires for pip
 python -m pip install $(echo ./dist/cuml*.whl)[test]
 
+EXITCODE=0
+trap "EXITCODE=1" ERR
+set +e
+
 # Run smoke tests for aarch64 pull requests
 if [[ "$(arch)" == "aarch64" && "${RAPIDS_BUILD_TYPE}" == "pull-request" ]]; then
     python ci/wheel_smoke_test_cuml.py
 else
-    python -m pytest ./python/cuml/tests -k 'not test_sparse_pca_inputs' -n 4 --ignore=python/cuml/tests/dask && python -m pytest ./python/cuml/tests -k 'test_sparse_pca_inputs' && python -m pytest ./python/cuml/tests/dask
+    rapids-logger "pytest cuml single GPU"
+    ./ci/run_cuml_singlegpu_pytests.sh \
+      --numprocesses=8 \
+      --dist=worksteal \
+      -k 'not test_sparse_pca_inputs' \
+      --junitxml="${RAPIDS_TESTS_DIR}/junit-cuml.xml"
+
+    # Run test_sparse_pca_inputs separately
+    ./ci/run_cuml_singlegpu_pytests.sh \
+      -k 'test_sparse_pca_inputs' \
+      --junitxml="${RAPIDS_TESTS_DIR}/junit-cuml-sparse-pca.xml"
+
+    rapids-logger "pytest cuml-dask"
+    ./ci/run_cuml_dask_pytests.sh \
+      --junitxml="${RAPIDS_TESTS_DIR}/junit-cuml-dask.xml"
 fi
+
+rapids-logger "Test script exiting with value: $EXITCODE"
+exit ${EXITCODE}