Repository: substratusai/kubeai
Pull request #416 (2025-02-17 07:15:34)

Ollama load model from PVC
Adding support for ollama PVC, specifically an e2e test for #392 to verify that it works. 


Patch:
diff --git a/Makefile b/Makefile
index 88dbfda..721a06c 100644
--- a/Makefile
+++ b/Makefile
@@ -113,6 +113,10 @@ test-e2e-engine-vllm-pvc: skaffold helm-dependency-build
 test-e2e-engine-infinity: skaffold helm-dependency-build
 	./test/e2e/run.sh engine-infinity --profile e2e-test-default
 
+.PHONY: test-e2e-engine-ollama-pvc
+test-e2e-engine-ollama-pvc: skaffold helm-dependency-build
+	./test/e2e/run.sh engine-ollama-pvc --profile e2e-test-default
+
 .PHONY: test-e2e-engine
 test-e2e-engine: skaffold helm-dependency-build
 	CACHE_PROFILE=$(CACHE_PROFILE) ./test/e2e/run.sh engine-$(ENGINE) --profile e2e-test-default
diff --git a/docs/how-to/install-models.md b/docs/how-to/install-models.md
index 1107ed1..15cad9b 100644
--- a/docs/how-to/install-models.md
+++ b/docs/how-to/install-models.md
@@ -94,22 +94,6 @@ spec:
   resourceProfile: nvidia-gpu-l4:1
 ```
 
-## Load Models from PVC
-
-You can store your models in a Persistent Volume Claim (PVC) and load them into KubeAI for serving. This guide will show you how to load models from a PVC.
-
-Currenly only vLLM supports loading models from PVCs.
-
-The following formats are supported to load models from a PVC:
-
-- `url: pvc://$PVC_NAME` - Loads the model from the PVC named `$PVC_NAME`.
-- `url: pvc://$PVC_NAME/$PATH` - Loads the model from the PVC named `$PVC_NAME` and mounts the subpath `$PATH` within the PVC.
-
-You need to make sure the model is preloaded into the PVC before you can use it in KubeAI.
-
-The Access Mode of the PVC should be `ReadOnlyMany` or `ReadWriteMany`, because otherwise
-KubeAI won't be able to spin up more than 1 replica of the model.
-
 ## Programmatically installing models
 
 See the [examples](https://github.com/substratusai/kubeai/tree/main/examples/k8s-api-clients).
diff --git a/docs/how-to/load-models-from-pvc.md b/docs/how-to/load-models-from-pvc.md
new file mode 100644
index 0000000..71a0eb3
--- /dev/null
+++ b/docs/how-to/load-models-from-pvc.md
@@ -0,0 +1,49 @@
+# Load Models from PVC
+
+You can store your models in a Persistent Volume Claim (PVC) and let KubeAI use them for serving.
+Both vLLM and Ollama engines support loading models from PVCs.
+
+You must ensure the model files are already present in the PVC before creating the Model resource.
+Alternatively you can use KubeAI's native caching mechanism which downloads the model for you:
+
+- [Cache Models with GCP Filestore](./cache-models-with-gcp-filestore.md)
+- [Cache Models with EFS](./cache-models-with-aws-efs.md)
+
+
+## vLLM
+
+For vLLM, use the following URL format:
+```yaml
+url: pvc://$PVC_NAME          # Loads the model from the PVC named $PVC_NAME
+url: pvc://$PVC_NAME/$PATH    # Loads from a specific path within the PVC
+```
+
+### PVC requirements
+
+vLLM supports both ReadWriteMany and ReadOnlyMany access modes. `Many` is used in order to support more than 1 vLLM replica.
+
+
+## Ollama
+
+For Ollama, use the following URL formats:
+```yaml
+url: pvc://$PVC_NAME?model=$MODEL_NAME    # Loads the model named $MODEL_NAME that's loaded on the disk
+url: pvc://$PVC_NAME/$PATH?model=$MODEL_NAME
+```
+
+### PVC Requirements
+Ollama requires using ReadWriteMany access mode because the rename operation `ollama cp` needs to write to the PVC.
+
+### Example: Loading Qwen 0.5b from PVC
+
+1. Create a PVC with ReadWriteMany named `model-pvc`. See [example](https://github.com/substratusai/kubeai/blob/main/examples/ollama-pvc/pvc.yaml).
+2. Create a K8s Job to load the model onto `model-pvc. See [example](https://github.com/substratusai/kubeai/blob/main/examples/ollama-pvc/job.yaml)
+
+    The PVC should now have a `blobs/` and `manifests/` directory after the loader completes.
+
+
+3. Create a Model to load from PVC:
+   
+   ```yaml
+   url: pvc://model-pvc?model=qwen:0.5b
+   ```
diff --git a/examples/ollama-pvc/job.yaml b/examples/ollama-pvc/job.yaml
new file mode 100644
index 0000000..914d7dc
--- /dev/null
+++ b/examples/ollama-pvc/job.yaml
@@ -0,0 +1,47 @@
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: ollama-load-model-to-pvc
+spec:
+  template:
+    spec:
+      containers:
+      - name: ollama
+        image: ollama/ollama:latest
+        env:
+        - name: OLLAMA_MODELS
+          value: /model
+        command: 
+        - /bin/sh
+        - -c
+        - |
+          /bin/ollama serve &
+          echo "Waiting for Ollama server to start..."
+          sleep 10
+          
+          # Pull the model and ensure it downloads successfully
+          echo "Pulling model qwen:0.5b..."
+          if ! /bin/ollama pull qwen:0.5b; then
+            echo "Failed to pull model"
+            exit 1
+          fi
+          
+          # Verify the model files exist
+          echo "Verifying model files..."
+          ls -R /model
+          if [ ! -d "/model/blobs" ] || [ ! -d "/model/manifests" ]; then
+            echo "Model directories not found"
+            exit 1
+          fi
+          
+          echo "Model setup completed successfully"
+          ls -la /model/manifests/registry.ollama.ai/library/qwen/0.5b
+        volumeMounts:
+        - name: models-volume
+          mountPath: /model
+      volumes:
+      - name: models-volume
+        persistentVolumeClaim:
+          claimName: model-pvc
+          readOnly: false
+      restartPolicy: OnFailure
diff --git a/examples/ollama-pvc/pvc.yaml b/examples/ollama-pvc/pvc.yaml
new file mode 100644
index 0000000..29c9ee6
--- /dev/null
+++ b/examples/ollama-pvc/pvc.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: model-pvc
+spec:
+  storageClassName: premium-rwx # replace with your actual storage class
+  accessModes:
+    - ReadWriteMany
+  resources:
+    requests:
+      storage: 10Gi
\ No newline at end of file
diff --git a/internal/modelcontroller/engine_ollama.go b/internal/modelcontroller/engine_ollama.go
index efc6c33..7464a77 100644
--- a/internal/modelcontroller/engine_ollama.go
+++ b/internal/modelcontroller/engine_ollama.go
@@ -33,6 +33,14 @@ func (r *ModelReconciler) oLlamaPodForModel(m *kubeaiv1.Model, c ModelConfig) *c
 			Value: "999999h",
 		},
 	}
+	// Adding an override variable for model location
+	if c.Source.url.scheme == "pvc" {
+		env = append(env, corev1.EnvVar{
+			Name:  "OLLAMA_MODELS",
+			Value: "/model",
+		})
+	}
+
 	var envKeys []string
 	for key := range m.Spec.Env {
 		envKeys = append(envKeys, key)
@@ -45,29 +53,12 @@ func (r *ModelReconciler) oLlamaPodForModel(m *kubeaiv1.Model, c ModelConfig) *c
 		})
 	}
 
-	ollamaModelRef := c.Source.url.ref
-
 	featuresMap := map[kubeaiv1.ModelFeature]struct{}{}
 	for _, f := range m.Spec.Features {
 		featuresMap[f] = struct{}{}
 	}
 
-	// Pull model and copy to rename it to Model.metadata.name.
-	// See Ollama issue for rename/copy workaround: https://github.com/ollama/ollama/issues/5914
-	// NOTE: The cp command should just create a pointer to the old model, not copy data
-	// (see https://github.com/ollama/ollama/issues/5914#issuecomment-2248168474).
-	// Use `ollama run` to send a single prompt to ollama to load the model into memory
-	// before the Pod becomes Ready. (by default it will load on the first prompt request).
-	startupProbeScript := fmt.Sprintf("/bin/ollama pull %s && /bin/ollama cp %s %s",
-		ollamaModelRef, ollamaModelRef, m.Name)
-	if _, ok := featuresMap[kubeaiv1.ModelFeatureTextGeneration]; ok {
-		// NOTE: Embedding text models do not support "ollama run":
-		//
-		// ollama run nomic-embed-text hey
-		// Error: "nomic-embed-text" does not support generate
-		//
-		startupProbeScript += fmt.Sprintf(" && /bin/ollama run %s hi", m.Name)
-	}
+	startupProbeScript := ollamaStartupProbeScript(m, c.Source.url)
 
 	pod := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -176,3 +167,38 @@ func (r *ModelReconciler) oLlamaPodForModel(m *kubeaiv1.Model, c ModelConfig) *c
 	return pod
 
 }
+
+func ollamaStartupProbeScript(m *kubeaiv1.Model, u modelURL) string {
+	// Pull model and copy to rename it to Model.metadata.name.
+	// See Ollama issue for rename/copy workaround: https://github.com/ollama/ollama/issues/5914
+	// NOTE: The cp command should just create a pointer to the old model, not copy data
+	// (see https://github.com/ollama/ollama/issues/5914#issuecomment-2248168474).
+	// Use `ollama run` to send a single prompt to ollama to load the model into memory
+	// before the Pod becomes Ready. (by default it will load on the first prompt request).
+	startupScript := ""
+	// If the model is using a pvc, we don't want to try to connect/pull a model
+
+	if u.scheme == "pvc" {
+		// There is a potential race condition when multiple pods try to rename/copy the same model.
+		startupScript = fmt.Sprintf("/bin/ollama cp %s %s",
+			u.modelParam, m.Name)
+	} else {
+		startupScript = fmt.Sprintf("/bin/ollama pull %s && /bin/ollama cp %s %s", u.ref, u.ref, m.Name)
+	}
+
+	// Only run the model if the model has features
+	featuresMap := map[kubeaiv1.ModelFeature]struct{}{}
+	for _, f := range m.Spec.Features {
+		featuresMap[f] = struct{}{}
+	}
+	if _, ok := featuresMap[kubeaiv1.ModelFeatureTextGeneration]; ok {
+		// NOTE: Embedding text models do not support "ollama run":
+		//
+		// ollama run nomic-embed-text hey
+		// Error: "nomic-embed-text" does not support generate
+		//
+		startupScript += fmt.Sprintf(" && /bin/ollama run %s hi", m.Name)
+	}
+
+	return startupScript
+}
diff --git a/internal/modelcontroller/model_source.go b/internal/modelcontroller/model_source.go
index 9cb1c31..5aad375 100644
--- a/internal/modelcontroller/model_source.go
+++ b/internal/modelcontroller/model_source.go
@@ -2,6 +2,7 @@ package modelcontroller
 
 import (
 	"fmt"
+	"net/url"
 	"regexp"
 	"strings"
 
@@ -214,28 +215,44 @@ func (r *ModelReconciler) pvcPodAdditions(url modelURL) *modelSourcePodAdditions
 	}
 }
 
-var modelURLRegex = regexp.MustCompile(`^([a-z]+):\/\/(\S+)$`)
+var modelURLRegex = regexp.MustCompile(`^([a-z]+):\/\/([^?]+)(\?.*)?$`)
 
 func parseModelURL(urlStr string) (modelURL, error) {
 	matches := modelURLRegex.FindStringSubmatch(urlStr)
-	if len(matches) != 3 {
+	if len(matches) != 3 && len(matches) != 4 {
 		return modelURL{}, fmt.Errorf("invalid model URL: %s", urlStr)
 	}
 	scheme, ref := matches[1], matches[2]
 	name, path, _ := strings.Cut(ref, "/")
+	var modelParam string
+
+	if len(matches) == 4 { // check for query parameter model, e.g. pvc://my-pvc?model=qwen2:0.5b
+		urlParser, err := url.ParseQuery(strings.TrimPrefix(matches[3], "?"))
+		if err == nil {
+			modelname := urlParser.Get("model")
+			if modelname != "" {
+				modelParam = modelname
+			}
+		}
+	}
+
 	return modelURL{
-		original: urlStr,
-		scheme:   scheme,
-		ref:      ref,
-		name:     name,
-		path:     path,
+		original:   urlStr,
+		scheme:     scheme,
+		ref:        ref,
+		name:       name,
+		path:       path,
+		modelParam: modelParam,
 	}, nil
 }
 
 type modelURL struct {
 	original string // e.g. "hf://username/model"
-	scheme   string // e.g. "hf", "s3", "gs", "oss"
+	scheme   string // e.g. "hf", "s3", "gs", "oss", "pvc"
 	ref      string // e.g. "username/model"
 	name     string // e.g. username or bucket-name
 	path     string // e.g. model or path/to/model
+	// e.g. "qwen2:0.5b" when ?model=qwen2:0.5b is part of the URL.
+	// This is used for Ollama where the PVC may have multiple models and we need to specify which one to load by name.
+	modelParam string
 }


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 9e703aa..f8ebb5c 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -35,6 +35,7 @@ jobs:
         - "autoscaler-restart-no-load"
         - "cache-shared-filesystem"
         - "engine-vllm-pvc"
+        - "engine-ollama-pvc"
         - "model-files"
     steps:
       - name: Checkout code
diff --git a/internal/modelcontroller/engine_ollama_test.go b/internal/modelcontroller/engine_ollama_test.go
new file mode 100644
index 0000000..d01fb08
--- /dev/null
+++ b/internal/modelcontroller/engine_ollama_test.go
@@ -0,0 +1,69 @@
+// TODO: write tests
+package modelcontroller
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	kubeaiv1 "github.com/substratusai/kubeai/api/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func Test_ollamaStartupProbeScript(t *testing.T) {
+	t.Parallel()
+
+	modelName := "model-name"
+	ollamaRef := "qwen2:0.5b"
+
+	cases := map[string]struct {
+		model       kubeaiv1.Model
+		modelURL    modelURL
+		featuresMap map[kubeaiv1.ModelFeature]struct{}
+		want        string
+	}{
+		"basic-model-no-pvc": {
+			model: kubeaiv1.Model{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: modelName,
+				},
+				Spec: kubeaiv1.ModelSpec{
+					Features: []kubeaiv1.ModelFeature{kubeaiv1.ModelFeatureTextGeneration},
+				},
+			},
+			modelURL: modelURL{
+				scheme: "ollama",
+				ref:    ollamaRef,
+				name:   "abc",
+			},
+			want: fmt.Sprintf("/bin/ollama pull %s && /bin/ollama cp %s %s && /bin/ollama run %s hi",
+				ollamaRef, ollamaRef, modelName, modelName),
+		},
+		"basic-model-with-pvc": {
+			model: kubeaiv1.Model{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: modelName,
+				},
+				Spec: kubeaiv1.ModelSpec{
+					Features: []kubeaiv1.ModelFeature{kubeaiv1.ModelFeatureTextGeneration},
+				},
+			},
+			modelURL: modelURL{
+				scheme:     "pvc",
+				ref:        "def",
+				name:       "abc",
+				modelParam: ollamaRef,
+			},
+			want: fmt.Sprintf("/bin/ollama cp %s %s && /bin/ollama run %s hi",
+				ollamaRef, modelName, modelName),
+		},
+	}
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			got := ollamaStartupProbeScript(&c.model, c.modelURL)
+			require.Equal(t, c.want, got)
+		})
+	}
+}
diff --git a/internal/modelcontroller/model_source_test.go b/internal/modelcontroller/model_source_test.go
index 2c3c453..f8f3beb 100644
--- a/internal/modelcontroller/model_source_test.go
+++ b/internal/modelcontroller/model_source_test.go
@@ -95,6 +95,26 @@ func Test_parseModelURL(t *testing.T) {
 				path:   "/",
 			},
 		},
+		"valid-pvc-with-modelname": {
+			input: "pvc://my-vpc?model=qwen2:0.5b",
+			want: modelURL{
+				scheme:     "pvc",
+				ref:        "my-vpc",
+				name:       "my-vpc",
+				path:       "",
+				modelParam: "qwen2:0.5b",
+			},
+		},
+		"valid-pvc-withpath-and-modelname": {
+			input: "pvc://my-vpc/path/to/model?model=qwen2:0.5b",
+			want: modelURL{
+				scheme:     "pvc",
+				ref:        "my-vpc/path/to/model",
+				name:       "my-vpc",
+				path:       "path/to/model",
+				modelParam: "qwen2:0.5b",
+			},
+		},
 	}
 
 	for name, c := range cases {
diff --git a/test/e2e/engine-ollama-pvc/ollama-hydrate-job.yaml b/test/e2e/engine-ollama-pvc/ollama-hydrate-job.yaml
new file mode 100644
index 0000000..fcb124c
--- /dev/null
+++ b/test/e2e/engine-ollama-pvc/ollama-hydrate-job.yaml
@@ -0,0 +1,47 @@
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: ollama-pvc-hydrate
+spec:
+  template:
+    spec:
+      containers:
+      - name: ollama
+        image: ollama/ollama:latest
+        env:
+        - name: OLLAMA_MODELS
+          value: /model
+        command: 
+        - /bin/sh
+        - -c
+        - |
+          /bin/ollama serve &
+          echo "Waiting for Ollama server to start..."
+          sleep 10
+          
+          # Pull the model and ensure it downloads successfully
+          echo "Pulling model qwen:0.5b..."
+          if ! /bin/ollama pull qwen:0.5b; then
+            echo "Failed to pull model"
+            exit 1
+          fi
+          
+          # Verify the model files exist
+          echo "Verifying model files..."
+          ls -R /model
+          if [ ! -d "/model/blobs" ] || [ ! -d "/model/manifests" ]; then
+            echo "Model directories not found"
+            exit 1
+          fi
+          
+          echo "Model setup completed successfully"
+          ls -la /model/manifests/registry.ollama.ai/library/qwen/0.5b
+        volumeMounts:
+        - name: models-volume
+          mountPath: /model
+      volumes:
+      - name: models-volume
+        persistentVolumeClaim:
+          claimName: model-pvc
+          readOnly: false
+      restartPolicy: OnFailure
diff --git a/test/e2e/engine-ollama-pvc/pv.yaml b/test/e2e/engine-ollama-pvc/pv.yaml
new file mode 100644
index 0000000..3b539b2
--- /dev/null
+++ b/test/e2e/engine-ollama-pvc/pv.yaml
@@ -0,0 +1,16 @@
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: kind-model-hostpath
+spec:
+  storageClassName: manual
+  capacity:
+    storage: 25Gi
+  accessModes:
+    - ReadWriteMany
+    - ReadOnlyMany
+    - ReadWriteOnce
+  hostPath:
+    path: $PV_HOST_PATH
+    type: DirectoryOrCreate
+  persistentVolumeReclaimPolicy: Retain
diff --git a/test/e2e/engine-ollama-pvc/pvc.yaml b/test/e2e/engine-ollama-pvc/pvc.yaml
new file mode 100644
index 0000000..6490f01
--- /dev/null
+++ b/test/e2e/engine-ollama-pvc/pvc.yaml
@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: model-pvc
+spec:
+  storageClassName: manual
+  accessModes:
+    - ReadWriteMany
+  resources:
+    requests:
+      storage: 10Gi
+  volumeName: kind-model-hostpath
\ No newline at end of file
diff --git a/test/e2e/engine-ollama-pvc/test.sh b/test/e2e/engine-ollama-pvc/test.sh
new file mode 100755
index 0000000..5bc6d05
--- /dev/null
+++ b/test/e2e/engine-ollama-pvc/test.sh
@@ -0,0 +1,49 @@
+#!/bin/bash
+
+source $REPO_DIR/test/e2e/common.sh
+
+models_release="kubeai-models"
+
+# Create PV_HOST_PATH inside the kind container
+kind_container=$(docker ps --filter "name=kind-control-plane" --format "{{.ID}}")
+export PV_HOST_PATH="/mnt/models"
+docker exec -i $kind_container mkdir -p $PV_HOST_PATH
+echo "PV_HOST_PATH: $PV_HOST_PATH"
+
+
+envsubst < $REPO_DIR/test/e2e/engine-ollama-pvc/pv.yaml | kubectl apply -f -
+kubectl apply -f $REPO_DIR/test/e2e/engine-ollama-pvc/pvc.yaml
+
+# Apply the Ollama hydrate job
+kubectl apply -f $REPO_DIR/test/e2e/engine-ollama-pvc/ollama-hydrate-job.yaml
+
+# Wait for job completion with timeout
+echo "Waiting for Ollama hydrate job to complete..."
+if ! kubectl wait --for=condition=complete --timeout=600s job/ollama-pvc-hydrate; then
+    echo "Ollama hydrate job failed or timed out"
+    kubectl logs job/ollama-pvc-hydrate
+    exit 1
+fi
+
+
+helm install $models_release $REPO_DIR/charts/models -f - <<EOF
+catalog:
+  qwen-500m-cpu:
+    enabled: true
+    url: pvc://model-pvc?model=qwen:0.5b
+    minReplicas: 2
+    engine: OLlama
+    resourceProfile: "cpu:1" 
+    features: [TextGeneration]
+EOF
+
+sleep 15
+
+# There are 2 replicas so send 10 requests to ensure that both replicas are used.
+for i in {1..10}; do
+  echo "Sending request $i"
+  curl http://localhost:8000/openai/v1/completions \
+    --max-time 600 \
+    -H "Content-Type: application/json" \
+    -d '{"model": "qwen-500m-cpu", "prompt": "Who was the first president of the United States?", "max_tokens": 40}'
+done