Fixing GPU resource to only be added if specified since Aurora scheduler by default will reject tasks containing GPU.
This commit is contained in:
parent
8d67d8c2f3
commit
9b3593e9d9
1 changed files with 9 additions and 7 deletions
16
task.go
16
task.go
|
@ -47,17 +47,12 @@ func NewTask() *AuroraTask {
|
||||||
numCpus := &aurora.Resource{}
|
numCpus := &aurora.Resource{}
|
||||||
ramMb := &aurora.Resource{}
|
ramMb := &aurora.Resource{}
|
||||||
diskMb := &aurora.Resource{}
|
diskMb := &aurora.Resource{}
|
||||||
numGpus := &aurora.Resource{}
|
|
||||||
|
|
||||||
numCpus.NumCpus = new(float64)
|
numCpus.NumCpus = new(float64)
|
||||||
ramMb.RamMb = new(int64)
|
ramMb.RamMb = new(int64)
|
||||||
diskMb.DiskMb = new(int64)
|
diskMb.DiskMb = new(int64)
|
||||||
|
|
||||||
resources := make(map[ResourceType]*aurora.Resource)
|
resources := map[ResourceType]*aurora.Resource{CPU: numCpus, RAM: ramMb, DISK: diskMb}
|
||||||
resources[CPU] = numCpus
|
|
||||||
resources[RAM] = ramMb
|
|
||||||
resources[DISK] = diskMb
|
|
||||||
resources[GPU] = numGpus
|
|
||||||
|
|
||||||
return &AuroraTask{task: &aurora.TaskConfig{
|
return &AuroraTask{task: &aurora.TaskConfig{
|
||||||
Job: &aurora.JobKey{},
|
Job: &aurora.JobKey{},
|
||||||
|
@ -247,7 +242,14 @@ func (t *AuroraTask) Disk(disk int64) *AuroraTask {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *AuroraTask) GPU(gpu 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
|
return t
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue