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. (#93)

Browse source 
* Adding support for setting GPU as a resource.
* Refactoring pulse update test.
This commit is contained in:
Renan DelValle 2019-01-08 15:11:52 -08:00 • committed by GitHub
parent 296af622d1
commit 2f7015571c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 83 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

40
job.go
View file

@ -58,10 +58,19 @@ type Job interface {
PartitionPolicy(policy *aurora.PartitionPolicy) 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
}
@ -88,15 +97,8 @@ func NewJob() Job {
ramMb := aurora.NewResource()
diskMb := aurora.NewResource()
resources := make(map[string]*aurora.Resource)
resources["cpu"] = numCpus
resources["ram"] = ramMb
resources["disk"] = diskMb
taskConfig.Resources = make(map[*aurora.Resource]bool)
taskConfig.Resources[numCpus] = true
taskConfig.Resources[ramMb] = true
taskConfig.Resources[diskMb] = true
resources := map[ResourceType]*aurora.Resource{CPU: numCpus, RAM: ramMb, DISK: diskMb}
taskConfig.Resources = map[*aurora.Resource]bool{numCpus: true, ramMb: true, diskMb: true}
numCpus.NumCpus = new(float64)
ramMb.RamMb = new(int64)
@ -155,20 +157,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
}