From 031d4efc2fb6e11a312438e98f71fe026ea260bc Mon Sep 17 00:00:00 2001 From: Kunjan Date: Tue, 12 Aug 2025 10:32:25 -0700 Subject: [PATCH 1/6] Update UnitTests.yml Modify to use new runner --- .github/workflows/UnitTests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index f7ee2166c..3435a44a1 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -31,7 +31,7 @@ jobs: strategy: fail-fast: false matrix: - tpu-type: ["v4-8"] + tpu-type: ["v5-8"] name: "TPU test (${{ matrix.tpu-type }})" runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}"] steps: From e284d88c3abf418a1ba0a5333ae7c307fb0995a6 Mon Sep 17 00:00:00 2001 From: Kunjan Date: Tue, 12 Aug 2025 10:43:31 -0700 Subject: [PATCH 2/6] Update UnitTests.yml Explicit, list runner --- .github/workflows/UnitTests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index 3435a44a1..6ac797224 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -33,7 +33,7 @@ jobs: matrix: tpu-type: ["v5-8"] name: "TPU test (${{ matrix.tpu-type }})" - runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}"] + runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}","maxdiffusion-github-runner-2"] steps: - uses: actions/checkout@v4 - name: Set up Python 3.12 From ce208229253509349c56044c20140cfd654e72a7 Mon Sep 17 00:00:00 2001 From: Kunjan Date: Tue, 12 Aug 2025 14:17:18 -0700 Subject: [PATCH 3/6] Update UnitTests.yml --- .github/workflows/UnitTests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index 6ac797224..be7890a94 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -31,9 +31,9 @@ jobs: strategy: fail-fast: false matrix: - tpu-type: ["v5-8"] + tpu-type: ["v5p-8"] name: "TPU test (${{ matrix.tpu-type }})" - runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}","maxdiffusion-github-runner-2"] + runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}", "maxdiffusion-github-runner-2"] steps: - uses: actions/checkout@v4 - name: Set up Python 3.12 From e727c11db98d9ea23e5f6ba0130e2fd41f4862f7 Mon Sep 17 00:00:00 2001 From: Kunjan Date: Tue, 12 Aug 2025 14:41:15 -0700 Subject: [PATCH 4/6] Update UnitTests.yml --- .github/workflows/UnitTests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index be7890a94..5ba508143 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -33,7 +33,7 @@ jobs: matrix: tpu-type: ["v5p-8"] name: "TPU test (${{ matrix.tpu-type }})" - runs-on: ["self-hosted", "tpu", "${{ matrix.tpu-type }}", "maxdiffusion-github-runner-2"] + runs-on: ["self-hosted","${{ matrix.tpu-type }}", "maxdiffusion-github-runner-2"] steps: - uses: actions/checkout@v4 - name: Set up Python 3.12 From ef515b1f2e89d219ac293604a2f96f13f9715fc3 Mon Sep 17 00:00:00 2001 From: Kunjan Date: Tue, 12 Aug 2025 14:43:56 -0700 Subject: [PATCH 5/6] Update UnitTests.yml --- .github/workflows/UnitTests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index 5ba508143..d2e9a5a2c 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -33,7 +33,7 @@ jobs: matrix: tpu-type: ["v5p-8"] name: "TPU test (${{ matrix.tpu-type }})" - runs-on: ["self-hosted","${{ matrix.tpu-type }}", "maxdiffusion-github-runner-2"] + runs-on: ["self-hosted","${{ matrix.tpu-type }}"] steps: - uses: actions/checkout@v4 - name: Set up Python 3.12 From d967bd4eec205e383529ffb98965dfa68b185e17 Mon Sep 17 00:00:00 2001 From: Kunjan patel Date: Tue, 12 Aug 2025 22:51:35 +0000 Subject: [PATCH 6/6] Fix unit test for v5p runner --- .github/workflows/XLML.yml | 2 +- tests/schedulers/test_scheduler_flax.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/XLML.yml b/.github/workflows/XLML.yml index c9a3bf69b..37f320787 100644 --- a/.github/workflows/XLML.yml +++ b/.github/workflows/XLML.yml @@ -2,7 +2,7 @@ name: Add Testgrid Link to PR on: pull_request: - types: [opened, synchronize] + types: [opened] jobs: add_testgrid_link: diff --git a/tests/schedulers/test_scheduler_flax.py b/tests/schedulers/test_scheduler_flax.py index a63e3a966..974ac3ab3 100644 --- a/tests/schedulers/test_scheduler_flax.py +++ b/tests/schedulers/test_scheduler_flax.py @@ -335,8 +335,8 @@ def test_full_loop_no_noise(self): result_mean = jnp.mean(jnp.abs(sample)) if jax_device == "tpu": - assert abs(result_sum - 257.2727) < 1e-2 - assert abs(result_mean - 0.3349905) < 1e-3 + assert abs(result_sum - 257.2727) < 1.5e-2 + assert abs(result_mean - 0.3349905) < 1e-5 else: assert abs(result_sum - 255.1113) < 1e-2 assert abs(result_mean - 0.332176) < 1e-3