From 6fa151e8082403c162fdbc496079c1da92dcfe2d Mon Sep 17 00:00:00 2001
From: Johan Euphrosine <proppy@google.com>
Date: Thu, 22 Sep 2022 22:24:30 +0900
Subject: [PATCH] silicon/design/main: add vm based notebook

---
 modules/silicon_design/main.tf    | 47 +++++++++++++++++++++++++++++--
 modules/silicon_design/outputs.tf |  2 +-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/modules/silicon_design/main.tf b/modules/silicon_design/main.tf
index 326e8a4e..c30ca0dc 100644
--- a/modules/silicon_design/main.tf
+++ b/modules/silicon_design/main.tf
@@ -262,10 +262,10 @@ resource "google_project_iam_member" "ai_notebook_user_role2" {
   role     = "roles/viewer"
 }
 
-resource "google_notebooks_instance" "ai_notebook" {
+resource "google_notebooks_instance" "ai_notebook_container" {
   count        = var.notebook_count
   project      = local.project.project_id
-  name         = local.notebook_names[count.index]
+  name         = "${local.notebook_names[count.index]}-container"
   location     = var.zone
   machine_type = var.machine_type
 
@@ -305,6 +305,49 @@ resource "google_notebooks_instance" "ai_notebook" {
   ]
 }
 
+resource "google_notebooks_instance" "ai_notebook_vm" {
+  count        = var.notebook_count
+  project      = local.project.project_id
+  name         = "${local.notebook_names[count.index]}-vm"
+  location     = var.zone
+  machine_type = var.machine_type
+
+  vm_image {
+    project      = local.project.project_id
+    image_name   = "${var.image_name}-${local.image_tag}"
+  }
+
+  service_account = google_service_account.sa_p_notebook.email
+
+  install_gpu_driver = false
+  boot_disk_type     = var.boot_disk_type
+  boot_disk_size_gb  = var.boot_disk_size_gb
+  data_disk_type     = var.data_disk_type
+  data_disk_size_gb  = var.data_disk_size_gb
+  no_remove_data_disk = true
+
+  no_public_ip    = false
+  no_proxy_access = false
+
+  network = local.network.self_link
+  subnet  = local.subnet.self_link
+
+  post_startup_script = "gs://${google_storage_bucket.staging_bucket.name}/copy-notebooks.sh"
+
+  labels = {
+    module = "silicon-design"
+  }
+
+  metadata = {
+    terraform  = "true"
+    proxy-mode = "service_account"
+  }
+  depends_on = [
+    time_sleep.wait_120_seconds,
+    null_resource.build_and_push_image,
+  ]
+}
+
 resource "null_resource" "ai_notebook_provisioning_state" {
   for_each = toset(google_notebooks_instance.ai_notebook[*].name)
   provisioner "local-exec" {
diff --git a/modules/silicon_design/outputs.tf b/modules/silicon_design/outputs.tf
index 407c201c..39062bf2 100644
--- a/modules/silicon_design/outputs.tf
+++ b/modules/silicon_design/outputs.tf
@@ -42,7 +42,7 @@ output "notebook_container_image" {
 
 output "notebook_instance_names" {
   description = "Notebook Instance Names"
-  value       = join(", ", google_notebooks_instance.ai_notebook[*].name)
+  value       = join(", ", concat(google_notebooks_instance.ai_notebook_container[*].name, google_notebooks_instance.ai_notebook_vm[*].name))
 }
 
 output "artifact_registry_repository_id" {