Skip to content

Commit

Permalink
feat: Add Util funcs, updating func names, logs, configs and ensure s…
Browse files Browse the repository at this point in the history 
…ervice requirement (#485)

**Reason for Change**:
Renaming functions, updating configs and services
  • Loading branch information
@ishaansehgal99
ishaansehgal99 authored Jun 28, 2024
1 parent f70dcc5 commit f534355
Show file tree
Hide file tree
Showing 9 changed files with 92 additions and 30 deletions.
17 changes: 8 additions & 9 deletions pkg/controllers/workspace_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,21 +109,20 @@ func (c *WorkspaceReconciler) addOrUpdateWorkspace(ctx context.Context, wObj *ka
return reconcile.Result{}, err
}

if err := c.ensureService(ctx, wObj); err != nil {
if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeReady, metav1.ConditionFalse,
"workspaceFailed", err.Error()); updateErr != nil {
klog.ErrorS(updateErr, "failed to update workspace status", "workspace", klog.KObj(wObj))
return reconcile.Result{}, updateErr
}
return reconcile.Result{}, err
}

if wObj.Tuning != nil {
if err = c.applyTuning(ctx, wObj); err != nil {
return reconcile.Result{}, err
}
}
if wObj.Inference != nil {
if err := c.ensureService(ctx, wObj); err != nil {
if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeReady, metav1.ConditionFalse,
"workspaceFailed", err.Error()); updateErr != nil {
klog.ErrorS(updateErr, "failed to update workspace status", "workspace", klog.KObj(wObj))
return reconcile.Result{}, updateErr
}
return reconcile.Result{}, err
}
if err = c.applyInference(ctx, wObj); err != nil {
if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeReady, metav1.ConditionFalse,
"workspaceFailed", err.Error()); updateErr != nil {
Expand Down
3 changes: 2 additions & 1 deletion pkg/inference/preset-inferences.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,8 @@ func updateTorchParamsForDistributedInference(ctx context.Context, kubeClient cl

func GetInferenceImageInfo(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, presetObj *model.PresetParam) (string, []corev1.LocalObjectReference) {
imagePullSecretRefs := []corev1.LocalObjectReference{}
if presetObj.ImageAccessMode == string(kaitov1alpha1.ModelImageAccessModePrivate) {
// Check if the workspace preset's access mode is private
if string(workspaceObj.Inference.Preset.AccessMode) == string(kaitov1alpha1.ModelImageAccessModePrivate) {
imageName := workspaceObj.Inference.Preset.PresetOptions.Image
for _, secretName := range workspaceObj.Inference.Preset.PresetOptions.ImagePullSecrets {
imagePullSecretRefs = append(imagePullSecretRefs, corev1.LocalObjectReference{Name: secretName})
Expand Down
7 changes: 6 additions & 1 deletion pkg/tuning/preset-tuning.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,8 @@ func getInstanceGPUCount(sku string) int {

func GetTuningImageInfo(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, presetObj *model.PresetParam) (string, []corev1.LocalObjectReference) {
imagePullSecretRefs := []corev1.LocalObjectReference{}
if presetObj.ImageAccessMode == string(kaitov1alpha1.ModelImageAccessModePrivate) {
// Check if the workspace preset's access mode is private
if string(workspaceObj.Tuning.Preset.AccessMode) == string(kaitov1alpha1.ModelImageAccessModePrivate) {
imageName := workspaceObj.Tuning.Preset.PresetOptions.Image
for _, secretName := range workspaceObj.Tuning.Preset.PresetOptions.ImagePullSecrets {
imagePullSecretRefs = append(imagePullSecretRefs, corev1.LocalObjectReference{Name: secretName})
Expand Down Expand Up @@ -401,7 +402,11 @@ func handleURLDataSource(ctx context.Context, workspaceObj *kaitov1alpha1.Worksp
Command: []string{"sh", "-c", `
for url in $DATA_URLS; do
filename=$(basename "$url" | sed 's/[?=&]/_/g')
echo "Downloading $url to $DATA_VOLUME_PATH/$filename"
curl -sSL $url -o $DATA_VOLUME_PATH/$filename
if [ $? -ne 0 ]; then
echo "Failed to download $url"
fi
done
`},
VolumeMounts: []corev1.VolumeMount{
Expand Down
2 changes: 1 addition & 1 deletion pkg/tuning/preset-tuning_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ func TestHandleURLDataSource(t *testing.T) {
},
expectedInitContainerName: "data-downloader",
expectedImage: "curlimages/curl",
expectedCommands: "filename=$(basename \"$url\" | sed 's/[?=&]/_/g')\ncurl -sSL $url -o $DATA_VOLUME_PATH/$filename",
expectedCommands: "curl -sSL $url -o $DATA_VOLUME_PATH/$filename",
expectedVolumeName: "data-volume",
expectedVolumeMountPath: utils.DefaultDataVolumePath,
},
Expand Down
1 change: 0 additions & 1 deletion presets/inference/text-generation/requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ numpy==1.22.4
bitsandbytes==0.42.0

# Less critical, can be latest
deepspeed
gputil
psutil
# For UTs
Expand Down
1 change: 0 additions & 1 deletion presets/tuning/text-generation/requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ peft==0.11.1
bitsandbytes==0.42.0

# Less critical, can be latest
deepspeed
loralib
einops
xformers
Expand Down
14 changes: 8 additions & 6 deletions test/e2e/preset_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,15 @@ const (
PresetMistral7BModel = "mistral-7b"
PresetMistral7BInstructModel = "mistral-7b-instruct"
PresetPhi2Model = "phi-2"
PresetPhi3Mini4kModel = "phi-3-mini-4k-instruct"
PresetPhi3Mini128kModel = "phi-3-mini-128k-instruct"
)

func createFalconWorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Workspace {
workspaceObj := &kaitov1alpha1.Workspace{}
By("Creating a workspace CR with Falcon 7B preset public mode", func() {
uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "", numOfNode, "Standard_NC12s_v3",
workspaceObj = utils.GenerateInferenceWorkspaceManifest(uniqueID, namespaceName, "", numOfNode, "Standard_NC12s_v3",
&metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-falcon"},
}, nil, PresetFalcon7BModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
Expand All @@ -53,7 +55,7 @@ func createMistralWorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Wo
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",
workspaceObj = utils.GenerateInferenceWorkspaceManifest(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)
Expand All @@ -67,7 +69,7 @@ func createPhi2WorkspaceWithPresetPublicMode(numOfNode int) *kaitov1alpha1.Works
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",
workspaceObj = utils.GenerateInferenceWorkspaceManifest(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)
Expand All @@ -81,7 +83,7 @@ func createLlama7BWorkspaceWithPresetPrivateMode(registry, registrySecret, image
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),
workspaceObj = utils.GenerateInferenceWorkspaceManifest(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-llama-2-7b"},
}, nil, PresetLlama2AChat, kaitov1alpha1.ModelImageAccessModePrivate, []string{registrySecret}, nil)
Expand All @@ -95,7 +97,7 @@ func createLlama13BWorkspaceWithPresetPrivateMode(registry, registrySecret, imag
workspaceObj := &kaitov1alpha1.Workspace{}
By("Creating a workspace CR with Llama 13B Chat preset private mode", func() {
uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, fmt.Sprintf("%s/%s:%s", registry, PresetLlama2BChat, imageVersion),
workspaceObj = utils.GenerateInferenceWorkspaceManifest(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-llama-2-13b"},
}, nil, PresetLlama2BChat, kaitov1alpha1.ModelImageAccessModePrivate, []string{registrySecret}, nil)
Expand All @@ -109,7 +111,7 @@ func createCustomWorkspaceWithPresetCustomMode(imageName string, numOfNode int)
workspaceObj := &kaitov1alpha1.Workspace{}
By("Creating a workspace CR with custom workspace mode", func() {
uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "",
workspaceObj = utils.GenerateInferenceWorkspaceManifest(uniqueID, namespaceName, "",
numOfNode, "Standard_D4s_v3", &metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test-custom"},
}, nil, "", utils.InferenceModeCustomTemplate, nil, utils.GeneratePodTemplate(uniqueID, namespaceName, imageName, nil))
Expand Down
69 changes: 63 additions & 6 deletions test/e2e/utils/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,22 @@
package utils

import (
"context"
"fmt"
"io"
"io/ioutil"
"math/rand"
"os"
"strings"
"time"

kaitov1alpha1 "github.com/azure/kaito/api/v1alpha1"
"github.com/samber/lo"
"gopkg.in/yaml.v2"
corev1 "k8s.io/api/core/v1"
v1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
)

const (
Expand All @@ -36,6 +42,13 @@ func GetEnv(envVar string) string {
return env
}

// GenerateRandomString generates a random number between 0 and 1000 and returns it as a string.
func GenerateRandomString() string {
rand.Seed(time.Now().UnixNano()) // Seed the random number generator
randomNumber := rand.Intn(1001) // Generate a random number between 0 and 1000
return fmt.Sprintf("%d", randomNumber)
}

func GetModelConfigInfo(configFilePath string) (map[string]interface{}, error) {
var data map[string]interface{}

Expand All @@ -52,6 +65,50 @@ func GetModelConfigInfo(configFilePath string) (map[string]interface{}, error) {
return data, nil
}

func GetPodNameForJob(coreClient *kubernetes.Clientset, namespace, jobName string) (string, error) {
podList, err := coreClient.CoreV1().Pods(namespace).List(context.TODO(), metav1.ListOptions{
LabelSelector: fmt.Sprintf("job-name=%s", jobName),
})
if err != nil {
return "", err
}

if len(podList.Items) == 0 {
return "", fmt.Errorf("no pods found for job %s", jobName)
}

return podList.Items[0].Name, nil
}

func GetPodLogs(coreClient *kubernetes.Clientset, namespace, podName, containerName string) (string, error) {
req := coreClient.CoreV1().Pods(namespace).GetLogs(podName, &v1.PodLogOptions{Container: containerName})
logs, err := req.Stream(context.Background())
if err != nil {
return "", err
}
defer logs.Close()

buf := new(strings.Builder)
_, err = io.Copy(buf, logs)
if err != nil {
return "", err
}

return buf.String(), nil
}

func CopySecret(original *corev1.Secret, targetNamespace string) *corev1.Secret {
newSecret := &corev1.Secret{
ObjectMeta: metav1.ObjectMeta{
Name: original.Name,
Namespace: targetNamespace,
},
Data: original.Data,
Type: original.Type,
}
return newSecret
}

func ExtractModelVersion(configs map[string]interface{}) (map[string]string, error) {
modelsInfo := make(map[string]string)
models, ok := configs["models"].([]interface{})
Expand Down Expand Up @@ -81,9 +138,9 @@ func ExtractModelVersion(configs map[string]interface{}) (map[string]string, err
return modelsInfo, nil
}

func GenerateWorkspaceManifest(name, namespace, imageName string, resourceCount int, instanceType string,
func GenerateInferenceWorkspaceManifest(name, namespace, imageName string, resourceCount int, instanceType string,
labelSelector *metav1.LabelSelector, preferredNodes []string, presetName kaitov1alpha1.ModelName,
inferenceMode kaitov1alpha1.ModelImageAccessMode, imagePullSecret []string,
accessMode kaitov1alpha1.ModelImageAccessMode, imagePullSecret []string,
podTemplate *corev1.PodTemplateSpec) *kaitov1alpha1.Workspace {

workspace := &kaitov1alpha1.Workspace{
Expand All @@ -100,20 +157,20 @@ func GenerateWorkspaceManifest(name, namespace, imageName string, resourceCount
}

var workspaceInference kaitov1alpha1.InferenceSpec
if inferenceMode == kaitov1alpha1.ModelImageAccessModePublic ||
inferenceMode == kaitov1alpha1.ModelImageAccessModePrivate {
if accessMode == kaitov1alpha1.ModelImageAccessModePublic ||
accessMode == kaitov1alpha1.ModelImageAccessModePrivate {
workspaceInference.Preset = &kaitov1alpha1.PresetSpec{
PresetMeta: kaitov1alpha1.PresetMeta{
Name: presetName,
AccessMode: inferenceMode,
AccessMode: accessMode,
},
PresetOptions: kaitov1alpha1.PresetOptions{
Image: imageName,
ImagePullSecrets: imagePullSecret,
},
}
}
if inferenceMode == InferenceModeCustomTemplate {
if accessMode == InferenceModeCustomTemplate {
workspaceInference.Template = podTemplate
}

Expand Down
8 changes: 4 additions & 4 deletions test/e2e/webhook_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ import (

var _ = Describe("Workspace Validation Webhook", func() {
It("should validate the workspace resource spec at creation ", func() {
workspaceObj := utils.GenerateWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_Bad",
workspaceObj := utils.GenerateInferenceWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_Bad",
&metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "webhook-e2e-test"},
}, nil, PresetFalcon7BModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
Expand All @@ -33,7 +33,7 @@ var _ = Describe("Workspace Validation Webhook", func() {
})

It("should validate the workspace inference spec at creation ", func() {
workspaceObj := utils.GenerateWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC6",
workspaceObj := utils.GenerateInferenceWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC6",
&metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "webhook-e2e-test"},
}, nil, "invalid-name", kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
Expand All @@ -51,7 +51,7 @@ var _ = Describe("Workspace Validation Webhook", func() {
//TODO custom template

It("should validate the workspace resource spec at update ", func() {
workspaceObj := utils.GenerateWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC12s_v3",
workspaceObj := utils.GenerateInferenceWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC12s_v3",
&metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "webhook-e2e-test"},
}, nil, PresetFalcon7BModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
Expand Down Expand Up @@ -94,7 +94,7 @@ var _ = Describe("Workspace Validation Webhook", func() {
})

It("should validate the workspace inference spec at update ", func() {
workspaceObj := utils.GenerateWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC12s_v3",
workspaceObj := utils.GenerateInferenceWorkspaceManifest(fmt.Sprint("webhook-", rand.Intn(1000)), namespaceName, "", 1, "Standard_NC12s_v3",
&metav1.LabelSelector{
MatchLabels: map[string]string{"kaito-workspace": "webhook-e2e-test"},
}, nil, PresetFalcon7BModel, kaitov1alpha1.ModelImageAccessModePublic, nil, nil)
Expand Down

0 comments on commit f534355

Please sign in to comment.