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,6 +2,10 @@
 import json
 import os
 
+cache_dir="../../datasets"
+if not os.path.exists(cache_dir):
+    cache_dir=""
+dataset = load_dataset("gsm8k", "main", cache_dir=cache_dir)
 
 def gsm8k_qa_tokens_template():
     dataset = load_dataset("gsm8k", "main")

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()"
@@ -125,7 +125,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\""
@@ -152,12 +152,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

@@ -601,6 +601,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.")
@@ -669,7 +672,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/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)
-}