only test hvd on GPU

tensorflow_recommenders_addons/dynamic_embedding/python/keras/callbacks_test.py

@@ -1,15 +1,13 @@
 import pytest
+from tensorflow.python.framework import test_util
+
 from tensorflow_recommenders_addons.dynamic_embedding.python.keras.callbacks import \
   DEHvdBroadcastGlobalVariablesCallback
 
 
-@pytest.fixture
-def broadcast_callback(root_rank=0, device='/gpu:0'):
-  # Instantiate with the corrected parameters
-  return DEHvdBroadcastGlobalVariablesCallback(root_rank=root_rank,
-                                               device=device)
-
-
-def test_on_batch_end_subsequent_calls(broadcast_callback):
-  broadcast_callback.broadcast_done = True
-  broadcast_callback.on_batch_end(1)
+def test_on_batch_end_subsequent_calls():
+  if test_util.is_gpu_available():
+    broadcast_callback = DEHvdBroadcastGlobalVariablesCallback(root_rank=0,
+                                                               device='/gpu:0')
+    broadcast_callback.broadcast_done = True
+    broadcast_callback.on_batch_end(1)

tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py

@@ -33,12 +33,10 @@
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework.errors_impl import NotFoundError
 from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 from tensorflow.python.training import adam
 from tensorflow.python.training import checkpoint_management
 from tensorflow.python.training import monitored_session
-from tensorflow.python.training.optimizer import Optimizer as tf1_opt
 from tensorflow.python.training import training_util
 try:
   from tensorflow.keras.optimizers.legacy import Adam
@@ -77,7 +75,8 @@ def test_adam_minimize_trainable(self):
       self.skipTest(
           "Apple silicon devices don't support synchronous training based on Horovod."
       )
-
+    if not is_gpu_available:
+      self.skipTest('Only test when gpu is available.')
     base_opt = adam.AdamOptimizer(1.0)
     test_opt = adam.AdamOptimizer(1.0)
     self.common_minimize_trainable_v1(base_opt, test_opt, name="adam")

tools/docker/cpu_tests.Dockerfile

@@ -4,7 +4,6 @@ FROM python:3.9 as build_wheel
 ARG TF_VERSION="2.15.1"
 ARG PY_VERSION="3.9"
 ARG MPI_VERSION="4.1.1"
-ARG HOROVOD_VERSION="0.28.1"
 
 RUN pip install --upgrade pip
 RUN pip install --default-timeout=1000 tensorflow==$TF_VERSION
@@ -16,9 +15,6 @@ RUN apt-get update && apt-get install -y sudo rsync cmake openmpi-bin libopenmpi
 COPY tools/install_deps/install_bazelisk.sh /install/
 RUN bash /install/install_bazelisk.sh
 
-COPY tools/docker/install/install_horovod.sh /install/
-RUN  /install/install_horovod.sh $HOROVOD_VERSION --only-cpu
-
 COPY requirements.txt ./
 RUN pip install -r requirements.txt
 

tools/testing/build_and_run_tests.sh

@@ -49,19 +49,22 @@ if ! [ -x "$(command -v nvidia-smi)" ]; then
   EXTRA_ARGS="-n auto"
 fi
 
-# Lack of HorovodJoin CPU kernels when install Horovod with NCCL
-if [ "$(uname)" != "Darwin" ]; then
-  # Mac only with MPI
-  python -m pip uninstall horovod -y
-  bash /install/install_horovod.sh $HOROVOD_VERSION --only-cpu
-fi
-# TODO(jamesrong): Test on GPU.
-CUDA_VISIBLE_DEVICES="" mpirun -np 2 -H localhost:2 --allow-run-as-root pytest -v ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
-# Reinstall Horovod after tests
-if [ "$(uname)" != "Darwin" ]; then
-  # Mac only with MPI
-  python -m pip uninstall horovod -y
-  bash /install/install_horovod.sh $HOROVOD_VERSION
+# only test and run horovod on GPU
+if [ "$TF_NEED_CUDA" -ne 0 ]; then
+  # Lack of HorovodJoin CPU kernels when install Horovod with NCCL
+  if [ "$(uname)" != "Darwin" ]; then
+    # Mac only with MPI
+    python -m pip uninstall horovod -y
+    bash /install/install_horovod.sh $HOROVOD_VERSION --only-cpu
+  fi
+  # TODO(jamesrong): Test on GPU.
+  CUDA_VISIBLE_DEVICES="" mpirun -np 2 -H localhost:2 --allow-run-as-root pytest -v ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
+  # Reinstall Horovod after tests
+  if [ "$(uname)" != "Darwin" ]; then
+    # Mac only with MPI
+    python -m pip uninstall horovod -y
+    bash /install/install_horovod.sh $HOROVOD_VERSION
+  fi
 fi
 
 IGNORE_HKV=""