Merge pull request #68 from AI-Hypercomputer/bvandermoon-tpu-recipes

Recommend Python 3.10 in MaxText instructions and other instruction c…

Author: bvandermoon

training/trillium/MAXTEXT_README.md

@@ -1,14 +1,14 @@
-# Prep for Maxtext workloads on GKE
+# Prep for MaxText workloads on GKE
 
 > **_NOTE:_** We recommend running these instructions and kicking off your recipe 
-workloads from a TPU VM.
+workloads from a VM in GCP using Python 3.10.
 
-1. Clone [Maxtext](https://github.com/google/maxtext) repo and move to its directory
+1. Clone [MaxText](https://github.com/google/maxtext) repo and move to its directory
 ```shell
 git clone https://github.com/google/maxtext.git
 cd maxtext
 # Checkout either the commit id or MaxText tag. 
-# Example: `git checkout tpu-recipes-v0.1.1`
+# Example: `git checkout tpu-recipes-v0.1.2`
 git checkout ${MAXTEXT_COMMIT_ID_OR_TAG}
 ```
 
@@ -18,7 +18,7 @@ bash setup.sh
 ```
 
 Optional: Use a virtual environment to setup and run your workloads. This can help with errors
-like `This environment is externally managed`.
+like `This environment is externally managed`:
 ```shell
 ## One time step of creating the venv
 VENV_DIR=~/venvp3
@@ -29,6 +29,10 @@ source $VENV_DIR/bin/activate
 bash setup.sh
 ```
 
+> **_NOTE:_** If you use a virtual environment, you must use the same one when running the 
+[XPK Installation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation) 
+steps linked in the [XPK_README](XPK_README.md) as well as your relevant tpu-recipe workloads.
+
 3. Run the following commands to build the docker image
 ```shell
 # Example BASE_IMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.5.2-rev1

training/trillium/XPK_README.md

@@ -1,7 +1,7 @@
 ## Initialization
 
 > **_NOTE:_** We recommend running these instructions and kicking off your recipe 
-workloads from a TPU VM.
+workloads from a VM in GCP using Python 3.10.
 
 1. Run the following commands to initialize the project and zone.
 ```shell
@@ -14,8 +14,10 @@ gcloud config set compute/zone $ZONE
 2. Install XPK by following the [prerequisites](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites) and [installation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation) 
 instructions. Also ensure you have the proper [GCP permissions](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation).
 
-* In order to run the tpu-recipes as-is, run the `git clone` command from your home directory:
+* In order to run the tpu-recipes as-is, run the `git clone` command from your home (~/) directory:
 ```shell
+# tpu-recipes requiring XPK will look for it in the home directory
+cd ~/
 git clone https://github.com/google/xpk.git
 ```
 
@@ -25,6 +27,11 @@ git clone https://github.com/google/xpk.git
 cd xpk # Should be equivalent to cd ~/xpk
 ```
 
+> **_NOTE:_** If you use a virtual environment in the 
+[XPK Installation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation)
+steps, you must use the same one to run the steps in the [MAXTEXT_README](MAXTEXT_README.md)
+as well as your relevant tpu-recipe workloads.
+
 ## GKE Cluster Creation 
 1. Specify your TPU GKE cluster configs.
 ```shell