Update CI to use TF 2.16 (#2378)

* Use TF 2.16 in CI * Use TF 2.16 in CI * Update to TF 2.16 * Increase time for tests to 90 mins * Increase time for tests to 90 mins
keras-team · Mar 20, 2024 · 301f7cd · 301f7cd
1 parent e21312a
commit 301f7cd
Show file tree

Hide file tree

Showing 13 changed files with 25 additions and 27 deletions.
diff --git a/.kokoro/github/ubuntu/gpu/build.sh b/.kokoro/github/ubuntu/gpu/build.sh
@@ -17,9 +17,7 @@ nvidia-smi
 nvcc --version
 
 cd "src/github/keras-cv"
-pip install -U pip setuptools
-# psutil is used by background log reader
-pip install -U psutil
+pip install -U pip setuptools psutil
 
 if [ "${KERAS2:-0}" == "1" ]
 then

diff --git a/.kokoro/github/ubuntu/gpu/jax/continuous.cfg b/.kokoro/github/ubuntu/gpu/jax/continuous.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "jax"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/jax/presubmit.cfg b/.kokoro/github/ubuntu/gpu/jax/presubmit.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "jax"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg b/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "1"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg b/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "1"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg b/.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "tensorflow"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg b/.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "tensorflow"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/torch/continuous.cfg b/.kokoro/github/ubuntu/gpu/torch/continuous.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "torch"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/.kokoro/github/ubuntu/gpu/torch/presubmit.cfg b/.kokoro/github/ubuntu/gpu/torch/presubmit.cfg
@@ -12,5 +12,5 @@ env_vars: {
    value: "torch"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 90 mins from default 180 mins
+timeout_mins: 90
diff --git a/requirements-jax-cuda.txt b/requirements-jax-cuda.txt
@@ -1,5 +1,5 @@
-# Tensorflow cpu-only version.
-tf-nightly-cpu==2.16.0.dev20240209  # Pin a working nightly until rc0.
+# TensorFlow cpu-only version.
+tensorflow-cpu>=2.16.1
 
 # Torch cpu-only version.
 --extra-index-url https://download.pytorch.org/whl/cpu

diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
@@ -1,5 +1,5 @@
-# Tensorflow with cuda support.
-tf-nightly[and-cuda]==2.16.0.dev20240209  # Pin a working nightly until rc0.
+# TensorFlow with cuda support.
+tensorflow[and-cuda]>=2.16.1
 
 # Torch cpu-only version.
 --extra-index-url https://download.pytorch.org/whl/cpu

diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
@@ -1,5 +1,5 @@
-# Tensorflow cpu-only version.
-tf-nightly-cpu==2.16.0.dev20240209  # Pin a working nightly until rc0.
+# TensorFlow cpu-only version.
+tensorflow-cpu>=2.16.1
 
 # Torch with cuda support.
 --extra-index-url https://download.pytorch.org/whl/cu121

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
-# Tensorflow.
-tf-nightly-cpu==2.16.0.dev20240209  # Pin a working nightly until rc0.
+# TensorFlow.
+tensorflow-cpu>=2.16.1
 
 # Torch.
 --extra-index-url https://download.pytorch.org/whl/cpu