Runner startup script (#81)

Author: DavidGOrtega

iterative/resource_runner.go

@@ -124,7 +124,9 @@ func resourceRunner() *schema.Resource {
 			},
 			"startup_script": &schema.Schema{
 				Type:     schema.TypeString,
-				Computed: true,
+				ForceNew: true,
+				Optional: true,
+				Default:  "",
 			},
 			"aws_security_group": &schema.Schema{
 				Type:     schema.TypeString,
@@ -222,6 +224,10 @@ sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
 
 sudo npm install -g git+https://github.com/iterative/cml.git
 
+{{if .runner_startup_script}}
+{{.runner_startup_script}}
+{{end}}
+
 sudo bash -c 'cat << EOF > /usr/bin/cml.sh
 #!/bin/sh
 
@@ -323,6 +329,13 @@ func provisionerCode(d *schema.ResourceData) (string, error) {
 	data["AZURE_TENANT_ID"] = os.Getenv("AZURE_TENANT_ID")
 	data["ami"] = isAMIAvailable(d.Get("cloud").(string), d.Get("region").(string))
 
+	script, err := base64.StdEncoding.DecodeString(d.Get("startup_script").(string))
+	if err != nil {
+		return "", err
+	}
+
+	data["runner_startup_script"] = string(script)
+
 	return renderScript(data)
 }
 

iterative/resource_runner_test.go

@@ -1,6 +1,7 @@
 package iterative
 
 import (
+	"encoding/base64"
 	"strings"
 	"testing"
 
@@ -39,4 +40,13 @@ func TestScript(t *testing.T) {
 		script, _ := renderScript(data)
 		assert.Equal(t, strings.Contains(script, "sudo ubuntu-drivers autoinstall"), true)
 	})
+
+	t.Run("Runner Startup Script", func(t *testing.T) {
+		data := make(map[string]interface{})
+		startupScript, _ := base64.StdEncoding.DecodeString("ZWNobyAiaGVsbG8gd29ybGQiCmVjaG8gImJ5ZSB3b3JsZCI=")
+		data["runner_startup_script"] = string(startupScript)
+
+		script, _ := renderScript(data)
+		assert.Equal(t, strings.Contains(script, "echo \"hello world\"\necho \"bye world\""), true)
+	})
 }