Adds Kokoro GPU Tests (#2224)

* Adds Kokoro tests

* Add Kokoro Tests

.kokoro/README.md

@@ -0,0 +1 @@
+CI to run on PR and merge to Master.

.kokoro/github/ubuntu/gpu/build.sh

@@ -0,0 +1,83 @@
+set -e
+set -x
+
+cd "${KOKORO_ROOT}/"
+
+sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
+
+PYTHON_BINARY="/usr/bin/python3.9"
+
+"${PYTHON_BINARY}" -m venv venv
+source venv/bin/activate
+# Check the python version
+python --version
+python3 --version
+
+export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:"
+# Check cuda
+nvidia-smi
+nvcc --version
+
+cd "src/github/keras-cv"
+pip install -U pip setuptools
+
+if [ "${KERAS2:-0}" == "1" ]
+then
+   echo "Keras2 detected."
+   pip install -r requirements-common.txt --progress-bar off
+   pip install tensorflow~=2.14
+   pip install --extra-index-url https://download.pytorch.org/whl/cpu torch==2.1.0+cpu
+   pip install torchvision~=0.16.0
+   pip install "jax[cpu]"
+
+elif [ "$KERAS_BACKEND" == "tensorflow" ]
+then
+   echo "TensorFlow backend detected."
+   pip install -r requirements-tensorflow-cuda.txt --progress-bar off
+
+elif [ "$KERAS_BACKEND" == "jax" ]
+then
+   echo "JAX backend detected."
+   pip install -r requirements-jax-cuda.txt --progress-bar off
+
+elif [ "$KERAS_BACKEND" == "torch" ]
+then
+   echo "PyTorch backend detected."
+   pip install -r requirements-torch-cuda.txt --progress-bar off
+fi
+
+pip install --no-deps -e "." --progress-bar off
+
+# Run Extra Large Tests for Continuous builds
+if [ "${RUN_XLARGE:-0}" == "1" ]
+then
+   pytest --check_gpu --run_large --run_extra_large --durations 0 \
+      keras_cv/bounding_box \
+      keras_cv/callbacks \
+      keras_cv/losses \
+      keras_cv/layers/object_detection \
+      keras_cv/layers/preprocessing \
+      keras_cv/models/backbones \
+      keras_cv/models/classification \
+      keras_cv/models/object_detection/retinanet \
+      keras_cv/models/object_detection/yolo_v8 \
+      keras_cv/models/object_detection_3d \
+      keras_cv/models/segmentation \
+      keras_cv/models/stable_diffusion \
+      --cov=keras-cv
+else
+   pytest --check_gpu --run_large --durations 0 \
+      keras_cv/bounding_box \
+      keras_cv/callbacks \
+      keras_cv/losses \
+      keras_cv/layers/object_detection \
+      keras_cv/layers/preprocessing \
+      keras_cv/models/backbones \
+      keras_cv/models/classification \
+      keras_cv/models/object_detection/retinanet \
+      keras_cv/models/object_detection/yolo_v8 \
+      keras_cv/models/object_detection_3d \
+      keras_cv/models/segmentation \
+      keras_cv/models/stable_diffusion \
+      --cov=keras-cv
+fi

.kokoro/github/ubuntu/gpu/jax/continuous.cfg

@@ -0,0 +1,18 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "jax"
+}
+
+env_vars: {
+   key: "RUN_XLARGE"
+   value: "1"
+}

.kokoro/github/ubuntu/gpu/jax/presubmit.cfg

@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "jax"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60

.kokoro/github/ubuntu/gpu/keras2/continuous.cfg

@@ -0,0 +1,18 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS2"
+   value: "1"
+}
+
+env_vars: {
+   key: "RUN_XLARGE"
+   value: "1"
+}

.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg

@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS2"
+   value: "1"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60

.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg

@@ -0,0 +1,18 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "tensorflow"
+}
+
+env_vars: {
+   key: "RUN_XLARGE"
+   value: "1"
+}

.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg

@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "tensorflow"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60

.kokoro/github/ubuntu/gpu/torch/continuous.cfg

@@ -0,0 +1,18 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "torch"
+}
+
+env_vars: {
+   key: "RUN_XLARGE"
+   value: "1"
+}

.kokoro/github/ubuntu/gpu/torch/presubmit.cfg

@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "torch"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60

keras_cv/conftest.py

@@ -17,6 +17,7 @@
 import tensorflow as tf
 from packaging import version
 
+from keras_cv.backend import config as backend_config
 from keras_cv.backend.config import keras_3
 
 
@@ -33,9 +34,35 @@ def pytest_addoption(parser):
         default=False,
         help="run extra_large tests",
     )
+    parser.addoption(
+        "--check_gpu",
+        action="store_true",
+        default=False,
+        help="fail if a gpu is not present",
+    )
 
 
 def pytest_configure(config):
+    # Verify that device has GPU and detected by backend
+    if config.getoption("--check_gpu"):
+        found_gpu = False
+        backend = backend_config.backend()
+        if backend == "jax":
+            import jax
+
+            try:
+                found_gpu = bool(jax.devices("gpu"))
+            except RuntimeError:
+                found_gpu = False
+        elif backend == "tensorflow":
+            found_gpu = bool(tf.config.list_logical_devices("GPU"))
+        elif backend == "torch":
+            import torch
+
+            found_gpu = bool(torch.cuda.device_count())
+        if not found_gpu:
+            pytest.fail(f"No GPUs discovered on the {backend} backend.")
+
     config.addinivalue_line(
         "markers", "large: mark test as being slow or requiring a network"
     )