Move installation to Github runner step, simplify command.

Author: RoshaniN

.github/workflows/RunTests.yml

@@ -82,8 +82,8 @@ jobs:
     with:
       device_type: tpu_pathways
       device_name: v4-8
-      # pytest_marker: 'not gpu_only and not integration_test'
-      # test_directory: 'tests'
+      pytest_marker: 'not gpu_only and not integration_test'
+      test_directory: 'tests'
       # xla_python_client_mem_fraction: 0.75
       # tf_force_gpu_allow_growth: false
       # container_resource_option: "--privileged"

.github/workflows/run_pathways_tests_internal.yml

@@ -25,13 +25,26 @@ on:
       device_name:
         required: true
         type: string
+      pytest_marker:
+        required: true
+        type: string
+      test_directory:
+        required: true
+        type: string
 
 jobs:
   start_pathways:
     name: Start Pathways tests on (${{ inputs.device_name }})
     runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
     steps:
+      - name: Install prerequisites on the self-hosted runner for Pathways.
+        run: |
+          echo "Setting up the prerequisites"
+          apt-get install docker
+          apt-get install docker-compose-plugin
+          docker compose version # To ensure docker compose is installed
+          gcloud auth configure-docker us-docker.pkg.dev --quiet
       - uses: actions/checkout@v4
       - name: Start Pathways and Maxtext containers using script
         run: |
-          bash docker_run_pathways_containers.sh maxtext_image=gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
+          bash docker_run_pathways_containers.sh maxtext_image=gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu command="cd MaxText ; python3 -m pytest ${{ inputs.test_directory }} -m '${{ inputs.pytest_marker }}' -s" 

docker_run_pathways_containers.sh

@@ -14,26 +14,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Example - # bash docker_run_pathways_containers.sh maxtext_image=<>
-
-# bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
+# Example - bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s" 
 
 # Stop execution if any command exits with error
 
 
 echo "Running docker_run_pathways_containers.sh"
 
-# Prerequisites on the self-hosted runner
-echo "Setting up the prerequisites"
-apt-get install docker
-apt-get install docker-compose-plugin
-docker compose version # To ensure docker compose is installed
-gcloud auth configure-docker us-docker.pkg.dev --quiet
-
 set -e
-# Default - 
-maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
 
+# Defaults - 
+maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
+command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
 
 # Parse input variables
 for ARGUMENT in "$@"; do
@@ -43,15 +35,8 @@ for ARGUMENT in "$@"; do
 done
 
 cd utils_pathways
-MAXTEXT_IMAGE=${maxtext_image} docker compose down
+MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose down
 sleep 10
-echo "MAXTEXT image ----------------"
-echo $MAXTEXT_IMAGE
-MAXTEXT_IMAGE=${maxtext_image} docker compose up --exit-code-from maxtext
+MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose up --exit-code-from maxtext
 sleep 10
-MAXTEXT_IMAGE=${maxtext_image} docker compose down
-# echo "Exit code from Maxtext..."
-# echo $?
-# sleep 20
-# docker compose ps
-# MAXTEXT_IMAGE=${maxtext_image} docker compose logs -f maxtext
+MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose down

utils_pathways/docker-compose.yml

@@ -1,5 +1,17 @@
+# Copyright 2024 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.
 
-#### LOCAL NETWORK
 services:
 
   resource_manager:
@@ -53,7 +65,6 @@ services:
 
   maxtext:
     image: ${MAXTEXT_IMAGE}
-    # us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
     privileged: true
     volumes: 
       - ../:/deps
@@ -62,8 +73,6 @@ services:
       - JAX_BACKEND_TARGET=grpc://proxy:29008
     entrypoint: ["bash", "-c"]
     command:
-      - cd MaxText ; python3 -m pytest tests -m "not gpu_only and not integration_test" -s
-      # - "cd MaxText ; python3 -m pytest -m tpu_only tests -s"
-      # - "cd MaxText ; python3 -m pytest -s tests/train_tests.py -m tpu_only"
+      - ${COMMAND}
     depends_on:
       - proxy