Fix ImportError by pinning tensorflow and tensorflow-text to 2.20.x

[chiajunglien]

Description

Fixes ImportError: /usr/local/lib/python3.12/site-packages/tensorflow_text/core/pybinds/libpybinds_library_common.so: undefined symbol: icudt77_dat when running multimodal SFT tests (e.g., test_qwen3_multimodal_sft.sh).

DAG Error Log

Root Cause

1. A recent requirement regeneration commit (bde759abd) automatically bumped tensorflow and tensorflow-text to >=2.21.0 in the post-train requirements.
2. Version 2.21.0 of tensorflow-text introduces a dependency on the ICU 77 library (icudt77_dat), which is currently missing from the base workload container environment.
3. (Dependency Generation) During local test regeneration, seed-env using extreme lowest resolution (--resolution=lowest) attempted to fetch markdown==2.6.8 (required by tensorboard). This ancient version fails to build on Python 3.12 due to the removal of the imp module.

Solution

1. Explicitly pinned tensorflow==2.20.0 and tensorflow-text==2.20.1 in src/dependencies/extra_deps/tpu_post_train_overrides.txt to rollback to the working versions.
2. Added markdown>=3.4.1 to the overrides to prevent the seed-env resolution failure on Python 3.12.
3. Regenerated the src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt lockfile.

Tests

Gemma3-4b Multimodal SFT: https://cloudlogging.app.goo.gl/TgW1NRpZNEGPHMk27
Qwen3-vl-2b Multimodal SFT: https://cloudlogging.app.goo.gl/N61H2VC3UUPNMiK59

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.

[gemini-code-assist]

Code Review

This pull request updates several dependency versions in the TPU post-train requirements, including attempting to pin tensorflow and tensorflow-text to version 2.20.x to prevent an ImportError caused by version 2.21.0. However, because the generation script automatically converts exact pins (==) to minimum bounds (>=), the generated requirements still allow installing the problematic 2.21.0 version. To resolve this, the generation script should be updated to preserve exact pins for these specific packages.

[gemini-code-assist]

The PR aims to pin tensorflow and tensorflow-text to 2.20.x to avoid the ImportError caused by version 2.21.0. However, because generate_requirements.sh automatically converts all exact pins (==) to minimum bounds (>=), the generated requirements file ends up with tensorflow>=2.20.0 and tensorflow-text>=2.20.1.

When users or CI environments install from tpu-post-train-requirements.txt, package managers (like pip or uv) will resolve to the latest available versions (e.g., 2.21.0 or newer), which re-introduces the ImportError and completely defeats the purpose of this PR.

To fix this, we should update generate_requirements.sh to preserve the exact pins (==) for tensorflow and tensorflow-text during the post-processing step. For example, you can add the following lines to the post-processing block in src/dependencies/scripts/generate_requirements.sh:

# Restore exact pins for tensorflow and tensorflow-text to prevent upgrading to 2.21.0+
sed -i 's/^\(tensorflow\)>=/\1==/' "$ARTIFACT_DIR/$GENERATED_REQUIREMENTS"
sed -i 's/^\(tensorflow-text\)>=/\1==/' "$ARTIFACT_DIR/$GENERATED_REQUIREMENTS"

After updating the script, please regenerate the requirements file.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!