WIP: Hotswap hero train job

ai-ml/hotswap-hero-train-job/deploy.sh

@@ -0,0 +1,46 @@
+#!/bin/sh
+
+# Set up env variables values
+
+# export PROJECT_ID=
+
+export REGION=us-west4
+export TPU_NODE_LOCATION=us-west4-a
+
+
+PROJECT_NUMBER=$(gcloud projects describe $PROJECT_ID --format="value(projectNumber)")
+
+gcloud services enable container.googleapis.com \
+    --project=$PROJECT_ID 
+
+# Create terraform.tfvars file 
+cat <<EOF >gke-platform/terraform.tfvars
+project_id                  = "$PROJECT_ID"
+enable_autopilot            = false
+enable_tpu                  = true
+region                      = "$REGION"
+tpu_node_location           = ["$TPU_NODE_LOCATION"]
+tpu_machine_type            = "ct5lp-hightpu-4t"
+tpu_topology                 = "2x4"
+tpu_node_pools_number       = 3
+EOF
+
+# Create clusters
+terraform -chdir=gke-platform init 
+terraform -chdir=gke-platform apply 
+
+# Get cluster credentials
+gcloud container clusters get-credentials llm-cluster \
+    --region=$REGION \
+    --project=$PROJECT_ID
+
+# Install JobSets
+kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.7.0/manifests.yaml
+
+sleep 60 # wait for jobset to install
+kubectl create -f workloads/priority.yaml
+
+kubectl create -f workloads/high-priority-job.yaml
+kubectl create -f workloads/low-priority-job.yaml
+
+

ai-ml/hotswap-hero-train-job/gke-platform/main.tf

@@ -0,0 +1,73 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+provider "google" {
+  project = var.project_id
+}
+
+provider "google-beta" {
+  project = var.project_id
+}
+
+resource "google_service_account" "service_account" {
+  account_id   = "gke-llm-sa"
+  display_name = "LLM clusters Service Account"
+}
+
+# Grant permissions to write metrics for monitoring purposes
+resource "google_project_iam_member" "project" {
+  project = var.project_id
+  role    = "roles/monitoring.metricWriter"
+  member  = "serviceAccount:${google_service_account.service_account.email}"
+}
+
+resource "google_project_iam_member" "logs_writer" {
+  project = var.project_id
+  role    = "roles/logging.logWriter"
+  member  = "serviceAccount:${google_service_account.service_account.email}"
+}
+
+module "gke_autopilot" {
+  source = "./modules/gke_autopilot"
+
+  project_id       = var.project_id
+  region           = var.region
+  cluster_name     = var.cluster_name
+  cluster_labels   = var.cluster_labels
+  enable_autopilot = var.enable_autopilot
+  service_account  = google_service_account.service_account.email
+  enable_fleet     = var.enable_fleet
+  fleet_project_id = var.fleet_project_id
+}
+
+
+
+module "gke_standard" {
+  source = "./modules/gke_standard"
+
+  project_id            = var.project_id
+  region                = var.region
+  cluster_name          = var.cluster_name
+  cluster_labels        = var.cluster_labels
+  enable_autopilot      = var.enable_autopilot
+  enable_tpu            = var.enable_tpu
+  tpu_node_location     = var.tpu_node_location
+  service_account       = google_service_account.service_account.email
+  enable_fleet          = var.enable_fleet
+  fleet_project_id      = var.fleet_project_id
+  gateway_api_channel   = var.gateway_api_channel
+  tpu_machine_type      = var.tpu_machine_type
+  tpu_node_pools_number = var.tpu_node_pools_number
+  tpu_topology          = var.tpu_topology
+}

ai-ml/hotswap-hero-train-job/gke-platform/modules/gke_autopilot/main.tf

@@ -0,0 +1,96 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+provider "google" {
+  project = var.project_id
+  region  = var.region
+}
+
+data "google_service_account" "default" {
+  account_id = var.service_account
+}
+
+# GKE cluster
+resource "google_container_cluster" "ml_cluster" {
+  name     = var.cluster_name
+  location = var.region
+  count    = var.enable_autopilot == true ? 1 : 0
+
+  deletion_protection = false
+
+  initial_node_count = 1
+
+  logging_config {
+    enable_components = ["SYSTEM_COMPONENTS", "WORKLOADS"]
+  }
+  node_config {
+    # Google recommends custom service accounts that have cloud-platform scope and permissions granted via IAM Roles.
+    service_account = data.google_service_account.default.email
+    oauth_scopes = [
+              "https://www.googleapis.com/auth/devstorage.read_only",
+              "https://www.googleapis.com/auth/logging.write",
+              "https://www.googleapis.com/auth/monitoring",
+              "https://www.googleapis.com/auth/service.management.readonly",
+              "https://www.googleapis.com/auth/servicecontrol",
+              "https://www.googleapis.com/auth/trace.append",
+    ]
+    reservation_affinity {
+      consume_reservation_type = "NO_RESERVATION"
+    }
+    gvnic {
+      enabled = true
+    }
+  }
+  cluster_autoscaling {
+    auto_provisioning_defaults {
+      service_account = data.google_service_account.default.email
+      oauth_scopes = [
+              "https://www.googleapis.com/auth/devstorage.read_only",
+              "https://www.googleapis.com/auth/logging.write",
+              "https://www.googleapis.com/auth/monitoring",
+              "https://www.googleapis.com/auth/service.management.readonly",
+              "https://www.googleapis.com/auth/servicecontrol",
+              "https://www.googleapis.com/auth/trace.append",
+      ]
+    }
+  }
+  monitoring_config {
+    enable_components = ["SYSTEM_COMPONENTS"]
+    managed_prometheus {
+      enabled = "true"
+    }
+  }
+
+  dynamic "fleet" {
+    for_each = var.enable_fleet ? [1] : []
+    content {
+      project = var.fleet_project_id
+    }
+  }
+
+  ip_allocation_policy {
+    cluster_ipv4_cidr_block  = ""
+    services_ipv4_cidr_block = ""
+  }
+
+  enable_autopilot = true
+
+  release_channel {
+    channel = "RAPID"
+  }
+
+  min_master_version = "1.31"
+
+  resource_labels = var.cluster_labels
+}

ai-ml/hotswap-hero-train-job/gke-platform/modules/gke_autopilot/output.tf

@@ -0,0 +1,38 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+output "project_id" {
+  description = "GCP project id"
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].project : null
+}
+
+output "region" {
+  description = "GCP region"
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].location : null
+}
+
+output "cluster_name" {
+  description = "The name of the GKE cluster"
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].name : null
+}
+
+output "kubernetes_host" {
+  description = "Kubernetes cluster host"
+  value       = var.enable_autopilot ? resource.google_container_cluster.ml_cluster[0].endpoint : null
+}
+
+output "cluster_certificate" {
+  description = "Kubernetes cluster CA certificate"
+  value       = var.enable_autopilot ? base64decode(resource.google_container_cluster.ml_cluster[0].master_auth[0].cluster_ca_certificate) : null
+}

ai-ml/hotswap-hero-train-job/gke-platform/modules/gke_autopilot/variables.tf

@@ -0,0 +1,64 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+variable "project_id" {
+  type        = string
+  description = "GCP project id"
+  default     = null
+}
+
+variable "region" {
+  type        = string
+  description = "GCP project region or zone"
+  default     = "us-central1"
+}
+
+variable "cluster_name" {
+  type        = string
+  description = "GKE cluster name"
+  default     = "ml-cluster"
+}
+
+variable "cluster_labels" {
+  type        = map(any)
+  description = "GKE cluster labels"
+  default = {
+    created-by = "ai-on-gke"
+  }
+}
+
+variable "num_gpu_nodes" {
+  description = "Number of GPU nodes in the cluster"
+  default     = 1
+}
+
+variable "enable_autopilot" {
+  type        = bool
+  description = "Set to true to enable GKE Autopilot clusters"
+  default     = false
+}
+
+variable "service_account" {
+  type = string
+}
+
+variable "enable_fleet" {
+  type    = bool
+  default = false
+}
+
+variable "fleet_project_id" {
+  type    = string
+  default = ""
+}