Submit and monitor jobs on managed compute clusters

This document explains how to prepare your environment and run distributed workloads on a managed compute cluster in Cluster Director. Because managed compute clusters provision Compute Engine instances in managed instance groups (MIGs) without a Slurm scheduler, you run workloads by connecting to your compute nodes by using SSH or to your own orchestrator.

By following the steps in this document, you can perform the following tasks:

Configure passwordless SSH and shared software across compute nodes that mount a shared /home directory.
Run a command across all active compute nodes in one or more MIGs.
Verify GPU drivers and run multi-node Open MPI and GROMACS workloads on MIGs with GPU instances.
Run distributed JAX and MaxText training workloads on MIGs with TPU instances.

To learn how Cluster Director orchestrates managed compute clusters, see Managed compute orchestration in Cluster Director.

Before you begin

If you haven't already, then connect to the first compute node in your MIG by using SSH. This node is named MIG_NAME-001 where MIG_NAME is the name of your MIG. For instructions, see Connect to a cluster's compute node.

Configure a MIG

To configure a MIG in a managed compute cluster to run distributed workloads, complete the following steps:

Configure passwordless SSH and a host file
Install shared software
Build Open MPI with CUDA support

Configure passwordless SSH and a host file

When all compute instances in a MIG mount the same shared file system in your /home directory, they share the ~/.ssh directory. You can generate an SSH key pair on your first compute node so that MPI launchers and distributed workloads can communicate across compute instances in the MIG without requiring a password.

To configure passwordless SSH and create a host file, complete the following steps:

If a ~/.ssh directory doesn't exist, then create it:
```
mkdir -p ~/.ssh
```
Restrict the permissions of the ~/.ssh directory so that only your user account can read, write, or access it:
```
chmod 700 ~/.ssh
```
Generate an Ed25519 SSH key pair without a passphrase, and save the key pair to ~/.ssh/id_ed25519:
```
ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519
```

Append your public key to ~/.ssh/authorized_keys:

cat ~/.ssh/id_ed25519.pub >> ~/.ssh/authorized_keys

Restrict the file permissions to your user account so that all compute instances that mount /home trust the key:
```
chmod 600 ~/.ssh/authorized_keys
```

Configure SSH to accept host keys between cluster nodes without interactive prompts:

cat << 'EOF' >> ~/.ssh/config
Host *
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    LogLevel ERROR
EOF

Restrict permissions on the ~/.ssh/config file to your user account:
```
chmod 600 ~/.ssh/config
```
Create a hosts file in /home that lists the hostnames of the compute instances in your MIG:
```
cat << 'EOF' > ~/hosts
MIG_NAME-001
MIG_NAME-002
EOF
```
Replace MIG_NAME with the name of your MIG.
Verify that you can connect to the other node without a password prompt—for example, connect to the second compute node in your MIG:
```
ssh MIG_NAME-002 hostname
```

Install shared software

OS images for GPU and CPU machines don't include a pre-installed MPI implementation. You can install software either into a shared directory within the /home directory, or on each compute instance's boot disk by using a startup script:

Install shared software in /home: compile or install software once in a shared directory (such as /home/software) so that all compute instances that mount /home can run the binaries without additional setup:
1. On one compute node, create a shared directory in /home:
```
sudo mkdir -p /home/software
```
2. Grant read and write permissions so that your user account can install software in /home/software:
```
sudo chmod a+rw /home/software
```
3. Install your software or create a Python virtual environment in /home/software, and add the binary directory to your PATH environment variable on the compute nodes that run your workload:
```
export PATH=/home/software/APPLICATION_DIR/bin:$PATH
```
  Replace APPLICATION_DIR with the directory name of your installed software. For a complete example that builds Open MPI in /home/software, see Build Open MPI with CUDA support.
Install software with a custom startup script: configure a startup script on your MIG to install packages from apt (such as mpich or nvidia-cuda-toolkit) on every compute instance when it boots.

To configure a startup script on an existing MIG, select one of the following options:
Console
1. In the Google Cloud console, go to the Cluster Director page.
  
  Go to Cluster Director
2. In the Clusters table, click the name of your cluster, and then click Edit.
3. Go to the Orchestration step, click the MIG that you want to update, and expand Advanced MIG settings.
4. In the Startup script field, enter your bash script—for example:
  
  #!/bin/bash apt-get update && apt-get install -y mpich nvidia-cuda-toolkit
5. Optional: If the script takes longer than five minutes to finish, then enter a larger timeout (in seconds) in the Startup script timeout field (for example, 600).
6. Click Done, and then click Save.
gcloud
1. Create a startup script file at an absolute path (for example, /tmp/startup.sh):
  
  cat << 'EOF' > /tmp/startup.sh #!/bin/bash apt-get update && apt-get install -y mpich nvidia-cuda-toolkit EOF
2. Run the gcloud alpha cluster-director clusters update command with the --update-managed-instance-groups flag and pass the absolute path to your startup script file:
  
  gcloud alpha cluster-director clusters update CLUSTER_NAME \ --location=REGION \ --update-managed-instance-groups=id=MIG_NAME,startupScript=/tmp/startup.sh
  
  Replace the following:
  
  CLUSTER_NAME: the name of your cluster.
  
  REGION: the region where your cluster exists.
  
  MIG_NAME: the name of your MIG.

Build Open MPI with CUDA support

To build and verify Open MPI 5.0 with CUDA support in the /home/software directory, complete the following steps:

If the /home/software directory doesn't exist, then create it:
```
sudo mkdir -p /home/software
```
Grant read and write permissions so that your user account can install software in it:
```
sudo chmod a+rw /home/software
```

Download and extract the Open MPI 5.0 source archive:

cd /home/software
wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.10.tar.gz
tar -xzf openmpi-5.0.10.tar.gz
cd openmpi-5.0.10

Configure the Open MPI build with CUDA support:

./configure --prefix=/home/software/openmpi \
    --with-cuda=/usr/local/cuda \
    --with-cuda-libdir=/usr/local/cuda/lib64/stubs

Compile the source code across all available CPU cores:
```
make -j$(nproc)
```
Install the binaries in the /home/software/openmpi directory:
```
make install
```
Add the Open MPI bin directory to your PATH environment variable:
```
export PATH=/home/software/openmpi/bin:$PATH
```
Verify the Open MPI installation across the first two compute nodes in your ~/hosts file:
```
mpirun -np 2 -N 1 --hostfile ~/hosts hostname
```

Run a command across compute nodes in a MIG

To run a non-interactive command across all running compute nodes in one or more MIGs within a managed compute cluster, use the gcloud alpha cluster-director clusters ssh command in the Google Cloud CLI. Select one of the following options:

Bash

gcloud alpha cluster-director clusters ssh CLUSTER_NAME \
    --location=REGION \
    --managed-instance-groups=MIG_NAMES \
    --command='COMMAND'

Powershell

gcloud alpha cluster-director clusters ssh CLUSTER_NAME `
    --location=REGION `
    --managed-instance-groups=MIG_NAMES `
    --command='COMMAND'

cmd.exe

gcloud alpha cluster-director clusters ssh CLUSTER_NAME ^
    --location=REGION ^
    --managed-instance-groups=MIG_NAMES ^
    --command="COMMAND"

Replace the following required placeholders:

CLUSTER_NAME: the name of the managed compute cluster.
REGION: the region where the cluster exists.
MIG_NAMES: a comma-separated list of MIGs in the cluster where you want to run the command—for example, mig-01,mig-02.
COMMAND: the command to run on each active compute instance in the specified MIGs—for example, echo 'Hello, World!'.

Optionally, you can include one or more of the following flags in the command:

To limit the total number of compute instances across your specified MIGs that run the command at the same time, include the --batch-size flag:
```
--batch-size=BATCH_SIZE
```
Replace BATCH_SIZE with the maximum number of compute instances that run the command at the same time. If the command fails on any compute instance in a batch, then the command finishes that batch and skips the remaining batches.
To save each compute instance's command output to a separate log file, include the --output-directory flag:
```
--output-directory=OUTPUT_DIRECTORY
```
Replace OUTPUT_DIRECTORY with the path to an existing directory. The command saves each compute instance's output to a file named INSTANCE_NAME.log in that directory, where INSTANCE_NAME is the name of the compute instance.
To run the command inside a container on each compute instance, if the compute instances run a Container-Optimized OS image, include the --container flag:
```
--container=CONTAINER_NAME
```
Replace CONTAINER_NAME with the name of the container.
To specify the username for the SSH connection, include the --user flag:
```
--user=USERNAME
```
Replace USERNAME with the username to use. If you omit this flag, then the command uses your default login user.
To specify how to route the SSH connection, include one of the following flags:
- To connect to compute instances by using their internal IP addresses instead of their external IP addresses, include the --internal-ip flag:
```
--internal-ip
```
- To tunnel the SSH connection through Identity-Aware Proxy (IAP) for TCP forwarding, include the --tunnel-through-iap flag:
```
--tunnel-through-iap
```
To pass additional flags to ssh(1), include the --ssh-flag flag:
```
--ssh-flag="SSH_FLAG"
```
Replace SSH_FLAG with the flags to pass to ssh(1)—for example, -o Hostname=nic0.%NAME%.%ZONE%.c.%PROJECT%.internal.gcpnode.com. You can include the following runtime placeholders in SSH_FLAG, which the gcloud CLI replaces for each target compute instance:
- %NAME%: the name of the compute instance.
- %ZONE%: the zone where the compute instance exists.
- %PROJECT%: the ID of the project where the compute instance exists.
- %USER%: the SSH login username on the compute instance.
- %INSTANCE%: the target IP address for the connection, which is the external IP address by default, or the internal IP address if you include the --internal-ip flag in the command.
- %INTERNAL%: the internal IP address for the compute instance.

Run workloads on GPU MIGs

GPU instances in managed compute clusters include pre-installed NVIDIA drivers. You can verify GPU health with the nvidia-smi command and run distributed GPU workloads such as GROMACS across the nodes in your MIG.

To verify your GPUs and run a two-node GROMACS benchmark, complete the following steps:

On your GPU compute node, verify that all attached GPUs are online:
```
nvidia-smi
```
Install the CUDA toolkit and nvcc compiler on each node in your MIG (or include this command in your MIG startup script):
```
sudo apt-get update && sudo apt-get install -y nvidia-cuda-toolkit cuda-nvcc-13-0
```

On one compute node, download and extract the CMake source archive in /home/software:

cd /home/software
wget https://github.com/Kitware/CMake/releases/download/v4.4.2/cmake-4.4.2.tar.gz
tar -xzf cmake-4.4.2.tar.gz
cd cmake-4.4.2

Configure the CMake build to install in /home/software/cmake:
```
./configure --prefix=/home/software/cmake
```
Compile CMake across all available CPU cores:
```
gmake -j$(nproc)
```
Install the CMake binaries in /home/software/cmake:
```
make install
```
Add the CMake and Open MPI bin directories to your PATH environment variable:
```
export PATH=/home/software/cmake/bin:/home/software/openmpi/bin:$PATH
```

Download and extract the GROMACS source archive in /home/software:

cd /home/software
wget https://ftp.gromacs.org/gromacs/gromacs-2026.3.tar.gz
tar -xzf gromacs-2026.3.tar.gz
cd gromacs-2026.3

Create and enter a build directory for GROMACS:
```
mkdir build && cd build
```

Configure GROMACS with MPI and CUDA support:

cmake .. \
    -DCMAKE_INSTALL_PREFIX=/home/software/gromacs \
    -DGMX_MPI=ON \
    -DGMX_GPU=CUDA \
    -DGMX_BUILD_OWN_FFTW=ON \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_C_COMPILER=mpicc \
    -DCMAKE_CXX_COMPILER=mpicxx \
    -DCMAKE_CUDA_COMPILER=/usr/local/cuda-13.0/bin/nvcc \
    -DGMX_PREFER_STATIC_LIBS=ON

Compile GROMACS across all available CPU cores:
```
make -j$(nproc)
```
Install the GROMACS binaries in /home/software/gromacs:
```
make install
```
Add the GROMACS bin directory to your PATH environment variable:
```
export PATH=/home/software/gromacs/bin:$PATH
```

Create a benchmark directory and copy your ~/hosts file into it:

mkdir -p ~/gromacs_bench
cp ~/hosts ~/gromacs_bench/hosts

Download and extract the GROMACS water benchmark dataset in ~/gromacs_bench:

cd ~/gromacs_bench
wget ftp://ftp.gromacs.org/pub/benchmarks/water_GMX50_bare.tar.gz
tar -xzf water_GMX50_bare.tar.gz

Generate the bench.tpr run input file for the 1536 water benchmark system:

cd ~/gromacs_bench/water-cut1.0_GMX50_bare/1536
gmx_mpi grompp -f pme.mdp -c conf.gro -p topol.top -o bench.tpr

Configure the shared library path, GPU-aware MPI communication variables, and OpenMP thread count:

export LD_LIBRARY_PATH=/home/software/openmpi/lib:/usr/local/cuda-13.0/lib64:$LD_LIBRARY_PATH
export GMX_ENABLE_DIRECT_GPU_COMM=1
export GMX_FORCE_GPU_AWARE_MPI=1
export GMX_GPU_PME_DECOMPOSITION=1
export OMP_NUM_THREADS=8

Launch the multi-node GROMACS benchmark across 16 GPUs (8 GPUs per node on two nodes):

mpirun \
    -np 16 \
    --map-by ppr:8:node \
    --hostfile ~/gromacs_bench/hosts \
    --mca btl_tcp_if_exclude lo,docker0 \
    -x PATH \
    -x LD_LIBRARY_PATH \
    -x GMX_ENABLE_DIRECT_GPU_COMM \
    -x GMX_FORCE_GPU_AWARE_MPI \
    -x GMX_GPU_PME_DECOMPOSITION \
    -x OMP_NUM_THREADS \
    gmx_mpi mdrun \
        -s bench.tpr \
        -deffnm water_1536_run \
        -nb gpu \
        -bonded cpu \
        -pme gpu \
        -pmefft gpu \
        -update gpu \
        -npme 1 \
        -nsteps 2000 \
        -v

Run workloads on TPU MIGs

On a managed compute cluster with TPU instances, you can verify TPU chip availability across your MIG with JAX and run multi-host distributed workloads such as MaxText.

To verify your TPUs and run a distributed workload on a TPU MIG, complete the following steps:

Verify TPU devices across a MIG with JAX
Run multi-host MaxText workloads on TPU v6e or TPU7x

Verify TPU devices across a MIG with JAX

To install jax[tpu] in a Python virtual environment and verify that JAX detects the TPU chips on every compute instance in your MIG, complete the following steps:

Define the command payload (MY_CMD) that creates a local Python virtual environment in /tmp/my_venv, installs jax[tpu], and prints the detected TPU devices:

read -r -d '' MY_CMD << 'EOF' || true
set -e
python3 -m venv /tmp/my_venv
source /tmp/my_venv/bin/activate
pip install --upgrade pip
pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
python3 -c 'import jax; print("TPU Devices:", jax.devices())'
EOF

Run the verification command across all compute instances in your TPU MIG by using the gcloud alpha cluster-director clusters ssh command:
```
gcloud alpha cluster-director clusters ssh CLUSTER_NAME \
    --location=REGION \
    --managed-instance-groups=MIG_NAME \
    --batch-size=8 \
    --command="$MY_CMD"
```
Replace the following:
- CLUSTER_NAME: the name of your managed compute cluster.
- REGION: the region where your cluster exists.
- MIG_NAME: the name of your TPU MIG.

Run multi-host MaxText workloads on TPU v6e or TPU7x

To run a distributed MaxText workload across a TPU7x or TPU v6e MIG that mounts a shared /home directory, complete the following steps:

Clone the MaxText repository into the shared /home directory:

git clone https://github.com/google/maxtext.git ~/maxtext

Install the uv Python package manager:

curl -LsSf https://astral.sh/uv/install.sh | sh

Add the uv Python package manager to your PATH environment variable:
```
export PATH="$HOME/.local/bin:$PATH"
```
Create a shared Python 3.12 virtual environment in ~/maxtext_venv:
```
uv venv --python 3.12 ~/maxtext_venv
```
Activate the ~/maxtext_venv Python virtual environment:
```
source ~/maxtext_venv/bin/activate
```

Install the MaxText TPU dependencies:

cd ~/maxtext
uv pip install --resolution=lowest -r src/dependencies/requirements/generated_requirements/tpu-requirements.txt

Install MaxText in editable mode:
```
uv pip install --no-deps -e .
```
Set the JAX coordinator IP address (the internal IP address of host 0), coordinator port, and host list from your ~/hosts file:
```
COORDINATOR_IP=$(hostname -I | awk '{print $1}')
PORT=1234
mapfile -t NODES < ~/hosts
NUM_HOSTS=${#NODES[@]}
```

Launch the MaxText workload across all TPU hosts in the slice with the coordinator address (JAX_COORDINATOR_ADDRESS), total host count (JAX_NUM_HOSTS), and each worker's zero-based JAX_HOST_ID:

for host_id in "${!NODES[@]}"; do
  node="${NODES[$host_id]}"
  ssh -n "$node" "
    source ~/maxtext_venv/bin/activate
    export JAX_PLATFORMS=tpu,cpu
    export JAX_COORDINATOR_ADDRESS=${COORDINATOR_IP}:${PORT}
    export JAX_NUM_HOSTS=${NUM_HOSTS}
    export JAX_HOST_ID=${host_id}
    cd ~/maxtext
    python3 -m maxtext.trainers.pre_train.train \
      src/maxtext/configs/base.yml \
      model_name=llama3.1-70b \
      run_name=maxtext_mig_run \
      base_output_directory=${HOME}/maxtext_output \
      dataset_type=synthetic \
      steps=20 \
      per_device_batch_size=2 \
      max_target_length=8192 \
      ici_fsdp_parallelism=-1 \
      remat_policy=custom \
      attention=flash \
      scan_layers=true \
      weight_dtype=bfloat16 \
      dtype=bfloat16 \
      enable_checkpointing=false
  " > "/tmp/worker_${host_id}.log" 2>&1 &
done

Stream the log file from worker 0 to monitor the workload:
```
tail -f /tmp/worker_0.log
```
As each step finishes, /tmp/worker_0.log outputs the step time, loss, TFLOP/s per device, and tokens per second per device.