Updated ray-finetune-test flow to reuse existing demo files from examples directory

Author: abhijeet-dhumal

examples/ray-finetune-llm-deepspeed/create_dataset.py

@@ -2,7 +2,10 @@
 import json
 import os
 
-dataset = load_dataset("gsm8k", "main", cache_dir="../../datasets")
+cache_dir="../../datasets"
+if not os.path.exists(cache_dir):
+    cache_dir=""
+dataset = load_dataset("gsm8k", "main", cache_dir=cache_dir)
 
 dataset_splits = {"train": dataset["train"], "test": dataset["test"]}
 

examples/ray-finetune-llm-deepspeed/ray_finetune_llm_deepspeed.ipynb

@@ -42,8 +42,8 @@
     "# On OpenShift, you can retrieve the token by running `oc whoami -t`,\n",
     "# and the server with `oc cluster-info`.\n",
     "auth = TokenAuthentication(\n",
-    "    token = \"\",\n",
-    "    server = \"\",\n",
+    "    token = '',\n",
+    "    server = '',\n",
     "    skip_tls=False\n",
     ")\n",
     "auth.login()"
@@ -69,7 +69,7 @@
     "    head_memory=128,\n",
     "    head_gpus=1,\n",
     "    num_gpus=1,\n",
-    "    image=\"quay.io/rhoai/ray:2.23.0-py39-cu121\",\n",
+    "    image='quay.io/rhoai/ray:2.23.0-py39-cu121',\n",
     "))"
    ]
   },
@@ -126,7 +126,7 @@
    "source": [
     "# The S3 bucket where to store checkpoint.\n",
     "# It can be set manually, otherwise it's retrieved from configured the data connection.\n",
-    "s3_bucket = \"\"\n",
+    "s3_bucket = ''\n",
     "if not s3_bucket:\n",
     "    s3_bucket = os.environ.get('AWS_S3_BUCKET')\n",
     "assert s3_bucket, \"An S3 bucket must be provided to store checkpoints\""
@@ -153,12 +153,12 @@
     "               \"--eval-batch-size-per-device=32 \",\n",
     "    runtime_env={\n",
     "        \"env_vars\": {\n",
-    "            \"AWS_ACCESS_KEY_ID\": os.environ.get('AWS_ACCESS_KEY_ID'),\n",
-    "            \"AWS_SECRET_ACCESS_KEY\": os.environ.get('AWS_SECRET_ACCESS_KEY'),\n",
-    "            \"AWS_DEFAULT_REGION\": os.environ.get('AWS_DEFAULT_REGION')\n",
+    "            'AWS_ACCESS_KEY_ID': os.environ.get('AWS_ACCESS_KEY_ID'),\n",
+    "            'AWS_SECRET_ACCESS_KEY': os.environ.get('AWS_SECRET_ACCESS_KEY'),\n",
+    "            'AWS_DEFAULT_REGION': os.environ.get('AWS_DEFAULT_REGION')\n",
     "        },\n",
-    "        \"pip\": \"requirements.txt\",\n",
-    "        \"working_dir\": \"./\",\n",
+    "        'pip': 'requirements.txt',\n",
+    "        'working_dir': './',\n",
     "        \"excludes\": [\"/docs/\", \"*.ipynb\", \"*.md\"]\n",
     "    },\n",
     ")\n",

examples/ray-finetune-llm-deepspeed/ray_finetune_llm_deepspeed.py

@@ -592,6 +592,9 @@ def parse_args():
 
     parser.add_argument("--lora", action="store_true", default=False,
                         help="If passed, will enable parameter efficient fine-tuning with LoRA.")
+    
+    parser.add_argument("--lora-config", type=str, default="./lora_configs/lora.json",
+                        help="Lora config json to use.")
 
     parser.add_argument("--num-epochs", type=int, default=1,
                         help="Number of epochs to train for.")
@@ -660,7 +663,7 @@ def main():
 
     # Add LoRA config if needed
     if args.lora:
-        with open("./lora_configs/lora.json", "r") as json_file:
+        with open(args.lora_config, "r") as json_file:
             lora_config = json.load(json_file)
         config["lora_config"] = lora_config
 

tests/odh/mnist_ray_test.go

@@ -81,7 +81,7 @@ func mnistRay(t *testing.T, numGpus int) {
 
 	// Test configuration
 	jupyterNotebookConfigMapFileName := "mnist_ray_mini.ipynb"
-	mnist := readMnistPy(test)
+	mnist := readMnistPy(test, "resources/mnist.py")
 	if numGpus > 0 {
 		mnist = bytes.Replace(mnist, []byte("accelerator=\"has to be specified\""), []byte("accelerator=\"gpu\""), 1)
 	} else {
@@ -91,7 +91,7 @@ func mnistRay(t *testing.T, numGpus int) {
 		// MNIST Ray Notebook
 		jupyterNotebookConfigMapFileName: ReadFile(test, "resources/mnist_ray_mini.ipynb"),
 		"mnist.py":                       mnist,
-		"requirements.txt":               readRequirementsTxt(test),
+		"requirements.txt":               ReadFile(test, "resources/requirements.txt"),
 	})
 
 	// Define the regular(non-admin) user
@@ -133,27 +133,7 @@ func mnistRay(t *testing.T, numGpus int) {
 		Should(HaveLen(0))
 }
 
-func readRequirementsTxt(test Test) []byte {
-	// Read the requirements.txt from resources and perform replacements for custom values using go template
-	props := struct {
-		PipIndexUrl    string
-		PipTrustedHost string
-	}{
-		PipIndexUrl: "--index " + string(GetPipIndexURL()),
-	}
-
-	// Provide trusted host only if defined
-	if len(GetPipTrustedHost()) > 0 {
-		props.PipTrustedHost = "--trusted-host " + GetPipTrustedHost()
-	}
-
-	template, err := files.ReadFile("resources/requirements.txt")
-	test.Expect(err).NotTo(HaveOccurred())
-
-	return ParseTemplate(test, template, props)
-}
-
-func readMnistPy(test Test) []byte {
+func readMnistPy(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()
@@ -184,7 +164,7 @@ func readMnistPy(test Test) []byte {
 		StorageBucketMnistDir:              storage_bucket_mnist_dir,
 		StorageBucketMnistDirExists:        storage_bucket_mnist_dir_exists,
 	}
-	template, err := files.ReadFile("resources/mnist.py")
+	template, err := files.ReadFile(filePath)
 	test.Expect(err).NotTo(HaveOccurred())
 
 	return ParseTemplate(test, template, props)

tests/odh/mnist_raytune_hpo_test.go

@@ -80,7 +80,7 @@ func mnistRayTuneHpo(t *testing.T, numGpus int) {
 
 	// Test configuration
 	jupyterNotebookConfigMapFileName := "mnist_hpo_raytune.ipynb"
-	mnist_hpo := ReadFile(test, "resources/mnist_hpo.py")
+	mnist_hpo := readMnistPy(test, "resources/mnist_hpo.py")
 
 	if numGpus > 0 {
 		mnist_hpo = bytes.Replace(mnist_hpo, []byte("gpu_value=\"has to be specified\""), []byte("gpu_value=\"1\""), 1)

tests/odh/notebook.go

@@ -45,23 +45,29 @@ type NotebookProps struct {
 	NotebookConfigMapFileName string
 	NotebookPVC               string
 	NumGpus                   int
+	PipIndexUrl               string
+	PipTrustedHost            string
 	S3BucketName              string
-	S3BucketNameExists        bool
 	S3AccessKeyId             string
-	S3AccessKeyIdExists       bool
 	S3SecretAccessKey         string
-	S3SecretAccessKeyExists   bool
 	S3DefaultRegion           string
 }
 
 func createNotebook(test Test, namespace *corev1.Namespace, notebookUserToken, jupyterNotebookConfigMapName, jupyterNotebookConfigMapFileName string, numGpus int) {
 	// Create PVC for Notebook
 	notebookPVC := CreatePersistentVolumeClaim(test, namespace.Name, "10Gi", corev1.ReadWriteOnce)
-	s3BucketName, _ := GetStorageBucketName()
+	s3BucketName, s3BucketNameExists := GetStorageBucketName()
 	s3AccessKeyId, _ := GetStorageBucketAccessKeyId()
 	s3SecretAccessKey, _ := GetStorageBucketSecretKey()
 	s3DefaultRegion, _ := GetStorageBucketDefaultRegion()
 
+	if !s3BucketNameExists {
+		s3BucketName = "''"
+		s3AccessKeyId = "''"
+		s3SecretAccessKey = "''"
+		s3DefaultRegion = "''"
+	}
+
 	// Read the Notebook CR from resources and perform replacements for custom values using go template
 	notebookProps := NotebookProps{
 		IngressDomain:             GetOpenShiftIngressDomain(test),
@@ -79,6 +85,8 @@ func createNotebook(test Test, namespace *corev1.Namespace, notebookUserToken, j
 		S3AccessKeyId:             s3AccessKeyId,
 		S3SecretAccessKey:         s3SecretAccessKey,
 		S3DefaultRegion:           s3DefaultRegion,
+		PipIndexUrl:               GetPipIndexURL(),
+		PipTrustedHost:            GetPipTrustedHost(),
 	}
 	notebookTemplate, err := files.ReadFile("resources/custom-nb-small.yaml")
 	test.Expect(err).NotTo(gomega.HaveOccurred())

tests/odh/ray_finetune_llm_deepspeed_test.go

@@ -17,39 +17,27 @@ limitations under the License.
 package odh
 
 import (
+	"fmt"
+	"os"
+	"strings"
 	"testing"
 
 	. "github.com/onsi/gomega"
 	. "github.com/project-codeflare/codeflare-common/support"
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 )
 
-func TestRayFinetuneDemo(t *testing.T) {
-	mnistRayLlmFinetune(t, 1)
+func TestRayFinetuneLlmDeepspeedDemo(t *testing.T) {
+	rayFinetuneLlmDeepspeed(t, 1)
 }
 
-func mnistRayLlmFinetune(t *testing.T, numGpus int) {
+func rayFinetuneLlmDeepspeed(t *testing.T, numGpus int) {
 	test := With(t)
 
 	// Create a namespace
 	namespace := test.NewTestNamespace()
-
-	// Test configuration
-	jupyterNotebookConfigMapFileName := "ray_finetune_llm_deepspeed.ipynb"
-
-	// Test configuration
-	configMap := map[string][]byte{
-		// MNIST Ray Notebook
-		jupyterNotebookConfigMapFileName: ReadFile(test, "resources/ray_finetune_demo/ray_finetune_llm_deepspeed.ipynb"),
-		"ray_finetune_llm_deepspeed.py":  ReadFile(test, "resources/ray_finetune_demo/ray_finetune_llm_deepspeed.py"),
-		"ray_finetune_requirements.txt":  ReadRayFinetuneRequirementsTxt(test),
-		"create_dataset.py":              ReadFile(test, "resources/ray_finetune_demo/create_dataset.py"),
-		"lora.json":                      ReadFile(test, "resources/ray_finetune_demo/lora.json"),
-		"zero_3_llama_2_7b.json":         ReadFile(test, "resources/ray_finetune_demo/zero_3_llama_2_7b.json"),
-		"utils.py":                       ReadFile(test, "resources/ray_finetune_demo/utils.py"),
-	}
-
-	config := CreateConfigMap(test, namespace.Name, configMap)
+	var workingDirectory, err = os.Getwd()
+	test.Expect(err).ToNot(HaveOccurred())
 
 	// Define the regular(non-admin) user
 	userName := GetNotebookUserName(test)
@@ -58,6 +46,53 @@ func mnistRayLlmFinetune(t *testing.T, numGpus int) {
 	// Create role binding with Namespace specific admin cluster role
 	CreateUserRoleBindingWithClusterRole(test, userName, namespace.Name, "admin")
 
+	// list changes required in llm-deepspeed-finetune-demo.ipynb file and update those
+	requiredChangesInNotebook := map[string]string{
+		"import os":  "import os,time,sys",
+		"import sys": "!cp /opt/app-root/notebooks/* ./",
+		"from codeflare_sdk.cluster.auth import TokenAuthentication": "from codeflare_sdk.cluster.auth import TokenAuthentication\\n\",\n\t\"from codeflare_sdk.job import RayJobClient",
+		"token = ''":                             fmt.Sprintf("token = '%s'", userToken),
+		"server = ''":                            fmt.Sprintf("server = '%s'", GetOpenShiftApiUrl(test)),
+		"namespace='ray-finetune-llm-deepspeed'": fmt.Sprintf("namespace='%s'", namespace.Name),
+		"head_cpus=16":                           "head_cpus=2",
+		"head_gpus=1":                            "head_gpus=0",
+		"num_workers=7":                          "num_workers=1",
+		"min_cpus=16":                            "min_cpus=4",
+		"max_cpus=16":                            "max_cpus=4",
+		"min_memory=128":                         "min_memory=48",
+		"max_memory=256":                         "max_memory=48",
+		"head_memory=128":                        "head_memory=48",
+		"num_gpus=1":                             fmt.Sprintf("worker_extended_resource_requests={'nvidia.com/gpu': %d},\\n\",\n\t\"    write_to_file=True,\\n\",\n\t\"    verify_tls=False", numGpus),
+		"image='quay.io/rhoai/ray:2.23.0-py39-cu121'":            fmt.Sprintf("image='%s'", GetRayImage()),
+		"client = cluster.job_client":                            "ray_dashboard = cluster.cluster_dashboard_uri()\\n\",\n\t\"header = {\\\"Authorization\\\": \\\"Bearer " + userToken + "\\\"}\\n\",\n\t\"client = RayJobClient(address=ray_dashboard, headers=header, verify=False)\\n",
+		"--num-devices=8":                                        fmt.Sprintf("--num-devices=%d", numGpus),
+		"--num-epochs=3":                                         fmt.Sprintf("--num-epochs=%d", 1),
+		"--ds-config=./deepspeed_configs/zero_3_llama_2_7b.json": "--ds-config=./zero_3_llama_2_7b.json \\\"\\n\",\n\t\"               \\\"--lora-config=./lora.json \\\"\\n\",\n\t\"               \\\"--as-test",
+		"'pip': 'requirements.txt'":                              "'pip': '/opt/app-root/src/requirements.txt'",
+		"'working_dir': './'":                                    "'working_dir': '/opt/app-root/src'",
+		"client.stop_job(submission_id)":                         "finished = False\\n\",\n\t\"while not finished:\\n\",\n\t\"    time.sleep(1)\\n\",\n\t\"    status = client.get_job_status(submission_id)\\n\",\n\t\"    finished = (status == \\\"SUCCEEDED\\\")\\n\",\n\t\"if finished:\\n\",\n\t\"    print(\\\"Job completed Successfully !\\\")\\n\",\n\t\"else:\\n\",\n\t\"    print(\\\"Job failed !\\\")\\n\",\n\t\"time.sleep(10)\\n",
+	}
+
+	updatedNotebookContent := string(ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/ray_finetune_llm_deepspeed.ipynb"))
+	for oldValue, newValue := range requiredChangesInNotebook {
+		updatedNotebookContent = strings.Replace(updatedNotebookContent, oldValue, newValue, -1)
+	}
+	updatedNotebook := []byte(updatedNotebookContent)
+
+	// Test configuration
+	jupyterNotebookConfigMapFileName := "ray_finetune_llm_deepspeed.ipynb"
+	configMap := map[string][]byte{
+		jupyterNotebookConfigMapFileName: updatedNotebook,
+		"ray_finetune_llm_deepspeed.py":  ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/ray_finetune_llm_deepspeed.py"),
+		"requirements.txt":               ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/requirements.txt"),
+		"create_dataset.py":              ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/create_dataset.py"),
+		"lora.json":                      ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/lora_configs/lora.json"),
+		"zero_3_llama_2_7b.json":         ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/deepspeed_configs/zero_3_llama_2_7b.json"),
+		"utils.py":                       ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/utils.py"),
+	}
+
+	config := CreateConfigMap(test, namespace.Name, configMap)
+
 	// Create Notebook CR
 	createNotebook(test, namespace, userToken, config.Name, jupyterNotebookConfigMapFileName, numGpus)
 
@@ -77,26 +112,6 @@ func mnistRayLlmFinetune(t *testing.T, numGpus int) {
 		)
 
 	// Make sure the RayCluster finishes and is deleted
-	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutGpuProvisioning).
+	test.Eventually(RayClusters(test, namespace.Name), TestTimeoutMedium).
 		Should(HaveLen(0))
 }
-
-func ReadRayFinetuneRequirementsTxt(test Test) []byte {
-	// Read the requirements.txt from resources and perform replacements for custom values using go template
-	props := struct {
-		PipIndexUrl    string
-		PipTrustedHost string
-	}{
-		PipIndexUrl: "--index " + string(GetPipIndexURL()),
-	}
-
-	// Provide trusted host only if defined
-	if len(GetPipTrustedHost()) > 0 {
-		props.PipTrustedHost = "--trusted-host " + GetPipTrustedHost()
-	}
-
-	template, err := files.ReadFile("resources/ray_finetune_demo/ray_finetune_requirements.txt")
-	test.Expect(err).NotTo(HaveOccurred())
-
-	return ParseTemplate(test, template, props)
-}

tests/odh/resources/custom-nb-small.yaml

@@ -50,8 +50,20 @@ spec:
           value: {{.NotebookImage}}
         - name: JUPYTER_NOTEBOOK_PORT
           value: "8888"
+        - name: AWS_ACCESS_KEY_ID
+          value: {{.S3AccessKeyId}}
+        - name: AWS_SECRET_ACCESS_KEY
+          value: {{.S3SecretAccessKey}}
+        - name: AWS_DEFAULT_REGION
+          value: {{.S3DefaultRegion}}
+        - name: AWS_S3_BUCKET
+          value: {{.S3BucketName}}
+        - name: PIP_INDEX_URL
+          value: {{.PipIndexUrl}}
+        - name: PIP_TRUSTED_HOST
+          value: {{.PipTrustedHost}}
         image: {{.NotebookImage}}
-        command: ["/bin/sh", "-c", "pip install papermill && papermill /opt/app-root/notebooks/{{.NotebookConfigMapFileName}} /opt/app-root/src/mcad-out.ipynb -p namespace {{.Namespace}} -p ray_image {{.RayImage}} -p openshift_api_url {{.OpenShiftApiUrl}} -p kubernetes_user_bearer_token {{.KubernetesUserBearerToken}} -p num_gpus {{ .NumGpus }} -p s3_bucket_name {{.S3BucketName}} -p s3_access_key_id {{.S3AccessKeyId}} -p s3_secret_access_key {{.S3SecretAccessKey}} -p s3_default_region {{.S3DefaultRegion}} --log-output && sleep infinity"]
+        command: ["/bin/sh", "-c", "pip install papermill && papermill /opt/app-root/notebooks/{{.NotebookConfigMapFileName}} /opt/app-root/src/mcad-out.ipynb -p namespace {{.Namespace}} -p ray_image {{.RayImage}} -p openshift_api_url {{.OpenShiftApiUrl}} -p kubernetes_user_bearer_token {{.KubernetesUserBearerToken}} -p num_gpus {{ .NumGpus }} --log-output && sleep infinity"]
         # args: ["pip install papermill && oc login --token=${OCP_TOKEN} --server=${OCP_SERVER} --insecure-skip-tls-verify=true && papermill /opt/app-root/notebooks/mcad.ipynb /opt/app-root/src/mcad-out.ipynb" ]
         imagePullPolicy: Always
         # livenessProbe:

tests/odh/resources/hpo_raytune_requirements.txt

@@ -1 +1 @@
-torchvision==0.18.0
+torchvision==0.18.0

tests/odh/resources/mnist_hpo.py

@@ -142,4 +142,4 @@ def train_mnist(config):
 
     print("Best hyperparameters config is:", results.get_best_result().config)
 
-    assert not results.errors
+    assert not results.errors

tests/odh/resources/mnist_hpo_raytune.ipynb

@@ -1,15 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5c8a5392",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%pip install codeflare-sdk -U"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -52,7 +42,6 @@
     "#parameters\n",
     "namespace = \"default\"\n",
     "ray_image = \"has to be specified\"\n",
-    "local_queue = \"has to be specified\"\n",
     "openshift_api_url = \"has to be specified\"\n",
     "kubernetes_user_bearer_token = \"has to be specified\"\n",
     "num_gpus = \"has to be specified\""