refactor shared logic into common module

Signed-off-by: Kevin <[email protected]>

Author: KPostOffice

tests/odh/environment.go renamed to tests/common/environment.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package odh
+package common
 
 import (
 	"os"

tests/odh/notebook.go renamed to tests/common/notebook.go

@@ -14,10 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package odh
+package common
 
 import (
 	"bytes"
+	"embed"
 
 	gomega "github.com/onsi/gomega"
 	. "github.com/project-codeflare/codeflare-common/support"
@@ -34,6 +35,16 @@ const (
 	NOTEBOOK_CONTAINER_NAME = "jupyter-nb-kube-3aadmin"
 )
 
+//go:embed resources/*
+var files embed.FS
+
+func readFile(t Test, fileName string) []byte {
+	t.T().Helper()
+	file, err := files.ReadFile(fileName)
+	t.Expect(err).NotTo(gomega.HaveOccurred())
+	return file
+}
+
 var notebookResource = schema.GroupVersionResource{Group: "kubeflow.org", Version: "v1", Resource: "notebooks"}
 
 type NotebookProps struct {
@@ -42,7 +53,7 @@ type NotebookProps struct {
 	KubernetesUserBearerToken string
 	Namespace                 string
 	OpenDataHubNamespace      string
-	RayImage                  string
+	Command                   []string
 	NotebookImage             string
 	NotebookConfigMapName     string
 	NotebookConfigMapFileName string
@@ -57,7 +68,7 @@ type NotebookProps struct {
 	S3DefaultRegion           string
 }
 
-func createNotebook(test Test, namespace *corev1.Namespace, notebookUserToken, rayImage string, jupyterNotebookConfigMapName, jupyterNotebookConfigMapFileName string, numGpus int) {
+func CreateNotebook(test Test, namespace *corev1.Namespace, notebookUserToken string, command []string, jupyterNotebookConfigMapName, jupyterNotebookConfigMapFileName string, numGpus int) {
 	// Create PVC for Notebook
 	notebookPVC := CreatePersistentVolumeClaim(test, namespace.Name, "10Gi", corev1.ReadWriteOnce)
 	s3BucketName, s3BucketNameExists := GetStorageBucketName()
@@ -81,7 +92,7 @@ func createNotebook(test Test, namespace *corev1.Namespace, notebookUserToken, r
 		KubernetesUserBearerToken: notebookUserToken,
 		Namespace:                 namespace.Name,
 		OpenDataHubNamespace:      GetOpenDataHubNamespace(test),
-		RayImage:                  rayImage,
+		Command:                   command,
 		NotebookImage:             GetNotebookImage(test),
 		NotebookConfigMapName:     jupyterNotebookConfigMapName,
 		NotebookConfigMapFileName: jupyterNotebookConfigMapFileName,
@@ -108,12 +119,12 @@ func createNotebook(test Test, namespace *corev1.Namespace, notebookUserToken, r
 	test.Expect(err).NotTo(gomega.HaveOccurred())
 }
 
-func deleteNotebook(test Test, namespace *corev1.Namespace) {
+func DeleteNotebook(test Test, namespace *corev1.Namespace) {
 	err := test.Client().Dynamic().Resource(notebookResource).Namespace(namespace.Name).Delete(test.Ctx(), "jupyter-nb-kube-3aadmin", metav1.DeleteOptions{})
 	test.Expect(err).NotTo(gomega.HaveOccurred())
 }
 
-func listNotebooks(test Test, namespace *corev1.Namespace) []*unstructured.Unstructured {
+func ListNotebooks(test Test, namespace *corev1.Namespace) []*unstructured.Unstructured {
 	ntbs, err := test.Client().Dynamic().Resource(notebookResource).Namespace(namespace.Name).List(test.Ctx(), metav1.ListOptions{})
 	test.Expect(err).NotTo(gomega.HaveOccurred())
 

tests/kfto/resources/custom-nb-small.yaml renamed to tests/common/resources/custom-nb-small.yaml

@@ -0,0 +1,76 @@
+/*
+Copyright 2024.
+
+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 (
+	"bytes"
+	"text/template"
+
+	"github.com/onsi/gomega"
+	"github.com/project-codeflare/codeflare-common/support"
+)
+
+func ParseAWSArgs(t support.Test, inputTemplate []byte) []byte {
+	storage_bucket_endpoint, storage_bucket_endpoint_exists := support.GetStorageBucketDefaultEndpoint()
+	storage_bucket_access_key_id, storage_bucket_access_key_id_exists := support.GetStorageBucketAccessKeyId()
+	storage_bucket_secret_key, storage_bucket_secret_key_exists := support.GetStorageBucketSecretKey()
+	storage_bucket_name, storage_bucket_name_exists := support.GetStorageBucketName()
+	storage_bucket_mnist_dir, storage_bucket_mnist_dir_exists := support.GetStorageBucketMnistDir()
+
+	props := struct {
+		StorageBucketDefaultEndpoint       string
+		StorageBucketDefaultEndpointExists bool
+		StorageBucketAccessKeyId           string
+		StorageBucketAccessKeyIdExists     bool
+		StorageBucketSecretKey             string
+		StorageBucketSecretKeyExists       bool
+		StorageBucketName                  string
+		StorageBucketNameExists            bool
+		StorageBucketMnistDir              string
+		StorageBucketMnistDirExists        bool
+	}{
+		StorageBucketDefaultEndpoint:       storage_bucket_endpoint,
+		StorageBucketDefaultEndpointExists: storage_bucket_endpoint_exists,
+		StorageBucketAccessKeyId:           storage_bucket_access_key_id,
+		StorageBucketAccessKeyIdExists:     storage_bucket_access_key_id_exists,
+		StorageBucketSecretKey:             storage_bucket_secret_key,
+		StorageBucketSecretKeyExists:       storage_bucket_secret_key_exists,
+		StorageBucketName:                  storage_bucket_name,
+		StorageBucketNameExists:            storage_bucket_name_exists,
+		StorageBucketMnistDir:              storage_bucket_mnist_dir,
+		StorageBucketMnistDirExists:        storage_bucket_mnist_dir_exists,
+	}
+
+	return ParseTemplate(t, inputTemplate, props)
+}
+
+func ParseTemplate(t support.Test, inputTemplate []byte, props interface{}) []byte {
+	t.T().Helper()
+
+	// Parse input template
+	parsedTemplate, err := template.New("template").Parse(string(inputTemplate))
+	t.Expect(err).NotTo(gomega.HaveOccurred())
+
+	// Filter template and store results to the buffer
+	buffer := new(bytes.Buffer)
+	err = parsedTemplate.Execute(buffer, props)
+	t.Expect(err).NotTo(gomega.HaveOccurred())
+	err = parsedTemplate.Execute(buffer, props) // NOTE: not sure if template package handles recursive case
+	t.Expect(err).NotTo(gomega.HaveOccurred())
+
+	return buffer.Bytes()
+}

tests/common/template.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	. "github.com/onsi/gomega"
+	. "github.com/opendatahub-io/distributed-workloads/tests/common"
 	. "github.com/project-codeflare/codeflare-common/support"
 
 	v1 "k8s.io/api/core/v1"
@@ -33,7 +34,7 @@ func TestMnistSDK(t *testing.T) {
 	userName := GetNotebookUserName(test)
 	userToken := GetNotebookUserToken(test)
 	jupyterNotebookConfigMapFileName := "mnist_kfto.ipynb"
-	mnist := readMnistScriptTemplate(test, "resources/kfto_sdk_mnist.py")
+	mnist := ParseAWSArgs(test, readFile(test, "resources/kfto_sdk_mnist.py"))
 
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
@@ -45,8 +46,8 @@ func TestMnistSDK(t *testing.T) {
 		"${namespace}": namespace.Name,
 	}
 
-	jupyterNotebook := string(ReadFile(test, "resources/mnist_kfto.ipynb"))
-	requirements := ReadFile(test, "resources/requirements.txt")
+	jupyterNotebook := string(readFile(test, "resources/mnist_kfto.ipynb"))
+	requirements := readFile(test, "resources/requirements.txt")
 	for oldValue, newValue := range requiredChangesInNotebook {
 		jupyterNotebook = strings.Replace(string(jupyterNotebook), oldValue, newValue, -1)
 	}
@@ -57,13 +58,21 @@ func TestMnistSDK(t *testing.T) {
 		"requirements.txt":               requirements,
 	})
 
+	notebookCommand := []string{
+		"bin/sh",
+		"-c",
+		"pip install papermill && papermill /opt/app-root/notebooks/{{.NotebookConfigMapFileName}}" +
+			" /opt/app-root/src/mcad-out.ipynb -p namespace {{.Namespace}} -p openshift_api_url {{.OpenShiftApiUrl}}" +
+			" -p kubernetes_user_bearer_token {{.KubernetesUserBearerToken}}" +
+			" -p num_gpus {{ .NumGpus }} --log-output && sleep infinity",
+	}
 	// Create Notebook CR
-	createNotebook(test, namespace, userToken, config.Name, jupyterNotebookConfigMapFileName, 0)
+	CreateNotebook(test, namespace, userToken, notebookCommand, config.Name, jupyterNotebookConfigMapFileName, 0)
 
 	// Gracefully cleanup Notebook
 	defer func() {
-		deleteNotebook(test, namespace)
-		test.Eventually(listNotebooks(test, namespace), TestTimeoutGpuProvisioning).Should(HaveLen(0))
+		DeleteNotebook(test, namespace)
+		test.Eventually(ListNotebooks(test, namespace), TestTimeoutGpuProvisioning).Should(HaveLen(0))
 	}()
 
 	// Make sure pytorch job is created
@@ -77,40 +86,3 @@ func TestMnistSDK(t *testing.T) {
 	// TODO: write torch job logs?
 	// time.Sleep(60 * time.Second)
 }
-
-func readMnistScriptTemplate(test Test, filePath string) []byte {
-	// Read the mnist.py from resources and perform replacements for custom values using go template
-	storage_bucket_endpoint, storage_bucket_endpoint_exists := GetStorageBucketDefaultEndpoint()
-	storage_bucket_access_key_id, storage_bucket_access_key_id_exists := GetStorageBucketAccessKeyId()
-	storage_bucket_secret_key, storage_bucket_secret_key_exists := GetStorageBucketSecretKey()
-	storage_bucket_name, storage_bucket_name_exists := GetStorageBucketName()
-	storage_bucket_mnist_dir, storage_bucket_mnist_dir_exists := GetStorageBucketMnistDir()
-
-	props := struct {
-		StorageBucketDefaultEndpoint       string
-		StorageBucketDefaultEndpointExists bool
-		StorageBucketAccessKeyId           string
-		StorageBucketAccessKeyIdExists     bool
-		StorageBucketSecretKey             string
-		StorageBucketSecretKeyExists       bool
-		StorageBucketName                  string
-		StorageBucketNameExists            bool
-		StorageBucketMnistDir              string
-		StorageBucketMnistDirExists        bool
-	}{
-		StorageBucketDefaultEndpoint:       storage_bucket_endpoint,
-		StorageBucketDefaultEndpointExists: storage_bucket_endpoint_exists,
-		StorageBucketAccessKeyId:           storage_bucket_access_key_id,
-		StorageBucketAccessKeyIdExists:     storage_bucket_access_key_id_exists,
-		StorageBucketSecretKey:             storage_bucket_secret_key,
-		StorageBucketSecretKeyExists:       storage_bucket_secret_key_exists,
-		StorageBucketName:                  storage_bucket_name,
-		StorageBucketNameExists:            storage_bucket_name_exists,
-		StorageBucketMnistDir:              storage_bucket_mnist_dir,
-		StorageBucketMnistDirExists:        storage_bucket_mnist_dir_exists,
-	}
-	template, err := files.ReadFile(filePath)
-	test.Expect(err).NotTo(HaveOccurred())
-
-	return ParseTemplate(test, template, props)
-}

tests/kfto/kfto_mnist_sdk_test.go

@@ -59,9 +59,9 @@ func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string,
 	// Create a namespace
 	namespace := test.NewTestNamespace()
 
-	mnist := ReadFile(test, "resources/mnist.py")
-	download_mnist_dataset := ReadFile(test, "resources/download_mnist_datasets.py")
-	requirementsFileName := ReadFile(test, requirementsFile)
+	mnist := readFile(test, "resources/mnist.py")
+	download_mnist_dataset := readFile(test, "resources/download_mnist_datasets.py")
+	requirementsFileName := readFile(test, requirementsFile)
 
 	if accelerator.isGpu() {
 		mnist = bytes.Replace(mnist, []byte("accelerator=\"has to be specified\""), []byte("accelerator=\"gpu\""), 1)

tests/kfto/kfto_mnist_training_test.go

@@ -70,7 +70,7 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpu Accelerator, numGpus, num
 
 	// Create a ConfigMap with training script
 	configData := map[string][]byte{
-		"hf_llm_training.py": ReadFile(test, "resources/hf_llm_training.py"),
+		"hf_llm_training.py": readFile(test, "resources/hf_llm_training.py"),
 	}
 	config := CreateConfigMap(test, namespace, configData)