Merge pull request #55 from dreadnode/ads/eng-1506-bug-dyana-bug-python-version-issue-in-ollama-loader54

chore: refactor ollama loader to fix pydantic errors

Author: GangGreenTemperTatum

dyana/loaders/ollama/Dockerfile

@@ -2,14 +2,20 @@ FROM ollama/ollama
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y build-essential python3-pip python3
+RUN apt-get update && apt-get install -y --no-install-recommends gnupg build-essential curl
+RUN echo "deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu focal main" > /etc/apt/sources.list.d/deadsnakes.list
+RUN apt-key adv --keyserver hkp://keyserver.ubuntu.com:80 --recv-keys f23c5a6cf475977595c89f51ba6932366a755776
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive TZ=America/New_York apt-get install -y --no-install-recommends python3.12
+
+RUN curl https://bootstrap.pypa.io/get-pip.py -o pip.py
+RUN python3.12 pip.py
+
 COPY dyana.py .
 COPY dyana-requirements-gpu.txt .
 RUN pip install --no-cache-dir -r dyana-requirements-gpu.txt
 
 COPY requirements.txt .
-COPY main.py .
-
 RUN pip install --no-cache-dir -r requirements.txt
+COPY dyana.py main.py ./
 
-ENTRYPOINT ["python3", "-W", "ignore", "main.py"]
+ENTRYPOINT ["python3.12", "-W", "ignore", "main.py"]

dyana/loaders/ollama/main.py

@@ -1,45 +1,144 @@
 import argparse
+import json
 import os
+import sys
 import time
+import traceback
+from typing import Any
 
 from ollama import Client
 
-from dyana import Profiler  # type: ignore[attr-defined]
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run an Ollama model")
-    parser.add_argument("--model", help="Name of the Ollama model to profile", required=True)
-    parser.add_argument("--input", help="The input sentence", default="This is an example sentence.")
-    args = parser.parse_args()
+def ensure_output(data: str) -> None:
+    """Write to both stdout and stderr to ensure output is captured"""
+    print(data)
+    print(data, file=sys.stderr)
+    sys.stdout.flush()
+    sys.stderr.flush()
 
-    # start ollama server
-    os.system("ollama serve > /dev/null 2>&1 &")
-    for _ in range(30):
-        # print(f"waiting for ollama to start... {i}")
-        if os.system("ollama ls > /dev/null 2>&1") == 0:
-            break
-        time.sleep(1)
 
-    # create profiler after the server is started
-    profiler: Profiler = Profiler(gpu=True)
+try:
+    ensure_output(json.dumps({"status": "script_started"}))
 
     try:
-        client = Client(
-            host="http://127.0.0.1:11434",
-        )
-        response = client.chat(
-            model=args.model,
-            messages=[
-                {
-                    "role": "user",
-                    "content": args.input,
-                },
-            ],
-        )
-
-        profiler.on_stage("after_inference")
-
-        print(response)
-
-    except Exception as e:
-        profiler.track_error("ollama", str(e))
+        from dyana import Profiler  # type: ignore[attr-defined]
+    except ImportError:
+        # Only define our own Profiler if the import fails
+        class Profiler:  # type: ignore
+            def __init__(self, gpu: bool = False) -> None:
+                self.gpu = gpu
+
+            def on_stage(self, stage: str) -> None:
+                pass
+
+            def track_error(self, source: str, error: str) -> None:
+                pass
+
+    if __name__ == "__main__":
+        parser = argparse.ArgumentParser(description="Run an Ollama model")
+        parser.add_argument("--model", help="Name of the Ollama model to profile", required=True)
+        parser.add_argument("--input", help="The input sentence", default="This is an example sentence.")
+        args = parser.parse_args()
+
+        result: dict[str, Any] = {
+            "status": "started",
+            "model": args.model,
+            "input": args.input,
+            "timestamp": time.time(),
+        }
+        ensure_output(json.dumps(result))
+
+        try:
+            # Create profiler
+            profiler = Profiler(gpu=True)
+
+            os.makedirs("/root/.ollama/manifests", exist_ok=True)
+            os.makedirs("/root/.ollama/cache", exist_ok=True)
+
+            try:
+                os.chmod("/root/.ollama", 0o755)
+                os.chmod("/root/.ollama/models", 0o755)
+                os.chmod("/root/.ollama/manifests", 0o755)
+                os.chmod("/root/.ollama/cache", 0o755)
+            except Exception as perm_error:
+                result["permission_warning"] = str(perm_error)
+
+            # Start ollama server
+            os.system("ollama serve > /dev/null 2>&1 &")
+
+            # Wait for server to start
+            server_started = False
+            for i in range(30):
+                if os.system("ollama ls > /dev/null 2>&1") == 0:
+                    server_started = True
+                    result["startup_time"] = i
+                    break
+                time.sleep(1)
+
+            if not server_started:
+                result["status"] = "error"
+                result["error"] = "Failed to start Ollama server after 30 seconds"
+                ensure_output(json.dumps(result))
+                sys.exit(1)
+
+            # Record initialization stage
+            profiler.on_stage("initialization")
+
+            # Connect to the ollama server
+            client = Client(host="http://127.0.0.1:11434")
+
+            # Check if model exists locally without trying to pull it first
+            models = client.list()
+            model_exists = any(m.get("name", "") == args.model for m in models.get("models", []))
+
+            result["model_found"] = model_exists
+
+            if model_exists:
+                # Skip pulling if the model already exists
+                result["status"] = "running_inference"
+                ensure_output(json.dumps(result))
+
+                # Run inference with existing model
+                chat_response = client.chat(
+                    model=args.model,
+                    messages=[{"role": "user", "content": args.input}],
+                )
+
+                # Mark completion of inference
+                profiler.on_stage("after_inference")
+
+                # Update result with success
+                result["status"] = "success"
+                if hasattr(chat_response, "model_dump"):
+                    result["response"] = chat_response.model_dump()
+                else:
+                    result["response"] = str(chat_response)
+            else:
+                # Can't pull models due to read-only filesystem
+                result["status"] = "error"
+                result["error"] = (
+                    "Model not found locally and cannot pull due to read-only filesystem. Please pull the model on your host with 'ollama pull "
+                    + args.model
+                    + "' before running dyana."
+                )
+                ensure_output(json.dumps(result))
+
+        except Exception as e:
+            # Handle any exceptions
+            result["status"] = "error"
+            result["error"] = str(e)
+            result["traceback"] = traceback.format_exc()
+            if "profiler" in locals():
+                profiler.track_error("ollama", str(e))
+
+        # Output final result
+        ensure_output(json.dumps(result, default=str))
+
+except Exception as outer_e:
+    # Last resort error handling
+    emergency_data: dict[str, str] = {
+        "status": "fatal_error",
+        "error": str(outer_e),
+        "traceback": traceback.format_exc(),
+    }
+    ensure_output(json.dumps(emergency_data))

dyana/loaders/ollama/requirements.txt

@@ -1 +1,2 @@
-ollama==0.4.7
+ollama==0.4.7
+pydantic>=2.0.0

dyana/loaders/ollama/settings.yml

@@ -3,12 +3,22 @@ description: Loads and profiles models via an Ollama server. Local models on the
 gpu: true
 
 volumes:
-  # on macOS
+  # on macOS - mount with write permissions
   - host: ~/.ollama/models
     guest: /root/.ollama/models
-  # on Linux
+    options: "rw"
+  # on Linux - mount with write permissions
   - host: /usr/share/ollama/.ollama/models
     guest: /root/.ollama/models
+    options: "rw"
+  # Add a writable directory for Ollama to store manifests
+  - host: ~/.ollama/manifests
+    guest: /root/.ollama/manifests
+    options: "rw"
+  # Add directory for model cache
+  - host: ~/.ollama/cache
+    guest: /root/.ollama/cache
+    options: "rw"
 
 args:
   - name: model