Support auto-downloading ollama models with docker-compose. (#419)

The way this works is if ollama is enabled in the cookiecutter,
it will automatically ask the user to specify a model. Once that
model is specified, it will set up the compose configuration to
automatically pull that model when ollama first starts up. This
eliminates the need for users to manually pull a model on the
first use of the compose configuration, which seemed to be a
major source of issues.

.gitignore

@@ -109,7 +109,6 @@ Output.txt
 research_outputs
 report.md
 src/data/*
-!src/data/.gitkeep
 # Logs
 *.log
 logs/
@@ -169,10 +168,6 @@ simulation_*summary.md
 !examples/benchmarks/**/run_*.sh
 !examples/benchmarks/**/*.py
 
-# Retain directory structure but ignore contents
-!data/benchmark_results/.gitkeep
-!data/optimization_results/.gitkeep
-
 **/.claude/settings.local.json
 
 # Folders created by Docker.
@@ -197,3 +192,7 @@ star_reviews_*.png
 web*.png
 tests/screenshots/
 screenshots/
+
+# Ignore cookiecutter-generated files.
+docker-compose.*.yml
+scripts/*.sh

cookiecutter-docker/cookiecutter.json

@@ -5,6 +5,5 @@
   "host_ip": "0.0.0.0",
   "host_network": false,
   "enable_gpu": true,
-  "enable_ollama": true,
   "enable_searxng": true
 }

cookiecutter-docker/hooks/post_gen_project.py

@@ -14,6 +14,10 @@ def main():
     output_dir = compose_path.parent.absolute()
     compose_path.rename(output_dir.parent / COMPOSE_FILE_NAME)
 
+    # Move the entrypoint script to the scripts directory.
+    entrypoint_path = output_dir / "ollama_entrypoint.sh"
+    entrypoint_path.rename(output_dir.parent / "scripts" / "ollama_entrypoint.sh")
+
     # Delete the directory.
     shutil.rmtree(output_dir)
 

cookiecutter-docker/hooks/pre_prompt.py

@@ -3,12 +3,32 @@ import os
 import subprocess
 from typing import Any, Dict
 
+import cookiecutter.prompt
+
 
 def run_command(command: str):
     result = subprocess.run(command, shell=True, capture_output=True, text=True)
     return result.stdout.strip()
 
 
+def config_ollama(context: Dict[str, Any]) -> None:
+    """
+    Prompts the user for questions that are specific to Ollama. It is in a hook
+    so that we can run it only if Ollama is enabled.
+
+    """
+    enable_ollama = cookiecutter.prompt.read_user_yes_no("enable_ollama", True)
+    ollama_model = "gemma3:12b"
+    if enable_ollama:
+        # Ask ollama-specific questions.
+        ollama_model = cookiecutter.prompt.read_user_variable(
+            "ollama_model", ollama_model
+        )
+
+    context["_enable_ollama"] = enable_ollama
+    context["_ollama_model"] = ollama_model
+
+
 def check_gpu(context: Dict[str, Any]) -> None:
     """
     Check if the system has an NVIDIA or AMD GPU and set flags in the
@@ -45,6 +65,8 @@ def main() -> None:
     # Check GPU information and update the context only if running on Linux.
     if os.name == "posix" and os.uname().sysname == "Linux":
         check_gpu(context)
+    # Ollama-specific config.
+    config_ollama(context)
 
     # Save the updated context back to cookiecutter.json.
     with open("cookiecutter.json", "w") as config:

cookiecutter-docker/{{cookiecutter.config_name}}/docker-compose.{{cookiecutter.config_name}}.yml

@@ -20,7 +20,7 @@ services:
       # Web Interface Settings
       - LDR_WEB_PORT={{cookiecutter.host_port}}
       - LDR_WEB_HOST={{cookiecutter.host_ip}}
-      {%- if cookiecutter.enable_ollama %}
+      {%- if cookiecutter._enable_ollama %}
       - LDR_LLM_PROVIDER=ollama
       - LDR_LLM_OLLAMA_URL=http://ollama:11434
       {% endif %}
@@ -33,25 +33,36 @@ services:
       - ./local_collections/project_docs:/local_collections/project_docs/
       - ./local_collections/research_papers:/local_collections/research_papers/
     restart: unless-stopped
-    {%- if cookiecutter.enable_ollama or cookiecutter.enable_searxng %}
+    {%- if cookiecutter._enable_ollama or cookiecutter.enable_searxng %}
     depends_on:
-      {%- if cookiecutter.enable_ollama %}
-      - ollama
+      {%- if cookiecutter._enable_ollama %}
+      ollama:
+        condition: service_healthy
       {% endif %}
       {%- if cookiecutter.enable_searxng %}
-      - searxng
+      searxng:
+        condition: service_started
       {% endif %}
     {% endif %}
 
-  {%- if cookiecutter.enable_ollama %}
+  {%- if cookiecutter._enable_ollama %}
   ollama:
     {%- if cookiecutter.enable_gpu and cookiecutter._amd_gpu %}
     image: ollama/ollama:rocm
     {% else %}
     image: ollama/ollama:latest
     {% endif %}
+    container_name: ollama_service
+    entrypoint: [ "/scripts/ollama_entrypoint.sh" ]
+    healthcheck:
+      test: [ "CMD", "ollama", "show", "{{ cookiecutter._ollama_model }}" ]
+      interval: 10s
+      timeout: 5s
+      start_period: 10m
+      retries: 2
     volumes:
       - ollama_data:/root/.ollama
+      - ./scripts/:/scripts/
     {%- if not cookiecutter.host_network %}
     networks:
       - ldr-network

cookiecutter-docker/{{cookiecutter.config_name}}/ollama_entrypoint.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+set -e
+
+# Start the main Ollama application
+ollama serve &
+
+# Wait for the Ollama application to be ready (optional, if necessary)
+while ! ollama ls; do
+  echo "Waiting for Ollama service to be ready..."
+  sleep 10
+done
+echo "Ollama service is ready."
+
+# Pull the model using ollama pull
+echo "Pulling the {{ cookiecutter._ollama_model }} with ollama pull..."
+ollama pull {{ cookiecutter._ollama_model }}
+# Check if the model was pulled successfully
+if [ $? -eq 0 ]; then
+  echo "Model pulled successfully."
+else
+  echo "Failed to pull model."
+  exit 1
+fi
+
+# Run ollama forever.
+sleep infinity