feat(model): implement ManagedModel controller

config/rbac/role.yaml

@@ -30,6 +30,18 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - batch
+  resources:
+  - jobs
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch
 - apiGroups:
   - locostack.com
   resources:

internal/controller/managedmodel_controller.go

@@ -18,13 +18,24 @@ package controller
 
 import (
 	"context"
+	"fmt"
+	"maps"
 
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	apimeta "k8s.io/apimachinery/pkg/api/meta"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/internal/reconciler"
+	"github.com/LocoStack/loco-operator/pkg/templates"
+	llamacpp "github.com/LocoStack/loco-operator/pkg/templates/llama_cpp"
 )
 
 // ManagedModelReconciler reconciles a ManagedModel object
@@ -36,28 +47,264 @@ type ManagedModelReconciler struct {
 // +kubebuilder:rbac:groups=locostack.com,resources=managedmodels,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=locostack.com,resources=managedmodels/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=locostack.com,resources=managedmodels/finalizers,verbs=update
+// +kubebuilder:rbac:groups=locostack.com,resources=components,verbs=get;list;watch;create;update;patch;delete
+// +kubebuilder:rbac:groups=locostack.com,resources=stacks,verbs=get;list;watch
+// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;update;patch;delete
 
-// Reconcile is part of the main kubernetes reconciliation loop which aims to
-// move the current state of the cluster closer to the desired state.
-// TODO(user): Modify the Reconcile function to compare the state specified by
-// the ManagedModel object against the actual cluster state, and then
-// perform operations to make the cluster state reflect the state specified by
-// the user.
-//
-// For more details, check Reconcile and its Result here:
-// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.23.3/pkg/reconcile
 func (r *ManagedModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
-	_ = logf.FromContext(ctx)
+	log := logf.FromContext(ctx)
 
-	// TODO(user): your logic here
+	mm := &v1alpha1.ManagedModel{}
+	if err := r.Get(ctx, req.NamespacedName, mm); err != nil {
+		if apierrors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		log.Error(err, "Failed to get ManagedModel", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
+		return ctrl.Result{}, err
+	}
 
+	patch := client.MergeFrom(mm.DeepCopy())
+
+	stack := &v1alpha1.Stack{}
+	if err := r.Get(ctx, client.ObjectKey{Namespace: mm.Namespace, Name: mm.Spec.StackRef.Name}, stack); err != nil {
+		if apierrors.IsNotFound(err) {
+			return ctrl.Result{}, nil
+		}
+		log.Error(err, "Failed to get Stack for ManagedModel", "namespace", mm.Namespace, "name", mm.Name, "stack", mm.Spec.StackRef.Name)
+		return ctrl.Result{}, err
+	}
+
+	chatTemplateConfigMapName, err := r.reconcileChatTemplate(ctx, stack, mm)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+
+	pvcName, weightsReady, weightsFailed, err := r.reconcileWeights(ctx, stack, mm)
+	if err != nil {
+		return ctrl.Result{}, err
+	}
+
+	if weightsFailed {
+		apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionFalse,
+			Reason:             "WeightsFailed",
+			Message:            "ManagedModel weights are not available",
+			ObservedGeneration: mm.Generation,
+		})
+	} else if !weightsReady {
+		apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+			Type:               "Available",
+			Status:             metav1.ConditionFalse,
+			Reason:             "WeightsNotReady",
+			Message:            "ManagedModel weights are not ready",
+			ObservedGeneration: mm.Generation,
+		})
+	} else {
+		componentAvailable, err := r.reconcileComponent(ctx, mm, pvcName, chatTemplateConfigMapName)
+		if err != nil {
+			return ctrl.Result{}, err
+		}
+		gatewayBound, err := r.reconcileGatewayStatus(ctx, mm)
+		if err != nil {
+			log.Error(err, "Failed to reconcile gateway status for ManagedModel", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
+			return ctrl.Result{}, err
+		}
+		if !componentAvailable {
+			apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+				Type:               "Available",
+				Status:             metav1.ConditionFalse,
+				Reason:             "ModelUnavailable",
+				Message:            "ManagedModel is not available",
+				ObservedGeneration: mm.Generation,
+			})
+		} else if !gatewayBound {
+			apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+				Type:               "Available",
+				Status:             metav1.ConditionFalse,
+				Reason:             "GatewayUnbound",
+				Message:            "ManagedModel is not bound to a gateway",
+				ObservedGeneration: mm.Generation,
+			})
+		} else {
+			apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+				Type:               "Available",
+				Status:             metav1.ConditionTrue,
+				Reason:             "ModelAvailable",
+				Message:            "ManagedModel is available",
+				ObservedGeneration: mm.Generation,
+			})
+		}
+	}
+
+	mm.Status.ObservedGeneration = mm.Generation
+	if err := r.Status().Patch(ctx, mm, patch); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	log.Info("Reconciled ManagedModel", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
 	return ctrl.Result{}, nil
 }
 
+func (r *ManagedModelReconciler) reconcileChatTemplate(ctx context.Context, stack *v1alpha1.Stack, mm *v1alpha1.ManagedModel) (string, error) {
+	if mm.Spec.Category != "language" {
+		return "undefined", nil
+	}
+	cmName := fmt.Sprintf("%s-chat-template", mm.Name)
+	if mm.Spec.RuntimeInferenceParams == nil || mm.Spec.RuntimeInferenceParams.ChatTemplate == "" {
+		return cmName, nil
+	}
+	cm := &corev1.ConfigMap{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      cmName,
+			Namespace: mm.Namespace,
+		},
+	}
+	err := r.Get(ctx, client.ObjectKey{Namespace: mm.Namespace, Name: cm.Name}, cm)
+	if err != nil && !apierrors.IsNotFound(err) {
+		return "", err
+	}
+
+	_, err = controllerutil.CreateOrUpdate(ctx, r.Client, cm, func() error {
+		if cm.CreationTimestamp.IsZero() {
+			cm.Labels = reconciler.ResourceLabels(stack.Name, "Artifact", cmName, "")
+			chatTemplate := ""
+			if mm.Spec.RuntimeInferenceParams != nil {
+				chatTemplate = mm.Spec.RuntimeInferenceParams.ChatTemplate
+			}
+			cm.Data = map[string]string{
+				"template": chatTemplate,
+			}
+		}
+		return controllerutil.SetControllerReference(mm, cm, r.Scheme)
+	})
+	if err != nil {
+		return "", fmt.Errorf("Failed to create or update chat template: %w", err)
+	}
+	return cmName, nil
+}
+
+func (r *ManagedModelReconciler) reconcileWeights(ctx context.Context, stack *v1alpha1.Stack, mm *v1alpha1.ManagedModel) (string, bool, bool, error) {
+	artifactsReconciler := reconciler.NewArtifactReconciler(r.Client, r.Scheme, stack, "model", mm, &mm.Spec.Weights)
+	pvc, completed, failed, err := artifactsReconciler.ReconcileArtifact(ctx)
+	pvcName := ""
+	if pvc != nil {
+		pvcName = pvc.Name
+	}
+	return pvcName, completed, failed, err
+}
+
+func (r *ManagedModelReconciler) reconcileComponent(ctx context.Context, mm *v1alpha1.ManagedModel, pvcName string, chatTemplateConfigMapName string) (bool, error) {
+	compName := fmt.Sprintf("model-%s", mm.Name)
+	comp := &v1alpha1.Component{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      compName,
+			Namespace: mm.Namespace,
+		},
+	}
+	err := r.Get(ctx, client.ObjectKey{Namespace: mm.Namespace, Name: compName}, comp)
+
+	tmpl, err := templates.Manager.ResolveTemplate(mm.Spec.Template, "")
+	if err != nil {
+		apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+			Type:               "ModelRuntimeResolved",
+			Status:             metav1.ConditionFalse,
+			Reason:             "TemplateNotResolved",
+			Message:            fmt.Sprintf("Failed to reconcile template for ManagedModel: %v", err),
+			ObservedGeneration: mm.Generation,
+		})
+		return false, nil
+	}
+	apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+		Type:               "ModelRuntimeResolved",
+		Status:             metav1.ConditionTrue,
+		Reason:             "TemplateResolved",
+		Message:            fmt.Sprintf("Successfully resolved template for ManagedModel: %v", tmpl.Name),
+		ObservedGeneration: mm.Generation,
+	})
+
+	_, err = controllerutil.CreateOrUpdate(ctx, r.Client, comp, func() error {
+		if comp.Labels == nil {
+			comp.Labels = make(map[string]string)
+		}
+		stackName := ""
+		if mm.Spec.StackRef != nil {
+			stackName = mm.Spec.StackRef.Name
+		}
+		maps.Copy(comp.Labels, reconciler.ResourceLabels(stackName, "Model", mm.Name, ""))
+		comp.Spec.Category = "Model"
+		comp.Spec.Template = &tmpl
+		variables := map[string]string{
+			"spec.chatTemplateConfigMapName": chatTemplateConfigMapName,
+		}
+		params := mm.Spec.RuntimeInferenceParams
+		if params != nil {
+			if params.ChatTemplate != "" {
+				variables["spec.runtimeInferenceParameters.chatTemplateFilePath"] = fmt.Sprintf("%s/template", llamacpp.LLAMA_CPP_TEMPLATE_PATH)
+			}
+			if params.ContextWindow != nil {
+				variables["spec.runtimeInferenceParameters.contextWindow"] = fmt.Sprintf("%d", *params.ContextWindow)
+			}
+			if params.Temperature != "" {
+				variables["spec.runtimeInferenceParameters.temperature"] = params.Temperature
+			}
+			if params.TopK != nil {
+				variables["spec.runtimeInferenceParameters.topK"] = fmt.Sprintf("%d", *params.TopK)
+			}
+			if params.TopP != "" {
+				variables["spec.runtimeInferenceParameters.topP"] = params.TopP
+			}
+		}
+		if mm.Spec.Weights.HuggingFace != nil {
+			variables["spec.modelName"] = mm.Spec.Weights.HuggingFace.FileName
+			variables["spec.pvcName"] = pvcName
+		}
+		comp.Spec.Variables = variables
+		if mm.Spec.StackRef != nil {
+			comp.Spec.StackRef = &corev1.LocalObjectReference{Name: mm.Spec.StackRef.Name}
+		}
+		return controllerutil.SetControllerReference(mm, comp, r.Scheme)
+	})
+
+	available := apimeta.IsStatusConditionTrue(comp.Status.Conditions, "Available")
+	if available {
+		mm.Status.Endpoint = comp.Status.Endpoint
+		apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionTrue,
+			Reason:             "ComponentAvailable",
+			Message:            "Component is available",
+			ObservedGeneration: mm.Generation,
+		})
+	} else {
+		apimeta.SetStatusCondition(&mm.Status.Conditions, metav1.Condition{
+			Type:               "Ready",
+			Status:             metav1.ConditionFalse,
+			Reason:             "ComponentUnavailable",
+			Message:            "Component is not available",
+			ObservedGeneration: mm.Generation,
+		})
+	}
+	return available, err
+}
+
+func (r *ManagedModelReconciler) reconcileGatewayStatus(ctx context.Context, mm *v1alpha1.ManagedModel) (bool, error) {
+	condition, err := reconciler.ReconcileGatewayStatus(ctx, r.Client, r.Scheme, mm.Spec.StackRef.Name, mm)
+	if err != nil {
+		return false, err
+	}
+	apimeta.SetStatusCondition(&mm.Status.Conditions, *condition)
+	return condition.Status == metav1.ConditionTrue, nil
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (r *ManagedModelReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha1.ManagedModel{}).
+		Owns(&batchv1.Job{}).
+		Owns(&corev1.PersistentVolumeClaim{}).
+		Owns(&v1alpha1.Component{}).
+		Watches(&v1alpha1.Component{}, ComponentEventHandler("ManagedModel")).
 		Named("managedmodel").
 		Complete(r)
 }