Few more changes.

Author: RoshaniN

.github/workflows/RunTests.yml

@@ -53,18 +53,6 @@ jobs:
   #     build_mode: stable_stack
   #     base_image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
 
-  tpu_pathways_installs:
-    runs-on: ["self-hosted", "tpu_pathways", "v4-8"]
-    needs: prelim
-    steps:
-    - name: Install Docker Compose and configure docker
-      run: |
-        sudo apt-get -y install docker
-        sudo apt-get -y install docker-compose-plugin
-        docker compose version
-        gcloud auth configure-docker us-docker.pkg.dev --quiet
-        
-
   # gpu_image:
   #   needs: prelim
   #   uses: ./.github/workflows/build_upload_internal.yml
@@ -86,7 +74,7 @@ jobs:
   #     container_resource_option: "--privileged"
 
   tpu_pathways_unit_tests:
-    needs: [tpu_pathways_installs]
+    needs: [prelim]
     uses: ./.github/workflows/run_tests_internal.yml
     with:
       device_type: tpu_pathways

.github/workflows/run_tests_internal.yml

@@ -42,31 +42,42 @@ on:
         type: string
 
 jobs:
-  pathways:
+  # pathways:
+  #   runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
+  #   steps:
+  #     - uses: actions/checkout@v4
+  #     - name: Start containers using script
+  #       run: |
+  #         bash docker_build_pathways_containers.sh
+  run:
     runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
+    container:
+      # image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
+      image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:latest
+      volumes:
+        - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
+      env:
+        XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
+        TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
+        # JAX_PLATFORMS: "proxy"
+        # JAX_BACKEND_TARGET: "grpc://localhost:29000"
+      options: ${{ inputs.container_resource_option }}
     steps:
       - uses: actions/checkout@v4
+      - name: Install Docker Compose and configure docker
+        run: |
+          apt-get -y install docker
+          apt-get -y install docker-compose-plugin
+          docker compose version
+          gcloud auth configure-docker us-docker.pkg.dev --quiet
       - name: Start containers using script
         run: |
           bash docker_build_pathways_containers.sh
-  # run:
-  #   runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
-  #   container:
-  #     image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
-  #     volumes:
-  #       - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
-  #     env:
-  #       XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
-  #       TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
-  #       # JAX_PLATFORMS: "proxy"
-  #       # JAX_BACKEND_TARGET: "grpc://localhost:29000"
-  #     options: ${{ inputs.container_resource_option }}
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #     - name: Run Tests
-  #       run: |
-  #         export JAX_PLATFORMS="proxy" 
-  #         export JAX_BACKEND_TARGET="grpc://localhost:29000" 
-  #         echo $JAX_BACKEND_TARGET 
-  #         echo "Set Pathways env variables"
-  #         python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"
+      - name: Run Tests with Pathways backend
+        run: |
+          export JAX_PLATFORMS="proxy" 
+          export JAX_BACKEND_TARGET="grpc://localhost:29000" 
+          echo "Set Pathways env variables"
+          echo $JAX_BACKEND_TARGET
+          cd MaxText/
+          python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"

utils_pathways/docker-compose.yml

@@ -1,9 +1,5 @@
 
 #### HOST NETWORK
-
-# networks:
-#   local: {}
-
 services:
 
   resource_manager:
@@ -61,14 +57,17 @@ services:
       - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
 
 
-  maxtext:
-    image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
-    network_mode: host
-    environment:
-      - JAX_PLATFORMS="proxy"
-      - JAX_BACKEND_TARGET="grpc://localhost:29000"
-    command:
-      - cd Maxtext ; python3 -m pytest -s tests -m tpu_only
+  # maxtext:
+  #   image: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
+  #   # us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
+  #   # gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:latest
+  #   network_mode: host
+  #   environment:
+  #     - JAX_PLATFORMS="proxy"
+  #     - JAX_BACKEND_TARGET="grpc://localhost:29000"
+  #   entrypoint: "/bin/bash"
+  #   command:
+  #     - "ls "
 
 #### LOCAL NETWORK