Add upgrade tests for Trainer v2

Author: ChughShilpa

tests/common/environment.go

@@ -27,6 +27,8 @@ import (
 )
 
 const (
+	// The environment variable referring to image simulating sleep condition in container
+	sleepImageEnvVar = "SLEEP_IMAGE"
 	// Name of the authenticated Notebook user
 	notebookUserName = "NOTEBOOK_USER_NAME"
 	// Token of the authenticated Notebook user
@@ -154,6 +156,10 @@ func GetHuggingFaceToken(t Test) string {
 	return token
 }
 
+func GetSleepImage() string {
+	return lookupEnvOrDefault(sleepImageEnvVar, "gcr.io/k8s-staging-perf-tests/sleep@sha256:8d91ddf9f145b66475efda1a1b52269be542292891b5de2a7fad944052bab6ea")
+}
+
 func init() {
 	flag.StringVar(&testTierParam, "testTier", "", "Test tier")
 }

tests/common/support/namespace.go

@@ -22,6 +22,7 @@ import (
 	"github.com/onsi/gomega"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -85,6 +86,23 @@ func GetNamespaceWithName(t Test, namespaceName string) *corev1.Namespace {
 	return namespace
 }
 
+// CreateOrGetTestNamespaceWithName creates a namespace with the given name if it doesn't exist,
+// or returns the existing namespace if it does. This is useful for scenarios where
+// the namespace needs to persist across multiple test phases.
+func CreateOrGetTestNamespaceWithName(t Test, name string, options ...Option[*corev1.Namespace]) *corev1.Namespace {
+	t.T().Helper()
+	namespace, err := t.Client().Core().CoreV1().Namespaces().Get(t.Ctx(), name, metav1.GetOptions{})
+	if err == nil {
+		return namespace
+	} else if errors.IsNotFound(err) {
+		t.T().Logf("%s namespace doesn't exist. Creating ...", name)
+		return CreateTestNamespaceWithName(t, name, options...)
+	} else {
+		t.T().Fatalf("Error retrieving namespace with name `%s`: %v", name, err)
+	}
+	return nil
+}
+
 func DeleteTestNamespace(t Test, namespace *corev1.Namespace) {
 	t.T().Helper()
 	propagationPolicy := metav1.DeletePropagationBackground

tests/kfto/environment.go

@@ -25,8 +25,6 @@ const (
 	bloomModelImageEnvVar = "BLOOM_MODEL_IMAGE"
 	// The environment variable referring to image containing Stanford Alpaca dataset
 	alpacaDatasetImageEnvVar = "ALPACA_DATASET_IMAGE"
-	// The environment variable referring to image simulating sleep condition in container
-	sleepImageEnvVar = "SLEEP_IMAGE"
 )
 
 func GetBloomModelImage() string {
@@ -37,10 +35,6 @@ func GetAlpacaDatasetImage() string {
 	return lookupEnvOrDefault(alpacaDatasetImageEnvVar, "quay.io/ksuta/alpaca-dataset@sha256:2e90f631180c7b2c916f9569b914b336b612e8ae86efad82546adc5c9fcbbb8d")
 }
 
-func GetSleepImage() string {
-	return lookupEnvOrDefault(sleepImageEnvVar, "gcr.io/k8s-staging-perf-tests/sleep@sha256:8d91ddf9f145b66475efda1a1b52269be542292891b5de2a7fad944052bab6ea")
-}
-
 func lookupEnvOrDefault(key, value string) string {
 	if v, ok := os.LookupEnv(key); ok {
 		return v

tests/kfto/kfto_kueue_mnist_upgrade_training_test.go

@@ -45,7 +45,7 @@ func TestSetupPytorchjob(t *testing.T) {
 	Tags(t, PreUpgrade)
 	test := With(t)
 
-	createOrGetUpgradeTestNamespace(test, namespaceName)
+	CreateOrGetTestNamespaceWithName(test, namespaceName)
 
 	// Create a ConfigMap with training dataset and configuration
 	mnist := readFile(test, "resources/mnist.py")
@@ -298,17 +298,3 @@ func createUpgradePyTorchJob(test Test, namespace, localQueueName string, config
 
 	return tuningJob
 }
-
-func createOrGetUpgradeTestNamespace(test Test, name string, options ...Option[*corev1.Namespace]) (namespace *corev1.Namespace) {
-	// Verify that the namespace really exists and return it, create it if doesn't exist yet
-	namespace, err := test.Client().Core().CoreV1().Namespaces().Get(test.Ctx(), name, metav1.GetOptions{})
-	if err == nil {
-		return
-	} else if errors.IsNotFound(err) {
-		test.T().Logf("%s namespace doesn't exists. Creating ...", name)
-		return CreateTestNamespaceWithName(test, name, options...)
-	} else {
-		test.T().Fatalf("Error retrieving namespace with name `%s`: %v", name, err)
-	}
-	return
-}

tests/kfto/kfto_upgrade_sleep_test.go

@@ -42,7 +42,7 @@ func TestSetupSleepPytorchjob(t *testing.T) {
 	test := With(t)
 
 	// Create a namespace
-	createOrGetUpgradeTestNamespace(test, sleepNamespaceName)
+	CreateOrGetTestNamespaceWithName(test, sleepNamespaceName)
 
 	// Create Kueue resources
 	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})