diff --git a/.github/workflows/run_pathways_tests.yml b/.github/workflows/run_pathways_tests.yml index 0777d33c91..13d28f2f05 100644 --- a/.github/workflows/run_pathways_tests.yml +++ b/.github/workflows/run_pathways_tests.yml @@ -104,6 +104,19 @@ jobs: python3 -m pip freeze - name: Copy test assets files run : gcloud storage cp gs://maxtext-test-assets/* tests/assets + - name: Wait for Pathways Proxy Readiness + shell: bash + run: | + echo "Waiting for Pathways Proxy at localhost:29000..." + for i in {1..30}; do + if python3 -c "import socket; socket.create_connection(('localhost', 29000), timeout=1)" 2>/dev/null; then + echo "Pathways Proxy port is open and listening!" + sleep 2 # Allow gRPC server to complete internal handshakes + exit 0 + fi + sleep 2 + done + echo "ERROR: Pathways Proxy failed to start within 60s" && exit 1 - name: Run Tests env: TOTAL_WORKERS: ${{ inputs.total_workers }} @@ -141,7 +154,7 @@ jobs: - "29001:29001" - "29002:29002" options: - --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29001, --node_type=resource_manager, --enforce_kernel_ipv6_support=false, --instance_count=1, --instance_type=tpuv6e:2x2, --gcs_scratch_location=gs://cloud-pathways-staging/tmp] + --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29001, --node_type=resource_manager, --enforce_kernel_ipv6_support=false, --instance_count=1, --instance_type=tpuv6e:2x2, --gcs_scratch_location=gs://cloud-pathways-staging/maxtext-ci/${{ github.run_id }}_${{ github.run_attempt }}_${{ inputs.worker_group }}] env: TPU_SKIP_MDS_QUERY: true @@ -153,7 +166,7 @@ jobs: - "8471:8471" - "8080:8080" options: - --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29005, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/tmp] + --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29005, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/maxtext-ci/${{ github.run_id }}_${{ github.run_attempt }}_${{ inputs.worker_group }}] --tpu=4 proxy: @@ -164,4 +177,4 @@ jobs: IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: true XLA_FLAGS: "--xla_dump_to=/tmp/aot_test_dump --xla_dump_hlo_as_text --xla_dump_hlo_module_re=jit_train_step" options: - --entrypoint=[/usr/pathways/run/cloud_proxy_server_sanitized, --server_port=29000, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/tmp, --xla_tpu_scoped_vmem_limit_kib=65536, --xla_tpu_spmd_rng_bit_generator_unsafe=true] + --entrypoint=[/usr/pathways/run/cloud_proxy_server_sanitized, --server_port=29000, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/maxtext-ci/${{ github.run_id }}_${{ github.run_attempt }}_${{ inputs.worker_group }}, --xla_tpu_scoped_vmem_limit_kib=65536, --xla_tpu_spmd_rng_bit_generator_unsafe=true]