Simplified flow test.

Author: RoshaniN

.github/workflows/RunTests.yml

@@ -134,17 +134,17 @@ jobs:
   #     container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
 
 
-  clean_up:
-    if: ${{ always() }}  # always execute, regardless of previous jobs or steps.
-    # needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_pathways_unit_tests, tpu_integration_tests]
-    needs: [tpu_pathways_unit_tests]
-    # needs: [tpu_unit_tests]
-    name: "Clean up"
-    runs-on: ["self-hosted", "tpu_pathways" ]
-    steps:
-    # - name: Delete GPU image
-    #   run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu --force-delete-tags --quiet
-    - name: Delete TPU image
-      run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu --force-delete-tags --quiet
+  # clean_up:
+  #   if: ${{ always() }}  # always execute, regardless of previous jobs or steps.
+  #   # needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_pathways_unit_tests, tpu_integration_tests]
+  #   needs: [tpu_pathways_unit_tests]
+  #   # needs: [tpu_unit_tests]
+  #   name: "Clean up"
+  #   runs-on: ["self-hosted", "tpu_pathways" ]
+  #   steps:
+  #   # - name: Delete GPU image
+  #   #   run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu --force-delete-tags --quiet
+  #   # - name: Delete TPU image
+  #   #   run: gcloud container images delete gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu --force-delete-tags --quiet
 
 

.github/workflows/run_tests_internal.yml

@@ -46,10 +46,9 @@ jobs:
     runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
     steps:
       - uses: actions/checkout@v4
-      - uses: ./.github/workflows/start_pathways_containers.yml
-        with:
-          device_type: tpu_pathways
-          device_name: v4-8
+      - name: Start containers using script
+        run: |
+          bash docker_build_pathways_containers.sh
   # run:
   #   runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
   #   container:

utils_pathways/docker-compose.yml

@@ -68,7 +68,7 @@ services:
       - JAX_PLATFORMS="proxy"
       - JAX_BACKEND_TARGET="grpc://localhost:29000"
     command:
-      - python3 -m pytest -s tests -m tpu_only
+      - cd Maxtext ; python3 -m pytest -s tests -m tpu_only
 
 #### LOCAL NETWORK