New changes to help run tests.

Author: RoshaniN

.github/workflows/RunTests.yml

@@ -35,14 +35,14 @@ jobs:
     - name: Cleanup old docker images
       run: docker system prune --all --force
 
-  tpu_image:
-    needs: prelim
-    uses: ./.github/workflows/build_upload_internal.yml
-    with:
-      device_type: tpu_pathways
-      device_name: v4-8
-      build_mode: stable_stack
-      base_image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
+  # tpu_image:
+  #   needs: prelim
+  #   uses: ./.github/workflows/build_upload_internal.yml
+  #   with:
+  #     device_type: tpu_pathways
+  #     device_name: v4-8
+  #     build_mode: stable_stack
+  #     base_image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
 
   tpu_pathways_image:
     needs: prelim

.github/workflows/run_tests_internal.yml

@@ -62,5 +62,6 @@ jobs:
           export JAX_PLATFORMS="proxy" ; export JAX_BACKEND_TARGET="grpc://localhost:29000" ; echo "Set Pathways env variables"
       - name: Run Tests
         run: |
+          export JAX_PLATFORMS="proxy" ; export JAX_BACKEND_TARGET="grpc://localhost:29000" ; echo "Set Pathways env variables" ;
           cd MaxText
           python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"

docker_build_pathways_containers.sh

@@ -1,3 +1,19 @@
+#!/bin/bash
+
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 apt-get install docker-compose-plugin
 docker compose version
 gcloud auth configure-docker us-docker.pkg.dev --quiet
@@ -6,7 +22,7 @@ cd utils_pathways
 docker compose down
 sleep 2
 docker compose up &
-sleep 5
+sleep 15
 docker compose ps
 
 export JAX_PLATFORMS="proxy"