Implement custom test tags for ODH KFTO tests

Author: sutaakar

README.md

@@ -24,6 +24,7 @@
 * `CODEFLARE_TEST_TIMEOUT_LONG` - Timeout duration for long tasks
 * `CODEFLARE_TEST_RAY_IMAGE` (Optional) - Ray image used for raycluster configuration 
 * `MINIO_CLI_IMAGE` (Optional) - Minio CLI image used for uploading/downloading data from/into s3 bucket
+* `TEST_TIER` (Optional) - Specifies test tier to run, skipping tests which don't belong to specified test tier. Supported test tiers: Smoke, Sanity, Tier1, Tier2, Tier3, Pre-Upgrade and Post-Upgrade.
 
     NOTE: `quay.io/modh/ray:2.35.0-py311-cu121` is the default image used for creating a RayCluster resource. If you have your own custom ray image which suits your purposes, specify it in `CODEFLARE_TEST_RAY_IMAGE` environment variable.
 

tests/common/environment.go

@@ -18,6 +18,7 @@ package common
 
 import (
 	"os"
+	"slices"
 
 	. "github.com/project-codeflare/codeflare-common/support"
 )
@@ -31,8 +32,24 @@ const (
 	notebookUserToken = "NOTEBOOK_USER_TOKEN"
 	// Image of the Notebook
 	notebookImage = "NOTEBOOK_IMAGE"
+	// Test tier to be invoked
+	testTierEnvVar = "TEST_TIER"
 )
 
+const (
+	tierSmoke   = "Smoke"
+	tierSanity  = "Sanity"
+	tier1       = "Tier1"
+	tier2       = "Tier2"
+	tier3       = "Tier3"
+	preUpgrade  = "Pre-Upgrade"
+	postUpgrade = "Post-Upgrade"
+	kftoCuda    = "KFTO-CUDA"
+	kftoRocm    = "KFTO-ROCm"
+)
+
+var testTiers = []string{tierSmoke, tierSanity, tier1, tier2, tier3, preUpgrade, postUpgrade, kftoCuda, kftoRocm}
+
 func GetOpenDataHubNamespace(t Test) string {
 	ns, ok := os.LookupEnv(odhNamespaceEnvVar)
 	if !ok {
@@ -64,3 +81,14 @@ func GetNotebookImage(t Test) string {
 	}
 	return notebook_image
 }
+
+func GetTestTier(t Test) (string, bool) {
+	tt, ok := os.LookupEnv(testTierEnvVar)
+	if ok {
+		if slices.Contains(testTiers, tt) {
+			return tt, true
+		}
+		t.T().Fatalf("Environment variable %s is defined and contains invalid value: '%s'. Valid values are: %v", testTierEnvVar, tt, testTiers)
+	}
+	return "", false
+}

tests/common/test_tag.go

@@ -0,0 +1,156 @@
+/*
+Copyright 2025.
+
+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.
+*/
+
+package common
+
+import (
+	"fmt"
+	"testing"
+
+	. "github.com/onsi/gomega"
+	. "github.com/project-codeflare/codeflare-common/support"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func Tags(t *testing.T, tags ...func(test Test) (runTest bool, skipReason string)) {
+	test := With(t)
+	for _, tag := range tags {
+		runTest, skipReason := tag(test)
+		if !runTest {
+			test.T().Skip(skipReason)
+		}
+	}
+}
+
+// Test tag list
+
+var Smoke = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, tierSmoke)
+}
+
+var Sanity = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, tierSanity)
+}
+
+var Tier1 = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, tier1)
+}
+
+var Tier2 = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, tier2)
+}
+
+var Tier3 = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, tier3)
+}
+
+var PreUpgrade = func(test Test) (runTest bool, skipReason string) {
+	return mandatoryTestTier(test, preUpgrade)
+}
+
+var PostUpgrade = func(test Test) (runTest bool, skipReason string) {
+	return mandatoryTestTier(test, postUpgrade)
+}
+
+var KftoCuda = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, kftoCuda)
+}
+
+var KftoRocm = func(test Test) (runTest bool, skipReason string) {
+	return testTier(test, kftoRocm)
+}
+
+func Gpu(accelerator Accelerator) func(test Test) (runTest bool, skipReason string) {
+	return func(test Test) (runTest bool, skipReason string) {
+		return isGpuCountAvailableForNodes(test, 1, accelerator.ResourceLabel, 1)
+	}
+}
+
+func MultiGpu(accelerator Accelerator, numberOfGpus int) func(test Test) (runTest bool, skipReason string) {
+	return func(test Test) (runTest bool, skipReason string) {
+		return isGpuCountAvailableForNodes(test, 1, accelerator.ResourceLabel, numberOfGpus)
+	}
+}
+
+func MultiNode(numberOfNodes int) func(test Test) (runTest bool, skipReason string) {
+	return func(test Test) (runTest bool, skipReason string) {
+		nodes, err := test.Client().Core().CoreV1().Nodes().List(test.Ctx(), v1.ListOptions{LabelSelector: "node-role.kubernetes.io/worker"})
+		test.Expect(err).NotTo(HaveOccurred())
+
+		if len(nodes.Items) < numberOfNodes {
+			return false, fmt.Sprintf("Detected number of nodes is %d, which is lower than expected %d.", len(nodes.Items), numberOfNodes)
+		}
+		return true, ""
+	}
+}
+
+func MultiNodeGpu(numberOfNodes int, accelerator Accelerator) func(test Test) (runTest bool, skipReason string) {
+	return func(test Test) (runTest bool, skipReason string) {
+		return isGpuCountAvailableForNodes(test, numberOfNodes, accelerator.ResourceLabel, 1)
+	}
+}
+
+func MultiNodeMultiGpu(numberOfNodes int, accelerator Accelerator, numberOfGpus int) func(test Test) (runTest bool, skipReason string) {
+	return func(test Test) (runTest bool, skipReason string) {
+		return isGpuCountAvailableForNodes(test, numberOfNodes, accelerator.ResourceLabel, numberOfGpus)
+	}
+}
+
+// util functions
+
+func testTier(test Test, expectedTestTier string) (runTest bool, skipReason string) {
+	actualTestTier, found := GetTestTier(test)
+	if !found || actualTestTier == expectedTestTier {
+		return true, ""
+	}
+	return false, fmt.Sprintf("Test tier '%s' doesn't match expected tier '%s'", actualTestTier, expectedTestTier)
+}
+
+func mandatoryTestTier(test Test, expectedTestTier string) (runTest bool, skipReason string) {
+	actualTestTier, found := GetTestTier(test)
+	if found && actualTestTier == expectedTestTier {
+		return true, ""
+	}
+	return false, fmt.Sprintf("Test tier '%s' doesn't match expected tier '%s'", actualTestTier, expectedTestTier)
+}
+
+func isGpuCountAvailableForNodes(test Test, expectedNodes int, gpuResourceName string, expectedGpus int) (runTest bool, skipReason string) {
+	nodes, err := test.Client().Core().CoreV1().Nodes().List(test.Ctx(), v1.ListOptions{LabelSelector: "node-role.kubernetes.io/worker"})
+	test.Expect(err).NotTo(HaveOccurred())
+
+	var gpuNodes []corev1.Node
+	for _, node := range nodes.Items {
+		if node.Status.Allocatable.Name(corev1.ResourceName(gpuResourceName), resource.DecimalSI).Value() != 0 {
+			gpuNodes = append(gpuNodes, node)
+		}
+	}
+
+	if len(gpuNodes) < expectedNodes {
+		return false, fmt.Sprintf("Detected number of nodes with resource '%s' is %d, which is lower than expected %d.", gpuResourceName, len(gpuNodes), expectedNodes)
+	}
+
+	for _, gpuNode := range gpuNodes {
+		gpuCount := int(gpuNode.Status.Allocatable.Name(corev1.ResourceName(gpuResourceName), resource.DecimalSI).Value())
+
+		if gpuCount < expectedGpus {
+			return false, fmt.Sprintf("Detected number of GPUs for nodes with resource '%s' is %d, which is lower than expected %d.", gpuResourceName, gpuCount, expectedGpus)
+		}
+	}
+	return true, ""
+}

tests/kfto/kfto_kueue_mnist_upgrade_training_test.go

@@ -29,6 +29,8 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	kueuev1beta1 "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	kueueacv1beta1 "sigs.k8s.io/kueue/client-go/applyconfiguration/kueue/v1beta1"
+
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 )
 
 var (
@@ -40,6 +42,7 @@ var (
 )
 
 func TestSetupPytorchjob(t *testing.T) {
+	Tags(t, PreUpgrade)
 	test := With(t)
 
 	createOrGetUpgradeTestNamespace(test, namespaceName)
@@ -102,6 +105,7 @@ func TestSetupPytorchjob(t *testing.T) {
 }
 
 func TestRunPytorchjob(t *testing.T) {
+	Tags(t, PostUpgrade)
 	test := With(t)
 	namespace := GetNamespaceWithName(test, namespaceName)
 

tests/kfto/kfto_mnist_sdk_test.go

@@ -30,6 +30,7 @@ import (
 )
 
 func TestMnistSDK(t *testing.T) {
+	Tags(t, Tier1)
 	test := With(t)
 	// Create a namespace
 	namespace := test.NewTestNamespace()

tests/kfto/kfto_mnist_training_test.go

@@ -28,28 +28,36 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 )
 
 func TestPyTorchJobMnistMultiNodeSingleCpu(t *testing.T) {
+	Tags(t, Sanity, MultiNode(3))
 	runKFTOPyTorchMnistJob(t, CPU, GetCudaTrainingImage(), "resources/requirements.txt", 2, 1)
 }
 func TestPyTorchJobMnistMultiNodeMultiCpu(t *testing.T) {
+	Tags(t, Tier1, MultiNode(3))
 	runKFTOPyTorchMnistJob(t, CPU, GetCudaTrainingImage(), "resources/requirements.txt", 2, 2)
 }
 
 func TestPyTorchJobMnistMultiNodeSingleGpuWithCuda(t *testing.T) {
+	Tags(t, KftoCuda)
 	runKFTOPyTorchMnistJob(t, NVIDIA, GetCudaTrainingImage(), "resources/requirements.txt", 1, 1)
 }
 
 func TestPyTorchJobMnistMultiNodeMultiGpuWithCuda(t *testing.T) {
+	Tags(t, KftoCuda)
 	runKFTOPyTorchMnistJob(t, NVIDIA, GetCudaTrainingImage(), "resources/requirements.txt", 1, 2)
 }
 
 func TestPyTorchJobMnistMultiNodeSingleGpuWithROCm(t *testing.T) {
+	Tags(t, KftoRocm)
 	runKFTOPyTorchMnistJob(t, AMD, GetROCmTrainingImage(), "resources/requirements-rocm.txt", 1, 1)
 }
 
 func TestPyTorchJobMnistMultiNodeMultiGpuWithROCm(t *testing.T) {
+	Tags(t, KftoRocm)
 	runKFTOPyTorchMnistJob(t, AMD, GetROCmTrainingImage(), "resources/requirements-rocm.txt", 1, 2)
 }
 

tests/kfto/kfto_pytorchjob_failed_test.go

@@ -10,13 +10,17 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 )
 
 func TestPyTorchJobFailureWithCuda(t *testing.T) {
+	Tags(t, Tier1)
 	runFailedPyTorchJobTest(t, GetCudaTrainingImage())
 }
 
 func TestPyTorchJobFailureWithROCm(t *testing.T) {
+	Tags(t, Tier1)
 	runFailedPyTorchJobTest(t, GetROCmTrainingImage())
 }
 

tests/kfto/kfto_training_test.go

@@ -28,37 +28,47 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 )
 
 func TestPyTorchJobSingleNodeSingleGpuWithCuda(t *testing.T) {
+	Tags(t, Tier1, Gpu(NVIDIA))
 	runKFTOPyTorchJob(t, GetCudaTrainingImage(), NVIDIA, 1, 0)
 }
 
 func TestPyTorchJobSingleNodeMultiGpuWithCuda(t *testing.T) {
+	Tags(t, KftoCuda)
 	runKFTOPyTorchJob(t, GetCudaTrainingImage(), NVIDIA, 2, 0)
 }
 
 func TestPyTorchJobMultiNodeSingleGpuWithCuda(t *testing.T) {
+	Tags(t, KftoCuda)
 	runKFTOPyTorchJob(t, GetCudaTrainingImage(), NVIDIA, 1, 1)
 }
 
 func TestPyTorchJobMultiNodeMultiGpuWithCuda(t *testing.T) {
+	Tags(t, KftoCuda)
 	runKFTOPyTorchJob(t, GetCudaTrainingImage(), NVIDIA, 2, 1)
 }
 
 func TestPyTorchJobSingleNodeSingleGpuWithROCm(t *testing.T) {
+	Tags(t, Tier1, Gpu(AMD))
 	runKFTOPyTorchJob(t, GetROCmTrainingImage(), AMD, 1, 0)
 }
 
 func TestPyTorchJobSingleNodeMultiGpuWithROCm(t *testing.T) {
+	Tags(t, KftoRocm)
 	runKFTOPyTorchJob(t, GetROCmTrainingImage(), AMD, 2, 0)
 }
 
 func TestPyTorchJobMultiNodeSingleGpuWithROCm(t *testing.T) {
+	Tags(t, KftoRocm)
 	runKFTOPyTorchJob(t, GetROCmTrainingImage(), AMD, 1, 1)
 }
 
 func TestPyTorchJobMultiNodeMultiGpuWithROCm(t *testing.T) {
+	Tags(t, KftoRocm)
 	runKFTOPyTorchJob(t, GetROCmTrainingImage(), AMD, 2, 1)
 }
 

tests/kfto/kfto_upgrade_sleep_test.go

@@ -26,6 +26,8 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 )
 
 var (
@@ -34,6 +36,7 @@ var (
 )
 
 func TestSetupSleepPytorchjob(t *testing.T) {
+	Tags(t, PreUpgrade)
 	test := With(t)
 
 	// Create a namespace
@@ -48,6 +51,7 @@ func TestSetupSleepPytorchjob(t *testing.T) {
 }
 
 func TestVerifySleepPytorchjob(t *testing.T) {
+	Tags(t, PostUpgrade)
 	test := With(t)
 	namespace := GetNamespaceWithName(test, sleepNamespaceName)