Upgrade dependencies (now supports and defaults to Python 3.12 and Ubuntu 24.04)

[SamuelMarks]

Description

Upgrade dependencies

• support 3.12 trivially
• support python3 -m pip install -r requirements.txt with no special flags

(I also sorted the dependencies so it's easier to debug and add/remove deps in the future)

FWIW Python 3.13 cannot yet be supported, but this patch did make the dependencies installable:

diff --git a/requirements.txt b/requirements.txt
index 37c759234..5592143d8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,9 +7,9 @@ datasets
 flax>=0.10.6
 gcsfs
 google-api-python-client
-google-cloud-aiplatform==1.61.0
+google-cloud-aiplatform
 google-cloud-monitoring
-google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
+#google-jetstream@git+https://github.com/AI-Hypercomputer/JetStream.git
 grain[parquet]>=0.2.6
 huggingface_hub
 jax>=0.4.30
@@ -31,11 +31,11 @@ pyink
 pylint
 pytest
 pytype
-sentencepiece==0.2.0
+#sentencepiece==0.2.0
 tensorboard-plugin-profile
 tensorboardx>=2.6.2.2
 tensorflow-datasets
-tensorflow-text>=2.17.0
-tensorflow>=2.16.0
+#tensorflow-text>=2.17.0
+#tensorflow>=2.16.0
 tiktoken
 transformers
diff --git a/setup.py b/setup.py
index 647025ba6..61f0ff1ed 100644
--- a/setup.py
+++ b/setup.py
@@ -80,7 +80,7 @@ def main():
     long_description_content_type="text/markdown",
     classifiers=[
       "Intended Audience :: Developers",
-      "License :: OSI Approved :: Apache Software License",
+#      "License :: OSI Approved :: Apache Software License",
       "Programming Language :: Python :: 3 :: Only",
       "Programming Language :: Python :: 3.10",
       "Programming Language :: Python :: 3.11",

Tests

N/A

Checklist

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

• I have performed a self-review of my code.
• 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.

[bvandermoon]

Thank you @SamuelMarks. I assume we will also need to make changes in constraints_gpu.txt and requirements_with_jax_ai_image.txt?

[bvandermoon]

Is it possible to test this setup with a MaxText workload/base model?

[SamuelMarks]

@bvandermoon Sure happy to update those also. @shralex was talking about making 3.12 the new default.

How about this:

• FROM python:3.10-slim-bullseye → FROM python:3.12-slim-bullseye
• FROM ubuntu:22.04 → FROM ubuntu:24.04

(not sure the process for upgrading us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack the candidate/gpu:latest and candidate/tpu:latest)

Can upgrade them all in this PR or separate it out into multiple. Let me know 🤗

[Rohan-Bierneni]

I tried building a maxtext image by forking this pr and found that the build fails when using MODE=stable or MODE=nightly. The commands I ran were

bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.7.0
bash docker_build_dependency_image.sh MODE=nightly

The build fails due to this line in setup.sh. We are pinning a version of setuptools that is incompatible with python 3.12.

I was able to get an image built successfully by upgrading the setuptools version with this command:

pip3 install --upgrade setuptools

@gobbleturk would it be fine to upgrade setuptools since it was pinned initially to fix the bug in cloud documentation: b/402501203

[SamuelMarks]

@Rohan-Bierneni Interesting. Yeah in #1954 I explicitly upgrade setuptools—also pip and wheel—to their latest versions for pytesting across 3.10, 3.11, and 3.12. Confirmed it all works fine there…

[bvandermoon]

Thank you @SamuelMarks. Please confirm there won't any recent requirements/dependency changes that need to be updated/rebased here too.

Have you been able to test the Docker images still build, MaxText still trains on a real TPU, etc.? The unit tests are good for a lot of that but they only run one of the images we provide. Please run these tests before merging if you haven't already

[Rohan-Bierneni]

For the file requirements_with_jax_stable_stack_0_6_1_pipreqs.txt, shouldn't we leave the setuptools pin as is since it was tested that it is compatible with 0.6.1. Essentially all the deps in that file should be pinned for reproducibility with jax 0.6.1.

[SamuelMarks]

For the file requirements_with_jax_stable_stack_0_6_1_pipreqs.txt, shouldn't we leave the setuptools pin as is since it was tested that it is compatible with 0.6.1. Essentially all the deps in that file should be pinned for reproducibility with jax 0.6.1.

I don't think it needs that pinning. It doesn't break 3.10, 3.11, 3.12. We should, IMHO, make the dependencies as widely supporting of underlying CPython version as possible.

[Rohan-Bierneni]

For the file requirements_with_jax_stable_stack_0_6_1_pipreqs.txt, shouldn't we leave the setuptools pin as is since it was tested that it is compatible with 0.6.1. Essentially all the deps in that file should be pinned for reproducibility with jax 0.6.1.

I don't think it needs that pinning. It doesn't break 3.10, 3.11, 3.12. We should, IMHO, make the dependencies as widely supporting of underlying CPython version as possible.

Sounds good that makes sense

[Rohan-Bierneni]

LGTM. Manually adding pull_ready tag since old CPUTests check is overriden in pr with new configuration. This causes a gap in the pr checklist where old test is never ran while new test passes. Once pr merges, this change should be fixed. Ref: https://screenshot.googleplex.com/s4M6YpsecC3UNrJ