Reused example finetune demo files and added provision for job status assert

Author: abhijeet-dhumal

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

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

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/notebook.go

@@ -57,11 +57,19 @@ type NotebookProps struct {
 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, exists := GetStorageBucketName()
 	s3AccessKeyId, _ := GetStorageBucketAccessKeyId()
 	s3SecretAccessKey, _ := GetStorageBucketSecretKey()
 	s3DefaultRegion, _ := GetStorageBucketDefaultRegion()
 
+	if !exists {
+		println("Storage bucket doesn't exists!")
+		s3BucketName = "\"\""
+		s3AccessKeyId = "\"\""
+		s3SecretAccessKey = "\"\""
+		s3DefaultRegion = "\"\""
+	}
+
 	// Read the Notebook CR from resources and perform replacements for custom values using go template
 	notebookProps := NotebookProps{
 		IngressDomain:             GetOpenShiftIngressDomain(test),

tests/odh/ray_finetune_llm_deepspeed_test.go

@@ -17,11 +17,19 @@ limitations under the License.
 package odh
 
 import (
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
 	"testing"
+	"time"
 
 	. "github.com/onsi/gomega"
 	. "github.com/project-codeflare/codeflare-common/support"
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 func TestRayFinetuneDemo(t *testing.T) {
@@ -33,20 +41,19 @@ func mnistRayLlmFinetune(t *testing.T, numGpus int) {
 
 	// Create a namespace
 	namespace := test.NewTestNamespace()
+	var workingDirectory, _ = os.Getwd()
 
 	// 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_llm_deepspeed.py":  ReadFileExt(test, workingDirectory+"/../../examples/ray-finetune-llm-deepspeed/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"),
+		"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)
@@ -75,9 +82,77 @@ func mnistRayLlmFinetune(t *testing.T, numGpus int) {
 				ContainElement(WithTransform(RayClusterState, Equal(rayv1.Ready))),
 			),
 		)
+	time.Sleep(30 * time.Second)
+
+	rayClusters, _ := test.Client().Ray().RayV1().RayClusters(namespace.Name).List(test.Ctx(), metav1.ListOptions{})
+	test.Expect(len(rayClusters.Items)).To(BeNumerically(">", 0))
+	dashboardName := "ray-dashboard-" + rayClusters.Items[0].Name
+	fmt.Printf("Raycluster created : %s\n", rayClusters.Items[0].Name)
+	route := GetRoute(test, namespace.Name, dashboardName)
+	hostname := route.Status.Ingress[0].Host
+
+	// Wait for expected HTTP code
+	fmt.Printf("Waiting for Route %s/%s to be available...\n", route.Namespace, route.Name)
+	tr := &http.Transport{
+		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		Proxy:           http.ProxyFromEnvironment,
+	}
+	client := &http.Client{Transport: tr}
+	req, err := http.NewRequest("GET", "https://"+hostname+"/api/jobs/", nil)
+	if err != nil {
+		test.T().Fatal(err)
+	}
+	req.Header.Add("Authorization", "Bearer "+test.Config().BearerToken)
+
+	resp, err := client.Do(req)
+	test.Expect(err).ToNot(HaveOccurred())
+	test.Expect(resp.StatusCode).ToNot(Equal(503))
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	test.Expect(err).ToNot(HaveOccurred())
+
+	var resp_json []map[string]interface{}
+	err = json.Unmarshal(body, &resp_json)
+	test.Expect(err).ToNot(HaveOccurred())
+	if len(resp_json) > 0 {
+		fmt.Printf("Job is submitted in the raycluster!\nSubmission-ID : %s\n", resp_json[0]["submission_id"])
+	}
+
+	var status string
+	var prevStatus string
+	fmt.Printf("Waiting for job to be Succeeded...\n")
+	for status != "SUCCEEDED" {
+		resp, err := client.Do(req)
+		test.Expect(err).ToNot(HaveOccurred())
+		body, err := io.ReadAll(resp.Body)
+		test.Expect(err).ToNot(HaveOccurred())
+		var result []map[string]interface{}
+		if err := json.Unmarshal(body, &result); err != nil {
+			time.Sleep(2 * time.Second)
+			break
+		}
+		if status, ok := result[0]["status"].(string); ok {
+			if prevStatus != status {
+				fmt.Printf("JobStatus : %s...\n", status)
+				prevStatus = status
+			}
+			if status == "SUCCEEDED" {
+				prevStatus = status
+				break
+			}
+			prevStatus = status
+		} else {
+			test.T().Logf("Status key not found or not a string")
+		}
+		time.Sleep(3 * time.Second)
+	}
+	if prevStatus != "SUCCEEDED" {
+		fmt.Printf("Job failed!")
+	}
+	test.Expect(prevStatus).To(Equal("SUCCEEDED"))
 
 	// 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))
 }
 

tests/odh/resources/mnist_ray_mini.ipynb

@@ -1,15 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "df737457",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%pip install codeflare-sdk -U"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -53,8 +43,7 @@
     "namespace = \"default\"\n",
     "openshift_api_url = \"has to be specified\"\n",
     "kubernetes_user_bearer_token = \"has to be specified\"\n",
-    "num_gpus = \"has to be specified\"\n",
-    "print(\"*\"*8, namespace,openshift_api_url,kubernetes_user_bearer_token, num_gpus)"
+    "num_gpus = \"has to be specified\""
    ]
   },
   {

tests/odh/resources/ray_finetune_demo/create_dataset.py

@@ -135,7 +135,7 @@
    "outputs": [],
    "source": [
     "directory_path = os.path.expanduser(\"~/.codeflare/resources/\")\n",
-    "outfile = os.path.join(directory_path, \"mnisttest.yaml\")\n",
+    "outfile = os.path.join(directory_path, \"ray-finetune-test.yaml\")\n",
     "cluster_yaml = None\n",
     "with open(outfile) as f:\n",
     "    cluster_yaml = yaml.load(f, yaml.FullLoader)\n",
@@ -201,12 +201,13 @@
     "    entrypoint=\"python ray_finetune_llm_deepspeed.py \"\n",
     "               \"--model-name=meta-llama/Llama-2-7b-chat-hf \"\n",
     "               \"--lora \"\n",
-    "               f\"--num-devices=1 \"\n",
+    "               \"--num-devices=1 \"\n",
     "               \"--num-epochs=1 \"\n",
-    "               f\"--ds-config=zero_3_llama_2_7b.json \"\n",
     "               f\"--storage-path=s3://{s3_bucket_name}/ray-finetune-llm-deepspeed3/\"\n",
     "               \"--batch-size-per-device=32 \"\n",
     "               \"--eval-batch-size-per-device=32 \"\n",
+    "               \"--ds-config=./zero_3_llama_2_7b.json \"\n",
+    "               \"--lora-config=./lora.json \"\n",
     "               \"--as-test \",\n",
     "    runtime_env={\n",
     "        \"env_vars\": {\n",
@@ -236,6 +237,7 @@
     "    finished = (status == \"SUCCEEDED\")\n",
     "if finished:\n",
     "    print(\"Job completed Successfully !\")\n",
+    "    sleep(10)\n",
     "else:\n",
     "    print(\"Job failed !\")"
    ]