Merge pull request #2462 from AI-Hypercomputer:pyproj_req_update

PiperOrigin-RevId: 817323126

Author: Google-ML-Automation

.github/workflows/run_tests_against_package.yml

@@ -58,6 +58,7 @@ 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' || '' }}
+        MAXTEXT_PACKAGE_EXTRA: ${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}
       options: ${{ inputs.container_resource_option }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -66,8 +67,12 @@ jobs:
         with:
           name: maxtext-wheel
       - name: Install the maxtext wheel
+        shell: bash
         run: |
-          python3 -m uv pip install maxtext-*-py3-none-any.whl --resolution=lowest --system
+          python3 -m uv venv --seed
+          source .venv/bin/activate
+          maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
+          uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
           install_maxtext_github_deps
           python3 --version
           python3 -m pip freeze
@@ -76,6 +81,7 @@ jobs:
       - name: Display contents of workspace # TODO: Remove
         run: ls -R
       - name: Run Tests
+        shell: bash
         run: |
           if [ "${{ inputs.is_scheduled_run }}" == "true" ]; then
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
@@ -88,4 +94,4 @@ jobs:
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
           # TODO: Fix the skipped tests and remove the deselect flags
-          python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"
+          .venv/bin/python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"

README.md

@@ -57,7 +57,7 @@ cd maxtext
 
 # 2. Install dependencies in editable mode
 pip install uv
-uv pip install -e . --resolution=lowest
+uv pip install -e .[tpu] --resolution=lowest
 install_maxtext_github_deps
 ```
 

base_requirements/requirements.txt

@@ -0,0 +1,42 @@
+absl-py
+aqtp
+array-record
+cloud-accelerator-diagnostics
+cloud-tpu-diagnostics
+datasets
+flax
+gcsfs
+google-api-python-client
+google-cloud-aiplatform
+google-cloud-monitoring
+grain[parquet]
+huggingface_hub
+jax!=0.7.1, !=0.7.2
+jaxlib!=0.7.1, !=0.7.2
+jaxtyping
+jsonlines
+ml-collections
+ml-goodput-measurement
+numpy
+omegaconf
+optax
+orbax-checkpoint
+pathwaysutils
+pillow
+pre-commit
+protobuf
+pyink
+pylint
+pytest
+pytype
+sentencepiece
+tensorboard-plugin-profile
+tensorboardx
+tensorflow-datasets
+tensorflow-text
+tensorflow
+tiktoken
+transformers
+qwix
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
+mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip

base_requirements/tpu-base-requirements.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+google-tunix

generated_requirements/tpu-requirements.txt

@@ -0,0 +1,235 @@
+# Generated by seed-env. Do not edit manually.
+# If you need to modify dependencies, please do so in the host requirements file and run seed-env again.
+
+absl-py>=2.3.1
+aiofiles>=24.1.0
+aiohappyeyeballs>=2.6.1
+aiohttp>=3.13.0
+aiosignal>=1.4.0
+annotated-types>=0.7.0
+antlr4-python3-runtime>=4.9.3
+anyio>=4.11.0
+aqtp>=0.9.0
+array-record>=0.8.1
+astroid>=3.3.11
+astunparse>=1.6.3
+attrs>=25.3.0
+auditwheel>=6.4.1
+black>=24.10.0
+blobfile>=3.1.0
+build>=1.2.2.post1
+cachetools>=6.2.0
+certifi>=2025.10.5
+cffi>=2.0.0 ; platform_python_implementation == 'PyPy'
+cfgv>=3.4.0
+charset-normalizer>=3.4.3
+chex>=0.1.90
+click>=8.3.0
+cloud-accelerator-diagnostics>=0.1.1
+cloud-tpu-diagnostics>=0.1.5
+cloudpickle>=3.1.1
+clu>=0.0.12
+colorama>=0.4.6
+contourpy>=1.3.2
+coverage>=7.10.7
+cycler>=0.12.1
+datasets>=4.2.0
+decorator>=5.2.1
+dill>=0.4.0
+distlib>=0.4.0
+dm-tree>=0.1.9
+docstring-parser>=0.17.0
+editdistance>=0.8.1
+einops>=0.8.1
+etils>=1.13.0
+evaluate>=0.4.6
+execnet>=2.1.1
+fastapi>=0.118.2
+filelock>=3.18.0
+flatbuffers>=25.2.10
+flax>=0.11.2
+fonttools>=4.59.0
+frozenlist>=1.8.0
+fsspec>=2025.7.0
+gast>=0.6.0
+gcsfs>=2025.7.0
+google-api-core>=2.26.0
+google-api-python-client>=2.184.0
+google-auth-httplib2>=0.2.0
+google-auth-oauthlib>=1.2.2
+google-auth>=2.41.1
+google-cloud-aiplatform>=1.120.0
+google-cloud-appengine-logging>=1.6.2
+google-cloud-audit-log>=0.3.3
+google-cloud-bigquery>=3.38.0
+google-cloud-core>=2.4.3
+google-cloud-logging>=3.12.1
+google-cloud-monitoring>=2.27.2
+google-cloud-resource-manager>=1.14.2
+google-cloud-storage>=2.19.0
+google-crc32c>=1.7.1
+google-genai>=1.42.0
+google-pasta>=0.2.0
+google-resumable-media>=2.7.2
+google-tunix>=0.1.1
+googleapis-common-protos>=1.70.0
+grain>=0.2.12
+grpc-google-iam-v1>=0.14.2
+grpcio-status>=1.71.2
+grpcio>=1.75.1
+gviz-api>=1.10.0
+h11>=0.16.0
+h5py>=3.14.0
+hf-xet>=1.1.10 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
+httpcore>=1.0.9
+httplib2>=0.31.0
+httpx>=0.28.1
+huggingface-hub>=0.35.3
+humanize>=4.13.0
+hypothesis>=6.136.4
+identify>=2.6.15
+idna>=3.10
+immutabledict>=4.2.1
+importlab>=0.8.1
+importlib-metadata>=8.7.0
+importlib-resources>=6.5.2
+iniconfig>=2.1.0
+isort>=6.1.0
+jax==0.7.0
+jaxlib==0.7.0
+jaxtyping>=0.3.3
+jinja2>=3.1.6
+joblib>=1.5.2
+jsonlines>=4.0.0
+kagglehub>=0.3.13
+keras>=3.11.3
+kiwisolver>=1.4.8
+libclang>=18.1.1
+libcst>=1.8.5
+libtpu>=0.0.19
+lxml>=6.0.2
+markdown-it-py>=3.0.0
+markdown>=3.9
+markupsafe>=3.0.3
+matplotlib>=3.10.3
+mccabe>=0.7.0
+mdurl>=0.1.2
+ml-collections>=1.1.0
+ml-dtypes>=0.5.3
+ml-goodput-measurement>=0.0.15
+more-itertools>=10.8.0
+mpmath>=1.3.0
+msgpack>=1.1.2
+msgspec>=0.19.0
+multidict>=6.7.0
+multiprocess>=0.70.16
+mypy-extensions>=1.1.0
+namex>=0.1.0
+nest-asyncio>=1.6.0
+networkx>=3.5
+ninja>=1.13.0
+nltk>=3.9.2
+nodeenv>=1.9.1
+numpy>=2.0.2
+oauthlib>=3.3.1
+omegaconf>=2.3.0
+opentelemetry-api>=1.37.0
+opt-einsum>=3.4.0
+optax>=0.2.6
+optree>=0.17.0
+orbax-checkpoint>=0.11.25
+packaging>=25.0
+pandas>=2.3.3
+parameterized>=0.9.0
+pathspec>=0.12.1
+pathwaysutils>=0.1.3
+pillow>=11.3.0
+platformdirs>=4.5.0
+pluggy>=1.6.0
+portpicker>=1.6.0
+pre-commit>=4.3.0
+prometheus-client>=0.23.1
+promise>=2.3
+propcache>=0.4.1
+proto-plus>=1.26.1
+protobuf>=5.29.5
+psutil>=7.0.0
+pyarrow>=21.0.0
+pyasn1-modules>=0.4.2
+pyasn1>=0.6.1
+pycnite>=2024.7.31
+pycparser>=2.23 ; implementation_name != 'PyPy' and platform_python_implementation == 'PyPy'
+pycryptodomex>=3.23.0
+pydantic-core>=2.41.1
+pydantic>=2.12.0
+pydot>=4.0.1
+pyelftools>=0.32
+pyglove>=0.4.5
+pygments>=2.19.2
+pyink>=24.10.1
+pylint>=3.3.9
+pyparsing>=3.2.3
+pyproject-hooks>=1.2.0
+pytest-xdist>=3.8.0
+pytest>=8.4.1
+python-dateutil>=2.9.0.post0
+python-dotenv>=1.1.1
+pytype>=2024.10.11
+pytz>=2025.2
+pyyaml>=6.0.3
+qwix>=0.1.1
+regex>=2025.9.18
+requests-oauthlib>=2.0.0
+requests>=2.32.5
+rich>=14.1.0
+rsa>=4.9.1
+safetensors>=0.6.2
+scipy>=1.16.0
+sentencepiece>=0.2.1
+seqio>=0.0.20
+setuptools>=80.9.0
+shapely>=2.1.2
+shortuuid>=1.0.13
+simple-parsing>=0.1.7
+simplejson>=3.20.2
+six>=1.17.0
+sniffio>=1.3.1
+sortedcontainers>=2.4.0
+starlette>=0.48.0
+sympy>=1.14.0
+tabulate>=0.9.0
+tenacity>=9.1.2
+tensorboard-data-server>=0.7.2
+tensorboard-plugin-profile>=2.13.0
+tensorboard>=2.19.0
+tensorboardx>=2.6.4
+tensorflow-datasets>=4.9.9
+tensorflow-metadata>=1.17.2
+tensorflow-text>=2.19.0
+tensorflow>=2.19.1
+tensorstore>=0.1.76
+termcolor>=3.1.0
+tiktoken>=0.12.0
+tokenizers>=0.22.1
+toml>=0.10.2
+tomlkit>=0.13.3
+toolz>=1.0.0
+tqdm>=4.67.1
+transformers>=4.57.0
+treescope>=0.1.10
+typing-extensions>=4.14.1
+typing-inspection>=0.4.2
+tzdata>=2025.2
+uritemplate>=4.2.0
+urllib3>=2.5.0
+uvicorn>=0.37.0
+virtualenv>=20.34.0
+wadler-lindig>=0.1.7
+websockets>=15.0.1
+werkzeug>=3.1.3
+wheel>=0.45.1
+wrapt>=1.17.3
+xxhash>=3.6.0
+yarl>=1.22.0
+zipp>=3.23.0
+zstandard>=0.23.0