From 6426511dceaebfe06d52f26c9ac34e6eb8637bb5 Mon Sep 17 00:00:00 2001
From: Alex Fraser <alex.d.fraser@gmail.com>
Date: Sat, 9 Nov 2024 07:56:54 +0000
Subject: [PATCH] fix(ollama): pull models on container creation

Pull configured models when the container is created, rather than when
the image is built. Otherwise they may be pulled as root and be
inaccessible to the user.
---
 LICENSE                              |  1 +
 src/ollama/devcontainer-feature.json |  8 +++++---
 src/ollama/install.sh                | 25 ++++++++++++++++---------
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/LICENSE b/LICENSE
index 6775bd7..430a663 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,7 @@
 MIT License
 
 Copyright (c) 2022 Microsoft Corporation
+Copyright (c) 2024 Alex Fraser
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/src/ollama/devcontainer-feature.json b/src/ollama/devcontainer-feature.json
index ae12e53..e26d9c8 100644
--- a/src/ollama/devcontainer-feature.json
+++ b/src/ollama/devcontainer-feature.json
@@ -1,16 +1,18 @@
 {
     "name": "ollama",
     "id": "ollama",
-    "version": "1.0.0",
+    "version": "1.0.1",
     "description": "Installs ollama",
     "documentationURL": "https://github.com/prulloac/devcontainer-features/tree/main/src/ollama",
     "options": {
         "pull": {
-            "type":"string",
+            "type": "string",
             "default": "",
-            "description": "comma separated list of models to pull"
+            "description": "Comma-separated list of models to pull",
+            "proposals": ["llama3.2", "llama3.2:3b,nomic-embed-text"]
         }
     },
+    "onCreateCommand": "/usr/local/share/ollama/init.sh",
     "entrypoint": "ollama serve",
     "installsAfter": [
         "ghcr.io/devcontainers/features/common-utils"
diff --git a/src/ollama/install.sh b/src/ollama/install.sh
index a63375c..d26d06d 100644
--- a/src/ollama/install.sh
+++ b/src/ollama/install.sh
@@ -20,12 +20,19 @@ check_packages ca-certificates curl
 
 curl https://ollama.ai/install.sh | sh
 
-# if PULL is not empty and not equal to "none", also ollama is serving, then pull the models
-if [ "$PULL" != "none" ] && [ "$PULL" != "" ]; then
-    # split PULL variable into array using comma as delimiter and then pass to ollama pull
-    ollama serve &
-    sleep 3
-    echo $PULL | tr ',' '\n' | xargs -I % sh -c "ollama pull %"
-fi
-
-echo 'Done!'
\ No newline at end of file
+# Make sure it can run
+ollama --version
+
+# Save requested models for later. They can't be installed now because the user may be incorrect.
+mkdir -p "/usr/local/share/ollama"
+echo "$PULL" | tr ',' '\n' > "/usr/local/share/ollama/models.txt"
+cat << "EOF" > "/usr/local/share/ollama/init.sh"
+#!/bin/sh
+set -e
+echo "Pulling models" 1>&2
+# Ollama should already be running
+cat /usr/local/share/ollama/models.txt | xargs -I % sh -c "ollama pull %"
+EOF
+chmod +x "/usr/local/share/ollama/init.sh"
+
+echo 'Done!'