Fix JobSet resource requirements

Signed-off-by: Antonin Stefanutti <[email protected]>
kubeflow · Feb 14, 2025 · 03d6f54 · 03d6f54
1 parent 4d7f02a
commit 03d6f54
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/pkg/runtime/framework/plugins/jobset/builder.go b/pkg/runtime/framework/plugins/jobset/builder.go
@@ -115,11 +115,17 @@ func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Build
 						if args := trainJob.Spec.Trainer.Args; args != nil {
 							b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Args = args
 						}
-						if resourcesPerNode := trainJob.Spec.Trainer.ResourcesPerNode; resourcesPerNode != nil {
+						if resourcesPerNode := trainJob.Spec.Trainer.ResourcesPerNode; resourcesPerNode != nil &&
+							(resourcesPerNode.Limits != nil || resourcesPerNode.Requests != nil) {
+							requirements := corev1ac.ResourceRequirements()
+							if limits := resourcesPerNode.Limits; limits != nil {
+								requirements.WithLimits(limits)
+							}
+							if requests := resourcesPerNode.Requests; requests != nil {
+								requirements.WithRequests(requests)
+							}
 							b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].
-								WithResources(corev1ac.ResourceRequirements().
-									WithRequests(resourcesPerNode.Requests).
-									WithLimits(resourcesPerNode.Limits))
+								WithResources(requirements)
 						}
 					}
 					// Update values from the Info object.