update workflows to clone main xla via yaml after test env removal

Author: nitins17

.github/workflows/bazel_cpu_rbe.yml

@@ -32,9 +32,7 @@ jobs:
                    (contains(matrix.runner, 'windows-x86') && null) }}
 
     env:
-      # Disable running `run_docker_container.sh`. JAX self-hosted runners already run in a Docker
-      # container.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
 
     steps:

.github/workflows/bazel_gpu_non_rbe.yml

@@ -26,8 +26,6 @@ jobs:
       image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
 
     env:
-      # GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
       JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: 3.11
 

.github/workflows/bazel_gpu_rbe.yml

@@ -26,8 +26,6 @@ jobs:
       image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
 
     env:
-      # Do not run Docker container for Linux runners. Linux runners already run in a Docker container.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
       JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: 3.12
 

.github/workflows/pytest_cpu.yml

@@ -33,9 +33,7 @@ jobs:
                 (contains(matrix.runner, 'windows-x86') && null) }}
 
     env:
-      # Disable running `run_docker_container.sh`. JAX self-hosted runners already run in a Docker
-      # container.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:

.github/workflows/pytest_gpu.yml

@@ -27,8 +27,7 @@ jobs:
       image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
 
     env:
-      # GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:

.github/workflows/pytest_tpu.yml

@@ -27,8 +27,7 @@ jobs:
       image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
 
     env:
-      # GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
-      JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps: