Tutorial: OSS Kubernetes on GCE with NVIDIA B200 GPUs, DRA for GPU, DRANET (OSS) and Gemma 4 (31B)

Google Cloud AI Infrastructure
, , , ,
1

gpu-oss-1
gpu-oss-11920Γ—1072 298 KB

Overview

This lab introduces you to building a high-performance, self-managed AI infrastructure directly on Google Compute Engine (GCE) A4 VMs using NVIDIA Blackwell B200 GPUs.

This example uses a reservation to get access to GPU capacity. Learn more about choosing a reservation.

You will bootstrap an unmanaged Kubernetes cluster on virtual machines (utilizing B200 GPUs) using Terraform and kubeadm, and configure Kubernetes Dynamic Resource Allocation (DRA) using open-source drivers. You will be working with the following:

Network and Storage Architecture

To achieve optimal throughput, you will build a topology-aware network layout consisting of 3 independent VPC networks:

  1. gVNIC Control Network 0 (oss-gpu-net-0): MTU 8896. Handles primary host-to-host control traffic, SSH, and the Kubernetes control plane.

  2. gVNIC Control Network 1 (oss-gpu-net-1): MTU 8896. Manages auxiliary inter-node orchestration traffic.

  3. RDMA Data Network (oss-gpu-mrdma): MTU 8896. Shares 8 subnets (10.2.0.0/16 through 10.9.0.0/16) using the Google-managed high-performance RoCE profile (XXXX-vpc-roce) to enable direct GPUDirect RDMA over 8 dedicated physical ConnectX-7 MRDMA network interfaces.

Setup Environment with Terraform

This example, will deploy the core cloud infrastructure in the region and zone using a capacity reservation.

Open Cloud Shell and ensure you see your target Project ID. Create your workspace directory and export your regional variables (replace the placeholders with your actual reservation details):

mkdir -p oss-kube-gpu-dra && cd oss-kube-gpu-dra

export PROJECT_ID=$(gcloud config get-value project)
export REGION="us-west3"
export ZONE="us-west3-c"
export RESERVATION_NAME="<YOUR_RESERVATION_NAME>" # Replace with your actual reservation name
export RESERVATION_BLOCK_NAME="<YOUR_RESERVATION_BLOCK_NAME>" # Replace with your actual reservation block name


echo "Project: $PROJECT_ID"
echo "Region: $REGION"
echo "Zone: $ZONE"
echo "Reservation: $RESERVATION_NAME"
echo "Reservation Block: $RESERVATION_BLOCK_NAME"

Create the terraform.tfvars:

cat << EOF > terraform.tfvars
project_id             = "${PROJECT_ID}"
region                 = "${REGION}"
zone                   = "${ZONE}"
reservation_name       = "${RESERVATION_NAME}"
reservation_block_name = "${RESERVATION_BLOCK_NAME}"
EOF

Create the variables.tf file to define resource variables and default machine profiles:

Terraform

cat << 'EOF' > variables.tf
variable "project_id" {
  type        = string
  description = "The Google Cloud Project ID"
}

variable "region" {
  type        = string
  description = "The region to deploy the resources"
}

variable "zone" {
  type        = string
  description = "The specific zone for the VMs"
}

variable "reservation_name" {
  type        = string
  description = "The name of the GCE reservation to consume"
}

variable "reservation_block_name" {
  type        = string
  description = "The name of the specific block within the GCE reservation to consume"
}

variable "control_plane_machine_type" {
  type        = string
  default     = "e2-standard-8"
  description = "Machine type for the Kubernetes control plane node (non-GPU)"
}

variable "gpu_worker_machine_type" {
  type        = string
  default     = "a4-highgpu-8g"
  description = "The machine type for the GPU workers (8x B200 GPUs)"
}
EOF

Create the vpc.tf file to establish the 3-VPC network topology with correct MTUs, subnets, and RoCE parameters:

Terraform

cat << 'EOF' > vpc.tf
terraform {
  required_version = ">= 1.5.0"
  required_providers {
    google = {
      source  = "hashicorp/google"
      version = "~> 7.32.0"
    }
  }
}

provider "google" {
  project = var.project_id
  region  = var.region
}

# 1. Two Regular VPC Networks for gVNICs
resource "google_compute_network" "gvnic_vpcs" {
  count                   = 2
  name                    = "oss-gpu-net-${count.index}"
  auto_create_subnetworks = false
  mtu                     = 8896
}

resource "google_compute_subnetwork" "gvnic_subnets" {
  count         = 2
  name          = "oss-gpu-sub-${count.index}"
  ip_cidr_range = "10.${count.index}.0.0/16" # Ranges: 10.0.0.0/16 and 10.1.0.0/16
  region        = var.region
  network       = google_compute_network.gvnic_vpcs[count.index].id
}

# 2. One RoCE VPC Network for RDMA NICs (GPU-to-GPU)
resource "google_compute_network" "roce_vpc" {
  name                    = "oss-gpu-mrdma"
  auto_create_subnetworks = false
  mtu                     = 8896
  network_profile         = "projects/${var.project_id}/global/networkProfiles/${var.zone}-vpc-roce"
}

# 8 Subnets inside the single RoCE VPC Network
resource "google_compute_subnetwork" "roce_subnets" {
  count         = 8
  name          = "oss-gpu-mrdma-sub-${count.index}"
  ip_cidr_range = "10.${count.index + 2}.0.0/16" # Ranges from 10.2.0.0/16 to 10.9.0.0/16
  region        = var.region
  network       = google_compute_network.roce_vpc.id
}

# 3. Outbound NAT Gateway for Primary Management
resource "google_compute_router" "router" {
  name    = "oss-gpu-router"
  network = google_compute_network.gvnic_vpcs[0].id
  region  = var.region
}

resource "google_compute_router_nat" "nat" {
  name                               = "oss-gpu-nat"
  router                             = google_compute_router.router.name
  region                             = var.region
  nat_ip_allocate_option             = "AUTO_ONLY"
  source_subnetwork_ip_ranges_to_nat = "ALL_SUBNETWORKS_ALL_IP_RANGES"
}

# 4. Internal Firewall Rules (ALLOW ALL within 10.0.0.0/8 range)
resource "google_compute_firewall" "gvnic_internal" {
  count   = 2
  name    = "oss-gpu-internal-${count.index}"
  network = google_compute_network.gvnic_vpcs[count.index].id
  allow {
    protocol = "tcp"
  }
  allow {
    protocol = "udp"
  }
  allow {
    protocol = "icmp"
  }
  source_ranges = ["10.0.0.0/8"]
}

resource "google_compute_firewall" "roce_internal" {
  name    = "oss-gpu-mrdma-internal"
  network = google_compute_network.roce_vpc.id
  allow {
    protocol = "tcp"
  }
  allow {
    protocol = "udp"
  }
  allow {
    protocol = "icmp"
  }
  source_ranges = ["10.0.0.0/8"]
}

# 5. External Access Firewalls (SSH & Ping on Net-0)
resource "google_compute_firewall" "gvnic_ssh" {
  name    = "oss-gpu-ssh"
  network = google_compute_network.gvnic_vpcs[0].id
  allow {
    protocol = "tcp"
    ports    = ["22"]
  }
  source_ranges = ["35.235.240.0/20"] # IAP range
}

resource "google_compute_firewall" "gvnic_ping_0" {
  name    = "oss-gpu-allow-ping-net-0"
  network = google_compute_network.gvnic_vpcs[0].id
  allow {
    protocol = "icmp"
  }
  source_ranges = ["0.0.0.0/0"] # Allows ping testing
}
EOF

Create the nodes.tf file to define your VMs. The node hostnames cleanly follow the k8s-gpu-control-plane pattern, and the reservation is targeted dynamically using your fully qualified resource URI:

Terraform

cat << 'EOF' > nodes.tf
# 1. K8s Control Plane VM (Standard Non-GPU Node)
resource "google_compute_instance" "control_plane" {
  name         = "k8s-gpu-control-plane"
  machine_type = var.control_plane_machine_type
  zone         = var.zone

  boot_disk {
    initialize_params {
      image = "projects/ubuntu-os-cloud/global/images/family/ubuntu-2204-lts"
      size  = 100
      type  = "pd-balanced"
    }
  }

  network_interface {
    network    = google_compute_network.gvnic_vpcs[0].id
    subnetwork = google_compute_subnetwork.gvnic_subnets[0].id
  }

  service_account {
    scopes = ["cloud-platform"]
  }
}

# 2. GPU Worker VMs (10 NICs total: 2x gVNIC, 8x MRDMA)
resource "google_compute_instance" "gpu_workers" {
  count        = 2
  name         = "k8s-gpu-worker-${count.index + 1}"
  machine_type = var.gpu_worker_machine_type
  zone         = var.zone

  boot_disk {
    initialize_params {
      image = "projects/ubuntu-os-cloud/global/images/family/ubuntu-2204-lts"
      size  = 200
      type  = "hyperdisk-balanced" # Mandatory for A4 VM platforms
    }
  }

  scheduling {
    on_host_maintenance = "TERMINATE"
    provisioning_model  = "RESERVATION_BOUND" # Consumes reservation
  }

  reservation_affinity {
    type = "SPECIFIC_RESERVATION"

    specific_reservation {
      key    = "compute.googleapis.com/reservation-name"
      values = ["projects/${var.project_id}/reservations/${var.reservation_name}/reservationBlocks/${var.reservation_block_name}"]
    }
  }

  # NIC 1: Primary Management (gVNIC) - net-0
  network_interface {
    network    = google_compute_network.gvnic_vpcs[0].id
    subnetwork = google_compute_subnetwork.gvnic_subnets[0].id
    nic_type   = "GVNIC"
  }

  # NIC 2: Secondary Control (gVNIC) - net-1
  network_interface {
    network    = google_compute_network.gvnic_vpcs[1].id
    subnetwork = google_compute_subnetwork.gvnic_subnets[1].id
    nic_type   = "GVNIC"
  }

  # NICs 3-10: ConnectX-7 MRDMA RoCE interfaces - roce-vpc (oss-gpu-mrdma)
  dynamic "network_interface" {
    for_each = range(8)
    content {
      network    = google_compute_network.roce_vpc.id
      subnetwork = google_compute_subnetwork.roce_subnets[network_interface.value].id
      nic_type   = "MRDMA"
    }
  }

  service_account {
    scopes = ["cloud-platform"]
  }

  lifecycle {
    ignore_changes = [
      boot_disk[0].initialize_params[0].image,
      guest_accelerator,
      metadata,
      scratch_disk,          
      reservation_affinity   
    ]
  }
}
EOF

Initialize, validate, and apply your configuration to deploy your cloud environment:

terraform init
terraform plan -out=tfplan
terraform apply tfplan

(This process may take between 5 - 10 minutes. Please confirm the successful creation of all subnets and instances before continuing).

Validate the Setup

echo -e "\n=== Verifying VPC Networks ==="
gcloud compute networks list --filter="name~oss-gpu-.*" --project=$PROJECT_ID

echo -e "\n=== Verifying Subnetworks ==="
gcloud compute networks subnets list --filter="name~oss-gpu-.*" --project=$PROJECT_ID

echo -e "\n=== Verifying Firewall Rules ==="
gcloud compute firewall-rules list --filter="name~oss-gpu-.*" --project=$PROJECT_ID

echo -e "\n=== Verifying Cloud NAT Gateway ==="
gcloud compute routers nats list --router=oss-gpu-router --router-region=$REGION --project=$PROJECT_ID

echo -e "\n=== Verifying Provisioned VM Instances ==="
gcloud compute instances list --filter="name~k8s-gpu-.*" --project=$PROJECT_ID

echo -e "\n=== Verifying Multi-NIC Layout (10 Interfaces: 2x gVNIC, 8x MRDMA) on GPU Workers ==="
for i in 1 2; do
  echo -e "\n--- Network Interfaces for k8s-gpu-worker-${i} ---"
  gcloud compute instances describe k8s-gpu-worker-${i} \
      --zone=$ZONE \
      --project=$PROJECT_ID \
      --format="table(networkInterfaces[].network.basename(), networkInterfaces[].nicType, networkInterfaces[].networkIP)"
done

echo -e "\n=== Verifying Capacity Block Reservation Affinity on GPU Workers ==="
gcloud compute instances describe k8s-gpu-worker-1 \
    --zone=$ZONE \
    --project=$PROJECT_ID \
    --format="yaml(reservationAffinity)"

Bootstrap Your Kubernetes Cluster Control Node

In this section, you will connect securely to the non-GPU control plane VM (k8s-gpu-control-plane), prepare host kernel modules, configure containerd, and bootstrap the unmanaged Kubernetes control plane services with strict traffic routing via Calico CNI.

SSH into the k8s-gpu-control-plane instance using Identity-Aware Proxy (IAP):

gcloud compute ssh k8s-gpu-control-plane \
    --zone=$ZONE \
    --tunnel-through-iap

Inside the control plane VM, write the init-control-plane.sh script:

cat << 'CONTROL_PLANE_EOF' > init-control-plane.sh
#!/bin/bash
set -e

echo "=== 1. Neutralizing Background Updates & Preparing Base OS ==="
sudo systemctl stop apt-daily.timer apt-daily-upgrade.timer || true
sudo systemctl disable apt-daily.timer apt-daily-upgrade.timer || true
sudo systemctl mask apt-daily.service apt-daily-upgrade.service || true

# Turn off swap
sudo swapoff -a
sudo sed -i '/ swap / s/^\(.*\)$/\#\1/g' /etc/fstab

# Load required kernel modules
cat << 'EOT' | sudo tee /etc/modules-load.d/k8s.conf
overlay
br_netfilter
EOT
sudo modprobe overlay
sudo modprobe br_netfilter

# Configure sysctl requirements for Kubernetes bridging
cat << 'EOT' | sudo tee /etc/sysctl.d/k8s.conf
net.bridge.bridge-nf-call-iptables  = 1
net.bridge.bridge-nf-call-ip6tables = 1
net.ipv4.ip_forward                 = 1
EOT
sudo sysctl --system

echo "=== 2. Installing Container Runtime (Containerd) ==="
sudo apt-get update
sudo apt-get install -y ca-certificates curl gnupg bash-completion
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor --yes -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo $VERSION_CODENAME) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt-get update
# FIX: Removed docker-ce and docker-ce-cli to keep the footprint minimal
sudo apt-get install -y containerd.io

echo "=== 3. Configuring Containerd with Systemd Cgroups ==="
sudo mkdir -p /etc/containerd
containerd config default | sudo tee /etc/containerd/config.toml >/dev/null
sudo sed -i 's/SystemdCgroup = false/SystemdCgroup = true/' /etc/containerd/config.toml
sudo systemctl daemon-reload
sudo systemctl restart containerd
sudo systemctl enable containerd

# Validation Step
if ! systemctl is-active --quiet containerd; then
    echo "❌ ERROR: Containerd failed to start."
    exit 1
fi
echo "βœ… Containerd is active and healthy."

echo "=== 4. Installing Kubernetes 1.37 Binaries (Auto-Detect Channel) ==="
K8S_VERSION="v1.37"

# STRICT FIX: Added -L to follow redirects so curl can properly detect the 403 Forbidden at the destination
if curl -fsSL -o /dev/null "https://pkgs.k8s.io/core:/stable:/${K8S_VERSION}/deb/Release"; then
    K8S_REPO="core:/stable:/${K8S_VERSION}"
    echo "πŸš€ Stable release found! Using channel: ${K8S_REPO}"
else
    K8S_REPO="core:/prerelease:/${K8S_VERSION}"
    echo "⚠️ Stable not found yet. Falling back to prerelease channel: ${K8S_REPO}"
fi

curl -fsSL "https://pkgs.k8s.io/${K8S_REPO}/deb/Release.key" | sudo gpg --dearmor --yes -o /etc/apt/keyrings/kubernetes-apt-keyring.gpg
echo "deb [signed-by=/etc/apt/keyrings/kubernetes-apt-keyring.gpg] https://pkgs.k8s.io/${K8S_REPO}/deb/ /" | sudo tee /etc/apt/sources.list.d/kubernetes.list

sudo apt-get update
sudo apt-get install -y kubelet kubeadm kubectl
sudo apt-mark hold kubelet kubeadm kubectl

# Autocomplete & Aliases
kubectl completion bash | sudo tee /etc/bash_completion.d/kubectl > /dev/null
kubeadm completion bash | sudo tee /etc/bash_completion.d/kubeadm > /dev/null
if ! grep -q 'alias k=kubectl' ~/.bashrc; then
  echo 'alias k=kubectl' >> ~/.bashrc
  echo 'complete -o default -F __start_kubectl k' >> ~/.bashrc
fi

echo "=== 5. Initializing Control Plane Engine ==="
# DYNAMIC VERSION FIX: Force kubeadm to use the exact version we just installed (e.g. 1.37.0-rc.0)
INSTALLED_K8S_VERSION=$(kubeadm version -o short)
echo "Bootstrapping cluster using explicit version: ${INSTALLED_K8S_VERSION}"
sudo kubeadm init --pod-network-cidr=192.168.0.0/16 --kubernetes-version="${INSTALLED_K8S_VERSION}"

echo "=== 6. Configuring Administrative Cluster Credentials ==="
mkdir -p $HOME/.kube
sudo cp -i /etc/kubernetes/admin.conf $HOME/.kube/config
sudo chown $(id -u):$(id -g) $HOME/.kube/config

echo "Waiting for local API server context..."
until kubectl cluster-info &>/dev/null; do
    sleep 2
done
echo "βœ… Kubernetes API server is responding locally."

echo "=== 7. Deploying Calico Network Operator ==="
kubectl create -f https://raw.githubusercontent.com/projectcalico/calico/v3.27.3/manifests/tigera-operator.yaml
echo "Waiting for Tigera Installation CRD to register on the API server..."
kubectl wait --for=condition=established crd/installations.operator.tigera.io --timeout=60s

echo "=== 8. Deploying Calico Custom Resources ==="
cat << 'CALICO_EOF' > custom-calico.yaml
apiVersion: operator.tigera.io/v1
kind: Installation
metadata:
  name: default
spec:
  calicoNetwork:
    nodeAddressAutodetectionV4:
      cidrs:
        - "10.0.0.0/16"
    ipPools:
    - blockSize: 26
      cidr: 192.168.0.0/16
      encapsulation: VXLANCrossSubnet
      natOutgoing: Enabled
      nodeSelector: all()
CALICO_EOF
kubectl apply -f custom-calico.yaml

echo "Waiting 10 seconds for Calico system namespaces to initialize..."
sleep 10
kubectl get pods -n calico-system

echo "=== 9. Exporting Worker Cluster Join Token ==="
sudo kubeadm token create --print-join-command > ~/join.sh
chmod +x ~/join.sh
echo "--------------------------------------------------------"
echo "βœ… CONTROL PLANE BOOTSTRAP COMPLETE!"
echo "Your cluster join command is saved below:"
echo "--------------------------------------------------------"
cat ~/join.sh
CONTROL_PLANE_EOF

Run the bootstrap script:

chmod +x init-control-plane.sh
./init-control-plane.sh

Confirm that the node is running in a Ready status and exit the SSH session:

kubectl get nodes
exit

Output:

NAME                    STATUS   ROLES           AGE     VERSION
k8s-gpu-control-plane   Ready    control-plane   2m24s   v1.37.0-rc.0

Add the GPU Worker Nodes

You will deploy an orchestration script from Cloud Shell that connects securely to the control plane VM (k8s-gpu-control-plane), extracts the registration credentials, and concurrently bootstraps both your B200 GPU workers to join them to the cluster.

In your Cloud Shell terminal, write the bootstrap-workers.sh script:

cat << 'WORKER_BOOTSTRAP_EOF' > bootstrap-workers.sh
#!/bin/bash
set -e

# 1. Fetch the join command safely from the control plane
echo "Fetching join command from Control Plane..."
JOIN_CMD=$(gcloud compute ssh k8s-gpu-control-plane --zone=$ZONE --tunnel-through-iap --command="cat ~/join.sh" 2>/dev/null)
if [ -z "$JOIN_CMD" ]; then
    echo "❌ ERROR: Failed to retrieve the join command. Ensure the control plane is reachable."
    exit 1
fi
echo "βœ… Successfully retrieved join command."

# 2. Generate the init-worker script to inject onto the hosts
cat << 'WORKER_INIT_EOF' > init-worker.sh
#!/bin/bash
set -e

echo "=== 1. Neutralizing Background Updates ==="
export DEBIAN_FRONTEND=noninteractive
sudo sed -i "s/\#\$nrconf{restart} = 'i';/\$nrconf{restart} = 'a';/g" /etc/needrestart/needrestart.conf 2>/dev/null || true
sudo systemctl stop apt-daily.timer apt-daily-upgrade.timer || true
sudo systemctl disable apt-daily.timer apt-daily-upgrade.timer || true
sudo systemctl mask apt-daily.service apt-daily-upgrade.service || true

echo "=== 2. Base OS Prep & Hardware Kernel Tuning ==="
sudo swapoff -a
sudo sed -i '/ swap / s/^\(.*\)$/\#\1/g' /etc/fstab

# Load required modules (including the vital GPUDirect nvidia_peermem module)
cat << 'EOT' | sudo tee /etc/modules-load.d/k8s.conf
overlay
br_netfilter
nvidia_peermem
EOT
sudo modprobe overlay
sudo modprobe br_netfilter
sudo modprobe nvidia_peermem || true

# Configure bridging and IP forwarding sysctls
cat << 'EOT' | sudo tee /etc/sysctl.d/k8s.conf
net.bridge.bridge-nf-call-iptables  = 1
net.bridge.bridge-nf-call-ip6tables = 1
net.ipv4.ip_forward                 = 1
EOT
sudo sysctl --system

# Tune all physical Mellanox ConnectX-7 interfaces to MTU 8896 (bypasses fragmentation)
echo "  πŸ“ Tuning host-level Mellanox network links..."
for dev in $(ls /sys/class/net/); do
  if [ -d /sys/class/net/$dev/device/driver ] && [ "$(basename $(readlink /sys/class/net/$dev/device/driver))" = "mlx5_core" ]; then
    sudo ip link set dev $dev mtu 8896
  fi
done

echo "=== 3. Installing & Tuning Containerd with CDI Support ==="
sudo apt-get update && sudo apt-get install -yq ca-certificates curl gnupg bash-completion
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor --yes -o /etc/apt/keyrings/docker.gpg
sudo chmod a+r /etc/apt/keyrings/docker.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo $VERSION_CODENAME) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

sudo apt-get update && sudo apt-get install -yq containerd.io
sudo mkdir -p /etc/containerd
containerd config default | sudo tee /etc/containerd/config.toml >/dev/null

# Set Systemd Cgroups
sudo sed -i 's/SystemdCgroup = false/SystemdCgroup = true/' /etc/containerd/config.toml

# Enable Container Device Interface (CDI) support so the DRA driver can inject GPUs
sudo sed -i "/\[plugins.\"io.containerd.grpc.v1.cri\"\]/a \  enable_cdi = true\n  cdi_spec_dirs = [\"/etc/cdi\", \"/var/run/cdi\"]" /etc/containerd/config.toml

sudo systemctl daemon-reload
sudo systemctl restart containerd
sudo systemctl enable containerd

# Validation: Check containerd status
if ! systemctl is-active --quiet containerd; then
    echo "❌ ERROR: Containerd failed to start."
    exit 1
fi

echo "=== 4. Installing Kubernetes 1.37 Binaries (Auto-Detect Channel) ==="
K8S_VERSION="v1.37"

if curl -s --head -f "https://pkgs.k8s.io/core:/stable:/${K8S_VERSION}/deb/Release" >/dev/null 2>&1; then
    K8S_REPO="core:/stable:/${K8S_VERSION}"
    echo "πŸš€ Stable release found! Using channel: ${K8S_REPO}"
else
    K8S_REPO="core:/prerelease:/${K8S_VERSION}"
    echo "⚠️ Stable not found yet. Falling back to prerelease channel: ${K8S_REPO}"
fi

curl -fsSL "https://pkgs.k8s.io/${K8S_REPO}/deb/Release.key" | sudo gpg --dearmor --yes -o /etc/apt/keyrings/kubernetes-apt-keyring.gpg
echo "deb [signed-by=/etc/apt/keyrings/kubernetes-apt-keyring.gpg] https://pkgs.k8s.io/${K8S_REPO}/deb/ /" | sudo tee /etc/apt/sources.list.d/kubernetes.list

sudo apt-get update && sudo apt-get install -yq kubelet kubeadm kubectl
sudo apt-mark hold kubelet kubeadm kubectl
WORKER_INIT_EOF

# 3. Append the cluster join command to the initialization script
echo "echo \"=== 5. Joining Cluster ===\"" >> init-worker.sh
echo "sudo $JOIN_CMD" >> init-worker.sh

# 4. Push and run on both Workers concurrently
echo "Starting concurrent, optimized bootstrap on both workers..."
(
    echo "[Worker 1] Copying script..."
    gcloud compute scp init-worker.sh k8s-gpu-worker-1:~ --zone=$ZONE --tunnel-through-iap --quiet
    echo "[Worker 1] Executing script..."
    gcloud compute ssh k8s-gpu-worker-1 --zone=$ZONE --tunnel-through-iap --command="bash ~/init-worker.sh"
    echo "βœ… [Worker 1] Bootstrap and Join complete!"
) &
(
    echo "[Worker 2] Copying script..."
    gcloud compute scp init-worker.sh k8s-gpu-worker-2:~ --zone=$ZONE --tunnel-through-iap --quiet
    echo "[Worker 2] Executing script..."
    gcloud compute ssh k8s-gpu-worker-2 --zone=$ZONE --tunnel-through-iap --command="bash ~/init-worker.sh"
    echo "βœ… [Worker 2] Bootstrap and Join complete!"
) &

wait
echo "--------------------------------------------------------"
echo "πŸŽ‰ BOTH B200 WORKERS BOOTSTRAPPED AT LINE-RATE SPEEDS!"
echo "--------------------------------------------------------"

# Final Validation Check from Control Plane
echo "Verifying cluster node status (waiting 5 seconds for nodes to register)..."
sleep 5
gcloud compute ssh k8s-gpu-control-plane --zone=$ZONE --tunnel-through-iap --command="kubectl get nodes -o wide"
WORKER_BOOTSTRAP_EOF

Execute the setup orchestration script:

chmod +x bootstrap-workers.sh
./bootstrap-workers.sh

Verification Check

Your terminal output should display all three nodes registered inside your Kubernetes cluster. Confirm that you see the following configuration before stopping:

NAME                    STATUS     ROLES           AGE     VERSION        INTERNAL-IP   EXTERNAL-IP   OS-IMAGE             KERNEL-VERSION           CONTAINER-RUNTIME
k8s-gpu-control-plane   Ready      control-plane   7m33s   v1.37.0-rc.0   10.0.0.2      <none>        Ubuntu 22.04.5 LTS   6.8.0-1064-gcp (amd64)   containerd://2.3.3
k8s-gpu-worker-1        NotReady   <none>          9s      v1.37.0-rc.0   10.0.0.3      <none>        Ubuntu 22.04.5 LTS   6.8.0-1064-gcp (amd64)   containerd://2.3.3
k8s-gpu-worker-2        NotReady   <none>          8s      v1.37.0-rc.0   10.0.0.4      <none>        Ubuntu 22.04.5 LTS   6.8.0-1064-gcp (amd64)   containerd://2.3.3

Install Helm on the Control Plane Node

Before deploying your GPU and network drivers, you must install the Kubernetes package manager Helm directly on the control plane node. Ensure you are connected to your secure SSH session on k8s-gpu-control-plane. If you have exited, reconnect first:

gcloud compute ssh k8s-gpu-control-plane \
    --zone=$ZONE \
    --tunnel-through-iap

Download and run the official Helm installation script:

curl -fsSL https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash

Confirm that Helm is installed and running:

helm version

Deploying OSS DRA NVIDIA GPU Driver

In this section, you will install the NVIDIA GPU Operator (disabling the standard device plugin so we can use DRA), label your GPU worker nodes with their specific accelerator details, and install the open-source NVIDIA GPU DRA driver using Helm. This driver is responsible for discovering the physical NVIDIA Blackwell B200 GPUs and mapping them natively to the Kubernetes API.

Add the NVIDIA Helm Repository and install the GPU Operator

Configure the repository and install the GPU Operator with the standard Kubernetes device plugin disabled (to let the DRA driver handle GPU allocation):

# Add official NVIDIA NGC Helm repository
helm repo add nvidia https://helm.ngc.nvidia.com/nvidia
helm repo update

# Install GPU Operator with standard device plugin disabled
helm upgrade --install gpu-operator nvidia/gpu-operator \
  --create-namespace \
  --namespace gpu-operator \
  --set devicePlugin.enabled=false \
  --set driver.manager.env[0].name=NODE_LABEL_FOR_GPU_POD_EVICTION \
  --set driver.manager.env[0].value="nvidia.com/dra-kubelet-plugin" \
  --wait

Label both GPU worker nodes

Label your workers to enable the DRA plugin, specify the Blackwell topology, and pre-apply the gpu.present scheduler label to prevent bootstrap deadlocks:

# Label Worker Node 1
kubectl label node k8s-gpu-worker-1 \
  nvidia.com/dra-kubelet-plugin=true \
  nvidia.com/gpu.present=true \
  cloud.google.com/gke-gpu-accelerator=b200 \
  cloud.google.com/gke-gpu-count=8 \
  --overwrite

# Label Worker Node 2
kubectl label node k8s-gpu-worker-2 \
  nvidia.com/dra-kubelet-plugin=true \
  nvidia.com/gpu.present=true \
  cloud.google.com/gke-gpu-accelerator=b200 \
  cloud.google.com/gke-gpu-count=8 \
  --overwrite

Clone and install the official NVIDIA GPU DRA driver

# Ensure the GPU Operator host driver compilation has completed
echo "=== Waiting for GPU Operator driver compilation to finish (takes ~3-5 mins) ==="
kubectl wait --for=condition=Ready pod -l app=nvidia-driver-daemonset -n gpu-operator --timeout=600s

# Deploy the stable driver directly from the official Kubernetes OCI Registry
helm upgrade --install dra-driver-nvidia-gpu oci://registry.k8s.io/dra-driver-nvidia/charts/dra-driver-nvidia-gpu \
  --version 0.4.1 \
  --create-namespace \
  --namespace dra-driver-nvidia-gpu \
  --set nvidiaDriverRoot=/run/nvidia/driver \
  --set 'kubeletPlugin.env[0].name=NODE_NAME' \
  --set 'kubeletPlugin.env[0].valueFrom.fieldRef.fieldPath=spec.nodeName' \
  --set-string 'kubeletPlugin.nodeSelector.nvidia\.com/dra-kubelet-plugin=true' \
  --set-string 'kubeletPlugin.nodeSelector.nvidia\.com/gpu\.present=true' \
  --set gpuResourcesEnabledOverride=true \
  --wait

Validate the GPU DRA driver setup

# Verify driver pod statuses (use the correct 'dra-driver-nvidia-gpu' namespace)
kubectl get pods -n dra-driver-nvidia-gpu -o wide

# Verify that GPU DeviceClasses are successfully registered to the cluster
kubectl get deviceclass

Deploying Open-Source DRANET & Device Classes

Run the following commands directly inside your SSH session on k8s-gpu-control-plane to configure the open-source DRANET network allocation layer and define your joint network-to-GPU claim templates:

Install the core components for DRANET

Apply the container runtime patch to filter out virtual interfaces:

# Install DRANET core components
kubectl apply -f https://raw.githubusercontent.com/kubernetes-sigs/dranet/refs/heads/main/install.yaml

# Patch DRANET DaemonSet to bypass non-physical networks (veth, vxlan, bridge)
kubectl patch daemonset dranet -n kube-system --type='json' -p='[
  {
    "op": "add",
    "path": "/spec/template/spec/containers/0/args/-",
    "value": "-filter=!( \"dra.net/type\" in attributes) || (attributes[\"dra.net/type\"].StringValue != \"veth\" && attributes[\"dra.net/type\"].StringValue != \"vxlan\" && attributes[\"dra.net/type\"].StringValue != \"bridge\")"
  }
]'

# Monitor deployment rollout progress
kubectl rollout status daemonset/dranet -n kube-system

Apply your topology-aligned DeviceClasses and joint ResourceClaimTemplates:

cat << 'EOF' | kubectl apply -f -
apiVersion: resource.k8s.io/v1
kind: DeviceClass
metadata:
  name: gpu-dranet
spec:
  selectors:
    - cel:
        expression: device.driver == "dra.net"
---
apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: gpu-net-interfaces
  namespace: default
spec:
  spec:
    devices:
      requests:
      - name: gpu-net-interface
        exactly:
          deviceClassName: gpu-dranet
          count: 8 # Mapped to our 8 physical MRDMA interfaces
          selectors:
          - cel:
              expression: device.attributes["gce.dra.net"].networkName.startsWith("oss-gpu-mrdma")
      config:
      - opaque:
          driver: dra.net
          parameters:
            interface:
              mtu: 8896
              gsoMaxSize: 65536
              groMaxSize: 65536
              gsoIPv4MaxSize: 65536
              groIPv4MaxSize: 65536
              disableEbpfPrograms: true
---
apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: gpu-device-template
  namespace: default
spec:
  spec:
    devices:
      requests:
      - name: gpu-devices
        exactly:
          deviceClassName: gpu.nvidia.com # Maps to our GPU DRA Driver
          allocationMode: ExactCount
          count: 8 # Demands all 8 B200 GPUs on a single host
EOF

Confirm that your resource classes and network layers are published correctly in the API server:

# Verify ResourceSlices exist and are actively serving both drivers
kubectl get resourceslices -o custom-columns=NAME:.metadata.name,NODE:.spec.nodeName,DRIVER:.spec.driver | grep -E "dra.net|gpu.nvidia.com"

# Verify the DRANET daemonset pods are Running across all nodes
kubectl get pods -n kube-system -l app=dranet -o wide

Output:

00000-gpu.nvidia.com-k8s-gpu-worker-1-dxsnl              k8s-gpu-worker-1        gpu.nvidia.com
00000-gpu.nvidia.com-k8s-gpu-worker-2-nhp5c              k8s-gpu-worker-2        gpu.nvidia.com
k8s-gpu-control-plane-dra.net-b25z2                      k8s-gpu-control-plane   dra.net
k8s-gpu-worker-1-dra.net-tn44x                           k8s-gpu-worker-1        dra.net
k8s-gpu-worker-2-dra.net-fvx8n                           k8s-gpu-worker-2        dra.net

NAME           READY   STATUS    RESTARTS   AGE    IP         NODE                    NOMINATED NODE   READINESS GATES
dranet-4zhzw   1/1     Running   0          117s   10.0.0.4   k8s-gpu-worker-2        <none>           <none>
dranet-8m4ng   1/1     Running   0          117s   10.0.0.3   k8s-gpu-worker-1        <none>           <none>
dranet-brqm8   1/1     Running   0          118s   10.0.0.2   k8s-gpu-control-plane   <none>           <none>

Verifying Drivers and Stress Testing the Cluster

Before we deploy our LLM, we need to guarantee that the Kubernetes Dynamic Resource Allocation (DRA) drivers are correctly exposing the GPUs and the RDMA network interfaces. Once verified, we will unleash the absolute maximum theoretical compute and network load on the cluster to prove its stability.

With the infrastructure previously verified, we will now push the Blackwell architecture to its limits. We are going to orchestrate a two-part stress test:

  1. Sustained 3.2 Tbps RDMA Network Test: Flooding the fabric with 1GB to 16GB payloads across 50 iterations to prove the GPUDirect line rate.

  2. Distributed BF16 Matrix Multiplication: Generating a Python script on the fly to engage the Blackwell Tensor Cores on all 16 GPUs concurrently.

Create the deployment script by pasting this entire block into your terminal:

cat << 'EOF' > native-test-nccl.sh
#!/bin/bash
set -e

echo "=== 1. Clean Up Previous State ==="
kubectl delete statefulset,job,svc,secret,pod -l app=nccl-benchmark --force --grace-period=0 2>/dev/null || true

echo "=== 2. Create SSH Key Secret ==="
rm -f nccl-ssh-key nccl-ssh-key.pub
ssh-keygen -t rsa -N "" -f nccl-ssh-key -q
kubectl create secret generic nccl-ssh-keys \
  --from-file=id_rsa=nccl-ssh-key \
  --from-file=id_rsa.pub=nccl-ssh-key.pub \
  --from-file=authorized_keys=nccl-ssh-key.pub
kubectl label secret nccl-ssh-keys app=nccl-benchmark

echo "=== 3. Generate Native Kubernetes Benchmark Manifest ==="
cat << 'YAML_EOF' > native-nccl.yaml
apiVersion: v1
kind: Service
metadata:
  name: nccl-workers
  labels:
    app: nccl-benchmark
spec:
  clusterIP: None
  selector:
    app: nccl-benchmark-worker
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: nccl-worker
  labels:
    app: nccl-benchmark
spec:
  serviceName: "nccl-workers"
  replicas: 2
  selector:
    matchLabels:
      app: nccl-benchmark-worker
  template:
    metadata:
      labels:
        app: nccl-benchmark-worker
    spec:
      hostIPC: true
      resourceClaims:
      - name: gpu-claim
        resourceClaimTemplateName: gpu-device-template
      - name: net-claim
        resourceClaimTemplateName: gpu-net-interfaces
      containers:
      - image: nvcr.io/nvidia/pytorch:25.01-py3
        name: worker
        imagePullPolicy: IfNotPresent
        securityContext:
          privileged: true
        command: ["/bin/bash", "-c"]
        args: 
        - |
          apt-get update -qq && apt-get install -y -qq openssh-server > /dev/null
          mkdir -p /var/run/sshd /root/.ssh
          cp /tmp/ssh/* /root/.ssh/
          chmod 600 /root/.ssh/id_rsa
          chmod 644 /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys
          
          # YAML-Safe limits generation
          echo "* soft memlock unlimited" > /etc/security/limits.conf
          echo "* hard memlock unlimited" >> /etc/security/limits.conf
          echo "root soft memlock unlimited" >> /etc/security/limits.conf
          echo "root hard memlock unlimited" >> /etc/security/limits.conf
          
          sed -i 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' /etc/pam.d/sshd
          echo "UsePAM yes" >> /etc/ssh/sshd_config
          
          echo "=== Compiling NVIDIA NCCL Tests on Worker ==="
          rm -rf /tmp/nccl-tests
          git clone https://github.com/NVIDIA/nccl-tests.git /tmp/nccl-tests
          cd /tmp/nccl-tests
          make MPI=1 MPI_HOME=/usr/local/mpi
          
          /usr/sbin/sshd -p 2222 -De
        resources:
          claims:
          - name: gpu-claim
          - name: net-claim
        volumeMounts:
        - mountPath: /dev/shm
          name: shm
        - mountPath: /tmp/ssh
          name: ssh-keys
      volumes:
      - name: shm
        emptyDir:
          medium: Memory
      - name: ssh-keys
        secret:
          secretName: nccl-ssh-keys
---
apiVersion: batch/v1
kind: Job
metadata:
  name: nccl-launcher
  labels:
    app: nccl-benchmark
spec:
  backoffLimit: 0
  template:
    metadata:
      labels:
        app: nccl-benchmark
    spec:
      restartPolicy: Never
      containers:
      - name: launcher
        image: nvcr.io/nvidia/pytorch:25.01-py3
        imagePullPolicy: IfNotPresent
        env:
        - name: OMPI_ALLOW_RUN_AS_ROOT
          value: "1"
        - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
          value: "1"
        command: ["/bin/bash", "-c"]
        args:
        - |
          apt-get update -qq && apt-get install -y -qq openssh-client > /dev/null
          mkdir -p /root/.ssh
          cp /tmp/ssh/* /root/.ssh/
          chmod 600 /root/.ssh/id_rsa
          
          echo "Host *" > /root/.ssh/config
          echo "    StrictHostKeyChecking no" >> /root/.ssh/config
          echo "    Port 2222" >> /root/.ssh/config
          echo "Host nccl-worker-0" >> /root/.ssh/config
          echo "    HostName nccl-worker-0.nccl-workers.default.svc.cluster.local" >> /root/.ssh/config
          echo "Host nccl-worker-1" >> /root/.ssh/config
          echo "    HostName nccl-worker-1.nccl-workers.default.svc.cluster.local" >> /root/.ssh/config
          
          echo "=== Compiling NVIDIA NCCL Tests on Launcher ==="
          rm -rf /tmp/nccl-tests
          git clone https://github.com/NVIDIA/nccl-tests.git /tmp/nccl-tests
          cd /tmp/nccl-tests
          make MPI=1 MPI_HOME=/usr/local/mpi
          
          # YAML-Safe generation of the Matrix Multiplication Script
          echo 'import torch, time, os' > /tmp/matmul.py
          echo 'rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", "0"))' >> /tmp/matmul.py
          echo 'local_rank = int(os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0"))' >> /tmp/matmul.py
          echo 'torch.cuda.set_device(local_rank)' >> /tmp/matmul.py
          echo 'size = 16384' >> /tmp/matmul.py
          echo "a = torch.randn(size, size, device='cuda', dtype=torch.bfloat16)" >> /tmp/matmul.py
          echo "b = torch.randn(size, size, device='cuda', dtype=torch.bfloat16)" >> /tmp/matmul.py
          echo 'for _ in range(20): torch.matmul(a, b)' >> /tmp/matmul.py
          echo 'torch.cuda.synchronize()' >> /tmp/matmul.py
          echo 'start = time.time()' >> /tmp/matmul.py
          echo 'iters = 1000' >> /tmp/matmul.py
          echo 'for _ in range(iters): torch.matmul(a, b)' >> /tmp/matmul.py
          echo 'torch.cuda.synchronize()' >> /tmp/matmul.py
          echo 'elapsed = time.time() - start' >> /tmp/matmul.py
          echo 'tflops = (2.0 * (size**3) * iters) / (elapsed * 1e12)' >> /tmp/matmul.py
          echo 'print(f"[Rank {rank:02d} | GPU {local_rank}] Matrix {size}x{size} BF16 GEMM: {tflops:.2f} TFLOPS")' >> /tmp/matmul.py
          
          echo "=== Waiting for GPU Workers to come online ==="
          echo "nccl-worker-0.nccl-workers.default.svc.cluster.local slots=8" > /tmp/hostfile
          echo "nccl-worker-1.nccl-workers.default.svc.cluster.local slots=8" >> /tmp/hostfile
          
          for host in nccl-worker-0.nccl-workers.default.svc.cluster.local nccl-worker-1.nccl-workers.default.svc.cluster.local; do
            while ! ssh -q -o ConnectTimeout=5 $host exit 2>/dev/null; do
              sleep 1
            done
            # Synchronize NCCL binaries and Python code to the remote workers
            scp -r /tmp/nccl-tests $host:/tmp/nccl-tests >/dev/null 2>&1
            scp /tmp/matmul.py $host:/tmp/matmul.py >/dev/null 2>&1
          done
          echo "βœ… All workers online and synchronized!"
          
          echo "========================================================"
          echo "=== PART 1: Intense 3.2 Tbps B200 RDMA Test ==="
          echo "========================================================"
          START_TIME=$SECONDS
          
          mpirun -np 16 --hostfile /tmp/hostfile \
            --bind-to none --map-by slot \
            -x LD_LIBRARY_PATH \
            -x NCCL_DEBUG=WARN \
            -x NCCL_NVLS_ENABLE=0 \
            -x NCCL_ALGO=Ring \
            -x NCCL_MIN_NCHANNELS=32 \
            -x NCCL_IB_PCI_RELAXED_ORDERING=1 \
            -x NCCL_IB_GDR_LEVEL=MAX \
            -x NCCL_IB_HCA=mlx5 \
            -x NCCL_SOCKET_IFNAME=eth0 \
            -x NCCL_IB_DISABLE=0 \
            bash -c 'ulimit -l unlimited && /tmp/nccl-tests/build/all_reduce_perf -b 1G -e 16G -f 2 -g 1 -w 10 -n 50'
            
          echo "βœ… Network Benchmark completed!"
          echo "========================================================"
          echo "=== PART 2: Distributed BF16 Matrix Multiplication ==="
          echo "========================================================"
          
          mpirun -np 16 --hostfile /tmp/hostfile \
            --bind-to none --map-by slot \
            -x LD_LIBRARY_PATH \
            python3 /tmp/matmul.py
            
          ELAPSED=$(( SECONDS - START_TIME ))
          echo "========================================================"
          echo "βœ… All tests completed successfully in $ELAPSED seconds!"
          echo "========================================================"
        volumeMounts:
        - mountPath: /tmp/ssh
          name: ssh-keys
      volumes:
      - name: ssh-keys
        secret:
          secretName: nccl-ssh-keys
YAML_EOF

echo "=== 4. Deploy Native StatefulSet & Launcher Job ==="
kubectl apply -f native-nccl.yaml

echo "=== 5. Wait for Launcher Pod to Provision ==="
until kubectl get pods -l app=nccl-benchmark | grep -q nccl-launcher; do
    sleep 2
done

LAUNCHER_POD=$(kubectl get pods -l app=nccl-benchmark -o jsonpath="{.items[?(@.metadata.labels.job-name=='nccl-launcher')].metadata.name}")
echo "βœ… Launcher pod scheduled: $LAUNCHER_POD"

echo "⏳ Waiting for pods to initialize..."
while true; do
    PHASE=$(kubectl get pod $LAUNCHER_POD -o jsonpath='{.status.phase}' 2>/dev/null || echo "Unknown")
    if [[ "$PHASE" == "Running" || "$PHASE" == "Succeeded" || "$PHASE" == "Failed" ]]; then
        break
    fi
    sleep 2
done

echo "--------------------------------------------------------"
echo "πŸ“Š Tailing real-time benchmark output..."
echo "--------------------------------------------------------"
kubectl logs -f $LAUNCHER_POD
EOF

Executing this script deploys a Kubernetes Job that dynamically compiles and runs a two-part hardware benchmark across all 16 Blackwell B200 GPUs. It first floods the 3.2 Tbps GPUDirect RoCE fabric using NVIDIA’s NCCL tests to verify network line rates, then orchestrates a distributed BF16 matrix multiplication to stress test the Tensor Cores. The entire process takes approximately 10 - 15 minutes on first run to provision the pods, compile the binaries, and stream the live TFLOPS performance metrics directly to your terminal.

Execute the Test:

bash native-test-nccl.sh

Output may look something similar to this:

#       size        count      type    redop     root      time   algbw   busbw  #wrong      time   algbw   busbw  #wrong 
#        (B)   (elements)                                  (us)  (GB/s)  (GB/s)              (us)  (GB/s)  (GB/s)         
  1073741824    268435456     float      sum       -1   5429.05  197.78  370.83       0   5456.98  196.76  368.93       0
  2147483648    536870912     float      sum       -1  10774.2   199.32  373.72       0  10790.7   199.01  373.15       0
  4294967296   1073741824     float      sum       -1  21448.2   200.25  375.47       0  21432.3   200.40  375.74       0
  8589934592   2147483648     float      sum       -1  42718.7   201.08  377.03       0  42735.3   201.00  376.88       0
 17179869184   4294967296     float      sum       -1  85289.9   201.43  377.68       0  85354.8   201.28  377.39       0
# Out of bounds values : 0 OK
# Avg bus bandwidth    : 374.682 
#
# Collective test concluded: all_reduce_perf
#

βœ… Network Benchmark completed!
========================================================
=== PART 2: Distributed BF16 Matrix Multiplication ===
========================================================
[Rank 03 | GPU 3] Matrix 16384x16384 BF16 GEMM: 1427.74 TFLOPS
[Rank 08 | GPU 0] Matrix 16384x16384 BF16 GEMM: 1432.65 TFLOPS
...

Orchestrating Gemma 4 with vLLM via Topology-Aware DRA (DRANET & NVIDIA GPU)

In this section, you will configure your secure Hugging Face API credentials as a Kubernetes secret, deploy your vLLM inference engine utilizing both your Dynamic Resource Allocation (DRA) network and hardware claims, and run an end-to-end test query against Google’s Gemma 4 model.

Prepare the Environment

Make sure you are logged into your secure SSH session on k8s-gpu-control-plane. Reconnect securely to the control plane VM from Cloud Shell. If you are already connected, skip this.

Bash

gcloud compute ssh k8s-gpu-control-plane \
    --zone=$ZONE \
    --tunnel-through-iap

Clean up previous benchmarking deployments and allow Kubernetes to release the hardware locks. Then, install the Google Managed Prometheus (GMP) Operator to enable vLLM telemetry:

Bash

cat << 'EOF' > install-gmp.sh
#!/bin/bash
set -e

echo "=== 1. Cleaning Up Benchmark State ==="
kubectl delete statefulset neper --ignore-not-found=true
kubectl delete statefulset,job,svc,pod -l app=nccl-benchmark --force --grace-period=0 2>/dev/null || true
echo "Waiting 10 seconds for the DRA controller to safely release GPU and Network locks..."
sleep 10

echo "=== 2. Installing GMP CRDs ==="
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/prometheus-engine/v0.17.2/manifests/setup.yaml

echo "=== 3. Configuring GMP for Unmanaged Kubernetes ==="
# Create the namespace manually first
kubectl create namespace gmp-public --dry-run=client -o yaml | kubectl apply -f -

# Fetch the exact Project ID and Zone this VM is running in
PROJECT_ID=$(curl -s "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
CLUSTER_ZONE=$(curl -s "http://metadata.google.internal/computeMetadata/v1/instance/zone" -H "Metadata-Flavor: Google" | awk -F/ '{print $4}')
echo "πŸ“ Detected Project: $PROJECT_ID"
echo "πŸ“ Detected Zone: $CLUSTER_ZONE"

# The project_id label is strictly required for unmanaged clusters to avoid internal RPC errors
cat << YAML > gmp-operator-config.yaml
apiVersion: monitoring.googleapis.com/v1
kind: OperatorConfig
metadata:
  namespace: gmp-public
  name: config
collection:
  externalLabels:
    project_id: "${PROJECT_ID}"
    cluster: "unmanaged-gpu-cluster"
    location: "${CLUSTER_ZONE}"
YAML

kubectl apply -f gmp-operator-config.yaml

echo "=== 4. Starting the GMP Operator ==="
kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/prometheus-engine/v0.17.2/manifests/operator.yaml

echo "=== 5. Injecting Cluster Identity to Operator Binary ==="
# Unmanaged clusters require these explicit flags on the operator container
kubectl patch deployment gmp-operator -n gmp-system --type='json' -p="[
  {\"op\": \"add\", \"path\": \"/spec/template/spec/containers/0/args/-\", \"value\": \"--cluster=unmanaged-gpu-cluster\"},
  {\"op\": \"add\", \"path\": \"/spec/template/spec/containers/0/args/-\", \"value\": \"--location=${CLUSTER_ZONE}\"}
]"

echo "⏳ Waiting for GMP operator pod to become available..."
kubectl wait --for=condition=Available --timeout=120s deployment/gmp-operator -n gmp-system

echo "--------------------------------------------------------"
echo "βœ… Google Managed Prometheus is fully installed & configured!"
echo "--------------------------------------------------------"
EOF

bash install-gmp.sh

Authenticate with Hugging Face

Store your Hugging Face Access Token. Replace <YOUR_ACTUAL_HUGGING_FACE_TOKEN> with your token.

export HF_TOKEN="<YOUR_ACTUAL_HUGGING_FACE_TOKEN>"

Create a secret:

kubectl create secret generic hf-token --from-literal=token="${HF_TOKEN}"

Deploy the DRA-Bound StatefulSet

This is where your custom infrastructure comes together. Instead of relying on legacy device plugins, you will deploy the Gemma-4-31B-it model using Kubernetes Dynamic Resource Allocation (DRA).

Your StatefulSet will natively request exact hardware topologies directly from the API server. By defining resourceClaims, you are simultaneously binding the 3.2 Tbps GPUDirect RoCE fabric (via the open-source DRANET driver) and 16 physical NVIDIA Blackwell B200 GPUs (via the NVIDIA GPU DRA driver) straight into your vLLM containers.

Run the following block:

Bash

cat << 'EOF' > gemma-distributed-inference.yaml
apiVersion: v1
kind: Service
metadata:
  name: gemma-ray-svc
  labels:
    app: gemma-distributed
spec:
  clusterIP: None
  selector:
    app: gemma-distributed
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-gemma-service
spec:
  selector:
    statefulset.kubernetes.io/pod-name: gemma-distributed-0
  ports:
  - protocol: TCP
    port: 8000
    targetPort: 8000
    name: vllm-http
  - protocol: TCP
    port: 8080
    targetPort: 8080
    name: ray-metrics
  type: ClusterIP
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: gemma-distributed
  labels:
    app: gemma-distributed
spec:
  podManagementPolicy: Parallel
  serviceName: "gemma-ray-svc"
  replicas: 2
  selector:
    matchLabels:
      app: gemma-distributed
  template:
    metadata:
      labels:
        app: gemma-distributed
    spec:
      hostIPC: true
      tolerations:
      - operator: "Exists"
      affinity:
        podAntiAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
          - labelSelector:
              matchExpressions:
              - key: app
                operator: In
                values:
                - gemma-distributed
            topologyKey: "kubernetes.io/hostname"
      
      initContainers:
      - name: pre-cache-model
        image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:gemma4
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-token
              key: token
        - name: HF_HOME
          value: "/models"
        command:
        - python3
        - -c
        - |
          from huggingface_hub import snapshot_download
          snapshot_download(repo_id="google/gemma-4-31B-it")
        volumeMounts:
        - mountPath: /models
          name: nvme-model-cache

      containers:
      - name: vllm-ray-node
        image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:gemma4
        securityContext:
          privileged: true
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-token
              key: token
        - name: HF_HOME
          value: "/models"
        - name: VLLM_DOWNLOAD_DIR
          value: "/models"
        - name: NCCL_SOCKET_IFNAME
          value: "eth0"
        - name: NCCL_IB_GDR_LEVEL
          value: "MAX"
        - name: NCCL_IB_HCA
          value: "mlx5"
        - name: NCCL_IB_DISABLE
          value: "0"
        command:
        - bash
        - -c
        - |
          export PYTHONUNBUFFERED=1
          export PATH="$HOME/.local/bin:$PATH"
          export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
          
          # Advanced NCCL & RDMA Tuning Variables
          export HF_HUB_OFFLINE=1
          export NCCL_TIMEOUT=120
          export VLLM_NCCL_TIMEOUT_S=120
          export NCCL_DEBUG=WARN
          export NCCL_NVLS_ENABLE=0
          export NCCL_ALGO=Ring
          export NCCL_MIN_NCHANNELS=32
          export NCCL_IB_PCI_RELAXED_ORDERING=1
          
          echo "=== Verifying Distributed Fabric (IP: $VLLM_HOST_IP) ==="
          export PIP_INDEX_URL="https://pypi.org/simple"
          pip install --user --no-cache-dir ray[default] || pip install --no-cache-dir ray[default]
          
          if [[ "$HOSTNAME" == "gemma-distributed-0" ]]; then
            echo "=== Starting Ray Head Node ==="
            ray start --head --node-ip-address=$VLLM_HOST_IP --port=6379 --num-gpus=8 --metrics-export-port=8080
            
            echo "=== Waiting for Worker Node to Join ==="
            echo "import ray, time" > /tmp/wait_gpus.py
            echo "ray.init(address='auto')" >> /tmp/wait_gpus.py
            echo "print('Connected to Ray. Waiting for 16 GPUs...')" >> /tmp/wait_gpus.py
            echo "while ray.cluster_resources().get('GPU', 0) < 16:" >> /tmp/wait_gpus.py
            echo "    time.sleep(5)" >> /tmp/wait_gpus.py
            echo "print('Success! Found 16 GPUs.')" >> /tmp/wait_gpus.py
            python3 /tmp/wait_gpus.py
            
            echo "βœ… 16 GPUs online and registered. Starting vLLM Server!"
            python3 -m vllm.entrypoints.openai.api_server \
              --model google/gemma-4-31B-it \
              --tensor-parallel-size 16 \
              --pipeline-parallel-size 1 \
              --distributed-executor-backend ray \
              --trust-remote-code \
              --max-model-len 32768 \
              --enforce-eager \
              --enable-prefix-caching \
              --disable-custom-all-reduce \
              --host 0.0.0.0 \
              --port 8000
          else
            echo "=== Starting Ray Worker Node ==="
            until ray start --address=gemma-distributed-0.gemma-ray-svc.default.svc.cluster.local:6379 --node-ip-address=$VLLM_HOST_IP --num-gpus=8 --metrics-export-port=8080 --block; do
              echo "⚠️ Ray Head not reachable yet. Retrying in 5 seconds..."
              sleep 5
            done
          fi
        ports:
        - containerPort: 8000
          name: vllm-http
        - containerPort: 8080
          name: ray-metrics
        - containerPort: 6379
          name: ray-redis
        
        resources:
          claims:
          - name: gpu-net-claim
          - name: gpu-hardware-claim
        volumeMounts:
        - name: dshm
          mountPath: /dev/shm
        - name: nvme-model-cache
          mountPath: /models
          
      volumes:
      - name: dshm
        hostPath:
          path: /dev/shm
      - name: nvme-model-cache
        hostPath:
          path: /mnt/models
          type: DirectoryOrCreate
          
      resourceClaims:
      - name: gpu-net-claim
        resourceClaimTemplateName: gpu-net-interfaces
      - name: gpu-hardware-claim
        resourceClaimTemplateName: gpu-device-template
---
apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
  name: vllm-metrics
  namespace: default
spec:
  selector:
    matchLabels:
      app: gemma-distributed
  endpoints:
  - port: 8000
    path: /metrics
    interval: 15s
    timeout: 10s
  - port: 8080
    path: /metrics
    interval: 15s
    timeout: 10s
EOF

kubectl apply -f gemma-distributed-inference.yaml

Chat Client Test

Create a script called start-chat.sh. This sends the mathematical stress-test prompt, and streams the response tokens back in real-time and allows an interactive interface.

Bash

cat << 'EOF' > start-chat.sh
#!/bin/bash
set -e

echo "=== Fetching Pod IP for Direct Cluster Routing ==="
POD_IP=$(kubectl get pod gemma-distributed-0 -o jsonpath='{.status.podIP}' 2>/dev/null || echo "")
if [ -z "$POD_IP" ]; then
    echo "❌ ERROR: Could not find pod gemma-distributed-0. Is the deployment running?"
    exit 1
fi
echo "Connecting directly to gemma-distributed-0 at http://${POD_IP}:8000"

python3 -c '
import requests, json, sys, time

pod_ip = "'"$POD_IP"'"
url = f"http://{pod_ip}:8000/v1/chat/completions"
headers = {"Content-Type": "application/json"}

# The conversational memory array
messages = []

def send_message(prompt, is_startup=False):
    if is_startup:
        print("\n=== πŸš€ STARTUP STRESS TEST: Deriving Schwarzschild Geodesics ===\n")
    else:
        print("\nπŸ€– Gemma: ", end="", flush=True)

    messages.append({"role": "user", "content": prompt})
    
    data = {
        "model": "google/gemma-4-31B-it",
        "messages": messages,
        "max_tokens": 8192,
        "temperature": 0.7,
        "stream": True
    }

    start = time.time()
    try:
        response = requests.post(url, headers=headers, json=data, stream=True, timeout=15)
        response.raise_for_status()
    except Exception as e:
        print(f"\n❌ Connection error: Make sure vLLM is running. Details: {e}", file=sys.stderr)
        sys.exit(1)

    tokens = 0
    full_text = ""
    for line in response.iter_lines():
        if line:
            line_str = line.decode("utf-8")
            if line_str.startswith("data: "):
                payload = line_str[6:]
                if payload != "[DONE]":
                    try:
                        chunk = json.loads(payload)
                        delta = chunk["choices"][0]["delta"].get("content", "")
                        if delta:
                            print(delta, end="", flush=True)
                            full_text += delta
                            tokens += 1
                    except Exception:
                        pass

    elapsed = time.time() - start
    speed = tokens / elapsed if elapsed > 0 else 0
    print(f"\n\n[⚑ Metrics | Total Time: {elapsed:.2f}s | Decode Speed: {speed:.2f} tok/s | Output: {tokens} tokens]")
    
    # Save the assistant response to memory
    messages.append({"role": "assistant", "content": full_text})


# 1. Run the Startup Stress Test
startup_prompt = "Derive the complete set of non-zero Christoffel symbols and the resulting geodesic equations for a massive particle orbiting a black hole using the Schwarzschild metric in spherical coordinates (t, r, ΞΈ, Ο†). You must: 1. Explicitly define the metric tensor and its inverse. 2. Show the step-by-step tensor calculus calculation for every non-zero Christoffel symbol using the standard formula. 3. Write out the four coupled second-order ordinary differential equations (the geodesic equations). 4. Provide a complete, production-ready Python script using numpy, matplotlib, and scipy.integrate.solve_ivp to numerically integrate these geodesic equations and plot the 2D orbital trajectory of the particle. Validate the code thoroughly."

send_message(startup_prompt, is_startup=True)

# 2. Drop into the Interactive Chat Loop
print("\n========================================================")
print("πŸ’¬ Interactive Chat Session Started. Type exit to quit.")
print("========================================================")

while True:
    try:
        user_input = input("\nπŸ‘€ You: ")
        if user_input.strip().lower() in ["exit", "quit"]:
            print("πŸ‘‹ Goodbye!")
            break
        if not user_input.strip():
            continue
            
        send_message(user_input)
        
    except (KeyboardInterrupt, EOFError):
        print("\nπŸ‘‹ Goodbye!")
        break
'
EOF

chmod +x start-chat.sh

Run Script to launch the interactive client. This runs a bunch of math calculations, returns the performance, and opens an interactive chat. Type exit to quit the chat session.

Bash

./start-chat.sh

View Real-Time Telemetry and Metrics

Because you installed the Google Managed Prometheus (GMP) Operator and configured the PodMonitoring custom resource, your cluster is actively scraping telemetry from the vLLM engine and pushing it to Google Cloud.

Now that you have run a heavy inference workload via the chat script, you can view your cluster’s performance data.

Option 1: Explore Pre-Built Dashboards Google Cloud Monitoring includes a predefined dashboard that is automatically installed when the vLLM integration is configured:

  1. In the Google Cloud console, navigate to Monitoring β†’ Dashboards.

  2. Select the Dashboard List tab.

  3. Choose the Integrations category.

  4. Click the name of the dashboard: vLLM Prometheus Overview.

Option 2: Run Custom PromQL Queries To manually inspect the hardware and token metrics you just generated:

  1. In the Google Cloud console, navigate to Monitoring β†’ Metrics Explorer.

  2. In the top right corner, ensure your time range is set to Last 1 hour.

  3. On the right side of the query builder, click the <> PromQL button to switch to raw text mode.

  4. Copy and paste any of the following queries into the text box and click Run Query:

Check GPU KV Cache Utilization: (If this consistently rides above 90%, you are running out of VRAM and need to increase tensor parallelism.)

Code snippet

vllm:gpu_cache_usage_perc

Track Token Generation Throughput (Speed): (Tracks the total number of tokens your cluster is generating per second.)

Code snippet

sum(rate(vllm:request_generation_tokens_sum[1m]))

Average Tokens Generated Per Request: (A chained query that calculates the average length of the model’s responses over a 5-minute window.)

Code snippet

sum(rate(vllm:request_generation_tokens_sum[5m])) / sum(rate(vllm:request_success_total[5m]))

Clean Up

Before tearing down the infrastructure, you must exit your secure SSH session on the control plane node and return to your Cloud Shell terminal. You also need to be inside your Terraform workspace directory.

# Exit the control plane SSH session
exit

# Navigate back to your Terraform workspace in Cloud Shell
cd ~/oss-kube-gpu-dra

Delete Orphan Firewalls

Kubernetes occasionally creates background firewall rules for load balancers or networking that Terraform doesn’t directly track. Run this script to cleanly remove them so Terraform can destroy the VPCs without hanging.

cat << 'EOF' > delete-orphan-firewalls.sh
#!/bin/bash

PROJECT_ID=$(gcloud config get-value project)

echo "=== Hunting for Auto-Generated Firewall Rules in Lab VPCs ==="
echo "Target Networks: oss-gpu-net-0, oss-gpu-net-1, oss-gpu-mrdma"
echo "-------------------------------------------------------------"

# 1. Fetch names of all firewall rules in lab networks
# 2. Exclude the exact 5 firewall rules that Terraform explicitly manages
ORPHAN_RULES=$(gcloud compute firewall-rules list \
  --project="${PROJECT_ID}" \
  --filter="network~oss-gpu-.*" \
  --format="value(name)" | \
  grep -vE "^(oss-gpu-internal-0|oss-gpu-internal-1|oss-gpu-mrdma-internal|oss-gpu-ssh|oss-gpu-allow-ping-net-0)$" || true)

# Calculate the count securely
if [ -z "$ORPHAN_RULES" ]; then
  COUNT=0
else
  COUNT=$(echo "$ORPHAN_RULES" | wc -l | tr -d ' ')
fi

echo "Found $COUNT orphan firewall rule(s)."
echo "-------------------------------------------------------------"

if [ "$COUNT" -gt 0 ]; then
  echo "Deleting the following auto-generated firewall rules..."
  
  # Print the rules being deleted for visibility
  echo "$ORPHAN_RULES"
  echo "..."
  
  # Convert newline-separated list to space-separated list for the delete command
  RULES_TO_DELETE=$(echo $ORPHAN_RULES | tr '\n' ' ')
  
  # Execute batch deletion
  gcloud compute firewall-rules delete $RULES_TO_DELETE \
    --project="${PROJECT_ID}" \
    --quiet
    
  echo "βœ… Successfully deleted $COUNT orphan firewall rules!"
else
  echo "βœ… No auto-generated firewall rules found. Your VPCs are clean."
fi

echo "-------------------------------------------------------------"
echo "Cleanup complete. You may now safely run 'terraform destroy'."
EOF

chmod +x delete-orphan-firewalls.sh
./delete-orphan-firewalls.sh

Destroy Infrastructure

Finally, tear down the virtual machines, disks, networks, and routers:

terraform destroy -auto-approve

Check out other related DRANET blogs

To learn more about DRANET checkout these other related experiments.

If you want to ask a question, find out more or share a thought? Please connect with me on LinkedIn or twitter @ammettw and send me a message.

I’ll be in touch

5 Likes