Preheat models with Pod initialization

Dockerfile.preheat

@@ -0,0 +1,37 @@
+ARG BASE_IMAGE
+ARG BUILDER_IMAGE
+
+# Build the manager binary
+FROM ${BUILDER_IMAGE} as builder
+ARG TARGETOS
+ARG TARGETARCH
+
+WORKDIR /workspace
+# Copy the Go Modules manifests
+COPY preheat/go.mod go.mod
+COPY preheat/go.sum go.sum
+# cache deps before building and copying source so that we don't need to re-download as much
+# and so that source changes don't invalidate our downloaded layer
+RUN go mod download
+
+# Copy the go source
+COPY preheat/main.go main.go
+
+# Build
+# the GOARCH has not a default value to allow the binary be built according to the host where the command
+# was called. For example, if we call make docker-build in a local env which has the Apple Silicon M1 SO
+# the docker BUILDPLATFORM arg will be linux/arm64 when for Apple x86 it will be linux/amd64. Therefore,
+# by leaving it empty we can ensure that the container and binary shipped on it will have the same platform.
+RUN CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build -a -o manager main.go
+
+# Use distroless as minimal base image to package the manager binary
+# Refer to https://github.com/GoogleContainerTools/distroless for more details
+FROM ${BASE_IMAGE}
+WORKDIR /
+COPY --from=builder /workspace/manager .
+USER 65532:65532
+
+# Expose the http server.
+EXPOSE 9090
+
+ENTRYPOINT ["/manager"]

Makefile

@@ -63,9 +63,12 @@ IMAGE_NAME ?= manta
 IMAGE_REPO := $(IMAGE_REGISTRY)/$(IMAGE_NAME)
 AGENT_IMAGE_NAME ?= manta-agent
 AGENT_IMAGE_REPO := $(IMAGE_REGISTRY)/$(AGENT_IMAGE_NAME)
+PREHEAT_IMAGE_NAME ?= manta-preheat
+PREHEAT_IMAGE_REPO := $(IMAGE_REGISTRY)/$(PREHEAT_IMAGE_NAME)
 GIT_TAG ?= $(shell git describe --tags --dirty --always)
 IMG ?= $(IMAGE_REPO):$(GIT_TAG)
 AGENT_IMG ?= $(AGENT_IMAGE_REPO):$(GIT_TAG)
+PREHEAT_IMG ?= $(PREHEAT_IMAGE_REPO):$(GIT_TAG)
 BUILDER_IMAGE ?= golang:$(GO_VERSION)
 KIND_CLUSTER_NAME ?= kind
 
@@ -190,6 +193,19 @@ agent-image-load: agent-image-build
 agent-image-push: IMAGE_BUILD_EXTRA_OPTS=--push
 agent-image-push: agent-image-build
 
+.PHONY: preheat-image-build
+preheat-image-build:
+	$(IMAGE_BUILD_CMD) -t $(PREHEAT_IMG) \
+		-f Dockerfile.preheat \
+		--build-arg BASE_IMAGE=$(BASE_IMAGE) \
+		--build-arg BUILDER_IMAGE=$(BUILDER_IMAGE) \
+		--build-arg CGO_ENABLED=$(CGO_ENABLED) \
+		$(IMAGE_BUILD_EXTRA_OPTS) ./
+preheat-image-load: IMAGE_BUILD_EXTRA_OPTS=--load
+preheat-image-load: preheat-image-build
+preheat-image-push: IMAGE_BUILD_EXTRA_OPTS=--push
+preheat-image-push: preheat-image-build
+
 KIND = $(shell pwd)/bin/kind
 .PHONY: kind
 kind:

agent/config/manager/kustomization.yaml

@@ -5,4 +5,4 @@ kind: Kustomization
 images:
 - name: controller
   newName: inftyai/manta-agent
-  newTag: v0.0.3
+  newTag: main

api/v1alpha1/torrent_types.go

@@ -25,6 +25,7 @@ import (
 const (
 	TorrentNameLabelKey        = "manta.io/torrent-name"
 	TorrentProtectionFinalizer = "manta.io/torrent-protect"
+	ParentPodNameAnnoKey       = "manta.io/parent-pod-name"
 
 	HUGGINGFACE_MODEL_HUB = "Huggingface"
 )

cmd/main.go

@@ -20,9 +20,6 @@ import (
 	"flag"
 	"os"
 
-	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
-	// to ensure that exec-entrypoint and run can make use of them.
-
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 
 	"k8s.io/apimachinery/pkg/runtime"
@@ -106,7 +103,6 @@ func main() {
 	// will block until the cert is ready before setting up the controllers.
 	// Controllers who register after manager starts will start directly.
 	go setupControllers(mgr, certsReady)
-
 	//+kubebuilder:scaffold:builder
 
 	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
@@ -161,6 +157,13 @@ func setupControllers(mgr ctrl.Manager, certsReady chan struct{}) {
 		setupLog.Error(err, "unable to create controller", "controller", "Torrent")
 		os.Exit(1)
 	}
+	if err := controller.NewPodReconciler(
+		mgr.GetClient(),
+		mgr.GetScheme(),
+	).SetupWithManager(mgr); err != nil {
+		setupLog.Error(err, "unable to create controller", "controller", "Pod")
+		os.Exit(1)
+	}
 
 	if os.Getenv("ENABLE_WEBHOOKS") != "false" {
 		if err := webhook.SetupTorrentWebhook(mgr); err != nil {
@@ -171,5 +174,9 @@ func setupControllers(mgr ctrl.Manager, certsReady chan struct{}) {
 			setupLog.Error(err, "unable to create webhook", "webhook", "Replication")
 			os.Exit(1)
 		}
+		if err := webhook.SetupPodWebhook(mgr); err != nil {
+			setupLog.Error(err, "unable to create webhook", "webhook", "Pod")
+			os.Exit(1)
+		}
 	}
 }

config/crd/bases/manta.io_replications.yaml

@@ -21,6 +21,9 @@ spec:
     - jsonPath: .status.phase
       name: phase
       type: string
+    - jsonPath: .metadata.creationTimestamp
+      name: Age
+      type: date
     name: v1alpha1
     schema:
       openAPIV3Schema:

config/manager/kustomization.yaml

@@ -5,4 +5,4 @@ kind: Kustomization
 images:
 - name: controller
   newName: inftyai/manta
-  newTag: v0.0.3
+  newTag: main

config/rbac/role.yaml

@@ -31,6 +31,14 @@ rules:
   - list
   - update
   - watch
+- apiGroups:
+  - ""
+  resources:
+  - pods
+  verbs:
+  - get
+  - list
+  - watch
 - apiGroups:
   - manta.io
   resources:

config/webhook/kustomization.yaml

@@ -4,3 +4,11 @@ resources:
 
 configurations:
 - kustomizeconfig.yaml
+
+patches:
+- path: mutating-patch.yaml
+  target:
+    group: admissionregistration.k8s.io
+    version: v1
+    kind: MutatingWebhookConfiguration
+    name: mutating-webhook-configuration

config/webhook/manifests.yaml

@@ -4,6 +4,25 @@ kind: MutatingWebhookConfiguration
 metadata:
   name: mutating-webhook-configuration
 webhooks:
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate--v1-pod
+  failurePolicy: Fail
+  name: mpod.kb.io
+  rules:
+  - apiGroups:
+    - ""
+    apiVersions:
+    - v1
+    operations:
+    - CREATE
+    resources:
+    - pods
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:

config/webhook/mutating-patch.yaml

@@ -0,0 +1,26 @@
+- op: replace
+  path: /webhooks/0
+  value:
+    admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: webhook-service
+        namespace: system
+        path: /mutate--v1-pod
+    failurePolicy: Fail
+    name: mpod.kb.io
+    objectSelector:
+      matchExpressions:
+      - key: manta.io/torrent-name
+        operator: Exists
+    rules:
+    - apiGroups:
+      - ""
+      apiVersions:
+      - v1
+      operations:
+      - CREATE
+      resources:
+      - pods
+    sideEffects: None

hack/e2e-test.sh

@@ -24,7 +24,8 @@ function cleanup {
     then
         $KIND delete cluster --name $KIND_CLUSTER_NAME
     fi
-    (cd $CWD/config/manager && $KUSTOMIZE edit set image controller=inftyai/llmaz:main)
+    (cd $CWD/config/manager && $KUSTOMIZE edit set image controller=inftyai/manta:main)
+    (cd $CWD/agent/config/manager && $KUSTOMIZE edit set image controller=inftyai/manta-agent:main)
 }
 function startup {
     if [ $USE_EXISTING_CLUSTER == 'false' ]
@@ -44,8 +45,8 @@ function deploy {
     # agent
 	cd $CWD/agent/config/manager && $KUSTOMIZE edit set image controller=$AGENT_IMAGE_TAG
     $KUSTOMIZE build $CWD/agent/config | $KUBECTL apply --server-side -f -
-
 }
+
 trap cleanup EXIT
 startup
 kind_load

pkg/controller/pod_controller.go

@@ -0,0 +1,127 @@
+/*
+Copyright 2024.
+
+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 controller
+
+import (
+	"context"
+	"time"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/klog/v2"
+	"k8s.io/utils/ptr"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	api "github.com/inftyai/manta/api/v1alpha1"
+)
+
+// PodReconciler reconciles a Torrent object
+type PodReconciler struct {
+	client.Client
+	Scheme *runtime.Scheme
+}
+
+func NewPodReconciler(client client.Client, scheme *runtime.Scheme) *PodReconciler {
+	return &PodReconciler{
+		Client: client,
+		Scheme: scheme,
+	}
+}
+
+//+kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch
+
+func (r *PodReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	logger := log.FromContext(ctx)
+
+	pod := &corev1.Pod{}
+	if err := r.Get(ctx, types.NamespacedName{Namespace: req.Namespace, Name: req.Name}, pod); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	logger.Info("reconcile Pod")
+	// This should not happen, double check here.
+	if pod.Labels == nil || pod.Labels[api.TorrentNameLabelKey] == "" {
+		return ctrl.Result{}, nil
+	}
+
+	torrentName := pod.Labels[api.TorrentNameLabelKey]
+
+	torrent := &api.Torrent{}
+	if err := r.Get(ctx, types.NamespacedName{Name: torrentName}, torrent); err != nil {
+		return ctrl.Result{}, err
+	}
+
+	newTorrent := constructTorrent(torrent, pod)
+	if err := r.Client.Create(ctx, &newTorrent); err != nil {
+		logger.Error(err, "failed to create Torrent", "Torrent", klog.KObj(&newTorrent))
+		return ctrl.Result{}, err
+	}
+
+	return ctrl.Result{}, nil
+}
+
+func (r *PodReconciler) Create(e event.CreateEvent) bool {
+	pod, match := e.Object.(*corev1.Pod)
+	if !match {
+		return false
+	}
+
+	// Pod should be managed by Manta.
+	if pod.Labels == nil || pod.Labels[api.TorrentNameLabelKey] == "" {
+		return false
+	}
+
+	return true
+}
+
+func (r *PodReconciler) Update(e event.UpdateEvent) bool {
+	return false
+}
+
+func (r *PodReconciler) Delete(e event.DeleteEvent) bool {
+	return false
+}
+
+func (r *PodReconciler) Generic(e event.GenericEvent) bool {
+	return false
+}
+
+func (r *PodReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&corev1.Pod{}).
+		WithEventFilter(r).
+		Complete(r)
+}
+
+func constructTorrent(torrent *api.Torrent, pod *corev1.Pod) api.Torrent {
+	newTorrent := api.Torrent{}
+	newTorrent.ObjectMeta.Name = torrent.Name + "--tmp--" + pod.Spec.NodeName
+	newTorrent.TypeMeta = torrent.TypeMeta
+	newTorrent.Annotations = map[string]string{api.ParentPodNameAnnoKey: pod.Namespace + "/" + pod.Name}
+	newTorrent.Spec = torrent.Spec
+	newTorrent.Spec.Preheat = ptr.To[bool](true)
+	newTorrent.Spec.Replicas = ptr.To[int32](1)
+	newTorrent.Spec.ReclaimPolicy = ptr.To[api.ReclaimPolicy](api.RetainReclaimPolicy)
+	newTorrent.Spec.TTLSecondsAfterReady = ptr.To[time.Duration](0)
+	newTorrent.Spec.NodeSelector = map[string]string{"kubernetes.io/hostname": pod.Spec.NodeName}
+
+	return newTorrent
+}