ob-metaflow-extensions 1.1.142__py2.py3-none-any.whl → 1.4.33__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/ollama/ollama.py

@@ -5,6 +5,18 @@ import socket
 import sys
 import os
 import functools
+import json
+import requests
+from enum import Enum
+import threading
+from datetime import datetime
+
+from .constants import OLLAMA_SUFFIX
+from .exceptions import (
+    EmptyOllamaManifestCacheException,
+    EmptyOllamaBlobCacheException,
+    UnspecifiedRemoteStorageRootException,
+)
 
 
 class ProcessStatus:
@@ -13,26 +25,514 @@ class ProcessStatus:
     SUCCESSFUL = "SUCCESSFUL"
 
 
-class OllamaManager:
+class CircuitBreakerState(Enum):
+    CLOSED = "CLOSED"
+    OPEN = "OPEN"
+    HALF_OPEN = "HALF_OPEN"
+
+
+class CircuitBreaker:
+    def __init__(
+        self,
+        failure_threshold,
+        recovery_timeout,
+        reset_timeout,
+        debug=False,
+        status_card=None,
+    ):
+        self.failure_threshold = failure_threshold
+        self.recovery_timeout = recovery_timeout
+        self.reset_timeout = reset_timeout
+        self.state = CircuitBreakerState.CLOSED
+        self.failure_count = 0
+        self.last_failure_time = None
+        self.last_open_time = None
+        self.debug = debug
+        self.status_card = status_card
+        self.lock = threading.Lock()
+        self.request_count = 0  # Track total requests for pattern detection
+
+        if self.debug:
+            print(
+                f"[@ollama] CircuitBreaker initialized: threshold={failure_threshold}, recovery={recovery_timeout}, reset={reset_timeout}"
+            )
+
+    def _log_state_change(self, new_state):
+        if self.debug:
+            print(
+                f"[@ollama] Circuit Breaker state change: {self.state.value} -> {new_state.value}"
+            )
+        self.state = new_state
+        self._update_status_card()
+
+    def _update_status_card(self):
+        """Update the status card with current circuit breaker state"""
+        if self.status_card:
+            self.status_card.update_status(
+                "circuit_breaker",
+                {
+                    "state": self.state.value,
+                    "failure_count": self.failure_count,
+                    "last_failure_time": self.last_failure_time,
+                    "last_open_time": self.last_open_time,
+                },
+            )
+
+    def record_success(self):
+        with self.lock:
+            self.request_count += 1
+            if self.state == CircuitBreakerState.HALF_OPEN:
+                self._log_state_change(CircuitBreakerState.CLOSED)
+                self.failure_count = 0
+            elif self.state == CircuitBreakerState.OPEN:
+                # Allow transition to HALF_OPEN on success - server might have recovered
+                self._log_state_change(CircuitBreakerState.HALF_OPEN)
+                if self.debug:
+                    print(
+                        f"[@ollama] Success recorded while circuit OPEN. Transitioning to HALF_OPEN for testing."
+                    )
+            self.failure_count = 0
+            self.last_failure_time = None
+
+            # Log request count milestone for pattern detection
+            if self.debug and self.request_count % 100 == 0:
+                print(f"[@ollama] Request count: {self.request_count}")
+
+            self._update_status_card()
+
+    def record_failure(self):
+        with self.lock:
+            self.request_count += 1
+            self.failure_count += 1
+            self.last_failure_time = time.time()
+            if (
+                self.failure_count >= self.failure_threshold
+                and self.state == CircuitBreakerState.CLOSED
+            ):
+                self._log_state_change(CircuitBreakerState.OPEN)
+                self.last_open_time = time.time()
+            elif self.state == CircuitBreakerState.HALF_OPEN:
+                # If we fail while testing recovery, go back to OPEN
+                self._log_state_change(CircuitBreakerState.OPEN)
+                self.last_open_time = time.time()
+            if self.debug:
+                print(
+                    f"[@ollama] Failure recorded. Count: {self.failure_count}, State: {self.state.value}, Total requests: {self.request_count}"
+                )
+            self._update_status_card()
+
+    def should_attempt_reset(self):
+        """Check if we should attempt to reset/restart Ollama based on reset_timeout"""
+        with self.lock:
+            if self.state == CircuitBreakerState.OPEN and self.last_open_time:
+                elapsed_time = time.time() - self.last_open_time
+                return elapsed_time > self.reset_timeout
+            return False
+
+    def is_request_allowed(self):
+        with self.lock:
+            if self.state == CircuitBreakerState.OPEN:
+                elapsed_time = time.time() - self.last_open_time
+                if elapsed_time > self.recovery_timeout:
+                    self._log_state_change(CircuitBreakerState.HALF_OPEN)
+                    if self.debug:
+                        print(
+                            f"[@ollama] Circuit Breaker transitioning to HALF_OPEN after {elapsed_time:.1f}s."
+                        )
+                    return True  # Allow a single request to test recovery
+                else:
+                    if self.debug:
+                        print(
+                            f"[@ollama] Circuit Breaker is OPEN. Not allowing request. Time until HALF_OPEN: {self.recovery_timeout - elapsed_time:.1f}s"
+                        )
+                    return False
+            elif self.state == CircuitBreakerState.HALF_OPEN:
+                # In HALF_OPEN, be more restrictive - only allow one request at a time
+                if self.debug:
+                    print(
+                        f"[@ollama] Circuit Breaker is HALF_OPEN. Allowing request to test recovery."
+                    )
+                return True
+            else:  # CLOSED
+                return True
+
+    def get_status(self):
+        with self.lock:
+            status = {
+                "state": self.state.value,
+                "failure_count": self.failure_count,
+                "last_failure_time": self.last_failure_time,
+                "last_open_time": self.last_open_time,
+            }
+            if self.debug:
+                print(f"[@ollama] Circuit Breaker status: {status}")
+            return status
+
+
+class TimeoutCommand:
+    def __init__(self, command, timeout, debug=False, **kwargs):
+        self.command = command
+        self.timeout = timeout
+        self.debug = debug
+        self.input_data = kwargs.pop("input", None)  # Remove input from kwargs
+        self.kwargs = kwargs
+
+    def run(self):
+        if self.debug:
+            print(
+                f"[@ollama] Executing command with timeout {self.timeout}s: {' '.join(self.command)}"
+            )
+        try:
+            process = subprocess.Popen(
+                self.command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                **self.kwargs,
+            )
+            stdout, stderr = process.communicate(
+                input=self.input_data, timeout=self.timeout
+            )
+            return process.returncode, stdout, stderr
+        except subprocess.TimeoutExpired:
+            if self.debug:
+                print(
+                    f"[@ollama] Command timed out after {self.timeout}s: {' '.join(self.command)}"
+                )
+            process.kill()
+            stdout, stderr = process.communicate()
+            return (
+                124,
+                stdout,
+                stderr,
+            )  # 124 is the standard exit code for `timeout` command
+        except Exception as e:
+            if self.debug:
+                print(
+                    f"[@ollama] Error executing command {' '.join(self.command)}: {e}"
+                )
+            return 1, "", str(e)
+
+
+class OllamaHealthChecker:
+    def __init__(self, ollama_url, circuit_breaker, ollama_manager, debug=False):
+        self.ollama_url = ollama_url
+        self.circuit_breaker = circuit_breaker
+        self.ollama_manager = ollama_manager
+        self.debug = debug
+        self._stop_event = threading.Event()
+        self._thread = None
+        self._interval = 30  # Check every 30 seconds (less aggressive)
+
+    def _check_health(self):
+        try:
+            health_timeout = self.ollama_manager.timeouts.get("health_check", 5)
+            if self.debug:
+                print(f"[@ollama] Health check: Pinging {self.ollama_url}/api/tags")
+            response = requests.get(
+                f"{self.ollama_url}/api/tags", timeout=health_timeout
+            )
+            if response.status_code == 200:
+                self.circuit_breaker.record_success()
+                self._update_server_health_status("Healthy")
+                if self.debug:
+                    print(
+                        f"[@ollama] Health check successful. Circuit state: {self.circuit_breaker.state.value}"
+                    )
+                return True
+            else:
+                if self.debug:
+                    print(
+                        f"[@ollama] Health check failed. Status code: {response.status_code}. Circuit state: {self.circuit_breaker.state.value}"
+                    )
+                self.circuit_breaker.record_failure()
+                self._update_server_health_status(
+                    f"Unhealthy (HTTP {response.status_code})"
+                )
+                return False
+        except requests.exceptions.RequestException as e:
+            if self.debug:
+                print(
+                    f"[@ollama] Health check exception: {e}. Circuit state: {self.circuit_breaker.state.value}"
+                )
+            self.circuit_breaker.record_failure()
+            self._update_server_health_status(f"Unhealthy ({str(e)[:50]})")
+            return False
+
+    def _update_server_health_status(self, status):
+        """Update server health status in the status card"""
+        if self.ollama_manager.status_card:
+            self.ollama_manager.status_card.update_status(
+                "server", {"health_status": status, "last_health_check": datetime.now()}
+            )
+
+    def _run_health_check_loop(self):
+        while not self._stop_event.is_set():
+            # Always perform health check to monitor server status
+            self._check_health()
+
+            # Check if we should attempt a restart based on reset_timeout
+            if self.circuit_breaker.should_attempt_reset():
+                try:
+                    if self.debug:
+                        print(
+                            "[@ollama] Circuit breaker reset timeout reached. Attempting restart..."
+                        )
+                    restart_success = self.ollama_manager._attempt_ollama_restart()
+                    if restart_success:
+                        if self.debug:
+                            print("[@ollama] Restart successful via health checker")
+                    else:
+                        if self.debug:
+                            print("[@ollama] Restart failed via health checker")
+                except Exception as e:
+                    if self.debug:
+                        print(
+                            f"[@ollama] Error during health checker restart attempt: {e}"
+                        )
+
+            self._stop_event.wait(self._interval)
+
+    def start(self):
+        if self._thread is None or not self._thread.is_alive():
+            self._stop_event.clear()
+            self._thread = threading.Thread(
+                target=self._run_health_check_loop, daemon=True
+            )
+            self._thread.start()
+            if self.debug:
+                print("[@ollama] OllamaHealthChecker started.")
+
+    def stop(self):
+        if self._thread and self._thread.is_alive():
+            self._stop_event.set()
+            self._thread.join(timeout=self._interval + 1)  # Wait for thread to finish
+            if self.debug:
+                print("[@ollama] OllamaHealthChecker stopped.")
+
+
+class OllamaRequestInterceptor:
+    def __init__(self, circuit_breaker, debug=False):
+        self.circuit_breaker = circuit_breaker
+        self.debug = debug
+        self.original_methods = {}
+        self._protection_installed = False
+
+    def install_protection(self):
+        """Install request protection by monkey-patching the ollama package"""
+        if self._protection_installed:
+            return
+
+        try:
+            import ollama  # Import the actual ollama package
+
+            # Store original methods
+            self.original_methods = {
+                "chat": getattr(ollama, "chat", None),
+                "generate": getattr(ollama, "generate", None),
+                "embeddings": getattr(ollama, "embeddings", None),
+            }
+
+            # Replace with protected versions
+            if hasattr(ollama, "chat"):
+                ollama.chat = self._protected_chat
+            if hasattr(ollama, "generate"):
+                ollama.generate = self._protected_generate
+            if hasattr(ollama, "embeddings"):
+                ollama.embeddings = self._protected_embeddings
+
+            self._protection_installed = True
+            if self.debug:
+                print(
+                    "[@ollama] Request protection installed on ollama package methods"
+                )
+
+        except ImportError:
+            if self.debug:
+                print(
+                    "[@ollama] Warning: Could not import ollama package for request protection"
+                )
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama] Error installing request protection: {e}")
+
+    def remove_protection(self):
+        """Remove request protection by restoring original methods"""
+        if not self._protection_installed:
+            return
+
+        try:
+            import ollama
+
+            # Restore original methods
+            for method_name, original_method in self.original_methods.items():
+                if original_method is not None and hasattr(ollama, method_name):
+                    setattr(ollama, method_name, original_method)
+
+            self._protection_installed = False
+            if self.debug:
+                print("[@ollama] Request protection removed")
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama] Error removing request protection: {e}")
+
+    def _protected_chat(self, *args, **kwargs):
+        if not self.circuit_breaker.is_request_allowed():
+            raise RuntimeError(
+                f"Ollama server is currently unavailable. Circuit Breaker is {self.circuit_breaker.state.value}. "
+                "Please wait or check Ollama server status. "
+                f"Current status: {self.circuit_breaker.get_status()}"
+            )
+        try:
+            if self.debug:
+                # Debug: log model being used in request
+                model_name = kwargs.get("model", "unknown")
+                if args and isinstance(args[0], dict) and "model" in args[0]:
+                    model_name = args[0]["model"]
+                print(f"[@ollama] DEBUG: Making chat request with model: {model_name}")
+                print(f"[@ollama] DEBUG: Request args: {args}")
+                print(f"[@ollama] DEBUG: Request kwargs keys: {list(kwargs.keys())}")
+
+            result = self.original_methods["chat"](*args, **kwargs)
+            self.circuit_breaker.record_success()
+            return result
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama] Protected chat call failed: {e}")
+                print(f"[@ollama] DEBUG: Exception type: {type(e)}")
+            self.circuit_breaker.record_failure()
+            raise
+
+    def _protected_generate(self, *args, **kwargs):
+        if not self.circuit_breaker.is_request_allowed():
+            raise RuntimeError(
+                f"Ollama server is currently unavailable. Circuit Breaker is {self.circuit_breaker.state.value}. "
+                "Please wait or check Ollama server status. "
+                f"Current status: {self.circuit_breaker.get_status()}"
+            )
+        try:
+            result = self.original_methods["generate"](*args, **kwargs)
+            self.circuit_breaker.record_success()
+            return result
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama] Protected generate call failed: {e}")
+            self.circuit_breaker.record_failure()
+            raise
+
+    def _protected_embeddings(self, *args, **kwargs):
+        if not self.circuit_breaker.is_request_allowed():
+            raise RuntimeError(
+                f"Ollama server is currently unavailable. Circuit Breaker is {self.circuit_breaker.state.value}. "
+                "Please wait or check Ollama server status. "
+                f"Current status: {self.circuit_breaker.get_status()}"
+            )
+        try:
+            result = self.original_methods["embeddings"](*args, **kwargs)
+            self.circuit_breaker.record_success()
+            return result
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama] Protected embeddings call failed: {e}")
+            self.circuit_breaker.record_failure()
+            raise
+
 
+class OllamaManager:
     """
     A process manager for Ollama runtimes.
-    This is run locally, e.g., whether @ollama has a local, remote, or managed backend.
+    Implements interface @ollama([models=...], ...) has a local, remote, or managed backend.
     """
 
-    def __init__(self, models, backend="local", debug=False):
+    def __init__(
+        self,
+        models,
+        backend="local",
+        flow_datastore_backend=None,
+        remote_storage_root=None,
+        force_pull=False,
+        cache_update_policy="auto",
+        force_cache_update=False,
+        debug=False,
+        circuit_breaker_config=None,
+        timeout_config=None,
+        status_card=None,
+    ):
         self.models = {}
         self.processes = {}
+        self.flow_datastore_backend = flow_datastore_backend
+        if self.flow_datastore_backend is not None:
+            self.remote_storage_root = self.get_ollama_storage_root(
+                self.flow_datastore_backend
+            )
+        elif remote_storage_root is not None:
+            self.remote_storage_root = remote_storage_root
+        else:
+            raise UnspecifiedRemoteStorageRootException(
+                "Can not determine the storage root, as both flow_datastore_backend and remote_storage_root arguments of OllamaManager are None."
+            )
+        self.force_pull = force_pull
+
+        # New cache logic
+        self.cache_update_policy = cache_update_policy
+        if force_cache_update:  # Simple override
+            self.cache_update_policy = "force"
+        self.cache_status = {}  # Track cache status per model
+
         self.debug = debug
         self.stats = {}
+        self.storage_info = {}
+        self.ollama_url = "http://localhost:11434"  # Ollama API base URL
+        self.status_card = status_card
+        self.initialization_start = time.time()
 
         if backend != "local":
             raise ValueError(
                 "OllamaManager only supports the 'local' backend at this time."
             )
 
+        # Validate and set up circuit breaker config
+        if circuit_breaker_config is None:
+            circuit_breaker_config = {
+                "failure_threshold": 3,
+                "recovery_timeout": 30,  # Reduced from 60s - faster testing
+                "reset_timeout": 60,  # Reduced from 300s - faster restart
+            }
+
+        # Set up timeout configuration
+        if timeout_config is None:
+            timeout_config = {
+                "pull": 600,  # 10 minutes for model pulls
+                "stop": 30,  # 30 seconds for model stops
+                "health_check": 5,  # 5 seconds for health checks
+                "install": 60,  # 1 minute for Ollama installation
+                "server_startup": 300,  # 5 minutes for server startup
+            }
+        self.timeouts = timeout_config
+
+        # Initialize Circuit Breaker and Health Checker
+        self.circuit_breaker = CircuitBreaker(
+            failure_threshold=circuit_breaker_config.get("failure_threshold", 3),
+            recovery_timeout=circuit_breaker_config.get("recovery_timeout", 30),
+            reset_timeout=circuit_breaker_config.get("reset_timeout", 60),
+            debug=self.debug,
+            status_card=self.status_card,
+        )
+        self.health_checker = OllamaHealthChecker(
+            self.ollama_url, self.circuit_breaker, self, self.debug
+        )
+
+        self._log_event("info", "Starting Ollama initialization")
         self._timeit(self._install_ollama, "install_ollama")
         self._timeit(self._launch_server, "launch_server")
+        self.health_checker.start()
+
+        # Collect version information
+        self._collect_version_info()
+
+        # Initialize cache status display
+        self._update_cache_status()
 
         # Pull models concurrently
         with ThreadPoolExecutor() as executor:
@@ -41,41 +541,237 @@ class OllamaManager:
                 try:
                     future.result()
                 except Exception as e:
-                    raise RuntimeError(f"Error pulling one or more models: {e}") from e
+                    raise RuntimeError(f"Error pulling one or more models. {e}") from e
+
+        # Update final cache status
+        self._update_cache_status()
 
         # Run models as background processes.
         for m in models:
             f = functools.partial(self._run_model, m)
             self._timeit(f, f"model_{m.lower()}")
 
+        # Record total initialization time
+        total_init_time = time.time() - self.initialization_start
+        self._update_performance("total_initialization_time", total_init_time)
+        self._log_event(
+            "success", f"Ollama initialization completed in {total_init_time:.1f}s"
+        )
+
+    def _collect_version_info(self):
+        """Collect version information for Ollama system and Python client"""
+        version_info = {}
+
+        # Get Ollama system version
+        try:
+            result = subprocess.run(
+                ["ollama", "--version"], capture_output=True, text=True, timeout=10
+            )
+            if result.returncode == 0:
+                # Extract version from output - handle different formats
+                version_line = result.stdout.strip()
+                # Common formats: "ollama version 0.1.0", "0.1.0", "v0.1.0"
+                if "version" in version_line.lower():
+                    # Extract everything after "version"
+                    parts = version_line.lower().split("version")
+                    if len(parts) > 1:
+                        version_info["ollama_system"] = parts[1].strip()
+                    else:
+                        version_info["ollama_system"] = version_line
+                elif version_line.startswith("v"):
+                    version_info["ollama_system"] = version_line[
+                        1:
+                    ]  # Remove 'v' prefix
+                else:
+                    version_info["ollama_system"] = version_line
+            else:
+                version_info["ollama_system"] = "Unknown"
+        except Exception as e:
+            version_info["ollama_system"] = "Error detecting"
+            if self.debug:
+                print(f"[@ollama] Error getting system version: {e}")
+
+        # Get Python ollama client version
+        try:
+            import ollama
+
+            if hasattr(ollama, "__version__"):
+                version_info["ollama_python"] = ollama.__version__
+            else:
+                # Try alternative methods to get version
+                try:
+                    import pkg_resources
+
+                    version_info["ollama_python"] = pkg_resources.get_distribution(
+                        "ollama"
+                    ).version
+                except:
+                    try:
+                        # Try importlib.metadata (Python 3.8+)
+                        from importlib import metadata
+
+                        version_info["ollama_python"] = metadata.version("ollama")
+                    except:
+                        version_info["ollama_python"] = "Unknown"
+        except ImportError:
+            version_info["ollama_python"] = "Not installed"
+        except Exception as e:
+            version_info["ollama_python"] = "Error detecting"
+            if self.debug:
+                print(f"[@ollama] Error getting Python client version: {e}")
+
+        # Update status card with version info
+        if self.status_card:
+            self.status_card.update_status("versions", version_info)
+            self._log_event(
+                "info",
+                f"Versions: System {version_info.get('ollama_system', 'Unknown')}, Python {version_info.get('ollama_python', 'Unknown')}",
+            )
+
+    def _check_cache_exists(self, m):
+        """Check if cache exists for the given model"""
+        if self.local_datastore:
+            # Local datastore - no remote cache
+            return False
+
+        if m not in self.storage_info:
+            # Storage not set up yet
+            return False
+
+        try:
+            from metaflow import S3
+            from metaflow.plugins.datatools.s3.s3 import MetaflowS3NotFound
+
+            with S3() as s3:
+                # Check if manifest exists in remote storage
+                manifest_exists = s3.get(self.storage_info[m]["manifest_remote"]).exists
+
+                if manifest_exists:
+                    self.cache_status[m] = "exists"
+                    self._update_cache_status()
+                    return True
+                else:
+                    self.cache_status[m] = "missing"
+                    self._update_cache_status()
+                    return False
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {m}] Error checking cache existence: {e}")
+            self.cache_status[m] = "error"
+            self._update_cache_status()
+            return False
+
+    def _should_update_cache(self, m):
+        """Determine if we should update cache for this model based on policy"""
+        if self.cache_update_policy == "never":
+            return False
+        elif self.cache_update_policy == "force":
+            return True
+        elif self.cache_update_policy == "auto":
+            # Only update if cache doesn't exist
+            cache_exists = self._check_cache_exists(m)
+            return not cache_exists
+        else:
+            # Unknown policy, default to auto behavior
+            cache_exists = self._check_cache_exists(m)
+            return not cache_exists
+
+    def _log_event(self, event_type, message):
+        """Log an event to the status card"""
+        if self.status_card:
+            self.status_card.add_event(event_type, message)
+        if self.debug:
+            print(f"[@ollama] {event_type.upper()}: {message}")
+
+    def _update_server_status(self, status, **kwargs):
+        """Update server status in the status card"""
+        if self.status_card:
+            update_data = {"status": status}
+            update_data.update(kwargs)
+            self.status_card.update_status("server", update_data)
+
+    def _update_model_status(self, model_name, **kwargs):
+        """Update model status in the status card"""
+        if self.status_card:
+            current_models = self.status_card.status_data.get("models", {})
+            if model_name not in current_models:
+                current_models[model_name] = {}
+            current_models[model_name].update(kwargs)
+            self.status_card.update_status("models", current_models)
+
+    def _update_performance(self, metric, value):
+        """Update performance metrics in the status card"""
+        if self.status_card:
+            self.status_card.update_status("performance", {metric: value})
+
+    def _update_circuit_breaker_status(self):
+        """Update circuit breaker status in the status card"""
+        if self.status_card:
+            cb_status = self.circuit_breaker.get_status()
+            self.status_card.update_status("circuit_breaker", cb_status)
+
+    def _update_cache_status(self):
+        """Update cache status in the status card"""
+        if self.status_card:
+            self.status_card.update_status(
+                "cache",
+                {
+                    "policy": self.cache_update_policy,
+                    "model_status": self.cache_status.copy(),
+                },
+            )
+
     def _timeit(self, f, name):
         t0 = time.time()
         f()
         tf = time.time()
-        self.stats[name] = {"process_runtime": tf - t0}
+        duration = tf - t0
+        self.stats[name] = {"process_runtime": duration}
 
-    def _install_ollama(self, max_retries=3):
+        # Update performance metrics for status card
+        if name == "install_ollama":
+            self._update_performance("install_time", duration)
+        elif name == "launch_server":
+            self._update_performance("server_startup_time", duration)
+
+    def _is_port_open(self, host, port, timeout=1):
+        """Check if a TCP port is open on a given host."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(timeout)
+            try:
+                sock.connect((host, port))
+                return True
+            except socket.error:
+                return False
 
+    def _install_ollama(self, max_retries=3):
+        self._log_event("info", "Checking for existing Ollama installation")
         try:
             result = subprocess.run(["which", "ollama"], capture_output=True, text=True)
             if result.returncode == 0:
-                if self.debug:
-                    print("[@ollama] is already installed.")
+                self._log_event("success", "Ollama is already installed")
+                print("[@ollama] Ollama is already installed.")
                 return
         except Exception as e:
-            print("[@ollama] Did not find Ollama installation: %s" % e)
+            if self.debug:
+                print(f"[@ollama] Did not find Ollama installation: {e}")
             if sys.platform == "darwin":
                 raise RuntimeError(
-                    "On macOS, please install Ollama manually from https://ollama.com/download"
+                    "On macOS, please install Ollama manually from https://ollama.com/download."
                 )
 
+        self._log_event("info", "Installing Ollama...")
+        if self.debug:
+            print("[@ollama] Installing Ollama...")
         env = os.environ.copy()
         env["CURL_IPRESOLVE"] = "4"
+
         for attempt in range(max_retries):
             try:
                 install_cmd = ["curl", "-fsSL", "https://ollama.com/install.sh"]
                 curl_proc = subprocess.run(
-                    install_cmd, capture_output=True, text=True, env=env
+                    install_cmd, capture_output=True, text=True, env=env, timeout=120
                 )
                 if curl_proc.returncode != 0:
                     raise RuntimeError(
@@ -87,41 +783,42 @@ class OllamaManager:
                     capture_output=True,
                     text=True,
                     env=env,
+                    timeout=self.timeouts.get("install", 60),
                 )
                 if sh_proc.returncode != 0:
                     raise RuntimeError(
                         f"Ollama installation script failed: stdout: {sh_proc.stdout}, stderr: {sh_proc.stderr}"
                     )
+                self._log_event("success", "Ollama installation completed successfully")
                 if self.debug:
-                    print("[@ollama] Installed successfully.")
-                    break
+                    print("[@ollama] Ollama installed successfully.")
+                break
             except Exception as e:
-                print(f"Installation attempt {attempt+1} failed: {e}")
+                self._log_event(
+                    "warning", f"Installation attempt {attempt+1} failed: {str(e)}"
+                )
+                if self.debug:
+                    print(f"[@ollama] Installation attempt {attempt+1} failed: {e}")
                 if attempt < max_retries - 1:
                     time.sleep(5)
                 else:
+                    self._log_event(
+                        "error",
+                        f"Ollama installation failed after {max_retries} attempts",
+                    )
                     raise RuntimeError(
                         f"Error installing Ollama after {max_retries} attempts: {e}"
                     ) from e
 
-    def _is_port_open(self, host, port, timeout=1):
-        """Check if a TCP port is open on a given host."""
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.settimeout(timeout)
-            try:
-                sock.connect((host, port))
-                return True
-            except socket.error:
-                return False
-
     def _launch_server(self):
         """
         Start the Ollama server process and ensure it's running.
-        This version waits until the server is listening on port 11434.
         """
+        self._update_server_status("Starting")
+        self._log_event("info", "Starting Ollama server...")
+
         try:
-            if self.debug:
-                print("[@ollama] Starting Ollama server...")
+            print("[@ollama] Starting Ollama server...")
             process = subprocess.Popen(
                 ["ollama", "serve"],
                 stdout=subprocess.PIPE,
@@ -133,23 +830,40 @@ class OllamaManager:
                 "properties": {"type": "api-server", "error_details": None},
                 "status": ProcessStatus.RUNNING,
             }
+
             if self.debug:
-                print(
-                    "[@ollama] Started Ollama server process with PID %s" % process.pid
-                )
+                print(f"[@ollama] Started server process with PID {process.pid}.")
 
-            # Wait until the server is ready (listening on 127.0.0.1:11434)
+            # Wait until the server is ready
             host, port = "127.0.0.1", 11434
             retries = 0
             max_retries = 10
             while (
                 not self._is_port_open(host, port, timeout=1) and retries < max_retries
             ):
-                print(
-                    "[@ollama] Waiting for server to be ready... (%d/%d)"
-                    % (retries + 1, max_retries)
-                )
-                time.sleep(1)
+                if retries == 0:
+                    print("[@ollama] Waiting for server to be ready...")
+                elif retries % 3 == 0:
+                    print(f"[@ollama] Still waiting... ({retries + 1}/{max_retries})")
+
+                # Check if process terminated unexpectedly during startup
+                returncode = process.poll()
+                if returncode is not None:
+                    # Process exited, get error details but don't call communicate() which can hang
+                    error_details = f"Return code: {returncode}"
+                    self.processes[process.pid]["properties"][
+                        "error_details"
+                    ] = error_details
+                    self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                    self._update_server_status("Failed", error_details=error_details)
+                    self._log_event(
+                        "error", f"Ollama server failed to start: {error_details}"
+                    )
+                    raise RuntimeError(
+                        f"Ollama server failed to start. {error_details}"
+                    )
+
+                time.sleep(5)
                 retries += 1
 
             if not self._is_port_open(host, port, timeout=1):
@@ -160,51 +874,546 @@ class OllamaManager:
                     "error_details"
                 ] = error_details
                 self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self._update_server_status("Failed", error_details=error_details)
+                self._log_event("error", f"Server startup timeout: {error_details}")
                 raise RuntimeError(f"Ollama server failed to start. {error_details}")
 
-            # Check if the process has unexpectedly terminated
+            # Final check if process terminated unexpectedly
             returncode = process.poll()
             if returncode is not None:
-                stdout, stderr = process.communicate()
-                error_details = f"Return code: {returncode}, Error: {stderr}"
+                error_details = f"Return code: {returncode}"
                 self.processes[process.pid]["properties"][
                     "error_details"
                 ] = error_details
                 self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                self._update_server_status("Failed", error_details=error_details)
+                self._log_event(
+                    "error", f"Server process died unexpectedly: {error_details}"
+                )
                 raise RuntimeError(f"Ollama server failed to start. {error_details}")
 
+            self._update_server_status("Running", uptime_start=datetime.now())
+            self._log_event("success", "Ollama server is ready and listening")
             print("[@ollama] Server is ready.")
 
         except Exception as e:
             if "process" in locals() and process.pid in self.processes:
                 self.processes[process.pid]["status"] = ProcessStatus.FAILED
                 self.processes[process.pid]["properties"]["error_details"] = str(e)
+            self._update_server_status("Failed", error_details=str(e))
+            self._log_event("error", f"Error starting Ollama server: {str(e)}")
             raise RuntimeError(f"Error starting Ollama server: {e}") from e
 
-    def _pull_model(self, m):
-        try:
-            if self.debug:
-                print("[@ollama] Pulling model: %s" % m)
-            result = subprocess.run(
-                ["ollama", "pull", m], capture_output=True, text=True
-            )
-            if result.returncode != 0:
-                raise RuntimeError(
-                    f"Failed to pull model {m}: stdout: {result.stdout}, stderr: {result.stderr}"
-                )
-            if self.debug:
-                print("[@ollama] Model %s pulled successfully." % m)
-        except Exception as e:
-            raise RuntimeError(f"Error pulling Ollama model {m}: {e}") from e
-
-    def _run_model(self, m):
+    def _setup_storage(self, m):
         """
-        Start the Ollama model as a subprocess and record its status.
+        Configure local and remote storage paths for an Ollama model.
         """
-        process = None
-        try:
+        # Parse model and tag name
+        ollama_model_name_components = m.split(":")
+        if len(ollama_model_name_components) == 1:
+            model_name = ollama_model_name_components[0]
+            tag = "latest"
+        elif len(ollama_model_name_components) == 2:
+            model_name = ollama_model_name_components[0]
+            tag = ollama_model_name_components[1]
+
+        # Find where Ollama actually stores models
+        possible_storage_roots = [
+            os.environ.get("OLLAMA_MODELS"),
+            "/usr/share/ollama/.ollama/models",
+            os.path.expanduser("~/.ollama/models"),
+            "/root/.ollama/models",
+        ]
+
+        ollama_local_storage_root = None
+        for root in possible_storage_roots:
+            if root and os.path.exists(root):
+                ollama_local_storage_root = root
+                break
+
+        if not ollama_local_storage_root:
+            # https://github.com/ollama/ollama/blob/main/docs/faq.md#where-are-models-stored
+            if sys.platform.startswith("linux"):
+                ollama_local_storage_root = "/usr/share/ollama/.ollama/models"
+            elif sys.platform == "darwin":
+                ollama_local_storage_root = os.path.expanduser("~/.ollama/models")
+
+        if self.debug:
+            print(
+                f"[@ollama {m}] Using Ollama storage root: {ollama_local_storage_root}."
+            )
+
+        blob_local_path = os.path.join(ollama_local_storage_root, "blobs")
+        manifest_base_path = os.path.join(
+            ollama_local_storage_root,
+            "manifests/registry.ollama.ai/library",
+            model_name,
+        )
+
+        # Create directories
+        try:
+            os.makedirs(blob_local_path, exist_ok=True)
+            os.makedirs(manifest_base_path, exist_ok=True)
+        except FileExistsError:
+            pass
+
+        # Set up remote paths
+        if not self.local_datastore and self.remote_storage_root is not None:
+            blob_remote_key = os.path.join(self.remote_storage_root, "blobs")
+            manifest_remote_key = os.path.join(
+                self.remote_storage_root,
+                "manifests/registry.ollama.ai/library",
+                model_name,
+                tag,
+            )
+        else:
+            blob_remote_key = None
+            manifest_remote_key = None
+
+        self.storage_info[m] = {
+            "blob_local_root": blob_local_path,
+            "blob_remote_root": blob_remote_key,
+            "manifest_local": os.path.join(manifest_base_path, tag),
+            "manifest_remote": manifest_remote_key,
+            "manifest_content": None,
+            "model_name": model_name,
+            "tag": tag,
+            "storage_root": ollama_local_storage_root,
+        }
+
+        if self.debug:
+            print(f"[@ollama {m}] Storage paths configured.")
+
+    def _fetch_manifest(self, m):
+        """
+        Load the manifest file and content, either from local storage or remote cache.
+        """
+        if self.debug:
+            print(f"[@ollama {m}] Checking for cached manifest...")
+
+        def _disk_to_memory():
+            with open(self.storage_info[m]["manifest_local"], "r") as f:
+                self.storage_info[m]["manifest_content"] = json.load(f)
+
+        if os.path.exists(self.storage_info[m]["manifest_local"]):
+            if self.storage_info[m]["manifest_content"] is None:
+                _disk_to_memory()
+            if self.debug:
+                print(f"[@ollama {m}] Manifest found locally.")
+        elif self.local_datastore:
+            if self.debug:
+                print(f"[@ollama {m}] No manifest found in local datastore.")
+            return None
+        else:
+            from metaflow import S3
+            from metaflow.plugins.datatools.s3.s3 import MetaflowS3NotFound
+
+            try:
+                with S3() as s3:
+                    s3obj = s3.get(self.storage_info[m]["manifest_remote"])
+                    if not s3obj.exists:
+                        raise EmptyOllamaManifestCacheException(
+                            f"No manifest in remote storage for model {m}"
+                        )
+
+                    if self.debug:
+                        print(f"[@ollama {m}] Downloaded manifest from cache.")
+                    os.rename(s3obj.path, self.storage_info[m]["manifest_local"])
+                    _disk_to_memory()
+
+                    if self.debug:
+                        print(
+                            f"[@ollama {m}] Manifest found in remote cache, downloaded locally."
+                        )
+            except (MetaflowS3NotFound, EmptyOllamaManifestCacheException):
+                if self.debug:
+                    print(
+                        f"[@ollama {m}] No manifest found locally or in remote cache."
+                    )
+                return None
+
+        return self.storage_info[m]["manifest_content"]
+
+    def _fetch_blobs(self, m):
+        """
+        Fetch missing blobs from remote cache.
+        """
+        if self.debug:
+            print(f"[@ollama {m}] Checking for cached blobs...")
+
+        manifest = self._fetch_manifest(m)
+        if not manifest:
+            raise EmptyOllamaBlobCacheException(f"No manifest available for model {m}")
+
+        blobs_required = [layer["digest"] for layer in manifest["layers"]]
+        missing_blob_info = []
+
+        # Check which blobs are missing locally
+        for blob_digest in blobs_required:
+            blob_filename = blob_digest.replace(":", "-")
+            local_blob_path = os.path.join(
+                self.storage_info[m]["blob_local_root"], blob_filename
+            )
+
+            if not os.path.exists(local_blob_path):
+                if self.debug:
+                    print(f"[@ollama {m}] Blob {blob_digest} not found locally.")
+
+                remote_blob_path = os.path.join(
+                    self.storage_info[m]["blob_remote_root"], blob_filename
+                )
+                missing_blob_info.append(
+                    {
+                        "digest": blob_digest,
+                        "filename": blob_filename,
+                        "remote_path": remote_blob_path,
+                        "local_path": local_blob_path,
+                    }
+                )
+
+        if not missing_blob_info:
+            if self.debug:
+                print(f"[@ollama {m}] All blobs found locally.")
+            return
+
+        if self.debug:
+            print(
+                f"[@ollama {m}] Downloading {len(missing_blob_info)} missing blobs from cache..."
+            )
+
+        remote_urls = [blob_info["remote_path"] for blob_info in missing_blob_info]
+
+        from metaflow import S3
+
+        try:
+            with S3() as s3:
+                if len(remote_urls) == 1:
+                    s3objs = [s3.get(remote_urls[0])]
+                else:
+                    s3objs = s3.get_many(remote_urls)
+
+                if not isinstance(s3objs, list):
+                    s3objs = [s3objs]
+
+                # Move each downloaded blob to correct location
+                for i, s3obj in enumerate(s3objs):
+                    if not s3obj.exists:
+                        blob_info = missing_blob_info[i]
+                        raise EmptyOllamaBlobCacheException(
+                            f"Blob {blob_info['digest']} not found in remote cache for model {m}"
+                        )
+
+                    blob_info = missing_blob_info[i]
+                    os.makedirs(os.path.dirname(blob_info["local_path"]), exist_ok=True)
+                    os.rename(s3obj.path, blob_info["local_path"])
+
+                    if self.debug:
+                        print(f"[@ollama {m}] Downloaded blob {blob_info['filename']}.")
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {m}] Error during blob fetch: {e}")
+            raise EmptyOllamaBlobCacheException(
+                f"Failed to fetch blobs for model {m}: {e}"
+            )
+
+        if self.debug:
+            print(
+                f"[@ollama {m}] Successfully downloaded all missing blobs from cache."
+            )
+
+    def _verify_model_available(self, m):
+        """
+        Verify model is available using Ollama API
+        """
+        try:
+            if self.debug:
+                print(f"[@ollama] DEBUG: Verifying model availability for: {m}")
+
+            response = requests.post(
+                f"{self.ollama_url}/api/show", json={"model": m}, timeout=10
+            )
+
+            available = response.status_code == 200
+
             if self.debug:
-                print("[@ollama] Running model: %s" % m)
+                if available:
+                    print(f"[@ollama {m}] ✓ Model is available via API.")
+                    # Also list all available models for debugging
+                    try:
+                        tags_response = requests.get(
+                            f"{self.ollama_url}/api/tags", timeout=10
+                        )
+                        if tags_response.status_code == 200:
+                            models = tags_response.json().get("models", [])
+                            model_names = [
+                                model.get("name", "unknown") for model in models
+                            ]
+                            print(
+                                f"[@ollama] DEBUG: All available models: {model_names}"
+                            )
+                    except Exception as e:
+                        print(f"[@ollama] DEBUG: Could not list models: {e}")
+                else:
+                    print(
+                        f"[@ollama {m}] ✗ Model not available via API (status: {response.status_code})."
+                    )
+                    try:
+                        error_detail = response.text
+                        print(f"[@ollama] DEBUG: Error response: {error_detail}")
+                    except:
+                        pass
+
+            return available
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {m}] Error verifying model: {e}")
+            return False
+
+    def _register_cached_model_with_ollama(self, m):
+        """
+        Register a cached model with Ollama using the API.
+        """
+        try:
+            show_response = requests.post(
+                f"{self.ollama_url}/api/show", json={"model": m}, timeout=10
+            )
+
+            if show_response.status_code == 200:
+                if self.debug:
+                    print(f"[@ollama {m}] Model already registered with Ollama.")
+                return True
+
+            # Try to create/register the model from existing files
+            if self.debug:
+                print(f"[@ollama {m}] Registering cached model with Ollama...")
+
+            create_response = requests.post(
+                f"{self.ollama_url}/api/create",
+                json={
+                    "model": m,
+                    "from": m,  # Use same name - should find existing files
+                    "stream": False,
+                },
+                timeout=60,
+            )
+
+            if create_response.status_code == 200:
+                result = create_response.json()
+                if result.get("status") == "success":
+                    if self.debug:
+                        print(f"[@ollama {m}] Successfully registered cached model.")
+                    return True
+                else:
+                    if self.debug:
+                        print(f"[@ollama {m}] Create response: {result}.")
+
+            # Fallback: try a pull which should be fast if files exist
+            if self.debug:
+                print(f"[@ollama {m}] Create failed, trying pull to register...")
+
+            pull_response = requests.post(
+                f"{self.ollama_url}/api/pull",
+                json={"model": m, "stream": False},
+                timeout=120,
+            )
+
+            if pull_response.status_code == 200:
+                result = pull_response.json()
+                if result.get("status") == "success":
+                    if self.debug:
+                        print(f"[@ollama {m}] Model registered via pull.")
+                    return True
+
+        except requests.exceptions.RequestException as e:
+            if self.debug:
+                print(f"[@ollama {m}] API registration failed: {e}")
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {m}] Error during registration: {e}")
+
+        return False
+
+    def _pull_model(self, m):
+        """
+        Pull/setup a model, using cache when possible.
+        """
+        self._update_model_status(m, status="Setting up storage")
+        self._log_event("info", f"Setting up model {m}")
+        pull_start_time = time.time()
+
+        self._setup_storage(m)
+
+        # Check cache existence and inform user about cache strategy
+        cache_exists = self._check_cache_exists(m)
+        will_update_cache = self._should_update_cache(m)
+
+        if cache_exists:
+            if will_update_cache:
+                self._log_event(
+                    "info",
+                    f"Cache exists for {m}, but will be updated due to {self.cache_update_policy} policy",
+                )
+                print(
+                    f"[@ollama {m}] Cache exists but will be updated ({self.cache_update_policy} policy)"
+                )
+            else:
+                self._log_event("info", f"Using existing cache for {m}")
+                print(f"[@ollama {m}] Using existing cache")
+        else:
+            if will_update_cache:
+                self._log_event(
+                    "info",
+                    f"No cache found for {m}, will populate after successful setup",
+                )
+                print(f"[@ollama {m}] No cache found, will populate cache after setup")
+            else:
+                self._log_event(
+                    "info",
+                    f"No cache found for {m}, but cache updates disabled ({self.cache_update_policy} policy)",
+                )
+                print(
+                    f"[@ollama {m}] No cache found, cache updates disabled ({self.cache_update_policy} policy)"
+                )
+
+        # Try to fetch manifest from cache first
+        manifest = None
+        try:
+            manifest = self._fetch_manifest(m)
+        except (EmptyOllamaManifestCacheException, Exception) as e:
+            if self.debug:
+                print(f"[@ollama {m}] No cached manifest found or error fetching: {e}")
+            manifest = None
+
+        # If we don't have a cached manifest or force_pull is True, pull the model
+        if self.force_pull or not manifest:
+            try:
+                self._update_model_status(m, status="Downloading")
+                self._log_event("info", f"Downloading model {m}...")
+                print(f"[@ollama {m}] Not using cache. Downloading model {m}...")
+                result = subprocess.run(
+                    ["ollama", "pull", m],
+                    capture_output=True,
+                    text=True,
+                    timeout=self.timeouts.get("pull", 600),
+                )
+                if result.returncode != 0:
+                    self._update_model_status(m, status="Failed")
+                    self._log_event(
+                        "error", f"Failed to pull model {m}: {result.stderr}"
+                    )
+                    raise RuntimeError(
+                        f"Failed to pull model {m}: stdout: {result.stdout}, stderr: {result.stderr}"
+                    )
+                pull_time = time.time() - pull_start_time
+                self._update_model_status(m, status="Downloaded", pull_time=pull_time)
+                self._log_event("success", f"Model {m} downloaded in {pull_time:.1f}s")
+                print(f"[@ollama {m}] Model downloaded successfully.")
+            except Exception as e:
+                self._update_model_status(m, status="Failed")
+                self._log_event("error", f"Error pulling model {m}: {str(e)}")
+                raise RuntimeError(f"Error pulling Ollama model {m}: {e}") from e
+        else:
+            # We have a cached manifest, try to fetch the blobs
+            try:
+                self._update_model_status(m, status="Loading from cache")
+                self._log_event("info", f"Loading model {m} from cache")
+                self._fetch_blobs(m)
+                print(f"[@ollama {m}] Using cached model.")
+
+                # Register the cached model with Ollama
+                if not self._verify_model_available(m):
+                    if not self._register_cached_model_with_ollama(m):
+                        self._update_model_status(m, status="Failed")
+                        self._log_event("error", f"Failed to register cached model {m}")
+                        raise RuntimeError(
+                            f"Failed to register cached model {m} with Ollama"
+                        )
+
+                pull_time = time.time() - pull_start_time
+                self._update_model_status(m, status="Cached", pull_time=pull_time)
+                self._log_event(
+                    "success", f"Model {m} loaded from cache in {pull_time:.1f}s"
+                )
+
+            except (EmptyOllamaBlobCacheException, Exception) as e:
+                if self.debug:
+                    print(f"[@ollama {m}] Cache failed, downloading model...")
+                    print(f"[@ollama {m}] Error: {e}")
+
+                # Fallback to pulling the model
+                try:
+                    self._update_model_status(m, status="Downloading (fallback)")
+                    self._log_event(
+                        "warning", f"Cache failed for {m}, downloading as fallback"
+                    )
+                    result = subprocess.run(
+                        ["ollama", "pull", m],
+                        capture_output=True,
+                        text=True,
+                        timeout=self.timeouts.get("pull", 600),
+                    )
+                    if result.returncode != 0:
+                        self._update_model_status(m, status="Failed")
+                        self._log_event("error", f"Fallback pull failed for model {m}")
+                        raise RuntimeError(
+                            f"Failed to pull model {m}: stdout: {result.stdout}, stderr: {result.stderr}"
+                        )
+                    pull_time = time.time() - pull_start_time
+                    self._update_model_status(
+                        m, status="Downloaded (fallback)", pull_time=pull_time
+                    )
+                    self._log_event(
+                        "success",
+                        f"Model {m} downloaded via fallback in {pull_time:.1f}s",
+                    )
+                    print(f"[@ollama {m}] Model downloaded successfully (fallback).")
+                except Exception as pull_e:
+                    self._update_model_status(m, status="Failed")
+                    self._log_event(
+                        "error",
+                        f"Fallback download failed for model {m}: {str(pull_e)}",
+                    )
+                    raise RuntimeError(
+                        f"Error pulling Ollama model {m} as fallback: {pull_e}"
+                    ) from pull_e
+
+        # Final verification that the model is available
+        if not self._verify_model_available(m):
+            self._update_model_status(m, status="Failed")
+            self._log_event("error", f"Model {m} verification failed")
+            raise RuntimeError(f"Model {m} is not available to Ollama after setup")
+
+        # Collect model metadata (size and blob count)
+        metadata = self._collect_model_metadata(m)
+        self._update_model_status(
+            m,
+            status="Ready",
+            size_formatted=metadata["size_formatted"],
+            blob_count=metadata["blob_count"],
+        )
+        self._log_event("success", f"Model {m} setup complete and verified")
+        if self.debug:
+            print(f"[@ollama {m}] Model setup complete and verified.")
+            if metadata["size_formatted"] != "Unknown":
+                print(
+                    f"[@ollama {m}] Model size: {metadata['size_formatted']}, Blobs: {metadata['blob_count']}"
+                )
+
+    def _run_model(self, m):
+        """
+        Start the Ollama model as a subprocess and record its status.
+        """
+        process = None
+        try:
+            self._update_model_status(m, status="Starting process")
+            self._log_event("info", f"Starting model process for {m}")
+            if self.debug:
+                print(f"[@ollama {m}] Starting model process...")
+
+            # For `ollama run`, we want it to stay running, so no timeout on Popen.
+            # The health checker will detect if it becomes unresponsive.
             process = subprocess.Popen(
                 ["ollama", "run", m],
                 stdout=subprocess.PIPE,
@@ -216,113 +1425,500 @@ class OllamaManager:
                 "properties": {"type": "model", "model": m, "error_details": None},
                 "status": ProcessStatus.RUNNING,
             }
+
             if self.debug:
-                print("[@ollama] Stored process %s for model %s." % (process.pid, m))
+                print(f"[@ollama {m}] Model process PID: {process.pid}.")
 
+            # We don't want to wait here indefinitely. Just check if it failed immediately.
+            # The health checker will monitor long-term responsiveness.
             try:
-                process.wait(timeout=1)
-            except subprocess.TimeoutExpired:
-                pass
-
-            returncode = process.poll()
-            if returncode is not None:
-                stdout, stderr = process.communicate()
-                if returncode == 0:
-                    self.processes[process.pid]["status"] = ProcessStatus.SUCCESSFUL
-                    if self.debug:
-                        print(
-                            "[@ollama] Process %s for model %s exited successfully."
-                            % (process.pid, m)
-                        )
-                else:
+                process.wait(timeout=1)  # Check if it exited immediately
+                returncode = process.poll()
+                if (
+                    returncode is not None and returncode != 0
+                ):  # If it exited immediately with an error
+                    stdout, stderr = process.communicate()
                     error_details = f"Return code: {returncode}, Error: {stderr}"
                     self.processes[process.pid]["properties"][
                         "error_details"
                     ] = error_details
                     self.processes[process.pid]["status"] = ProcessStatus.FAILED
+                    self._update_model_status(m, status="Failed")
+                    self._log_event(
+                        "error",
+                        f"Model {m} process failed immediately: {error_details}",
+                    )
+                    if self.debug:
+                        print(
+                            f"[@ollama {m}] Process {process.pid} failed immediately: {error_details}."
+                        )
+                    raise RuntimeError(
+                        f"Ollama model {m} failed to start immediately: {error_details}"
+                    )
+                elif returncode == 0:
+                    # This case should ideally not happen for a long-running model
                     if self.debug:
                         print(
-                            "[@ollama] Process %s for model %s failed: %s"
-                            % (process.pid, m, error_details)
+                            f"[@ollama {m}] Process {process.pid} exited immediately with success. This might be unexpected for a model process."
                         )
+                    self.processes[process.pid]["status"] = ProcessStatus.SUCCESSFUL
+
+            except subprocess.TimeoutExpired:
+                # This is the expected case: process is running and hasn't exited
+                self._update_model_status(m, status="Running")
+                self._log_event("success", f"Model {m} process started successfully")
+                if self.debug:
+                    print(
+                        f"[@ollama {m}] Model process {process.pid} is running in background."
+                    )
+                pass  # Process is still running, which is good
+
         except Exception as e:
             if process and process.pid in self.processes:
                 self.processes[process.pid]["status"] = ProcessStatus.FAILED
                 self.processes[process.pid]["properties"]["error_details"] = str(e)
+            self._update_model_status(m, status="Failed")
+            self._log_event("error", f"Error running model {m}: {str(e)}")
             raise RuntimeError(f"Error running Ollama model {m}: {e}") from e
 
-    def terminate_models(self):
+    def terminate_models(self, skip_push_check=None):
         """
-        Terminate all processes gracefully.
-        First, stop model processes using 'ollama stop <model>'.
-        Then, shut down the API server process.
+        Terminate all processes gracefully and update cache.
         """
+        shutdown_start_time = time.time()
+        self._log_event("info", "Starting Ollama shutdown sequence")
+        print("[@ollama] Shutting down models...")
+
+        # Stop the health checker first
+        self.health_checker.stop()
+
+        # Handle backward compatibility for skip_push_check parameter
+        if skip_push_check is not None:
+            # Legacy parameter provided
+            if skip_push_check:
+                self.cache_update_policy = "never"
+                self._log_event(
+                    "warning",
+                    "Using legacy skip_push_check=True, setting cache policy to 'never'",
+                )
+            else:
+                self.cache_update_policy = "force"
+                self._log_event(
+                    "warning",
+                    "Using legacy skip_push_check=False, setting cache policy to 'force'",
+                )
 
+        # Shutdown models
+        model_shutdown_results = {}
         for pid, process_info in list(self.processes.items()):
             if process_info["properties"].get("type") == "model":
                 model_name = process_info["properties"].get("model")
+                model_shutdown_start = time.time()
+                shutdown_cause = "graceful"
+
+                self._update_model_status(model_name, status="Stopping")
+                self._log_event("info", f"Stopping model {model_name}")
                 if self.debug:
-                    print(
-                        "[@ollama] Stopping model %s using 'ollama stop'" % model_name
-                    )
+                    print(f"[@ollama {model_name}] Stopping model process...")
+
                 try:
                     result = subprocess.run(
-                        ["ollama", "stop", model_name], capture_output=True, text=True
+                        ["ollama", "stop", model_name],
+                        capture_output=True,
+                        text=True,
+                        timeout=self.timeouts.get("stop", 30),
                     )
                     if result.returncode == 0:
                         process_info["status"] = ProcessStatus.SUCCESSFUL
+                        self._update_model_status(model_name, status="Stopped")
+                        self._log_event(
+                            "success", f"Model {model_name} stopped gracefully"
+                        )
                         if self.debug:
-                            print(
-                                "[@ollama] Model %s stopped successfully." % model_name
-                            )
+                            print(f"[@ollama {model_name}] Stopped successfully.")
                     else:
                         process_info["status"] = ProcessStatus.FAILED
+                        shutdown_cause = "force_kill"
+                        self._update_model_status(model_name, status="Force stopped")
+                        self._log_event(
+                            "warning",
+                            f"Model {model_name} stop command failed, killing process",
+                        )
                         if self.debug:
                             print(
-                                "[@ollama] Model %s failed to stop gracefully. Return code: %s, Error: %s"
-                                % (model_name, result.returncode, result.stderr)
+                                f"[@ollama {model_name}] Stop failed: {result.stderr}. Attempting to kill process directly."
+                            )
+                        # Fallback: if 'ollama stop' fails, try to kill the process directly
+                        try:
+                            process_info["p"].terminate()
+                            process_info["p"].wait(timeout=5)
+                            if process_info["p"].poll() is None:
+                                process_info["p"].kill()
+                                process_info["p"].wait()
+                            process_info["status"] = ProcessStatus.SUCCESSFUL
+                            self._update_model_status(model_name, status="Killed")
+                            self._log_event(
+                                "warning", f"Model {model_name} process killed directly"
+                            )
+                            if self.debug:
+                                print(
+                                    f"[@ollama {model_name}] Process killed directly."
+                                )
+                        except Exception as kill_e:
+                            process_info["status"] = ProcessStatus.FAILED
+                            shutdown_cause = "failed"
+                            self._update_model_status(
+                                model_name, status="Failed to stop"
+                            )
+                            self._log_event(
+                                "error",
+                                f"Model {model_name} failed to stop: {str(kill_e)}",
+                            )
+                            print(
+                                f"[@ollama {model_name}] Error killing process directly: {kill_e}"
                             )
+
                 except Exception as e:
                     process_info["status"] = ProcessStatus.FAILED
-                    print("[@ollama] Error stopping model %s: %s" % (model_name, e))
+                    shutdown_cause = "failed"
+                    self._update_model_status(model_name, status="Failed to stop")
+                    self._log_event(
+                        "error", f"Model {model_name} shutdown error: {str(e)}"
+                    )
+                    print(f"[@ollama {model_name}] Error stopping: {e}")
+
+                # Record model shutdown timing
+                model_shutdown_time = time.time() - model_shutdown_start
+                model_shutdown_results[model_name] = {
+                    "shutdown_time": model_shutdown_time,
+                    "shutdown_cause": shutdown_cause,
+                }
+
+                # Smart cache update logic
+                should_update = self._should_update_cache(model_name)
+                if should_update:
+                    self._log_event(
+                        "info",
+                        f"Updating cache for {model_name} ({self.cache_update_policy} policy)",
+                    )
+                    self._update_model_cache(model_name)
+                else:
+                    cache_reason = f"policy is '{self.cache_update_policy}'"
+                    if (
+                        self.cache_update_policy == "auto"
+                        and self.cache_status.get(model_name) == "exists"
+                    ):
+                        cache_reason = "cache already exists"
+                    self._log_event(
+                        "info",
+                        f"Skipping cache update for {model_name} ({cache_reason})",
+                    )
+                    if self.debug:
+                        print(
+                            f"[@ollama {model_name}] Skipping cache update: {cache_reason}"
+                        )
 
-        # Then, stop the API server
+        # Stop the API server
+        server_shutdown_cause = "graceful"
+        server_shutdown_start = time.time()
         for pid, process_info in list(self.processes.items()):
             if process_info["properties"].get("type") == "api-server":
+                self._update_server_status("Stopping")
+                self._log_event("info", "Stopping Ollama API server")
                 if self.debug:
-                    print(
-                        "[@ollama] Stopping API server process with PID %s using process.terminate()"
-                        % pid
-                    )
+                    print(f"[@ollama] Stopping API server process PID {pid}.")
+
                 process = process_info["p"]
                 try:
                     process.terminate()
                     try:
                         process.wait(timeout=5)
                     except subprocess.TimeoutExpired:
+                        server_shutdown_cause = "force_kill"
+                        self._log_event(
+                            "warning",
+                            "API server did not terminate gracefully, killing...",
+                        )
                         print(
-                            "[@ollama] API server process %s did not terminate in time; killing it."
-                            % pid
+                            f"[@ollama] API server PID {pid} did not terminate, killing..."
                         )
                         process.kill()
                         process.wait()
-                    returncode = process.poll()
-                    if returncode is None or returncode != 0:
-                        process_info["status"] = ProcessStatus.FAILED
-                        print(
-                            "[@ollama] API server process %s terminated with error code %s."
-                            % (pid, returncode)
+
+                    process_info["status"] = ProcessStatus.SUCCESSFUL
+                    self._update_server_status("Stopped")
+                    self._log_event(
+                        "success", f"API server stopped ({server_shutdown_cause})"
+                    )
+                    if self.debug:
+                        print(f"[@ollama] API server terminated successfully.")
+                except Exception as e:
+                    process_info["status"] = ProcessStatus.FAILED
+                    server_shutdown_cause = "failed"
+                    self._update_server_status("Failed to stop")
+                    self._log_event("error", f"API server shutdown error: {str(e)}")
+                    print(f"[@ollama] Warning: Error terminating API server: {e}")
+
+        # Record total shutdown time and performance metrics
+        total_shutdown_time = time.time() - shutdown_start_time
+        server_shutdown_time = time.time() - server_shutdown_start
+
+        # Update performance metrics
+        self._update_performance("server_shutdown_time", server_shutdown_time)
+        self._update_performance("total_shutdown_time", total_shutdown_time)
+        self._update_performance("shutdown_cause", server_shutdown_cause)
+
+        # Log individual model shutdown times
+        for model_name, results in model_shutdown_results.items():
+            self._update_performance(
+                f"{model_name}_shutdown_time", results["shutdown_time"]
+            )
+            self._update_performance(
+                f"{model_name}_shutdown_cause", results["shutdown_cause"]
+            )
+
+        self._log_event(
+            "success", f"Ollama shutdown completed in {total_shutdown_time:.1f}s"
+        )
+        print("[@ollama] All models stopped.")
+
+        # Show performance summary
+        if self.debug:
+            if hasattr(self, "stats") and self.stats:
+                print("[@ollama] Performance summary:")
+                for operation, stats in self.stats.items():
+                    runtime = stats.get("process_runtime", 0)
+                    if runtime > 1:  # Only show operations that took meaningful time
+                        print(f"[@ollama]   {operation}: {runtime:.1f}s")
+
+    def _update_model_cache(self, model_name):
+        """
+        Update the remote cache with model files if needed.
+        """
+        try:
+            manifest = self._fetch_manifest(model_name)
+            if not manifest:
+                if self.debug:
+                    print(
+                        f"[@ollama {model_name}] No manifest available for cache update."
+                    )
+                return
+
+            from metaflow import S3
+
+            cache_up_to_date = True
+            key_paths = [
+                (
+                    self.storage_info[model_name]["manifest_remote"],
+                    self.storage_info[model_name]["manifest_local"],
+                )
+            ]
+
+            with S3() as s3:
+                # Check if blobs need updating
+                s3objs = s3.list_paths(
+                    [self.storage_info[model_name]["blob_remote_root"]]
+                )
+                for layer in manifest["layers"]:
+                    expected_blob_sha = layer["digest"]
+                    if expected_blob_sha not in s3objs:
+                        cache_up_to_date = False
+                        break
+
+                if not cache_up_to_date:
+                    blob_count = len(manifest.get("layers", []))
+                    print(
+                        f"[@ollama {model_name}] Uploading {blob_count} files to cache..."
+                    )
+
+                    # Add blob paths to upload
+                    for layer in manifest["layers"]:
+                        blob_filename = layer["digest"].replace(":", "-")
+                        key_paths.append(
+                            (
+                                os.path.join(
+                                    self.storage_info[model_name]["blob_remote_root"],
+                                    blob_filename,
+                                ),
+                                os.path.join(
+                                    self.storage_info[model_name]["blob_local_root"],
+                                    blob_filename,
+                                ),
+                            )
                         )
-                    else:
+
+                    s3.put_files(key_paths)
+                    print(f"[@ollama {model_name}] Cache updated.")
+                else:
+                    if self.debug:
+                        print(f"[@ollama {model_name}] Cache is up to date.")
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {model_name}] Error updating cache: {e}")
+
+    def get_ollama_storage_root(self, backend):
+        """
+        Return the path to the root of the datastore.
+        """
+        if backend.TYPE == "s3":
+            from metaflow.metaflow_config import DATASTORE_SYSROOT_S3
+
+            self.local_datastore = False
+            return os.path.join(DATASTORE_SYSROOT_S3, OLLAMA_SUFFIX)
+        elif backend.TYPE == "azure":
+            from metaflow.metaflow_config import DATASTORE_SYSROOT_AZURE
+
+            self.local_datastore = False
+            return os.path.join(DATASTORE_SYSROOT_AZURE, OLLAMA_SUFFIX)
+        elif backend.TYPE == "gs":
+            from metaflow.metaflow_config import DATASTORE_SYSROOT_GS
+
+            self.local_datastore = False
+            return os.path.join(DATASTORE_SYSROOT_GS, OLLAMA_SUFFIX)
+        else:
+            self.local_datastore = True
+            return None
+
+    def _attempt_ollama_restart(self):
+        """Attempt to restart Ollama when circuit breaker suggests it"""
+        try:
+            print("[@ollama] Attempting Ollama restart due to circuit breaker...")
+
+            # Stop existing server processes
+            server_stopped = False
+            for pid, process_info in list(self.processes.items()):
+                if process_info["properties"].get("type") == "api-server":
+                    process = process_info["p"]
+                    try:
+                        process.terminate()
+                        process.wait(timeout=10)
+                        if process.poll() is None:
+                            process.kill()
+                            process.wait()
                         process_info["status"] = ProcessStatus.SUCCESSFUL
+                        server_stopped = True
                         if self.debug:
                             print(
-                                "[@ollama] API server process %s terminated successfully."
-                                % pid
+                                f"[@ollama] Stopped server process {pid} during restart"
                             )
-                except Exception as e:
-                    process_info["status"] = ProcessStatus.FAILED
-                    print(
-                        "[@ollama] Warning: Error while terminating API server process %s: %s"
-                        % (pid, e)
-                    )
+                    except Exception as e:
+                        if self.debug:
+                            print(
+                                f"[@ollama] Error stopping server {pid} during restart: {e}"
+                            )
+
+            if not server_stopped:
+                if self.debug:
+                    print("[@ollama] No server process found to stop during restart")
+
+            # Small delay to ensure cleanup
+            time.sleep(2)
+
+            # Restart server
+            self._launch_server()
+
+            # Verify health with multiple attempts
+            health_attempts = 3
+            for attempt in range(health_attempts):
+                if self._verify_server_health():
+                    print("[@ollama] Restart successful")
+                    return True
+                else:
+                    if attempt < health_attempts - 1:
+                        if self.debug:
+                            print(
+                                f"[@ollama] Health check failed, attempt {attempt + 1}/{health_attempts}"
+                            )
+                        time.sleep(5)
+
+            print(
+                "[@ollama] Restart failed - server not healthy after multiple attempts"
+            )
+            return False
+
+        except Exception as e:
+            print(f"[@ollama] Restart failed: {e}")
+            return False
+
+    def _verify_server_health(self):
+        """Quick health check for server availability"""
+        try:
+            response = requests.get(
+                f"{self.ollama_url}/api/tags",
+                timeout=self.timeouts.get("health_check", 5),
+            )
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def _collect_model_metadata(self, m):
+        """
+        Collect model metadata including size and blob count from manifest and API
+        """
+        metadata = {"size_bytes": None, "size_formatted": "Unknown", "blob_count": 0}
+
+        try:
+            # First try to get info from manifest (works for cached models)
+            manifest = self._fetch_manifest(m)
+            if manifest and "layers" in manifest:
+                metadata["blob_count"] = len(manifest["layers"])
+
+                # Calculate total size from manifest layers
+                total_size = 0
+                for layer in manifest["layers"]:
+                    if "size" in layer:
+                        total_size += layer["size"]
+
+                if total_size > 0:
+                    metadata["size_bytes"] = total_size
+                    metadata["size_formatted"] = self._format_bytes(total_size)
+
+            # Try to get more detailed info from Ollama API if available
+            try:
+                response = requests.post(
+                    f"{self.ollama_url}/api/show", json={"model": m}, timeout=10
+                )
+                if response.status_code == 200:
+                    model_info = response.json()
+
+                    # Extract size if available in the response
+                    if (
+                        "details" in model_info
+                        and "parameter_size" in model_info["details"]
+                    ):
+                        # Sometimes the API returns parameter size info
+                        param_size = model_info["details"]["parameter_size"]
+                        if self.debug:
+                            print(f"[@ollama {m}] Parameter size: {param_size}")
+
+                    # If we get model_info but didn't have manifest info, try to get layer count
+                    if metadata["blob_count"] == 0 and "details" in model_info:
+                        details = model_info["details"]
+                        if "families" in details or "family" in details:
+                            # API response structure varies, estimate blob count
+                            metadata["blob_count"] = "API"
+
+            except Exception as api_e:
+                if self.debug:
+                    print(f"[@ollama {m}] Could not get API metadata: {api_e}")
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@ollama {m}] Error collecting model metadata: {e}")
+
+        return metadata
+
+    def _format_bytes(self, bytes_count):
+        """Format bytes into human-readable string"""
+        if bytes_count is None:
+            return "Unknown"
+
+        for unit in ["B", "KB", "MB", "GB", "TB"]:
+            if bytes_count < 1024.0:
+                if unit == "B":
+                    return f"{int(bytes_count)} {unit}"
+                else:
+                    return f"{bytes_count:.1f} {unit}"
+            bytes_count /= 1024.0
+        return f"{bytes_count:.1f} PB"