Automate the update image flow

.cloudbuild/README.md

@@ -25,44 +25,18 @@ To authorize yourself with `keras-team-test`, run:
 gcloud config set project keras-team-test
 ```
 
-To add/update dependency for GPU tests:
-- Add/update dependency to `requirements.txt`
-- Create a `Dockerfile` with the following contents:
-  ```
-  FROM tensorflow/tensorflow:2.12.0-gpu
-  RUN apt-get -y update
-  RUN apt-get -y install git
-  RUN git clone https://github.com/keras-team/keras-nlp.git
-  RUN cd keras-nlp
-  RUN pip install -r keras-nlp/requirements.txt
-  ```
-- Run the following command from the directory with your `Dockerfile`:
-  ```
-  gcloud builds submit --region=us-west1 --tag us-west1-docker.pkg.dev/keras-team-test/keras-nlp-test/keras-nlp-image-tensorflow:deps --timeout=30m
-  ```
-- Repeat the last two steps for Jax and Torch (replacing "tensorflow" with "jax"
- or "torch" in the docker image target name). `Dockerfile` for jax:
-  ```
-  FROM nvidia/cuda:11.7.1-base-ubuntu20.04
-  RUN apt-get update
-  RUN apt-get install -y python3 python3-pip
-  RUN apt-get install -y git
-  RUN git clone https://github.com/keras-team/keras-nlp.git
-  RUN cd keras-nlp
-  RUN pip install -r keras-nlp/requirements.txt
-  RUN pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-  ```
-  and for torch:
-  ```
-  FROM nvidia/cuda:11.7.1-base-ubuntu20.04
-  RUN apt-get update
-  RUN apt-get install -y python3 python3-pip
-  RUN apt-get install -y git
-  RUN git clone https://github.com/keras-team/keras-nlp.git
-  RUN cd keras-nlp
-  RUN pip install -r keras-nlp/requirements.txt
-  RUN pip install torch
-  ```
+To add/update dependency for GPU tests for a given framework:
+- Add/update dependencies in `requirements.txt`.
+- Add/update dependencies in `.cloudbuild/{framework}/Dockerfile`.
+- Run the following:
+```
+gcloud builds submit --region=us-west1 --tag us-west1-docker.pkg.dev/keras-team-test/keras-nlp-test/keras-nlp-image-{framework}:deps --timeout=30m
+```
+
+Alternately, to update all docker images at once, just run:
+```
+./cloudbuild/update_images.sh
+```
 
 ### Run TPU Testing
 

.cloudbuild/jax/Dockerfile

@@ -0,0 +1,8 @@
+FROM nvidia/cuda:11.7.1-base-ubuntu20.04
+RUN apt-get update
+RUN apt-get install -y python3 python3-pip
+RUN apt-get install -y git
+RUN git clone https://github.com/keras-team/keras-nlp.git
+RUN cd keras-nlp
+RUN pip install -r keras-nlp/requirements.txt
+RUN pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

.cloudbuild/tensorflow/Dockerfile

@@ -0,0 +1,6 @@
+FROM tensorflow/tensorflow:2.12.0-gpu
+RUN apt-get -y update
+RUN apt-get -y install git
+RUN git clone https://github.com/keras-team/keras-nlp.git
+RUN cd keras-nlp
+RUN pip install -r keras-nlp/requirements.txt

.cloudbuild/torch/Dockerfile

@@ -0,0 +1,8 @@
+FROM nvidia/cuda:11.7.1-base-ubuntu20.04
+RUN apt-get update
+RUN apt-get install -y python3 python3-pip
+RUN apt-get install -y git
+RUN git clone https://github.com/keras-team/keras-nlp.git
+RUN cd keras-nlp
+RUN pip install -r keras-nlp/requirements.txt
+RUN pip install torch

.cloudbuild/update_images.sh

@@ -0,0 +1,12 @@
+#!/bin/bash -ex
+
+base_dir=$(dirname $0)
+
+for plaform in "jax" "tensorflow" "torch"; do
+    pushd "${base_dir}/${plaform}" > /dev/null
+    gcloud builds submit \
+        --region=us-west1\
+        --tag "us-west1-docker.pkg.dev/keras-team-test/keras-nlp-test/keras-nlp-image-${plarform}:deps" \
+        --timeout=30m
+    popd
+done