aws az runner

Author: DavidGOrtega

iterative/azure/provider.go

@@ -22,8 +22,6 @@ func ResourceMachineCreate(ctx context.Context, d *schema.ResourceData, m interf
 	subscriptionID := os.Getenv("AZURE_SUBSCRIPTION_ID")
 
 	username := "ubuntu"
-	//username := d.Get("ssh_user").(string)
-
 	customData := d.Get("startup_script").(string)
 	region := getRegion(d.Get("region").(string))
 	instanceType := getInstanceType(d.Get("instance_type").(string), d.Get("instance_gpu").(string))

iterative/resource_machine.go

@@ -74,12 +74,6 @@ func machineSchema() *map[string]*schema.Schema {
 			Type:     schema.TypeString,
 			Computed: true,
 		},
-		"ssh_public": &schema.Schema{
-			Type:     schema.TypeString,
-			ForceNew: true,
-			Optional: true,
-			Default:  "",
-		},
 		"ssh_private": &schema.Schema{
 			Type:     schema.TypeString,
 			ForceNew: true,
@@ -88,9 +82,7 @@ func machineSchema() *map[string]*schema.Schema {
 		},
 		"ssh_name": &schema.Schema{
 			Type:     schema.TypeString,
-			ForceNew: true,
-			Optional: true,
-			Default:  "ubuntu",
+			Computed: true,
 		},
 		"startup_script": &schema.Schema{
 			Type:     schema.TypeString,
@@ -112,17 +104,34 @@ func resourceMachineCreate(ctx context.Context, d *schema.ResourceData, m interf
 
 	utils.SetId(d)
 
-	script64 := base64.StdEncoding.EncodeToString([]byte(d.Get("startup_script").(string)))
-	d.Set("startup_script", script64)
+	if len(d.Get("ssh_private").(string)) == 0 {
+		private, err := utils.PrivatePEM()
+		if err != nil {
+			diags = append(resourceMachineDelete(ctx, d, m), diag.Diagnostic{
+				Severity: diag.Error,
+				Summary:  fmt.Sprintf("Failed creating the private PEM: %v", err),
+			})
 
-	keyPublic := d.Get("ssh_public").(string)
-	if len(keyPublic) == 0 {
-		public, private, _ := utils.SSHKeyPair()
+			return diags
+		}
 
-		d.Set("ssh_public", public)
 		d.Set("ssh_private", private)
 	}
 
+	public, err := utils.PublicFromPrivatePEM(d.Get("ssh_private").(string))
+	if err != nil {
+		diags = append(resourceMachineDelete(ctx, d, m), diag.Diagnostic{
+			Severity: diag.Error,
+			Summary:  fmt.Sprintf("Failed creating the public key: %v", err),
+		})
+
+		return diags
+	}
+	d.Set("ssh_public", public)
+
+	script64 := base64.StdEncoding.EncodeToString([]byte(d.Get("startup_script").(string)))
+	d.Set("startup_script", script64)
+
 	cloud := d.Get("cloud").(string)
 	if cloud == "aws" {
 		err := aws.ResourceMachineCreate(ctx, d, m)

iterative/resource_runner.go

@@ -99,22 +99,14 @@ func resourceRunner() *schema.Resource {
 			},
 			"ssh_public": &schema.Schema{
 				Type:     schema.TypeString,
-				ForceNew: true,
-				Optional: true,
-				Default:  "",
+				Computed: true,
 			},
 			"ssh_private": &schema.Schema{
 				Type:     schema.TypeString,
 				ForceNew: true,
 				Optional: true,
 				Default:  "",
 			},
-			"ssh_name": &schema.Schema{
-				Type:     schema.TypeString,
-				ForceNew: true,
-				Optional: true,
-				Default:  "ubuntu",
-			},
 			"startup_script": &schema.Schema{
 				Type:     schema.TypeString,
 				Computed: true,
@@ -152,7 +144,6 @@ func resourceRunnerCreate(ctx context.Context, d *schema.ResourceData, m interfa
 		})
 	} else {
 		diags = resourceMachineCreate(ctx, d, m)
-		//d.SetId("local")
 	}
 
 	return diags
@@ -204,6 +195,7 @@ func provisionerCode(d *schema.ResourceData) (string, error) {
 	}
 
 	data := make(map[string]string)
+	data["cloud"] = d.Get("cloud").(string)
 	data["token"] = d.Get("token").(string)
 	data["repo"] = d.Get("repo").(string)
 	data["driver"] = d.Get("driver").(string)
@@ -219,40 +211,34 @@ func provisionerCode(d *schema.ResourceData) (string, error) {
 	data["AZURE_TENANT_ID"] = os.Getenv("AZURE_TENANT_ID")
 
 	tmpl, err := template.New("deploy").Parse(`#!/bin/bash
-DEBIAN_FRONTEND=noninteractive
+export DEBIAN_FRONTEND=noninteractive
 
-{{if .cloud eq "azure"}}
+{{if eq .cloud "azure"}}
 echo "APT::Get::Assume-Yes \"true\";" | sudo tee -a /etc/apt/apt.conf.d/90assumeyes
-
 sudo apt update
-sudo curl -fsSL https://get.docker.com -o get-docker.sh && sudo sh get-docker.sh && \
+sudo curl -fsSL https://get.docker.com -o get-docker.sh && sudo sh get-docker.sh
 sudo usermod -aG docker ubuntu
 sudo setfacl --modify user:ubuntu:rw /var/run/docker.sock
-
 curl -fsSL https://apt.releases.hashicorp.com/gpg | sudo apt-key add -
 sudo apt-add-repository "deb [arch=amd64] https://apt.releases.hashicorp.com $(lsb_release -cs) main"
 sudo apt update && sudo apt-get install -y terraform
-
 curl -sL https://deb.nodesource.com/setup_12.x | sudo bash
 sudo apt update && sudo apt-get install -y nodejs
-
 sudo apt install -y ubuntu-drivers-common git
 sudo ubuntu-drivers autoinstall 
 sudo rmmod nvidia && sudo nvidia-smi
-curl -s -L https://nvidia.GitHub.io/nvidia-docker/gpgkey | sudo apt-key add - && \
+curl -s -L https://nvidia.GitHub.io/nvidia-docker/gpgkey | sudo apt-key add -
 curl -s -L https://nvidia.GitHub.io/nvidia-docker/ubuntu18.04/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
 sudo apt update && sudo apt install -y nvidia-container-toolkit
 {{end}}
-
 sudo npm install -g git+https://github.com/iterative/cml.git#cml-runner
-
 export AWS_SECRET_ACCESS_KEY={{.AWS_SECRET_ACCESS_KEY}}
 export AWS_ACCESS_KEY_ID={{.AWS_ACCESS_KEY_ID}}
 export AZURE_CLIENT_ID={{.AZURE_CLIENT_ID}}
 export AZURE_CLIENT_SECRET={{.AZURE_CLIENT_SECRET}}
 export AZURE_SUBSCRIPTION_ID={{.AZURE_SUBSCRIPTION_ID}}
 export AZURE_TENANT_ID={{.AZURE_TENANT_ID}}
-nohup cml-runner{{if .name}} --name {{.name}}{{end}}c --labels {{.labels}}{{end}}{{if .idle_timeout}} --idle-timeout {{.idle_timeout}}{{end}}{{if .driver}} --driver {{.driver}}{{end}}{{if .repo}} --repo {{.repo}}{{end}}{{if .token}} --token {{.token}}{{end}}{{if .tf_resource}} --tf_resource={{.tf_resource}}{{end}} < /dev/null > std.out 2> std.err &
+nohup cml-runner{{if .name}} --name {{.name}}{{end}}{{if .labels}} --labels {{.labels}}{{end}}{{if .idle_timeout}} --idle-timeout {{.idle_timeout}}{{end}}{{if .driver}} --driver {{.driver}}{{end}}{{if .repo}} --repo {{.repo}}{{end}}{{if .token}} --token {{.token}}{{end}}{{if .tf_resource}} --tf_resource={{.tf_resource}}{{end}} < /dev/null > std.out 2> std.err &
 sleep 10
 `)
 	var customDataBuffer bytes.Buffer

iterative/utils/ssh.go

@@ -10,29 +10,35 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-//SSHKeyPair generates an ssh keyppair
-func SSHKeyPair() (string, string, error) {
+func PrivatePEM() (string, error) {
 	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
 	if err != nil {
-		return "", "", err
+		return "", err
 	}
 
-	// generate and write private key as PEM
 	var privKeyBuf strings.Builder
-
 	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
 	if err := pem.Encode(&privKeyBuf, privateKeyPEM); err != nil {
-		return "", "", err
+		return "", err
+	}
+
+	return privKeyBuf.String(), nil
+}
+
+func PublicFromPrivatePEM(privateKey string) (string, error) {
+	block, _ := pem.Decode([]byte(privateKey))
+	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
+	if err != nil {
+		return "", err
 	}
 
-	// generate and write public key
-	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
+	pub, err := ssh.NewPublicKey(&key.PublicKey)
 	if err != nil {
-		return "", "", err
+		return "", err
 	}
 
 	var pubKeyBuf strings.Builder
 	pubKeyBuf.Write(ssh.MarshalAuthorizedKey(pub))
 
-	return pubKeyBuf.String(), privKeyBuf.String(), nil
+	return pubKeyBuf.String(), nil
 }