fix ci

Author: gulsumgudukbay

.github/workflows/run_tests_against_package.yml

@@ -97,9 +97,11 @@ jobs:
         XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
         TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
         TPU_SKIP_MDS_QUERY: ${{ inputs.device_type == 'cpu' && '1' || '' }}
+        # ROCm installs the wheel with --no-deps plus a pinned requirements file; never use TPU wheel extras.
         MAXTEXT_PACKAGE_EXTRA: >-
           ${{ 
-            !contains(inputs.pytest_marker, 'not post_training') && 'tpu-post-train' 
+            inputs.device_type == 'rocm' && 'rocm'
+            || !contains(inputs.pytest_marker, 'not post_training') && 'tpu-post-train' 
             || (inputs.device_type == 'cpu' && 'tpu' || inputs.device_type) 
           }}
         ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency)
@@ -134,7 +136,11 @@ jobs:
               uv pip install -r "${{ inputs.requirements_file }}"
             fi
             uv pip install ${maxtext_wheel} --no-deps
-            uv pip install -r src/dependencies/github_deps/pre_train_deps.txt
+            # When a requirements file is set (e.g. decoupled or rocm-unit JAX stacks), it already
+            # carries test/runtime pins; otherwise add GitHub-sourced pins (no TPU extra / scripts).
+            if [ -z "${{ inputs.requirements_file }}" ]; then
+              uv pip install -r src/dependencies/extra_deps/pre_train_github_deps.txt --no-deps
+            fi
           else
             if [ -n "${{ inputs.requirements_file }}" ]; then
               echo "Installing requirements from ${{ inputs.requirements_file }}"

.github/workflows/run_tests_coordinator.yml

@@ -125,8 +125,8 @@ jobs:
             "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training",
             "cpu-unit": "cpu_only and not post_training",
             "cpu-post-training-unit": "cpu_only and post_training",
-            "rocm-unit": "not cpu_only and not tpu_only and not integration_test and not post_training and not decoupled",
-            "rocm-decoupled": "decoupled"
+            "rocm-unit": "not cpu_only and not tpu_only and not decoupled",
+            "rocm-decoupled": "not cpu_only and not tpu_only and decoupled"
           }')[inputs.flavor] }}
 
       pytest_addopts: >-
@@ -153,7 +153,7 @@ jobs:
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
             "cpu-post-training-unit": "",
-            "rocm-unit": "--ignore=tests/post_training",
+            "rocm-unit": "",
             "rocm-decoupled": ""
           }')[inputs.flavor] }}
         ${{ inputs.additional_pytest_args }}