KEP-2628: Support KAI Scheduler in Kubeflow Trainer #2663
Conversation
Signed-off-by: Harshal292004 <malaniharshal95@gmail.com>
|
@Harshal292004: GitHub didn't allow me to request PR reviews from the following users: romanbaron, kannon92. Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs. In response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository. |
|
[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.
Needs approval from an approver in each of these files:
Approvers can indicate their approval by writing |
Pull Request Test Coverage Report for Build 15726042285Details
💛 - Coveralls |
Signed-off-by: Harshal292004 <malaniharshal95@gmail.com>
| @@ -0,0 +1,89 @@ | |||
| # KEP-2628: Support KAI Scheduler in Kubeflow Trainer | |||
|
|
|||
| This KEP proposes integrating NVIDIA's KAI Scheduler into Kubeflow Trainer V2 to enable gang-scheduling capabilities for TrainJob resources, extending the existing PodGroupPolicy API to support KAI alongside current schedulers like Co-Scheduling. The integration will leverage KAI's PodGrouper service to create scheduling queues and apply appropriate labels for efficient resource allocation in AI workloads. | |||
There was a problem hiding this comment.
Small comment here: PodGrouper Service will create the podGroup object, queues will have to be created by a different entity, probably by a trainer plugin.
This point is perfectly explained below so the rest is accurate.
There was a problem hiding this comment.
I have resolved this , you may have a look @romanbaron
| } | ||
|
|
||
| type KAIConfig struct { | ||
| Queue string `json:"queue,omitempty"` |
There was a problem hiding this comment.
This is going to be really confusing to someone who uses KAI + Kueue.
cc @romanbaron.
Is there anyway to use gang scheduling in KAI without needing a KAI queue?
Volcano is used for PodGroup functionality in Trainer but I don't think one has to use Volcano Queues to use this.
There was a problem hiding this comment.
Each KAI PodGroup must reference a designated KAI scheduling queue to clarify which queue’s resources are allocated to it.
One approach to integrating KAI with Kueue is to configure a single, unlimited queue for KAI, allowing Kueue to govern resources and dispatch workloads. Under this model, KAI schedules only the workloads that have been dispatched (i.e., pods are created and ungated), while fairness enforcement remains outside its scope.
@kannon92 - does it sounds like a possible solution? If not, we can discuss the possible options further.
There was a problem hiding this comment.
I think that makes sense to me.
Signed-off-by: Harshal292004 <malaniharshal95@gmail.com>
There was a problem hiding this comment.
@Harshal292004, starting from v0.6.0 release KAI-Scheduler uses a different label key for queue, now it is kai.scheduler/queue instead of runai/queue.
Some more details are available here:
https://github.com/NVIDIA/KAI-Scheduler/tree/main/docs/migrationguides/v0.6.0#scheduling-queue-label-key
I think it will be best to align to the new version from now on.
Signed-off-by: Harshal292004 <malaniharshal95@gmail.com>
|
@romanbaron Does this make sense ?: type PodGroupPolicySource struct {
Coscheduling *CoschedulingPodGroupPolicySource `json:"coscheduling,omitempty"`
Kaischeduling *KaischedulingPodGroupPolicySource `json:"kaischeduling,omitempty"`
}
type KaischedulingPodGroupPolicySource struct {
MinMember *int32 `json:"minMember,omitempty"`
Queue *string `json:"queue,omitempty"`
PriorityClassName *string `json:"priorityClassName,omitempty"`
MarkUnschedulable *bool `json:"markUnschedulable,omitempty"`
SchedulingBackoff *int32 `json:"schedulingBackoff,omitempty"`
} |
This is the Kubeflow Enhancement Proposal for supporting KAI Scheduler in Kubeflow Trainer V2:
Related: #2628
Any suggestions or feedback are welcome !
@romanbaron @kannon92
/cc @andreyvelich @tenzen-y @Electronic-Waste