From 28160055afdec0e04fb440307d30eb0d491fbd42 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wiktor=20Niesiob=C4=99dzki?= <wiktorn@google.com>
Date: Sun, 20 Oct 2024 17:11:57 +0000
Subject: [PATCH] replace deprecated google_notebooks_instance with
 google_workbench_instance

---
 .../data-solutions/vertex-mlops/main.tf       |  9 +-
 .../data-solutions/vertex-mlops/outputs.tf    |  2 +-
 .../data-solutions/vertex-mlops/vertex.tf     | 87 ++++++++++---------
 3 files changed, 52 insertions(+), 46 deletions(-)

diff --git a/blueprints/data-solutions/vertex-mlops/main.tf b/blueprints/data-solutions/vertex-mlops/main.tf
index 7e3d08fd4..453392f8e 100644
--- a/blueprints/data-solutions/vertex-mlops/main.tf
+++ b/blueprints/data-solutions/vertex-mlops/main.tf
@@ -141,10 +141,10 @@ module "vpc-local" {
   name       = "vertex"
   subnets = [
     {
-      "name" : "subnet-${var.region}",
-      "region" : "${var.region}",
-      "ip_cidr_range" : "10.4.0.0/24",
-      "secondary_ip_range" : null
+      name                  = "subnet-${var.region}",
+      region                = var.region,
+      ip_cidr_range         = "10.4.0.0/24"
+      enable_private_access = true
     }
   ]
   psa_configs = [{
@@ -264,6 +264,7 @@ module "project" {
     "bigquery.googleapis.com",
     "bigquerystorage.googleapis.com",
     "cloudbuild.googleapis.com",
+    "containerfilesystem.googleapis.com",
     "compute.googleapis.com",
     "datacatalog.googleapis.com",
     "dataflow.googleapis.com",
diff --git a/blueprints/data-solutions/vertex-mlops/outputs.tf b/blueprints/data-solutions/vertex-mlops/outputs.tf
index 503d6168e..5cc1c70ba 100644
--- a/blueprints/data-solutions/vertex-mlops/outputs.tf
+++ b/blueprints/data-solutions/vertex-mlops/outputs.tf
@@ -36,7 +36,7 @@ output "notebook" {
   description = "Vertex AI notebooks ids."
   value = merge(
     { for k, v in resource.google_notebooks_runtime.runtime : k => v.id },
-    { for k, v in resource.google_notebooks_instance.playground : k => v.id }
+    { for k, v in resource.google_workbench_instance.playground : k => v.id }
   )
 }
 
diff --git a/blueprints/data-solutions/vertex-mlops/vertex.tf b/blueprints/data-solutions/vertex-mlops/vertex.tf
index c23de7804..32fdf8463 100644
--- a/blueprints/data-solutions/vertex-mlops/vertex.tf
+++ b/blueprints/data-solutions/vertex-mlops/vertex.tf
@@ -37,6 +37,13 @@ module "service-account-notebook" {
   source     = "../../../modules/iam-service-account"
   project_id = module.project.project_id
   name       = "notebook-sa"
+  iam_project_roles = {
+    (module.project.project_id) = [
+      "roles/logging.logWriter",
+      "roles/monitoring.metricWriter",
+      "roles/serviceusage.serviceUsageConsumer",
+    ]
+  }
 }
 
 resource "google_notebooks_runtime" "runtime" {
@@ -79,54 +86,52 @@ resource "google_notebooks_runtime" "runtime" {
   }
 }
 
-resource "google_notebooks_instance" "playground" {
-  for_each     = { for k, v in var.notebooks : k => v if v.type == "USER_MANAGED" }
-  name         = "${var.prefix}-${each.key}"
-  location     = "${var.region}-b"
-  machine_type = var.notebooks[each.key].machine_type
-  project      = module.project.project_id
+resource "google_workbench_instance" "playground" {
+  for_each = { for k, v in var.notebooks : k => v if v.type == "USER_MANAGED" }
+  project  = module.project.project_id
+  name     = "${var.prefix}-${each.key}"
+  location = "${var.region}-b"
 
-  container_image {
-    repository = "gcr.io/deeplearning-platform-release/base-cpu"
-    tag        = "latest"
-  }
-
-  install_gpu_driver = true
-  boot_disk_type     = "PD_SSD"
-  boot_disk_size_gb  = 110
-  disk_encryption    = var.service_encryption_keys.notebooks != null ? "CMEK" : null
-  kms_key            = var.service_encryption_keys.notebooks
-
-  no_public_ip    = var.notebooks[each.key].internal_ip_only
-  no_proxy_access = false
-
-  network = local.vpc
-  subnet  = local.subnet
-
-  instance_owners = try(tolist(var.notebooks[each.key].owner), null)
-  service_account = module.service-account-notebook.email
-  service_account_scopes = [
-    "https://www.googleapis.com/auth/cloud-platform",
-    "https://www.googleapis.com/auth/userinfo.email",
-  ]
-
-
-  metadata = {
-    notebook-disable-nbconvert = "false"
-    notebook-disable-downloads = "false"
-    notebook-disable-terminal  = "false"
-    notebook-disable-root      = "true"
+  gce_setup {
+    machine_type = var.notebooks[each.key].machine_type
+    container_image {
+      repository = "gcr.io/deeplearning-platform-release/workbench-container"
+      tag        = "latest"
+    }
+    boot_disk {
+      disk_size_gb    = 150
+      disk_type       = "PD_SSD"
+      disk_encryption = var.service_encryption_keys.notebooks != null ? "CMEK" : null
+      kms_key         = var.service_encryption_keys.notebooks
+    }
+
+    disable_public_ip = var.notebooks[each.key].internal_ip_only
+
+    network_interfaces {
+      network = local.vpc
+      subnet  = local.subnet
+    }
+    service_accounts {
+      email = module.service-account-notebook.email
+    }
+    metadata = {
+      notebook-disable-nbconvert = "false"
+      notebook-disable-downloads = "false"
+      notebook-disable-terminal  = "false"
+      notebook-disable-root      = "true"
+    }
+    tags = ["ssh"]
   }
+  disable_proxy_access = true
+  instance_owners      = try(tolist(var.notebooks[each.key].owner), null)
 
   # Remove once terraform-provider-google/issues/9164 is fixed
-  lifecycle {
-    ignore_changes = [disk_encryption, kms_key]
-  }
+  # lifecycle {
+  #   ignore_changes = [disk_encryption, kms_key]
+  # }
 
   #TODO Uncomment once terraform-provider-google/issues/9273 is fixed
-  # tags = ["ssh"]
   depends_on = [
     google_project_iam_member.shared_vpc,
   ]
 }
-