Skip to content

test: add TAS simple level and constraint tests #349

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Ronkahn21 merged 11 commits into from
Jan 26, 2026
Merged

test: add TAS simple level and constraint tests #349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
c3d90fb
test: add TAS e2e test infrastructure and basic tests
Ronkahn21 Jan 18, 2026
114c7a9
test: add TAS simple level and constraint tests
Ronkahn21 Jan 18, 2026
19e460f
refactor: apply CR review fixes to TAS tests
Ronkahn21 Jan 21, 2026
2850010
fix: correct PCSG SubGroup expectations for PR #357
Ronkahn21 Jan 21, 2026
8f8dd09
refactor: align TAS1/TAS2 with upstream and revert helper changes
Ronkahn21 Jan 21, 2026
17de004
test: improve TAS6 naming and add helper functions
Ronkahn21 Jan 24, 2026
3feb71d
refactor: update PodGroup retrieval to use dynamic workload names
Ronkahn21 Jan 24, 2026
2edc57d
refactor: rename YAML file and update workload configuration for TAS
Ronkahn21 Jan 24, 2026
2cd477f
feat: add utility functions for PodGroup retrieval and topology verif…
Ronkahn21 Jan 25, 2026
c5bab0b
refactor: extract naming logic into utility functions for PodClique s…
Ronkahn21 Jan 25, 2026
d232b22
refactor: update GetBasePodGangName to use PCS terminology
Ronkahn21 Jan 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions operator/e2e/tests/setup.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,10 @@ const (
defaultPollTimeout = 4 * time.Minute
// defaultPollInterval is the interval for most polling conditions
defaultPollInterval = 5 * time.Second

// Grove label keys
LabelPodClique = "grove.io/podclique"
LabelPodCliqueScalingGroup = "grove.io/podcliquescalinggroup"
)

// TestContext holds common test parameters that are shared across many utility functions.
Expand Down
350 changes: 296 additions & 54 deletions operator/e2e/tests/topology_test.go
View file
Open in desktop

Large diffs are not rendered by default.

90 changes: 85 additions & 5 deletions operator/e2e/utils/kai_topology.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,14 @@
package utils

import (
"fmt"

kaischedulingv2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"

"context"
"fmt"
"time"

kaischedulingv2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

"k8s.io/client-go/dynamic"
"k8s.io/utils/ptr"
)

// ExpectedSubGroup defines the expected structure of a KAI PodGroup SubGroup for verification
Expand All @@ -40,6 +38,43 @@ type ExpectedSubGroup struct {
PreferredTopologyLevel string
}

// CreateExpectedStandalonePCLQSubGroup creates an ExpectedSubGroup for a standalone PodClique (not in PCSG)
// Name format: <pcs-name>-<pcs-replica>-<clique-name>
func CreateExpectedStandalonePCLQSubGroup(pcsName string, pcsReplica int, cliqueName string, minMember int32, topologyLevel string) ExpectedSubGroup {
name := GetStandalonePCLQSubGroupName(pcsName, pcsReplica, cliqueName)
return ExpectedSubGroup{
Name: name,
MinMember: minMember,
Parent: nil,
RequiredTopologyLevel: topologyLevel,
}
}

// CreateExpectedPCSGParentSubGroup creates an ExpectedSubGroup for a PCSG parent (scaling group replica)
// Name format: <pcs-name>-<pcs-replica>-<sg-name>-<sg-replica>
func CreateExpectedPCSGParentSubGroup(pcsName string, pcsReplica int, sgName string, sgReplica int, topologyLevel string) ExpectedSubGroup {
name := GetPCSGParentSubGroupName(pcsName, pcsReplica, sgName, sgReplica)
return ExpectedSubGroup{
Name: name,
MinMember: 0,
Parent: nil,
RequiredTopologyLevel: topologyLevel,
}
}

// CreateExpectedPCLQInPCSGSubGroup creates an ExpectedSubGroup for a PodClique within a PCSG
// Name format: <pcs-name>-<pcs-replica>-<sg-name>-<sg-replica>-<clique-name>
func CreateExpectedPCLQInPCSGSubGroup(pcsName string, pcsReplica int, sgName string, sgReplica int, cliqueName string, minMember int32, topologyLevel string) ExpectedSubGroup {
name := GetPCLQInPCSGSubGroupName(pcsName, pcsReplica, sgName, sgReplica, cliqueName)
parentName := GetPCSGParentSubGroupName(pcsName, pcsReplica, sgName, sgReplica)
return ExpectedSubGroup{
Name: name,
MinMember: minMember,
Parent: ptr.To(parentName),
RequiredTopologyLevel: topologyLevel,
}
}

// GetKAIPodGroupsForPCS retrieves all KAI PodGroups for a given PodCliqueSet by label selector
// KAI scheduler creates PodGroups with label: app.kubernetes.io/part-of=<pcs-name>
// Returns a list of PodGroups that tests can filter by owner reference if needed
Expand Down Expand Up @@ -179,3 +214,48 @@ func VerifyKAIPodGroupSubGroups(podGroup *kaischedulingv2alpha2.PodGroup, expect
logger.Infof("KAI PodGroup %s verified with %d SubGroups", podGroup.Name, len(expectedSubGroups))
return nil
}

// GetPodGroupForBasePodGangReplica retrieves the KAI PodGroup of the corresponding PodGang
// which is the base PodGang of specific PodGangSet replica.
// For a PodGangSet workload "my-workload", replica 0's base PodGang is "my-workload-0".
func GetPodGroupForBasePodGangReplica(
ctx context.Context,
dynamicClient dynamic.Interface,
namespace string,
workloadName string,
pgsReplica int,
timeout time.Duration,
interval time.Duration,
logger *Logger,
) (*kaischedulingv2alpha2.PodGroup, error) {
podGroups, err := WaitForKAIPodGroups(ctx, dynamicClient, namespace, workloadName, timeout, interval, logger)
if err != nil {
return nil, fmt.Errorf("failed to get KAI PodGroups: %w", err)
}

basePodGangName := GetBasePodGangName(workloadName, pgsReplica)
basePodGroup, err := FilterPodGroupByOwner(podGroups, basePodGangName)
if err != nil {
return nil, fmt.Errorf("failed to find PodGroup for PodGang %s: %w", basePodGangName, err)
}

return basePodGroup, nil
}

// VerifyPodGroupTopology verifies both top-level topology constraint and SubGroups structure.
func VerifyPodGroupTopology(
podGroup *kaischedulingv2alpha2.PodGroup,
requiredLevel, preferredLevel string,
expectedSubGroups []ExpectedSubGroup,
logger *Logger,
) error {
if err := VerifyKAIPodGroupTopologyConstraint(podGroup, requiredLevel, preferredLevel, logger); err != nil {
return fmt.Errorf("top-level constraint verification failed: %w", err)
}

if err := VerifyKAIPodGroupSubGroups(podGroup, expectedSubGroups, logger); err != nil {
return fmt.Errorf("SubGroups verification failed: %w", err)
}

return nil
}
45 changes: 45 additions & 0 deletions operator/e2e/utils/naming.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
//go:build e2e

package utils

// /*
// Copyright 2025 The Grove Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// */

import "fmt"

// GetBasePodGangName constructs the base PodGang name for a specific PCS replica.
// Format: <pgs-name>-<replica-index>
func GetBasePodGangName(workloadName string, pcsReplica int) string {
return fmt.Sprintf("%s-%d", workloadName, pcsReplica)
}

// GetStandalonePCLQSubGroupName constructs the SubGroup name for a standalone PodClique.
// Format: <pcs-name>-<pcs-replica>-<clique-name>
func GetStandalonePCLQSubGroupName(pcsName string, pcsReplica int, cliqueName string) string {
return fmt.Sprintf("%s-%d-%s", pcsName, pcsReplica, cliqueName)
}

// GetPCSGParentSubGroupName constructs the SubGroup name for a PCSG parent (scaling group replica).
// Format: <pcs-name>-<pcs-replica>-<sg-name>-<sg-replica>
func GetPCSGParentSubGroupName(pcsName string, pcsReplica int, sgName string, sgReplica int) string {
return fmt.Sprintf("%s-%d-%s-%d", pcsName, pcsReplica, sgName, sgReplica)
}

// GetPCLQInPCSGSubGroupName constructs the SubGroup name for a PodClique within a PCSG.
// Format: <pcs-name>-<pcs-replica>-<sg-name>-<sg-replica>-<clique-name>
func GetPCLQInPCSGSubGroupName(pcsName string, pcsReplica int, sgName string, sgReplica int, cliqueName string) string {
return fmt.Sprintf("%s-%d-%s-%d-%s", pcsName, pcsReplica, sgName, sgReplica, cliqueName)
}
21 changes: 21 additions & 0 deletions operator/e2e/utils/topology.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,3 +182,24 @@ func VerifyPodsInSameTopologyDomain(ctx context.Context, clientset kubernetes.In
logger.Infof("Verified %d pods are in same topology domain %s=%s", len(pods), topologyKey, expectedValue)
return nil
}

// VerifyLabeledPodsInTopologyDomain filters pods by label, verifies count, and checks topology domain.
func VerifyLabeledPodsInTopologyDomain(
ctx context.Context,
clientset kubernetes.Interface,
allPods []v1.Pod,
labelKey, labelValue string,
expectedCount int,
topologyKey string,
logger *Logger,
) error {
filteredPods := FilterPodsByLabel(allPods, labelKey, labelValue)
if len(filteredPods) != expectedCount {
return fmt.Errorf(
"expected %d pods with %s=%s, got %d",
expectedCount, labelKey, labelValue, len(filteredPods),
)
}

return VerifyPodsInSameTopologyDomain(ctx, clientset, filteredPods, topologyKey, logger)
}
44 changes: 44 additions & 0 deletions operator/e2e/yaml/tas-host-level.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
# Workload: Host-Level Packing - PCS with single clique at host level
# Test scenario: PCS with single clique constrained to host level
---
apiVersion: grove.io/v1alpha1
kind: PodCliqueSet
metadata:
name: tas-host-level
labels:
app: tas-host-level
spec:
replicas: 1
template:
cliques:
- name: worker
labels:
kai.scheduler/queue: test
topologyConstraint:
packDomain: host
spec:
roleName: worker
replicas: 2
minAvailable: 2
podSpec:
schedulerName: kai-scheduler
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node_role.e2e.grove.nvidia.com
operator: In
values:
- agent
tolerations:
- key: node_role.e2e.grove.nvidia.com
operator: Equal
value: agent
effect: NoSchedule
containers:
- name: worker
image: registry:5001/nginx:alpine-slim
resources:
requests:
memory: 30Mi
48 changes: 48 additions & 0 deletions operator/e2e/yaml/tas-no-constraint.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# Workload: SL3 - No Topology Constraints
# Test scenario: No constraints at PCS, PCSG, or PCLQ levels (pure gang scheduling)
---
apiVersion: grove.io/v1alpha1
kind: PodCliqueSet
metadata:
name: tas-no-constraint
labels:
app: tas-no-constraint
spec:
replicas: 1
template:
podCliqueScalingGroups:
- name: workers
replicas: 2
minAvailable: 2
cliqueNames:
- worker
cliques:
- name: worker
labels:
kai.scheduler/queue: test
spec:
roleName: worker
replicas: 2
minAvailable: 2
podSpec:
schedulerName: kai-scheduler
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node_role.e2e.grove.nvidia.com
operator: In
values:
- agent
tolerations:
- key: node_role.e2e.grove.nvidia.com
operator: Equal
value: agent
effect: NoSchedule
containers:
- name: worker
image: registry:5001/nginx:alpine-slim
resources:
requests:
memory: 30Mi
79 changes: 79 additions & 0 deletions operator/e2e/yaml/tas-sl-pcs-only.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# Workload: PCS-Only Topology - PCS with rack constraint
# Test scenario: PCS with rack constraint, PCSG without constraint, standalone clique
---
apiVersion: grove.io/v1alpha1
kind: PodCliqueSet
metadata:
name: tas-sl-pcs-only
labels:
app: tas-sl-pcs-only
spec:
replicas: 1
template:
topologyConstraint:
packDomain: rack
podCliqueScalingGroups:
- name: workers
replicas: 2
minAvailable: 2
cliqueNames:
- worker
cliques:
- name: worker
labels:
kai.scheduler/queue: test
spec:
roleName: worker
replicas: 1
minAvailable: 1
podSpec:
schedulerName: kai-scheduler
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node_role.e2e.grove.nvidia.com
operator: In
values:
- agent
tolerations:
- key: node_role.e2e.grove.nvidia.com
operator: Equal
value: agent
effect: NoSchedule
containers:
- name: worker
image: registry:5001/nginx:alpine-slim
resources:
requests:
memory: 30Mi
- name: router
labels:
kai.scheduler/queue: test
spec:
roleName: router
replicas: 2
minAvailable: 2
podSpec:
schedulerName: kai-scheduler
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node_role.e2e.grove.nvidia.com
operator: In
values:
- agent
tolerations:
- key: node_role.e2e.grove.nvidia.com
operator: Equal
value: agent
effect: NoSchedule
containers:
- name: router
image: registry:5001/nginx:alpine-slim
resources:
requests:
memory: 30Mi
Loading
Loading