Skip to content

Commit dda5ba8

Browse files
romanbaronromanbaron
authored
Design doc for hierarchical PodGroup structure (#324)
Design doc for hierarchical PodGroup structure
1 parent 5702f4e commit dda5ba8

File tree

1 file changed

+259
-0
lines changed
  • docs/developer/designs/hierarchical-podgroup

1 file changed

+259
-0
lines changed

‎docs/developer/designs/hierarchical-podgroup/README.md‎

Lines changed: 259 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,259 @@
1+
# Hierarchical PodGroup structure
2+
3+
## Overview
4+
This document outlines an enhancement to the PodGroup architecture by introducing SubGroups within a PodGroup, enabling fine-grained control over gang scheduling and topology-aware placement for subsets of pods within a workload. While the PodGroup continues to enforce gang scheduling for the workload as a whole, this extension allows workloads to specify minMember and topology constraints on individual SubGroups, providing greater resolution over how smaller groups of pods are placed and scheduled. For example, a SubGroup can specify that only a subset of its pods must be scheduled together (minMember) and define specific network topology constraints aligned with that subset’s requirements. This design preserves atomic workload scheduling semantics while enabling advanced, granular placement policies that align closely with topology-aware scheduling strategies.
5+
6+
## Motivation
7+
Modern inference workloads consist of components with varying network sensitivity and resource requirements, necessitating the ability to express fine-grained scheduling and placement preferences within the context of gang-scheduled workloads. Uniform constraints at the PodGroup level may impose overly restrictive placement requirements, failing to capture the nuanced needs of distinct workload components. By introducing SubGroups within a PodGroup, workloads can specify per-SubGroup minMember requirements and topology constraints while maintaining overall gang scheduling semantics. This enables advanced placement strategies, improves scheduling flexibility and efficiency, and ensures differentiated placement and subgroup-level gang constraints can be expressed without compromising the workload’s atomic scheduling guarantees.
8+
9+
## Goals
10+
- Enable fine-grained scheduling specification: Allow PodGroups to define how smaller SubGroups of pods should be scheduled, specifying minMember and topology constraints per SubGroup.
11+
- Preserve atomicity while enabling flexibility: Maintain gang scheduling semantics at the workload level while enabling per-SubGroup constraints for granular control over subsets of pods.
12+
- Support advanced topology-aware placement: Allow SubGroups within a PodGroup to define specific topology constraints for precise placement within the cluster.
13+
- Enhance workload-aligned resource utilization: Improve cluster utilization by supporting differentiated placement strategies without imposing overly restrictive uniform constraints.
14+
- Maintain backward compatibility: Ensure existing workloads using PodGroups without SubGroups continue to function without modification.
15+
16+
## Non-Goals
17+
- Partial Workload Execution: This design does not support executing subsets of pods independently before the PodGroup’s gang scheduling conditions are met.
18+
- Topology Discovery and Optimization: The design does not include enhancements for topology discovery or automated topology optimization.
19+
- Changes to Pod-Level Scheduling: This enhancement does not alter pod-level scheduling mechanisms beyond the SubGroup minMember and topology constraints applied within the PodGroup context.
20+
21+
## User Stories
22+
23+
### User Story 1: Fine-Grained Gang Scheduling
24+
I want to specify a PodGroup requiring 10 pods to be scheduled atomically, while defining that the decode-workers SubGroup requires at least 4 pods to be scheduled as a gang, enabling better resource alignment for my workload’s architecture.
25+
26+
### User Story 2: SubGroup-Specific Topology Constraints
27+
I want one SubGroup to be placed within the same rack for low latency, while allowing another SubGroup pods to be placed flexibly across the cluster, aligning network-sensitive components with workload requirements.
28+
29+
### User Story 3: Multi-SubGroup Workload with Co-Located Leaders and Workers
30+
I want to define decode-leaders, decode-workers, prefill-leaders, and prefill-workers SubGroups where leaders and workers are co-located within each group, but prefill and decode groups may be scheduled in different locations to optimize resource placement.
31+
32+
### User Story 4: Backward Compatibility
33+
I want my existing PodGroups without SubGroups to continue working without modification, ensuring a seamless migration path for advanced scheduling capabilities.
34+
35+
## Design Proposal
36+
This proposal extends the PodGroup CRD and scheduling flow to introduce SubGroups within a PodGroup, enabling workloads to define fine-grained gang scheduling (minMember) and topology constraints for subsets of pods while preserving gang scheduling semantics at the workload level.
37+
38+
## Key Concepts
39+
- PodGroup represents a workload scheduled atomically using gang scheduling.
40+
- The PodGroup is divided into distinct SubGroups, enabling structured subdivision of the workload while preserving atomic scheduling boundaries.
41+
- Scheduling Semantics:
42+
- The PodGroup is scheduled atomically (gang scheduling) to guarantee coordinated workload orchestration and execution consistency.
43+
- SubGroups enable fine-grained placement and policy specification while maintaining atomicity; partial execution of the workload is not permitted.
44+
- SubGroup - A logical subset within a PodGroup, enabling scoped placement and scheduling requirements:
45+
- name: Unique identifier within the parent PodGroup.
46+
- minMember: Specifies the minimum number of pods required within the SubGroup to satisfy scheduling constraints.
47+
- Pods are assigned to SubGroups with `kai.scheduler/subgroup-name` label, where the value is the name of the respective SubGroup.
48+
- TopologyConstraints – Provides hierarchical placement control across three levels:
49+
- A global constraint applied to all pods in the PodGroup when no more specific constraint is defined.
50+
- Specific constraints applied to explicitly named SubGroups for targeted placement control.
51+
- Shared constraints applied to sets of SubGroups collectively, enabling coordinated placement policies across related SubGroups.
52+
53+
## API Changes
54+
To support SubGroups within a PodGroup, the PodGroupSpec API is extended as follows:
55+
```go
56+
// PodGroupSpec defines the desired state of a PodGroup, representing a workload to be scheduled as a gang.
57+
type PodGroupSpec struct {
58+
// MinMember defines the minimal number of members/tasks to run the pod group;
59+
// if there's not enough resources to start all tasks, the scheduler will not start anyone.
60+
MinMember int32 `json:"minMember,omitempty"`
61+
62+
// Queue defines the queue to allocate resource for PodGroup; if queue does not exist,
63+
// the PodGroup will not be scheduled.
64+
Queue string `json:"queue,omitempty"`
65+
66+
// If specified, indicates the PodGroup's priority. "system-node-critical" and
67+
// "system-cluster-critical" are two special keywords which indicate the
68+
// highest priorities with the former being the highest priority. Any other
69+
// name must be defined by creating a PriorityClass object with that name.
70+
// If not specified, the PodGroup priority will be default or zero if there is no
71+
// default.
72+
// +optional
73+
PriorityClassName string `json:"priorityClassName,omitempty"`
74+
75+
// Should add "Unschedulable" event to the pods or not.
76+
MarkUnschedulable *bool `json:"markUnschedulable,omitempty"`
77+
78+
// The number of scheduling cycles to try before marking the pod group as UnschedulableOnNodePool. Currently only supporting -1 and 1
79+
SchedulingBackoff *int32 `json:"schedulingBackoff,omitempty"`
80+
81+
// TopologyConstraints define the topology constraints for this PodGroup
82+
TopologyConstraints TopologyConstraints `json:"topologyConstraints,omitempty"`
83+
84+
// SubGroups defines finer-grained subsets of pods within the PodGroup with individual scheduling constraints
85+
SubGroups []SubGroup `json:"subGroups,omitempty"`
86+
}
87+
88+
type SubGroup struct {
89+
// Name uniquely identifies the SubGroup within the parent PodGroup.
90+
Name string `json:"name"`
91+
92+
// MinMember defines the minimal number of members/tasks to run this SubGroup;
93+
// if there's not enough resources to start all tasks, the scheduler will not start anyone.
94+
MinMember int32 `json:"minMember,omitempty"`
95+
}
96+
97+
// TopologyConstraints defines topology constraints at group, subgroup, and subgroup-set levels.
98+
type TopologyConstraints struct {
99+
// Global applies the constraint to all pods in the PodGroup if no more specific subgroup constraint applies.
100+
Global *TopologyConstraint `json:"global,omitempty"`
101+
102+
// SubGroups defines topology constraints for specific named subgroups.
103+
// The key is the subgroup name, and the constraint applies to all pods in that subgroup.
104+
SubGroups map[string]TopologyConstraint `json:"subGroups,omitempty"`
105+
106+
// SubGroupSets allows defining constraints that apply to multiple subgroups collectively.
107+
// Each entry specifies a set of subgroup names to which the constraint will be applied.
108+
// Useful for scenarios where several subgroups need to share the same topology policy.
109+
SubGroupSets []SubGroupSetTopologyConstraint `json:"subGroupSets,omitempty"`
110+
}
111+
112+
// SubGroupSetTopologyConstraint defines a topology constraint for a set of subgroup names.
113+
type SubGroupSetTopologyConstraint struct {
114+
// SubGroups is the list of subgroup names that this constraint applies to collectively.
115+
SubGroups []string `json:"subGroups"`
116+
117+
// Constraint defines the topology constraint to apply to the listed subgroups.
118+
Constraint TopologyConstraint `json:"constraint"`
119+
}
120+
121+
type TopologyConstraint struct {
122+
// PreferredTopologyLevel defines the preferred level in the topology hierarchy
123+
// that this constraint applies to (e.g., "rack", "zone", "datacenter").
124+
// Jobs will be scheduled to maintain locality at this level when possible.
125+
PreferredTopologyLevel string `json:"preferredTopologyLevel,omitempty"`
126+
127+
// RequiredTopologyLevel defines the maximal level in the topology hierarchy
128+
// that all pods must be scheduled within.
129+
// If set, all pods of the job must be scheduled within a single domain at this level.
130+
RequiredTopologyLevel string `json:"requiredTopologyLevel,omitempty"`
131+
132+
// Topology specifies the name of the topology CRD that defines the
133+
// physical layout to use for this constraint. This allows for supporting
134+
// multiple different topology configurations in the same cluster.
135+
Topology string `json:"topology,omitempty"`
136+
}
137+
```
138+
139+
### Validation
140+
The following validations will be enforced via a Validating Webhook:
141+
- Unique SubGroup name validation - ensure that all SubGroups within a PodGroup have unique names, preventing conflicts and enabling reliable hierarchical processing during scheduling.
142+
- Validate that, if SubGroups are defined, the PodGroup’s global minMember equals the sum of all SubGroups’ minMember values to ensure scheduling consistency and atomicity.
143+
- SubGroup membership validation - the Validating Webhook can enforce that each pod associated with the workload is assigned to exactly one SubGroup and each SubGroup is part in a single SubGroupSet, ensuring clear, non-overlapping SubGroup membership across all pods within the PodGroup.
144+
145+
## Examples
146+
147+
### Example 1: PodGroup with Two SubGroups
148+
This example demonstrates a PodGroup with two SubGroups, each defining distinct minMember requirements and enforcing rack-level topology constraints to ensure efficient, localized scheduling.
149+
```mermaid
150+
graph TD
151+
PG[PodGroup]
152+
153+
Decode[decoders<br>minMember: 4<br>Topology: rack]
154+
Prefill[prefills<br>minMember: 1<br>Topology: rack]
155+
156+
PG --> Decode
157+
PG --> Prefill
158+
```
159+
```yaml
160+
spec:
161+
minMember: 5 # 4 (decoders) + 1 (prefills), ensuring the gang scheduling invariant
162+
subGroups:
163+
- name: decoders
164+
minMember: 4
165+
166+
- name: prefills
167+
minMember: 1
168+
169+
topologyConstraints:
170+
subGroups:
171+
decoders:
172+
requiredTopologyLevel: rack
173+
prefills:
174+
requiredTopologyLevel: rack
175+
```
176+
177+
### Example 2: Multi-SubGroup Leaders and Workers
178+
This example illustrates a hierarchical PodGroup structure with multiple SubGroups representing leaders and workers, applying both rack-level and block-level topology constraints to achieve co-located placement within each logical group while allowing separation between groups.
179+
```mermaid
180+
graph TD
181+
PG[PodGroup]
182+
183+
Decode[decode<br>Topology: block]
184+
DecodeWorkers[decode-workers<br>minMember: 4<br>Topology: rack]
185+
DecodeLeaders[decode-leaders<br>minMember: 1<br>Topology: rack]
186+
187+
Prefill[prefill<br>Topology: block]
188+
PrefillWorkers[prefill-workers<br>minMember: 4<br>Topology: rack]
189+
PrefillLeaders[prefill-leaders<br>minMember: 1<br>Topology: rack]
190+
191+
PG --> Decode
192+
Decode --> DecodeWorkers
193+
Decode --> DecodeLeaders
194+
195+
PG --> Prefill
196+
Prefill --> PrefillWorkers
197+
Prefill --> PrefillLeaders
198+
```
199+
```yaml
200+
spec:
201+
minMember: 10
202+
subGroups:
203+
- name: decode-workers
204+
minMember: 4
205+
206+
- name: decode-leaders
207+
minMember: 1
208+
209+
- name: prefill-workers
210+
minMember: 4
211+
212+
- name: prefill-leaders
213+
minMember: 1
214+
215+
topologyConstraints:
216+
subGroups:
217+
decode-workers:
218+
requiredTopologyLevel: rack
219+
decode-leaders:
220+
requiredTopologyLevel: rack
221+
prefill-workers:
222+
requiredTopologyLevel: rack
223+
prefill-leaders:
224+
requiredTopologyLevel: rack
225+
226+
subGroupSets:
227+
- subGroups:
228+
- decode-workers
229+
- decode-leaders
230+
constraint:
231+
requiredTopologyLevel: block
232+
- subGroups:
233+
- prefill-workers
234+
- prefill-leaders
235+
constraint:
236+
requiredTopologyLevel: block
237+
```
238+
239+
## Development
240+
To ensure a controlled and backward-compatible rollout of PodGroup SubGroups with fine-grained scheduling, the following phased development plan is proposed:
241+
242+
### Phase 1: API Definition and Validation
243+
- Extend the PodGroup [CRD](https://github.com/NVIDIA/KAI-Scheduler/blob/main/pkg/apis/scheduling/v2alpha2/podgroup_types.go) to support SubGroups with name and minMember.
244+
- (Optional) Implement optional validation to ensure:
245+
- Unique SubGroup names within a PodGroup.
246+
- minMember consistency with global level.
247+
- Every SubGroup belongs to a single SubGroupSets in topology constraints.
248+
- Guarantee backward compatibility with existing PodGroups that do not define SubGroups.
249+
250+
### Phase 2: Grove Plugin Integration
251+
Update the PodGrouper Grove plugin to:
252+
- Parse Grove CRDs defining hierarchical PodClique structures.
253+
- Construct matching PodGroup CRDs with SubGroups, mapping PodCliques to SubGroups and topology constraints.
254+
255+
### Phase 3: Scheduler Algorithm Adjustments
256+
- Extend the scheduler to support multi-level PodGroup scheduling:
257+
- Enforce constraints hierarchically, validating higher-level PodGroup constraints before evaluating subordinate SubGroup constraints.
258+
- Proceed hierarchically, ensuring constraints are met before advancing to deeper levels.
259+
- Preserve gang scheduling semantics at the PodGroup level while enabling fine-grained SubGroup placement.

0 commit comments

Comments
 (0)