Mix training and inference infra and manifests

ai-ml/mix-train-and-inference/deploy.sh

@@ -0,0 +1,91 @@
+#!/bin/sh
+# Copyright 2024 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.
+
+# Set up env variables values
+# export HF_TOKEN=HF_TOKEN
+# export PROJECT_ID=
+export REGION=europe-west6
+export GPU_POOL_MACHINE_TYPE="g2-standard-24"
+export GPU_POOL_ACCELERATOR_TYPE="nvidia-l4"
+export TRAINING_DATA_BUCKET="data-bucket-$PROJECT_ID"
+export MODEL_BUCKET="model-bucket-$PROJECT_ID"
+
+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            = true
+region                      = "$REGION"
+gpu_pool_machine_type       = "$GPU_POOL_MACHINE_TYPE"
+gpu_pool_accelerator_type   = "$GPU_POOL_ACCELERATOR_TYPE"
+gpu_pool_node_locations     = $(gcloud compute accelerator-types list --filter="zone ~ $REGION AND name=$GPU_POOL_ACCELERATOR_TYPE" --limit=2 --format=json | jq -sr 'map(.[].zone|split("/")|.[8])|tojson')
+
+enable_fleet                = false
+gateway_api_channel         = "CHANNEL_STANDARD"
+training_data_bucket        = "$TRAINING_DATA_BUCKET"
+model_bucket                = "$MODEL_BUCKET"
+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
+
+# upload training dataset to bucket
+gcloud storage rsync -r training_data/ gs://$TRAINING_DATA_BUCKET/
+
+NAMESPACE=llm
+
+kubectl create ns $NAMESPACE
+kubectl create secret generic hf-secret \
+--from-literal=hf_api_token=$HF_TOKEN \
+--dry-run=client -o yaml | kubectl apply -n $NAMESPACE -f -
+
+# kubectl apply --server-side -f manifests.yaml
+kubectl kustomize kueue/ |kubectl apply --server-side -f - 
+
+sleep 180 # wait for kueue deployment
+
+cd workloads
+kubectl apply -f flavors.yaml
+kubectl apply -f default-priorityclass.yaml
+kubectl apply -f high-priorityclass.yaml
+kubectl apply -f low-priorityclass.yaml
+kubectl apply -f cluster-queue.yaml
+kubectl apply -f local-queue.yaml -n $NAMESPACE
+
+gcloud storage buckets add-iam-policy-binding "gs://$MODEL_BUCKET" \
+    --role=roles/storage.objectAdmin \
+    --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/$NAMESPACE/sa/default \
+    --condition=None
+gcloud storage buckets add-iam-policy-binding "gs://$TRAINING_DATA_BUCKET" \
+    --role=roles/storage.objectViewer \
+    --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/$NAMESPACE/sa/default \
+    --condition=None
+
+kubectl create -f tgi-gemma-2-9b-it-hp.yaml -n $NAMESPACE
+
+# deploy fine-tuning job
+sed -e "s/<TRAINING_BUCKET>/$TRAINING_DATA_BUCKET/g" \
+-e "s/<MODEL_BUCKET>/$MODEL_BUCKET/g" \
+fine-tune-l4.yaml |kubectl apply -f - -n $NAMESPACE

ai-ml/mix-train-and-inference/gke-platform/main.tf

@@ -0,0 +1,83 @@
+# Copyright 2024 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
+  gpu_pool_machine_type     = var.gpu_pool_machine_type
+  gpu_pool_accelerator_type = var.gpu_pool_accelerator_type
+  gpu_pool_node_locations   = var.gpu_pool_node_locations
+  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
+}
+
+resource "google_storage_bucket" "train_data" {
+  name          = var.training_data_bucket
+  location      = var.region
+  uniform_bucket_level_access = true
+}
+
+resource "google_storage_bucket" "model_data" {
+  name          = var.model_bucket
+  location      = var.region
+  uniform_bucket_level_access = true
+}

ai-ml/mix-train-and-inference/gke-platform/modules/gke_autopilot/main.tf

@@ -0,0 +1,84 @@
+# Copyright 2024 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
+
+  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/cloud-platform"
+    ]
+    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/cloud-platform"
+      ]
+    }
+  }
+  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/mix-train-and-inference/gke-platform/modules/gke_autopilot/output.tf

@@ -0,0 +1,38 @@
+# Copyright 2024 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/mix-train-and-inference/gke-platform/modules/gke_autopilot/variables.tf

@@ -0,0 +1,64 @@
+# Copyright 2024 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 = ""
+}