One more change.

AI-Hypercomputer · Feb 8, 2025 · e8a31d7 · e8a31d7
1 parent 574e6fe
commit e8a31d7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 33 deletions.
diff --git a/.github/workflows/run_tests_internal.yml b/.github/workflows/run_tests_internal.yml
@@ -49,24 +49,25 @@ jobs:
       - 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 }}"
+  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 "Set Pathways env variables"
+          echo $JAX_BACKEND_TARGET
+          cd MaxText/
+          python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"
diff --git a/utils_pathways/docker-compose.yml b/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