Merge pull request #3585 from AI-Hypercomputer:bvandermoon-uxr-fixes

PiperOrigin-RevId: 896056051

Author: Google-ML-Automation

docs/build_maxtext.md

@@ -60,7 +60,9 @@ 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
+# This enables Docker image building and workload scheduling via XPK.
+# Once installed, you will have access to the `build_maxtext_docker_image`
+# and `upload_maxtext_docker_image` commands.
 uv pip install maxtext[runner]==0.2.1 --resolution=lowest
 ```
 
@@ -80,8 +82,11 @@ 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
+# Install MaxText with the [runner] extra in editable mode.
+# This enables Docker image building and workload scheduling via XPK.
+# Once installed, you will have access to the `build_maxtext_docker_image`
+# and `upload_maxtext_docker_image` commands.
+uv pip install -e .[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.

docs/install_maxtext.md

@@ -45,7 +45,7 @@ install_tpu_pre_train_extra_deps
 
 # Option 2: Installing maxtext[cuda12]
 uv pip install maxtext[cuda12]==0.2.1 --resolution=lowest
-install_cuda12_pre_train_extra_dep
+install_cuda12_pre_train_extra_deps
 
 # Option 3: Installing maxtext[tpu-post-train]
 uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
@@ -55,6 +55,8 @@ install_tpu_post_train_extra_deps
 uv pip install maxtext[runner]==0.2.1 --resolution=lowest
 ```
 
+> **Note:** The `maxtext[runner]` extra is used for building MaxText Docker images and scheduling workloads through XPK. Once installed, you will have access to the `build_maxtext_docker_image`, `upload_maxtext_docker_image`, and `xpk` commands. For more details on building and uploading Docker images, see the [Build MaxText Docker Image](https://maxtext.readthedocs.io/en/latest/build_maxtext.html) guide.
+
 > **Note:** The `install_tpu_pre_train_extra_deps`, `install_cuda12_pre_train_extra_deps`, and
 > `install_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
 > that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.