Refactor Machine pool functions

Author: sutaakar

test/e2e/instascale.go renamed to test/e2e/instascale_app_wrapper.go

@@ -9,55 +9,10 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
-	ocmsdk "github.com/openshift-online/ocm-sdk-go"
-
 	. "github.com/project-codeflare/codeflare-operator/test/support"
 )
 
-func TestConfig(test Test, namespace string) (*corev1.ConfigMap, error) {
-	// Test configuration
-	configMap := &corev1.ConfigMap{
-		TypeMeta: metav1.TypeMeta{
-			APIVersion: corev1.SchemeGroupVersion.String(),
-			Kind:       "ConfigMap",
-		},
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      "mnist-mcad",
-			Namespace: namespace,
-		},
-		BinaryData: map[string][]byte{
-			// pip requirements
-			"requirements.txt": ReadFile(test, "mnist_pip_requirements.txt"),
-			// MNIST training script
-			"mnist.py": ReadFile(test, "mnist.py"),
-		},
-		Immutable: Ptr(true),
-	}
-
-	config, err := test.Client().Core().CoreV1().ConfigMaps(namespace).Create(test.Ctx(), configMap, metav1.CreateOptions{})
-	test.Expect(err).NotTo(HaveOccurred())
-	test.T().Logf("Created ConfigMap %s/%s successfully", config.Namespace, config.Name)
-
-	return configMap, err
-}
-
-func CreateConnection(test Test) (*ocmsdk.Connection, error) {
-	instascaleOCMSecret, err := test.Client().Core().CoreV1().Secrets("default").Get(test.Ctx(), "instascale-ocm-secret", metav1.GetOptions{})
-	if err != nil {
-		test.T().Errorf("unable to retrieve instascale-ocm-secret - Error : %v", err)
-	}
-	test.Expect(err).NotTo(HaveOccurred())
-	ocmToken := string(instascaleOCMSecret.Data["token"])
-	test.T().Logf("Retrieved Secret %s successfully", instascaleOCMSecret.Name)
-
-	connection, err := CreateOCMConnection(ocmToken)
-	if err != nil {
-		test.T().Errorf("Unable to create ocm connection - Error : %v", err)
-	}
-	return connection, err
-}
-
-func JobAppwrapperSetup(test Test, namespace *corev1.Namespace, config *corev1.ConfigMap) (*batchv1.Job, *mcadv1beta1.AppWrapper, error) {
+func createInstaScaleJobAppWrapper(test Test, namespace *corev1.Namespace, config *corev1.ConfigMap) (*batchv1.Job, *mcadv1beta1.AppWrapper, error) {
 	// Batch Job
 	job := &batchv1.Job{
 		TypeMeta: metav1.TypeMeta{
@@ -78,7 +33,7 @@ func JobAppwrapperSetup(test Test, namespace *corev1.Namespace, config *corev1.C
 							Name:  "job",
 							Image: GetPyTorchImage(),
 							Env: []corev1.EnvVar{
-								corev1.EnvVar{Name: "PYTHONUSERBASE", Value: "/test2"},
+								{Name: "PYTHONUSERBASE", Value: "/workdir"},
 							},
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
 							Args:    []string{"$PYTHONUSERBASE"},
@@ -92,7 +47,7 @@ func JobAppwrapperSetup(test Test, namespace *corev1.Namespace, config *corev1.C
 									MountPath: "/workdir",
 								},
 							},
-							WorkingDir: "workdir",
+							WorkingDir: "/workdir",
 						},
 					},
 					Volumes: []corev1.Volume{
@@ -125,7 +80,7 @@ func JobAppwrapperSetup(test Test, namespace *corev1.Namespace, config *corev1.C
 			Name:      "test-instascale",
 			Namespace: namespace.Name,
 			Labels: map[string]string{
-				"orderedinstance": "m5.xlarge_g4dn.xlarge",
+				"orderedinstance": "g4dn.xlarge",
 			},
 		},
 		Spec: mcadv1beta1.AppWrapperSpec{
@@ -170,7 +125,7 @@ func JobAppwrapperSetup(test Test, namespace *corev1.Namespace, config *corev1.C
 	test.Expect(err).NotTo(HaveOccurred())
 	test.T().Logf("AppWrapper created successfully %s/%s", aw.Namespace, aw.Name)
 
-	test.Eventually(AppWrapper(test, namespace, aw.Name), TestTimeoutShort).
+	test.Eventually(AppWrapper(test, namespace, aw.Name), TestTimeoutGpuProvisioning).
 		Should(WithTransform(AppWrapperState, Equal(mcadv1beta1.AppWrapperStateActive)))
 
 	return job, aw, err

test/e2e/instascale_machinepool_test.go

@@ -2,7 +2,6 @@ package e2e
 
 import (
 	"testing"
-	"time"
 
 	. "github.com/onsi/gomega"
 	mcadv1beta1 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/apis/controller/v1beta1"
@@ -21,33 +20,31 @@ func TestInstascaleMachinePool(t *testing.T) {
 	namespace := test.NewTestNamespace()
 
 	// Test configuration
-	config, err := TestConfig(test, namespace.Name)
-	test.Expect(err).NotTo(HaveOccurred())
+	testConfigData := map[string][]byte{
+		// pip requirements
+		"requirements.txt": ReadFile(test, "mnist_pip_requirements.txt"),
+		// MNIST training script
+		"mnist.py": ReadFile(test, "mnist.py"),
+	}
+	cm := CreateConfigMap(test, namespace.Name, testConfigData)
 
 	//create OCM connection
-	connection, err := CreateConnection(test)
-	test.Expect(err).NotTo(HaveOccurred())
+	connection := CreateOCMConnection(test)
 
 	defer connection.Close()
 
 	// check existing cluster machine pool resources
 	// look for machine pool with aw name - expect not to find it
-	foundMachinePool, err := CheckMachinePools(connection, TestName)
-	test.Expect(err).NotTo(HaveOccurred())
-	test.Expect(foundMachinePool).To(BeFalse())
+	test.Expect(GetMachinePools(test, connection)).
+		ShouldNot(ContainElement(WithTransform(MachinePoolId, Equal("test-instascale-g4dn-xlarge"))))
 
 	// Setup batch job and AppWrapper
-	job, aw, err := JobAppwrapperSetup(test, namespace, config)
+	job, aw, err := createInstaScaleJobAppWrapper(test, namespace, cm)
 	test.Expect(err).NotTo(HaveOccurred())
 
-	// time.Sleep is used twice throughout the test, each for 30 seconds. Can look into using sync package waitGroup instead if that makes more sense
-	// wait for required resources to scale up before checking them again
-	time.Sleep(TestTimeoutMedium)
-
 	// look for machine pool with aw name - expect to find it
-	foundMachinePool, err = CheckMachinePools(connection, TestName)
-	test.Expect(err).NotTo(HaveOccurred())
-	test.Expect(foundMachinePool).To(BeTrue())
+	test.Eventually(MachinePools(test, connection), TestTimeoutLong).
+		Should(ContainElement(WithTransform(MachinePoolId, Equal("test-instascale-g4dn-xlarge"))))
 
 	// Assert that the job has completed
 	test.T().Logf("Waiting for Job %s/%s to complete", job.Namespace, job.Name)
@@ -64,12 +61,8 @@ func TestInstascaleMachinePool(t *testing.T) {
 	test.Eventually(AppWrapper(test, namespace, aw.Name), TestTimeoutShort).
 		Should(WithTransform(AppWrapperState, Equal(mcadv1beta1.AppWrapperStateCompleted)))
 
-	// allow time for the resources to scale down before checking them again
-	time.Sleep(TestTimeoutMedium)
-
 	// look for machine pool with aw name - expect not to find it
-	foundMachinePool, err = CheckMachinePools(connection, TestName)
-	test.Expect(err).NotTo(HaveOccurred())
-	test.Expect(foundMachinePool).To(BeFalse())
+	test.Eventually(MachinePools(test, connection), TestTimeoutLong).
+		ShouldNot(ContainElement(WithTransform(MachinePoolId, Equal("test-instascale-g4dn-xlarge"))))
 
 }

test/support/clusterpools.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/onsi/gomega"
 	ocmsdk "github.com/openshift-online/ocm-sdk-go"
 	cmv1 "github.com/openshift-online/ocm-sdk-go/clustersmgmt/v1"
 )
@@ -15,42 +16,21 @@ var (
 	TestName  string = "test-instascale"
 )
 
-func CreateOCMConnection(secret string) (*ocmsdk.Connection, error) {
-	logger, err := ocmsdk.NewGoLoggerBuilder().
-		Debug(false).
-		Build()
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Can't build logger: %v\n", err)
-		return nil, err
-	}
-	connection, err := ocmsdk.NewConnectionBuilder().
-		Logger(logger).
-		Tokens(string(secret)).
-		Build()
-	if err != nil || connection == nil {
-		fmt.Fprintf(os.Stderr, "Can't build connection: %v\n", err)
-		return nil, err
+func MachinePools(t Test, connection *ocmsdk.Connection) func(g gomega.Gomega) []*cmv1.MachinePool {
+	return func(g gomega.Gomega) []*cmv1.MachinePool {
+		machinePoolsListResponse, err := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).MachinePools().List().Send()
+		g.Expect(err).NotTo(gomega.HaveOccurred())
+		return machinePoolsListResponse.Items().Slice()
 	}
-
-	return connection, nil
 }
 
-func CheckMachinePools(connection *ocmsdk.Connection, awName string) (foundMachinePool bool, err error) {
-	machinePoolsConnection := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).MachinePools().List()
-	machinePoolsListResponse, err := machinePoolsConnection.Send()
-	if err != nil {
-		return false, fmt.Errorf("unable to send request, error: %v", err)
-	}
-	machinePoolsList := machinePoolsListResponse.Items()
-	machinePoolsList.Range(func(index int, item *cmv1.MachinePool) bool {
-		instanceName, _ := item.GetID()
-		if strings.Contains(instanceName, awName) {
-			foundMachinePool = true
-		}
-		return true
-	})
+func GetMachinePools(t Test, connection *ocmsdk.Connection) []*cmv1.MachinePool {
+	t.T().Helper()
+	return MachinePools(t, connection)(t)
+}
 
-	return foundMachinePool, err
+func MachinePoolId(machinePool *cmv1.MachinePool) string {
+	return machinePool.ID()
 }
 
 func CheckNodePools(connection *ocmsdk.Connection, awName string) (foundNodePool bool, err error) {

test/support/codeflare.go

@@ -30,8 +30,12 @@ const (
 	CodeFlareTestPyTorchImage = "CODEFLARE_TEST_PYTORCH_IMAGE"
 
 	// The testing output directory, to write output files into.
-
 	CodeFlareTestOutputDir = "CODEFLARE_TEST_OUTPUT_DIR"
+
+	// The name of a secret containing InstaScale OCM token.
+	InstaScaleOcmSecretName = "INSTASCALE_OCM_SECRET_NAME"
+	// The namespace where a secret containing InstaScale OCM token is stored.
+	InstaScaleOcmSecretNamespace = "INSTASCALE_OCM_SECRET_NAMESPACE"
 )
 
 func GetCodeFlareSDKVersion() string {
@@ -50,6 +54,14 @@ func GetPyTorchImage() string {
 	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")
 }
 
+func GetInstaScaleOcmSecretName() string {
+	return lookupEnvOrDefault(InstaScaleOcmSecretName, "instascale-ocm-secret")
+}
+
+func GetInstaScaleOcmSecretNamespace() string {
+	return lookupEnvOrDefault(InstaScaleOcmSecretNamespace, "default")
+}
+
 func lookupEnvOrDefault(key, value string) string {
 	if v, ok := os.LookupEnv(key); ok {
 		return v

test/support/config_map.go

@@ -0,0 +1,44 @@
+/*
+Copyright 2023.
+
+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 support
+
+import (
+	"github.com/onsi/gomega"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func CreateConfigMap(t Test, namespace string, content map[string][]byte) *corev1.ConfigMap {
+	configMap := &corev1.ConfigMap{
+		TypeMeta: metav1.TypeMeta{
+			APIVersion: corev1.SchemeGroupVersion.String(),
+			Kind:       "ConfigMap",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			GenerateName: "config-",
+			Namespace:    namespace,
+		},
+		BinaryData: content,
+		Immutable:  Ptr(true),
+	}
+
+	configMap, err := t.Client().Core().CoreV1().ConfigMaps(namespace).Create(t.Ctx(), configMap, metav1.CreateOptions{})
+	t.Expect(err).NotTo(gomega.HaveOccurred())
+	t.T().Logf("Created ConfigMap %s/%s successfully", configMap.Namespace, configMap.Name)
+
+	return configMap
+}

test/support/ocm.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2023.
+
+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 support
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/onsi/gomega"
+	ocmsdk "github.com/openshift-online/ocm-sdk-go"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func CreateOCMConnection(test Test) *ocmsdk.Connection {
+	instascaleOCMSecret, err := test.Client().Core().CoreV1().Secrets(GetInstaScaleOcmSecretNamespace()).Get(test.Ctx(), GetInstaScaleOcmSecretName(), metav1.GetOptions{})
+	test.Expect(err).NotTo(gomega.HaveOccurred())
+
+	ocmToken := string(instascaleOCMSecret.Data["token"])
+	test.T().Logf("Retrieved Secret %s/%s successfully", instascaleOCMSecret.Namespace, instascaleOCMSecret.Name)
+
+	connection, err := buildOCMConnection(ocmToken)
+	test.Expect(err).NotTo(gomega.HaveOccurred())
+	return connection
+}
+
+func buildOCMConnection(secret string) (*ocmsdk.Connection, error) {
+	connection, err := ocmsdk.NewConnectionBuilder().
+		Tokens(secret).
+		Build()
+	if err != nil || connection == nil {
+		fmt.Fprintf(os.Stderr, "Can't build connection: %v\n", err)
+		return nil, err
+	}
+
+	return connection, nil
+}

test/support/support.go

@@ -30,9 +30,10 @@ import (
 var (
 	ApplyOptions = metav1.ApplyOptions{FieldManager: "codeflare-test", Force: true}
 
-	TestTimeoutShort  = 1 * time.Minute
-	TestTimeoutMedium = 2 * time.Minute
-	TestTimeoutLong   = 5 * time.Minute
+	TestTimeoutShort           = 1 * time.Minute
+	TestTimeoutMedium          = 2 * time.Minute
+	TestTimeoutLong            = 5 * time.Minute
+	TestTimeoutGpuProvisioning = 30 * time.Minute
 )
 
 func init() {