feat: Add additional e2e tests (#256)

Add e2e tests for mistral and phi-2 Note: This PR also: 1. Checks deletion timestamp - to ensure nodes being deleted aren't considered qualified 2. Corrects mistral version from 0.0.1 -> 0.0.2. 0.0.2 is actually the latest/correct version --------- Signed-off-by: Ishaan Sehgal <[email protected]>
kaito-project · Feb 28, 2024 · 7bd5641 · 7bd5641
1 parent de2b45c
commit 7bd5641
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 7 deletions.
diff --git a/pkg/controllers/workspace_controller.go b/pkg/controllers/workspace_controller.go
@@ -270,13 +270,18 @@ func (c *WorkspaceReconciler) getAllQualifiedNodes(ctx context.Context, wObj *ka
 	if err != nil {
 		return nil, err
 	}
+
 	if len(nodeList.Items) == 0 {
 		klog.InfoS("no current nodes match the workspace resource spec", "workspace", klog.KObj(wObj))
 		return nil, nil
 	}
 
 	for index := range nodeList.Items {
 		nodeObj := nodeList.Items[index]
+		// Skip nodes that are being deleted
+		if nodeObj.DeletionTimestamp != nil {
+			continue
+		}
 		foundInstanceType := c.validateNodeInstanceType(ctx, wObj, lo.ToPtr(nodeObj))
 		_, statusRunning := lo.Find(nodeObj.Status.Conditions, func(condition corev1.NodeCondition) bool {
 			return condition.Type == corev1.NodeReady && condition.Status == corev1.ConditionTrue

diff --git a/pkg/controllers/workspace_controller_test.go b/pkg/controllers/workspace_controller_test.go
@@ -8,6 +8,7 @@ import (
 	"reflect"
 	"sort"
 	"testing"
+	"time"
 
 	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
 	"github.com/azure/kaito/api/v1alpha1"
@@ -492,6 +493,17 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 		"Gets all qualified nodes": {
 			callMocks: func(c *utils.MockClient) {
 				nodeList := utils.MockNodeList
+				deletedNode := corev1.Node{
+					ObjectMeta: v1.ObjectMeta{
+						Name: "node4",
+						Labels: map[string]string{
+							corev1.LabelInstanceTypeStable: "Standard_NC12s_v3",
+						},
+						DeletionTimestamp: &v1.Time{Time: time.Now()},
+					},
+				}
+				nodeList.Items = append(nodeList.Items, deletedNode)
+
 				relevantMap := c.CreateMapWithType(nodeList)
 				//insert node objects into the map
 				for _, obj := range utils.MockNodeList.Items {

diff --git a/presets/models/mistral/model.go b/presets/models/mistral/model.go
@@ -27,8 +27,8 @@ var (
 	PresetMistral7BInstructModel  = PresetMistral7BModel + "-instruct"
 
 	PresetMistralTagMap = map[string]string{
-		"Mistral7B":          "0.0.1",
-		"Mistral7BInstruct":  "0.0.1",
+		"Mistral7B":          "0.0.2",
+		"Mistral7BInstruct":  "0.0.2",
 	}
 
 	baseCommandPresetMistral = "accelerate launch"

diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go
@@ -30,6 +30,9 @@ const (
 	PresetLlama2BChat    = "llama-2-13b-chat"
 	PresetFalcon7BModel  = "falcon-7b"
 	PresetFalcon40BModel = "falcon-40b"
+	PresetMistral7BModel = "mistral-7b"
+	PresetMistral7BInstructModel = "mistral-7b-instruct"
+	PresetPhi2Model = "phi-2"
 )
 
 func createFalconWorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Workspace {
@@ -38,21 +41,49 @@ func createFalconWorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Wor
 		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
 		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "", numOfNode, "Standard_NC12s_v3",
 			&metav1.LabelSelector{
-				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test"},
+				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-falcon"},
 			}, nil, PresetFalcon7BModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
 
 		createAndValidateWorkspace(workspaceObj)
 	})
 	return workspaceObj
 }
 
+func createMistralWorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Workspace {
+	workspaceObj := &kaitov1alpha1.Workspace{}
+	By("Creating a workspace CR with Mistral 7B preset public mode", func() {
+		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
+		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "", numOfNode, "Standard_NC12s_v3",
+			&metav1.LabelSelector{
+				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-mistral"},
+			}, nil, PresetMistral7BInstructModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
+
+		createAndValidateWorkspace(workspaceObj)
+	})
+	return workspaceObj
+}
+
+func createPhi2WorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Workspace {
+	workspaceObj := &kaitov1alpha1.Workspace{}
+	By("Creating a workspace CR with Phi 2 preset public mode", func() {
+		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
+		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "", numOfNode, "Standard_NC6s_v3",
+			&metav1.LabelSelector{
+				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-phi-2"},
+			}, nil, PresetPhi2Model, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
+
+		createAndValidateWorkspace(workspaceObj)
+	})
+	return workspaceObj
+}
+
 func createLlama7BWorkspaceWithPresetPrivateMode(registry, registrySecret, imageVersion string, numOfNode int) *kaitov1alpha1.Workspace {
 	workspaceObj := &kaitov1alpha1.Workspace{}
 	By("Creating a workspace CR with Llama 7B Chat preset private mode", func() {
 		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
 		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, fmt.Sprintf("%s/%s:%s", registry, PresetLlama2AChat, imageVersion),
 			numOfNode, "Standard_NC12s_v3", &metav1.LabelSelector{
-				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test"},
+				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test-llama-2-7b"},
 			}, nil, PresetLlama2AChat, kaitov1alpha1.ModelImageAccessModePrivate, []string{registrySecret}, nil)
 
 		createAndValidateWorkspace(workspaceObj)
@@ -66,7 +97,7 @@ func createLlama13BWorkspaceWithPresetPrivateMode(registry, registrySecret, imag
 		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
 		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, fmt.Sprintf("%s/%s:%s", registry, PresetLlama2BChat, imageVersion),
 			numOfNode, "Standard_NC12s_v3", &metav1.LabelSelector{
-				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test"},
+				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test-llama-2-13b"},
 			}, nil, PresetLlama2BChat, kaitov1alpha1.ModelImageAccessModePrivate, []string{registrySecret}, nil)
 
 		createAndValidateWorkspace(workspaceObj)
@@ -80,7 +111,7 @@ func createCustomWorkspaceWithPresetCustomMode(imageName string, numOfNode int)
 		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
 		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "",
 			numOfNode, "Standard_D4s_v3", &metav1.LabelSelector{
-				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test"},
+				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test-custom"},
 			}, nil, "", utils.InferenceModeCustomTemplate, nil, utils.GeneratePodTemplate(uniqueID, namespaceName, imageName, nil))
 
 		createAndValidateWorkspace(workspaceObj)
@@ -335,7 +366,7 @@ var _ = Describe("Workspace Preset", func() {
 		}
 	})
 
-	It("should create a workspace with preset public mode successfully", func() {
+	It("should create a falcon workspace with preset public mode successfully", func() {
 		numOfNode := 1
 		workspaceObj := createFalconWorkspaceWithPresetPublicMode(numOfNode)
 
@@ -354,6 +385,45 @@ var _ = Describe("Workspace Preset", func() {
 		validateWorkspaceReadiness(workspaceObj)
 	})
 
+	It("should create a mistral workspace with preset public mode successfully", func() {
+		numOfNode := 1
+		workspaceObj := createMistralWorkspaceWithPresetPublicMode(numOfNode)
+
+		defer cleanupResources(workspaceObj)
+		time.Sleep(30 * time.Second)
+
+		validateMachineCreation(workspaceObj, numOfNode)
+		validateResourceStatus(workspaceObj)
+
+		time.Sleep(30 * time.Second)
+
+		validateAssociatedService(workspaceObj)
+
+		validateInferenceResource(workspaceObj, int32(numOfNode), false)
+
+		validateWorkspaceReadiness(workspaceObj)
+	})
+
+
+	It("should create a Phi-2 workspace with preset public mode successfully", func() {
+		numOfNode := 1
+		workspaceObj := createPhi2WorkspaceWithPresetPublicMode(numOfNode)
+
+		defer cleanupResources(workspaceObj)
+		time.Sleep(30 * time.Second)
+
+		validateMachineCreation(workspaceObj, numOfNode)
+		validateResourceStatus(workspaceObj)
+
+		time.Sleep(30 * time.Second)
+
+		validateAssociatedService(workspaceObj)
+
+		validateInferenceResource(workspaceObj, int32(numOfNode), false)
+
+		validateWorkspaceReadiness(workspaceObj)
+	})
+
 	It("should create a llama 7b workspace with preset private mode successfully", func() {
 		numOfNode := 1
 		modelVersion, ok := modelInfo[PresetLlama2AChat]