update docker scripts

Author: nitins17

ci/envs/docker.env

@@ -15,24 +15,41 @@
 # This file contains all the docker specifc envs that are needed by the
 # ci/utilities/run_docker_container.sh script.
 
-# Inherit default JAXCI environment variables.
-source ci/envs/default.env
-
 os=$(uname -s | awk '{print tolower($0)}')
 arch=$(uname -m)
 
-# TODO: Set GPU Docker args and GPU Docker images
-# Linux x86 specifc settings
+# The path to the JAX git repository.
+export JAXCI_JAX_GIT_DIR=$(pwd)
+
+export JAXCI_DOCKER_WORK_DIR="/jax"
+export JAXCI_DOCKER_ARGS=""
+
+# Linux x86 image for building JAX artifacts, running Pytests CPU/TPU tests, and Bazel tests
 if [[ $os == "linux" ]] && [[ $arch == "x86_64" ]]; then
   export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
 fi
 
-# Linux Aarch64 specifc settings
+# Linux Aarch64 image for building JAX artifacts, running Pytests CPU tests, and Bazel tests
 if [[ $os == "linux" ]] && [[ $arch == "aarch64" ]]; then
-  export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest"
+  export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest"
 fi
 
-# Windows specific settings
+# Windows image for building JAX artifacts, running Pytests CPU tests, and Bazel tests
 if [[ $os =~ "msys_nt" ]]; then
-  export JAXCI_DOCKER_IMAGE="gcr.io/tensorflow-testing/tf-win2019-rbe@sha256:1082ef4299a72e44a84388f192ecefc81ec9091c146f507bc36070c089c0edcc"
-fi
+  export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/tf-test-windows:latest"
+fi
+
+# Uncomment the following lines if you want to run the GPU tests with Pytest.
+# Note that GPU Pytests, as a prequisite, require that the following JAX artifacts be
+# present in the $JAXCI_OUTPUT_DIR: jaxlib, jax-cuda-plugin, jax-cuda-pjrt. If you don't
+# have these wheels stored there, either build them from source via ci/build_artifacts.sh or
+# download them from PyPI into that folder.
+#
+# Linux x86 image for running Pytest GPU tests
+# if [[ $os == "linux" ]] && [[ $arch == "x86_64" ]]; then
+#  # Choose one of: 12.3, 12.1
+#  export JAXCI_DOCKER_CUDA_VERSION=${JAXCI_DOCKER_CUDA_VERSION:-12.3}
+#  export JAXCI_DOCKER_IMAGE="gcr.io/tensorflow-testing/nosla-cuda${JAXCI_DOCKER_CUDA_VERSION}-cudnn9.1-ubuntu20.04-manylinux2014-multipython"
+#
+#  export JAXCI_DOCKER_ARGS="--gpus all --shm-size=16g"
+# fi

ci/utilities/run_docker_container.sh

@@ -13,7 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Set up the Docker container and start it for JAX CI jobs.
+# Sets up a Docker container for JAX CI.
+
+# This script creates and starts a Docker container named "jax" for internal
+# JAX CI jobs. These jobs primarily handle building and publishing JAX artifacts
+# to PyPI and/or GCS.
+
+# Note: GitHub Actions workflows do not utilize this script, as they leverage
+# built-in containerization features to run jobs within a container. However,
+# they use the same Docker image to maintain consistency. This script also helps
+# ensure that local build environments mirror the behavior of CI builds.
+# Usage:
+#     ./ci/utilities/run_docker_container.sh
+#     docker exec -it jax <build-script>
+#     E.g: docker exec -it jax ./ci/build_artifacts.sh jaxlib
 #
 # -e: abort script if one command fails
 # -u: error if undefined variable used
@@ -46,13 +59,13 @@ if ! docker container inspect jax >/dev/null 2>&1 ; then
     JAXCI_DOCKER_ARGS="$JAXCI_DOCKER_ARGS -v $HOME/.config/gcloud:/root/.config/gcloud"
   fi
 
-  # Start the container. `user_set_jaxci_envs` is read after `jax_ci_envs` to
-  # allow the user to override any environment variables set by JAXCI_ENV_FILE.
+  # Start the container.
   docker run $JAXCI_DOCKER_ARGS --name jax \
-      -w "$JAXCI_DOCKER_WORK_DIR" -itd --rm \
-      -v "$JAXCI_JAX_GIT_DIR:$JAXCI_DOCKER_WORK_DIR" \
-      "$JAXCI_DOCKER_IMAGE" \
-    bash
+          --env-file <(env | grep JAXCI_) \
+          -w "$JAXCI_DOCKER_WORK_DIR" -itd --rm \
+          -v "$JAXCI_JAX_GIT_DIR:$JAXCI_DOCKER_WORK_DIR" \
+          "$JAXCI_DOCKER_IMAGE" \
+          bash
 
   if [[ "$(uname -s)" =~ "MSYS_NT" ]]; then
     # Allow requests from the container.