Repository: kaito-project/kaito

Issue:
Add gpt-oss-20b and gpt-oss-120b as KAITO Presets
**Is your feature request related to a problem? Please describe.**

https://openai.com/open-models/

According to https://huggingface.co/openai/gpt-oss-120b, we need to use a pre-released version of vLLM so we might need to wait until the official release of 0.10.1.

```
uv pip install --pre vllm==0.10.1+gptoss \
    --extra-index-url https://wheels.vllm.ai/gpt-oss/ \
    --extra-index-url https://download.pytorch.org/whl/nightly/cu128 \
    --index-strategy unsafe-best-match

vllm serve openai/gpt-oss-120b
```

**Describe the solution you'd like**

**Describe alternatives you've considered**

**Additional context**

Discussion:
@sdesai345: While vLLM v0.10.1 awaits release, let's investigate an upgrade to HF Transformers v4.55 to support gpt-oss as presets.@sdesai345: Update: https://github.com/vllm-project/vllm/releases/tag/v0.10.1 pre-release is out

Solution patch:
diff --git a/.github/actions/e2e-base-setup/action.yaml b/.github/actions/e2e-base-setup/action.yaml
index 79e9be6..eab7da0 100644
--- a/.github/actions/e2e-base-setup/action.yaml
+++ b/.github/actions/e2e-base-setup/action.yaml
@@ -114,11 +114,32 @@ runs:
       run: |
         make generate
 
+    - name: build KAITO base image
+      shell: bash
+      if: ${{ inputs.build_kaito_base_image == 'true' }}
+      run: |
+        docker buildx build \
+          --builder img-builder \
+          --build-arg VERSION=$VERSION \
+          --build-arg MODEL_TYPE=text-generation \
+          --output type=image,name=$REGISTRY/kaito-base:$VERSION,push=true,compression=zstd \
+          -f ./docker/presets/models/tfs/Dockerfile .
+
+        # Override base image properties in supported_models.yaml for testing purposes.
+        # The updated supported_models.yaml will be read by the Workspace controller via
+        # go:embed in presets/workspace/models/metadata.go.
+        yq '(.models[0].registry = "${{ env.REGISTRY }}")' -i presets/workspace/models/supported_models.yaml
+        yq '(.models[0].tag = "${{ env.VERSION }}")' -i presets/workspace/models/supported_models.yaml
+
+        cat presets/workspace/models/supported_models.yaml
+
     - name: build KAITO image
       if: ${{ inputs.isRelease == 'false' }}
       shell: bash
       run: |
         make docker-build-workspace
+      env:
+        ARCH: amd64
 
     - name: create cluster
       shell: bash
diff --git a/.github/determine_models.py b/.github/determine_models.py
index db57b45..ca3f17e 100644
--- a/.github/determine_models.py
+++ b/.github/determine_models.py
@@ -107,17 +107,10 @@ def detect_changes_in_yaml(yaml_main, yaml_pr):
 
 def models_to_build(files_changed):
     """Models to build based on changed files."""
-    models, seen_model_types = set(), set()
+    models = set()
     if supp_models_yaml in files_changed:
         yaml_main = get_yaml_from_branch("main", supp_models_yaml)
         models.update(detect_changes_in_yaml(yaml_main, YAML_PR))
-    for model, model_info in MODELS.items():
-        if model_info["type"] not in seen_model_types and any(
-            file.startswith(f"presets/workspace/inference/{model_info['type']}")
-            for file in files_changed
-        ):
-            models.add(model)
-            seen_model_types.add(model_info["type"])
     return list(models)
 
 
diff --git a/.github/workflows/e2e-workflow.yaml b/.github/workflows/e2e-workflow.yaml
index 99e8437..3b506cc 100644
--- a/.github/workflows/e2e-workflow.yaml
+++ b/.github/workflows/e2e-workflow.yaml
@@ -24,6 +24,8 @@ on:
       k8s_version:
         type: string
         default: "1.31.10"
+      pull_request_number:
+        type: string
 
 jobs:
   e2e-tests:
@@ -53,6 +55,33 @@ jobs:
         with:
           ref: ${{ inputs.git_sha }}
 
+      - name: Install gh
+        run: |
+          (type -p wget >/dev/null || (sudo apt update && sudo apt install wget -y)) \
+            && sudo mkdir -p -m 755 /etc/apt/keyrings \
+            && out=$(mktemp) && wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg \
+            && cat $out | sudo tee /etc/apt/keyrings/githubcli-archive-keyring.gpg > /dev/null \
+            && sudo chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg \
+            && sudo mkdir -p -m 755 /etc/apt/sources.list.d \
+            && echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null \
+            && sudo apt update \
+            && sudo apt install gh -y
+
+      - name: Set BUILD_KAITO_BASE_IMAGE
+        run: |
+          CHANGED_FILES=$(gh pr diff "${{ inputs.pull_request_number }}" --name-only)
+          # Check if specific files in the following directory are modified to determine if base image needs to be built
+          # - presets/workspace/inference/**
+          # - presets/workspace/tuning/**
+          # - docker/presets/models/tfs/Dockerfile
+          if echo "$CHANGED_FILES" | grep -E "(presets/workspace/inference|presets/workspace/tuning|docker/presets/models/tfs/Dockerfile)"; then
+            echo "BUILD_KAITO_BASE_IMAGE=true" >> $GITHUB_ENV
+          else
+            echo "BUILD_KAITO_BASE_IMAGE=false" >> $GITHUB_ENV
+          fi
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
       - name: Run e2e-base-setup composite action
         id: base-setup
         uses: ./.github/actions/e2e-base-setup
@@ -64,6 +93,7 @@ jobs:
           region: ${{ inputs.region }}
           k8s_version: ${{ inputs.k8s_version }}
           registry: ${{ inputs.registry }}
+          build_kaito_base_image: ${{ env.BUILD_KAITO_BASE_IMAGE }}
 
       - name: Set Registry
         if: ${{ inputs.isRelease }}
@@ -106,7 +136,7 @@ jobs:
           kubectl get pods -n kaito-workspace -o name | grep "^pod/kaito-workspace" | sed 's/^pod\///' | xargs -I {} kubectl logs -n kaito-workspace {}
 
       - name: Run full e2e test
-        if: ${{ inputs.isRelease }}
+        if: ${{ inputs.isRelease || env.BUILD_KAITO_BASE_IMAGE == 'true' }}
         run: |
           AI_MODELS_REGISTRY=$E2E_ACR_AMRT_USERNAME.azurecr.io \
           AI_MODELS_REGISTRY_SECRET=$E2E_AMRT_SECRET_NAME \
@@ -118,7 +148,7 @@ jobs:
           E2E_ACR_REGISTRY_SECRET: ${{ env.CLUSTER_NAME }}-acr-secret
 
       - name: Run fast e2e test
-        if: ${{ !inputs.isRelease }}
+        if: ${{ !inputs.isRelease && env.BUILD_KAITO_BASE_IMAGE != 'true' }}
         run: |
           AI_MODELS_REGISTRY=$E2E_ACR_AMRT_USERNAME.azurecr.io \
           AI_MODELS_REGISTRY_SECRET=$E2E_AMRT_SECRET_NAME \
@@ -130,7 +160,6 @@ jobs:
           E2E_ACR_REGISTRY: ${{ env.CLUSTER_NAME }}.azurecr.io
           E2E_ACR_REGISTRY_SECRET: ${{ env.CLUSTER_NAME }}-acr-secret
 
-
       - name: Cleanup e2e resources
         if: ${{ always() }}
         env:
@@ -139,4 +168,4 @@ jobs:
         with:
           inlineScript: |
             set +e
-            az group delete --name "${AZURE_RESOURCE_GROUP}" --yes --no-wait || true
\ No newline at end of file
+            az group delete --name "${AZURE_RESOURCE_GROUP}" --yes --no-wait || true
diff --git a/.github/workflows/preset-image-build-1ES.yaml b/.github/workflows/preset-image-build-1ES.yaml
index dcd9fe2..58424c3 100644
--- a/.github/workflows/preset-image-build-1ES.yaml
+++ b/.github/workflows/preset-image-build-1ES.yaml
@@ -5,12 +5,6 @@ concurrency:
   cancel-in-progress: true
 
 on:
-  pull_request:
-    branches:
-      - main
-    paths:
-      - 'presets/workspace/inference/**'
-      - 'presets/workspace/models/supported_models.yaml'
   push:
     branches:
       - main
diff --git a/.github/workflows/preset-image-build.yaml b/.github/workflows/preset-image-build.yaml
index e7ed603..f85718e 100644
--- a/.github/workflows/preset-image-build.yaml
+++ b/.github/workflows/preset-image-build.yaml
@@ -5,12 +5,6 @@ concurrency:
   cancel-in-progress: true
 
 on:
-  pull_request:
-    branches:
-      - main
-    paths:
-      - 'presets/workspace/inference/**'
-      - 'presets/workspace/models/supported_models.yaml'
   push:
     branches:
       - main
diff --git a/.github/workflows/workspace-e2e.yaml b/.github/workflows/workspace-e2e.yaml
index d8c398d..43e28ab 100644
--- a/.github/workflows/workspace-e2e.yaml
+++ b/.github/workflows/workspace-e2e.yaml
@@ -27,5 +27,6 @@ jobs:
       statuses: write
     uses: ./.github/workflows/e2e-workflow.yaml
     with:
+      pull_request_number: ${{ github.event.pull_request.number }}
       git_sha: ${{ github.event.pull_request.head.sha }}
       node_provisioner: ${{ matrix.node-provisioner }}
diff --git a/Makefile b/Makefile
index 96f98d5..0ecd9ca 100644
--- a/Makefile
+++ b/Makefile
@@ -209,7 +209,7 @@ create-acr: ## Create Azure container registry and login into it.
 create-aks-cluster: ## Create an AKS cluster with MSI, OIDC, and workload identity enabled.
 	az aks create  --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) \
 	--location $(AZURE_LOCATION) --attach-acr $(AZURE_ACR_NAME) \
-	--kubernetes-version $(AKS_K8S_VERSION) --node-count 1 --generate-ssh-keys  \
+	--kubernetes-version $(AKS_K8S_VERSION) --generate-ssh-keys  \
 	--enable-managed-identity --enable-workload-identity --enable-oidc-issuer --node-vm-size Standard_D2d_v4 -o none
 	az aks get-credentials --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) --overwrite-existing
 
@@ -217,7 +217,7 @@ create-aks-cluster: ## Create an AKS cluster with MSI, OIDC, and workload identi
 create-aks-cluster-with-kaito: ## Create an AKS cluster with MSI, OIDC, and Kaito enabled.
 	az aks create  --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) \
 	--location $(AZURE_LOCATION) --attach-acr $(AZURE_ACR_NAME) \
-	--kubernetes-version $(AKS_K8S_VERSION) --node-count 1 --generate-ssh-keys  \
+	--kubernetes-version $(AKS_K8S_VERSION) --generate-ssh-keys  \
 	--enable-managed-identity --enable-workload-identity --enable-oidc-issuer -o none
 	az aks get-credentials --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) --overwrite-existing
 
@@ -225,7 +225,7 @@ create-aks-cluster-with-kaito: ## Create an AKS cluster with MSI, OIDC, and Kait
 create-aks-cluster-for-karpenter: ## Create an AKS cluster with MSI, Cillium, OIDC, and workload identity enabled.
 	az aks create --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) \
     --location $(AZURE_LOCATION) --attach-acr $(AZURE_ACR_NAME) --node-vm-size "Standard_D2d_v4" \
-    --kubernetes-version $(AKS_K8S_VERSION) --node-count 3 --generate-ssh-keys \
+    --kubernetes-version $(AKS_K8S_VERSION) --generate-ssh-keys \
     --network-plugin azure --network-plugin-mode overlay --network-dataplane cilium \
     --enable-managed-identity --enable-oidc-issuer --enable-workload-identity -o none
 	az aks get-credentials --name $(AZURE_CLUSTER_NAME) --resource-group $(AZURE_RESOURCE_GROUP) --overwrite-existing
diff --git a/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml b/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
index 6f7b5a3..ac4a1b1 100644
--- a/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
+++ b/charts/kaito/workspace/crds/kaito.sh_workspaces.yaml
@@ -702,22 +702,28 @@ spec:
                   It is only set when the InferenceSpec is specified in the Workspace.
                 properties:
                   perReplicaNodeCount:
-                    description: PerReplicaNodeCount is used for recording the number
-                      of gpu nodes for one replica of workspace workload.
+                    description: |-
+                      PerReplicaNodeCount is used for recording the number of gpu nodes for one replica of workspace workload.
+                      This field remains immutable after being set.
                     format: int32
                     type: integer
                   replicas:
-                    description: Total number of running inference workloads of the
-                      workspace.
+                    description: |-
+                      Total number of running inference workloads of the workspace.
+                      This field is dynamic and changes based on scaling actions. It will set to Deployement.Status.Replicas
+                      when underlay workload is Deployment for Workspace.
                     format: int32
                     type: integer
                   selector:
-                    description: Selector is used to select the pods for the inference
-                      workload.
+                    description: |-
+                      Selector is used to select the pods that provide metrics for making scaling action decisions.
+                      This field must be set when HPA and VPA is used for scaling.
                     type: string
                   targetNodeCount:
-                    description: TargetNodeCount is used for recording the total number
-                      of gpu nodes for the workspace.
+                    description: |-
+                      TargetNodeCount is used for recording the desired number of gpu nodes that needed for all replicas of workspace.
+                      This field is dynamic and changes based on scaling actions. It has a specific calculation formula:
+                      TargetNodeCount = Resource.Replicas * InferenceStatus.PerReplicaNodeCount
                     format: int32
                     type: integer
                 type: object
diff --git a/charts/kaito/workspace/templates/supported-models-configmap.yaml b/charts/kaito/workspace/templates/supported-models-configmap.yaml
index 21c8f11..0213fe5 100644
--- a/charts/kaito/workspace/templates/supported-models-configmap.yaml
+++ b/charts/kaito/workspace/templates/supported-models-configmap.yaml
@@ -9,7 +9,7 @@ data:
       - name: base
         type: text-generation
         runtime: tfs
-        tag: 0.0.5
+        tag: 0.0.6
       - name: llama-3.1-8b-instruct
         type: text-generation
         version: https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/commit/0e9e39f249a16976918f6564b8830bc894c89659
diff --git a/pkg/model/interface.go b/pkg/model/interface.go
index da42ff1..3d69f58 100644
--- a/pkg/model/interface.go
+++ b/pkg/model/interface.go
@@ -16,14 +16,12 @@ package model
 import (
 	"fmt"
 	"maps"
-	"math"
 	"path"
 	"strconv"
 	"strings"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/kaito-project/kaito/pkg/sku"
@@ -56,20 +54,11 @@ const (
 
 	DefaultTuningMainFile = "/workspace/tfs/fine_tuning.py"
 	ConfigfileNameVLLM    = "inference_config.yaml"
-	DefaultMemoryUtilVLLM = 0.9 // Default gpu memory utilization for VLLM runtime
-	UpperMemoryUtilVLLM   = 0.95
 
 	// PortRayCluster is the default port for communication between the head and worker nodes in a Ray cluster.
 	PortRayCluster = 6379
 )
 
-var (
-	// vLLM will do kvcache pre-allocation,
-	// We need to reserve enough memory for other ephemeral operations to avoid OOM.
-	// This is an empirical value.
-	ReservedNonKVCacheMemory = resource.MustParse("1.5Gi")
-)
-
 // Metadata defines the metadata for a model.
 type Metadata struct {
 	// Name is the name of the model, which serves as a unique identifier.
@@ -102,6 +91,11 @@ type Metadata struct {
 	// If the model uses the Kaito base image, the tag field can be ignored
 	// +optional
 	Tag string `yaml:"tag,omitempty"`
+
+	// Registry is the container registry where the model is stored.
+	// If this is empty, os.Getenv("PRESET_REGISTRY_NAME") will be used.
+	// +optional
+	Registry string `yaml:"registry,omitempty"`
 }
 
 // Validate checks if the Metadata is valid.
@@ -282,8 +276,6 @@ func (p *PresetParam) buildVLLMInferenceCommand(rc RuntimeContext) []string {
 		}
 		p.VLLM.ModelRunParams["download-dir"] = utils.DefaultWeightsVolumePath
 	}
-	gpuMemUtil := getGPUMemoryUtilForVLLM(rc.GPUConfig)
-	p.VLLM.ModelRunParams["gpu-memory-utilization"] = strconv.FormatFloat(gpuMemUtil, 'f', 2, 64)
 	if rc.ConfigVolume != nil {
 		p.VLLM.ModelRunParams["kaito-config-file"] = path.Join(rc.ConfigVolume.MountPath, ConfigfileNameVLLM)
 	}
@@ -344,23 +336,6 @@ func (p *PresetParam) Validate(rc RuntimeContext) error {
 	return nil
 }
 
-func getGPUMemoryUtilForVLLM(gpuConfig *sku.GPUConfig) float64 {
-	if gpuConfig == nil || gpuConfig.GPUMemGB <= 0 || gpuConfig.GPUCount <= 0 {
-		return DefaultMemoryUtilVLLM
-	}
-
-	gpuMem := resource.MustParse(fmt.Sprintf("%dGi", gpuConfig.GPUMemGB))
-	gpuMemPerGPU := float64(gpuMem.Value()) / float64(gpuConfig.GPUCount)
-
-	if float64(ReservedNonKVCacheMemory.Value()) >= gpuMemPerGPU {
-		// looks impossible, just prevent this case
-		return DefaultMemoryUtilVLLM
-	}
-
-	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 {
diff --git a/pkg/utils/common_preset.go b/pkg/utils/common_preset.go
index 33b8f47..cb1060f 100644
--- a/pkg/utils/common_preset.go
+++ b/pkg/utils/common_preset.go
@@ -220,8 +220,9 @@ func ConfigAdapterVolume() (corev1.Volume, corev1.VolumeMount) {
 	return volume, volumeMount
 }
 
-func GetPresetImageName(name, tag string) string {
-	return fmt.Sprintf("%s/kaito-%s:%s",
-		os.Getenv("PRESET_REGISTRY_NAME"),
-		name, tag)
+func GetPresetImageName(registry, name, tag string) string {
+	if registry == "" {
+		registry = os.Getenv("PRESET_REGISTRY_NAME")
+	}
+	return fmt.Sprintf("%s/kaito-%s:%s", registry, name, tag)
 }
diff --git a/pkg/workspace/inference/preset_inferences.go b/pkg/workspace/inference/preset_inferences.go
index 1925645..5fe1466 100644
--- a/pkg/workspace/inference/preset_inferences.go
+++ b/pkg/workspace/inference/preset_inferences.go
@@ -298,7 +298,7 @@ func getDistributedInferenceProbe(probeType probeType, wObj *v1beta1.Workspace,
 
 func GetBaseImageName() string {
 	presetObj := metadata.MustGet("base")
-	return utils.GetPresetImageName(presetObj.Name, presetObj.Tag)
+	return utils.GetPresetImageName(presetObj.Registry, presetObj.Name, presetObj.Tag)
 }
 
 func GenerateInferencePodSpec(gpuConfig *sku.GPUConfig, numNodes int) func(*generator.WorkspaceGeneratorContext, *corev1.PodSpec) error {
diff --git a/pkg/workspace/manifests/manifests.go b/pkg/workspace/manifests/manifests.go
index dc082e2..c40b797 100644
--- a/pkg/workspace/manifests/manifests.go
+++ b/pkg/workspace/manifests/manifests.go
@@ -349,7 +349,7 @@ func GenerateDeploymentManifestWithPodTemplate(workspaceObj *kaitov1beta1.Worksp
 }
 
 func GetModelImageName(presetObj *pkgmodel.PresetParam) string {
-	return utils.GetPresetImageName(presetObj.Name, presetObj.Tag)
+	return utils.GetPresetImageName(presetObj.Registry, presetObj.Name, presetObj.Tag)
 }
 
 // GenerateModelPullerContainer creates an init container that pulls model images using ORAS
diff --git a/pkg/workspace/tuning/preset_tuning.go b/pkg/workspace/tuning/preset_tuning.go
index 9fd1391..2627eda 100644
--- a/pkg/workspace/tuning/preset_tuning.go
+++ b/pkg/workspace/tuning/preset_tuning.go
@@ -15,6 +15,7 @@ package tuning
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"path/filepath"
 	"strings"
@@ -44,6 +45,9 @@ const (
 )
 
 var (
+	//go:embed scripts/data-downloader.sh
+	dataDownloaderScript string
+
 	containerPorts = []corev1.ContainerPort{{
 		ContainerPort: consts.PortInferenceServer,
 	}}
@@ -68,7 +72,7 @@ var (
 
 func GetTuningImageInfo() string {
 	presetObj := metadata.MustGet("base")
-	return utils.GetPresetImageName(presetObj.Name, presetObj.Tag)
+	return utils.GetPresetImageName(presetObj.Registry, presetObj.Name, presetObj.Tag)
 }
 
 func GetDataSrcImageInfo(ctx context.Context, wObj *kaitov1beta1.Workspace) (string, []corev1.LocalObjectReference) {
@@ -371,7 +375,7 @@ func prepareDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace
 		return pullerContainer, []corev1.Volume{imagePullSecretVolume, dataVolume}, []corev1.VolumeMount{imagePullSecretVolumeMount, dataVolumeMount}
 
 	case len(input.URLs) > 0:
-		initContainer, volume, volumeMount := handleURLDataSource(ctx, workspaceObj)
+		initContainer, volume, volumeMount := handleURLDataSource(workspaceObj)
 		return initContainer, []corev1.Volume{volume}, []corev1.VolumeMount{volumeMount}
 
 	case input.Volume != nil:
@@ -383,39 +387,12 @@ func prepareDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace
 	}
 }
 
-func handleURLDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace) (*corev1.Container, corev1.Volume, corev1.VolumeMount) {
+func handleURLDataSource(workspaceObj *kaitov1beta1.Workspace) (*corev1.Container, corev1.Volume, corev1.VolumeMount) {
+	volume, volumeMount := utils.ConfigDataVolume(nil)
 	initContainer := &corev1.Container{
-		Name:  "data-downloader",
-		Image: "curlimages/curl",
-		Command: []string{"sh", "-c", `
-			if [ -z "$DATA_URLS" ]; then
-				echo "No URLs provided in DATA_URLS."
-				exit 1
-			fi
-			for url in $DATA_URLS; do
-				filename=$(basename "$url" | sed 's/[?=&]/_/g')
-				echo "Downloading $url to $DATA_VOLUME_PATH/$filename"
-				retry_count=0
-				while [ $retry_count -lt 3 ]; do
-					http_status=$(curl -sSL -w "%{http_code}" -o "$DATA_VOLUME_PATH/$filename" "$url")
-					curl_exit_status=$?  # Save the exit status of curl immediately
-					if [ "$http_status" -eq 200 ] && [ -s "$DATA_VOLUME_PATH/$filename" ] && [ $curl_exit_status -eq 0 ]; then
-						echo "Successfully downloaded $url"
-						break
-					else
-						echo "Failed to download $url, HTTP status code: $http_status, retrying..."
-						retry_count=$((retry_count + 1))
-						rm -f "$DATA_VOLUME_PATH/$filename" # Remove incomplete file
-						sleep 2
-					fi
-				done
-				if [ $retry_count -eq 3 ]; then
-					echo "Failed to download $url after 3 attempts"
-					exit 1  # Exit with a non-zero status to indicate failure
-				fi
-			done
-			echo "All downloads completed successfully"
-		`},
+		Name:    "data-downloader",
+		Image:   "curlimages/curl",
+		Command: []string{"sh", "-c", dataDownloaderScript},
 		Env: []corev1.EnvVar{
 			{
 				Name:  "DATA_URLS",
@@ -426,7 +403,7 @@ func handleURLDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspa
 				Value: utils.DefaultDataVolumePath,
 			},
 		},
+		VolumeMounts: []corev1.VolumeMount{volumeMount},
 	}
-	volume, volumeMount := utils.ConfigDataVolume(nil)
 	return initContainer, volume, volumeMount
 }
diff --git a/pkg/workspace/tuning/scripts/data-downloader.sh b/pkg/workspace/tuning/scripts/data-downloader.sh
new file mode 100644
index 0000000..9d7adc9
--- /dev/null
+++ b/pkg/workspace/tuning/scripts/data-downloader.sh
@@ -0,0 +1,29 @@
+#!/bin/sh
+
+if [ -z "$DATA_URLS" ]; then
+    echo "No URLs provided in DATA_URLS."
+    exit 1
+fi
+for url in $DATA_URLS; do
+    filename=$(basename "$url" | sed 's/[?#].*//')
+    echo "Downloading $url to $DATA_VOLUME_PATH/$filename"
+    retry_count=0
+    while [ $retry_count -lt 3 ]; do
+        http_status=$(curl -sSL -w "%{http_code}" -o "$DATA_VOLUME_PATH/$filename" "$url")
+        curl_exit_status=$?  # Save the exit status of curl immediately
+        if [ "$http_status" -eq "200" ] && [ -s "$DATA_VOLUME_PATH/$filename" ] && [ $curl_exit_status -eq 0 ]; then
+            echo "Successfully downloaded $url"
+            break
+        else
+            echo "Failed to download $url, HTTP status code: $http_status, retrying..."
+            retry_count=$((retry_count + 1))
+            rm -f "$DATA_VOLUME_PATH/$filename" # Remove incomplete file
+            sleep 2
+        fi
+    done
+    if [ $retry_count -eq 3 ]; then
+        echo "Failed to download $url after 3 attempts"
+        exit 1  # Exit with a non-zero status to indicate failure
+    fi
+done
+echo "All downloads completed successfully"
diff --git a/presets/ragengine/vector_store/base.py b/presets/ragengine/vector_store/base.py
index 9239178..8079892 100644
--- a/presets/ragengine/vector_store/base.py
+++ b/presets/ragengine/vector_store/base.py
@@ -379,7 +379,7 @@ class BaseVectorStore(ABC):
             max_tokens = self.llm.metadata.context_window - prompt_len
 
         logger.info(
-            f"Creating chat engine for index '{request.get('index_name')}' with prompt: {prompt}"
+            f"Creating chat engine for index '{request.get('index_name')}' with prompt size: {prompt_len}"
         )
         chat_engine = self.index_map[
             request.get("index_name")
@@ -387,7 +387,6 @@ class BaseVectorStore(ABC):
             llm=self.llm,
             similarity_top_k=100,  # Might want to make this a function of avg doc node size in an index but this should be a wide enough default
             chat_mode=ChatMode.CONDENSE_PLUS_CONTEXT,
-            verbose=True,
             node_postprocessors=[
                 ContextSelectionProcessor(
                     rag_context_token_fill_ratio=request.get(
diff --git a/presets/workspace/dependencies/requirements.txt b/presets/workspace/dependencies/requirements.txt
index 69ba063..4518369 100644
--- a/presets/workspace/dependencies/requirements.txt
+++ b/presets/workspace/dependencies/requirements.txt
@@ -1,10 +1,10 @@
 # Dependencies for TFS
 
 # Core Dependencies
-vllm==0.9.1
-transformers == 4.51.3
-torch==2.7.0
-accelerate==1.0.0
+vllm==0.10.1.1
+transformers==4.55.2
+torch==2.7.1
+accelerate==1.3.0
 fastapi==0.116.1
 pydantic>=2.9
 uvicorn[standard]>=0.29.0,<0.30.0  # Allow patch updates
@@ -14,15 +14,14 @@ numpy<3.0,>=1.25.0
 sentencepiece==0.2.0
 jinja2>=3.1.0
 starlette==0.47.2
-lmcache==0.3.2
+lmcache==0.3.3
 
 # For accessing Ray dashboard. ray[default] version should be consistent with the one in vllm/vllm-openai:<vllm-version>.
 # Check with `docker run --entrypoint "" vllm/vllm-openai:<vllm-version> pip freeze | grep ray`
-ray[default]==2.43.0
+ray[default]==2.48.0
 
 # Utility libraries
 datasets==2.19.1
-peft==0.11.1
 bitsandbytes==0.45.3
 
 # Less critical, can be latest
@@ -30,3 +29,4 @@ gputil
 setuptools
 psutil
 trl
+nvidia-ml-py3
diff --git a/presets/workspace/inference/vllm/inference_api.py b/presets/workspace/inference/vllm/inference_api.py
index 7bf9175..74949b0 100644
--- a/presets/workspace/inference/vllm/inference_api.py
+++ b/presets/workspace/inference/vllm/inference_api.py
@@ -21,6 +21,7 @@ from dataclasses import dataclass
 from typing import Any
 
 import psutil
+import pynvml
 import torch
 import uvloop
 import vllm.entrypoints.openai.api_server as api_server
@@ -90,7 +91,7 @@ class KAITOArgumentParser(argparse.ArgumentParser):
         engine_default_args = {
             "model": "/workspace/vllm/weights",
             "cpu_offload_gb": 0,
-            "gpu_memory_utilization": 0.95,
+            "gpu_memory_utilization": get_max_gpu_memory_utilization(),
             "swap_space": 4,
             "disable_log_stats": False,
             "uvicorn_log_level": "error",
@@ -291,6 +292,28 @@ def try_get_max_available_seq_len(args: argparse.Namespace) -> int | None:
         return None
 
 
+def get_max_gpu_memory_utilization(device_index: int = 0) -> float:
+    # Calculate gpu_memory_utilization based on available GPU memory.
+    # This ensures vLLM only uses currently free memory to avoid OOM errors.
+    # See https://github.com/kaito-project/kaito/issues/1374.
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+    pynvml.nvmlShutdown()
+
+    # Reserve an additional 600MiB for pytorch memory fragments, calculated based on profiling
+    free_memory = info.free - 600 * 1024**2
+
+    # Floor to 2 decimal places
+    gpu_memory_utilization = (free_memory * 100 // info.total) / 100
+
+    # The value is capped at 0.95 to maintain compatibility with previous behavior
+    gpu_memory_utilization = min(0.95, gpu_memory_utilization)
+
+    logger.info(f"Set default gpu_memory_utilization to {gpu_memory_utilization}")
+    return gpu_memory_utilization
+
+
 def set_kv_cache_offloading_if_appliable(args: argparse.Namespace) -> None:
     """
     Set KV cache offloading to CPU RAM if applicable.
@@ -318,10 +341,14 @@ def set_kv_cache_offloading_if_appliable(args: argparse.Namespace) -> None:
             psutil.virtual_memory().total - psutil.virtual_memory().used
         ) / (1024**3)
         logger.info(
-            f"VLLM_USE_V1 is set, Offload KV cache to CPU RAM, size limit: {available_memory_gb} * {args.kaito_kv_cache_cpu_memory_utilization} GB"
+            f"VLLM_USE_V1 is set, Offload KV cache to CPU RAM, size limit: {available_memory_gb} * {args.kaito_kv_cache_cpu_memory_utilization} GB split among {args.tensor_parallel_size} GPUs"
         )
+
+        # When using tensor parallelism, the KV cache CPU memory allocation must be divided evenly
+        # across all GPUs. Each GPU should only allocate its portion (1/tensor_parallel_size) of the
+        # total available CPU memory to prevent OOM.
         os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = (
-            f"{available_memory_gb * args.kaito_kv_cache_cpu_memory_utilization}"
+            f"{available_memory_gb * args.kaito_kv_cache_cpu_memory_utilization / args.tensor_parallel_size}"
         )
 
         if args.kv_transfer_config is None:
diff --git a/presets/workspace/models/deepseek/model.go b/presets/workspace/models/deepseek/model.go
index 576eace..fd5e7b5 100644
--- a/presets/workspace/models/deepseek/model.go
+++ b/presets/workspace/models/deepseek/model.go
@@ -56,7 +56,6 @@ var (
 	}
 	deepseekLlama8bRunParamsVLLM = map[string]string{
 		"dtype":            "float16",
-		"enable-reasoning": "",
 		"reasoning-parser": "deepseek_r1",
 	}
 	deepseekQwen14bRunParams = map[string]string{
@@ -65,7 +64,6 @@ var (
 	}
 	deepseekQwen14bRunParamsVLLM = map[string]string{
 		"dtype":            "float16",
-		"enable-reasoning": "",
 		"reasoning-parser": "deepseek_r1",
 	}
 	deepseekR1RunParams = map[string]string{
@@ -74,7 +72,6 @@ var (
 		"allow_remote_files": "",
 	}
 	deepseekR1RunParamsVLLM = map[string]string{
-		"enable-reasoning":        "",
 		"reasoning-parser":        "deepseek_r1",
 		"chat-template":           "/workspace/chat_templates/tool-chat-deepseekr1.jinja",
 		"tool-call-parser":        "deepseek_v3",
diff --git a/presets/workspace/models/supported_models.yaml b/presets/workspace/models/supported_models.yaml
index e8f2084..b82676c 100644
--- a/presets/workspace/models/supported_models.yaml
+++ b/presets/workspace/models/supported_models.yaml
@@ -3,8 +3,9 @@ models:
   - name: base
     type: text-generation
     runtime: tfs
-    tag: 0.0.5
+    tag: 0.0.6
     # Tag history:
+    # 0.0.6 - bump vLLM to 0.10.1.1, transformers to 4.55.2, torch to 2.7.1, lmcache to 0.3.3, ray to 2.48.0, accelerate to 1.3.0
     # 0.0.5 - update chat-template for new models
     # 0.0.4 - bump vLLM to 0.9.0.
     # 0.0.3 - add depdenencies to support vLLM multi-node distributed inference



Test patch:
diff --git a/pkg/model/interface_test.go b/pkg/model/interface_test.go
deleted file mode 100644
index 20058d7..0000000
--- a/pkg/model/interface_test.go
+++ /dev/null
@@ -1,91 +0,0 @@
-// Copyright (c) KAITO authors.
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package model
-
-import (
-	"testing"
-
-	"github.com/kaito-project/kaito/pkg/sku"
-)
-
-func TestGetGPUMemoryUtilForVLLM(t *testing.T) {
-	tests := []struct {
-		name      string
-		gpuConfig *sku.GPUConfig
-		expected  float64
-	}{
-		{
-			name:      "Nil GPUConfig",
-			gpuConfig: nil,
-			expected:  0.9,
-		},
-		{
-			name: "Zero GPU Memory",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 0,
-				GPUCount: 1,
-			},
-			expected: 0.9,
-		},
-		{
-			name: "Zero GPU Count",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 16,
-				GPUCount: 0,
-			},
-			expected: 0.9,
-		},
-		{
-			name: "V100",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 16,
-				GPUCount: 1,
-			},
-			expected: 0.9,
-		},
-		{
-			name: "V100 x 2",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 32,
-				GPUCount: 2,
-			},
-			expected: 0.9,
-		},
-		{
-			name: "A100",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 80,
-				GPUCount: 1,
-			},
-			expected: 0.95,
-		},
-		{
-			name: "Invalid",
-			gpuConfig: &sku.GPUConfig{
-				GPUMemGB: 1,
-				GPUCount: 1,
-			},
-			expected: 0.9,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := getGPUMemoryUtilForVLLM(tt.gpuConfig)
-			if result != tt.expected {
-				t.Errorf("expected %v, got %v", tt.expected, result)
-			}
-		})
-	}
-}
diff --git a/pkg/utils/common_preset_test.go b/pkg/utils/common_preset_test.go
new file mode 100644
index 0000000..7395a28
--- /dev/null
+++ b/pkg/utils/common_preset_test.go
@@ -0,0 +1,98 @@
+// Copyright (c) KAITO authors.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils
+
+import (
+	"os"
+	"testing"
+)
+
+func TestGetPresetImageName(t *testing.T) {
+	type args struct {
+		registry string
+		name     string
+		tag      string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "all parameters provided",
+			args: args{
+				registry: "myregistry.com",
+				name:     "phi-3",
+				tag:      "v1.0",
+			},
+			want: "myregistry.com/kaito-phi-3:v1.0",
+		},
+		{
+			name: "empty registry uses environment variable",
+			args: args{
+				registry: "",
+				name:     "llama-2",
+				tag:      "latest",
+			},
+			want: "test-registry.io/kaito-llama-2:latest",
+		},
+		{
+			name: "registry with different format",
+			args: args{
+				registry: "docker.io/kaito",
+				name:     "falcon",
+				tag:      "7b",
+			},
+			want: "docker.io/kaito/kaito-falcon:7b",
+		},
+		{
+			name: "complex model name with hyphens",
+			args: args{
+				registry: "azurecr.io",
+				name:     "phi-3-mini-128k-instruct",
+				tag:      "v2.1.3",
+			},
+			want: "azurecr.io/kaito-phi-3-mini-128k-instruct:v2.1.3",
+		},
+		{
+			name: "minimal case",
+			args: args{
+				registry: "reg",
+				name:     "m",
+				tag:      "t",
+			},
+			want: "reg/kaito-m:t",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set up environment variable for the test case that needs it
+			if tt.name == "empty registry uses environment variable" {
+				originalValue := os.Getenv("PRESET_REGISTRY_NAME")
+				os.Setenv("PRESET_REGISTRY_NAME", "test-registry.io")
+				defer func() {
+					if originalValue == "" {
+						os.Unsetenv("PRESET_REGISTRY_NAME")
+					} else {
+						os.Setenv("PRESET_REGISTRY_NAME", originalValue)
+					}
+				}()
+			}
+
+			if got := GetPresetImageName(tt.args.registry, tt.args.name, tt.args.tag); got != tt.want {
+				t.Errorf("GetPresetImageName() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/pkg/workspace/inference/preset_inferences_test.go b/pkg/workspace/inference/preset_inferences_test.go
index 1237623..d46ab48 100644
--- a/pkg/workspace/inference/preset_inferences_test.go
+++ b/pkg/workspace/inference/preset_inferences_test.go
@@ -67,7 +67,7 @@ func TestGeneratePresetInference(t *testing.T) {
 			expectedModelImage: "test-registry/kaito-test-model:1.0.0",
 			// No BaseCommand, AccelerateParams, 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",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --served-model-name=mymodel --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
 		},
 
@@ -82,7 +82,7 @@ func TestGeneratePresetInference(t *testing.T) {
 			expectedModelImage: "test-registry/kaito-test-no-tensor-parallel-model:1.0.0",
 			// No BaseCommand, AccelerateParams, 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",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
 		},
 
@@ -97,7 +97,7 @@ func TestGeneratePresetInference(t *testing.T) {
 			expectedModelImage: "test-registry/kaito-test-no-lora-support-model:1.0.0",
 			// No BaseCommand, AccelerateParams, 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",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
 		},
 
@@ -110,7 +110,7 @@ func TestGeneratePresetInference(t *testing.T) {
 			},
 			workload:           "Deployment",
 			expectedModelImage: "test-registry/kaito-test-model:1.0.0",
-			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",
+			expectedCmd:        "/bin/sh -c python3 /workspace/vllm/inference_api.py --enable-lora --tensor-parallel-size=2 --served-model-name=mymodel --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters:        true,
 			expectedVolume:     "adapter-volume",
 			expectedEnvVars: []corev1.EnvVar{{
@@ -160,7 +160,7 @@ func TestGeneratePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.ConfigMap{}), mock.Anything).Return(nil)
 			},
 			workload:    "Deployment",
-			expectedCmd: `/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --model=test-repo/test-model --code-revision=test-revision --gpu-memory-utilization=0.90 --download-dir=/workspace/weights --kaito-config-file=/mnt/config/inference_config.yaml`,
+			expectedCmd: `/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --model=test-repo/test-model --code-revision=test-revision --download-dir=/workspace/weights --kaito-config-file=/mnt/config/inference_config.yaml`,
 			expectedEnvVars: []corev1.EnvVar{{
 				Name: "HF_TOKEN",
 				ValueFrom: &corev1.EnvVarSource{
@@ -183,7 +183,7 @@ func TestGeneratePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.Service{}), mock.Anything).Return(nil)
 			},
 			workload:    "StatefulSet",
-			expectedCmd: `/bin/sh -c if [ "${POD_INDEX}" = "0" ]; then  --ray_cluster_size=2 --ray_port=6379; python3 /workspace/vllm/inference_api.py --model=test-repo/test-model --code-revision=test-revision --gpu-memory-utilization=0.90 --download-dir=/workspace/weights --kaito-config-file=/mnt/config/inference_config.yaml --pipeline-parallel-size=2 --tensor-parallel-size=2; else  --ray_address=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local --ray_port=6379; fi`,
+			expectedCmd: `/bin/sh -c if [ "${POD_INDEX}" = "0" ]; then  --ray_cluster_size=2 --ray_port=6379; python3 /workspace/vllm/inference_api.py --model=test-repo/test-model --code-revision=test-revision --download-dir=/workspace/weights --kaito-config-file=/mnt/config/inference_config.yaml --pipeline-parallel-size=2 --tensor-parallel-size=2; else  --ray_address=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local --ray_port=6379; fi`,
 			expectedEnvVars: []corev1.EnvVar{{
 				Name: "HF_TOKEN",
 				ValueFrom: &corev1.EnvVarSource{
@@ -215,7 +215,7 @@ func TestGeneratePresetInference(t *testing.T) {
 				c.On("Get", mock.IsType(context.TODO()), mock.Anything, mock.IsType(&corev1.Service{}), mock.Anything).Return(nil)
 			},
 			workload:    "StatefulSet",
-			expectedCmd: `/bin/sh -c if [ "${POD_INDEX}" = "0" ]; then  --ray_cluster_size=2 --ray_port=6379; python3 /workspace/vllm/inference_api.py --model=test-repo/test-model --code-revision=test-revision --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml --download-dir=/workspace/weights --pipeline-parallel-size=2 --tensor-parallel-size=2; else  --ray_address=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local --ray_port=6379; fi`,
+			expectedCmd: `/bin/sh -c if [ "${POD_INDEX}" = "0" ]; then  --ray_cluster_size=2 --ray_port=6379; python3 /workspace/vllm/inference_api.py --model=test-repo/test-model --code-revision=test-revision --kaito-config-file=/mnt/config/inference_config.yaml --download-dir=/workspace/weights --pipeline-parallel-size=2 --tensor-parallel-size=2; else  --ray_address=testWorkspace-0.testWorkspace-headless.kaito.svc.cluster.local --ray_port=6379; fi`,
 			expectedEnvVars: []corev1.EnvVar{{
 				Name: "HF_TOKEN",
 				ValueFrom: &corev1.EnvVarSource{
diff --git a/pkg/workspace/tuning/preset_tuning_test.go b/pkg/workspace/tuning/preset_tuning_test.go
index 9ba2f4a..4e249e0 100644
--- a/pkg/workspace/tuning/preset_tuning_test.go
+++ b/pkg/workspace/tuning/preset_tuning_test.go
@@ -160,6 +160,7 @@ func TestHandleURLDataSource(t *testing.T) {
 		expectedCommands          string
 		expectedVolumeName        string
 		expectedVolumeMountPath   string
+		expectedVolumeMounts      []corev1.VolumeMount
 	}{
 		"Handle URL Data Source": {
 			workspaceObj: &kaitov1beta1.Workspace{
@@ -174,20 +175,19 @@ func TestHandleURLDataSource(t *testing.T) {
 			expectedCommands:          "curl -sSL -w \"%{http_code}\" -o \"$DATA_VOLUME_PATH/$filename\" \"$url\"",
 			expectedVolumeName:        "data-volume",
 			expectedVolumeMountPath:   utils.DefaultDataVolumePath,
+			expectedVolumeMounts:      []corev1.VolumeMount{{Name: "data-volume", MountPath: utils.DefaultDataVolumePath}},
 		},
 	}
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			initContainer, volume, volumeMount := handleURLDataSource(context.Background(), tc.workspaceObj)
-
+			initContainer, volume, volumeMount := handleURLDataSource(tc.workspaceObj)
 			assert.Equal(t, tc.expectedInitContainerName, initContainer.Name)
 			assert.Equal(t, tc.expectedImage, initContainer.Image)
 			assert.Contains(t, normalize(initContainer.Command[2]), normalize(tc.expectedCommands))
-
 			assert.Equal(t, tc.expectedVolumeName, volume.Name)
-
 			assert.Equal(t, tc.expectedVolumeMountPath, volumeMount.MountPath)
+			assert.Equal(t, tc.expectedVolumeMounts, initContainer.VolumeMounts)
 		})
 	}
 }
diff --git a/presets/workspace/inference/text-generation/tests/test_inference_api.py b/presets/workspace/inference/text-generation/tests/test_inference_api.py
index b4df8a0..e1ecdea 100644
--- a/presets/workspace/inference/text-generation/tests/test_inference_api.py
+++ b/presets/workspace/inference/text-generation/tests/test_inference_api.py
@@ -302,7 +302,7 @@ def test_generation_with_max_length(configured_app):
         "prompt": prompt,
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"max_length": max_length},
+        "generate_kwargs": {"max_length": max_length, "max_new_tokens": None},
     }
 
     response = client.post("/chat", json=request_data)
@@ -346,7 +346,11 @@ def test_generation_with_min_length(configured_app):
         "prompt": prompt,
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"min_length": min_length, "max_length": max_length},
+        "generate_kwargs": {
+            "min_length": min_length,
+            "max_length": max_length,
+            "max_new_tokens": None,
+        },
     }
 
     response = client.post("/chat", json=request_data)