From aa541480ba27f2dfec85cbe4f07c07743bb901b2 Mon Sep 17 00:00:00 2001
From: Peter Amstutz <pamstutz@veritasgenetics.com>
Date: Wed, 27 Feb 2019 16:38:33 -0500
Subject: [PATCH] 14291: Support EBS attached storage and preemptible instances

Arvados-DCO-1.1-Signed-off-by: Peter Amstutz <pamstutz@veritasgenetics.com>
---
 lib/cloud/ec2/ec2.go      | 30 ++++++++++++++++--
 lib/cloud/ec2/ec2_test.go | 65 ++++++++++++++++++++++++++++++++++++++-
 sdk/go/arvados/config.go  |  1 +
 3 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/lib/cloud/ec2/ec2.go b/lib/cloud/ec2/ec2.go
index 115c1a0c09..664c51cd48 100644
--- a/lib/cloud/ec2/ec2.go
+++ b/lib/cloud/ec2/ec2.go
@@ -7,6 +7,7 @@ package ec2
 import (
 	"encoding/base64"
 	"encoding/json"
+	"fmt"
 	"log"
 	"strings"
 
@@ -83,6 +84,10 @@ func (instanceSet *ec2InstanceSet) Create(
 			Key:   aws.String(ARVADOS_DISPATCH_ID),
 			Value: aws.String(string(instanceSet.dispatcherID)),
 		},
+		&ec2.Tag{
+			Key:   aws.String("arvados-class"),
+			Value: aws.String("dynamic-compute"),
+		},
 	}
 	for k, v := range newTags {
 		ec2tags = append(ec2tags, &ec2.Tag{
@@ -91,7 +96,7 @@ func (instanceSet *ec2InstanceSet) Create(
 		})
 	}
 
-	rsv, err := instanceSet.client.RunInstances(&ec2.RunInstancesInput{
+	rii := ec2.RunInstancesInput{
 		ImageId:      aws.String(string(imageID)),
 		InstanceType: &instanceType.ProviderType,
 		MaxCount:     aws.Int64(1),
@@ -114,7 +119,28 @@ func (instanceSet *ec2InstanceSet) Create(
 				ResourceType: aws.String("instance"),
 				Tags:         ec2tags,
 			}},
-	})
+	}
+
+	if instanceType.ExtraScratch > 0 {
+		rii.BlockDeviceMappings = []*ec2.BlockDeviceMapping{&ec2.BlockDeviceMapping{
+			DeviceName: aws.String("/dev/xvdt"),
+			Ebs: &ec2.EbsBlockDevice{
+				DeleteOnTermination: aws.Bool(true),
+				VolumeSize:          aws.Int64((int64(instanceType.ExtraScratch) / 1000000000) + 1),
+				VolumeType:          aws.String("gp2"),
+			}}}
+	}
+
+	if instanceType.Preemptible {
+		rii.InstanceMarketOptions = &ec2.InstanceMarketOptionsRequest{
+			MarketType: aws.String("spot"),
+			SpotOptions: &ec2.SpotMarketOptions{
+				InstanceInterruptionBehavior: aws.String("terminate"),
+				MaxPrice:                     aws.String(fmt.Sprintf("%v", instanceType.Price)),
+			}}
+	}
+
+	rsv, err := instanceSet.client.RunInstances(&rii)
 
 	if err != nil {
 		return nil, err
diff --git a/lib/cloud/ec2/ec2_test.go b/lib/cloud/ec2/ec2_test.go
index cd469af1d0..ab6c6b597b 100644
--- a/lib/cloud/ec2/ec2_test.go
+++ b/lib/cloud/ec2/ec2_test.go
@@ -65,6 +65,25 @@ func GetInstanceSet() (cloud.InstanceSet, cloud.ImageID, arvados.Cluster, error)
 				Price:        .02,
 				Preemptible:  false,
 			},
+			"tiny-with-extra-scratch": arvados.InstanceType{
+				Name:         "tiny",
+				ProviderType: "m1.small",
+				VCPUs:        1,
+				RAM:          4000000000,
+				Scratch:      10000000000,
+				ExtraScratch: 20000000000,
+				Price:        .02,
+				Preemptible:  false,
+			},
+			"tiny-preemptible": arvados.InstanceType{
+				Name:         "tiny",
+				ProviderType: "m1.small",
+				VCPUs:        1,
+				RAM:          4000000000,
+				Scratch:      10000000000,
+				Price:        .02,
+				Preemptible:  true,
+			},
 		})}
 	if *live != "" {
 		var exampleCfg testConfig
@@ -108,10 +127,54 @@ func (*EC2InstanceSetSuite) TestCreate(c *check.C) {
 
 }
 
+func (*EC2InstanceSetSuite) TestCreateWithExtraScratch(c *check.C) {
+	ap, img, cluster, err := GetInstanceSet()
+	if err != nil {
+		c.Fatal("Error making provider", err)
+	}
+
+	pk, _, _, _, err := ssh.ParseAuthorizedKey(testKey)
+	c.Assert(err, check.IsNil)
+
+	inst, err := ap.Create(cluster.InstanceTypes["tiny-with-extra-scratch"],
+		img, map[string]string{
+			"TestTagName": "test tag value",
+		}, "umask 0600; echo -n test-file-data >/var/run/test-file", pk)
+
+	c.Assert(err, check.IsNil)
+
+	tags := inst.Tags()
+	c.Check(tags["TestTagName"], check.Equals, "test tag value")
+	c.Logf("inst.String()=%v Address()=%v Tags()=%v", inst.String(), inst.Address(), tags)
+
+}
+
+func (*EC2InstanceSetSuite) TestCreatePreemptible(c *check.C) {
+	ap, img, cluster, err := GetInstanceSet()
+	if err != nil {
+		c.Fatal("Error making provider", err)
+	}
+
+	pk, _, _, _, err := ssh.ParseAuthorizedKey(testKey)
+	c.Assert(err, check.IsNil)
+
+	inst, err := ap.Create(cluster.InstanceTypes["tiny-preemptible"],
+		img, map[string]string{
+			"TestTagName": "test tag value",
+		}, "umask 0600; echo -n test-file-data >/var/run/test-file", pk)
+
+	c.Assert(err, check.IsNil)
+
+	tags := inst.Tags()
+	c.Check(tags["TestTagName"], check.Equals, "test tag value")
+	c.Logf("inst.String()=%v Address()=%v Tags()=%v", inst.String(), inst.Address(), tags)
+
+}
+
 func (*EC2InstanceSetSuite) TestListInstances(c *check.C) {
 	ap, _, _, err := GetInstanceSet()
 	if err != nil {
-		c.Fatal("Error making provider", err)
+		c.Fatal("Error making provider: ", err)
 	}
 
 	l, err := ap.Instances(nil)
diff --git a/sdk/go/arvados/config.go b/sdk/go/arvados/config.go
index c2154d0f29..5f6235cf0a 100644
--- a/sdk/go/arvados/config.go
+++ b/sdk/go/arvados/config.go
@@ -99,6 +99,7 @@ type InstanceType struct {
 	VCPUs        int
 	RAM          ByteSize
 	Scratch      ByteSize
+	ExtraScratch ByteSize
 	Price        float64
 	Preemptible  bool
 }
-- 
2.30.2