From d03a1beee5dd90e33ec18e181b0657b30f4c1959 Mon Sep 17 00:00:00 2001
From: bcreddy-gcp <123543489+bcreddy-gcp@users.noreply.github.com>
Date: Wed, 21 Feb 2024 00:18:21 +0530
Subject: [PATCH] Add support for custom containers for workbench instances
 (#9973)

* Add support for custom containers for workbench instances

* Use conflicts instead of exactly_one_of for backward compatibility
---
 mmv1/products/workbench/Instance.yaml         | 28 +++++++++++++++++++
 .../workbench_instance_basic_container.tf.erb | 11 ++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 mmv1/templates/terraform/examples/workbench_instance_basic_container.tf.erb

diff --git a/mmv1/products/workbench/Instance.yaml b/mmv1/products/workbench/Instance.yaml
index 1c0faeab3958..7b3e0d2b47c0 100644
--- a/mmv1/products/workbench/Instance.yaml
+++ b/mmv1/products/workbench/Instance.yaml
@@ -43,6 +43,15 @@ examples:
     region_override: 'us-west1-a'
     vars:
       instance_name: 'workbench-instance'
+  - !ruby/object:Provider::Terraform::Examples
+    name: 'workbench_instance_basic_container'
+    primary_resource_id: 'instance'
+    primary_resource_name: "fmt.Sprintf(\"tf-test-workbench-instance%s\",
+      context[\"\
+      random_suffix\"])"
+    region_override: 'us-west1-a'
+    vars:
+      instance_name: 'workbench-instance'
   - !ruby/object:Provider::Terraform::Examples
     name: 'workbench_instance_basic_gpu'
     primary_resource_id: 'instance'
@@ -213,6 +222,8 @@ properties:
         name: vmImage
         custom_flatten: templates/terraform/custom_flatten/workbench_instance_vm_image_flatten.go.erb
         default_from_api: true
+        conflicts:
+          - gce_setup.0.container_image
         immutable: true
         description: |
           Definition of a custom Compute Engine virtual machine image for starting
@@ -240,6 +251,23 @@ properties:
             exactly_one_of:
               - vm_image.0.name
               - vm_image.0.family
+      - !ruby/object:Api::Type::NestedObject
+        name: containerImage
+        conflicts:
+          - gce_setup.0.vm_image
+        description: |
+          Use a container image to start the workbench instance.
+        properties:
+          - !ruby/object:Api::Type::String
+            name: 'repository'
+            description: |
+              The path to the container image repository.
+              For example: gcr.io/{project_id}/{imageName}
+            required: true
+          - !ruby/object:Api::Type::String
+            name: 'tag'
+            description: |
+              The tag of the container image. If not specified, this defaults to the latest tag.
       - !ruby/object:Api::Type::NestedObject
         name: bootDisk
         default_from_api: true
diff --git a/mmv1/templates/terraform/examples/workbench_instance_basic_container.tf.erb b/mmv1/templates/terraform/examples/workbench_instance_basic_container.tf.erb
new file mode 100644
index 000000000000..55ba370ffee1
--- /dev/null
+++ b/mmv1/templates/terraform/examples/workbench_instance_basic_container.tf.erb
@@ -0,0 +1,11 @@
+resource "google_workbench_instance" "<%= ctx[:primary_resource_id] %>" {
+  name = "<%= ctx[:vars]['instance_name'] %>"
+  location = "us-west1-a"
+
+  gce_setup {
+    container_image {
+      repository = "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/base-cu113.py310"
+      tag = "latest"
+    }
+  }
+}