Gcluster Job Submission Guide

This guide provides a step-by-step process to deploy a GKE cluster, submit a sample Python script as a job using gcluster job submit with on-the-fly image building, and then destroy the cluster.

1. Prerequisites

Before you begin, ensure you have installed gcluster and set up your environment by following the instructions in the cluster-toolkit README.md.

Conditional Prerequisites (For On-the-Fly Builds)

If you use --build-context to build images on-the-fly, you must set:

• GCLUSTER_IMAGE_REPO: The name of your Artifact Registry repository only (e.g., gcluster-repo). The tool will automatically construct the full path using the cluster's region and project ID.
• USER or USERNAME: Used for unique image tagging (usually set automatically by your OS).

1.1 Multi-Tier Checkpointing (MTC) Prerequisites

If you plan to use Multi-Tier Checkpointing (--gke-mtc-enabled flag), ensure that the Multi-Tier Checkpointing feature is enabled on your GKE cluster.

To use this feature, the cluster administrator must also ensure that the required CheckpointConfiguration Custom Resource is deployed to the cluster, specifying the target cloud storage bucket for checkpoints. Once the cluster is configured, job submitters simply pass --gke-mtc-enabled to their jobs. For more details on cluster configuration, see the GKE Multi-Tier Checkpointing documentation.

2. Prepare Sample Application Code

Create a directory named job_details and place your application files inside it. This will serve as your build context for the job. The tool will package all files in this directory and add them to the image.

cluster-toolkit/job_details/app.py

# app.py
print("Hello from the gcluster job submit application!")
print("This is a sample application running on GKE.")

3. Deploy a GKE Cluster

For this example, we'll deploy a basic GKE cluster using the hpc-gke.yaml blueprint.

gcloud config set project <PROJECT_ID>
gcloud config set compute/region <REGION/ZONE> # Or your preferred region

3.1 Create the Deployment Directory

Create the deployment directory using the hpc-gke.yaml blueprint:

./gcluster create examples/hpc-gke.yaml --vars="project_id=<PROJECT_ID>,deployment_name=<CLUSTER_NAME>,region=<REGION/ZONE>,gcp_public_cidrs_access_enabled=false,authorized_cidr=$(curl -s ifconfig.me)/32"

Note: Please ensure that you replace all placeholders enclosed in angle brackets (such as <PROJECT_ID>, <CLUSTER_NAME>, and <REGION/ZONE>) with your actual environment values before executing the command.

3.2 Deploy the GKE Cluster

Deploy the GKE cluster:

./gcluster deploy <CLUSTER_NAME>

This command will show a Terraform plan. You will be prompted to confirm the changes (type a and press Enter).

This deployment process can take a significant amount of time (e.g., 10-20 minutes or more) as it provisions cloud resources. Wait for the command to complete successfully.

4. Submit the Sample Job

Now that the cluster is deployed and your application code is prepared, you can submit your sample Python script as a JobSet job. gcluster job submit will automatically build your container image and push it to Artifact Registry in your project.

• You must set the GCLUSTER_IMAGE_REPO environment variable to specify the name of the Artifact Registry repository when using --build-context for on-the-fly builds (e.g., export GCLUSTER_IMAGE_REPO=gcluster-repo). The tool will automatically construct the full path using the cluster's region and project ID. The command will fail fast if this variable is not set. The repository must exist before submitting the job. If it does not exist, you can create it with:

  gcloud artifacts repositories create <REPOSITORY_NAME> \
      --repository-format=docker \
      --location=<REGION>
  

• You must have either USER or USERNAME environment variable set when using --build-context (usually set automatically by your OS). gcluster uses this to ensure unique image tagging (e.g., my-user-runner:tag). The command will fail if both are missing.

4.1 Unified Job Submission

By specifying the --compute-type flag, you can use the exact same command to target a standard CPU cluster (using a full GCE machine type like n2-standard-32), an accelerated GPU cluster (using a GKE accelerator type like nvidia-l4), or a TPU cluster (using a shorthand string representing total chips/cores like v6e-8). The tool will automatically resolve the machine type, calculate num-nodes, and deduce the correct TPU topology if needed.

./gcluster job config set project <PROJECT_ID>
./gcluster job config set cluster <CLUSTER_NAME>
./gcluster job config set location <REGION/ZONE>

./gcluster job config list

4.2 Submit the Job

./gcluster job submit \
  --project <PROJECT_ID> \
  --cluster <CLUSTER_NAME> \
  --location <REGION/ZONE> \
  --base-image python:3.9-slim \
  --build-context job_details \
  --command "python app.py" \
  --name my-python-app-job \
  --compute-type n2-standard-32

Note: Please ensure that you replace all placeholders enclosed in angle brackets (such as <PROJECT_ID>, <CLUSTER_NAME>, and <REGION/ZONE>) with your actual environment values before executing the command.

This command will:

1. Verify/install the JobSet CRD on your cluster.
2. Auto-discover the Kueue LocalQueue name from the cluster.
3. Use the compute type installed on the cluster nodes and map the necessary resource requests.
4. Build a container image from the job_details directory using python:3.9-slim as the base, and push it to Artifact Registry.
5. Generate and apply an intelligently configured Kubernetes JobSet manifest to your cluster.

Note: The following examples assume you have configured your default project, cluster, and location using ./gcluster job config set.

4.3 Example for Multi-Slice GPU Job

If you want to run a job across multiple groups of GPU nodes (e.g., 2 groups of 4 nodes each), you can use --num-slices and --num-nodes:

./gcluster job submit \
  --image us-docker.pkg.dev/my-project/my-repo/my-image:latest \
  --command "python train.py" \
  --name my-gpu-job \
  --compute-type l4-1 \
  --num-slices 2 \
  --num-nodes 4

This creates a JobSet with 2 replicas, each having 4 pods, totaling 8 nodes.

4.4 Example: Submit Job with Persistent Storage

You can mount Cloud Storage buckets, Filestore instances, existing PVCs (e.g., for Lustre), or host paths using the --mount flag.

Mounts must use the format: --mount "<src>;<dest>[;<mode>][;options=<options>]"

• mode is optional and defaults to ro (read-only). To allow writes, append ;rw.
• options is optional and allows passing custom mount options (currently only supported for GCS fuse volumes).

Supported volume sources (<src>):

• Cloud Storage: gs://<bucket-name> (mounts via GCS Fused Driver)
• Filestore: filestore://<instance-name-or-ip>/<share-name> (auto-provisions PV and PVC)
• Existing PVC: <pvc-name> (the PersistentVolumeClaim must already exist in the target Kubernetes namespace)
• Host Path: /host/path/on/node (mounts a directory directly from the host node)

Examples:

Mounting a GCS bucket (read-write):

./gcluster job submit \
  --name my-storage-job \
  --command "python app.py" \
  --compute-type n2-standard-32 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --mount "gs://<YOUR_BUCKET_NAME>;/data;rw;options=logging:severity:info,enable-atomic-rename-object:true"

Mounting an existing PVC named lustre-pvc (read-only):

./gcluster job submit \
  --name my-lustre-job \
  --command "python app.py" \
  --compute-type n2-standard-32 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --mount "lustre-pvc;/data"

4.5 Example: Submit Job with Custom Environment Variables

You can pass custom environment variables to the container using the --env flag:

./gcluster job submit \
  --name my-env-job \
  --command "python app.py" \
  --compute-type n2-standard-32 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --env "TRAINING_EPOCHS=10" \
  --env "DEBUG=true"

5. Verify the Job

Verify that the Kubernetes JobSet ran successfully on your GKE cluster.

• Check Job Status: You can check the status of your submitted job directly with gcluster job list:

  ./gcluster job list
  

  Look for my-python-app-job with a Succeeded status.

• Get Job Logs: You can view the logs of your submitted job directly with gcluster job logs:

  ./gcluster job logs my-python-app-job
  

  You should see the output:

  Hello from the gcluster job submit application!
  This is a sample application running on GKE.
  

• Cancel Jobs: You can clean up a specific job without destroying the entire cluster:

  ./gcluster job cancel my-python-app-job
  

  Verify it's gone by running gcluster job list again.

• Inspect Cluster and Workload Health: If you encounter scheduling delays, errors, or suspect resource exhaustion, you can run gcluster job inspect to capture a comprehensive diagnostic sweep of your cluster state and active workloads.

  ./gcluster job inspect
  

  To debug a specific workload, provide its name to fetch details like Kueue Workload status and JobSet configurations:

  ./gcluster job inspect --name my-python-app-job
  

  To view logs output in the console in addition to saving them to the diagnostic file, add the --show (or -s) flag:

  ./gcluster job inspect --name my-python-app-job --show
  

  The tool will create a timestamped log file gcluster-inspect-<cluster>-<timestamp>.log in your current working directory containing:

  • Local Setup: Gcloud version and active configuration.
  • GKE Infra: GKE cluster descriptions, node-pool listings, and metadata/resources ConfigMaps.
  • Node Status: Wide listing of nodes, along with Go-calculated counts of total and healthy nodes per node pool.
  • Kueue / JobSet: Configurations and logs for Kueue and JobSet controller managers.
  • Slice Controller: Slice controller deployment details and manager logs (if GKE Kueue dynamic slicing is active).
  • Workloads: Overview of all workloads in the cluster, and specific JobSet/Workload descriptors if a name is targeted.
  • Console Links: Direct links to GKE clusters, GKE workloads, IAM permissions, and Quota administration consoles.

./gcluster job submit ... --gke-namespace custom-namespace
./gcluster job list --gke-namespace custom-namespace
./gcluster job logs my-python-app-job --gke-namespace custom-namespace
./gcluster job cancel my-python-app-job --gke-namespace custom-namespace
./gcluster job inspect --gke-namespace custom-namespace
./gcluster job inspect --name my-python-app-job --gke-namespace custom-namespace

6. Advanced Workloads

Note: The following examples assume you have configured your default project, cluster, and location using ./gcluster job config set.

6.1 Run with Advanced Scheduling Flags

Try running a job with advanced scheduling options.

Example 1: Target a Specific Node Pool (Node Constraint) Use --node-constraint to target a specific node pool. This maps to node labels in GKE and aligns with SLURM's --constraint flag for future compatibility.

./gcluster job submit \
  --name my-nodepool-job \
  --command "python app.py" \
  --compute-type c2-standard-60 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --node-constraint "cloud.google.com/gke-nodepool=my-custom-nodepool"

Example 2: Target Multiple Node Pools (Fallback Scheduling) Use a pipe separator (|) in --node-constraint to specify multiple allowed values for standard constraints. This generates a nodeAffinity block instead of a strict nodeSelector, allowing the workload to schedule on any of the specified pools.

./gcluster job submit \
  --name my-fallback-job \
  --command "python app.py" \
  --compute-type c2-standard-60 \
  --node-constraint "cloud.google.com/gke-nodepool=nodepool-a|nodepool-b"

Example 3: TPU Sub-Slicing If you want to schedule a job with a smaller TPU topology on a node pool with larger TPU slices (e.g., running a logical 2x2 workload on a physical 4x4 node pool):

• Specify the logical slice size in the --topology flag: --topology 2x2 (or let it auto-discover from shorthand labels like --compute-type v6e-4)
• Specify the single physical topology size of the target node pool in --node-constraint: --node-constraint "cloud.google.com/gke-tpu-topology=4x4"

./gcluster job submit \
  --name my-subslicing-job \
  --command "python app.py" \
  --compute-type v6e-4 \
  --node-constraint "cloud.google.com/gke-tpu-topology=4x4"

Example 4: Use Placement Policy Use --placement-policy to specify a GCE Placement Policy (e.g., for compact placement to reduce latency).

./gcluster job submit \
  ... \
  --name my-compact-job \
  --placement-policy "compact-placement"

(Note: requires a PlacementPolicy resource named compact-placement to exist on the cluster)

Example 5: Pod Failure Policy Use --restart-on-exit-codes to specify retriable exit codes at the pod level (these do not count against the restarts budget).

./gcluster job submit \
  ... \
  --name my-robust-job \
  --restart-on-exit-codes 1,137

Example 6: Private Registry & Service Account Use --image-pull-secret and --service-account for secure jobs.

./gcluster job submit \
  ... \
  --name my-secure-job \
  --image-pull-secret "my-private-registry-secret" \
  --service-account "my-workload-sa"

Example 7: Explicit Kueue Queue Selection Use --queue to submit the job to a specific Kueue LocalQueue.

./gcluster job submit \
  --name my-kueue-job \
  --command "python app.py" \
  --compute-type n2-standard-32 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --queue "my-local-queue"

(Note: You would need to ensure a Kueue LocalQueue named my-local-queue is configured on your cluster.)

Example 8: Targeting Provisioning Models (Spot/Reservation) on GKE NAP Clusters

Use the --gke-nap-provisioning spot option to run your workload on Spot VMs:

./gcluster job submit \
  --name my-spot-job \
  --command "python app.py" \
  --compute-type v6e-4 \
  --gke-nap-provisioning spot

Use the --gke-nap-provisioning reservation option and target a GCE reservation via the --gke-nap-reservation flag:

./gcluster job submit \
  --name my-reservation-job \
  --command "python app.py" \
  --compute-type v6e-4 \
  --gke-nap-provisioning reservation \
  --gke-nap-reservation my-tpu-reservation

6.3 Job Retention (TTL)

By default, finished jobs are kept for 1 hour. You can change this using --gke-ttl-after-finished and pass flexible durations.

./gcluster job submit ... --gke-ttl-after-finished 10m # Keep for only 10 minutes
./gcluster job submit ... --gke-ttl-after-finished 2h  # Keep for 2 hours

6.4 Graceful Termination (Grace Period)

You can give your workloads a buffer period to save checkpoints or perform cleanups before they are forcefully killed using --grace-period.

./gcluster job submit ... --grace-period 2m # Allow 2 minutes for cleanup

6.5 Topology & Scheduler

Example 1: Topology Awareness Request a specific TPU slice topology using --topology.

./gcluster job submit \
  --project <PROJECT_ID> \
  --cluster <CLUSTER_NAME> \
  --location <REGION/ZONE> \
  --name my-topology-job \
  --base-image python:3.9-slim \
  --build-context job_details \
  --command "python app.py" \
  --compute-type tpu-v6e-slice \
  --topology 4x4

Example 2: Scheduler Selection Use a specific GKE scheduler (e.g., gke.io/topology-aware-auto) using --gke-scheduler.

./gcluster job submit \
  --project <PROJECT_ID> \
  --cluster <CLUSTER_NAME> \
  --location <REGION/ZONE> \
  --name my-scheduler-job \
  --command "python app.py" \
  --compute-type n2-standard-32 \
  --base-image python:3.9-slim \
  --build-context job_details \
  --gke-scheduler gke.io/topology-aware-auto

7. Sophisticated Workloads: MaxText

7.1 Llama3.1-8B on TPU v6e

This section describes how to deploy a more complex workload, specifically training a Llama3.1-8B model using MaxText on a TPU v6e cluster.

7.1.1 Prepare MaxText Workload Directory

Create a directory named maxtext_workload_v6e and place the following files inside it.

cluster-toolkit/maxtext_workload_v6e/requirements.txt

# This is your requirements.txt file

cluster-toolkit/maxtext_workload_v6e/Dockerfile

# Use the recommended base image
FROM us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.6.1-rev1

# Set the working directory
WORKDIR /deps

# Install system dependencies
RUN apt-get update && apt-get install -y dnsutils

# Install MaxText dependencies
RUN pip install google-cloud-monitoring

# Clone MaxText
RUN git clone https://github.com/AI-Hypercomputer/maxtext.git /app \
    && cd /app \
    && git checkout tpu-recipes-v0.1.4

# Skip explicit pip install because base image has prerequisites

# Set working directory to MaxText root for the runner
WORKDIR /app

# Copy the wrapper script
COPY run_maxtext.sh /app/run_maxtext.sh
RUN chmod +x /app/run_maxtext.sh

# Entrypoint is left to default or overridden by gcluster

cluster-toolkit/maxtext_workload_v6e/run_maxtext.sh

#!/bin/bash

# Exit on error
set -e

echo "Starting MaxText Workload..."

# 1. Set environment variables
# Combine all the required XLA flags into LIBTPU_INIT_ARGS
export LIBTPU_INIT_ARGS=" --xla_tpu_scoped_vmem_limit_kib=98304 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_assign_all_reduce_scatter_layout=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_all_reduce_offload_tracing=true --xla_tpu_use_tc_device_shape_on_sc=true --xla_sc_enable_instruction_fusion=false --xsc_disjoint_spmem=false --xla_sc_disable_megacore_partitioning=true --2a886c8_chip_config_name=megachip_tccontrol --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_tpu_enable_ag_backward_pipelining=true --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_jf_spmd_threshold_for_windowed_einsum_mib=1000000"

export JAX_PLATFORMS="tpu,cpu"
export ENABLE_PJRT_COMPATIBILITY=true

# 2. Extract arguments
OUTPUT_DIR=${1}
if [ -z "$OUTPUT_DIR" ]; then
  echo "Error: Output directory argument missing."
  echo "Usage: \$0 <output_gcs_bucket>"
  exit 1
fi

MODEL_NAME=${2:-"llama3.1-8b"}

echo "LIBTPU_INIT_ARGS=$LIBTPU_INIT_ARGS"
echo "OUTPUT_DIR=$OUTPUT_DIR"
echo "MODEL_NAME=$MODEL_NAME"

# 3. Run training
python3 -m MaxText.train MaxText/configs/base.yml \
    per_device_batch_size=3 \
    ici_fsdp_parallelism=-1 \
    remat_policy=custom \
    decoder_layer_input=offload \
    out_proj=offload \
    query_proj=offload \
    key_proj=offload \
    value_proj=offload \
    max_target_length=8192 \
    attention=flash \
    use_iota_embed=True \
    dataset_path=gs://max-datasets-rogue \
    dataset_type=synthetic \
    enable_checkpointing=False \
    sa_block_q=2048 \
    sa_block_kv=2048 \
    sa_block_kv_compute=2048 \
    sa_block_q_dkv=2048 \
    sa_block_kv_dkv=2048 \
    sa_block_kv_dkv_compute=2048 \
    sa_block_q_dq=2048 \
    sa_block_kv_dq=2048 \
    sa_use_fused_bwd_kernel=True \
    profiler=xplane \
    skip_first_n_steps_for_profiler=10 \
    profiler_steps=5 \
    steps=20 \
    monitor_goodput=True \
    enable_goodput_recording=True \
    model_name=$MODEL_NAME \
    base_output_directory=$OUTPUT_DIR \
    use_vertex_tensorboard=false

cluster-toolkit/maxtext_workload_v6e/build.sh

#!/bin/bash

# Get current project
PROJECT=$(gcloud config get-value project)

if [ -z "$PROJECT" ]; then
  echo "Error: Could not determine GCP project. Please run 'gcloud config set project <PROJECT_ID>'"
  exit 1
fi

IMAGE_NAME=gcr.io/$PROJECT/maxtext-runner:latest

echo "Building image $IMAGE_NAME using Cloud Build..."
gcloud builds submit --tag $IMAGE_NAME .

echo "Image built successfully!"
echo "You can now submit the job with:"
echo "  gcluster job submit --image $IMAGE_NAME --command 'bash run_maxtext.sh <OUTPUT_DIR>'"

cluster-toolkit/maxtext_workload_v6e/submit.sh

#!/bin/bash

# Configuration - UPDATE THESE
CLUSTER_NAME="v6e-cluster"
LOCATION="us-central1"
OUTPUT_DIR="gs://$PROJECT/maxtext_output"
SA_NAME="v6e-cluster-gke-wl-sa"

# Look up project
PROJECT=$(gcloud config get-value project)

if [ -z "$PROJECT" ]; then
  echo "Error: Could not determine GCP project. Please run 'gcloud config set project <PROJECT_ID>'"
  exit 1
fi

IMAGE_NAME=gcr.io/$PROJECT/maxtext-runner:latest

echo "Ensuring permissions for $SA_NAME..."
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/logging.logWriter" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/monitoring.metricWriter" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/logging.viewer" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.objectViewer" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:$SA_NAME@${PROJECT}.iam.gserviceaccount.com" --role="roles/artifactregistry.reader" --quiet > /dev/null

echo "Submitting MaxText job to cluster $CLUSTER_NAME..."

# Navigate to cluster-toolkit root if running from maxtext_workload_v6e
if [ -f "../gcluster" ]; then
  GCLUSTER="../gcluster"
else
  GCLUSTER="./gcluster"
fi

$GCLUSTER job submit \
    --name maxtext-llama3-1-final-v6e8-2 \
    --cluster $CLUSTER_NAME \
    --location $LOCATION \
    --image $IMAGE_NAME \
    --command "cd /app && pip install psutil jaxtyping tiktoken sentencepiece ray fastapi uvicorn portpicker pydantic ninja Pillow gcsfs omegaconf jsonlines PyYAML safetensors tabulate tensorstore transformers datasets evaluate nltk pandas ml_collections ml_dtypes pathwaysutils orbax grain tensorflow_text tensorflow_datasets tqdm && sed -i 's/use_vertex_tensorboard=false/use_vertex_tensorboard=false run_name=llama3-1-v6e8-test1/g' run_maxtext.sh && bash run_maxtext.sh $OUTPUT_DIR" \
    --compute-type v6e-8 \
    --num-slices 1 \
    --topology 2x4 \
    --priority medium \
    --service-account workload-identity-k8s-sa

7.1.2 Build and Submit

cd maxtext_workload_v6e
./build.sh
./submit.sh

7.1.3 Verify Job and Logs

You can verify the job status and check logs using gcluster or kubectl.

Using gcluster:

# List jobs
./gcluster job list --project $PROJECT --cluster $CLUSTER_NAME --location $LOCATION

# View logs
./gcluster job logs maxtext-llama3-1-final-v6e8-2 --project $PROJECT --cluster $CLUSTER_NAME --location $LOCATION

Using kubectl:

# Get pods
kubectl get pods --namespace default -l jobset.sigs.k8s.io/jobset-name=maxtext-llama3-1-final-v6e8-2

# View logs for a specific pod
kubectl logs <POD_NAME> --namespace default

7.2 Llama3.1-8B on TPU v7x

This section describes how to deploy the MaxText workload specifically optimized for TPU v7x (Ironwood) hardware.

7.2.1 Prepare MaxText v7x Workload Directory

Create a directory named maxtext_workload_v7x and place the following files inside it.

cluster-toolkit/maxtext_workload_v7x/requirements.txt

psutil
jaxtyping
tiktoken
sentencepiece
ray
fastapi
uvicorn
portpicker
pydantic
ninja
Pillow
gcsfs
omegaconf
jsonlines
PyYAML
safetensors
tabulate
tensorstore
transformers
datasets
evaluate
nltk
pandas
ml_collections
ml_dtypes
pathwaysutils
orbax
grain
tensorflow_text
tensorflow_datasets
tqdm

cluster-toolkit/maxtext_workload_v7x/Dockerfile

# Use the recommended base image for TPU7x
FROM us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.8.2-rev1

# Set the working directory
WORKDIR /deps

# Install system dependencies
RUN apt-get update && apt-get install -y dnsutils

# Install MaxText dependencies
RUN pip install google-cloud-monitoring

# Install Python requirements
COPY requirements.txt .
RUN pip install -r requirements.txt

# Clone MaxText
RUN git clone https://github.com/AI-Hypercomputer/maxtext.git /app \
    && cd /app \
    && git checkout maxtext-tutorial-v1.0.0

# Set working directory to MaxText root for the runner
WORKDIR /app

# Copy the wrapper script
COPY run_maxtext.sh /app/run_maxtext.sh
RUN chmod +x /app/run_maxtext.sh

# Entrypoint is left to default or overridden by gcluster

cluster-toolkit/maxtext_workload_v7x/run_maxtext.sh

#!/bin/bash

# Exit on error
set -e

echo "Starting MaxText Workload..."

# 1. Set environment variables
export ENABLE_PATHWAYS_PERSISTENCE='1'
# Combine all the required XLA flags into LIBTPU_INIT_ARGS
export LIBTPU_INIT_ARGS=" --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true "

export JAX_PLATFORMS="tpu,cpu"
export ENABLE_PJRT_COMPATIBILITY=true
export PYTHONPATH=$PYTHONPATH:$(pwd)/src
export JAX_TRACEBACK_FILTERING=off

# 2. Extract arguments
OUTPUT_DIR=${1}
if [ -z "$OUTPUT_DIR" ]; then
  echo "Error: Output directory argument missing."
  echo "Usage: \$0 <output_gcs_bucket>"
  exit 1
fi

MODEL_NAME=${2:-"llama3.1-8b"}

echo "LIBTPU_INIT_ARGS=$LIBTPU_INIT_ARGS"
echo "OUTPUT_DIR=$OUTPUT_DIR"
echo "MODEL_NAME=$MODEL_NAME"

# 3. Run training
python3 src/MaxText/train.py src/MaxText/configs/base.yml \
    model_name=$MODEL_NAME \
    skip_jax_distributed_system=False \
    dtype=bfloat16 \
    per_device_batch_size=1 \
    ici_fsdp_parallelism=64 \
    max_target_length=4096 \
    profiler=xplane \
    profile_periodically_period=10000 \
    async_checkpointing=False \
    enable_checkpointing=False \
    use_iota_embed=True \
    remat_policy=custom \
    decoder_layer_input=offload \
    query_proj=offload \
    key_proj=offload \
    value_proj=offload \
    out_proj=offload \
    dataset_type=synthetic \
    opt_type=adamw \
    mu_dtype=bfloat16 \
    tokenizer_type=tiktoken \
    tokenizer_path=assets/tokenizer_llama3.tiktoken \
    sa_use_fused_bwd_kernel=True \
    attention=flash \
    steps=30 \
    base_output_directory=$OUTPUT_DIR \
    use_vertex_tensorboard=false

cluster-toolkit/maxtext_workload_v7x/build.sh

#!/bin/bash

# Get current project
PROJECT=$(gcloud config get-value project)

if [ -z "$PROJECT" ]; then
  echo "Error: Could not determine GCP project. Please run 'gcloud config set project <PROJECT_ID>'"
  exit 1
fi

IMAGE_NAME=gcr.io/$PROJECT/maxtext-runner:latest

echo "Building image $IMAGE_NAME using Cloud Build..."
gcloud builds submit --tag $IMAGE_NAME .

echo "Image built successfully!"
echo "You can now submit the job with:"
echo "  gcluster job submit --image $IMAGE_NAME --command 'bash run_maxtext.sh <OUTPUT_DIR>'"

cluster-toolkit/maxtext_workload_v7x/submit.sh

#!/bin/bash

# Configuration - UPDATE THESE
CLUSTER_NAME="tpu7x-cluster"
LOCATION="<REGION/ZONE>"
OUTPUT_DIR="gs://$PROJECT/maxtext_output_7x"

# Look up project
PROJECT=$(gcloud config get-value project)

if [ -z "$PROJECT" ]; then
  echo "Error: Could not determine GCP project. Please run 'gcloud config set project <PROJECT_ID>'"
  exit 1
fi

IMAGE_NAME=gcr.io/$PROJECT/maxtext-runner:latest

echo "Ensuring permissions for tpu7x-cluster-gke-wl-sa..."
# Note: Ensure the SA matches the one created by the 7x blueprint (tpu7x-cluster-gke-wl-sa)
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/logging.logWriter" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/monitoring.metricWriter" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/logging.viewer" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.objectViewer" --quiet > /dev/null

echo "Ensuring permissions for node pool service account..."
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/artifactregistry.reader" --quiet > /dev/null
gcloud projects add-iam-policy-binding $PROJECT --member="serviceAccount:tpu7x-cluster-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.objectViewer" --quiet > /dev/null

echo "Submitting MaxText job to cluster $CLUSTER_NAME..."

# Navigate to cluster-toolkit root if running from maxtext_workload_v7x
if [ -f "../gcluster" ]; then
  GCLUSTER="../gcluster"
else
  GCLUSTER="./gcluster"
fi

$GCLUSTER job submit \
    --name maxtext-llama3-1-final-tpu7x-32 \
    --cluster $CLUSTER_NAME \
    --location $LOCATION \
    --image $IMAGE_NAME \
    --command "cd /app && sed -i 's/use_vertex_tensorboard=false/use_vertex_tensorboard=false run_name=llama3-1-7x-test1/g' run_maxtext.sh && bash run_maxtext.sh $OUTPUT_DIR" \
    --compute-type tpu7x-32 \
    --num-slices 1 \
    --topology 2x4x4 \
    --priority medium \
    --service-account workload-identity-k8s-sa

To disable parallel containers and use a single container per VM, add the --gke-disable-parallel-containers flag:

$GCLUSTER job submit \
    --name maxtext-llama3-1-final-tpu7x-32 \
    --cluster $CLUSTER_NAME \
    --location $LOCATION \
    --image $IMAGE_NAME \
    --command "cd /app && sed -i 's/use_vertex_tensorboard=false/use_vertex_tensorboard=false run_name=llama3-1-7x-test1/g' run_maxtext.sh && bash run_maxtext.sh $OUTPUT_DIR" \
    --compute-type tpu7x-32 \
    --num-slices 1 \
    --topology 2x4x4 \
    --priority medium \
    --service-account workload-identity-k8s-sa \
    --gke-disable-parallel-containers

7.2.2 Build and Submit

cd maxtext_workload_v7x
./build.sh
./submit.sh

7.2.3 Verify Job and Logs

TPU v7x utilizes Megacore, which initializes 2 logical devices per chip. For a 32-chip slice (2x4x4 topology), you should see 64 logical devices in the logs.

Using gcluster:

# List jobs
./gcluster job list --project $PROJECT --cluster $CLUSTER_NAME --location $LOCATION

# View logs
./gcluster job logs maxtext-llama3-1-final-tpu7x-32 --project $PROJECT --cluster $CLUSTER_NAME --location $LOCATION

Verification Highlights: In the logs, look for successful JAX initialization and step processing:

System Information: Jax Version: 0.8.2
System Information: Jax Backend: PJRT C API
TFRT TPU7x
Num_devices: 64, shape (1, 1, 64, 1, 1, 1, 1, 1, 1, 1, 1, 1)
...
completed step: 4, seconds: 1.182, TFLOP/s/device: 167.154, Tokens/s/device: 3464.384, loss: 10.187
completed step: 5, seconds: 1.186, TFLOP/s/device: 166.597, Tokens/s/device: 3452.840, loss: 9.184

8. Advanced GKE infrastructure features

This section details how the gcluster job submit command orchestrates advanced Google Kubernetes Engine (GKE) hardware and scheduling capabilities, such as dynamic slice reconfiguration, the Pathways distributed AI framework, and GKE Node Auto-Provisioning (NAP).

8.1 Dynamic slicing (TPU v7x and future generations)

GKE Dynamic Slicing provides flexible TPU capacity scheduling by letting you logically group or slice physical nodes dynamically. In Cluster Toolkit and GKE, dynamic slicing is supported starting with TPU v7x (Ironwood) and future TPU generations (earlier generations like TPU v4, TPU v5e, TPU v5p, and TPU v6e require static slice topologies). Dynamic slicing supports both superslicing (aggregating multiple physical cubes into larger topologies) and subslicing (partitioning a single cube into smaller fractional topologies) dynamically at job submission time.

Capabilities and scheduling benefits

Dynamic slicing provides the following scheduling capabilities:

• Dynamic superslicing: Aggregate multiple physical TPU v7x cubes together into a larger logical slice (such as combining multiple 4x4x4 cubes into 4x4x8 or 4x4x16 topologies) dynamically for large-scale distributed training.
• Dynamic subslicing: Partition a single physical TPU cube into smaller fractional topologies (such as slicing a 4x4x4 cube into 2x2x4 or 2x4x4 sub-slices) dynamically, enabling efficient bin-packing and co-tenancy for smaller workloads.
• Latency optimization: Kueue Topology-Aware Scheduling (TAS) places TPU pods with minimal network hop latency across the physical TPU interconnect mesh, maximizing distributed training throughput.
• Multi-slice synchronization: For multi-slice jobs (--num-slices > 1), Cluster Toolkit coordinates slice-level reservations to ensure that all slices are acquired concurrently, preventing mismatched scaling.
• Internal scheduling annotations: Cluster Toolkit automatically translates high-level TPU requests into partition-level requirements (cloud.google.com/gke-tpu-partition-TOPOLOGY-id) and switches between single-slice (kueue.x-k8s.io/podset-required-topology) and multi-slice (kueue.x-k8s.io/podset-slice-required-topology) admission annotation keys.

Example CLI command

To submit a dynamic slicing workload targeting TPU v7x nodes, enter the following command:

./gcluster job submit \
  --name my-dynamic-slice-job \
  --command "python train.py" \
  --compute-type tpu-v7x-slice \
  --topology 4x4x4

GKE documentation reference

• TPU Dynamic Slicing on GKE Concepts
• Scheduling Dynamic Slices with Kueue and TAS on GKE

8.2 Pathways distributed AI framework

Pathways is a distributed execution framework for large-scale multi-slice machine learning workloads. The gcluster job submit command provides end-to-end integration for deploying Pathways-enabled jobs on GKE.

Capabilities and workload orchestration roles

Pathways integration provides the following orchestration capabilities:

• Multi-role JobSet orchestration: When you specify the --pathways flag, Cluster Toolkit generates and coordinates a multi-role JobSet comprising the Resource Manager (pathways-rm), Proxy server (pathways-proxy), and JAX worker pods.
• Kueue resource management: Cluster Toolkit automatically maps multi-slice resource limits and unifies Kueue ClusterQueue resource groups so that coordinator and worker pods schedule without resource deadlocks.
• Headless execution mode: With the --pathways-headless flag, Cluster Toolkit provisions the Pathways server infrastructure without requiring a client workload container, letting external clients (such as local notebooks) connect to the cluster.

Example CLI commands

To submit a standard multi-slice Pathways training workload, enter the following command:

./gcluster job submit \
  --name my-pathways-job \
  --command "python train_pathways.py" \
  --compute-type v6e-16 \
  --pathways \
  --pathways-gcs-location gs://BUCKET_NAME/pathways-artifacts

To submit a headless Pathways server cluster, enter the following command:

./gcluster job submit \
  --name my-pathways-headless \
  --compute-type v6e-16 \
  --pathways \
  --pathways-gcs-location gs://BUCKET_NAME/pathways-artifacts \
  --pathways-headless

8.3 Node auto-provisioning (NAP) and compute consumption

GKE Node Auto-Provisioning (NAP) automatically provisions and manages node pools based on unschedulable pod requirements.

Capabilities, scheduling, and pre-flight verification

NAP integration provides the following capabilities:

• Consumption targeting: The --gke-nap-provisioning flag lets you target spot, on-demand, or reservation compute models.
• Reservation integration: The --gke-nap-reservation RESERVATION_NAME flag injects reservation affinity selectors, letting NAP spawn node pools directly in specific Compute Engine reservations.
• Pre-flight limit verification: Cluster Toolkit validates requested machine types against cluster NAP limits before job submission, failing quickly with a diagnostic error if limits would be exceeded.

Example CLI commands

To target Spot VMs, enter the following command:

./gcluster job submit \
  --name my-nap-spot-job \
  --command "python app.py" \
  --compute-type v6e-4 \
  --gke-nap-provisioning spot

To target a Compute Engine reservation, enter the following command:

./gcluster job submit \
  --name my-nap-reservation-job \
  --command "python app.py" \
  --compute-type v6e-4 \
  --gke-nap-provisioning reservation \
  --gke-nap-reservation RESERVATION_NAME

8.4 Custom Kubernetes manifest templates

GCluster uses embedded Go templates to generate the Kubernetes manifests (JobSets, Kueue configurations, etc.) deployed to your GKE cluster. If you need to deeply customize the generated manifests beyond what the CLI flags provide (e.g., adding custom sidecars, specific annotations, or modifying the core structure), you can override these embedded templates with your own local copies.

1. Extract the default templates

First, use the gke-template-extract command to extract the embedded default templates into a local directory so you have a starting point:

./gcluster job gke-template-extract --output-dir=/path/to/my-templates

(Use --overwrite if the directory already contains files you want to replace).

2. Modify the templates

Open the extracted .tmpl files in your preferred editor. These are standard Go text templates (text/template and yamltemplate). You can inject any static Kubernetes YAML or modify the existing template variables.

3. Use your custom templates

When submitting your job, you can tell gcluster to use your custom templates instead of the embedded ones.

Pass the --gke-custom-templates-path flag to the submit command:

./gcluster job submit ... --gke-custom-templates-path=/path/to/my-templates

9. gcluster job Command Reference

9.1 Common Flags

These flags are common to almost all gcluster job subcommands (except config). project, cluster, and location can be set as defaults via config set.

9.2 Configuration Commands

Use these commands to manage persistent defaults for your job submissions, avoiding the need to pass common flags repeatedly.

gcluster job config set [key] [value]

Sets a persistent configuration property.

• Supported Keys:
  • project: Google Cloud Project ID
  • cluster: GKE Cluster Name
  • location: GKE Cluster Location (region or zone)

Example:

./gcluster job config set project my-awesome-project

gcluster job config list

Lists all persistent configuration properties currently set.

9.3 submit Flags

The gcluster job submit command deploys a container image as a job (Kubernetes JobSet) on a GKE cluster, integrated with Kueue for advanced queuing. It can use pre-built images or build images on-the-fly without a local Docker daemon (powered internally by the Crane container utility).

9.3.1 General Flags

These flags control core workload execution, identity, and basic build context options regardless of the underlying hardware paradigm.

(Note: --cluster, --location, and --project are also supported as common flags, see 9.1)

9.3.2 TPU Related Flags

Use these flags to orchestrate specialized TPU multi-slice topologies or leverage the Pathways distributed execution framework.

9.3.3 GPU Related Flags

Use these flags to tune specialized multi-GPU topologies and related node parameters.

9.3.4 GKE & Advanced Orchestration Flags

These flags unlock advanced GKE capabilities, including Kueue queue priorities, workload identities, network placements, and robust lifecycle failure policies.

9.4 list Flags

Use these flags to filter the list of jobs.

9.5 logs Flags

Use these flags when fetching logs.

10. Troubleshooting: ImagePullBackOff

If your job status remains Pending and the underlying pods show ImagePullBackOff or ErrImagePull, the GKE node pool service account may lack permission to read from the Artifact Registry repository.

A project administrator can grant the necessary access manually by running:

gcloud artifacts repositories add-iam-policy-binding <REPOSITORY_NAME> \
    --location <REGION/ZONE> \
    --project <PROJECT_ID> \
    --member "serviceAccount:<GKE_NODE_SERVICE_ACCOUNT>" \
    --role "roles/artifactregistry.reader"

11. Cleanup

To avoid incurring unnecessary costs, destroy the deployed GKE cluster and its resources:

./gcluster destroy <CLUSTER_NAME>

You will be prompted to confirm the destruction (type a and press Enter).