Run selected tests in Decoupled mode on CPU in CI by shralex · Pull Request #4642 · AI-Hypercomputer/maxtext

shralex · 2026-07-28T12:14:29Z

This PR adds a targeted decoupled test suite executed with DECOUPLE_GCLOUD=TRUE on CPU in CI to validate fragile offline components (checkpointing, quantization, GRPO NNX loss dispatch).

Changes:

pytest.ini: Registered decoupled_target marker ("tests to be additionally validated in decoupled mode (DECOUPLE_GCLOUD=TRUE) on CPU in CI.").
Test Modules: Tagged 9 target test files across tests/unit/ and tests/integration/ with pytestmark = [pytest.mark.decoupled_target].
CI Pipeline (.github/workflows/run_tests_against_package.yml): Added a step under inputs.flavor == 'cpu-unit' executing $PYTHON_EXE -m pytest -v -n auto -m "decoupled_target".

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

- Register decoupled_target marker in pytest.ini for tests to be additionally validated in decoupled mode on CPU in CI. - Mark targeted test modules in unit/ and integration/ with pytestmark = [pytest.mark.decoupled_target]. - Add step in .github/workflows/run_tests_against_package.yml to execute pytest -m decoupled_target -n auto on CPU runners.

gemini-code-assist · 2026-07-28T12:14:32Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

SurbhiJainUSC · 2026-07-28T16:39:23Z

    tpu_backend: marks tests that require a TPU-enabled JAX install (TPU PJRT plugin), but not TPU hardware
    gpu_only: marks tests to be run on GPUs only
    cpu_only: marks tests to be run on CPUs only
    decoupled: tests that validate offline / DECOUPLE_GCLOUD=TRUE mode.


How about we reuse this marker instead of creating a new one decoupled_target?

The goal is different -- decoupled says which tests are able to execute without GCP dependencies. I believe its also applied automatically for tests not explicitly excluded. Decoupled_target says which tests should we run in decoupled mode as part of CI testing (currently none of them run). If you see a way to clarify this please let me know

I updated the comment for decoupled_target

- Register decoupled_target marker in pytest.ini for tests to be additionally validated in decoupled mode on CPU in CI. - Mark targeted test modules in unit/ and integration/ with pytestmark = [pytest.mark.decoupled_target]. - Add step in .github/workflows/run_tests_against_package.yml to execute pytest -m decoupled_target -n auto on CPU runners.

gulsumgudukbay · 2026-07-28T17:58:11Z

+            --junitxml=test-results-decoupled-targeted.xml
+
+
        env:


This step landed between the previous step's run: and its env:, so that env: now belongs to this step. Run Tests loses all ten INPUTS_* vars for every flavor, it would run -m "", unsharded, with no --ignore=tests/post_training. I think this step should be moved below the env: block.

gulsumgudukbay · 2026-07-28T17:58:52Z

          fi

+      - name: Run Targeted Decoupled Tests
+        if: inputs.flavor == 'cpu-unit'


cpu-unit is a 4-way matrix (CPU_UNIT_TOTAL_WORKERS=4), so this runs in all four shards with no --splits/--group, the same 48 tests four times per PR.

gulsumgudukbay · 2026-07-28T18:00:11Z

+        if: inputs.flavor == 'cpu-unit'
+        run: |
+          echo "Running targeted tests in decoupled mode (DECOUPLE_GCLOUD=TRUE)..."
+          export DECOUPLE_GCLOUD=TRUE


This doesn't reproduce decoupled mode here: cpu-unit installs the tpu extra, which pins google-cloud-storage>=3.13.0, and DECOUPLE_GCLOUD only changes what gcloud_stub returns, no meta_path hook exists.

a module-scope from google.cloud import storage still resolves, so #4003, #4180 and #4637 would all have passed this gate, as would today's main.

I strongly believe here decoupled requirements should be installed.

gulsumgudukbay · 2026-07-28T18:01:34Z

+          echo "Running targeted tests in decoupled mode (DECOUPLE_GCLOUD=TRUE)..."
+          export DECOUPLE_GCLOUD=TRUE
+
+          $PYTHON_EXE -m pytest -v -n auto -m "decoupled_target" \


Is $PYTHON_EXE empty here? since it's a shell variable from the previous step's run: body and it doesn't cross steps. So this becomes -m pytest ...
and can cause the bash: -m: command not found error. we may also see this behavior for source .venv/bin/activate, PYTHONPATH, and the cd /deps branch.

gulsumgudukbay · 2026-07-28T18:02:53Z

+          export DECOUPLE_GCLOUD=TRUE
+
+          $PYTHON_EXE -m pytest -v -n auto -m "decoupled_target" \
+            --junitxml=test-results-decoupled-targeted.xml


I see arg issues here: no paths, so addopts --ignores apply and 4 of the 9 tagged files never collect (48 tests from 5 files); no --ignore=tests/post_training, and -m filters after collection, so there are 14 errors during collection; -n auto only works since when there are more than one worker xdist is installed

gulsumgudukbay · 2026-07-28T18:04:06Z

    decoupled: tests that validate offline / DECOUPLE_GCLOUD=TRUE mode.
 	       NOTE: this marker is not to be used manually, it is auto-
-	       applied to tests with external_* or tpu_only marker. 
+	       applied to tests without external_* or tpu_only marker. 


tpu_only tests do get the decoupled marker

Suggested change

applied to tests without external_* or tpu_only marker.

all tests except external_serving / external_training.

gulsumgudukbay · 2026-07-28T18:06:16Z

+                      with DECOUPLE_GCLOUD=TRUE, conftest.py automatically
+                      selects at runtime only those tests that can run in


conftest.py only deselects external_serving / external_training; it doesn't judge whether a test can run decoupled. this PR selects -m "decoupled_target", not decoupled and decoupled_target, so that marker isn't involved at all.

gulsumgudukbay · 2026-07-28T18:07:07Z

+import pytest
 import safetensors.numpy

+pytestmark = [pytest.mark.decoupled_target]


#4637 is a prerequisite for this to run without failures.

gulsumgudukbay · 2026-07-28T18:07:48Z

 from maxtext.checkpoint_conversion import to_maxtext

-pytestmark = [pytest.mark.integration_test]
+pytestmark = [pytest.mark.integration_test, pytest.mark.decoupled_target]


cpu-unit's marker deliberately excludes integration tests so we don't need this.

shralex added 2 commits July 28, 2026 12:09

shralex requested review from bvandermoon, darisoy, gobbleturk, huytransformer, igorts-git, khatwanimohit, parambole, richjames0 and xibinliu as code owners July 28, 2026 12:14

shralex requested review from A9isha, NuojCheng, RissyRan, SurbhiJainUSC, abhinavclemson, aireenmei, dipannita08, gagika, hengtaoguo, jiangjy1982, suexu1025 and vipannalla as code owners July 28, 2026 12:14

SurbhiJainUSC reviewed Jul 28, 2026

View reviewed changes

shralex added 4 commits July 28, 2026 16:40

SurbhiJainUSC approved these changes Jul 28, 2026

View reviewed changes

gulsumgudukbay suggested changes Jul 28, 2026

View reviewed changes

	applied to tests without external_* or tpu_only marker.
	all tests except external_serving / external_training.

		with DECOUPLE_GCLOUD=TRUE, conftest.py automatically
		selects at runtime only those tests that can run in

Uh oh!

Conversation

shralex commented Jul 28, 2026

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

SurbhiJainUSC Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

SurbhiJainUSC Jul 28, 2026 •

edited

Loading