ridv/gorealis-v2
1
0
Fork 0
Code Issues Pull requests Projects Releases 13 Packages Wiki Activity Actions

Initial support for Thermos and GPU resources.

Browse source
This commit is contained in:
Renan DelValle 2019-01-07 14:38:08 -08:00
parent afcdaa84b8
commit e13349db26
No known key found for this signature in database
GPG key ID: C240AD6D6F443EC9
9 changed files with 359 additions and 176 deletions
Download patch file Download diff file Expand all files Collapse all files

58
task.go
View file

@ -15,6 +15,7 @@
package realis
import (
"encoding/json"
"strconv"
"git.apache.org/thrift.git/lib/go/thrift"
@ -27,6 +28,7 @@ const (
CPU ResourceType = iota
RAM
DISK
GPU
)
const (
@ -38,12 +40,14 @@ type AuroraTask struct {
task *aurora.TaskConfig
resources map[ResourceType]*aurora.Resource
portCount int
thermos *ThermosExecutor
}
func NewTask() *AuroraTask {
numCpus := &aurora.Resource{}
ramMb := &aurora.Resource{}
diskMb := &aurora.Resource{}
numGpus := &aurora.Resource{}
numCpus.NumCpus = new(float64)
ramMb.RamMb = new(int64)
@ -53,6 +57,7 @@ func NewTask() *AuroraTask {
resources[CPU] = numCpus
resources[RAM] = ramMb
resources[DISK] = diskMb
resources[GPU] = numGpus
return &AuroraTask{task: &aurora.TaskConfig{
Job: &aurora.JobKey{},
@ -122,7 +127,7 @@ func TaskFromThrift(config *aurora.TaskConfig) *AuroraTask {
// Copy all ports
for _, resource := range config.Resources {
// Copy only ports, skip CPU, RAM, and DISK
// Copy only ports. Set CPU, RAM, DISK, and GPU
if resource != nil {
if resource.NamedPort != nil {
newTask.task.Resources = append(newTask.task.Resources, &aurora.Resource{NamedPort: thrift.StringPtr(*resource.NamedPort)})
@ -140,6 +145,10 @@ func TaskFromThrift(config *aurora.TaskConfig) *AuroraTask {
if resource.DiskMb != nil {
newTask.Disk(*resource.DiskMb)
}
if resource.NumGpus != nil {
newTask.GPU(*resource.NumGpus)
}
}
}
@ -237,6 +246,11 @@ func (t *AuroraTask) Disk(disk int64) *AuroraTask {
return t
}
func (t *AuroraTask) GPU(gpu int64) *AuroraTask {
*t.resources[GPU].NumGpus = gpu
return t
}
func (t *AuroraTask) Tier(tier string) *AuroraTask {
t.task.Tier = &tier
return t
@ -363,5 +377,47 @@ func (t *AuroraTask) JobKey() aurora.JobKey {
func (t *AuroraTask) Clone() *AuroraTask {
newTask := TaskFromThrift(t.task)
if t.thermos != nil {
newTask.ThermosExecutor(*t.thermos.Clone())
}
return newTask
}
func (t *AuroraTask) ThermosExecutor(thermos ThermosExecutor) *AuroraTask {
t.thermos = &thermos
return t
}
func (t *AuroraTask) BuildThermosPayload() error {
if t.thermos != nil {
// Set the correct resources
if t.resources[CPU].NumCpus != nil {
t.thermos.cpu(*t.resources[CPU].NumCpus)
}
if t.resources[RAM].RamMb != nil {
t.thermos.ram(*t.resources[RAM].RamMb)
}
if t.resources[DISK].DiskMb != nil {
t.thermos.disk(*t.resources[DISK].DiskMb)
}
if t.resources[GPU].NumGpus != nil {
t.thermos.gpu(*t.resources[GPU].NumGpus)
}
payload, err := json.Marshal(t.thermos)
if err != nil {
return err
}
t.ExecutorName(aurora.AURORA_EXECUTOR_NAME)
t.ExecutorData(string(payload))
}
return nil
}