feat(api): implement DRA support for TrainJob (KEP-2782) - #4027
feat(api): implement DRA support for TrainJob (KEP-2782)#4027Sridhar1030 wants to merge 2 commits into
Conversation
Implement Phase 1 of the approved KEP from kubeflow#3540. - Add trainer.resourceClaimsPerNode and runtimePatches DRA fields - Wire pod-level resourceClaims and node container resources.claims - Resolve DRA GPU count for numProcPerNode in torch/MPI/XGBoost/Flux plugins - Gate DRA behind DynamicResourceAllocation alpha feature flag - Add RBAC, validation, admission warning, and operator docs SDK changes from the KEP remain a follow-up in the Kubeflow SDK repo. E2E tests with a DRA-capable cluster are deferred per the KEP. Refs: kubeflow#3540 Closes: kubeflow#2782 Signed-off-by: Sridhar1030 <sridharpillai75@gmail.com>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
Sridhar1030
changed the title
|
One deviation from the KEP: this PR puts DRA behind a |
Signed-off-by: Sridhar1030 <sridharpillai75@gmail.com>
|
/retest |
robert-bell
left a comment
robert-bell
left a comment
There was a problem hiding this comment.
Thanks @Sridhar1030!
I've left some initial comments. I'll follow up with a more detailed review next week.
| // The controller adds these claims to the trainer node Pod's resourceClaims and wires | ||
| // container-level resources.claims on the node container. To attach a claim to other |
There was a problem hiding this comment.
I think we can document the outcome. Wdyt?
| // The controller adds these claims to the trainer node Pod's resourceClaims and wires | |
| // container-level resources.claims on the node container. To attach a claim to other | |
| // These claims are added to the trainer node Pod's resourceClaims and automatically | |
| // referenced in the node container's resources.claims. To attach a claim to another |
| // resourceClaimTemplateName is the name of a ResourceClaimTemplate in the TrainJob namespace. | ||
| // A separate ResourceClaim is created from it for every training node Pod. | ||
| // +kubebuilder:validation:MinLength=1 | ||
| // +kubebuilder:validation:MaxLength=253 |
There was a problem hiding this comment.
Are these max lengths the same as the upstream DRA?
| // +optional | ||
| Volumes []corev1.Volume `json:"volumes,omitempty"` | ||
|
|
||
| // resourceClaims patches the Pod's resourceClaims. Containers consume a claim by |
There was a problem hiding this comment.
Can we document that this is for targeting sidecar and init containers? And that users should prefer the resourceClaimsPerNode?
| VolumeMounts []corev1.VolumeMount `json:"volumeMounts,omitempty"` | ||
|
|
||
| // resources patches the container's compute resources, including the resources.claims | ||
| // that reference the Pod's resourceClaims. For the node container, trainer.resourcesPerNode |
There was a problem hiding this comment.
Same here- point users to the top level api.
|
|
||
| func (r *TrainingRuntime) newRuntimeInfo( | ||
| trainJob *trainer.TrainJob, jobSetTemplateSpec trainer.JobSetTemplateSpec, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy, | ||
| ctx context.Context, trainJob *trainer.TrainJob, jobSetTemplateSpec trainer.JobSetTemplateSpec, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy, |
There was a problem hiding this comment.
Why do we need ctx here? I'd have expected we'd only need to look up ResourceClaimTemplates when resolving the gpu count.
Summary
Implements Phase 1 of KEP-2782, approved in #3540.
spec.trainer.resourceClaimsPerNodeso users can request DRA devices at the same level asresourcesPerNoderesources.claimson thenodecontainerruntimePatcheswithPodSpecPatch.resourceClaimsandContainerPatch.resourcesfor sidecars/init containersResourceClaimTemplatefornumProcPerNodeauto-detection in torch, torchtune, MPI, XGBoost, and FluxresourcesPerNode.claimsvia CRD CEL validation; validate dangling container claims in the JobSet pluginDynamicResourceAllocationalpha feature gate (defaultfalse), RBAC forresourceclaimtemplates/resourceclaims, and operator docsOut of scope for this PR: Kubeflow SDK changes (
list_resource_claim_templates, runtime/step introspection) and DRA E2E in CI — both called out in the KEP as follow-up/deferred work.Motivation
Kubernetes DRA is GA in 1.34+. Trainer users today must nest DRA claims under deep
runtimePatcheswhile GPU counts live inresourcesPerNode. This KEP unifies the UX with a top-levelresourceClaimsPerNodefield and keeps GPU auto-detection working without extended resources.Test plan
make test— unit tests for runtime, plugins, webhooksmake test-integration— controller DRA integration tests + webhook validationResourceClaimTemplate→ TrainJob withresourceClaimsPerNoderesourceClaims+nodecontainerresources.claimsPET_NPROC_PER_NODE=autowhen RCT resolves 2 GPUsResourceClaimreachesallocated,reserved; podRunningRelated