Skip to content

Commit a688c60

Browse files
RoshaniNRoshaniN
committed
Few more changes.
1 parent 574e6fe commit a688c60

File tree

3 files changed

+58
-49
lines changed

3 files changed

+58
-49
lines changed

‎.github/workflows/RunTests.yml

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -53,18 +53,6 @@ jobs:
5353
# build_mode: stable_stack
5454
# base_image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
5555

56-
tpu_pathways_installs:
57-
runs-on: ["self-hosted", "tpu_pathways", "v4-8"]
58-
needs: prelim
59-
steps:
60-
- name: Install Docker Compose and configure docker
61-
run: |
62-
sudo apt-get -y install docker
63-
sudo apt-get -y install docker-compose-plugin
64-
docker compose version
65-
gcloud auth configure-docker us-docker.pkg.dev --quiet
66-
67-
6856
# gpu_image:
6957
# needs: prelim
7058
# uses: ./.github/workflows/build_upload_internal.yml
@@ -86,7 +74,7 @@ jobs:
8674
# container_resource_option: "--privileged"
8775

8876
tpu_pathways_unit_tests:
89-
needs: [tpu_pathways_installs]
77+
needs: [prelim]
9078
uses: ./.github/workflows/run_tests_internal.yml
9179
with:
9280
device_type: tpu_pathways

‎.github/workflows/run_tests_internal.yml

Lines changed: 46 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -42,31 +42,53 @@ on:
4242
type: string
4343

4444
jobs:
45-
pathways:
46-
runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
47-
steps:
48-
- uses: actions/checkout@v4
49-
- name: Start containers using script
50-
run: |
51-
bash docker_build_pathways_containers.sh
52-
# run:
45+
# pathways:
5346
# runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
54-
# container:
55-
# image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
56-
# volumes:
57-
# - /home/runner/actions-runner/_work/maxtext/maxtext:/deps
58-
# env:
59-
# XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
60-
# TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
61-
# # JAX_PLATFORMS: "proxy"
62-
# # JAX_BACKEND_TARGET: "grpc://localhost:29000"
63-
# options: ${{ inputs.container_resource_option }}
6447
# steps:
6548
# - uses: actions/checkout@v4
66-
# - name: Run Tests
49+
# - name: Start containers using script
6750
# run: |
68-
# export JAX_PLATFORMS="proxy"
69-
# export JAX_BACKEND_TARGET="grpc://localhost:29000"
70-
# echo $JAX_BACKEND_TARGET
71-
# echo "Set Pathways env variables"
72-
# python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"
51+
# bash docker_build_pathways_containers.sh
52+
run:
53+
runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
54+
container:
55+
# image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu
56+
image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:latest
57+
volumes:
58+
- /home/runner/actions-runner/_work/maxtext/maxtext:/deps
59+
env:
60+
XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
61+
TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
62+
# JAX_PLATFORMS: "proxy"
63+
# JAX_BACKEND_TARGET: "grpc://localhost:29000"
64+
options: ${{ inputs.container_resource_option }}
65+
steps:
66+
- uses: actions/checkout@v4
67+
# - name: Install Docker Compose and configure docker
68+
# run: |
69+
# # apt-get -y install docker
70+
# # apt-get -y install docker-compose-plugin
71+
# docker compose version
72+
# gcloud auth configure-docker us-docker.pkg.dev --quiet
73+
# - name: Start containers using script
74+
# run: |
75+
# bash docker_build_pathways_containers.sh
76+
- name: Set up Docker Compose
77+
uses: docker/setup-buildx-action@v2
78+
with:
79+
driver-opts: image=moby/buildkit:v0.11.2-rootless
80+
81+
- name: Build and run Docker Compose
82+
uses: docker/compose-action@v3
83+
with:
84+
compose-file: /utils_pathways/docker-compose.yml # Path to your docker-compose.yml
85+
up: true # Start the containers
86+
down: true # Ensure containers are stopped after the job (important!)
87+
- name: Run Tests with Pathways backend
88+
run: |
89+
export JAX_PLATFORMS="proxy"
90+
export JAX_BACKEND_TARGET="grpc://localhost:29000"
91+
echo "Set Pathways env variables"
92+
echo $JAX_BACKEND_TARGET
93+
cd MaxText/
94+
python3 -m pytest -s ${{ inputs.test_directory }} -m "${{ inputs.pytest_marker }}"

‎utils_pathways/docker-compose.yml

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11

22
#### HOST NETWORK
3-
4-
# networks:
5-
# local: {}
6-
73
services:
84

95
resource_manager:
@@ -61,14 +57,17 @@ services:
6157
- --gcs_scratch_location=gs://cloud-pathways-staging/tmp
6258

6359

64-
maxtext:
65-
image: us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest
66-
network_mode: host
67-
environment:
68-
- JAX_PLATFORMS="proxy"
69-
- JAX_BACKEND_TARGET="grpc://localhost:29000"
70-
command:
71-
- cd Maxtext ; python3 -m pytest -s tests -m tpu_only
60+
# maxtext:
61+
# image: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
62+
# # us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
63+
# # gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:latest
64+
# network_mode: host
65+
# environment:
66+
# - JAX_PLATFORMS="proxy"
67+
# - JAX_BACKEND_TARGET="grpc://localhost:29000"
68+
# entrypoint: "/bin/bash"
69+
# command:
70+
# - "ls "
7271

7372
#### LOCAL NETWORK
7473

0 commit comments

Comments
 (0)