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

Fixing GPU resource to only be added if specified since Aurora scheduler by default will reject tasks containing GPU.

Browse source
This commit is contained in:
Renan DelValle 2019-01-08 17:39:22 -08:00
parent 8d67d8c2f3
commit 9b3593e9d9
No known key found for this signature in database
GPG key ID: C240AD6D6F443EC9
1 changed files with 9 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

16
task.go
View file

@ -47,17 +47,12 @@ func NewTask() *AuroraTask {
numCpus := &aurora.Resource{}
ramMb := &aurora.Resource{}
diskMb := &aurora.Resource{}
numGpus := &aurora.Resource{}
numCpus.NumCpus = new(float64)
ramMb.RamMb = new(int64)
diskMb.DiskMb = new(int64)
resources := make(map[ResourceType]*aurora.Resource)
resources[CPU] = numCpus
resources[RAM] = ramMb
resources[DISK] = diskMb
resources[GPU] = numGpus
resources := map[ResourceType]*aurora.Resource{CPU: numCpus, RAM: ramMb, DISK: diskMb}
return &AuroraTask{task: &aurora.TaskConfig{
Job: &aurora.JobKey{},
@ -247,7 +242,14 @@ func (t *AuroraTask) Disk(disk int64) *AuroraTask {
}
func (t *AuroraTask) GPU(gpu int64) *AuroraTask {
*t.resources[GPU].NumGpus = gpu
// GPU resource must be set explicitly since the scheduler by default
// rejects jobs with GPU resources attached to it.
if _, ok := t.resources[GPU]; !ok {
t.resources[GPU] = &aurora.Resource{}
t.task.Resources = append(t.task.Resources, t.resources[GPU])
}
t.resources[GPU].NumGpus = &gpu
return t
}