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 🤗