Merge pull request #648 from travis-ci/epic-vm

added vm_size parameter for setting instance size from amqp (#646)

amqp_job.go

@@ -184,6 +184,10 @@ func (j *amqpJob) createStateUpdateBody(ctx gocontext.Context, state string) map
 		body["meta"].(map[string]interface{})["uuid"] = uuid
 	}
 
+	if j.startAttributes.VMSize != "" {
+		body["vm_size"] = j.startAttributes.VMSize
+	}
+
 	if j.Payload().Job.QueuedAt != nil {
 		body["queued_at"] = j.Payload().Job.QueuedAt.UTC().Format(time.RFC3339)
 	}

amqp_job_queue.go

@@ -222,6 +222,7 @@ func (q *AMQPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err erro
 
 				buildJob.startAttributes = startAttrs.Config
 				buildJob.startAttributes.VMType = buildJob.payload.VMType
+				buildJob.startAttributes.VMSize = buildJob.payload.VMSize
 				buildJob.startAttributes.VMConfig = buildJob.payload.VMConfig
 				buildJob.startAttributes.Warmer = buildJob.payload.Warmer
 				buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultArch, q.DefaultGroup, q.DefaultOS, VMTypeDefault, VMConfigDefault)

backend/ec2.go

@@ -54,6 +54,13 @@ chown -R travis:travis ~travis/.ssh/
 
 {{ .UserData }}
 `))
+
+	ec2VMSizeMapping = map[string]string{
+		"medium": "m6g.large",
+		"large": "m6g.xlarge",
+		"x-large": "m6g.2xlarge",
+		"2x-large": "m6g.4xlarge",
+	}
 )
 
 type ec2StartupScriptData struct {
@@ -438,11 +445,19 @@ func (p *ec2Provider) Start(ctx gocontext.Context, startAttributes *StartAttribu
 	if p.keyName != "" {
 		keyName = aws.String(p.keyName)
 	}
-	//RequestSpotInstances
 
+	instanceType := p.instanceType
+
+	if startAttributes.VMSize != "" {
+		if mtype, ok:= ec2VMSizeMapping[startAttributes.VMSize]; ok {
+			instanceType = mtype
+			startAttributes.VMSize = instanceType // convert provided vm_size to EC2 machine size
+		}
+	}
+	//RequestSpotInstances
 	runOpts := &ec2.RunInstancesInput{
 		ImageId:             aws.String(imageID),
-		InstanceType:        aws.String(p.instanceType),
+		InstanceType:        aws.String(instanceType),
 		MaxCount:            aws.Int64(1),
 		MinCount:            aws.Int64(1),
 		KeyName:             keyName,

backend/gce.go

@@ -145,6 +145,13 @@ Set-LocalUser -Name travis -Password $pw
 	// FIXME: get rid of the need for this global goop
 	gceCustomHTTPTransport     http.RoundTripper
 	gceCustomHTTPTransportLock sync.Mutex
+
+	gceVMSizeMapping = map[string]string{
+		"medium": "n2-standard-2",
+		"large": "n2-standard-4",
+		"x-large": "n2-standard-8",
+		"2x-large": "n2-standard-16",
+	}
 )
 
 type gceStartupScriptData struct {
@@ -818,8 +825,12 @@ func (p *gceProvider) Setup(ctx gocontext.Context) error {
 		zoneNames[i] = zone.Name
 	}
 
+	machineTypes := []string{p.ic.MachineType, p.ic.PremiumMachineType}
+	for _, machineType := range gceVMSizeMapping {
+		machineTypes = append(machineTypes, machineType);
+	}
 	for _, zoneName := range append(zoneNames, p.alternateZones...) {
-		for _, machineType := range []string{p.ic.MachineType, p.ic.PremiumMachineType} {
+		for _, machineType := range machineTypes {
 			if zoneName == "" || machineType == "" {
 				continue
 			}
@@ -1497,7 +1508,14 @@ func (p *gceProvider) buildInstance(ctx gocontext.Context, c *gceStartContext) (
 
 	machineType := p.ic.MachineType
 	if c.startAttributes.VMType == "premium" {
+		c.startAttributes.VMSize = "premium"
 		machineType = p.ic.PremiumMachineType
+	} else if c.startAttributes.VMSize != "" {
+		if mtype, ok := gceVMSizeMapping[c.startAttributes.VMSize]; ok {
+			machineType = mtype;
+			//storing converted machine type for instance size identification
+			c.startAttributes.VMSize = machineType
+		}
 	}
 
 	var ok bool

backend/start_attributes.go

@@ -29,6 +29,12 @@ type StartAttributes struct {
 	// the job payload, see the worker.JobPayload struct.
 	VMConfig VmConfig `json:"-"`
 
+
+	// The VMSize isn't stored in the config directly, but in the top level of
+	// the job payload, see the worker.JobPayload struct.
+	VMSize string `json:"-"`
+
+
 	// Warmer isn't stored in the config directly, but in the top level of
 	// the job payload, see the worker.JobPayload struct.
 	Warmer bool `json:"-"`

job.go

@@ -36,6 +36,7 @@ type JobPayload struct {
 	Timeouts   TimeoutsPayload        `json:"timeouts,omitempty"`
 	VMType     string                 `json:"vm_type"`
 	VMConfig   backend.VmConfig       `json:"vm_config"`
+	VMSize     string                 `json:"vm_size"`
 	Meta       JobMetaPayload         `json:"meta"`
 	Queue      string                 `json:"queue"`
 	Trace      bool                   `json:"trace"`