Repository: kaito-project/kaito

Issue:
headlamp gui


Discussion:
@ritazh: We should add docs to point to https://kaito-project.github.io/headlamp-kaito/ and https://github.com/kaito-project/headlamp-kaito 

Solution patch:
diff --git a/.github/workflows/deploy-docs.yml b/.github/workflows/deploy-docs.yml
index bdf3dbd..786b343 100644
--- a/.github/workflows/deploy-docs.yml
+++ b/.github/workflows/deploy-docs.yml
@@ -39,7 +39,7 @@ jobs:
           github.com:443 registry.yarnpkg.com:443 registry.npmjs.org:443 *.githubusercontent.com:443 *.blob.core.windows.net:443
 
     - name: Setup Node
-      uses: actions/setup-node@cdca7365b2dadb8aad0a33bc7601856ffabcc48e # v4.3.0
+      uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
       with:
         node-version: 20.x
 
diff --git a/charts/kaito/ragengine/templates/knative-logging-configmap.yaml b/charts/kaito/ragengine/templates/knative-logging-configmap.yaml
new file mode 100644
index 0000000..ad06d09
--- /dev/null
+++ b/charts/kaito/ragengine/templates/knative-logging-configmap.yaml
@@ -0,0 +1,35 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: config-logging
+  namespace: {{ .Release.Namespace }}
+data:
+  loglevel.controller: {{ .Values.logging.level | quote }}
+  loglevel.webhook: {{ .Values.logging.level | quote }}
+  loglevel.autoscaler: {{ .Values.logging.level | quote }}
+  loglevel.queueproxy: {{ .Values.logging.level | quote }}
+  loglevel.activator: {{ .Values.logging.level | quote }}
+  loglevel.hpaautoscaler: {{ .Values.logging.level | quote }}
+  loglevel.certcontroller: {{ .Values.logging.level | quote }}
+  loglevel.istiocontroller: {{ .Values.logging.level | quote }}
+  zap-logger-config: |
+    {
+      "level": {{ .Values.logging.level | quote }},
+      "development": false,
+      "outputPaths": ["stdout"],
+      "errorOutputPaths": ["stderr"],
+      "encoding": "json",
+      "encoderConfig": {
+        "timeKey": "ts",
+        "levelKey": "level",
+        "nameKey": "logger",
+        "callerKey": "caller",
+        "messageKey": "msg",
+        "stacktraceKey": "stacktrace",
+        "lineEnding": "",
+        "levelEncoder": "",
+        "timeEncoder": "iso8601",
+        "durationEncoder": "",
+        "callerEncoder": ""
+      }
+    }
diff --git a/charts/kaito/ragengine/values.yaml b/charts/kaito/ragengine/values.yaml
index eaad3b6..9d14c03 100644
--- a/charts/kaito/ragengine/values.yaml
+++ b/charts/kaito/ragengine/values.yaml
@@ -17,6 +17,9 @@ securityContext:
       - "ALL"
 webhook:
   port: 9443
+# Knative logging configuration
+logging:
+  level: "error"
 resources:
   limits:
     cpu: 500m
diff --git a/charts/kaito/workspace/templates/knative-logging-configmap.yaml b/charts/kaito/workspace/templates/knative-logging-configmap.yaml
new file mode 100644
index 0000000..ad06d09
--- /dev/null
+++ b/charts/kaito/workspace/templates/knative-logging-configmap.yaml
@@ -0,0 +1,35 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: config-logging
+  namespace: {{ .Release.Namespace }}
+data:
+  loglevel.controller: {{ .Values.logging.level | quote }}
+  loglevel.webhook: {{ .Values.logging.level | quote }}
+  loglevel.autoscaler: {{ .Values.logging.level | quote }}
+  loglevel.queueproxy: {{ .Values.logging.level | quote }}
+  loglevel.activator: {{ .Values.logging.level | quote }}
+  loglevel.hpaautoscaler: {{ .Values.logging.level | quote }}
+  loglevel.certcontroller: {{ .Values.logging.level | quote }}
+  loglevel.istiocontroller: {{ .Values.logging.level | quote }}
+  zap-logger-config: |
+    {
+      "level": {{ .Values.logging.level | quote }},
+      "development": false,
+      "outputPaths": ["stdout"],
+      "errorOutputPaths": ["stderr"],
+      "encoding": "json",
+      "encoderConfig": {
+        "timeKey": "ts",
+        "levelKey": "level",
+        "nameKey": "logger",
+        "callerKey": "caller",
+        "messageKey": "msg",
+        "stacktraceKey": "stacktrace",
+        "lineEnding": "",
+        "levelEncoder": "",
+        "timeEncoder": "iso8601",
+        "durationEncoder": "",
+        "callerEncoder": ""
+      }
+    }
diff --git a/charts/kaito/workspace/values.yaml b/charts/kaito/workspace/values.yaml
index 7f034c9..d467fb1 100644
--- a/charts/kaito/workspace/values.yaml
+++ b/charts/kaito/workspace/values.yaml
@@ -19,6 +19,9 @@ featureGates:
   vLLM: "true"
 webhook:
   port: 9443
+# Knative logging configuration
+logging:
+  level: "error"
 presetRegistryName: mcr.microsoft.com/aks/kaito
 resources:
   limits:
diff --git a/pkg/model/interface.go b/pkg/model/interface.go
index b695df5..da42ff1 100644
--- a/pkg/model/interface.go
+++ b/pkg/model/interface.go
@@ -54,6 +54,7 @@ const (
 	RuntimeNameHuggingfaceTransformers RuntimeName = "transformers"
 	RuntimeNameVLLM                    RuntimeName = "vllm"
 
+	DefaultTuningMainFile = "/workspace/tfs/fine_tuning.py"
 	ConfigfileNameVLLM    = "inference_config.yaml"
 	DefaultMemoryUtilVLLM = 0.9 // Default gpu memory utilization for VLLM runtime
 	UpperMemoryUtilVLLM   = 0.95
@@ -359,3 +360,15 @@ func getGPUMemoryUtilForVLLM(gpuConfig *sku.GPUConfig) float64 {
 	util := math.Floor((1.0-float64(ReservedNonKVCacheMemory.Value())/gpuMemPerGPU)*100) / 100
 	return math.Min(util, UpperMemoryUtilVLLM)
 }
+
+// Only support Huggingface for now
+func (p *PresetParam) GetTuningCommand(rc RuntimeContext) []string {
+	if p.Transformers.AccelerateParams == nil {
+		p.Transformers.AccelerateParams = make(map[string]string)
+	}
+
+	p.Transformers.AccelerateParams["num_processes"] = strconv.Itoa(rc.SKUNumGPUs)
+	torchCommand := utils.BuildCmdStr(p.Transformers.BaseCommand, p.Transformers.AccelerateParams)
+	modelCommand := utils.BuildCmdStr(DefaultTuningMainFile, p.Transformers.ModelRunParams)
+	return utils.ShellCmd(torchCommand + " " + modelCommand)
+}
diff --git a/pkg/utils/common-preset.go b/pkg/utils/common-preset.go
index bd6fd02..33b8f47 100644
--- a/pkg/utils/common-preset.go
+++ b/pkg/utils/common-preset.go
@@ -58,6 +58,17 @@ func ConfigResultsVolume(outputPath string, outputVolume *corev1.VolumeSource) (
 	return sharedWorkspaceVolume, sharedVolumeMount
 }
 
+func FindResultsVolumeMount(spec *corev1.PodSpec) *corev1.VolumeMount {
+	for _, container := range spec.Containers {
+		for _, volumeMount := range container.VolumeMounts {
+			if volumeMount.Name == "results-volume" {
+				return &volumeMount
+			}
+		}
+	}
+	return nil
+}
+
 func ConfigImagePullSecretVolume(nameSuffix string, imagePullSecrets []string) (corev1.Volume, corev1.VolumeMount) {
 	name := fmt.Sprintf("docker-config-%s", nameSuffix)
 
diff --git a/pkg/utils/generator/generator.go b/pkg/utils/generator/generator.go
index 71051a9..ba75a44 100644
--- a/pkg/utils/generator/generator.go
+++ b/pkg/utils/generator/generator.go
@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	appsv1 "k8s.io/api/apps/v1"
+	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
@@ -37,7 +38,7 @@ type GeneratorContext interface {
 }
 
 type ManifestType interface {
-	appsv1.StatefulSet | corev1.PodSpec | appsv1.Deployment
+	appsv1.StatefulSet | appsv1.Deployment | batchv1.Job | corev1.PodSpec
 }
 
 type TypedManifestModifier[C GeneratorContext, T ManifestType] func(ctx *C, obj *T) error
diff --git a/pkg/workspace/controllers/workspace_controller.go b/pkg/workspace/controllers/workspace_controller.go
index 9566888..8409eb0 100644
--- a/pkg/workspace/controllers/workspace_controller.go
+++ b/pkg/workspace/controllers/workspace_controller.go
@@ -307,9 +307,12 @@ func (c *WorkspaceReconciler) syncControllerRevision(ctx context.Context, wObj *
 		annotations[kaitov1beta1.WorkspaceRevisionAnnotation] = strconv.FormatInt(controllerRevision.Revision, 10)
 	}
 	annotations[WorkspaceHashAnnotation] = currentHash
-	wObj.SetAnnotations(annotations)
 
-	if err := c.Update(ctx, wObj); err != nil {
+	err = updateWorkspaceWithRetry(ctx, c.Client, wObj, func(ws *kaitov1beta1.Workspace) error {
+		ws.SetAnnotations(annotations)
+		return nil
+	})
+	if err != nil {
 		return fmt.Errorf("failed to update Workspace annotations: %w", err)
 	}
 	return nil
@@ -543,6 +546,7 @@ func (c *WorkspaceReconciler) ensureNodePlugins(ctx context.Context, wObj *kaito
 	tick := timeClock.NewTicker(consts.NodePluginInstallTimeout)
 	defer tick.Stop()
 
+	hasUpdatedLabel := false
 	for {
 		select {
 		case <-ctx.Done():
@@ -562,15 +566,18 @@ func (c *WorkspaceReconciler) ensureNodePlugins(ctx context.Context, wObj *kaito
 				return nil
 			}
 
-			err = resources.UpdateNodeWithLabel(ctx, freshNode, resources.LabelKeyNvidia, resources.LabelValueNvidia, c.Client)
-			if apierrors.IsNotFound(err) {
-				klog.ErrorS(err, "nvidia plugin cannot be installed, node not found", "node", freshNode.Name)
-				if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1beta1.ConditionTypeNodeClaimStatus, metav1.ConditionFalse,
-					"checkNodeClaimStatusFailed", err.Error()); updateErr != nil {
-					klog.ErrorS(updateErr, "failed to update workspace status", "workspace", klog.KObj(wObj))
-					return updateErr
+			if !hasUpdatedLabel {
+				err = resources.UpdateNodeWithLabel(ctx, freshNode, resources.LabelKeyNvidia, resources.LabelValueNvidia, c.Client)
+				if apierrors.IsNotFound(err) {
+					klog.ErrorS(err, "nvidia plugin cannot be installed, node not found", "node", freshNode.Name)
+					if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1beta1.ConditionTypeNodeClaimStatus, metav1.ConditionFalse,
+						"checkNodeClaimStatusFailed", err.Error()); updateErr != nil {
+						klog.ErrorS(updateErr, "failed to update workspace status", "workspace", klog.KObj(wObj))
+						return updateErr
+					}
+					return err
 				}
-				return err
+				hasUpdatedLabel = true
 			}
 
 			time.Sleep(1 * time.Second)
@@ -655,7 +662,7 @@ func (c *WorkspaceReconciler) applyTuning(ctx context.Context, wObj *kaitov1beta
 					}
 
 					var workloadObj client.Object
-					workloadObj, err = tuning.CreatePresetTuning(ctx, wObj, revisionNum, tuningParam, c.Client)
+					workloadObj, err = tuning.CreatePresetTuning(ctx, wObj, revisionNum, model, c.Client)
 					if err != nil {
 						return
 					}
@@ -668,7 +675,7 @@ func (c *WorkspaceReconciler) applyTuning(ctx context.Context, wObj *kaitov1beta
 			} else if apierrors.IsNotFound(err) {
 				var workloadObj client.Object
 				// Need to create a new workload
-				workloadObj, err = tuning.CreatePresetTuning(ctx, wObj, revisionNum, tuningParam, c.Client)
+				workloadObj, err = tuning.CreatePresetTuning(ctx, wObj, revisionNum, model, c.Client)
 				if err != nil {
 					return
 				}
@@ -805,6 +812,20 @@ func (c *WorkspaceReconciler) applyInference(ctx context.Context, wObj *kaitov1b
 	return nil
 }
 
+// updateWorkspaceWithRetry gets the latest workspace object, applies the modify function, and retries on conflict
+func updateWorkspaceWithRetry(ctx context.Context, c client.Client, wObj *kaitov1beta1.Workspace, modifyFn func(*kaitov1beta1.Workspace) error) error {
+	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
+		latestWorkspace := &kaitov1beta1.Workspace{}
+		if err := c.Get(ctx, client.ObjectKeyFromObject(wObj), latestWorkspace); err != nil {
+			return err
+		}
+		if err := modifyFn(latestWorkspace); err != nil {
+			return err
+		}
+		return c.Update(ctx, latestWorkspace)
+	})
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (c *WorkspaceReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	c.Recorder = mgr.GetEventRecorderFor("Workspace")
diff --git a/pkg/workspace/controllers/workspace_gc_finalizer.go b/pkg/workspace/controllers/workspace_gc_finalizer.go
index 839a90e..9e19bd5 100644
--- a/pkg/workspace/controllers/workspace_gc_finalizer.go
+++ b/pkg/workspace/controllers/workspace_gc_finalizer.go
@@ -16,6 +16,7 @@ package controllers
 import (
 	"context"
 
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -47,14 +48,19 @@ func (c *WorkspaceReconciler) garbageCollectWorkspace(ctx context.Context, wObj
 		}
 	}
 
-	if controllerutil.RemoveFinalizer(wObj, consts.WorkspaceFinalizer) {
-		if updateErr := c.Update(ctx, wObj, &client.UpdateOptions{}); updateErr != nil {
-			klog.ErrorS(updateErr, "failed to remove the finalizer from the workspace",
-				"workspace", klog.KObj(wObj))
-			return ctrl.Result{}, updateErr
+	updateErr := updateWorkspaceWithRetry(ctx, c.Client, wObj, func(ws *kaitov1beta1.Workspace) error {
+		controllerutil.RemoveFinalizer(ws, consts.WorkspaceFinalizer)
+		return nil
+	})
+	if updateErr != nil {
+		if apierrors.IsNotFound(updateErr) {
+			return ctrl.Result{}, nil
 		}
-		klog.InfoS("successfully removed the workspace finalizers", "workspace", klog.KObj(wObj))
+		klog.ErrorS(updateErr, "failed to update the workspace to remove finalizer", "workspace", klog.KObj(wObj))
+		return ctrl.Result{}, updateErr
 	}
 
+	klog.InfoS("successfully removed the workspace finalizers", "workspace", klog.KObj(wObj))
+
 	return ctrl.Result{}, nil
 }
diff --git a/pkg/workspace/controllers/workspace_status.go b/pkg/workspace/controllers/workspace_status.go
index 317556a..dda453b 100644
--- a/pkg/workspace/controllers/workspace_status.go
+++ b/pkg/workspace/controllers/workspace_status.go
@@ -33,7 +33,7 @@ import (
 func (c *WorkspaceReconciler) updateWorkspaceStatus(ctx context.Context, name *client.ObjectKey, condition *metav1.Condition, workerNodes []string) error {
 	return retry.OnError(retry.DefaultRetry,
 		func(err error) bool {
-			return apierrors.IsServiceUnavailable(err) || apierrors.IsServerTimeout(err) || apierrors.IsTooManyRequests(err)
+			return apierrors.IsServiceUnavailable(err) || apierrors.IsServerTimeout(err) || apierrors.IsTooManyRequests(err) || apierrors.IsConflict(err)
 		},
 		func() error {
 			// Read the latest version to avoid update conflict.
diff --git a/pkg/workspace/inference/preset-inferences.go b/pkg/workspace/inference/preset-inferences.go
index 6377b48..9e9ed27 100644
--- a/pkg/workspace/inference/preset-inferences.go
+++ b/pkg/workspace/inference/preset-inferences.go
@@ -129,7 +129,6 @@ func GenerateModelWeightsCacheVolume(ctx context.Context, workspaceObj *v1beta1.
 	}
 }
 
-// TODO: refactor this function
 func GeneratePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace, revisionNum string,
 	model pkgmodel.Model, kubeClient client.Client) (client.Object, error) {
 
diff --git a/pkg/workspace/manifests/manifests.go b/pkg/workspace/manifests/manifests.go
index 1349ab1..813bd7a 100644
--- a/pkg/workspace/manifests/manifests.go
+++ b/pkg/workspace/manifests/manifests.go
@@ -159,77 +159,42 @@ func SetStatefulSetPodSpec(podSpec *corev1.PodSpec) func(*generator.WorkspaceGen
 	}
 }
 
-func GenerateTuningJobManifest(wObj *kaitov1beta1.Workspace, revisionNum string, imageName string,
-	imagePullSecretRefs []corev1.LocalObjectReference, replicas int, commands []string, containerPorts []corev1.ContainerPort,
-	livenessProbe, readinessProbe *corev1.Probe, resourceRequirements corev1.ResourceRequirements, tolerations []corev1.Toleration,
-	initContainers []corev1.Container, sidecarContainers []corev1.Container, volumes []corev1.Volume, volumeMounts []corev1.VolumeMount,
-	envVars []corev1.EnvVar) *batchv1.Job {
-	labels := map[string]string{
-		kaitov1beta1.LabelWorkspaceName: wObj.Name,
-	}
-
-	// TODO: make containers only mount the volumes they need
-
-	for i := range initContainers {
-		initContainers[i].VolumeMounts = utils.DedupVolumeMounts(append(initContainers[i].VolumeMounts, volumeMounts...))
-	}
-
-	for i := range sidecarContainers {
-		sidecarContainers[i].VolumeMounts = utils.DedupVolumeMounts(append(sidecarContainers[i].VolumeMounts, volumeMounts...))
-	}
-
-	// Construct the complete list of containers (main and sidecars)
-	containers := append([]corev1.Container{
-		{
-			Name:           wObj.Name,
-			Image:          imageName,
-			Command:        commands,
-			Resources:      resourceRequirements,
-			LivenessProbe:  livenessProbe,
-			ReadinessProbe: readinessProbe,
-			Ports:          containerPorts,
-			VolumeMounts:   volumeMounts,
-			Env:            envVars,
-		},
-	}, sidecarContainers...)
-
-	shouldShareProcessNamespace := ptr.To(true)
-	if len(sidecarContainers) == 0 {
-		shouldShareProcessNamespace = ptr.To(false)
-	}
+func GenerateTuningJobManifest(revisionNum string) func(*generator.WorkspaceGeneratorContext, *batchv1.Job) error {
+	return func(ctx *generator.WorkspaceGeneratorContext, j *batchv1.Job) error {
+		labels := map[string]string{
+			kaitov1beta1.LabelWorkspaceName: ctx.Workspace.Name,
+		}
 
-	return &batchv1.Job{
-		TypeMeta: v1.TypeMeta{
-			APIVersion: "batch/v1",
-			Kind:       "Job",
-		},
-		ObjectMeta: v1.ObjectMeta{
-			Name:      wObj.Name,
-			Namespace: wObj.Namespace,
+		j.ObjectMeta = v1.ObjectMeta{
+			Name:      ctx.Workspace.Name,
+			Namespace: ctx.Workspace.Namespace,
 			Labels:    labels,
 			Annotations: map[string]string{
 				kaitov1beta1.WorkspaceRevisionAnnotation: revisionNum,
 			},
 			OwnerReferences: []v1.OwnerReference{
-				*v1.NewControllerRef(wObj, kaitov1beta1.GroupVersion.WithKind("Workspace")),
+				*v1.NewControllerRef(ctx.Workspace, kaitov1beta1.GroupVersion.WithKind("Workspace")),
 			},
-		},
-		Spec: batchv1.JobSpec{
+		}
+
+		j.Spec = batchv1.JobSpec{
 			Template: corev1.PodTemplateSpec{
 				ObjectMeta: v1.ObjectMeta{
 					Labels: labels,
 				},
-				Spec: corev1.PodSpec{
-					InitContainers:        initContainers,
-					Containers:            containers,
-					RestartPolicy:         corev1.RestartPolicyNever,
-					ShareProcessNamespace: shouldShareProcessNamespace,
-					Volumes:               volumes,
-					Tolerations:           tolerations,
-					ImagePullSecrets:      imagePullSecretRefs,
-				},
 			},
-		},
+		}
+		return nil
+	}
+}
+
+func SetJobPodSpec(podSpec *corev1.PodSpec) func(*generator.WorkspaceGeneratorContext, *batchv1.Job) error {
+	return func(ctx *generator.WorkspaceGeneratorContext, j *batchv1.Job) error {
+		if len(podSpec.Containers) > 1 {
+			podSpec.ShareProcessNamespace = ptr.To(true)
+		}
+		j.Spec.Template.Spec = *podSpec
+		return nil
 	}
 }
 
@@ -296,8 +261,12 @@ func GeneratePullerContainers(wObj *kaitov1beta1.Workspace, volumeMounts []corev
 		source := adapter.Source
 		sourceName := source.Name
 
-		volume, volumeMount := utils.ConfigImagePullSecretVolume(sourceName+"-inference-adapter", source.ImagePullSecrets)
-		volumes = append(volumes, volume)
+		var volumeMount corev1.VolumeMount
+		var volume corev1.Volume
+		if len(source.ImagePullSecrets) > 0 {
+			volume, volumeMount = utils.ConfigImagePullSecretVolume(sourceName+"-inference-adapter", source.ImagePullSecrets)
+			volumes = append(volumes, volume)
+		}
 
 		if adapter.Strength != nil {
 			envVar := corev1.EnvVar{
diff --git a/pkg/workspace/tuning/preset-tuning.go b/pkg/workspace/tuning/preset-tuning.go
index acfe818..d5d02da 100644
--- a/pkg/workspace/tuning/preset-tuning.go
+++ b/pkg/workspace/tuning/preset-tuning.go
@@ -27,9 +27,11 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
-	"github.com/kaito-project/kaito/pkg/model"
+	pkgmodel "github.com/kaito-project/kaito/pkg/model"
+	"github.com/kaito-project/kaito/pkg/sku"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/generator"
 	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/kaito-project/kaito/pkg/workspace/image"
 	"github.com/kaito-project/kaito/pkg/workspace/manifests"
@@ -37,8 +39,8 @@ import (
 )
 
 const (
-	PortInferenceServer     = int32(5000)
-	TuningFile              = "/workspace/tfs/fine_tuning.py"
+	PortInferenceServer = int32(5000)
+
 	DefaultBaseDir          = "/mnt"
 	DefaultOutputVolumePath = "/mnt/output"
 )
@@ -66,14 +68,6 @@ var (
 	}
 )
 
-func getInstanceGPUCount(sku string) int {
-	skuHandler, _ := utils.GetSKUHandler()
-	if gpuConfig := skuHandler.GetGPUConfigBySKU(sku); gpuConfig != nil {
-		return gpuConfig.GPUCount
-	}
-	return 1
-}
-
 func GetTuningImageInfo() string {
 	presetObj := metadata.MustGet("base")
 	return utils.GetPresetImageName(presetObj.Name, presetObj.Tag)
@@ -147,151 +141,172 @@ func SetupTrainingOutputVolume(ctx context.Context, configMap *corev1.ConfigMap,
 	return resultsVolume, resultsVolumeMount, outputDir
 }
 
-func setupDefaultSharedVolumes(workspaceObj *kaitov1beta1.Workspace, cmName string) ([]corev1.Volume, []corev1.VolumeMount) {
-	var volumes []corev1.Volume
-	var volumeMounts []corev1.VolumeMount
-
-	shmVolume, shmVolumeMount := utils.ConfigSHMVolume()
-	volumes = append(volumes, shmVolume)
-	volumeMounts = append(volumeMounts, shmVolumeMount)
-
-	// Add shared volume for tuning parameters
-	cmVolume, cmVolumeMount := utils.ConfigCMVolume(cmName)
-	volumes = append(volumes, cmVolume)
-	volumeMounts = append(volumeMounts, cmVolumeMount)
+func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, revisionNum string,
+	model pkgmodel.Model, kubeClient client.Client) (client.Object, error) {
 
-	return volumes, volumeMounts
-}
+	var skuNumGPUs int
+	gpuConfig, err := utils.GetGPUConfigBySKU(workspaceObj.Resource.InstanceType)
+	if err != nil {
+		gpuConfig, err = utils.TryGetGPUConfigFromNode(ctx, kubeClient, workspaceObj.Status.WorkerNodes)
+		if err != nil {
+			defaultNumGPU := resource.MustParse(model.GetTuningParameters().GPUCountRequirement)
+			skuNumGPUs = int(defaultNumGPU.Value())
+		}
+	}
+	if gpuConfig != nil {
+		skuNumGPUs = gpuConfig.GPUCount
+	}
 
-func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, revisionNum string,
-	tuningObj *model.PresetParam, kubeClient client.Client) (client.Object, error) {
+	gctx := &generator.WorkspaceGeneratorContext{
+		Ctx:        ctx,
+		Workspace:  workspaceObj,
+		Model:      model,
+		KubeClient: kubeClient,
+	}
 
-	var defaultConfigName string
-	if workspaceObj.Tuning.Method == kaitov1beta1.TuningMethodLora {
-		defaultConfigName = kaitov1beta1.DefaultLoraConfigMapTemplate
-	} else if workspaceObj.Tuning.Method == kaitov1beta1.TuningMethodQLora {
-		defaultConfigName = kaitov1beta1.DefaultQloraConfigMapTemplate
+	podSpec, err := generator.GenerateManifest(gctx,
+		GenerateBasicTuningPodSpec(gpuConfig, skuNumGPUs),
+		SetTrainingResultVolume,
+		SetTrainingInput,
+		SetTrainingOutputImagePush,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate pod spec: %w", err)
 	}
-	configVolume, err := resources.EnsureConfigOrCopyFromDefault(ctx, kubeClient,
-		client.ObjectKey{
-			Namespace: workspaceObj.Namespace,
-			Name:      workspaceObj.Tuning.Config,
-		},
-		client.ObjectKey{Name: defaultConfigName},
+
+	jobObj, err := generator.GenerateManifest(gctx,
+		manifests.GenerateTuningJobManifest(revisionNum),
+		manifests.SetJobPodSpec(podSpec),
 	)
 	if err != nil {
+		return nil, fmt.Errorf("failed to generate job manifest: %w", err)
+	}
+
+	err = resources.CreateResource(ctx, jobObj, kubeClient)
+	if client.IgnoreAlreadyExists(err) != nil {
 		return nil, err
 	}
+	return jobObj, nil
+}
 
-	var imagePullSecrets []corev1.LocalObjectReference
-	var initContainers, sidecarContainers []corev1.Container
-	volumes, volumeMounts := setupDefaultSharedVolumes(workspaceObj, configVolume.Name)
+func GenerateBasicTuningPodSpec(gpuConfig *sku.GPUConfig, skuNumGPUs int) func(*generator.WorkspaceGeneratorContext, *corev1.PodSpec) error {
+	return func(ctx *generator.WorkspaceGeneratorContext, spec *corev1.PodSpec) error {
+		// additional volume
+		var volumes []corev1.Volume
+		var volumeMounts []corev1.VolumeMount
+		var initContainers []corev1.Container
+
+		// add share memory for cross process communication
+		shmVolume, shmVolumeMount := utils.ConfigSHMVolume()
+		volumes = append(volumes, shmVolume)
+		volumeMounts = append(volumeMounts, shmVolumeMount)
+
+		// Add volume for model weights access
+		volumes = append(volumes, utils.DefaultModelWeightsVolume)
+		volumeMounts = append(volumeMounts, utils.DefaultModelWeightsVolumeMount)
+		initContainers = append(initContainers, manifests.GenerateModelPullerContainer(ctx.Ctx, ctx.Workspace, ctx.Model.GetTuningParameters())...)
+
+		// resource requirements
+		resourceRequirements := corev1.ResourceRequirements{
+			Requests: corev1.ResourceList{
+				corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
+			},
+			Limits: corev1.ResourceList{
+				corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
+			},
+		}
 
-	// TODO: make containers only mount the volumes they need
+		// default env
+		var envVars []corev1.EnvVar
 
-	// Add volume for training output
-	outputVolume := workspaceObj.Tuning.Output.Volume
-	trainingOutputVolume, trainingOutputVolumeMount, outputDir := SetupTrainingOutputVolume(ctx, configVolume, outputVolume)
-	volumes = append(volumes, trainingOutputVolume)
-	volumeMounts = append(volumeMounts, trainingOutputVolumeMount)
+		// Append environment variable for default target modules if using Phi3 model
+		presetName := strings.ToLower(string(ctx.Workspace.Tuning.Preset.Name))
+		if strings.HasPrefix(presetName, "phi-3") {
+			envVars = append(envVars, corev1.EnvVar{
+				Name:  "DEFAULT_TARGET_MODULES",
+				Value: "k_proj,q_proj,v_proj,o_proj,gate_proj,down_proj,up_proj",
+			})
+		}
+		// Add Expandable Memory Feature to reduce Peak GPU Mem Usage
+		envVars = append(envVars, corev1.EnvVar{
+			Name:  "PYTORCH_CUDA_ALLOC_CONF",
+			Value: "expandable_segments:True",
+		})
 
-	// Add volume for training input
-	initContainer, dataSourceVolumes, dataSourceVolumeMounts := prepareDataSource(ctx, workspaceObj)
-	volumes = append(volumes, dataSourceVolumes...)
-	volumeMounts = append(volumeMounts, dataSourceVolumeMounts...)
-	if initContainer != nil && initContainer.Name != "" {
-		initContainers = append(initContainers, *initContainer)
-	}
+		// tuning commands
+		tuningParam := ctx.Model.GetInferenceParameters().DeepCopy()
+		commands := tuningParam.GetTuningCommand(pkgmodel.RuntimeContext{
+			SKUNumGPUs: skuNumGPUs,
+		})
 
-	// Add volume for model weights access
-	volumes = append(volumes, utils.DefaultModelWeightsVolume)
-	volumeMounts = append(volumeMounts, utils.DefaultModelWeightsVolumeMount)
-	initContainers = append(initContainers, manifests.GenerateModelPullerContainer(ctx, workspaceObj, tuningObj)...)
-
-	sidecarContainer, imagePushSecret, imagePushSecretVolume, imagePushSecretVolumeMount := prepareDataDestination(ctx, workspaceObj, outputDir)
-	volumes = append(volumes, imagePushSecretVolume...)
-	volumeMounts = append(volumeMounts, imagePushSecretVolumeMount...)
-	if sidecarContainer != nil {
-		pauseContainer := corev1.Container{
-			Name:            "pause",
-			Image:           "registry.k8s.io/pause:latest",
-			ImagePullPolicy: corev1.PullAlways,
+		spec.Tolerations = tolerations
+		spec.InitContainers = append(spec.InitContainers, initContainers...)
+		spec.Containers = []corev1.Container{
+			{
+				Name:         ctx.Workspace.Name,
+				Image:        GetTuningImageInfo(),
+				Command:      commands,
+				Resources:    resourceRequirements,
+				Ports:        containerPorts,
+				Env:          envVars,
+				VolumeMounts: volumeMounts,
+			},
 		}
-		sidecarContainers = append(sidecarContainers, pauseContainer, *sidecarContainer)
-	}
-	if imagePushSecret != nil {
-		imagePullSecrets = append(imagePullSecrets, *imagePushSecret)
+		spec.Volumes = volumes
+		spec.RestartPolicy = corev1.RestartPolicyNever
+		return nil
 	}
+}
 
-	modelCommand, err := prepareModelRunParameters(ctx, tuningObj)
-	if err != nil {
-		return nil, err
+func SetTrainingResultVolume(ctx *generator.WorkspaceGeneratorContext, spec *corev1.PodSpec) error {
+	var defaultConfigName string
+	if ctx.Workspace.Tuning.Method == kaitov1beta1.TuningMethodLora {
+		defaultConfigName = kaitov1beta1.DefaultLoraConfigMapTemplate
+	} else if ctx.Workspace.Tuning.Method == kaitov1beta1.TuningMethodQLora {
+		defaultConfigName = kaitov1beta1.DefaultQloraConfigMapTemplate
 	}
-
-	var skuNumGPUs int
-	gpuConfig, err := utils.GetGPUConfigBySKU(workspaceObj.Resource.InstanceType)
+	configVolume, err := resources.EnsureConfigOrCopyFromDefault(ctx.Ctx, ctx.KubeClient,
+		client.ObjectKey{
+			Namespace: ctx.Workspace.Namespace,
+			Name:      ctx.Workspace.Tuning.Config,
+		},
+		client.ObjectKey{Name: defaultConfigName},
+	)
 	if err != nil {
-		gpuConfig, err = utils.TryGetGPUConfigFromNode(ctx, kubeClient, workspaceObj.Status.WorkerNodes)
-		if err != nil {
-			defaultNumGPU := resource.MustParse(tuningObj.GPUCountRequirement)
-			skuNumGPUs = int(defaultNumGPU.Value())
-		}
-	}
-	if gpuConfig != nil {
-		skuNumGPUs = gpuConfig.GPUCount
+		return err
 	}
 
-	commands, resourceReq := prepareTuningParameters(ctx, workspaceObj, modelCommand, tuningObj, skuNumGPUs)
-	tuningImage := GetTuningImageInfo()
+	// Add shared volume for tuning parameters
+	cmVolume, cmVolumeMount := utils.ConfigCMVolume(configVolume.Name)
 
-	var envVars []corev1.EnvVar
-	presetName := strings.ToLower(string(workspaceObj.Tuning.Preset.Name))
-	// Append environment variable for default target modules if using Phi3 model
-	if strings.HasPrefix(presetName, "phi-3") {
-		envVars = append(envVars, corev1.EnvVar{
-			Name:  "DEFAULT_TARGET_MODULES",
-			Value: "k_proj,q_proj,v_proj,o_proj,gate_proj,down_proj,up_proj",
-		})
+	// Add results volume for training output
+	outputDir, err := GetTrainingOutputDir(ctx.Ctx, configVolume)
+	if err != nil {
+		return fmt.Errorf("failed to get training output directory from config: %w", err)
 	}
-	// Add Expandable Memory Feature to reduce Peak GPU Mem Usage
-	envVars = append(envVars, corev1.EnvVar{
-		Name:  "PYTORCH_CUDA_ALLOC_CONF",
-		Value: "expandable_segments:True",
-	})
-
-	// remove duplicate volumes
-	seen := make(map[string]bool)
-	var deduplicatedVolumes []corev1.Volume
-	for _, vol := range volumes {
-		if !seen[vol.Name] {
-			deduplicatedVolumes = append(deduplicatedVolumes, vol)
-			seen[vol.Name] = true
+	resultsVolume, resultsVolumeMount := utils.ConfigResultsVolume(outputDir, ctx.Workspace.Tuning.Output.Volume)
+
+	volumes := []corev1.Volume{cmVolume, resultsVolume}
+	volumeMounts := []corev1.VolumeMount{cmVolumeMount, resultsVolumeMount}
+	for i := range spec.Containers {
+		if spec.Containers[i].Name == ctx.Workspace.Name {
+			spec.Containers[i].VolumeMounts = append(spec.Containers[i].VolumeMounts, volumeMounts...)
+			spec.Volumes = append(spec.Volumes, volumes...)
+			break
 		}
 	}
-	volumes = deduplicatedVolumes
-
-	jobObj := manifests.GenerateTuningJobManifest(workspaceObj, revisionNum, tuningImage, imagePullSecrets, *workspaceObj.Resource.Count, commands,
-		containerPorts, nil, nil, resourceReq, tolerations, initContainers, sidecarContainers, volumes, volumeMounts, envVars)
 
-	err = resources.CreateResource(ctx, jobObj, kubeClient)
-	if client.IgnoreAlreadyExists(err) != nil {
-		return nil, err
-	}
-	return jobObj, nil
+	return nil
 }
 
 // Now there are two options for data destination 1. Volume - 2. Image
-func prepareDataDestination(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, inputDirectory string) (*corev1.Container, *corev1.LocalObjectReference, []corev1.Volume, []corev1.VolumeMount) {
-	tuning := workspaceObj.Tuning
+// notes: this modifier requires the results volume to be set in the pod spec
+func SetTrainingOutputImagePush(ctx *generator.WorkspaceGeneratorContext, spec *corev1.PodSpec) error {
+	tuning := ctx.Workspace.Tuning
 	output := tuning.Output
 
-	if output.Volume != nil {
-		return nil, nil, nil, nil
-	}
-
-	outputImage := output.Image
-	if outputImage == "" {
-		return nil, nil, nil, nil
+	if output.Volume != nil || output.Image == "" {
+		return nil
 	}
 
 	var annotationsData map[string]map[string]string
@@ -303,15 +318,46 @@ func prepareDataDestination(ctx context.Context, workspaceObj *kaitov1beta1.Work
 		}
 	}
 
-	pusherContainer := image.NewPusherContainer(inputDirectory, outputImage, annotationsData, nil)
-
+	// additional secret volume for image push
 	imagePushSecretRef := corev1.LocalObjectReference{
 		Name: output.ImagePushSecret,
 	}
+	spec.ImagePullSecrets = append(spec.ImagePullSecrets, imagePushSecretRef)
+	secretVolume, secretVolumeMount := utils.ConfigImagePushSecretVolume(imagePushSecretRef.Name)
+	spec.Volumes = append(spec.Volumes, secretVolume)
+
+	// additional sidecar container for uploading image
+	resultVolumeMount := utils.FindResultsVolumeMount(spec)
+	if resultVolumeMount == nil {
+		return fmt.Errorf("results volume mount not found in pod spec")
+	}
+	inputDirectory := resultVolumeMount.MountPath
+	pusherContainer := image.NewPusherContainer(inputDirectory, output.Image, annotationsData, nil)
+	pusherContainer.VolumeMounts = append(pusherContainer.VolumeMounts, secretVolumeMount, *resultVolumeMount)
+	pauseContainer := corev1.Container{
+		Name:            "pause",
+		Image:           "registry.k8s.io/pause:latest",
+		ImagePullPolicy: corev1.PullAlways,
+	}
+	spec.Containers = append(spec.Containers, *pusherContainer, pauseContainer)
+	return nil
+}
+
+func SetTrainingInput(ctx *generator.WorkspaceGeneratorContext, spec *corev1.PodSpec) error {
+	initContainer, dataSourceVolumes, dataSourceVolumeMounts := prepareDataSource(ctx.Ctx, ctx.Workspace)
+	if initContainer != nil && initContainer.Name != "" {
+		spec.InitContainers = append(spec.InitContainers, *initContainer)
+	}
 
-	volume, volumeMount := utils.ConfigImagePushSecretVolume(imagePushSecretRef.Name)
+	spec.Volumes = append(spec.Volumes, dataSourceVolumes...)
 
-	return pusherContainer, &imagePushSecretRef, []corev1.Volume{volume}, []corev1.VolumeMount{volumeMount}
+	for _, volumeMount := range dataSourceVolumeMounts {
+		for i := range spec.Containers {
+			spec.Containers[i].VolumeMounts = append(spec.Containers[i].VolumeMounts, volumeMount)
+		}
+	}
+
+	return nil
 }
 
 // Now there are three options for DataSource: 1. URL - 2. Volume - 3. Image
@@ -323,6 +369,7 @@ func prepareDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace
 		dataVolume, dataVolumeMount := utils.ConfigDataVolume(nil)
 		imagePullSecretVolume, imagePullSecretVolumeMount := utils.ConfigImagePullSecretVolume(input.Name+"-tuning-input", input.ImagePullSecrets)
 		pullerContainer := image.NewPullerContainer(input.Image, utils.DefaultDataVolumePath)
+		pullerContainer.VolumeMounts = append(pullerContainer.VolumeMounts, imagePullSecretVolumeMount, dataVolumeMount)
 		return pullerContainer, []corev1.Volume{imagePullSecretVolume, dataVolume}, []corev1.VolumeMount{imagePullSecretVolumeMount, dataVolumeMount}
 
 	case len(input.URLs) > 0:
@@ -385,36 +432,3 @@ func handleURLDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspa
 	volume, volumeMount := utils.ConfigDataVolume(nil)
 	return initContainer, volume, volumeMount
 }
-
-func prepareModelRunParameters(ctx context.Context, tuningObj *model.PresetParam) (string, error) {
-	modelCommand := utils.BuildCmdStr(TuningFile, tuningObj.Transformers.ModelRunParams)
-	return modelCommand, nil
-}
-
-// prepareTuningParameters builds a PyTorch command:
-// accelerate launch <TORCH_PARAMS> baseCommand <MODEL_PARAMS>
-// and sets the GPU resources required for tuning.
-// Returns the command and resource configuration.
-func prepareTuningParameters(ctx context.Context, wObj *kaitov1beta1.Workspace, modelCommand string,
-	tuningObj *model.PresetParam, skuNumGPUs int) ([]string, corev1.ResourceRequirements) {
-	hfParam := tuningObj.Transformers // Only support Huggingface for now
-	if hfParam.AccelerateParams == nil {
-		hfParam.AccelerateParams = make(map[string]string)
-	}
-	// Set # of processes to GPU Count
-	numProcesses := getInstanceGPUCount(wObj.Resource.InstanceType)
-	hfParam.AccelerateParams["num_processes"] = fmt.Sprintf("%d", numProcesses)
-	torchCommand := utils.BuildCmdStr(hfParam.BaseCommand, hfParam.AccelerateParams)
-	commands := utils.ShellCmd(torchCommand + " " + modelCommand)
-
-	resourceRequirements := corev1.ResourceRequirements{
-		Requests: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
-		},
-		Limits: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
-		},
-	}
-
-	return commands, resourceRequirements
-}
diff --git a/presets/workspace/inference/text-generation/inference_api.py b/presets/workspace/inference/text-generation/inference_api.py
index 20b0102..ff851b6 100644
--- a/presets/workspace/inference/text-generation/inference_api.py
+++ b/presets/workspace/inference/text-generation/inference_api.py
@@ -297,7 +297,7 @@ class UnifiedRequestModel(BaseModel):
     # Field for conversational model
     messages: Optional[List[Message]] = Field(None, description="Messages for conversational model. Required for conversational pipeline. Do not use with 'prompt'.")
     def messages_to_dict_list(self):
-        return [message.dict() for message in self.messages] if self.messages else []
+        return [message.model_dump() for message in self.messages] if self.messages else []
 
 class ErrorResponse(BaseModel):
     detail: str
@@ -365,7 +365,7 @@ def generate_text(
                             "clean_up_tokenization_spaces": False,
                             "prefix": None,
                             "handle_long_generation": None,
-                            "generate_kwargs": GenerateKwargs().dict(),
+                            "generate_kwargs": GenerateKwargs().model_dump(),
                         },
                     },
                     "conversation_example": {
@@ -381,7 +381,7 @@ def generate_text(
                             "clean_up_tokenization_spaces": False,
                             "prefix": None,
                             "handle_long_generation": None,
-                            "generate_kwargs": GenerateKwargs().dict(),
+                            "generate_kwargs": GenerateKwargs().model_dump(),
                         },
                     },
                 },
@@ -392,7 +392,7 @@ def generate_text(
     Processes chat requests, generating text based on the specified pipeline (text generation or conversational).
     Validates required parameters based on the pipeline and returns the generated text.
     """
-    user_generate_kwargs = request_model.generate_kwargs.dict() if request_model.generate_kwargs else {}
+    user_generate_kwargs = request_model.generate_kwargs.model_dump() if request_model.generate_kwargs else {}
     generate_kwargs = {**default_generate_config, **user_generate_kwargs}
 
     if args.pipeline == "text-generation":
diff --git a/website/docs/headlamp-kaito.md b/website/docs/headlamp-kaito.md
new file mode 100644
index 0000000..d93d300
--- /dev/null
+++ b/website/docs/headlamp-kaito.md
@@ -0,0 +1,36 @@
+---
+title: Headlamp KAITO
+---
+
+![Headlamp](/img/Headlamp.png)
+
+Headlamp plugin for KAITO is a web-based dashboard extension for managing and monitoring KAITO Workspaces through the Headlamp Kubernetes dashboard. This extension provides an intuitive interface for operators to deploy, monitor, and manage AI workloads.
+
+## Overview
+
+Headlamp plugin for KAITO is a plugin for [Headlamp](https://headlamp.dev/), the Kubernetes web UI, that adds specialized views and management capabilities for KAITO resources. Key features include:
+
+- **Visual Workspace Management**: Create, view, and manage KAITO workspaces through a user-friendly interface
+- **Resource Monitoring**: Real-time visibility into GPU utilization, inference metrics, and system health
+- **Model Deployment**: Simplified deployment of pre-configured AI models from KAITO's preset library
+
+![Headlamp-KAITO plugin installation](/img/headlamp-kaito-plugin-installation.png)
+
+## Getting Started
+
+For detailed installation instructions, usage guides, and configuration options, please refer to the official documentation:
+
+![Headlamp-KAITO Documentation](/img/headlamp-kaito-docs.png)
+
+**[Headlamp-KAITO Documentation](https://kaito-project.github.io/headlamp-kaito/)**
+
+## Resources
+
+- **[GitHub Repository](https://github.com/kaito-project/headlamp-kaito)** - Source code, issues, and contributions
+- **[Documentation](https://kaito-project.github.io/headlamp-kaito/)** - Complete setup and usage guide
+- **[Headlamp](https://headlamp.dev/)** - The underlying Kubernetes dashboard
+
+## Support
+
+For issues:
+- Report bugs and request features on [GitHub Issues](https://github.com/kaito-project/headlamp-kaito/issues)
diff --git a/website/docs/proposals/20250611-workspace-subresource-scale-api.md b/website/docs/proposals/20250611-workspace-subresource-scale-api.md
index 216ac21..63990c0 100644
--- a/website/docs/proposals/20250611-workspace-subresource-scale-api.md
+++ b/website/docs/proposals/20250611-workspace-subresource-scale-api.md
@@ -49,7 +49,7 @@ To ensure different auto-scaler solutions can integrate with Kaito to manage inf
 
 Workspace CRD will support scale subresource API, so different auto-scaler solutions can adjust Workspace replicas dynamically. and Workspace controller/webhook will update the underlay workloads(like Deployment, LWS, etc.) which are used for managing inference workloads. The detailed workflow is shown in the following figure:
 
-![auto-scaler](/img/workspace-scale-api.png)
+![auto-scaler](../../static/img/workspace-scale-api.png)
 
 - Workspace CRD: upgrade workspace CRD for supporting scale subresource api.
 
@@ -120,11 +120,11 @@ type InferenceStatus struct {
 
 Workspace webhook should be improved for initializing workspace status and validate scale request is allowed or not.
 
-![webhook-for-ws-create-req](/img/workspace-webhook-for-create-req.png)
+![webhook-for-ws-create-req](../../static/img/workspace-webhook-for-create-req.png)
 
 - gpu nodes estimator will be used to calculate the `PerReplicaNodeCount` for workspace.
 
-![webhook-for-scale-req](/img/workspace-webhook-for-scale-req.png)
+![webhook-for-scale-req](../../static/img/workspace-webhook-for-scale-req.png)
 
 - For workspace without inference settings or workspace with prefered nodes, the scale request will be rejected.
 
@@ -132,7 +132,7 @@ Workspace webhook should be improved for initializing workspace status and valid
 
 The workspace controller should reconcile TargetNodeCount in order to keep consistent with scaling action. then scale underlay workloads(Deployment) and update workspace status. The workflow should be improved as following:
 
-![workspace-controller-workflow](/img/workspace-controller-workflow.png)
+![workspace-controller-workflow](../../static/img/workspace-controller-workflow.png)
 
 ### Workspace CRD Subresource Change
 
diff --git a/website/docs/proposals/20250704-keda-scaler-for-inference-workloads.md b/website/docs/proposals/20250704-keda-scaler-for-inference-workloads.md
index 94a46a6..bf14135 100644
--- a/website/docs/proposals/20250704-keda-scaler-for-inference-workloads.md
+++ b/website/docs/proposals/20250704-keda-scaler-for-inference-workloads.md
@@ -30,7 +30,7 @@ From the technical perspective, we don't want to make a new wheel for auto-scale
 
 The auto-scaler solution for Kaito should be shown as follows:
 
-![keda-kaito-scaler](/img/keda-kaito-scaler.png)
+![keda-kaito-scaler](../../static/img/keda-kaito-scaler.png)
 
 We will divide this auto-scaler feature into two parts as follows:
 
@@ -53,7 +53,7 @@ To ensure ease of use, the specialized kaito scaler is hosted in an independent
 
 Keda can provides both metric-based and time-based scaler capability, but for metric-based scaler it should work together with prometheus. The detailed auto-scaler architecture is shown in the following figure:
 
-![keda-prometheus-cron-auto-scaler](/img/keda-prometheus-cron-auto-scaler.png)
+![keda-prometheus-cron-auto-scaler](../../static/img/keda-prometheus-cron-auto-scaler.png)
 
 - **Keda**: includes two components: metrics-adapter and keda-core. metrics-adpter is used for exposing external metrics on kube-apiserver wihch will be used by native HPA. keda-core includes the core logic of keda system, like the generation of HPA resource according to ScaledObject, and supporting different scalers(like prometheus, cron scalers).
 - **Scalers**: Keda providers more than 100+ scalers, but these scalers are only used for providing metrics for native HPA, and the scaling logic and actions are taken by native HPA.
@@ -353,7 +353,7 @@ In order to address the above challenges, we consider to create a new component
 - kaito scaler: works as an external scaler of keda, and collects metrics from inference pods according to ScaledObject configurations. This means that `PodMonitor`, `TriggerAuthentication`, `Secret` configurations aren't needed, and users only need to configure `ScaledObject`. at the same time, users don't need to maintain the Prometheus stack.
 - kaito scaler manager: works as a deployment, includes webhooks for configuring default values for `ScaledObject`, so users don't need to dive into the details of HPA behavior or vllm metrics. and controllers for ensuring secret which used by grpc connection between keda core and kaito scaler(external scaler).
 
-![keda-kaito-scaler-arch](/img/keda-kaito-scaler-arch.png)
+![keda-kaito-scaler-arch](../../static/img/keda-kaito-scaler-arch.png)
 
 #### Trigger Specification of Kaito Scaler
 
diff --git a/website/sidebars.js b/website/sidebars.js
index ebfd102..8e67322 100644
--- a/website/sidebars.js
+++ b/website/sidebars.js
@@ -47,6 +47,7 @@ const sidebars = {
                 'rag',
                 'custom-model',
                 'tool-calling',
+                'headlamp-kaito',
             ],
         },
         {
diff --git a/website/static/img/Headlamp.png b/website/static/img/Headlamp.png
new file mode 100644
index 0000000..68eaa7e
Binary files /dev/null and b/website/static/img/Headlamp.png differ
diff --git a/website/static/img/headlamp-kaito-docs.png b/website/static/img/headlamp-kaito-docs.png
new file mode 100644
index 0000000..4797179
Binary files /dev/null and b/website/static/img/headlamp-kaito-docs.png differ
diff --git a/website/static/img/headlamp-kaito-plugin-installation.png b/website/static/img/headlamp-kaito-plugin-installation.png
new file mode 100644
index 0000000..22d4b75
Binary files /dev/null and b/website/static/img/headlamp-kaito-plugin-installation.png differ


Test patch:
diff --git a/pkg/workspace/controllers/workspace_controller_test.go b/pkg/workspace/controllers/workspace_controller_test.go
index 7b2702a..9e7c0c4 100644
--- a/pkg/workspace/controllers/workspace_controller_test.go
+++ b/pkg/workspace/controllers/workspace_controller_test.go
@@ -863,7 +863,7 @@ func TestApplyWorkspaceResource(t *testing.T) {
 	}
 }
 
-func TestUpdateControllerRevision1(t *testing.T) {
+func TestSyncControllerRevision(t *testing.T) {
 	testcases := map[string]struct {
 		callMocks     func(c *test.MockClient)
 		workspace     v1beta1.Workspace
@@ -883,9 +883,17 @@ func TestUpdateControllerRevision1(t *testing.T) {
 									WorkspaceHashAnnotation: "1171dc5d15043c92e684c8f06689eb241763a735181fdd2b59c8bd8fd6eecdd4",
 								},
 							},
+							Revision: 1,
 						}
 					}).
 					Return(nil)
+				// Add mock for workspace retrieval in updateWorkspaceWithRetry
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1beta1.Workspace{}), mock.Anything).
+					Run(func(args mock.Arguments) {
+						ws := args.Get(2).(*v1beta1.Workspace)
+						*ws = test.MockWorkspaceWithComputeHash
+					}).
+					Return(nil)
 				c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1beta1.Workspace{}), mock.Anything).
 					Return(nil)
 			},
@@ -894,7 +902,7 @@ func TestUpdateControllerRevision1(t *testing.T) {
 			verifyCalls: func(c *test.MockClient) {
 				c.AssertNumberOfCalls(t, "List", 1)
 				c.AssertNumberOfCalls(t, "Create", 0)
-				c.AssertNumberOfCalls(t, "Get", 1)
+				c.AssertNumberOfCalls(t, "Get", 2) // 1 for ControllerRevision, 1 for Workspace
 				c.AssertNumberOfCalls(t, "Delete", 0)
 				c.AssertNumberOfCalls(t, "Update", 1)
 			},
@@ -906,8 +914,6 @@ func TestUpdateControllerRevision1(t *testing.T) {
 				c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(errors.New("failed to create ControllerRevision"))
 				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).
 					Return(apierrors.NewNotFound(appsv1.Resource("ControllerRevision"), test.MockWorkspaceFailToCreateCR.Name))
-				c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1beta1.Workspace{}), mock.Anything).
-					Return(nil)
 			},
 			workspace:     test.MockWorkspaceFailToCreateCR,
 			expectedError: errors.New("failed to create new ControllerRevision: failed to create ControllerRevision"),
@@ -926,6 +932,13 @@ func TestUpdateControllerRevision1(t *testing.T) {
 				c.On("Create", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil)
 				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).
 					Return(apierrors.NewNotFound(appsv1.Resource("ControllerRevision"), test.MockWorkspaceFailToCreateCR.Name))
+				// Add mock for workspace retrieval in updateWorkspaceWithRetry
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1beta1.Workspace{}), mock.Anything).
+					Run(func(args mock.Arguments) {
+						ws := args.Get(2).(*v1beta1.Workspace)
+						*ws = test.MockWorkspaceSuccessful
+					}).
+					Return(nil)
 				c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1beta1.Workspace{}), mock.Anything).
 					Return(nil)
 			},
@@ -934,7 +947,7 @@ func TestUpdateControllerRevision1(t *testing.T) {
 			verifyCalls: func(c *test.MockClient) {
 				c.AssertNumberOfCalls(t, "List", 1)
 				c.AssertNumberOfCalls(t, "Create", 1)
-				c.AssertNumberOfCalls(t, "Get", 1)
+				c.AssertNumberOfCalls(t, "Get", 2) // 1 for ControllerRevision, 1 for Workspace
 				c.AssertNumberOfCalls(t, "Delete", 0)
 				c.AssertNumberOfCalls(t, "Update", 1)
 			},
@@ -967,6 +980,13 @@ func TestUpdateControllerRevision1(t *testing.T) {
 				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).
 					Return(apierrors.NewNotFound(appsv1.Resource("ControllerRevision"), test.MockWorkspaceFailToCreateCR.Name))
 				c.On("Delete", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil)
+				// Add mock for workspace retrieval in updateWorkspaceWithRetry
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1beta1.Workspace{}), mock.Anything).
+					Run(func(args mock.Arguments) {
+						ws := args.Get(2).(*v1beta1.Workspace)
+						*ws = test.MockWorkspaceWithDeleteOldCR
+					}).
+					Return(nil)
 				c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1beta1.Workspace{}), mock.Anything).
 					Return(nil)
 			},
@@ -975,7 +995,7 @@ func TestUpdateControllerRevision1(t *testing.T) {
 			verifyCalls: func(c *test.MockClient) {
 				c.AssertNumberOfCalls(t, "List", 1)
 				c.AssertNumberOfCalls(t, "Create", 1)
-				c.AssertNumberOfCalls(t, "Get", 1)
+				c.AssertNumberOfCalls(t, "Get", 2) // 1 for ControllerRevision, 1 for Workspace
 				c.AssertNumberOfCalls(t, "Delete", 1)
 				c.AssertNumberOfCalls(t, "Update", 1)
 			},
@@ -1008,6 +1028,13 @@ func TestUpdateControllerRevision1(t *testing.T) {
 				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).
 					Return(apierrors.NewNotFound(appsv1.Resource("ControllerRevision"), test.MockWorkspaceFailToCreateCR.Name))
 				c.On("Delete", mock.IsType(context.Background()), mock.IsType(&appsv1.ControllerRevision{}), mock.Anything).Return(nil)
+				// Add mock for workspace retrieval in updateWorkspaceWithRetry
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&v1beta1.Workspace{}), mock.Anything).
+					Run(func(args mock.Arguments) {
+						ws := args.Get(2).(*v1beta1.Workspace)
+						*ws = test.MockWorkspaceUpdateCR
+					}).
+					Return(nil)
 				c.On("Update", mock.IsType(context.Background()), mock.IsType(&v1beta1.Workspace{}), mock.Anything).
 					Return(fmt.Errorf("failed to update Workspace annotations"))
 			},
@@ -1016,7 +1043,7 @@ func TestUpdateControllerRevision1(t *testing.T) {
 			verifyCalls: func(c *test.MockClient) {
 				c.AssertNumberOfCalls(t, "List", 1)
 				c.AssertNumberOfCalls(t, "Create", 1)
-				c.AssertNumberOfCalls(t, "Get", 1)
+				c.AssertNumberOfCalls(t, "Get", 2) // 1 for ControllerRevision, 1 for Workspace
 				c.AssertNumberOfCalls(t, "Delete", 1)
 				c.AssertNumberOfCalls(t, "Update", 1)
 			},
diff --git a/pkg/workspace/inference/preset-inferences_test.go b/pkg/workspace/inference/preset-inferences_test.go
index a105543..0d3ccbf 100644
--- a/pkg/workspace/inference/preset-inferences_test.go
+++ b/pkg/workspace/inference/preset-inferences_test.go
@@ -638,6 +638,405 @@ func TestGetGPUConfig(t *testing.T) {
 		})
 	}
 }
+func TestSetAdapterPuller(t *testing.T) {
+	testcases := map[string]struct {
+		workspace       *v1beta1.Workspace
+		spec            *corev1.PodSpec
+		expectedVolumes []string
+		expectedEnvVars []string
+		expectError     bool
+	}{
+		"no adapters": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Adapters: []v1beta1.AdapterSpec{}, // empty adapters
+				},
+			},
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-container",
+					},
+				},
+			},
+			expectedVolumes: []string{},
+			expectedEnvVars: []string{},
+			expectError:     false,
+		},
+		"single adapter": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Adapters: []v1beta1.AdapterSpec{
+						{
+							Source: &v1beta1.DataSource{
+								Name:  "adapter-1",
+								Image: "test-registry/adapter:v1",
+							},
+							Strength: &ValidStrength,
+						},
+					},
+				},
+			},
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-container",
+					},
+				},
+			},
+			expectedVolumes: []string{"adapter-volume"},
+			expectedEnvVars: []string{"adapter-1"},
+			expectError:     false,
+		},
+		"multiple adapters": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Adapters: []v1beta1.AdapterSpec{
+						{
+							Source: &v1beta1.DataSource{
+								Name:  "adapter-1",
+								Image: "test-registry/adapter1:v1",
+							},
+							Strength: &ValidStrength,
+						},
+						{
+							Source: &v1beta1.DataSource{
+								Name:  "adapter-2",
+								Image: "test-registry/adapter2:v1",
+							},
+							Strength: &ValidStrength,
+						},
+					},
+				},
+			},
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-container",
+					},
+				},
+			},
+			expectedVolumes: []string{"adapter-volume"},
+			expectedEnvVars: []string{"adapter-1", "adapter-2"},
+			expectError:     false,
+		},
+		"adapters with image pull secrets": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Adapters: []v1beta1.AdapterSpec{
+						{
+							Source: &v1beta1.DataSource{
+								Name:             "adapter-1",
+								Image:            "private-registry/adapter:v1",
+								ImagePullSecrets: []string{"my-secret"},
+							},
+							Strength: &ValidStrength,
+						},
+					},
+				},
+			},
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-container",
+					},
+				},
+			},
+			expectedVolumes: []string{"adapter-volume", "docker-config-adapter-1-inference-adapter"},
+			expectedEnvVars: []string{"adapter-1"},
+			expectError:     false,
+		},
+		"multiple containers": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Adapters: []v1beta1.AdapterSpec{
+						{
+							Source: &v1beta1.DataSource{
+								Name:  "adapter-1",
+								Image: "test-registry/adapter:v1",
+							},
+							Strength: &ValidStrength,
+						},
+					},
+				},
+			},
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "container-1",
+					},
+					{
+						Name: "container-2",
+					},
+				},
+			},
+			expectedVolumes: []string{"adapter-volume"},
+			expectedEnvVars: []string{"adapter-1"},
+			expectError:     false,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			ctx := &generator.WorkspaceGeneratorContext{
+				Ctx:       context.TODO(),
+				Workspace: tc.workspace,
+			}
+
+			err := SetAdapterPuller(ctx, tc.spec)
+
+			if tc.expectError && err == nil {
+				t.Errorf("expected error but got nil")
+			}
+			if !tc.expectError && err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			// Check volumes
+			actualVolumes := make([]string, 0)
+			for _, volume := range tc.spec.Volumes {
+				actualVolumes = append(actualVolumes, volume.Name)
+			}
+			if !reflect.DeepEqual(actualVolumes, tc.expectedVolumes) {
+				t.Errorf("volumes mismatch: expected %v, got %v", tc.expectedVolumes, actualVolumes)
+			}
+
+			// Check volume mounts on containers
+			if len(tc.expectedVolumes) > 0 {
+				for _, container := range tc.spec.Containers {
+					foundAdapterMount := false
+					for _, mount := range container.VolumeMounts {
+						if mount.Name == "adapter-volume" {
+							foundAdapterMount = true
+							if mount.MountPath != "/mnt/adapter" {
+								t.Errorf("unexpected mount path for adapter volume: %s", mount.MountPath)
+							}
+							break
+						}
+					}
+					if !foundAdapterMount {
+						t.Errorf("adapter volume mount not found in container %s", container.Name)
+					}
+				}
+			}
+
+			// Check environment variables
+			if len(tc.expectedEnvVars) > 0 {
+				for _, container := range tc.spec.Containers {
+					actualEnvVarNames := make([]string, 0)
+					for _, env := range container.Env {
+						actualEnvVarNames = append(actualEnvVarNames, env.Name)
+					}
+
+					// Check if all expected env vars are present
+					for _, expectedVar := range tc.expectedEnvVars {
+						found := false
+						for _, actualVar := range actualEnvVarNames {
+							if actualVar == expectedVar {
+								found = true
+								break
+							}
+						}
+						if !found {
+							t.Errorf("expected env var %s not found in container %s", expectedVar, container.Name)
+						}
+					}
+				}
+			}
+
+			// Check init containers
+			if len(tc.workspace.Inference.Adapters) > 0 {
+				if len(tc.spec.InitContainers) == 0 {
+					t.Errorf("expected init containers for adapter pulling but found none")
+				}
+			}
+		})
+	}
+}
+
+func TestSetModelDownloadInfo(t *testing.T) {
+	test.RegisterTestModel()
+
+	testcases := map[string]struct {
+		workspace             *v1beta1.Workspace
+		modelName             string
+		spec                  *corev1.PodSpec
+		expectedEnvVars       []corev1.EnvVar
+		expectedInitContainer int
+		expectError           bool
+		expectedErrorMsg      string
+	}{
+		"download at runtime - add HF_TOKEN": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Preset: &v1beta1.PresetSpec{
+						PresetMeta: v1beta1.PresetMeta{
+							Name: "test-model-download",
+						},
+						PresetOptions: v1beta1.PresetOptions{
+							ModelAccessSecret: "hf-secret",
+						},
+					},
+				},
+			},
+			modelName: "test-model-download",
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "container-1",
+					},
+				},
+			},
+			expectedEnvVars: []corev1.EnvVar{
+				{
+					Name: "HF_TOKEN",
+					ValueFrom: &corev1.EnvVarSource{
+						SecretKeyRef: &corev1.SecretKeySelector{
+							LocalObjectReference: corev1.LocalObjectReference{
+								Name: "hf-secret",
+							},
+							Key: "HF_TOKEN",
+						},
+					},
+				},
+			},
+			expectedInitContainer: 0,
+			expectError:           false,
+		},
+		"model puller - add init containers": {
+			workspace: &v1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Inference: &v1beta1.InferenceSpec{
+					Preset: &v1beta1.PresetSpec{
+						PresetMeta: v1beta1.PresetMeta{
+							Name: "test-model",
+						},
+					},
+				},
+			},
+			modelName: "test-model",
+			spec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-container",
+					},
+				},
+			},
+			expectedEnvVars:       []corev1.EnvVar{},
+			expectedInitContainer: 1, // Expecting model-weights-downloader
+			expectError:           false,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			model := plugin.KaitoModelRegister.MustGet(tc.modelName)
+
+			ctx := &generator.WorkspaceGeneratorContext{
+				Ctx:       context.TODO(),
+				Workspace: tc.workspace,
+				Model:     model,
+			}
+
+			// Store original state
+			originalInitContainerCount := len(tc.spec.InitContainers)
+
+			err := SetModelDownloadInfo(ctx, tc.spec)
+
+			// Check error
+			if tc.expectError {
+				if err == nil {
+					t.Errorf("expected error but got nil")
+				} else if tc.expectedErrorMsg != "" && err.Error() != tc.expectedErrorMsg {
+					t.Errorf("expected error message %q, got %q", tc.expectedErrorMsg, err.Error())
+				}
+			} else if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			// Skip further checks if error is expected
+			if tc.expectError {
+				return
+			}
+
+			// Check environment variables if expected
+			if len(tc.expectedEnvVars) > 0 {
+				for i, container := range tc.spec.Containers {
+					found := false
+					for _, env := range container.Env {
+						if env.Name == "HF_TOKEN" {
+							found = true
+							if !reflect.DeepEqual(env, tc.expectedEnvVars[0]) {
+								t.Errorf("container %d: HF_TOKEN env var mismatch: expected %+v, got %+v",
+									i, tc.expectedEnvVars[0], env)
+							}
+						}
+					}
+					if !found {
+						t.Errorf("container %d: HF_TOKEN env var not found", i)
+					}
+				}
+			} else {
+				// Verify no HF_TOKEN env vars were added
+				for i, container := range tc.spec.Containers {
+					for _, env := range container.Env {
+						if env.Name == "HF_TOKEN" {
+							t.Errorf("container %d: unexpected HF_TOKEN env var found", i)
+						}
+					}
+				}
+			}
+
+			// Check init containers
+			actualInitContainerCount := len(tc.spec.InitContainers) - originalInitContainerCount
+			if actualInitContainerCount != tc.expectedInitContainer {
+				t.Errorf("init container count mismatch: expected %d new containers, got %d",
+					tc.expectedInitContainer, actualInitContainerCount)
+			}
+
+			if tc.expectedInitContainer > 0 {
+				// Look for model-weights-downloader container
+				found := false
+				for _, initContainer := range tc.spec.InitContainers {
+					if initContainer.Name == "model-weights-downloader" {
+						found = true
+						break
+					}
+				}
+				if !found {
+					t.Errorf("expected model-weights-downloader init container not found")
+				}
+			}
+		})
+	}
+}
 
 func toParameterMap(in []string) map[string]string {
 	ret := make(map[string]string)
diff --git a/pkg/workspace/tuning/preset-tuning_test.go b/pkg/workspace/tuning/preset-tuning_test.go
index 1615e18..9ba2f4a 100644
--- a/pkg/workspace/tuning/preset-tuning_test.go
+++ b/pkg/workspace/tuning/preset-tuning_test.go
@@ -20,12 +20,14 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	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/client/fake"
 
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
-	"github.com/kaito-project/kaito/pkg/model"
 	"github.com/kaito-project/kaito/pkg/utils"
-	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/generator"
 	"github.com/kaito-project/kaito/pkg/workspace/image"
 )
 
@@ -33,35 +35,6 @@ func normalize(s string) string {
 	return strings.Join(strings.Fields(s), " ")
 }
 
-func TestGetInstanceGPUCount(t *testing.T) {
-	t.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
-
-	testcases := map[string]struct {
-		sku              string
-		expectedGPUCount int
-	}{
-		"SKU Exists With Multiple GPUs": {
-			sku:              "Standard_NC24s_v3",
-			expectedGPUCount: 4,
-		},
-		"SKU Exists With One GPU": {
-			sku:              "Standard_NC6s_v3",
-			expectedGPUCount: 1,
-		},
-		"SKU Does Not Exist": {
-			sku:              "sku_unknown",
-			expectedGPUCount: 1,
-		},
-	}
-
-	for name, tc := range testcases {
-		t.Run(name, func(t *testing.T) {
-			result := getInstanceGPUCount(tc.sku)
-			assert.Equal(t, tc.expectedGPUCount, result)
-		})
-	}
-}
-
 func TestGetDataSrcImageInfo(t *testing.T) {
 	testcases := map[string]struct {
 		wObj            *kaitov1beta1.Workspace
@@ -219,88 +192,6 @@ func TestHandleURLDataSource(t *testing.T) {
 	}
 }
 
-func TestPrepareTuningParameters(t *testing.T) {
-	ctx := context.TODO()
-
-	testcases := map[string]struct {
-		name                 string
-		workspaceObj         *kaitov1beta1.Workspace
-		modelCommand         string
-		tuningObj            *model.PresetParam
-		expectedCommands     []string
-		expectedRequirements corev1.ResourceRequirements
-	}{
-		"Basic Tuning Parameters Setup": {
-			workspaceObj: &kaitov1beta1.Workspace{
-				Resource: kaitov1beta1.ResourceSpec{
-					InstanceType: "gpu-instance-type",
-				},
-			},
-			modelCommand: "model-command",
-			tuningObj: &model.PresetParam{
-				RuntimeParam: model.RuntimeParam{
-					Transformers: model.HuggingfaceTransformersParam{
-						BaseCommand:      "python train.py",
-						AccelerateParams: map[string]string{},
-					},
-				},
-				GPUCountRequirement: "2",
-			},
-			expectedCommands: []string{"/bin/sh", "-c", "python train.py --num_processes=1 model-command"},
-			expectedRequirements: corev1.ResourceRequirements{
-				Requests: corev1.ResourceList{
-					corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("2"),
-				},
-				Limits: corev1.ResourceList{
-					corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("2"),
-				},
-			},
-		},
-	}
-
-	for name, tc := range testcases {
-		t.Run(name, func(t *testing.T) {
-			t.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
-
-			commands, resources := prepareTuningParameters(ctx, tc.workspaceObj, tc.modelCommand, tc.tuningObj, 2)
-			assert.Equal(t, tc.expectedCommands, commands)
-			assert.True(t, tc.expectedRequirements.Requests.Name("nvidia.com/gpu", resource.DecimalSI).Equal(*resources.Requests.Name("nvidia.com/gpu", resource.DecimalSI)))
-			assert.True(t, tc.expectedRequirements.Limits.Name("nvidia.com/gpu", resource.DecimalSI).Equal(*resources.Limits.Name("nvidia.com/gpu", resource.DecimalSI)))
-		})
-	}
-}
-
-func TestPrepareDataDestination_ImageDestination(t *testing.T) {
-	ctx := context.TODO()
-
-	workspaceObj := &kaitov1beta1.Workspace{
-		Tuning: &kaitov1beta1.TuningSpec{
-			Output: &kaitov1beta1.DataDestination{
-				Image:           "custom/data-loader-image",
-				ImagePushSecret: "image-push-secret",
-			},
-		},
-	}
-
-	expectedImagePushSecret := corev1.LocalObjectReference{
-		Name: workspaceObj.Tuning.Output.ImagePushSecret,
-	}
-
-	expectedVolume, expectedVolumeMount := utils.ConfigImagePushSecretVolume(expectedImagePushSecret.Name)
-
-	outDir := "/mnt/results"
-
-	expectedSidecarContainer := image.NewPusherContainer(outDir, workspaceObj.Tuning.Output.Image, nil, nil)
-
-	sidecarContainer, imagePushSecret, volume, volumeMount := prepareDataDestination(ctx, workspaceObj, outDir)
-
-	// Assertions
-	assert.Equal(t, expectedSidecarContainer, sidecarContainer)
-	assert.Equal(t, []corev1.Volume{expectedVolume}, volume)
-	assert.Equal(t, []corev1.VolumeMount{expectedVolumeMount}, volumeMount)
-	assert.Equal(t, expectedImagePushSecret, *imagePushSecret)
-}
-
 func TestPrepareDataSource_ImageSource(t *testing.T) {
 	ctx := context.TODO()
 
@@ -360,6 +251,7 @@ func TestPrepareDataSource_ImageSource(t *testing.T) {
 			MountPath: "/mnt/data",
 		},
 	}
+	expectedInitContainer.VolumeMounts = expectedVolumeMounts
 
 	initContainer, volumes, volumeMounts := prepareDataSource(ctx, workspaceObj)
 
@@ -368,3 +260,430 @@ func TestPrepareDataSource_ImageSource(t *testing.T) {
 	assert.Equal(t, expectedVolumes, volumes)
 	assert.Equal(t, expectedVolumeMounts, volumeMounts)
 }
+
+func TestSetTrainingResultVolume(t *testing.T) {
+	ctx := context.Background()
+
+	testcases := map[string]struct {
+		workspace      *kaitov1beta1.Workspace
+		existingConfig *corev1.ConfigMap
+		expectedError  bool
+		expectedVolume string
+		expectedMount  string
+	}{
+		"LoRA config success": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Method: kaitov1beta1.TuningMethodLora,
+					Config: "lora-config",
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+								ClaimName: "output-pvc",
+							},
+						},
+					},
+				},
+			},
+			existingConfig: &corev1.ConfigMap{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "lora-config",
+					Namespace: "default",
+				},
+				Data: map[string]string{
+					"training_config.yaml": `
+training_config:
+  TrainingArguments:
+    output_dir: "lora-output"
+`,
+				},
+			},
+			expectedError:  false,
+			expectedVolume: "results-volume",
+			expectedMount:  "/mnt/lora-output",
+		},
+		"QLoRA config success": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Method: kaitov1beta1.TuningMethodQLora,
+					Config: "qlora-config",
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			existingConfig: &corev1.ConfigMap{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "qlora-config",
+					Namespace: "default",
+				},
+				Data: map[string]string{
+					"training_config.yaml": `
+training_config:
+  TrainingArguments:
+    output_dir: "qlora-output"
+`,
+				},
+			},
+			expectedError:  false,
+			expectedVolume: "results-volume",
+			expectedMount:  "/mnt/qlora-output",
+		},
+		"Config not found": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Method: kaitov1beta1.TuningMethodLora,
+					Config: "missing-config",
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			existingConfig: nil,
+			expectedError:  true,
+		},
+		"Invalid config data": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Method: kaitov1beta1.TuningMethodLora,
+					Config: "invalid-config",
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			existingConfig: &corev1.ConfigMap{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "invalid-config",
+					Namespace: "default",
+				},
+				Data: map[string]string{
+					"training_config.yaml": "invalid yaml data",
+				},
+			},
+			expectedError: true,
+		},
+		"Default output directory": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Method: kaitov1beta1.TuningMethodLora,
+					Config: "default-output-config",
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			existingConfig: &corev1.ConfigMap{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "default-output-config",
+					Namespace: "default",
+				},
+				Data: map[string]string{
+					"training_config.yaml": `
+training_config:
+  TrainingArguments: {}
+`,
+				},
+			},
+			expectedError:  false,
+			expectedVolume: "results-volume",
+			expectedMount:  DefaultOutputVolumePath,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			// Create a fake client
+			scheme := runtime.NewScheme()
+			_ = corev1.AddToScheme(scheme)
+			_ = kaitov1beta1.AddToScheme(scheme)
+
+			var kubeClient client.Client
+			if tc.existingConfig != nil {
+				kubeClient = fake.NewClientBuilder().
+					WithScheme(scheme).
+					WithObjects(tc.existingConfig).
+					Build()
+			} else {
+				kubeClient = fake.NewClientBuilder().
+					WithScheme(scheme).
+					Build()
+			}
+
+			// Create generator context
+			gctx := &generator.WorkspaceGeneratorContext{
+				Ctx:        ctx,
+				Workspace:  tc.workspace,
+				KubeClient: kubeClient,
+			}
+
+			// Create a pod spec with a container
+			podSpec := &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: tc.workspace.Name,
+					},
+				},
+			}
+
+			// Call the function
+			err := SetTrainingResultVolume(gctx, podSpec)
+
+			// Check error
+			if tc.expectedError {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+
+			// Check volume was added
+			volumeFound := false
+			for _, v := range podSpec.Volumes {
+				if v.Name == tc.expectedVolume {
+					volumeFound = true
+					break
+				}
+			}
+			assert.True(t, volumeFound, "Expected volume %s not found", tc.expectedVolume)
+
+			// Check volume mount was added to the correct container
+			mountFound := false
+			for _, c := range podSpec.Containers {
+				if c.Name == tc.workspace.Name {
+					for _, vm := range c.VolumeMounts {
+						if vm.Name == tc.expectedVolume && vm.MountPath == tc.expectedMount {
+							mountFound = true
+							break
+						}
+					}
+				}
+			}
+			assert.True(t, mountFound, "Expected volume mount not found on container")
+		})
+	}
+}
+
+func TestPrepareTrainingOutput(t *testing.T) {
+	ctx := context.Background()
+
+	testcases := map[string]struct {
+		workspace          *kaitov1beta1.Workspace
+		initialPodSpec     *corev1.PodSpec
+		expectedError      bool
+		expectedContainers int
+		validateFunc       func(*testing.T, *corev1.PodSpec)
+	}{
+		"Output volume specified - early return": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Output: &kaitov1beta1.DataDestination{
+						Volume: &corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			initialPodSpec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-workspace",
+					},
+				},
+			},
+			expectedError:      false,
+			expectedContainers: 1, // No additional containers added
+		},
+		"No output image - early return": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Output: &kaitov1beta1.DataDestination{
+						Image: "",
+					},
+				},
+			},
+			initialPodSpec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-workspace",
+					},
+				},
+			},
+			expectedError:      false,
+			expectedContainers: 1, // No additional containers added
+		},
+		"Image output with preset": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Preset: &kaitov1beta1.PresetSpec{
+						PresetMeta: kaitov1beta1.PresetMeta{
+							Name: "phi-3-mini-128k-instruct",
+						},
+					},
+					Output: &kaitov1beta1.DataDestination{
+						Image:           "registry.example.com/output:latest",
+						ImagePushSecret: "push-secret",
+					},
+				},
+			},
+			initialPodSpec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-workspace",
+						VolumeMounts: []corev1.VolumeMount{
+							{
+								Name:      "results-volume",
+								MountPath: "/mnt/output",
+							},
+						},
+					},
+				},
+				Volumes: []corev1.Volume{
+					{
+						Name: "results-volume",
+						VolumeSource: corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+				},
+			},
+			expectedError:      false,
+			expectedContainers: 3, // Original + pusher + pause
+			validateFunc: func(t *testing.T, spec *corev1.PodSpec) {
+				// Check image pull secrets
+				assert.Equal(t, 1, len(spec.ImagePullSecrets))
+				assert.Equal(t, "push-secret", spec.ImagePullSecrets[0].Name)
+
+				// Check docker config volume
+				volumeFound := false
+				for _, v := range spec.Volumes {
+					if v.Name == "docker-config" {
+						volumeFound = true
+						break
+					}
+				}
+				assert.True(t, volumeFound, "docker-config volume not found")
+
+				// Check pusher container
+				pusherFound := false
+				for _, c := range spec.Containers {
+					if c.Name == "pusher" {
+						pusherFound = true
+						// Verify pusher has the docker config volume mount
+						mountFound := false
+						for _, vm := range c.VolumeMounts {
+							if vm.Name == "docker-config" {
+								mountFound = true
+								break
+							}
+						}
+						assert.True(t, mountFound, "docker-config volume mount not found in pusher container")
+					}
+				}
+				assert.True(t, pusherFound, "pusher container not found")
+
+				// Check pause container
+				pauseFound := false
+				for _, c := range spec.Containers {
+					if c.Name == "pause" {
+						pauseFound = true
+					}
+				}
+				assert.True(t, pauseFound, "pause container not found")
+			},
+		},
+		"Missing results volume mount": {
+			workspace: &kaitov1beta1.Workspace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-workspace",
+					Namespace: "default",
+				},
+				Tuning: &kaitov1beta1.TuningSpec{
+					Output: &kaitov1beta1.DataDestination{
+						Image:           "registry.example.com/output:latest",
+						ImagePushSecret: "push-secret",
+					},
+				},
+			},
+			initialPodSpec: &corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: "test-workspace",
+						// No volume mounts
+					},
+				},
+			},
+			expectedError: true,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			// Create generator context
+			gctx := &generator.WorkspaceGeneratorContext{
+				Ctx:       ctx,
+				Workspace: tc.workspace,
+			}
+
+			// Deep copy the initial pod spec to avoid mutations
+			podSpec := tc.initialPodSpec.DeepCopy()
+
+			// Call the function
+			err := SetTrainingOutputImagePush(gctx, podSpec)
+
+			// Check error
+			if tc.expectedError {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+
+			// Check number of containers
+			assert.Equal(t, tc.expectedContainers, len(podSpec.Containers))
+
+			// Run additional validation if provided
+			if tc.validateFunc != nil {
+				tc.validateFunc(t, podSpec)
+			}
+		})
+	}
+}
diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go
index 5cc529b..eb9d3b9 100644
--- a/test/e2e/preset_test.go
+++ b/test/e2e/preset_test.go
@@ -340,7 +340,8 @@ func updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj *kaitov1beta1.Wo
 			}
 		} else {
 			workspaceObj.Tuning.Input = &kaitov1beta1.DataSource{
-				Image: datasetImageName,
+				Image:            datasetImageName,
+				ImagePullSecrets: []string{e2eACRSecret},
 			}
 		}
 		if outputVolume != nil {