feat: Added front end to kaito

pkg/controllers/workspace_controller.go

@@ -116,6 +116,17 @@ func (c *WorkspaceReconciler) addOrUpdateWorkspace(ctx context.Context, wObj *ka
 			return reconcile.Result{}, updateErr
 		}
 		return reconcile.Result{}, err
+	} else {
+		// Given service and inference were successfully applied
+		wAnnotation := wObj.GetAnnotations()
+		if len(wAnnotation) != 0 {
+			val, found := wAnnotation[kaitov1alpha1.AnnotationEnableSampleFrontEnd]
+			if found && val == "True" {
+				if c.createSampleFrontEnd(ctx, wObj); err != nil {
+					klog.ErrorS(err, "failed to create sample front end", "workspace", klog.KObj(wObj))
+				}
+			}
+		}
 	}
 
 	// TODO apply TrainingSpec
@@ -450,7 +461,6 @@ func (c *WorkspaceReconciler) applyInference(ctx context.Context, wObj *kaitov1a
 				existingObj = &appsv1.StatefulSet{}
 			} else {
 				existingObj = &appsv1.Deployment{}
-
 			}
 
 			if err = resources.GetResource(ctx, wObj.Name, wObj.Namespace, c.Client, existingObj); err == nil {
@@ -491,6 +501,27 @@ func (c *WorkspaceReconciler) applyInference(ctx context.Context, wObj *kaitov1a
 	return nil
 }
 
+func (c *WorkspaceReconciler) createSampleFrontEnd(ctx context.Context, wObj *kaitov1alpha1.Workspace) error {
+	deploymentName := wObj.Name + "-chainlit-sample-frontend"
+	if err := resources.GetResource(ctx, deploymentName, wObj.Namespace, c.Client, &appsv1.Deployment{}); err == nil {
+		klog.InfoS("A sample front end already exists for workspace", "workspace", klog.KObj(wObj))
+		if err = resources.CheckResourceStatus(&appsv1.Deployment{}, c.Client, time.Duration(10)*time.Minute); err != nil {
+			return err
+		}
+	} else if apierrors.IsNotFound(err) {
+		var frontEndObj client.Object
+		// Need to create a new workload
+		frontEndObj, err = inference.CreatePresetFrontEnd(ctx, wObj, c.Client)
+		if err != nil {
+			return err
+		}
+		if err = resources.CheckResourceStatus(frontEndObj, c.Client, time.Duration(10)*time.Minute); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (c *WorkspaceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	c.Recorder = mgr.GetEventRecorderFor("Workspace")

pkg/inference/preset-inferences.go

@@ -137,6 +137,15 @@ func CreatePresetInference(ctx context.Context, workspaceObj *kaitov1alpha1.Work
 	return depObj, nil
 }
 
+func CreatePresetFrontEnd(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, kubeClient client.Client) (client.Object, error) {
+	depObj := resources.GenerateFrontEndManifest(ctx, workspaceObj)
+	err := resources.CreateResource(ctx, depObj, kubeClient)
+	if client.IgnoreAlreadyExists(err) != nil {
+		return nil, err
+	}
+	return depObj, nil
+}
+
 // prepareInferenceParameters builds a PyTorch command:
 // torchrun <TORCH_PARAMS> <OPTIONAL_RDZV_PARAMS> baseCommand <MODEL_PARAMS>
 // and sets the GPU resources required for inference.

pkg/resources/manifests.go

@@ -5,6 +5,7 @@ package resources
 import (
 	"context"
 	"fmt"
+	"k8s.io/utils/pointer"
 
 	"k8s.io/apimachinery/pkg/util/intstr"
 
@@ -53,6 +54,66 @@ func GenerateHeadlessServiceManifest(ctx context.Context, workspaceObj *kaitov1a
 	}
 }
 
+func GenerateFrontEndManifest(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace) *appsv1.Deployment {
+	deploymentName := workspaceObj.Name + "-chainlit-sample-frontend"
+	serviceDNSName := workspaceObj.Name + "." + workspaceObj.Namespace + ".svc.cluster.local"
+
+	return &appsv1.Deployment{
+		ObjectMeta: v1.ObjectMeta{
+			Name:      deploymentName,
+			Namespace: workspaceObj.Namespace,
+			OwnerReferences: []v1.OwnerReference{
+				{
+					APIVersion: kaitov1alpha1.GroupVersion.String(),
+					Kind:       "Workspace",
+					Name:       workspaceObj.Name,
+					UID:        workspaceObj.UID,
+					Controller: pointer.BoolPtr(true),
+				},
+			},
+		},
+		Spec: appsv1.DeploymentSpec{
+			Replicas: pointer.Int32Ptr(1),
+			Selector: &v1.LabelSelector{
+				MatchLabels: map[string]string{
+					"app": deploymentName,
+				},
+			},
+			Template: corev1.PodTemplateSpec{
+				ObjectMeta: v1.ObjectMeta{
+					Labels: map[string]string{
+						"app": deploymentName,
+					},
+				},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name:  "chainlit-frontend",
+							Image: "python:3.10-slim",
+							Command: []string{
+								"/bin/sh", "-c",
+							},
+							Args: []string{
+								"apt-get update && apt-get install -y wget && " +
+									"pip install chainlit requests && " +
+									"wget -O /app/frontend/chainlit.py https://raw.githubusercontent.com/Azure/kaito/main/presets/inference/text-generation/frontend/chainlit.py && " +
+									"chainlit run frontend/chainlit.py -w",
+							},
+							Env: []corev1.EnvVar{
+								{
+									Name:  "WORKSPACE_SERVICE_URL",
+									Value: "http://" + serviceDNSName + ":80",
+								},
+							},
+							WorkingDir: "/app",
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
 func GenerateServiceManifest(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, serviceType corev1.ServiceType, isStatefulSet bool) *corev1.Service {
 	selector := map[string]string{
 		kaitov1alpha1.LabelWorkspaceName: workspaceObj.Name,

presets/inference/text-generation/frontend/chainlit.py

@@ -1,11 +1,11 @@
 import chainlit as cl
 import requests
+import os
 
-
+URL = os.environ.get('WORKSPACE_SERVICE_URL')
 @cl.step
 def inference(prompt):
     # Endpoint URL
-    url = "http://localhost:5000/chat"
     data = {
         "prompt": prompt,
         "return_full_text": False,
@@ -15,15 +15,15 @@ def inference(prompt):
         }
     }
 
-    response = requests.post(url, json=data)
+    response = requests.post(URL, json=data)
 
     if response.status_code == 200:
         response_data = response.json()
         return response_data.get("Result", "No result found")
     else:
         return f"Error: Received response code {response.status_code}"
 
-@cl.on_message  # this function will be called every time a user inputs a message in the UI
+@cl.on_message
 async def main(message: cl.Message):
     """
     This function is called every time a user inputs a message in the UI.

presets/test/manifests/frontend/frontend.yaml

@@ -0,0 +1,28 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: my-workspace-chainlit-frontend-deployment
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: my-workspace-chainlit-frontend
+  template:
+    metadata:
+      labels:
+        app: my-workspace-chainlit-frontend
+    spec:
+      containers:
+        - name: chainlit-frontend
+          image: python:3.10-slim
+          command: ["/bin/sh"]
+          args:
+            - -c
+            - |
+              pip install chainlit requests && \
+              wget -O /app/frontend/chainlit.py https://raw.githubusercontent.com/Azure/kaito/main/presets/inference/text-generation/frontend/chainlit.py && \
+              chainlit run frontend/chainlit.py -w
+          env:
+            - name: WORKSPACE_SERVICE_URL
+              value: http://my-workspace.my-namespace.svc.cluster.local:80
+          workingDir: /app