From 0500297dde9a16e181d0e2228e06b835d72ad325 Mon Sep 17 00:00:00 2001
From: bcreddy-gcp <123543489+bcreddy-gcp@users.noreply.github.com>
Date: Fri, 31 Jan 2025 16:41:28 -0800
Subject: [PATCH] Support Post sart up script on runtime templates (#12929)

---
 mmv1/products/colab/RuntimeTemplate.yaml      | 37 +++++++++++++++++++
 .../colab_runtime_template_full.tf.tmpl       | 13 +++++++
 2 files changed, 50 insertions(+)

diff --git a/mmv1/products/colab/RuntimeTemplate.yaml b/mmv1/products/colab/RuntimeTemplate.yaml
index 01528df0facd..2d3edddd7378 100644
--- a/mmv1/products/colab/RuntimeTemplate.yaml
+++ b/mmv1/products/colab/RuntimeTemplate.yaml
@@ -179,3 +179,40 @@ properties:
       - name: 'kmsKeyName'
         type: String
         description: 'The Cloud KMS encryption key (customer-managed encryption key) used to protect the runtime.'
+  - name: softwareConfig
+    type: NestedObject
+    description: 'The notebook software configuration of the notebook runtime.'
+    properties:
+      - name: 'env'
+        type: Array
+        description: 'Environment variables to be passed to the container.'
+        item_type:
+          type: NestedObject
+          properties:
+            - name: 'name'
+              type: String
+              description: 'Name of the environment variable. Must be a valid C identifier.'
+            - name: 'value'
+              type: String
+              description: 'Variables that reference a $(VAR_NAME) are expanded using the previous defined
+                             environment variables in the container and any service environment variables.
+                            If a variable cannot be resolved, the reference in the input string will be unchanged.
+                            The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped
+                            references will never be expanded, regardless of whether the variable exists or not.'
+      - name: 'postStartupScriptConfig'
+        type: NestedObject
+        description: 'Post startup script config.'
+        properties:
+          - name: 'postStartupScript'
+            type: String
+            description: 'Post startup script to run after runtime is started.'
+          - name: 'postStartupScriptUrl'
+            type: String
+            description: 'Post startup script url to download. Example: https://bucket/script.sh.'
+          - name: 'postStartupScriptBehavior'
+            type: Enum
+            description: 'Post startup script behavior that defines download and execution behavior.'
+            enum_values:
+              - 'RUN_ONCE'
+              - 'RUN_EVERY_START'
+              - 'DOWNLOAD_AND_RUN_EVERY_START'
diff --git a/mmv1/templates/terraform/examples/colab_runtime_template_full.tf.tmpl b/mmv1/templates/terraform/examples/colab_runtime_template_full.tf.tmpl
index 3078041236d9..d5a9b79539ff 100644
--- a/mmv1/templates/terraform/examples/colab_runtime_template_full.tf.tmpl
+++ b/mmv1/templates/terraform/examples/colab_runtime_template_full.tf.tmpl
@@ -53,4 +53,17 @@ resource "google_colab_runtime_template" "{{$.PrimaryResourceId}}" {
   encryption_spec {
     kms_key_name = "{{index $.Vars "key_name"}}"
   }
+
+  software_config {
+    env {
+      name         = "TEST"
+      value   = 1
+    }
+
+    post_startup_script_config {
+      post_startup_script = "echo 'hello world'"
+      post_startup_script_url = "gs://colab-enterprise-pss-secure/secure_pss.sh"
+      post_startup_script_behavior = "RUN_ONCE"
+    }
+  }
 }