Update CI to use TF 2.16 (#2378)

sampathweb · web-flow · commit 301f7cd565ec · 2024-03-20T10:18:39.000-07:00
* 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
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

Original file line number	Diff line number	Diff line change
`@@ -12,5 +12,5 @@ env_vars: {`
`12`	`12`	`value: "jax"`
`13`	`13`	`}`
`14`	`14`
`15`		`-# Set timeout to 60 mins from default 180 mins`
`16`		`-timeout_mins: 60`
	`15`	`+# Set timeout to 90 mins from default 180 mins`
	`16`	`+timeout_mins: 90`