From 34211553d84c957b05ee143a74c5754be03a4d97 Mon Sep 17 00:00:00 2001 From: khatwanimohit Date: Wed, 17 Jun 2026 04:17:51 +0000 Subject: [PATCH] upgrade post training reqs --- .../maxtext_tpu_dependencies.Dockerfile | 2 +- .../extra_deps/post_train_github_deps.txt | 6 +- .../tpu-post-train-requirements.txt | 2 + .../tpu-post-train-requirements.txt | 62 ++++++++----------- 4 files changed, 33 insertions(+), 39 deletions(-) diff --git a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile index 8980d5b3be..7ad2d69bde 100644 --- a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile @@ -1,6 +1,6 @@ # syntax=docker/dockerfile:experimental -ARG BASEIMAGE=python:3.12-slim-bullseye +ARG BASEIMAGE=python:3.12-slim-bookworm FROM $BASEIMAGE # Install system dependencies diff --git a/src/dependencies/extra_deps/post_train_github_deps.txt b/src/dependencies/extra_deps/post_train_github_deps.txt index 02730e98e7..4a8fa7d37b 100644 --- a/src/dependencies/extra_deps/post_train_github_deps.txt +++ b/src/dependencies/extra_deps/post_train_github_deps.txt @@ -1,3 +1,3 @@ -google-tunix @ https://github.com/google/tunix/archive/683256db1a0919b5cfd46cee52cebc96331494fb.zip -tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/4d08971683a64fd796a1b9fd0bb71128188882d5.zip -vllm @ git+https://github.com/vllm-project/vllm@2131b597b18d051dced4c4a605d362fa37f46ed1 +google-tunix @ https://github.com/google/tunix/archive/b02231b8102a3c36d681f339f0bc058b8d197c11.zip +tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/dc7889800a47f8e03028c3910684f03bbb1aceab.zip +vllm @ git+https://github.com/vllm-project/vllm@a30addc7548a9a8b9b3323a7bc3eb7d7c4895d1c diff --git a/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt index 3af5f4967e..ac88d8636d 100644 --- a/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt +++ b/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt @@ -4,6 +4,8 @@ astor blake3>=1.0 cachetools cbor2 +fastapi[standard]>=0.136.3 +fastar>=0.9.0 gepa gguf google-tunix diff --git a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt index a436b1e86d..85726d69dc 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt @@ -2,7 +2,7 @@ # See https://maxtext.readthedocs.io/en/latest/development/update_dependencies.html for details. absl-py>=2.4.0 -accelerate>=1.13.0 +accelerate>=1.14.0 aiofiles>=25.1.0 aiohappyeyeballs>=2.6.2 aiohttp>=3.14.1 @@ -10,9 +10,9 @@ aiohttp-cors>=0.1.0 aiosignal>=1.4.0 annotated-doc>=0.0.4 annotated-types>=0.7.0 -anthropic>=0.109.1 +anthropic>=0.109.2 antlr4-python3-runtime>=4.9.3 -anyio>=4.13.0 +anyio>=4.14.0 apache-tvm-ffi>=0.1.12 appnope>=0.1.4 ; sys_platform == 'darwin' aqtp>=0.9.0 @@ -36,7 +36,7 @@ cffi>=2.0.0 ; implementation_name == 'pypy' or platform_python_implementation != cfgv>=3.5.0 charset-normalizer>=3.4.7 cheroot>=11.1.2 -chex>=0.1.91 +chex>=0.1.92 click>=8.4.1 cloud-accelerator-diagnostics>=0.1.1 cloudpickle>=3.1.2 @@ -47,7 +47,7 @@ colorful>=0.0.0 comm>=0.2.3 compressed-tensors>=0.17.0 contourpy>=1.3.3 -cryptography>=48.0.1 +cryptography>=49.0.0 cycler>=0.12.1 dataclasses>=0.5 dataclasses-json>=0.0.1 @@ -57,16 +57,14 @@ decorator>=5.3.1 depyf>=0.20.0 dill>=0.4.1 diskcache>=5.6.3 -distlib>=0.4.2 +distlib>=0.4.3 distro>=1.9.0 dm-tree>=0.1.10 -dnspython>=2.0.0 docstring-parser>=0.18.0 -drjax>=0.1.4 +drjax>=0.2.0 editdistance>=0.8.1 einops>=0.8.2 einshape>=1.0 -email-validator>=2.0.0 entrypoints>=0.4 etils>=1.14.0 execnet>=2.1.2 @@ -88,19 +86,19 @@ gepa>=0.1.1 gguf>=0.19.0 google-api-core>=2.31.0 google-api-python-client>=2.197.0 -google-auth>=2.53.0 +google-auth>=2.55.0 google-auth-httplib2>=0.4.0 google-auth-oauthlib>=1.4.0 google-cloud-aiplatform>=1.157.0 google-cloud-appengine-logging>=1.10.0 google-cloud-audit-log>=0.6.0 -google-cloud-bigquery>=3.41.0 +google-cloud-bigquery>=3.42.0 google-cloud-core>=2.6.0 google-cloud-logging>=3.16.0 google-cloud-mldiagnostics>=1.0.3 google-cloud-monitoring>=2.31.0 google-cloud-resource-manager>=1.17.0 -google-cloud-storage>=3.11.0 +google-cloud-storage>=3.12.0 google-cloud-storage-control>=1.12.0 google-crc32c>=1.8.0 google-genai>=2.8.0 @@ -121,10 +119,9 @@ hf-transfer>=0.1.9 hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' httpcore>=1.0.9 httplib2>=0.31.2 -httptools>=0.8.0 httpx>=0.28.1 httpx-sse>=0.1.0 -huggingface-hub>=1.18.0 +huggingface-hub>=1.19.0 humanize>=4.15.0 hypothesis>=6.151.9 identify>=2.6.19 @@ -142,7 +139,7 @@ isort>=8.0.1 jaraco-functools>=4.5.0 jax>=0.10.1 jaxlib>=0.10.1 -jaxtyping>=0.3.10 +jaxtyping>=0.3.11 jedi>=0.20.0 jinja2>=3.1.6 jiter>=0.15.0 @@ -154,7 +151,7 @@ jupyter-client>=8.9.1 jupyter-core>=5.9.1 jupyterlab-widgets>=3.0.16 kagglehub>=1.0.2 -kagglesdk>=0.1.28 +kagglesdk>=0.1.30 keras>=3.14.1 kiwisolver>=1.5.0 lark>=1.2.2 @@ -179,11 +176,11 @@ mdurl>=0.1.2 mistral-common>=1.11.3 ml-collections>=1.1.0 ml-dtypes>=0.5.4 -ml-goodput-measurement>=0.0.16 -model-hosting-container-standards>=0.1.15 +ml-goodput-measurement>=0.2.0 +model-hosting-container-standards>=0.1.16 more-itertools>=11.1.0 mpmath>=1.3.0 -msgpack>=1.1.2 +msgpack>=1.2.0 msgspec>=0.21.1 multidict>=6.7.1 multiprocess>=0.70.19 @@ -200,8 +197,8 @@ numba>=0.65.1 numpy>=2.1.3 numpy-typing-compat>=20251206.2.1 oauthlib>=3.3.1 -omegaconf>=2.3.0 -openai>=2.41.1 +omegaconf>=2.3.1 +openai>=2.42.0 openai-harmony>=0.0.8 opencensus>=0.0.1 opencv-python-headless>=4.13.0.90 @@ -229,7 +226,7 @@ parameterized>=0.9.0 parso>=0.8.7 partial-json-parser>=0.2.1.1.post7 pathspec>=1.1.1 -pathwaysutils>=0.1.8 +pathwaysutils>=0.1.9 peft>=0.19.1 perfetto>=0.56.0 pexpect>=4.9.0 ; sys_platform != 'emscripten' and sys_platform != 'win32' @@ -260,38 +257,34 @@ pycparser>=3.0 ; (implementation_name != 'PyPy' and platform_python_implementati pydantic>=2.13.4 pydantic-core>=2.46.4 pydantic-extra-types>=2.11.1 -pydantic-settings>=2.0.0 pydot>=4.0.1 pyelftools>=0.32 pyglove>=0.4.5 pygments>=2.20.0 pyink>=25.12.0 -pylint>=4.0.5 +pylint>=4.0.6 pynvml>=13.0.1 -pyopenssl>=26.2.0 +pyopenssl>=26.3.0 pyparsing>=3.3.2 pyproject-hooks>=1.2.0 pytest>=8.4.2 pytest-mock>=3.15.1 pytest-xdist>=3.8.0 python-dateutil>=2.9.0.post0 -python-discovery>=1.4.0 +python-discovery>=1.4.2 python-dotenv>=1.2.2 python-json-logger>=4.1.0 -python-multipart>=0.0.18 pytokens>=0.4.1 pytype>=2024.10.11 pyyaml>=6.0.3 pyzmq>=27.1.0 -qwix>=0.1.6 +qwix>=0.1.7 ray>=2.55.1 referencing>=0.37.0 regex>=2026.5.9 requests>=2.34.2 requests-oauthlib>=2.0.0 rich>=15.0.0 -rich-toolkit>=0.14.8 -rignore>=0.5.1 rpds-py>=2026.5.1 runai-model-streamer>=0.15.9 runai-model-streamer-gcs>=0.15.9 @@ -302,7 +295,6 @@ scipy>=1.17.1 scipy-stubs>=1.17.1.4 semantic-version>=2.8.2 sentencepiece>=0.2.1 -sentry-sdk>=2.20.0 seqio>=0.0.20 setproctitle>=1.3.7 setuptools>=78.1.0 @@ -317,7 +309,7 @@ sniffio>=1.3.1 sortedcontainers>=2.4.0 sse-starlette>=0.1.0 stack-data>=0.6.3 -starlette>=1.2.1 +starlette>=1.3.1 supervisor>=4.3.0 sympy>=1.14.0 tabulate>=0.10.0 @@ -345,7 +337,7 @@ tornado>=6.5.7 tpu-info>=0.7.1 tqdm>=4.68.2 traitlets>=5.15.1 -transformers>=5.11.0 +transformers>=5.12.1 treescope>=0.1.10 triton>=3.6.0 ; sys_platform == 'linux' typeguard>=2.13.3 @@ -357,7 +349,7 @@ uritemplate>=4.2.0 urllib3>=2.7.0 uvicorn>=0.49.0 uvloop>=0.22.1 -virtualenv>=21.4.2 +virtualenv>=21.5.1 wadler-lindig>=0.1.7 watchfiles>=1.2.0 wcwidth>=0.8.1 @@ -367,7 +359,7 @@ wheel>=0.47.0 widgetsnbextension>=4.0.15 win32-setctime>=1.2.0 ; sys_platform == 'win32' wrapt>=2.2.1 -xgrammar>=0.2.1 +xgrammar>=0.2.2 xprof>=2.22.3 xxhash>=3.7.0 yapf>=0.43.0