Parcourir la source

feat: add artifact reconciler

Thomas Zhang il y a 2 mois
Parent
commit
bf5b12f6f1
3 fichiers modifiés avec 258 ajouts et 0 suppressions
  1. 139 0
      internal/reconciler/artifact.go
  2. 68 0
      pkg/templates/downloader/hf.go
  3. 51 0
      pkg/templates/downloader/url.go

+ 139 - 0
internal/reconciler/artifact.go

@@ -0,0 +1,139 @@
+/*
+Copyright 2026 LocoStack.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package reconciler
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	"github.com/LocoStack/loco-operator/pkg/templates/downloader"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+)
+
+type ArtifactReconciler struct {
+	*JobReconciler
+	client        client.Client
+	scheme        *runtime.Scheme
+	stack         *v1alpha1.Stack
+	owner         metav1.Object
+	artifactsSpec *v1alpha1.ArtifactSpec
+}
+
+func NewArtifactReconciler(client client.Client, scheme *runtime.Scheme, stack *v1alpha1.Stack, component string, owner metav1.Object, artifactsSpec *v1alpha1.ArtifactSpec) *ArtifactReconciler {
+	return &ArtifactReconciler{
+		JobReconciler: NewJobReconciler(
+			client, scheme, stack, "Artifact", owner.GetName(), owner,
+		),
+		client:        client,
+		scheme:        scheme,
+		stack:         stack,
+		owner:         owner,
+		artifactsSpec: artifactsSpec,
+	}
+}
+
+func (r *ArtifactReconciler) ReconcileArtifact(ctx context.Context) (*corev1.PersistentVolumeClaim, bool, bool, error) {
+	var pvcName string
+	if r.artifactsSpec.PVC != nil {
+		pvcName = r.artifactsSpec.PVC.ClaimName
+	} else {
+		pvcName = fmt.Sprintf("%s-artifacts", r.owner.GetName())
+	}
+	pvc, err := r.reconcilePVC(ctx, pvcName)
+	if err != nil {
+		return nil, false, true, err
+	}
+	if r.artifactsSpec.HuggingFace != nil || r.artifactsSpec.URL != nil {
+		variables := map[string]string{
+			"spec.pvcName": pvcName,
+		}
+		var tmpl *v1alpha1.Template
+		if r.artifactsSpec.HuggingFace != nil {
+			hf := r.artifactsSpec.HuggingFace
+			tmpl = downloader.HFDownloaderTemplate.DeepCopy()
+			variables["spec.hfRepo"] = hf.Repo
+			variables["spec.hfFilename"] = hf.FileName
+			if hf.Revision == "" {
+				variables["spec.hfRevision"] = "main"
+			} else {
+				variables["spec.hfRevision"] = hf.Revision
+			}
+			if hf.Endpoint != "" {
+				variables["spec.hfEndpoint"] = hf.Endpoint
+			} else {
+				variables["spec.hfEndpoint"] = "https://huggingface.co"
+			}
+			if hf.TokenSecretRef != nil {
+				variables["spec.hfToken.name"] = hf.TokenSecretRef.Name
+				variables["spec.hfToken.key"] = hf.TokenSecretRef.Key
+			} else {
+				variables["spec.hfToken.name"] = "not-set"
+				variables["spec.hfToken.key"] = "not-set"
+			}
+		} else if r.artifactsSpec.URL != nil {
+			tmpl = downloader.URLDownloaderTemplate.DeepCopy()
+			variables["spec.download"] = r.artifactsSpec.URL.URL
+			variables["spec.fileName"] = r.artifactsSpec.URL.FileName
+		}
+		completed, failed, err := r.JobReconciler.ReconcileJob(ctx, tmpl, variables)
+		if err != nil {
+			return nil, false, true, err
+		}
+		return pvc, completed, failed, nil
+	}
+	return pvc, true, false, nil
+}
+
+func (r *ArtifactReconciler) reconcilePVC(ctx context.Context, pvcName string) (*corev1.PersistentVolumeClaim, error) {
+	pvc := &corev1.PersistentVolumeClaim{ObjectMeta: metav1.ObjectMeta{
+		Name:      pvcName,
+		Namespace: r.owner.GetNamespace(),
+	}}
+	pvcNotFound := false
+	if err := r.client.Get(ctx, client.ObjectKeyFromObject(pvc), pvc); err != nil {
+		if client.IgnoreNotFound(err) != nil {
+			return nil, fmt.Errorf("Failed to get PVC: %w", err)
+		}
+		pvcNotFound = true
+	}
+	if pvcNotFound {
+		_, err := controllerutil.CreateOrUpdate(ctx, r.client, pvc, func() error {
+			if pvc.CreationTimestamp.IsZero() {
+				pvc.Labels = ResourceLabels(r.stack.Name, "Artifact", pvcName, "")
+				pvc.Spec = corev1.PersistentVolumeClaimSpec{
+					AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
+					Resources: corev1.VolumeResourceRequirements{
+						Requests: corev1.ResourceList{
+							corev1.ResourceStorage: resource.MustParse("30Gi"),
+						},
+					},
+				}
+			}
+			return controllerutil.SetControllerReference(r.owner, pvc, r.scheme)
+		})
+		if err != nil {
+			return nil, fmt.Errorf("Failed to create or update PVC: %w", err)
+		}
+	}
+	return pvc, nil
+}

+ 68 - 0
pkg/templates/downloader/hf.go

@@ -0,0 +1,68 @@
+/*
+Copyright 2026 LocoStack.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package downloader
+
+import (
+	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/utils/ptr"
+)
+
+const (
+	DOWNLOADER_DOWNLOAD_DIR = "/data"
+	command                 = `pip install -q huggingface_hub && \
+hf download $HF_REPO $HF_REPO_FILENAME --local-dir $DOWNLOAD_DIR --revision $HF_REPO_REVISION`
+)
+
+var HFDownloaderTemplate = v1alpha1.Template{
+	Name: "hf-downloader",
+	Spec: v1alpha1.TemplateSpec{
+		Runtime: v1alpha1.RuntimeSpec{
+			Name:    "downloader",
+			Image:   "python:3.11-slim",
+			Command: []string{"/bin/sh", "-c", command},
+			Args:    []string{},
+			Env: []corev1.EnvVar{
+				{Name: "DOWNLOAD_DIR", Value: DOWNLOADER_DOWNLOAD_DIR},
+				{Name: "HF_REPO", Value: "$(spec.hfRepo)"},
+				{Name: "HF_REPO_FILENAME", Value: "$(spec.hfFilename)"},
+				{Name: "HF_REPO_REVISION", Value: "$(spec.hfRevision)"},
+				{Name: "HF_ENDPOINT", Value: "$(spec.hfEndpoint)"},
+				{Name: "HF_TOKEN", ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						LocalObjectReference: corev1.LocalObjectReference{Name: "$(spec.hfToken.name)"},
+						Key:                  "$(spec.hfToken.key)",
+						Optional:             ptr.To(true),
+					},
+				}},
+			},
+			VolumeMounts: []corev1.VolumeMount{
+				{Name: "artifacts", MountPath: DOWNLOADER_DOWNLOAD_DIR},
+			},
+		},
+		Volumes: []corev1.Volume{
+			{
+				Name: "artifacts",
+				VolumeSource: corev1.VolumeSource{
+					PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+						ClaimName: "$(spec.pvcName)",
+					},
+				},
+			},
+		},
+	},
+}

+ 51 - 0
pkg/templates/downloader/url.go

@@ -0,0 +1,51 @@
+/*
+Copyright 2026 LocoStack.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package downloader
+
+import (
+	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
+)
+
+var URLDownloaderTemplate = v1alpha1.Template{
+	Name: "url-downloader",
+	Spec: v1alpha1.TemplateSpec{
+		Runtime: v1alpha1.RuntimeSpec{
+			Name:    "downloader",
+			Image:   "curlimages/curl:8.20.0",
+			Command: []string{"curl", "-L", "$DOWNLOAD_TARGET", "-o", "$DOWNLOAD_DIR/$DOWNLOAD_FILENAME"},
+			Env: []corev1.EnvVar{
+				{Name: "DOWNLOAD_TARGET", Value: "$(spec.download)"},
+				{Name: "DOWNLOAD_DIR", Value: DOWNLOADER_DOWNLOAD_DIR},
+				{Name: "DOWNLOAD_FILENAME", Value: "$(spec.fileName)"},
+			},
+			VolumeMounts: []corev1.VolumeMount{
+				{Name: "artifacts", MountPath: DOWNLOADER_DOWNLOAD_DIR},
+			},
+		},
+		Volumes: []corev1.Volume{
+			{
+				Name: "artifacts",
+				VolumeSource: corev1.VolumeSource{
+					PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+						ClaimName: "$(metadata.name)-artifacts",
+					},
+				},
+			},
+		},
+	},
+}