diff --git a/amqp_job.go b/amqp_job.go index e87f56d26..8e7196aa8 100644 --- a/amqp_job.go +++ b/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) } diff --git a/amqp_job_queue.go b/amqp_job_queue.go index 39025277d..e2b7df9ad 100644 --- a/amqp_job_queue.go +++ b/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) diff --git a/backend/ec2.go b/backend/ec2.go index 028d175c1..84328a703 100644 --- a/backend/ec2.go +++ b/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, diff --git a/backend/gce.go b/backend/gce.go index f2dace457..2d80c66ed 100644 --- a/backend/gce.go +++ b/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 diff --git a/backend/start_attributes.go b/backend/start_attributes.go index 0bfacce56..075efbdb9 100644 --- a/backend/start_attributes.go +++ b/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:"-"` diff --git a/job.go b/job.go index f6b0b0532..00b4b79c7 100644 --- a/job.go +++ b/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"`