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

Adding support for setting GPU as a resource.

Browse source
This commit is contained in:
Renan DelValle 2019-01-07 17:12:17 -08:00
parent 9a835631b2
commit 25f5ff9f07
No known key found for this signature in database
GPG key ID: C240AD6D6F443EC9
3 changed files with 86 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

37
job.go
View file

@ -57,10 +57,19 @@ type Job interface {
Container(container Container) Job
}
type ResourceType int
const (
CPU ResourceType = iota
RAM
DISK
GPU
)
// Structure to collect all information pertaining to an Aurora job.
type AuroraJob struct {
jobConfig *aurora.JobConfiguration
resources map[string]*aurora.Resource
resources map[ResourceType]*aurora.Resource
portCount int
}
@ -87,10 +96,10 @@ func NewJob() Job {
ramMb := aurora.NewResource()
diskMb := aurora.NewResource()
resources := make(map[string]*aurora.Resource)
resources["cpu"] = numCpus
resources["ram"] = ramMb
resources["disk"] = diskMb
resources := make(map[ResourceType]*aurora.Resource)
resources[CPU] = numCpus
resources[RAM] = ramMb
resources[DISK] = diskMb
taskConfig.Resources = make(map[*aurora.Resource]bool)
taskConfig.Resources[numCpus] = true
@ -154,20 +163,28 @@ func (j *AuroraJob) ExecutorData(data string) Job {
}
func (j *AuroraJob) CPU(cpus float64) Job {
*j.resources["cpu"].NumCpus = cpus
*j.resources[CPU].NumCpus = cpus
return j
}
func (j *AuroraJob) RAM(ram int64) Job {
*j.resources["ram"].RamMb = ram
*j.resources[RAM].RamMb = ram
return j
}
func (j *AuroraJob) Disk(disk int64) Job {
*j.resources["disk"].DiskMb = disk
*j.resources[DISK].DiskMb = disk
return j
}
func (j *AuroraJob) GPU(gpus int64) Job {
if _,ok := j.resources[GPU]; !ok {
numGPUs := &aurora.Resource{NumGpus: new(int64)}
j.resources[GPU] = numGPUs
j.TaskConfig().Resources[numGPUs] = true
}
*j.resources[GPU].NumGpus = gpus
return j
}