Repository: kaito-project/kaito

Issue:
vLLM sometimes doesn't work when `PYTORCH_CUDA_ALLOC_CONF` is set to `expandable_segments:True`
**Is your feature request related to a problem? Please describe.**

vLLM deployment crashed with the following error message:

![Image](https://github.com/user-attachments/assets/69329358-9efd-45c3-b459-4719f377a92f)

Found this from vLLM repository: https://github.com/vllm-project/vllm/issues/9046.

It works fine for tensor parallelism = 1 but anything >1 would crash vLLM.

**Describe the solution you'd like**

Remove the environment variable from vLLM deployment but not sure what effect it would have.

**Describe alternatives you've considered**

**Additional context**

Discussion:
@zhuangqh: The main idea of adding this env was to reduce memory Fragments.

```yaml
apiVersion: kaito.sh/v1beta1
kind: Workspace
metadata:
  name: workspace-phi-4-mini
resource:
  instanceType: "Standard_NC12s_v3"
  labelSelector:
    matchLabels:
      apps: phi-4
inference:
  preset:
    name: phi-4-mini-instruct 
  config: "inference-params"
---
apiVersion: v1
kind: ConfigMap
metadata:
  name: inference-params
data:
  inference_config.yaml: |
    # Maximum number of steps to find the max available seq len fitting in the GPU memory.
    max_probe_steps: 6

    vllm:
      cpu-offload-gb: 0
      gpu-memory-utilization: 0.95
      swap-space: 4
      max-model-len: 2048
```

I can successfully run this workspace with 2 tensor parallelism. In fact, we have a e2e test to cover this case.
Can you confirm how to reproduce this issue? do you run it on multi nodes?@chewong: I am using one pre-provisioned `Standard_NC48ads_A100_v4` (2 x A100) instance in my setup. The rest of the setup is the same as yours. Let me give it a try again.@chewong: Tested with 1 `Standard_NC12s_v3` + TP = 2 + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` and it works with no problem. I am starting to suspect that only certain GPUs are affected.@chewong: ![Image](https://github.com/user-attachments/assets/154ec801-03e8-4c9e-b37b-439b2f543805)

I got the same error when running TP=2 with two H100s on a node, which confirms my previous observation that this issue only affects certain GPU models. Current affected models are A100s and H100s.

Edit: again, it works after removing the environment variable/setting the value to False:

![Image](https://github.com/user-attachments/assets/0e1e55c4-4509-42d1-af55-394a25924ee9)

![Image](https://github.com/user-attachments/assets/e35cd1da-2bd8-4af6-9466-6f0ca37ad460)



Solution patch:
diff --git a/.github/workflows/trivy.yml b/.github/workflows/trivy.yml
new file mode 100644
index 0000000..2150dc7
--- /dev/null
+++ b/.github/workflows/trivy.yml
@@ -0,0 +1,40 @@
+name: Trivy vulnerability scanner
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+jobs:
+  build:
+    name: Build
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Go 1.x
+        uses: actions/setup-go@v5
+        with:
+          go-version: 1.24.2
+        id: go
+
+      - name: Checkout code
+        uses: actions/checkout@v4
+      
+      - name: Build workspace image from Dockerfile
+        run: |
+          export REGISTRY=test
+          export VERSION=latest
+          export DOCKER_CLI_EXPERIMENTAL=enabled
+          export OUTPUT_TYPE="type=docker"
+          export ARCH="amd64"
+          make docker-build-workspace
+      
+      - name: Run Trivy vulnerability scanner
+        uses: aquasecurity/trivy-action@master
+        env:
+          TRIVY_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-db:2"
+        with:
+          image-ref: 'test/workspace:latest'
+          format: 'table'
+          exit-code: '1'
+          ignore-unfixed: true
+          vuln-type: 'os,library'
+          severity: 'CRITICAL,HIGH,MEDIUM,LOW,UNKNOWN'
diff --git a/api/v1beta1/workspace_types.go b/api/v1beta1/workspace_types.go
index 32e3a93..8960977 100644
--- a/api/v1beta1/workspace_types.go
+++ b/api/v1beta1/workspace_types.go
@@ -62,6 +62,9 @@ type PresetOptions struct {
 	// ImagePullSecrets is a list of secret names in the same namespace used for pulling the model image.
 	// +optional
 	ImagePullSecrets []string `json:"imagePullSecrets,omitempty"`
+	// ModelAccessSecret is the name of the secret that contains the huggingface access token.
+	// +optional
+	ModelAccessSecret string `json:"modelAccessSecret,omitempty"`
 }
 
 // PresetSpec provides the information for rendering preset configurations to run the model inference service.
diff --git a/api/v1beta1/workspace_validation.go b/api/v1beta1/workspace_validation.go
index a2b9e3f..23f78c2 100644
--- a/api/v1beta1/workspace_validation.go
+++ b/api/v1beta1/workspace_validation.go
@@ -465,6 +465,13 @@ func (i *InferenceSpec) validateCreate(ctx context.Context, namespace string, in
 			errs = errs.Also(apis.ErrGeneric("When AccessMode is private, an image must be provided in PresetOptions"))
 		}
 		// Note: we don't enforce private access mode to have image secrets, in case anonymous pulling is enabled
+
+		// For models that require downloading at runtime, we need to check if the modelAccessSecret is provided
+		if params.DownloadAtRuntime && i.Preset.PresetOptions.ModelAccessSecret == "" {
+			errs = errs.Also(apis.ErrGeneric("This preset requires a modelAccessSecret with HF_TOKEN key under presetOptions to download the model"))
+		} else if !params.DownloadAtRuntime && i.Preset.PresetOptions.ModelAccessSecret != "" {
+			errs = errs.Also(apis.ErrGeneric("This preset does not require a modelAccessSecret with HF_TOKEN key under presetOptions"))
+		}
 	}
 	if len(i.Adapters) > MaxAdaptersNumber {
 		errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("Number of Adapters exceeds the maximum limit, maximum of %s allowed", strconv.Itoa(MaxAdaptersNumber))))
diff --git a/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml b/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
index 973e178..86fcebe 100644
--- a/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
+++ b/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
@@ -512,6 +512,10 @@ spec:
                         items:
                           type: string
                         type: array
+                      modelAccessSecret:
+                        description: ModelAccessSecret is the name of the secret that
+                          contains the huggingface access token.
+                        type: string
                     type: object
                 required:
                 - name
@@ -774,6 +778,10 @@ spec:
                         items:
                           type: string
                         type: array
+                      modelAccessSecret:
+                        description: ModelAccessSecret is the name of the secret that
+                          contains the huggingface access token.
+                        type: string
                     type: object
                 required:
                 - name
diff --git a/config/crd/bases/kaito.sh_workspaces.yaml b/config/crd/bases/kaito.sh_workspaces.yaml
index 973e178..86fcebe 100644
--- a/config/crd/bases/kaito.sh_workspaces.yaml
+++ b/config/crd/bases/kaito.sh_workspaces.yaml
@@ -512,6 +512,10 @@ spec:
                         items:
                           type: string
                         type: array
+                      modelAccessSecret:
+                        description: ModelAccessSecret is the name of the secret that
+                          contains the huggingface access token.
+                        type: string
                     type: object
                 required:
                 - name
@@ -774,6 +778,10 @@ spec:
                         items:
                           type: string
                         type: array
+                      modelAccessSecret:
+                        description: ModelAccessSecret is the name of the secret that
+                          contains the huggingface access token.
+                        type: string
                     type: object
                 required:
                 - name
diff --git a/pkg/model/interface.go b/pkg/model/interface.go
index a948b9a..d418448 100644
--- a/pkg/model/interface.go
+++ b/pkg/model/interface.go
@@ -78,6 +78,18 @@ type Metadata struct {
 	DownloadAtRuntime bool `yaml:"downloadAtRuntime,omitempty"`
 }
 
+// Validate checks if the Metadata is valid.
+func (m *Metadata) Validate() error {
+	// Some private models may not have a version URL, so we allow it to be empty until
+	// we remove support for private preset models.
+	if m.Version == "" {
+		return nil
+	}
+
+	_, _, err := utils.ParseHuggingFaceModelVersion(m.Version)
+	return err
+}
+
 // PresetParam defines the preset inference parameters for a model.
 type PresetParam struct {
 	Metadata
@@ -197,6 +209,13 @@ func (p *PresetParam) GetInferenceCommand(rc RuntimeContext) []string {
 }
 
 func (p *PresetParam) buildHuggingfaceInferenceCommand() []string {
+	if p.DownloadAtRuntime {
+		repoId, revision, _ := utils.ParseHuggingFaceModelVersion(p.Version)
+		p.Transformers.ModelRunParams["pretrained_model_name_or_path"] = repoId
+		if revision != "" {
+			p.Transformers.ModelRunParams["revision"] = revision
+		}
+	}
 	torchCommand := utils.BuildCmdStr(
 		p.Transformers.BaseCommand,
 		p.Transformers.TorchRunParams,
@@ -219,6 +238,13 @@ func (p *PresetParam) buildVLLMInferenceCommand(rc RuntimeContext) []string {
 	if !p.VLLM.DisallowLoRA && rc.AdaptersEnabled {
 		p.VLLM.ModelRunParams["enable-lora"] = ""
 	}
+	if p.DownloadAtRuntime {
+		repoId, revision, _ := utils.ParseHuggingFaceModelVersion(p.Version)
+		p.VLLM.ModelRunParams["model"] = repoId
+		if revision != "" {
+			p.VLLM.ModelRunParams["code-revision"] = revision
+		}
+	}
 	gpuMemUtil := getGPUMemoryUtilForVLLM(rc.GPUConfig)
 	p.VLLM.ModelRunParams["gpu-memory-utilization"] = strconv.FormatFloat(gpuMemUtil, 'f', 2, 64)
 	if rc.ConfigVolume != nil {
diff --git a/pkg/utils/common.go b/pkg/utils/common.go
index 881752f..b1abec8 100644
--- a/pkg/utils/common.go
+++ b/pkg/utils/common.go
@@ -6,6 +6,7 @@ package utils
 import (
 	"context"
 	"fmt"
+	"net/url"
 	"os"
 	"sort"
 	"strings"
@@ -27,6 +28,10 @@ import (
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 )
 
+const (
+	errInvalidModelVersionURL = "invalid model version URL: %s. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>"
+)
+
 var (
 	karpenterSchemeGroupVersion = schema.GroupVersion{Group: karpenterapis.Group, Version: "v1"}
 	awsSchemeGroupVersion       = schema.GroupVersion{Group: awsapis.Group, Version: "v1beta1"}
@@ -268,3 +273,44 @@ func SelectNodes(qualified []*v1.Node, preferred []string, previous []string, co
 
 	return qualified[0:count]
 }
+
+// ParseHuggingFaceModelVersion parses the model version in the format of https://huggingface.co/<org>/<model>/commit/<revision>
+// and returns the repoId and revision. If the commit is not specified, it returns an empty string for revision,
+// and the main branch HEAD commit is used.
+//
+// Example 1:
+//
+//	Version: "https://huggingface.co/tiiuae/falcon-7b/commit/ec89142b67d748a1865ea4451372db8313ada0d8"
+//	RepoId: "tiiuae/falcon-7b"
+//	Revision: "ec89142b67d748a1865ea4451372db8313ada0d8"
+//
+// Example 2:
+//
+//	Version: https://huggingface.co/tiiuae/falcon-7b
+//	RepoId: "tiiuae/falcon-7b"
+//	Revision: "" (main branch HEAD commit is used)
+func ParseHuggingFaceModelVersion(version string) (repoId string, revision string, err error) {
+	parsedURL, err := url.Parse(version)
+	if err != nil {
+		return "", "", err
+	}
+
+	if parsedURL.Host != "huggingface.co" {
+		return "", "", fmt.Errorf(errInvalidModelVersionURL, version)
+	}
+
+	parts := strings.Split(strings.Trim(parsedURL.Path, "/"), "/")
+	switch len(parts) {
+	case 2: // Expected path: "<org>/<model>"
+		repoId, revision = parts[0]+"/"+parts[1], ""
+		return
+	case 4: // Expected path: "<org>/<model>/commit/<revision>"
+		if parts[2] != "commit" {
+			break
+		}
+		repoId, revision = parts[0]+"/"+parts[1], parts[3]
+		return
+	}
+
+	return "", "", fmt.Errorf(errInvalidModelVersionURL, version)
+}
diff --git a/pkg/workspace/inference/preset-inferences.go b/pkg/workspace/inference/preset-inferences.go
index 6c311f4..bb18e3c 100644
--- a/pkg/workspace/inference/preset-inferences.go
+++ b/pkg/workspace/inference/preset-inferences.go
@@ -21,6 +21,7 @@ import (
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/kaito-project/kaito/pkg/workspace/manifests"
+	metadata "github.com/kaito-project/kaito/presets/workspace/models"
 )
 
 const (
@@ -113,20 +114,33 @@ func GetInferenceImageInfo(ctx context.Context, workspaceObj *v1beta1.Workspace,
 			}
 		}
 	}
+
+	// Three possible cases for inference workload image selection:
+	// 1. If the preset's access mode is private, use the user-specified image.
+	// 2. If the preset is set to download at runtime, use the 'kaito-base' image.
+	// 3. Otherwise, use the preset image, which has the model weights packaged in.
+	var imageName, imageTag string
 	if string(workspaceObj.Inference.Preset.AccessMode) == string(v1beta1.ModelImageAccessModePrivate) {
-		imageName := workspaceObj.Inference.Preset.PresetOptions.Image
+		imageName = workspaceObj.Inference.Preset.PresetOptions.Image
 		for _, secretName := range workspaceObj.Inference.Preset.PresetOptions.ImagePullSecrets {
 			imagePullSecretRefs = append(imagePullSecretRefs, corev1.LocalObjectReference{Name: secretName})
 		}
 		return imageName, imagePullSecretRefs
+	} else if presetObj.DownloadAtRuntime {
+		// Force the use of kaito-base image if the preset is set to download at runtime.
+		// The kaito-base image is the same as other preset images but without the model
+		// files packaged in.
+		imageName = "base"
+		imageTag = metadata.MustGet(imageName).Tag
 	} else {
-		imageName := string(workspaceObj.Inference.Preset.Name)
-		imageTag := presetObj.Tag
-		registryName := os.Getenv("PRESET_REGISTRY_NAME")
-		imageName = fmt.Sprintf("%s/kaito-%s:%s", registryName, imageName, imageTag)
-
-		return imageName, imagePullSecretRefs
+		imageName = string(workspaceObj.Inference.Preset.Name)
+		imageTag = presetObj.Tag
 	}
+
+	registryName := os.Getenv("PRESET_REGISTRY_NAME")
+	imageName = fmt.Sprintf("%s/kaito-%s:%s", registryName, imageName, imageTag)
+
+	return imageName, imagePullSecretRefs
 }
 
 func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace, revisionNum string,
@@ -178,6 +192,7 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 	// additional volume
 	var volumes []corev1.Volume
 	var volumeMounts []corev1.VolumeMount
+	var envVars []corev1.EnvVar
 
 	// Add config volume mount
 	cmVolume, cmVolumeMount := utils.ConfigCMVolume(configVolume.Name)
@@ -197,6 +212,19 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 		volumes = append(volumes, adapterVolume)
 		volumeMounts = append(volumeMounts, adapterVolumeMount)
 	}
+	if inferenceParam.DownloadAtRuntime {
+		envVars = append(envVars, corev1.EnvVar{
+			Name: "HF_TOKEN",
+			ValueFrom: &corev1.EnvVarSource{
+				SecretKeyRef: &corev1.SecretKeySelector{
+					LocalObjectReference: corev1.LocalObjectReference{
+						Name: workspaceObj.Inference.Preset.PresetOptions.ModelAccessSecret,
+					},
+					Key: "HF_TOKEN",
+				},
+			},
+		})
+	}
 
 	// inference command
 	runtimeName := v1beta1.GetWorkspaceRuntimeName(workspaceObj)
@@ -215,10 +243,10 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 	var depObj client.Object
 	if model.SupportDistributedInference() {
 		depObj = manifests.GenerateStatefulSetManifest(workspaceObj, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
-			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts)
+			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts, envVars)
 	} else {
 		depObj = manifests.GenerateDeploymentManifest(workspaceObj, revisionNum, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
-			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts)
+			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts, envVars)
 	}
 	err = resources.CreateResource(ctx, depObj, kubeClient)
 	if client.IgnoreAlreadyExists(err) != nil {
diff --git a/pkg/workspace/manifests/manifests.go b/pkg/workspace/manifests/manifests.go
index 24cbf5c..b9d0685 100644
--- a/pkg/workspace/manifests/manifests.go
+++ b/pkg/workspace/manifests/manifests.go
@@ -97,7 +97,7 @@ func GenerateServiceManifest(workspaceObj *kaitov1beta1.Workspace, serviceType c
 func GenerateStatefulSetManifest(workspaceObj *kaitov1beta1.Workspace, imageName string,
 	imagePullSecretRefs []corev1.LocalObjectReference, replicas int, commands []string, containerPorts []corev1.ContainerPort,
 	livenessProbe, readinessProbe *corev1.Probe, resourceRequirements corev1.ResourceRequirements,
-	tolerations []corev1.Toleration, volumes []corev1.Volume, volumeMount []corev1.VolumeMount) *appsv1.StatefulSet {
+	tolerations []corev1.Toleration, volumes []corev1.Volume, volumeMount []corev1.VolumeMount, envVars []corev1.EnvVar) *appsv1.StatefulSet {
 
 	nodeRequirements := make([]corev1.NodeSelectorRequirement, 0, len(workspaceObj.Resource.LabelSelector.MatchLabels))
 	for key, value := range workspaceObj.Resource.LabelSelector.MatchLabels {
@@ -115,12 +115,10 @@ func GenerateStatefulSetManifest(workspaceObj *kaitov1beta1.Workspace, imageName
 		MatchLabels: selector,
 	}
 	// Add PYTORCH_CUDA_ALLOC_CONF environment variable
-	envVars := []corev1.EnvVar{
-		{
-			Name:  "PYTORCH_CUDA_ALLOC_CONF",
-			Value: "expandable_segments:True",
-		},
-	}
+	envVars = append(envVars, corev1.EnvVar{
+		Name:  "PYTORCH_CUDA_ALLOC_CONF",
+		Value: "expandable_segments:True",
+	})
 
 	ss := &appsv1.StatefulSet{
 		ObjectMeta: v1.ObjectMeta{
@@ -247,7 +245,7 @@ func GenerateTuningJobManifest(wObj *kaitov1beta1.Workspace, revisionNum string,
 func GenerateDeploymentManifest(workspaceObj *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, volumes []corev1.Volume, volumeMount []corev1.VolumeMount) *appsv1.Deployment {
+	tolerations []corev1.Toleration, volumes []corev1.Volume, volumeMount []corev1.VolumeMount, envVars []corev1.EnvVar) *appsv1.Deployment {
 
 	nodeRequirements := make([]corev1.NodeSelectorRequirement, 0, len(workspaceObj.Resource.LabelSelector.MatchLabels))
 	for key, value := range workspaceObj.Resource.LabelSelector.MatchLabels {
@@ -264,15 +262,14 @@ func GenerateDeploymentManifest(workspaceObj *kaitov1beta1.Workspace, revisionNu
 	labelselector := &v1.LabelSelector{
 		MatchLabels: selector,
 	}
-	envs := []corev1.EnvVar{}
 	// Add PYTORCH_CUDA_ALLOC_CONF environment variable
-	envs = append(envs, corev1.EnvVar{
+	envVars = append(envVars, corev1.EnvVar{
 		Name:  "PYTORCH_CUDA_ALLOC_CONF",
 		Value: "expandable_segments:True",
 	})
 
 	pullerContainers, pullerEnvVars, pullerVolumes := GeneratePullerContainers(workspaceObj, volumeMount)
-	envs = append(envs, pullerEnvVars...)
+	envVars = append(envVars, pullerEnvVars...)
 	volumes = append(volumes, pullerVolumes...)
 
 	return &appsv1.Deployment{
@@ -330,7 +327,7 @@ func GenerateDeploymentManifest(workspaceObj *kaitov1beta1.Workspace, revisionNu
 							ReadinessProbe: readinessProbe,
 							Ports:          containerPorts,
 							VolumeMounts:   volumeMount,
-							Env:            envs,
+							Env:            envVars,
 						},
 					},
 					Tolerations: tolerations,
diff --git a/presets/workspace/models/metadata.go b/presets/workspace/models/metadata.go
index a929a8e..5084edf 100644
--- a/presets/workspace/models/metadata.go
+++ b/presets/workspace/models/metadata.go
@@ -35,6 +35,7 @@ func init() {
 	utilruntime.Must(yaml.Unmarshal(supportedModelsYAML, &catalog))
 
 	for _, m := range catalog.Models {
+		utilruntime.Must(m.Validate())
 		supportedModels.Store(m.Name, &m)
 	}
 }



Test patch:
diff --git a/api/v1beta1/workspace_validation_test.go b/api/v1beta1/workspace_validation_test.go
index 40ab201..6fc011e 100644
--- a/api/v1beta1/workspace_validation_test.go
+++ b/api/v1beta1/workspace_validation_test.go
@@ -110,10 +110,38 @@ func (*testModelPrivate) SupportTuning() bool {
 	return true
 }
 
+type testModelDownload struct{}
+
+func (*testModelDownload) GetInferenceParameters() *model.PresetParam {
+	return &model.PresetParam{
+		Metadata: model.Metadata{
+			Version:           "https://huggingface.co/test-repo/test-model/commit/test-revision",
+			DownloadAtRuntime: true,
+		},
+		GPUCountRequirement:       gpuCountRequirement,
+		TotalGPUMemoryRequirement: totalGPUMemoryRequirement,
+		PerGPUMemoryRequirement:   perGPUMemoryRequirement,
+	}
+}
+func (*testModelDownload) GetTuningParameters() *model.PresetParam {
+	return &model.PresetParam{
+		GPUCountRequirement:       gpuCountRequirement,
+		TotalGPUMemoryRequirement: totalGPUMemoryRequirement,
+		PerGPUMemoryRequirement:   perGPUMemoryRequirement,
+	}
+}
+func (*testModelDownload) SupportDistributedInference() bool {
+	return false
+}
+func (*testModelDownload) SupportTuning() bool {
+	return false
+}
+
 func RegisterValidationTestModels() {
 	var test testModel
 	var testPrivate testModelPrivate
 	var testStatic testModelStatic
+	var testDownload testModelDownload
 	plugin.KaitoModelRegister.Register(&plugin.Registration{
 		Name:     "test-validation",
 		Instance: &test,
@@ -126,6 +154,10 @@ func RegisterValidationTestModels() {
 		Name:     "test-validation-static",
 		Instance: &testStatic,
 	})
+	plugin.KaitoModelRegister.Register(&plugin.Registration{
+		Name:     "test-validation-download",
+		Instance: &testDownload,
+	})
 }
 
 func pointerToInt(i int) *int {
@@ -837,6 +869,46 @@ func TestInferenceSpecValidateCreate(t *testing.T) {
 			},
 			runtimeName: model.RuntimeNameHuggingfaceTransformers,
 		},
+		{
+			name: "download model at runtime with access secret",
+			inferenceSpec: &InferenceSpec{
+				Preset: &PresetSpec{
+					PresetMeta: PresetMeta{
+						Name: ModelName("test-validation-download"),
+					},
+					PresetOptions: PresetOptions{
+						ModelAccessSecret: "test-secret",
+					},
+				},
+			},
+		},
+		{
+			name: "download model at runtime but no access secret",
+			inferenceSpec: &InferenceSpec{
+				Preset: &PresetSpec{
+					PresetMeta: PresetMeta{
+						Name: ModelName("test-validation-download"),
+					},
+				},
+			},
+			errContent: "This preset requires a modelAccessSecret with HF_TOKEN key under presetOptions to download the model",
+			expectErrs: true,
+		},
+		{
+			name: "Preset with model weights packaged but with access secret",
+			inferenceSpec: &InferenceSpec{
+				Preset: &PresetSpec{
+					PresetMeta: PresetMeta{
+						Name: ModelName("test-validation"),
+					},
+					PresetOptions: PresetOptions{
+						ModelAccessSecret: "test-secret",
+					},
+				},
+			},
+			errContent: "This preset does not require a modelAccessSecret with HF_TOKEN key under presetOptions",
+			expectErrs: true,
+		},
 	}
 
 	for _, tc := range tests {
diff --git a/pkg/utils/common_test.go b/pkg/utils/common_test.go
index 29a6b05..178c325 100644
--- a/pkg/utils/common_test.go
+++ b/pkg/utils/common_test.go
@@ -5,6 +5,7 @@ package utils
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -120,3 +121,113 @@ func TestFetchGPUCountFromNodes(t *testing.T) {
 		})
 	}
 }
+
+func TestParseHuggingFaceModelVersion(t *testing.T) {
+	tests := []struct {
+		name             string
+		version          string
+		expectedRepoId   string
+		expectedRevision string
+		expectErr        bool
+		expectedErrMsg   string // Use Contains for url.Parse errors, Exact for custom errors
+	}{
+		{
+			name:             "Valid URL with commit revision",
+			version:          "https://huggingface.co/tiiuae/falcon-7b/commit/ec89142b67d748a1865ea4451372db8313ada0d8",
+			expectedRepoId:   "tiiuae/falcon-7b",
+			expectedRevision: "ec89142b67d748a1865ea4451372db8313ada0d8",
+			expectErr:        false,
+		},
+		{
+			name:             "Valid URL without commit revision",
+			version:          "https://huggingface.co/tiiuae/falcon-7b",
+			expectedRepoId:   "tiiuae/falcon-7b",
+			expectedRevision: "",
+			expectErr:        false,
+		},
+		{
+			name:           "Invalid URL path structure - too few parts",
+			version:        "https://huggingface.co/tiiuae",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://huggingface.co/tiiuae. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>",
+		},
+		{
+			name:           "Invalid URL path structure - incorrect middle part",
+			version:        "https://huggingface.co/tiiuae/falcon-7b/blob/main",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://huggingface.co/tiiuae/falcon-7b/blob/main. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>",
+		},
+		{
+			name:           "Invalid URL path structure - too many parts",
+			version:        "https://huggingface.co/tiiuae/falcon-7b/commit/rev/extra",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://huggingface.co/tiiuae/falcon-7b/commit/rev/extra. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>",
+		},
+		{
+			name:           "Invalid URL format - parsing error",
+			version:        "://invalid-url",
+			expectErr:      true,
+			expectedErrMsg: "missing protocol scheme", // Contains check
+		},
+		{
+			name:             "Valid URL with trailing slash",
+			version:          "https://huggingface.co/org/model/",
+			expectedRepoId:   "org/model",
+			expectedRevision: "",
+			expectErr:        false,
+		},
+		{
+			name:             "Valid URL with commit and trailing slash",
+			version:          "https://huggingface.co/org/model/commit/revision123/",
+			expectedRepoId:   "org/model",
+			expectedRevision: "revision123",
+			expectErr:        false,
+		},
+		{
+			name:           "Invalid host",
+			version:        "https://github.com/org/model",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://github.com/org/model. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>",
+		},
+		{
+			name:           "Empty input string",
+			version:        "",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: . Expected format: https://huggingface.co/<org>/<model>/commit/<revision>", // url.Parse("") returns empty URL, host check fails
+		},
+		{
+			name:           "URL with only host",
+			version:        "https://huggingface.co",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://huggingface.co. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>", // Path is empty, len(parts) is 0 or 1 depending on trailing slash
+		},
+		{
+			name:           "URL with path /org/model/tree/branch",
+			version:        "https://huggingface.co/org/model/tree/branch",
+			expectErr:      true,
+			expectedErrMsg: "invalid model version URL: https://huggingface.co/org/model/tree/branch. Expected format: https://huggingface.co/<org>/<model>/commit/<revision>", // len(parts) is 4, but parts[2] is not "commit"
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			repoId, revision, err := ParseHuggingFaceModelVersion(tt.version)
+
+			if tt.expectErr {
+				require.Error(t, err)
+				if tt.expectedErrMsg != "" {
+					// Use Contains for url.Parse errors which might vary slightly, exact match for our custom errors
+					if strings.Contains(tt.name, "parsing error") || strings.Contains(tt.name, "Empty input string") {
+						assert.Contains(t, err.Error(), tt.expectedErrMsg)
+					} else {
+						assert.Equal(t, tt.expectedErrMsg, err.Error())
+					}
+				}
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.expectedRepoId, repoId)
+				assert.Equal(t, tt.expectedRevision, revision)
+			}
+		})
+	}
+}
diff --git a/pkg/utils/test/testModel.go b/pkg/utils/test/testModel.go
index 5eb8d9e..68369d2 100644
--- a/pkg/utils/test/testModel.go
+++ b/pkg/utils/test/testModel.go
@@ -16,6 +16,9 @@ var emptyParams = map[string]string{}
 
 func (*baseTestModel) GetInferenceParameters() *model.PresetParam {
 	return &model.PresetParam{
+		Metadata: model.Metadata{
+			Tag: "base-test-model",
+		},
 		GPUCountRequirement: "1",
 		RuntimeParam: model.RuntimeParam{
 			VLLM: model.VLLMParam{
@@ -64,6 +67,9 @@ type testNoTensorParallelModel struct {
 
 func (*testNoTensorParallelModel) GetInferenceParameters() *model.PresetParam {
 	return &model.PresetParam{
+		Metadata: model.Metadata{
+			Tag: "test-no-tensor-parallel-model",
+		},
 		GPUCountRequirement: "1",
 		RuntimeParam: model.RuntimeParam{
 			DisableTensorParallelism: true,
@@ -87,8 +93,41 @@ type testNoLoraSupportModel struct {
 	baseTestModel
 }
 
+type testModelDownload struct {
+	baseTestModel
+}
+
+func (*testModelDownload) SupportDistributedInference() bool {
+	return false
+}
+
+func (*testModelDownload) GetInferenceParameters() *model.PresetParam {
+	return &model.PresetParam{
+		Metadata: model.Metadata{
+			Version:           "https://huggingface.co/test-repo/test-model/commit/test-revision",
+			DownloadAtRuntime: true,
+		},
+		GPUCountRequirement: "1",
+		RuntimeParam: model.RuntimeParam{
+			VLLM: model.VLLMParam{
+				BaseCommand:    "python3 /workspace/vllm/inference_api.py",
+				ModelRunParams: emptyParams,
+			},
+			Transformers: model.HuggingfaceTransformersParam{
+				BaseCommand:       "accelerate launch",
+				InferenceMainFile: "/workspace/tfs/inference_api.py",
+				ModelRunParams:    emptyParams,
+			},
+		},
+		ReadinessTimeout: time.Duration(30) * time.Minute,
+	}
+}
+
 func (*testNoLoraSupportModel) GetInferenceParameters() *model.PresetParam {
 	return &model.PresetParam{
+		Metadata: model.Metadata{
+			Tag: "test-no-lora-support-model",
+		},
 		GPUCountRequirement: "1",
 		RuntimeParam: model.RuntimeParam{
 			DisableTensorParallelism: true,
@@ -129,4 +168,9 @@ func RegisterTestModel() {
 		Name:     "test-no-lora-support-model",
 		Instance: &testNoLoraSupportModel{},
 	})
+
+	plugin.KaitoModelRegister.Register(&plugin.Registration{
+		Name:     "test-model-download",
+		Instance: &testModelDownload{},
+	})
 }
diff --git a/pkg/utils/test/testUtils.go b/pkg/utils/test/testUtils.go
index 2b2ac3d..59bace9 100644
--- a/pkg/utils/test/testUtils.go
+++ b/pkg/utils/test/testUtils.go
@@ -197,6 +197,59 @@ var (
 			},
 		},
 	}
+	MockWorkspaceWithPresetDownloadVLLM = &v1beta1.Workspace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "testWorkspace",
+			Namespace: "kaito",
+		},
+		Resource: v1beta1.ResourceSpec{
+			Count:        &gpuNodeCount,
+			InstanceType: "Standard_NC12s_v3",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"apps": "test",
+				},
+			},
+		},
+		Inference: &v1beta1.InferenceSpec{
+			Preset: &v1beta1.PresetSpec{
+				PresetMeta: v1beta1.PresetMeta{
+					Name: "test-model-download",
+				},
+				PresetOptions: v1beta1.PresetOptions{
+					ModelAccessSecret: "test-secret",
+				},
+			},
+		},
+	}
+	MockWorkspaceWithPresetDownloadTransformers = &v1beta1.Workspace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "testWorkspace",
+			Namespace: "kaito",
+			Annotations: map[string]string{
+				v1beta1.AnnotationWorkspaceRuntime: string(model.RuntimeNameHuggingfaceTransformers),
+			},
+		},
+		Resource: v1beta1.ResourceSpec{
+			Count:        &gpuNodeCount,
+			InstanceType: "Standard_NC12s_v3",
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"apps": "test",
+				},
+			},
+		},
+		Inference: &v1beta1.InferenceSpec{
+			Preset: &v1beta1.PresetSpec{
+				PresetMeta: v1beta1.PresetMeta{
+					Name: "test-model-download",
+				},
+				PresetOptions: v1beta1.PresetOptions{
+					ModelAccessSecret: "test-secret",
+				},
+			},
+		},
+	}
 )
 
 var MockWorkspaceWithPresetHash = "89ae127050ec264a5ce84db48ef7226574cdf1299e6bd27fe90b927e34cc8adb"
diff --git a/pkg/workspace/inference/preset-inferences_test.go b/pkg/workspace/inference/preset-inferences_test.go
index b5a6774..a3583a2 100644
--- a/pkg/workspace/inference/preset-inferences_test.go
+++ b/pkg/workspace/inference/preset-inferences_test.go
@@ -25,14 +25,16 @@ var ValidStrength string = "0.5"
 func TestCreatePresetInference(t *testing.T) {
 	test.RegisterTestModel()
 	testcases := map[string]struct {
-		workspace      *v1beta1.Workspace
-		nodeCount      int
-		modelName      string
-		callMocks      func(c *test.MockClient)
-		workload       string
-		expectedCmd    string
-		hasAdapters    bool
-		expectedVolume string
+		workspace       *v1beta1.Workspace
+		nodeCount       int
+		modelName       string
+		callMocks       func(c *test.MockClient)
+		workload        string
+		expectedCmd     string
+		hasAdapters     bool
+		expectedImage   string
+		expectedVolume  string
+		expectedEnvVars []corev1.EnvVar
 	}{
 
 		"test-model/vllm": {
@@ -43,11 +45,16 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
-			workload: "Deployment",
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-test-model:base-test-model",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --served-model-name=mymodel --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 
 		"test-model-no-parallel/vllm": {
@@ -58,11 +65,16 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
-			workload: "Deployment",
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-test-model:test-no-tensor-parallel-model",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 
 		"test-model-no-lora-support/vllm": {
@@ -73,11 +85,16 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
-			workload: "Deployment",
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-test-model:test-no-lora-support-model",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 
 		"test-model-with-adapters/vllm": {
@@ -89,9 +106,17 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
 			workload:       "Deployment",
+			expectedImage:  "test-registry/kaito-test-model:base-test-model",
 			expectedCmd:    "/bin/sh -c python3 /workspace/vllm/inference_api.py --enable-lora --tensor-parallel-size=2 --served-model-name=mymodel --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters:    true,
 			expectedVolume: "adapter-volume",
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}, {
+				Name:  "Adapter-1",
+				Value: "0.5",
+			}},
 		},
 
 		"test-model/transformers": {
@@ -102,11 +127,16 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
-			workload: "Deployment",
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-test-model:base-test-model",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
 			expectedCmd: "/bin/sh -c accelerate launch /workspace/tfs/inference_api.py",
 			hasAdapters: false,
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 
 		"test-distributed-model/transformers": {
@@ -118,9 +148,14 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.Service{}), mock.Anything).Return(nil)
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.StatefulSet{}), mock.Anything).Return(nil)
 			},
-			workload:    "StatefulSet",
-			expectedCmd: "/bin/sh -c accelerate launch --nnodes=1 --nproc_per_node=0 --max_restarts=3 --rdzv_id=job --rdzv_backend=c10d --rdzv_endpoint=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local:29500 /workspace/tfs/inference_api.py",
-			hasAdapters: false,
+			workload:      "StatefulSet",
+			expectedImage: "test-registry/kaito-test-distributed-model:base-test-model",
+			expectedCmd:   "/bin/sh -c accelerate launch --nnodes=1 --nproc_per_node=0 --max_restarts=3 --rdzv_id=job --rdzv_backend=c10d --rdzv_endpoint=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local:29500 /workspace/tfs/inference_api.py",
+			hasAdapters:   false,
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 
 		"test-model-with-adapters": {
@@ -132,15 +167,78 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
 			workload:       "Deployment",
+			expectedImage:  "test-registry/kaito-test-model:base-test-model",
 			expectedCmd:    "/bin/sh -c accelerate launch /workspace/tfs/inference_api.py",
 			hasAdapters:    true,
 			expectedVolume: "adapter-volume",
+			expectedEnvVars: []corev1.EnvVar{{
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}, {
+				Name:  "Adapter-1",
+				Value: "0.5",
+			}},
+		},
+
+		"test-model-download/vllm": {
+			workspace: test.MockWorkspaceWithPresetDownloadVLLM,
+			nodeCount: 1,
+			modelName: "test-model-download",
+			callMocks: func(c *test.MockClient) {
+				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
+				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
+			},
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-base:0.0.1",
+			expectedCmd:   "/bin/sh -c python3 /workspace/vllm/inference_api.py --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml --model=test-repo/test-model --code-revision=test-revision --tensor-parallel-size=2",
+			expectedEnvVars: []corev1.EnvVar{{
+				Name: "HF_TOKEN",
+				ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						Key: "HF_TOKEN",
+						LocalObjectReference: corev1.LocalObjectReference{
+							Name: "test-secret",
+						},
+					},
+				},
+			}, {
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
+		},
+
+		"test-model-download/transformers": {
+			workspace: test.MockWorkspaceWithPresetDownloadTransformers,
+			nodeCount: 1,
+			modelName: "test-model-download",
+			callMocks: func(c *test.MockClient) {
+				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
+				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
+			},
+			workload:      "Deployment",
+			expectedImage: "test-registry/kaito-base:0.0.1",
+			expectedCmd:   "/bin/sh -c accelerate launch /workspace/tfs/inference_api.py --pretrained_model_name_or_path=test-repo/test-model --revision=test-revision",
+			expectedEnvVars: []corev1.EnvVar{{
+				Name: "HF_TOKEN",
+				ValueFrom: &corev1.EnvVarSource{
+					SecretKeyRef: &corev1.SecretKeySelector{
+						Key: "HF_TOKEN",
+						LocalObjectReference: corev1.LocalObjectReference{
+							Name: "test-secret",
+						},
+					},
+				},
+			}, {
+				Name:  "PYTORCH_CUDA_ALLOC_CONF",
+				Value: "expandable_segments:True",
+			}},
 		},
 	}
 
 	for k, tc := range testcases {
 		t.Run(k, func(t *testing.T) {
 			t.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
+			t.Setenv("PRESET_REGISTRY_NAME", "test-registry")
 
 			mockClient := test.NewClient()
 			tc.callMocks(mockClient)
@@ -178,15 +276,27 @@ func TestCreatePresetInference(t *testing.T) {
 
 			createdObject, _ := CreatePresetInference(context.TODO(), workspace, test.MockWorkspaceWithPresetHash, model, mockClient)
 			createdWorkload := ""
-			switch createdObject.(type) {
+			image := ""
+			envVars := []corev1.EnvVar{}
+			switch t := createdObject.(type) {
 			case *appsv1.Deployment:
 				createdWorkload = "Deployment"
+				image = t.Spec.Template.Spec.Containers[0].Image
+				envVars = t.Spec.Template.Spec.Containers[0].Env
 			case *appsv1.StatefulSet:
 				createdWorkload = "StatefulSet"
+				image = t.Spec.Template.Spec.Containers[0].Image
+				envVars = t.Spec.Template.Spec.Containers[0].Env
 			}
 			if tc.workload != createdWorkload {
 				t.Errorf("%s: returned workload type is wrong", k)
 			}
+			if image != tc.expectedImage {
+				t.Errorf("%s: Image is not expected, got %s, expect %s", k, image, tc.expectedImage)
+			}
+			if !reflect.DeepEqual(envVars, tc.expectedEnvVars) {
+				t.Errorf("%s: EnvVars are not expected, got %v, expect %v", k, envVars, tc.expectedEnvVars)
+			}
 
 			var workloadCmd string
 			if tc.workload == "Deployment" {
diff --git a/pkg/workspace/manifests/manifests_test.go b/pkg/workspace/manifests/manifests_test.go
index ddce817..5d0628e 100644
--- a/pkg/workspace/manifests/manifests_test.go
+++ b/pkg/workspace/manifests/manifests_test.go
@@ -33,6 +33,7 @@ func TestGenerateStatefulSetManifest(t *testing.T) {
 			nil, //tolerations
 			nil, //volumes
 			nil, //volumeMount
+			nil, //envVars
 		)
 
 		assert.Len(t, obj.OwnerReferences, 1, "Expected 1 OwnerReference")
@@ -85,6 +86,7 @@ func TestGenerateDeploymentManifest(t *testing.T) {
 			nil, //tolerations
 			nil, //volumes
 			nil, //volumeMount
+			nil, //envVars
 		)
 
 		assert.Len(t, obj.OwnerReferences, 1, "Expected 1 OwnerReference")
diff --git a/presets/workspace/models/metadata_test.go b/presets/workspace/models/metadata_test.go
new file mode 100644
index 0000000..83990e7
--- /dev/null
+++ b/presets/workspace/models/metadata_test.go
@@ -0,0 +1,3 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+package models