feat: Add Dataset Image Use for E2E Tests (#548)

**Reason for Change**:
This PR ensures we use an image for dataset that is locally built during
e2e as opposed to an external URL.

This prevents flakey network errors that can waste entire pipeline runs.

Addresses - #531

.github/workflows/e2e-workflow.yml

@@ -134,6 +134,13 @@ jobs:
           make docker-build-adapter
         env:
           REGISTRY: ${{ env.CLUSTER_NAME }}.azurecr.io
+
+      - name: build dataset image
+        shell: bash
+        run: |
+          make docker-build-dataset
+        env:
+          REGISTRY: ${{ env.CLUSTER_NAME }}.azurecr.io
 
       - name: create cluster
         shell: bash
@@ -210,7 +217,7 @@ jobs:
           REGISTRY: ${{ env.REGISTRY }}
           AI_MODELS_REGISTRY: ${{ secrets.E2E_ACR_AMRT_USERNAME }}.azurecr.io
           AI_MODELS_REGISTRY_SECRET: ${{ secrets.E2E_AMRT_SECRET_NAME }}
-          ADAPTER_REGISTRY: ${{ env.CLUSTER_NAME }}.azurecr.io
+          E2E_ACR_REGISTRY: ${{ env.CLUSTER_NAME }}.azurecr.io
 
       - name: Cleanup e2e resources
         if: ${{ always() }}

Makefile

@@ -199,6 +199,15 @@ docker-build-adapter: docker-buildx
 		--pull \
 		--tag $(REGISTRY)/e2e-adapter:0.0.1 .
 
+.PHONY: docker-build-dataset
+docker-build-dataset: docker-buildx
+	docker buildx build \
+		--file ./docker/dataset/Dockerfile \
+		--output=$(OUTPUT_TYPE) \
+		--platform="linux/$(ARCH)" \
+		--pull \
+		--tag $(REGISTRY)/e2e-dataset:0.0.1 .
+
 ##@ Deployment
 
 ifndef ignore-not-found

docker/dataset/Dockerfile

@@ -0,0 +1,5 @@
+FROM busybox:latest
+
+RUN mkdir -p /data
+
+COPY docker/dataset/dataset.parquet /data/

docker/dataset/README.md

@@ -0,0 +1,21 @@
+# E2E Fine-Tuning Dataset Files
+
+## Overview
+
+This dataset file is used for conducting end-to-end (E2E) testing for fine-tuning. The Dockerfile builds an image incorporating the [dolly-15k-oai-style](https://huggingface.co/datasets/philschmid/dolly-15k-oai-style) dataset which is then used within an init container specifically for fine-tuning. 
+
+## Files
+
+- **Dockerfile**: Builds the Docker image for the E2E tests.
+
+- **dataset.parquet**: The dataset itself, downloaded from [dolly-15k-oai-style](https://huggingface.co/datasets/philschmid/dolly-15k-oai-style)
+
+
+## Usage
+
+Build the Docker image with the following command:
+
+```bash
+
+make docker-build-dataset
+

test/e2e/inference_with_adapters.go

@@ -19,7 +19,7 @@ import (
 var DefaultStrength = "1.0"
 
 var imageName = "e2e-adapter"
-var fullImageName = utils.GetEnv("ADAPTER_REGISTRY") + "/" + imageName + ":0.0.1"
+var fullImageName = utils.GetEnv("E2E_ACR_REGISTRY") + "/" + imageName + ":0.0.1"
 
 var validAdapters = []kaitov1alpha1.AdapterSpec{
 	{

test/e2e/preset_test.go

@@ -40,6 +40,11 @@ const (
 	PresetPhi3Mini128kModel      = "phi-3-mini-128k-instruct"
 )
 
+var (
+	datasetImageName     = "e2e-dataset"
+	fullDatasetImageName = utils.GetEnv("E2E_ACR_REGISTRY") + "/" + datasetImageName + ":0.0.1"
+)
+
 func loadTestEnvVars() {
 	var err error
 	runLlama13B, err = strconv.ParseBool(os.Getenv("RUN_LLAMA_13B"))
@@ -216,7 +221,7 @@ func createPhi3TuningWorkspaceWithPresetPublicMode(configMapName string, numOfNo
 		uniqueID = fmt.Sprint("preset-", rand.Intn(1000))
 		outputRegistryUrl := fmt.Sprintf("%s.azurecr.io/%s:%s", azureClusterName, e2eOutputImageName, e2eOutputImageTag)
 		workspaceObj = utils.GenerateE2ETuningWorkspaceManifest(uniqueID, namespaceName, "",
-			outputRegistryUrl, numOfNode, "Standard_NC6s_v3", &metav1.LabelSelector{
+			fullDatasetImageName, outputRegistryUrl, numOfNode, "Standard_NC6s_v3", &metav1.LabelSelector{
 				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-tuning-falcon"},
 			}, nil, PresetPhi3Mini128kModel, kaitov1alpha1.ModelImageAccessModePublic, []string{aiModelsRegistrySecret}, configMapName)
 
@@ -287,11 +292,12 @@ func validateMachineCreation(workspaceObj *kaitov1alpha1.Workspace, expectedCoun
 		Eventually(func() bool {
 			machineList, err := getAllValidMachines(workspaceObj)
 			if err != nil {
-				fmt.Printf("Failed to get all valid machines: %v", err)
+				fmt.Printf("Failed to get all valid machines: %v\n", err)
 				return false
 			}
 
 			if len(machineList.Items) != expectedCount {
+				fmt.Printf("Expected %d machines, but found %d machines\n", expectedCount, len(machineList.Items))
 				return false
 			}
 

test/e2e/utils/utils.go

@@ -271,7 +271,7 @@ func GenerateTuningWorkspaceManifest(name, namespace, imageName string, resource
 	return workspace
 }
 
-func GenerateE2ETuningWorkspaceManifest(name, namespace, imageName, outputRegistry string,
+func GenerateE2ETuningWorkspaceManifest(name, namespace, imageName, datasetImageName, outputRegistry string,
 	resourceCount int, instanceType string, labelSelector *metav1.LabelSelector,
 	preferredNodes []string, presetName kaitov1alpha1.ModelName, accessMode kaitov1alpha1.ModelImageAccessMode,
 	imagePullSecret []string, customConfigMapName string) *kaitov1alpha1.Workspace {
@@ -306,7 +306,7 @@ func GenerateE2ETuningWorkspaceManifest(name, namespace, imageName, outputRegist
 	workspace.Tuning = &workspaceTuning
 	workspace.Tuning.Method = kaitov1alpha1.TuningMethodQLora
 	workspace.Tuning.Input = &kaitov1alpha1.DataSource{
-		URLs: []string{ExampleDatasetURL},
+		Image: datasetImageName,
 	}
 	workspace.Tuning.Output = &kaitov1alpha1.DataDestination{
 		Image:           outputRegistry,