Merge branch 'AI-Hypercomputer:main' into ai-gsutil-migration-4779320b60d44c658f872e9b583efb19

Author: gurusai-voleti

.github/workflows/build_and_test_maxtext.yml

@@ -262,7 +262,7 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
+      extra_pip_deps_file: 'src/dependencies/github_deps/post_train_base_deps.txt'
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
   maxtext_post_training_tpu_unit_tests:
@@ -284,7 +284,7 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
+      extra_pip_deps_file: 'src/dependencies/github_deps/post_train_base_deps.txt'
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
   maxtext_gpu_integration_tests:

.github/workflows/run_pathways_tests.yml

@@ -85,7 +85,7 @@ jobs:
           source .venv/bin/activate
           maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
           uv pip install ${maxtext_wheel}[tpu] --resolution=lowest
-          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
+          uv pip install -r src/dependencies/github_deps/pre_train_deps.txt
           python3 --version
           python3 -m pip freeze
       - name: Copy test assets files

.github/workflows/run_tests_against_package.yml

@@ -96,7 +96,7 @@ jobs:
           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
-          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
+          uv pip install -r src/dependencies/github_deps/pre_train_deps.txt
           python3 --version
           python3 -m pip freeze
           uv pip install pytest-cov

PREFLIGHT.md

@@ -1,35 +1,35 @@
 # Optimization 1: Multihost recommended network settings
-We included all the recommended network settings in [rto_setup.sh](https://github.com/google/maxtext/blob/main/rto_setup.sh). 
+We included all the recommended network settings in [rto_setup.sh](https://github.com/google/maxtext/blob/main/src/dependencies/scripts/rto_setup.sh). 
 
-[preflight.sh](https://github.com/google/maxtext/blob/main/preflight.sh) will help you apply them based on GCE or GKE platform.
+[preflight.sh](https://github.com/google/maxtext/blob/main/src/dependencies/scripts/preflight.sh) will help you apply them based on GCE or GKE platform.
 
 Before you run ML workload on Multihost with GCE or GKE, simply apply `bash preflight.sh PLATFORM=[GCE or GKE]` to leverage the best DCN network performance.
 
 Here is an example for GCE:
 ```
-bash preflight.sh PLATFORM=GCE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
+bash src/dependencies/scripts/preflight.sh PLATFORM=GCE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 Here is an example for GKE:
 ```
-bash preflight.sh PLATFORM=GKE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
+bash src/dependencies/scripts/preflight.sh PLATFORM=GKE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 # Optimization 2: Numa binding (You can only apply this to v4 and v5p)
 NUMA binding is recommended for enhanced performance, as it reduces memory latency and maximizes data throughput, ensuring that your high-performance applications operate more efficiently and effectively.
 
 For GCE, 
-[preflight.sh](https://github.com/google/maxtext/blob/main/preflight.sh) will help you install `numactl` dependency, so you can use it directly, here is an example:
+[preflight.sh](https://github.com/google/maxtext/blob/main/src/dependencies/scripts/preflight.sh) will help you install `numactl` dependency, so you can use it directly, here is an example:
 
 ```
-bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
+bash src/dependencies/scripts/preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 For GKE,
 `numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
 
 ```
-bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
+bash src/dependencies/scripts/preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 1. `numactl`: This is the command-line tool used for controlling NUMA policy for processes or shared memory. It's particularly useful on multi-socket systems where memory locality can impact performance.

docs/build_maxtext.md

@@ -0,0 +1,137 @@
+<!--
+ Copyright 2023-2026 Google LLC
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      https://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ -->
+
+# Build and Upload MaxText Docker Images
+
+This guide covers setting up a MaxText development environment and building container images for TPU and GPU workloads. These images can be used to run MaxText on GKE clusters with TPUs or GPUs, and are also required for running MaxText through XPK.
+
+## Prerequisites
+
+Before starting, ensure you have the following tools installed and configured:
+
+1. Environment Prep: Install and configure all [XPK prerequisites](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites).
+
+2. Docker Permissions: Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/) to run Docker without `sudo`.
+
+3. Artifact Registry Access: Authenticate with [Google Artifact Registry](https://docs.cloud.google.com/artifact-registry/docs/docker/authentication#gcloud-helper) for permission to push your images and other access.
+
+4. Authentication & Access: Run the following commands to authenticate your account and configure Docker:
+
+```bash
+# Authenticate your user account for gcloud CLI access
+gcloud auth login
+
+# Configure application default credentials for Docker and other tools
+gcloud auth application-default login
+
+# Configure Docker credentials and test your access
+gcloud auth configure-docker
+docker run hello-world
+```
+
+## Installation Modes
+
+We recommend building MaxText inside a Python virtual environment using `uv` for speed and dependency management.
+
+### Option 1: From PyPI (Recommended)
+
+This is the easiest way to get started with the latest stable version.
+
+```bash
+# Install uv, a fast Python package installer
+pip install uv
+
+# Create virtual environment
+export VENV_NAME=<your virtual env name> # e.g., docker_venv
+uv venv --python 3.12 --seed ${VENV_NAME?}
+source ${VENV_NAME?}/bin/activate
+
+# Install MaxText with the [runner] extra
+# This enables Docker image building and workload scheduling via XPK
+uv pip install maxtext[runner] --resolution=lowest
+```
+
+> **Note:** The `maxtext[runner]` extra includes all necessary dependencies for building MaxText Docker images and running workloads through XPK. It automatically installs XPK, so you do not need to install it separately to manage your clusters and workloads.
+
+### Option 2: From Source
+
+If you plan to contribute to MaxText or need the latest unreleased features, install from source.
+
+```bash
+# Clone the repository
+git clone https://github.com/AI-Hypercomputer/maxtext.git
+cd maxtext
+
+# Create virtual environment
+export VENV_NAME=<your virtual env name> # e.g., docker_venv
+uv venv --python 3.12 --seed ${VENV_NAME?}
+source ${VENV_NAME?}/bin/activate
+
+# Install MaxText with the [runner] extra in editable mode
+uv pip install .[runner] --resolution=lowest
+```
+
+> **Note:** The `maxtext[runner]` extra includes all necessary dependencies for building MaxText Docker images and running workloads through XPK. It automatically installs XPK, so you do not need to install it separately to manage your clusters and workloads.
+
+## Build MaxText Docker Image
+
+Select the appropriate build commands based on your hardware (`TPU` or `GPU`) and your specific workflow (`pre-training` or `post-training`). Each of these commands will generate a local Docker image named `maxtext_base_image`.
+
+### TPU Pre-Training Docker Image
+
+```bash
+# Option 1: Build with the stable versions of dependencies (default)
+build_maxtext_docker_image
+
+# Option 2: Build with latest nightly versions of jax/jaxlib
+build_maxtext_docker_image MODE=nightly
+
+# Option 3: Build with the specified jax/jaxlib version
+build_maxtext_docker_image MODE=nightly JAX_VERSION=$JAX_VERSION
+```
+
+### GPU Pre-Training Docker Image
+
+```bash
+# Option 1: Build with the stable versions of dependencies (default)
+build_maxtext_docker_image DEVICE=gpu
+
+# Option 2: Build with latest nightly versions of jax/jaxlib
+build_maxtext_docker_image DEVICE=gpu MODE=nightly
+
+# Option 3: Build with base image as `ghcr.io/nvidia/jax:base-2024-12-04`
+build_maxtext_docker_image DEVICE=gpu MODE=pinned
+
+# Option 4: Build with the specified jax/jaxlib version
+build_maxtext_docker_image DEVICE=gpu MODE=nightly JAX_VERSION=$JAX_VERSION
+```
+
+### TPU Post-Training Docker Image
+
+```bash
+# This build process takes approximately 10 to 15 minutes.
+build_maxtext_docker_image WORKFLOW=post-training
+```
+
+## Upload MaxText Docker Image to Artifact Registry
+
+> **Note:** You will need the [**Artifact Registry Writer**](https://docs.cloud.google.com/artifact-registry/docs/access-control#permissions) role to push Docker images to your project's Artifact Registry and to allow the cluster to pull them during workload execution. If you don't have this permission, contact your project administrator to grant you this role through "Google Cloud Console -> IAM -> Grant access".
+
+```bash
+# Make sure to replace <Docker Image Name> with your desired image name.
+export CLOUD_IMAGE_NAME=<Docker Image Name>
+upload_maxtext_docker_image CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME?}
+```

docs/conf.py

@@ -120,8 +120,8 @@
     os.path.join("run_maxtext", "run_maxtext_via_multihost_runner.md"),
     os.path.join("reference", "core_concepts", "llm_calculator.ipynb"),
     os.path.join("reference", "api_generated", "modules.rst"),
-    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.rst"),
-    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.install_github_deps.rst"),
+    os.path.join("reference", "api_generated", "dependencies.github_deps.rst"),
+    os.path.join("reference", "api_generated", "dependencies.github_deps.install_pre_train_deps.rst"),
 ]